Python: Isolate provider-owned state in agent tools

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

@@ -524,7 +524,10 @@ that manually replay messages own the equivalent rule: do not resend an approval
   `propagate_session=True`, child application-state changes merge back into the parent while framework approval and
   invocation-budget state remain isolated, including approval queues stored under custom child middleware
   `source_id` values. Parent and child `ToolApprovalMiddleware` instances must use distinct `source_id` values; an
-  overlap fails before the child runs.
+  overlap fails before the child runs. The automatic function loop also supplies the parent agent/client's declared
+  provider-owned state keys. A custom loop that directly invokes the agent tool must set
+  `FunctionInvocationContext.parent_service_session_state_keys` explicitly; when that ownership is unavailable, a
+  non-empty parent session fails before the child runs rather than propagating unknown provider state.
 
 ### Approval control content
 
@@ -594,6 +597,7 @@ that manually replay messages own the equivalent rule: do not resend an approval
 | Declaration-only mixed batch | Known local siblings are never converted into Host pauses. With an authoritative session they remain deferred until every Host pause occurrence is answered, survive session serialization with a process-portable logical-run budget (including a reclassified approval continuation), reclassify against the current tool policy, execute exactly once using the configured concurrent or sequential invocation mode, replay results in original model order, and remain inert on replay; without one they execute on the same turn under the same configured invocation mode. Streaming updates and finalization preserve one Host pause and one result per local sibling. | `test_mixed_declaration_only_batch_defers_siblings_until_all_host_results`, `test_mixed_declaration_only_deferred_call_keeps_ordinary_middleware_semantics`, `test_mixed_declaration_only_deferred_call_reclassifies_changed_approval_policy`, `test_mixed_declaration_only_budget_survives_reclassification_to_approval`, `test_mixed_declaration_only_deferred_call_does_not_execute_removed_tool`, `test_mixed_declaration_only_budget_is_process_portable_without_approval_middleware`, `test_sessionless_mixed_declaration_only_batch_executes_siblings_sequentially`, `test_sessionless_mixed_declaration_only_batch_executes_siblings_concurrently`, `test_streaming_sessionless_mixed_declaration_only_batch_returns_pause_and_result` |
 | Function invocation disabled | The client bypasses the invocation loop without losing invocation kwargs. | `test_function_invocation_config_enabled_false`, `test_function_invocation_config_enabled_false_preserves_invocation_kwargs`, `test_streaming_function_invocation_config_enabled_false` |
 | Runtime tool changes | Added tools become available on the next iteration and retain approval behavior. | `test_add_tools_available_next_iteration`, `test_add_tools_with_approval_required_tool` |
+| Direct delegated session propagation | A custom loop must explicitly identify parent provider-owned state before sharing a non-empty session; missing ownership fails before the child runs, while declared parent keys remain isolated and application state still propagates. | `packages/core/tests/core/test_agents.py::test_chat_agent_as_tool_direct_propagation_requires_parent_provider_ownership`, `test_chat_agent_as_tool_direct_propagation_uses_explicit_parent_provider_ownership` |
 | In-run compaction summaries | Summaries inserted for tool groups excluded by in-run compaction are returned in the final non-streaming response transcript before the group they replace, so history loaded with `skip_excluded` keeps the summarized content; summaries of caller-owned input messages are not added. | `packages/core/tests/core/test_clients.py::test_function_loop_returns_compaction_summaries_in_final_response`, `test_function_loop_returns_compaction_summaries_when_iteration_budget_exhausted`, `test_function_loop_reconciles_nested_compaction_summaries`, `test_function_loop_returns_compacted_transcript_on_early_terminal_exit`, `packages/core/tests/core/test_agents.py::test_agent_run_returns_and_persists_compaction_summaries` |
 | Per-service-call provider context | When framework-managed history is loaded around each model call, tools and instructions contributed by `HistoryProvider.before_run` remain available to the model and function loop without duplication in streaming or non-streaming execution. | `packages/core/tests/core/test_agents.py::test_vector_history_search_tool_is_available_with_per_service_call_persistence` |
 

python/packages/core/AGENTS.md

@@ -57,6 +57,12 @@ agent_framework/
 - **`RawAgent.open()` / `close()`** (inherited by `Agent`) - Explicitly enter the client's and configured MCP tools'
   async contexts, then release them along with lazily connected MCP tools. `async with agent` delegates to these
   methods; a partially failed `open()` closes resources already entered.
+- Providers that store service-owned continuation handles in `AgentSession.state` declare their keys through
+  `service_session_state_keys`. `as_tool(propagate_session=True)` combines declarations from the parent and child
+  agents and their clients to isolate those handles in both directions while ordinary application-owned state
+  propagates. The automatic function loop places parent declarations on `FunctionInvocationContext`; direct or
+  custom-loop invocation must set `parent_service_session_state_keys` explicitly for a non-empty shared session
+  and otherwise fails closed.
 
 ### Chat Clients (`_clients.py`)
 

python/packages/core/agent_framework/_agents.py

@@ -81,6 +81,8 @@
     from typing_extensions import Self, TypedDict  # pragma: no cover
 
 if TYPE_CHECKING:
+    from collections.abc import Collection
+
     from mcp import types
     from mcp.server.lowlevel import Server
     from pydantic import BaseModel
@@ -113,6 +115,17 @@ def _tool_approval_source_ids(middleware: Sequence[MiddlewareTypes] | None) -> f
     return frozenset(item.source_id for item in middleware if isinstance(item, ToolApprovalMiddleware))
 
 
+def _provider_service_session_state_keys(agent: object) -> frozenset[str]:
+    """Return provider-owned session-state keys declared by an agent or its client."""
+    keys: set[str] = set()
+    # A generic Agent can use a provider client that persists its own continuation state.
+    for owner in (agent, getattr(agent, "client", None)):
+        declared = getattr(owner, "service_session_state_keys", ())
+        if isinstance(declared, (list, tuple, set, frozenset)):
+            keys.update(key for key in cast("Collection[Any]", declared) if isinstance(key, str))
+    return frozenset(keys)
+
+
 def _merge_delegated_session_state(
     parent_state: MutableMapping[str, Any],
     initial_child_state: Mapping[str, Any],
@@ -665,9 +678,10 @@ def as_tool(
             propagate_session: If True, the parent agent's session is forwarded
                 to this sub-agent's ``run()`` call. Application-state changes
                 propagate back to the parent, while framework approval continuation
-                state remains isolated. Defaults to False. The sub-agent always
-                receives an AgentSession so session-backed middleware can run.
-                When False, that session is private to this invocation.
+                state and provider-owned service session state remain isolated.
+                Defaults to False. The sub-agent always receives an AgentSession
+                so session-backed middleware can run. When False, that session is
+                private to this invocation.
 
         Returns:
             A FunctionTool that can be used as a tool by other agents.
@@ -683,6 +697,10 @@ def as_tool(
             values. The delegated call raises ToolExecutionException before running
             the child when their shared session-state keys overlap.
 
+            Direct or custom-loop invocation with ``propagate_session=True`` must set
+            :attr:`FunctionInvocationContext.parent_service_session_state_keys` when
+            the parent session is non-empty. Automatic function calling supplies it.
+
         Examples:
             .. code-block:: python
 
@@ -733,10 +751,24 @@ async def _agent_wrapper(ctx: FunctionInvocationContext, **kwargs: Any) -> str:
             session = AgentSession()
             child_approval_source_ids = _tool_approval_source_ids(self.middleware)
             parent_approval_source_ids: frozenset[str] = frozenset()
+            parent_service_session_state_keys: frozenset[str] = frozenset()
 
             if propagate_session and parent_session is not None:
                 from ._tools import _PARENT_TOOL_APPROVAL_SOURCE_IDS_CONTEXT_KEY  # pyright: ignore[reportPrivateUsage]
 
+                # A custom loop can bypass the framework seam that identifies the parent provider's state.
+                # Refuse a non-empty shared session rather than guessing which keys are safe to delegate.
+                if ctx.parent_service_session_state_keys is None:
+                    if parent_session.state:
+                        raise ToolExecutionException(
+                            f"Agent tool {tool_name!r} cannot safely propagate a non-empty parent session because "
+                            "the invocation path did not provide parent provider-owned session state keys. "
+                            "Use the automatic function-calling loop, set "
+                            "FunctionInvocationContext.parent_service_session_state_keys in a custom loop, "
+                            "or set propagate_session=False."
+                        )
+                else:
+                    parent_service_session_state_keys = ctx.parent_service_session_state_keys
                 raw_parent_approval_source_ids = ctx.metadata.get(_PARENT_TOOL_APPROVAL_SOURCE_IDS_CONTEXT_KEY)
                 parent_approval_source_ids = (
                     cast("frozenset[str]", raw_parent_approval_source_ids)
@@ -770,6 +802,10 @@ async def _agent_wrapper(ctx: FunctionInvocationContext, **kwargs: Any) -> str:
                     _TOOL_APPROVAL_STATE_KEY,
                     _FUNCTION_INVOCATION_BUDGET_STATE_KEY,
                     _FUNCTION_RESULT_PAYLOAD_BUDGET_STATE_KEY,
+                    # Service handles belong to one agent's remote resources, not shared application state.
+                    # Exclude them in both directions so later children cannot inherit an earlier child's handles.
+                    *_provider_service_session_state_keys(self),
+                    *parent_service_session_state_keys,
                     *child_approval_source_ids,
                     *parent_approval_source_ids,
                 })
@@ -1589,7 +1625,10 @@ async def _prepare_run_context(
 
         agent_name = self._get_agent_name()
         from ._mcp import MCPTool
-        from ._tools import _PARENT_TOOL_APPROVAL_SOURCE_IDS_CONTEXT_KEY  # pyright: ignore[reportPrivateUsage]
+        from ._tools import (
+            _PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY,  # pyright: ignore[reportPrivateUsage]
+            _PARENT_TOOL_APPROVAL_SOURCE_IDS_CONTEXT_KEY,  # pyright: ignore[reportPrivateUsage]
+        )
 
         base_tools = _normalize_tools(chat_options.pop("tools", None))
         mcp_duplicate_message = "Tool names must be unique. Consider setting `tool_name_prefix` on the MCPTool."
@@ -1637,6 +1676,10 @@ async def _prepare_run_context(
         additional_function_arguments[_PARENT_TOOL_APPROVAL_SOURCE_IDS_CONTEXT_KEY] = _tool_approval_source_ids(
             self.middleware
         )
+        # Recompute ownership for this invoking agent; caller kwargs must not replace its declarations.
+        additional_function_arguments[_PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY] = (
+            _provider_service_session_state_keys(self)
+        )
 
         model = opts.pop("model", None)
 

python/packages/core/agent_framework/_middleware.py

@@ -457,6 +457,10 @@ class FunctionInvocationContext:
                 bare ``Content`` of type ``function_approval_request``, which the
                 invocation layer passes through untouched to drive the approval flow.
         kwargs: Additional runtime keyword arguments forwarded to the function invocation.
+        parent_service_session_state_keys: Provider-owned keys in ``session.state`` declared by the invoking
+                parent agent and its client. The automatic function-calling loop supplies this value. Custom loops
+                that directly invoke an agent tool with session propagation must provide it explicitly, including
+                an empty collection when the parent declares no such keys.
         tools: The live, mutable list of tools available to the model for the current
                 agent run, or ``None`` when the function is invoked outside of a
                 function-calling loop (for example via ``FunctionTool.invoke`` directly).
@@ -504,6 +508,7 @@ def __init__(
         result: Any = None,
         kwargs: Mapping[str, Any] | None = None,
         tools: list[ToolTypes] | None = None,
+        parent_service_session_state_keys: Collection[str] | None = None,
     ) -> None:
         """Initialize the FunctionInvocationContext.
 
@@ -520,13 +525,18 @@ def __init__(
             tools: The live, mutable list of tools for the current agent run. When provided,
                 this is the same list object the model sees on the next iteration, so
                 appending or removing tools changes the model's available tools.
+            parent_service_session_state_keys: Provider-owned keys declared by the invoking parent agent and client,
+                or ``None`` when the invocation path cannot establish that ownership.
         """
         self.function = function
         self.arguments = arguments
         self.session = session
         self.metadata: dict[str, Any] = dict(metadata) if metadata is not None else {}
         self.result = result
         self.kwargs: dict[str, Any] = dict(kwargs) if kwargs is not None else {}
+        self.parent_service_session_state_keys: frozenset[str] | None = (
+            frozenset(parent_service_session_state_keys) if parent_service_session_state_keys is not None else None
+        )
         self.tools = tools
 
     @experimental(feature_id=ExperimentalFeature.PROGRESSIVE_TOOLS)

python/packages/core/agent_framework/_tools.py

@@ -107,6 +107,7 @@ def _generate_function_call_occurrence_id() -> str:
 _TOOL_APPROVAL_STATE_KEY: Final[str] = "tool_approval"
 _APPROVAL_SESSION_IS_AUTHORITATIVE_KEY: Final[str] = "_approval_session_is_authoritative"
 _PARENT_TOOL_APPROVAL_SOURCE_IDS_CONTEXT_KEY: Final[str] = "_parent_tool_approval_source_ids"
+_PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY: Final[str] = "_parent_service_session_state_keys"
 
 
 def _has_authoritative_approval_session(invocation_session: AgentSession | None) -> bool:
@@ -2099,6 +2100,7 @@ async def _auto_invoke_function(
             "middleware",
             "conversation_id",
             _PARENT_TOOL_APPROVAL_SOURCE_IDS_CONTEXT_KEY,
+            _PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY,
         }
     }
     raw_parent_approval_source_ids = (custom_args or {}).get(_PARENT_TOOL_APPROVAL_SOURCE_IDS_CONTEXT_KEY)
@@ -2108,6 +2110,12 @@ async def _auto_invoke_function(
         if isinstance(raw_parent_approval_source_ids, frozenset)
         else frozenset()
     )
+    raw_parent_service_keys = (custom_args or {}).get(_PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY)
+    parent_service_session_state_keys: frozenset[str] = (
+        cast("frozenset[str]", raw_parent_service_keys)
+        if isinstance(raw_parent_service_keys, frozenset)
+        else frozenset()
+    )
     if invocation_session is not None:
         runtime_kwargs["session"] = invocation_session
     args = dict(parsed_args)
@@ -2124,6 +2132,7 @@ async def _auto_invoke_function(
                     arguments=args,
                     session=invocation_session,
                     kwargs=runtime_kwargs.copy(),
+                    parent_service_session_state_keys=parent_service_session_state_keys,
                     tools=live_tools,
                 )
                 direct_context.metadata[_PARENT_TOOL_APPROVAL_SOURCE_IDS_CONTEXT_KEY] = parent_approval_source_ids
@@ -2163,6 +2172,7 @@ async def _auto_invoke_function(
         arguments=args,
         session=invocation_session,
         kwargs=runtime_kwargs.copy(),
+        parent_service_session_state_keys=parent_service_session_state_keys,
         tools=live_tools,
     )
     middleware_context.metadata[_PARENT_TOOL_APPROVAL_SOURCE_IDS_CONTEXT_KEY] = parent_approval_source_ids

python/packages/core/tests/core/test_agent_hooks.py

@@ -2194,7 +2194,10 @@ async def _run() -> AgentResponse[Any]:
                 # The custom loop invokes the sub-agent tool directly, bypassing
                 # _execute_single_function_call (like GitHubCopilotAgent's loop).
                 tool_context = FunctionInvocationContext(
-                    function=sub_tool, arguments={"task": "delegated task"}, session=session
+                    function=sub_tool,
+                    arguments={"task": "delegated task"},
+                    session=session,
+                    parent_service_session_state_keys=(),
                 )
                 await sub_tool.invoke(arguments={"task": "delegated task"}, context=tool_context)
                 response: AgentResponse[Any] = AgentResponse(