feat(integrations): per-tool agent approvals for remote mcp servers#2855
feat(integrations): per-tool agent approvals for remote mcp servers#2855jordan-umusu wants to merge 1 commit into
Conversation
|
✅ No security or compliance issues detected. Reviewed everything up to d16921f. Security Overview
Detected Code Changes
|
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d16921f158
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
daryllimyt
left a comment
daryllimyt
left a comment
There was a problem hiding this comment.
Found two functional issues that should be addressed before this patch is considered correct.
| if not server or not remote_tool: | ||
| return None | ||
| tool_name = action_name | ||
| return action_name_to_mcp_tool_name(tool_name) |
There was a problem hiding this comment.
This path is lossy for user MCP names that contain dots: the approval key uses mcp.{server}.{tool}, then action_name_to_mcp_tool_name() converts every dot to __, so mcp.Jira.issue.get becomes mcp__Jira__issue__get rather than preserving the remote segment as issue.get.
If dotted remote tool names are intentionally unsupported because provider tool registration rejects them, we should enforce that during discovery/verification and surface a clear validation error. Otherwise, approved execution needs a non-lossy mapping from provider-safe tool names back to the original mcp__{server}__{remote_tool} name.
| if policy_update.requires_approval is not None: | ||
| update_fields["requires_approval"] = policy_update.requires_approval |
There was a problem hiding this comment.
For organizations without agent_addons, this endpoint still persists requires_approval=true for custom MCP tools; later agent compilation sees the stored policy and _check_tool_approval_entitlement() rejects the run, so a user with only integration:update can save a setting that makes the MCP-backed agent unusable. Reject approval-enabling updates without the entitlement while still allowing false/enabled changes.
|
I think this is worth adding as a real Temporal integration test. The PR already has good unit coverage for policy merge/update, build-time approval mapping, runtime hook denial, and name conversion, but it does not currently exercise the exact new workflow branch:
Here is a compact version that should fit in # Add to existing imports:
# from tracecat_ee.agent.activities import ExecuteRemoteMCPToolArgs
# from tracecat.agent.tokens import UserMCPServerClaim
@pytest.mark.anyio
@pytest.mark.integration
async def test_agent_workflow_routes_approved_user_mcp_tool_to_remote_activity_and_replays(
svc_role: Role,
temporal_client: Client,
agent_worker_factory,
mock_session_id: uuid.UUID,
) -> None:
queue = f"test-agent-queue-{mock_session_id}"
integration_id = uuid.uuid4()
approval_request_recorded = asyncio.Event()
agent_inputs: list[AgentExecutorInput] = []
remote_activity_inputs: list[ExecuteRemoteMCPToolArgs] = []
registry_activity_inputs: list[RunActionInput] = []
reconcile_inputs: list[ReconcileToolResultsInput] = []
@activity.defn(name="load_session_activity")
async def mock_load_session_activity(
input: LoadSessionInput,
) -> LoadSessionResult:
assert input.session_id == mock_session_id
return LoadSessionResult(
found=bool(agent_inputs),
sdk_session_id="sdk-session" if agent_inputs else None,
sdk_session_data=None,
is_fork=False,
)
@activity.defn(name="build_agent_tool_definitions")
async def mock_build_tool_definitions(
args: BuildAgentToolDefsArgs,
) -> BuildAgentToolDefsResult:
tool_result = BuildToolDefsResult(
tool_definitions={
"mcp__Jira__getIssue": MCPToolDefinition(
name="mcp__Jira__getIssue",
description="Get an issue",
parameters_json_schema={
"type": "object",
"properties": {"issue_key": {"type": "string"}},
"required": ["issue_key"],
"additionalProperties": False,
},
)
},
registry_lock=RegistryLock(origins={}, actions={}),
user_mcp_claims=[UserMCPServerClaim(name="Jira", id=integration_id)],
tool_approvals={"mcp.Jira.getIssue": True},
)
return BuildAgentToolDefsResult(
scopes={scope.scope: tool_result for scope in args.scopes}
)
@activity.defn(name="run_agent_activity")
async def mock_run_agent_activity(
input: AgentExecutorInput,
) -> AgentExecutorResult:
agent_inputs.append(input)
if len(agent_inputs) == 1:
assert input.is_approval_continuation is False
return AgentExecutorResult(
success=True,
approval_requested=True,
approval_items=[
ToolCallContent(
id="call-user-mcp",
name="mcp__tracecat-registry__mcp__Jira__getIssue",
input={"issue_key": "ISSUE-1"},
)
],
)
assert input.is_approval_continuation is True
return AgentExecutorResult(
success=True,
approval_requested=False,
output={"continued": True},
)
@activity.defn(name="record_approval_requests")
async def mock_record_approval_requests(
input: PersistApprovalsActivityInputs,
) -> None:
assert [approval.tool_call_id for approval in input.approvals] == [
"call-user-mcp"
]
approval_request_recorded.set()
@activity.defn(name="apply_approval_decisions")
async def mock_apply_approval_decisions(
input: ApplyApprovalResultsActivityInputs,
) -> None:
assert [decision.tool_call_id for decision in input.decisions] == [
"call-user-mcp"
]
assert input.decisions[0].approved is True
@activity.defn(name="execute_remote_mcp_tool")
async def mock_execute_remote_mcp_tool(args: ExecuteRemoteMCPToolArgs) -> str:
remote_activity_inputs.append(args)
assert args.tool_name == "mcp__Jira__getIssue"
assert args.args == {"issue_key": "ISSUE-1"}
assert args.mcp_auth_token
return '{"ok": true}'
@activity.defn(name="execute_action_activity")
async def mock_execute_action_activity(
input: RunActionInput,
role: Role,
) -> InlineObject[dict[str, str]]:
del role
registry_activity_inputs.append(input)
return InlineObject(data={"unexpected": "registry path"})
@activity.defn(name="reconcile_tool_results_activity")
async def mock_reconcile_tool_results_activity(
input: ReconcileToolResultsInput,
) -> ReconcileToolResultsResult:
reconcile_inputs.append(input)
return ReconcileToolResultsResult(
results=[
ToolExecutionResult(
tool_call_id=pending.tool_call_id,
tool_name=pending.tool_name,
result=pending.raw_result,
is_error=pending.is_error,
)
for pending in input.pending_results
]
)
workflow_args = AgentWorkflowArgs(
role=svc_role,
agent_args=RunAgentArgs(
session_id=mock_session_id,
user_prompt="Validate remote MCP approval continuation",
config=AgentConfig(
model_name="claude-3-5-sonnet-20241022",
model_provider="anthropic",
actions=[],
tool_approvals={"mcp.Jira.getIssue": True},
),
),
entity_type=AgentSessionEntity.WORKFLOW,
entity_id=uuid.uuid4(),
)
activities = [
create_mock_create_session_activity(),
mock_load_session_activity,
create_mock_load_session_messages_activity(),
mock_build_tool_definitions,
mock_run_agent_activity,
mock_record_approval_requests,
mock_apply_approval_decisions,
mock_execute_remote_mcp_tool,
mock_execute_action_activity,
mock_reconcile_tool_results_activity,
]
async with agent_worker_factory(
temporal_client,
task_queue=queue,
custom_activities=activities,
):
wf_handle = await temporal_client.start_workflow(
DurableAgentWorkflow.run,
workflow_args,
id=AgentWorkflowID(mock_session_id),
task_queue=queue,
retry_policy=RETRY_POLICIES["workflow:fail_fast"],
execution_timeout=timedelta(seconds=30),
)
await asyncio.wait_for(approval_request_recorded.wait(), timeout=10)
suspended_history = await fetch_history_after_completed_workflow_task(wf_handle)
await replay_durable_agent_workflow_history(temporal_client, suspended_history)
await wf_handle.execute_update(
DurableAgentWorkflow.set_approvals,
WorkflowApprovalSubmission(
approvals={"call-user-mcp": True},
approved_by=svc_role.user_id,
),
)
result = await wf_handle.result()
completed_history = await wf_handle.fetch_history()
await replay_durable_agent_workflow_history(temporal_client, completed_history)
assert result.output == {"continued": True}
assert len(remote_activity_inputs) == 1
assert remote_activity_inputs[0].tool_name == "mcp__Jira__getIssue"
assert remote_activity_inputs[0].args == {"issue_key": "ISSUE-1"}
assert registry_activity_inputs == []
assert len(reconcile_inputs) == 1
assert len(reconcile_inputs[0].pending_results) == 1
assert reconcile_inputs[0].pending_results[0].raw_result == '{"ok": true}'
assert [input.is_approval_continuation for input in agent_inputs] == [False, True]I validated this shape locally with a standalone Temporal time-skipping run. The important part is not the mock remote MCP call itself; it is proving the workflow schedules the new |
| """Whether the tool was present in the latest successful discovery.""" | ||
|
|
||
| @classmethod | ||
| def validate_stored( |
There was a problem hiding this comment.
nit: why do we need this? is it just to correlate the validation error to the ID? seems redundant otherwise
|
Small clarification on the Temporal test suggestion above: that test covers replay of histories generated by the updated workflow code. It does not prove deploy compatibility for an execution that already passed approval under the pre-PR workflow and recorded If we want to explicitly cover that risk, we should add a separate replay fixture generated from
and replay that fixture against this branch. That is the test that would catch a post-approval command mismatch between old |
2 participants
Core change
mcpintegration.toolscolumn schema extended — each stored tool entry now carriesper-tool policy:
enabled(defaulttrue),requires_approval(defaultfalse), andstatus(available|missing). No migration needed: old[{name, description}]rowsself-heal on read via
MCPToolSummarydefaults, and the next successful verificationrewrites them in the new shape.
stored rows own user policy. Tools that disappear from the server are kept as
status: "missing"(policy preserved) instead of silently dropped, and a failedverification no longer clears the stored tool snapshot.
Summary by cubic
Adds per-tool enable/approval policies for remote MCP integrations. Agents now enforce these policies and route approved user MCP tool calls through the trusted MCP router; the UI adds toggles and clearer tool counts.
/workspaces/{workspace_id}/mcp-integrations/{mcp_integration_id}/toolsto update per-tool policies; addsMCPToolPolicyUpdate,MCPToolPolicyUpdateRequest;MCPToolSummarynow includesenabled,requires_approval, andstatus.tool_approvals, checks entitlements, blocks unapproved calls, and executes approved user MCP tools via newexecute_remote_mcp_toolactivity (token-scoped trusted router).useUpdateMcpIntegrationToolPolicieshook and generated client types.Written for commit d16921f. Summary will update on new commits.