diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 84820bba408..6dd1799a49d 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -40,6 +40,7 @@ from agent_framework._tools import ( _ALREADY_APPROVED_APPROVAL_REQUEST_GROUPS_KEY, # type: ignore _collect_approval_responses, # type: ignore + _get_tool_map, # type: ignore _replace_approval_contents_with_results, # type: ignore _TOOL_APPROVAL_STATE_KEY, # type: ignore _try_execute_function_calls, # type: ignore @@ -1308,8 +1309,21 @@ async def _resolve_approval_responses( approved_function_results: list[Any] = [] - # Execute approved tool calls - if approved_responses and tools: + # Partition approved responses into static (execute now) and deferred (execute during run) + tool_map = _get_tool_map(tools) if tools else {} + static_approved = [] + deferred_approval_ids = set() + + for approval in approved_responses: + tool_name = approval.function_call.name if approval.function_call else None + if tool_name not in tool_map: + # Provider-injected tool — defer to harness for execution during run + deferred_approval_ids.add(approval.id or "") + else: + static_approved.append(approval) + + # Execute only statically-available approved tool calls + if static_approved and tools: client = getattr(agent, "client", None) config = normalize_function_invocation_configuration(getattr(client, "function_invocation_configuration", None)) middleware_pipeline = FunctionMiddlewarePipeline( @@ -1322,7 +1336,7 @@ async def _resolve_approval_responses( results, _ = await _try_execute_function_calls( custom_args=tool_kwargs, attempt_idx=0, - function_calls=approved_responses, + function_calls=static_approved, tools=tools, middleware_pipeline=middleware_pipeline, config=config, @@ -1332,9 +1346,10 @@ async def _resolve_approval_responses( logger.exception("Failed to execute approved tool calls; injecting error results: %s", e) approved_function_results = [] - # Build results for approved responses (used for TOOL_CALL_RESULT event emission) + # Build results for executed (static) approved responses (used for TOOL_CALL_RESULT event emission) + # Deferred provider-injected approvals are left in messages for ToolApprovalMiddleware to process. approved_results: list[Content] = [] - for idx, approval in enumerate(approved_responses): + for idx, approval in enumerate(static_approved): if ( idx < len(approved_function_results) and getattr(approved_function_results[idx], "type", None) == "function_result" diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index ac5ba412565..a205dbf9f1e 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -5160,3 +5160,208 @@ def factory(thread_id: str) -> Any: runner.clear_thread_workflow("thread-1") assert runner._resolve_workflow("thread-1", "tenant-b") is not workflow_b + + +async def test_endpoint_agent_approval_defers_provider_injected_tools() -> None: + """Approving provider-injected tools should not produce rejection errors. + + When approval responses include tools not in the static tool map, + they should be deferred (not executed by transport) rather than + producing "Tool call invocation was rejected" errors. + + This prevents the scenario where approval mode loops with tool invocation failures. + Fixes issue #7043. + """ + executed_tools: list[str] = [] + + def static_tool() -> str: + executed_tools.append("static") + return "static result" + + # Create approval responses: one for static tool, one for provider tool (not in static list) + static_approval = Content.from_function_approval_response( + id="static_id", + approved=True, + function_call=Content.from_function_call( + call_id="call_static", + name="static_tool", + arguments="{}", + ), + ) + provider_approval = Content.from_function_approval_response( + id="provider_id", + approved=True, + function_call=Content.from_function_call( + call_id="call_provider", + name="provider_file_access", # Not in agent's tool list + arguments="{}", + ), + ) + + # User message containing both approval responses + message = Message( + role="user", + contents=[static_approval, provider_approval], + ) + + # Stub agent that returns the approval responses embedded in a message + agent = StubAgent( + updates=[ + AgentResponseUpdate(contents=[static_approval, provider_approval], role="user"), + ] + ) + + app = FastAPI() + wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") + client = TestClient(app) + + # Submit approval responses (static and provider-injected) + response = client.post( + "/approval", + json={ + "runId": "run-approval", + "threadId": "thread-approval", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "function_approval_response", + "id": "static_id", + "approved": True, + "function_call": { + "call_id": "call_static", + "name": "static_tool", + "arguments": "{}", + }, + }, + { + "type": "function_approval_response", + "id": "provider_id", + "approved": True, + "function_call": { + "call_id": "call_provider", + "name": "provider_file_access", + "arguments": "{}", + }, + }, + ], + } + ], + }, + ) + + assert response.status_code == 200 + events = _decode_sse_events(response) + response_text = json.dumps(events) + + # Verify no rejection error for provider tool + assert "Tool call invocation was rejected" not in response_text, ( + "Deferred provider tool should not produce rejection error" + ) + + +async def test_endpoint_agent_approval_deferred_provider_tool_executes(streaming_chat_client_stub) -> None: + """A provider-injected tool approved via AG-UI executes in-run instead of being rejected. + + Regression for #7043. A tool registered by a context provider during ``before_run`` is + absent from the transport's static tool map, so ``_resolve_approval_responses`` must defer + it (not execute or reject it) and leave it for the in-run ``ToolApprovalMiddleware`` to run. + This drives the full pause -> approve -> resume flow with a real provider-injected tool and + asserts the approved side effect actually happens without any rejection/failure result. + + Note: the deferred tool executes inside ``agent.run`` and its ``function_result`` is consumed + there, so (unlike a statically executed approval) no ``TOOL_CALL_RESULT`` event is streamed + for it. Surfacing that result to the client is tracked separately in #7241. + """ + side_effects: list[str] = [] + state = {"phase": "pause"} + + def provider_write() -> str: + side_effects.append("wrote") + return "wrote to disk" + + provider_tool = FunctionTool( + name="provider_write", + description="Write to disk (provider-injected)", + func=provider_write, + approval_mode="always_require", + ) + + class ToolInjectingProvider(ContextProvider): + """Registers a tool during before_run, mimicking FileAccessProvider/CodeInterpreterProvider.""" + + async def before_run(self, *, agent, session, context, state) -> None: # type: ignore[override] # pyrefly: ignore # ty: ignore + del agent, session, state + context.extend_tools(self.source_id, [provider_tool]) + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + if state["phase"] == "pause": + yield ChatResponseUpdate( + contents=[Content.from_function_call(call_id="call_provider", name="provider_write", arguments="{}")], + role="assistant", + ) + return + yield ChatResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant") + + # provider_write is intentionally NOT in the static tools list -- it is only injected via before_run. + agent = Agent( + name="test_agent", + instructions="Test", + client=streaming_chat_client_stub(stream_fn), + tools=[], + middleware=[ToolApprovalMiddleware()], + context_providers=[ToolInjectingProvider(source_id="tool_injector")], + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + AgentFrameworkAgent(agent=agent, require_confirmation=False), + path="/approval", + ) + client = TestClient(app) + + # Pause: the harness surfaces the provider-injected tool for approval, nothing executes yet. + pause_response = client.post( + "/approval", + json={ + "runId": "run-pause", + "threadId": "thread-provider", + "messages": [{"role": "user", "content": "Write something"}], + }, + ) + assert pause_response.status_code == 200 + pause_finished = [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_FINISHED"] + assert [interrupt["id"] for interrupt in _run_finished_interrupts(pause_finished[-1])] == ["call_provider"] + assert side_effects == [] + + # Resume with approval: the deferred provider tool runs during agent.run. + state["phase"] = "resume" + resume_response = client.post( + "/approval", + json={ + "runId": "run-resume", + "threadId": "thread-provider", + "messages": [], + "resume": [{"interruptId": "call_provider", "status": "resolved", "payload": {"accepted": True}}], + }, + ) + assert resume_response.status_code == 200 + resume_events = _decode_sse_events(resume_response) + resume_text = json.dumps(resume_events) + + # The approved provider tool actually executed -- its side effect fired. + assert side_effects == ["wrote"] + # And it was neither rejected nor reported as a transport failure (the #7043 bug). + assert "Tool call invocation was rejected" not in resume_text + assert "Tool call invocation failed" not in resume_text + assert not [event for event in resume_events if event.get("type") == "RUN_ERROR"] + + +