Python: preserve approval context across resume paths

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.

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

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

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

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] = []
 

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,

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"