Python: preserve approval context across resume paths

docs/specs/004-python-function-calling-loop.md

@@ -191,9 +191,12 @@ Before acting on a model function-call batch, the loop classifies every actionab
 call aborts the complete batch before approval state changes or execution. Otherwise approval-required and Host-owned
 calls are returned together in model order, while session-backed executable siblings remain deferred. An incomplete
 session-backed mixed approval/Host response remains pending without executing a deferred call; a stateless incomplete
-response is rejected. Correlation is scoped to the active mixed batch so completed or abandoned historical Host calls
-remain unchanged. `ToolApprovalMiddleware` may resolve approval requests through standing or automatic policies, but
-it preserves non-approval user-input requests and does not split manual approvals away from their Host-owned siblings.
+response is rejected, including when no response is supplied, message roles vary within the model output, or a newer
+standalone request follows the incomplete batch. Stateless discovery selects the latest unresolved mixed batch, so a
+completed batch remains inert on later turns. Correlation is scoped to the active mixed batch so completed or abandoned
+historical Host calls remain unchanged.
+`ToolApprovalMiddleware` may resolve approval requests through standing or automatic policies, but it preserves
+non-approval user-input requests and does not split or reorder manual approvals relative to their Host-owned siblings.
 
 ### Reasoning-bound function-call groups
 
@@ -612,15 +615,15 @@ that manually replay messages own the equivalent rule: do not resend an approval
 | Scenario | Required invariant | Primary regression test |
 |---|---|---|
 | Fatal call mixed with pauses | Complete-batch classification raises before approval or execution, independent of call order. | `packages/core/tests/core/test_function_invocation_logic.py::test_mixed_batch_fatal_unknown_precedes_every_pause` |
-| Approval and Host-owned calls | Both pause types are returned in model order; a session-backed partial response remains pending across serialization; historical Host calls do not participate; a complete response executes the exact approved arguments once. | `test_mixed_batch_returns_approval_and_host_pause_in_model_order`, `test_mixed_batch_requires_complete_responses_before_execution`, `test_active_mixed_pause_ignores_historical_host_requests` |
+| Approval and Host-owned calls | Both pause types are returned in model order; a session-backed partial response remains pending across serialization; a stateless zero-response or partial response fails closed across message roles and cannot be hidden by a newer standalone request; completed mixed batches remain inert; historical Host calls do not participate; a complete response executes the exact approved arguments once. | `test_mixed_batch_returns_approval_and_host_pause_in_model_order`, `test_mixed_batch_requires_complete_responses_before_execution`, `test_stateless_split_mixed_batch_rejects_incomplete_replay_before_execution`, `test_stateless_mixed_batch_across_message_roles_requires_complete_responses`, `test_later_standalone_request_does_not_hide_incomplete_stateless_mixed_batch`, `test_completed_split_stateless_mixed_batch_is_inert_on_later_turn`, `test_active_mixed_pause_ignores_historical_host_requests` |
 | Safe and approval-required calls in one batch | Hidden safe calls replay only with the matching visible approval. | `packages/core/tests/core/test_harness_tool_approval.py::test_mixed_batch_hides_already_approved_request_until_approval_replay` |
 | Restored approval state | Serialized `ToolApprovalState` restores mixed-batch behavior. | `test_mixed_batch_accepts_restored_tool_approval_state` |
 | Unrelated turn before approval | Hidden calls do not execute on an unrelated turn. | `test_hidden_mixed_batch_requests_do_not_replay_on_unrelated_turn` |
 | Multiple abandoned batches | Hidden calls replay only for the matching batch. | `test_hidden_mixed_batch_requests_replay_only_for_matching_visible_approval` |
 | Queued approvals | One unresolved approval is surfaced per run without premature execution. | `test_tool_approval_middleware_queues_multiple_approval_requests`, `test_tool_approval_middleware_queues_streamed_approval_requests` |
 | Middleware state plus hidden core state | State saves do not discard hidden mixed-batch calls. | `test_tool_approval_middleware_preserves_hidden_mixed_batch_requests` |
 | Auto-approval callback | Callback receives the original function call and executes the approved set once. | `test_tool_approval_middleware_auto_approval_rule_receives_function_call` |
-| Approval policy with Host-owned sibling | Reordering the same calls has the same outcome; auto-approved and safe calls remain deferred until Host input arrives; manual approvals stay in the mixed batch; an approved Host tool returns to Host-owned handling before local execution. | `test_tool_approval_middleware_mixed_batch_is_order_independent`, `test_tool_approval_middleware_auto_approves_with_host_pause_and_cached_safe_call`, `test_tool_approval_middleware_keeps_manual_approvals_together_with_host_pause`, `test_tool_approval_middleware_policy_approval_reclassifies_host_tool` |
+| Approval policy with Host-owned sibling | Reordering the same calls has the same outcome; auto-approved and safe calls remain deferred until Host input arrives; separately streamed Host input stays immediately visible while later manual approvals remain in model order and in the same batch; an approved Host tool returns to Host-owned handling before local execution. | `test_tool_approval_middleware_mixed_batch_is_order_independent`, `test_tool_approval_middleware_auto_approves_with_host_pause_and_cached_safe_call`, `test_tool_approval_middleware_keeps_manual_approvals_together_with_host_pause`, `test_tool_approval_middleware_preserves_split_streamed_mixed_batch_order`, `test_tool_approval_middleware_streams_standalone_user_input_immediately`, `test_tool_approval_middleware_policy_approval_reclassifies_host_tool` |
 | Shared call budget | Auto-approved re-entry does not reset `max_function_calls`, and every executed approval group counts even when it pauses for input. | `test_tool_approval_middleware_auto_approved_loops_share_function_call_budget`, `test_approval_resume_user_input_counts_toward_function_call_budget` |
 | Standing tool rule | Tool-level approval applies only to later matching tools. | `test_tool_approval_middleware_always_approve_tool_rule` |
 | Forged standing rule | An unbound or substituted hosted response cannot create a standing middleware approval rule for caller-selected metadata. | `test_tool_approval_middleware_drops_forged_standing_approval`, `test_tool_approval_middleware_rebinds_hosted_standing_approval` |

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
+  merges only that mapping, the wrapped agent's default `additional_function_arguments`, and the current run options
+  (in increasing 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,22 @@ 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 {}
+        )
+        default_options = getattr(agent, "default_options", None)
+        if isinstance(default_options, Mapping) and (
+            default_function_arguments := default_options.get("additional_function_arguments")
+        ):
+            tool_kwargs.update(cast(Mapping[str, Any], default_function_arguments))
+        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 +2683,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 +2697,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 +2710,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 +2730,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 +3125,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,10 +88,18 @@ 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,
+    default_additional_function_arguments: 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]},
+        default_options={
+            "tools": [tool],
+            "additional_function_arguments": default_additional_function_arguments,
+        },
     )
     store = InMemoryAGUIApprovalStateStore()
     store.register(
@@ -114,6 +122,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 +143,30 @@ 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_merges_function_invocation_kwargs_with_agent_defaults() -> None:
+    """Static approval execution uses normal runtime/default precedence without run-kwarg leakage."""
+    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={"runtime_only": "runtime", "shared": "runtime"},
+        default_additional_function_arguments={"default_only": "default", "shared": "default"},
+    )
+
+    results = [event for event in events if isinstance(event, ToolCallResultEvent)]
+    assert observed_kwargs[0] == {
+        "default_only": "default",
+        "runtime_only": "runtime",
+        "session": observed_kwargs[0]["session"],
+        "shared": "default",
+    }
+    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

@@ -1405,8 +1405,10 @@ async def _prepare_run_context(
             self.context_providers.append(InMemoryHistoryProvider())
 
         active_session = session
+        framework_created_session = False
         if active_session is None and self.context_providers:
             active_session = AgentSession()
+            framework_created_session = True
 
         per_service_call_history_providers = self._resolve_per_service_call_history_providers(
             session=active_session,
@@ -1529,8 +1531,12 @@ async def _prepare_run_context(
         session_messages: list[Message] = session_context.get_messages(include_input=True)
 
         effective_client_kwargs = dict(client_kwargs) if client_kwargs is not None else {}
+        from ._tools import _APPROVAL_SESSION_IS_AUTHORITATIVE_KEY  # pyright: ignore[reportPrivateUsage]
+
+        effective_client_kwargs.pop(_APPROVAL_SESSION_IS_AUTHORITATIVE_KEY, None)
         if active_session is not None:
             effective_client_kwargs["session"] = active_session
+            effective_client_kwargs[_APPROVAL_SESSION_IS_AUTHORITATIVE_KEY] = not framework_created_session
         per_service_call_history_middleware: PerServiceCallHistoryPersistingMiddleware | None = None
         if per_service_call_history_providers and active_session is not None:
             per_service_call_history_middleware = PerServiceCallHistoryPersistingMiddleware(