diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index 25233eaf56d..6e5c9d93a74 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -41,6 +41,10 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. `Agent.run` path, including Agent middleware, provider `before_run` preparation, and function middleware. `_approval_execution.py` tracks execution through existing function/chat middleware interfaces; it does not prepare providers or invoke the tools itself. `MiddlewareFailure` remains fatal. +- Programmatic `AgentFrameworkAgent.run(...)` and `run_agent_stream(...)` calls accept + `function_invocation_kwargs` and preserve that tool-runtime context across approval resume. Adapter-local execution + forwards only that mapping plus `options["additional_function_arguments"]` (with options taking precedence), not + unrelated AG-UI run kwargs. - Queued approvals follow the wrapped Agent's scheduling. With `ToolApprovalMiddleware`, collected decisions remain pending until that middleware releases the batch; AG-UI does not execute a tool ahead of it. Collected server-side grants do not require the client to submit the same approval again. diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent.py index 9d127f2bbff..ab179da165c 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent.py @@ -3,7 +3,7 @@ """AgentFrameworkAgent wrapper for AG-UI protocol.""" import warnings -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Mapping from typing import Any, cast from ag_ui.core import BaseEvent @@ -175,11 +175,14 @@ def snapshot_store(self) -> AGUIThreadSnapshotStore | None: async def run( self, input_data: dict[str, Any], + *, + function_invocation_kwargs: Mapping[str, Any] | None = None, ) -> AsyncGenerator[BaseEvent, None]: """Run the wrapped agent and yield AG-UI events. Args: input_data: The AG-UI run input containing messages, state, etc. + function_invocation_kwargs: Keyword arguments forwarded only to tool invocation. Yields: AG-UI events @@ -190,5 +193,6 @@ async def run( self.agent, self.config, approval_state_store=self._approval_state_store, + function_invocation_kwargs=function_invocation_kwargs, ): yield event 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 825a60ec621..78dea958d5e 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 @@ -1938,7 +1938,17 @@ async def forward_hosted_decision(approval: Content = approval) -> list[Content] raise TypeError("Local approval execution with context providers requires a core Agent execution path.") client = getattr(agent, "client", None) config = normalize_function_invocation_configuration(getattr(client, "function_invocation_configuration", None)) - tool_kwargs = {k: v for k, v in run_kwargs.items() if k != "options"} + raw_function_invocation_kwargs = run_kwargs.get("function_invocation_kwargs") + tool_kwargs = ( + dict(cast(Mapping[str, Any], raw_function_invocation_kwargs)) + if raw_function_invocation_kwargs is not None + else {} + ) + raw_options = run_kwargs.get("options") + if isinstance(raw_options, Mapping) and ( + additional_function_arguments := raw_options.get("additional_function_arguments") + ): + tool_kwargs.update(cast(Mapping[str, Any], additional_function_arguments)) for approval in static_approved: function_call = approval.function_call call_id = (function_call.call_id if function_call else None) or approval.id or "" @@ -2668,6 +2678,8 @@ async def run_agent_stream( agent: SupportsAgentRun, config: AgentConfig, approval_state_store: InMemoryAGUIApprovalStateStore | None = None, + *, + function_invocation_kwargs: Mapping[str, Any] | None = None, ) -> AsyncGenerator[BaseEvent]: """Run agent and yield AG-UI events. @@ -2680,6 +2692,7 @@ async def run_agent_stream( config: Agent configuration approval_state_store: Optional server-side Approval State store used to preserve approval-only middleware state across AG-UI requests. + function_invocation_kwargs: Keyword arguments forwarded only to tool invocation. Yields: AG-UI events @@ -2692,6 +2705,7 @@ async def run_agent_stream( config, state_store, authorized_executions=authorized_executions, + function_invocation_kwargs=function_invocation_kwargs, ) try: async for event in stream: @@ -2711,6 +2725,7 @@ async def _run_agent_stream( approval_state_store: InMemoryAGUIApprovalStateStore, *, authorized_executions: dict[ApprovalOccurrenceIdentity, AuthorizedExecution], + function_invocation_kwargs: Mapping[str, Any] | None, ) -> AsyncGenerator[BaseEvent]: # Parse IDs supplied_thread_id = input_data.get("thread_id") or input_data.get("threadId") @@ -3105,6 +3120,8 @@ async def _run_agent_stream( # Build run kwargs (Feature #6: Azure store flag when metadata present) run_kwargs: dict[str, Any] = {"session": session} + if function_invocation_kwargs is not None: + run_kwargs["function_invocation_kwargs"] = dict(function_invocation_kwargs) if tools: run_kwargs["tools"] = tools # Hand the forwarded AG-UI context to the A2UI runner PER REQUEST (not just at diff --git a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py index b32cea9af8c..2a01c03a4b4 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py +++ b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py @@ -8,7 +8,16 @@ import pytest from ag_ui.core import RunErrorEvent, ToolCallResultEvent -from agent_framework import Agent, ChatOptions, ChatResponseUpdate, Content, Message +from agent_framework import ( + Agent, + ChatOptions, + ChatResponseUpdate, + Content, + ContextProvider, + FunctionInvocationContext, + Message, + tool, +) from pydantic import BaseModel from agent_framework_ag_ui._approval_lifecycle import ApprovalExecutionOwner @@ -23,6 +32,82 @@ def _approval_request_id(events: list[Any]) -> str: return str(event.value["id"]) +async def test_approval_resume_with_context_provider_preserves_function_invocation_kwargs( + streaming_chat_client_stub: Any, +) -> None: + """Context-provider preparation and runtime tool kwargs survive an AG-UI resume.""" + from agent_framework.ag_ui import AgentFrameworkAgent + + observed_kwargs: list[dict[str, Any]] = [] + phase = "pause" + + class NoOpContextProvider(ContextProvider): + pass + + @tool(name="sensitive_action", approval_mode="always_require") + def sensitive_action(context: FunctionInvocationContext) -> str: + observed_kwargs.append(dict(context.kwargs)) + return "executed" + + async def stream_fn( + messages: MutableSequence[Message], + options: ChatOptions, + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del messages, options, kwargs + if phase == "pause": + yield ChatResponseUpdate( + contents=[ + Content.from_function_call( + name="sensitive_action", + call_id="call-sensitive", + arguments="{}", + ) + ] + ) + return + yield ChatResponseUpdate(contents=[Content.from_text(text="Done")]) + + wrapper = AgentFrameworkAgent( + agent=Agent( + client=streaming_chat_client_stub(stream_fn), + name="test_agent", + instructions="Test", + tools=[sensitive_action], + context_providers=[NoOpContextProvider("context")], + ) + ) + runtime_kwargs = {"user_id": "user-123"} + thread_id = "context-provider-runtime-kwargs" + pause_events = [ + event + async for event in wrapper.run( + {"thread_id": thread_id, "messages": [{"role": "user", "content": "do it"}]}, + function_invocation_kwargs=runtime_kwargs, + ) + ] + approval_id = _approval_request_id(pause_events) + + phase = "resume" + resume_events = [ + event + async for event in wrapper.run( + { + "thread_id": thread_id, + "messages": [], + "resume": [{"interruptId": approval_id, "status": "resolved", "payload": {"accepted": True}}], + }, + function_invocation_kwargs=runtime_kwargs, + ) + ] + + assert observed_kwargs[0]["user_id"] == runtime_kwargs["user_id"] + assert set(observed_kwargs[0]) == {"session", "user_id"} + assert [ + (event.tool_call_id, event.content) for event in resume_events if isinstance(event, ToolCallResultEvent) + ] == [("call-sensitive", "executed")] + + async def test_agent_initialization_basic(streaming_chat_client_stub): """Test basic agent initialization without state schema.""" from agent_framework.ag_ui import AgentFrameworkAgent diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py index a43dde00d19..133562061c9 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py @@ -8,7 +8,7 @@ from typing import Any from ag_ui.core import ToolCallResultEvent -from agent_framework import AgentResponseUpdate, Content, FunctionTool +from agent_framework import AgentResponseUpdate, Content, FunctionInvocationContext, FunctionTool from agent_framework.exceptions import UserInputRequiredException from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] @@ -88,7 +88,11 @@ async def _run_resume( return events -async def _run_custom_approval(tool: FunctionTool) -> tuple[list[Any], StubAgent]: +async def _run_custom_approval( + tool: FunctionTool, + *, + function_invocation_kwargs: dict[str, Any] | None = None, +) -> tuple[list[Any], StubAgent]: agent = StubAgent( updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], default_options={"tools": [tool]}, @@ -114,6 +118,7 @@ async def _run_custom_approval(tool: FunctionTool) -> tuple[list[Any], StubAgent agent, AgentConfig(), approval_state_store=store, + function_invocation_kwargs=function_invocation_kwargs, ) ] return events, agent @@ -134,6 +139,25 @@ async def test_approved_call_emits_one_live_result_under_original_identity() -> assert [(event.tool_call_id, event.content) for event in results] == [("call-weather", "Sunny in Seattle")] +async def test_approved_call_receives_only_function_invocation_kwargs() -> None: + """Static approval execution forwards only the documented tool-runtime kwargs.""" + observed_kwargs: list[dict[str, Any]] = [] + + def inspect_context(context: FunctionInvocationContext) -> str: + observed_kwargs.append(dict(context.kwargs)) + return "inspected" + + events, _ = await _run_custom_approval( + FunctionTool(name="inspect_context", description="Inspect context", func=inspect_context), + function_invocation_kwargs={"user_id": "user-123"}, + ) + + results = [event for event in events if isinstance(event, ToolCallResultEvent)] + assert observed_kwargs[0]["user_id"] == "user-123" + assert set(observed_kwargs[0]) == {"session", "user_id"} + assert [(event.tool_call_id, event.content) for event in results] == [("call-custom", "inspected")] + + async def test_rejected_call_does_not_execute_or_emit_live_result() -> None: executions: list[str] = [] diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index c4f5da1c0de..74dfbeb07b7 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -1407,6 +1407,9 @@ async def _prepare_run_context( active_session = session if active_session is None and self.context_providers: active_session = AgentSession() + from ._tools import _RUN_LOCAL_MIDDLEWARE_SESSION_ATTR # pyright: ignore[reportPrivateUsage] + + setattr(active_session, _RUN_LOCAL_MIDDLEWARE_SESSION_ATTR, True) per_service_call_history_providers = self._resolve_per_service_call_history_providers( session=active_session, diff --git a/python/packages/core/tests/core/test_harness_tool_approval.py b/python/packages/core/tests/core/test_harness_tool_approval.py index 95f65e5a467..bda3d23b2ca 100644 --- a/python/packages/core/tests/core/test_harness_tool_approval.py +++ b/python/packages/core/tests/core/test_harness_tool_approval.py @@ -24,6 +24,7 @@ ChatResponse, ChatResponseUpdate, Content, + ContextProvider, FileHistoryProvider, FunctionInvocationContext, FunctionMiddleware, @@ -134,6 +135,57 @@ def safe_tool() -> str: ] +async def test_sessionless_approval_resume_with_context_provider_uses_message_authority( + chat_client_base: MockBaseChatClient, +) -> None: + """A framework-created context-provider session must not become approval authority.""" + calls: list[str] = [] + + class NoOpContextProvider(ContextProvider): + pass + + @tool(name="approved_tool", approval_mode="always_require") + def approved_tool() -> str: + calls.append("approved") + return "approved" + + agent = Agent( + client=chat_client_base, + tools=[approved_tool], + context_providers=[NoOpContextProvider("context")], + ) + chat_client_base.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="approved-call", + name="approved_tool", + arguments="{}", + id="approved-occurrence", + ) + ], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + first = await agent.run("request approval") + approval_request = first.user_input_requests[0] + resumed = await agent.run([ + Message(role="user", contents=["request approval"]), + *first.messages, + Message(role="user", contents=[approval_request.to_function_approval_response(True)]), + ]) + + assert calls == ["approved"] + assert [(message.role, [content.type for content in message.contents]) for message in resumed.messages] == [ + ("tool", ["function_result"]), + ("assistant", ["text"]), + ] + + class _StringApprovalValue(str, Enum): ALPHA = "alpha"