From 79bd4f9412f075e994a3fe6d17ed87eeac651a84 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Tue, 29 Sep 2026 16:33:56 +0100 Subject: [PATCH 1/4] fix(python): isolate delegated session state Keep provider-owned state with its agent or client while retaining shared application state. --- python/packages/core/AGENTS.md | 3 + .../packages/core/agent_framework/_agents.py | 23 +++++-- .../packages/core/tests/core/test_agents.py | 61 +++++++++++++++++++ python/packages/foundry/README.md | 4 ++ .../foundry/agent_framework_foundry/_agent.py | 5 +- .../tests/foundry/test_foundry_agent.py | 61 +++++++++++++++++++ 6 files changed, 152 insertions(+), 5 deletions(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 8e9d3a5e393..1e4403cb727 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -57,6 +57,9 @@ 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 child agent and + its client to isolate those handles in both directions while ordinary application-owned state propagates. ### Chat Clients (`_clients.py`) diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index db264535aa1..6f82de013e8 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -5,7 +5,7 @@ import logging import re import sys -from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence +from collections.abc import Awaitable, Callable, Collection, Mapping, MutableMapping, Sequence from contextlib import AbstractAsyncContextManager, AsyncExitStack from copy import deepcopy from functools import partial @@ -110,6 +110,17 @@ def _tool_approval_source_ids(middleware: Sequence[MiddlewareTypes] | None) -> f return frozenset(item.source_id for item in middleware or () 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], @@ -662,9 +673,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. @@ -767,6 +779,9 @@ 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), *child_approval_source_ids, *parent_approval_source_ids, }) diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index e0e79153288..bb2d83f8be5 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/python/packages/core/tests/core/test_agents.py @@ -2585,6 +2585,67 @@ def capturing_run(*args: Any, **kwargs: Any) -> Any: assert parent_session.state["counter"] == 1 +@pytest.mark.parametrize("parent_handle", [None, "parent-provider-session"]) +async def test_chat_agent_as_tool_propagate_session_isolates_provider_owned_state( + client: SupportsChatGetResponse, parent_handle: str | None, monkeypatch: pytest.MonkeyPatch +) -> None: + """Test that provider-owned state stays isolated across sequential delegated calls.""" + + class ProviderStateAgent(Agent): + service_session_state_keys = frozenset({"provider_session"}) + + monkeypatch.setattr(client, "service_session_state_keys", frozenset({"client_provider_session"}), raising=False) + parent_session = AgentSession(session_id="shared-session", service_session_id="parent-service-session") + parent_session.state.update({ + "counter": 0, + "ordinary": "parent-value", + }) + if parent_handle is not None: + parent_session.state["provider_session"] = parent_handle + parent_session.state["client_provider_session"] = parent_handle + + observed_child_states: list[dict[str, Any]] = [] + agents = [ + ProviderStateAgent(client=client, name="FirstChild"), + ProviderStateAgent(client=client, name="SecondChild"), + ] + + for agent in (agents[0], agents[1], agents[0]): + original_run = agent.run + + def capturing_run(*args: Any, _run: Callable[..., Any] = original_run, **kwargs: Any) -> Any: + child_session = cast(AgentSession, kwargs["session"]) + observed_child_states.append(dict(child_session.state)) + assert child_session.service_session_id is None + child_session.state["counter"] += 1 + child_session.state["ordinary"] = f"child-value-{len(observed_child_states)}" + child_session.state["provider_session"] = f"child-provider-session-{len(observed_child_states)}" + child_session.state["client_provider_session"] = f"client-provider-session-{len(observed_child_states)}" + child_session.service_session_id = f"child-service-session-{len(observed_child_states)}" + return _run(*args, **kwargs) + + delegated_tool = agent.as_tool(propagate_session=True) + with patch.object(agent, "run", side_effect=capturing_run): + await delegated_tool.invoke( + context=FunctionInvocationContext( + function=delegated_tool, + arguments={"task": "Run child"}, + session=parent_session, + ) + ) + + assert [state.get("provider_session") for state in observed_child_states] == [None, None, None] + assert [state.get("client_provider_session") for state in observed_child_states] == [None, None, None] + assert [state["counter"] for state in observed_child_states] == [0, 1, 2] + assert [state["ordinary"] for state in observed_child_states] == ["parent-value", "child-value-1", "child-value-2"] + expected_state: dict[str, Any] = {"counter": 3, "ordinary": "child-value-3"} + if parent_handle is not None: + expected_state["provider_session"] = parent_handle + expected_state["client_provider_session"] = parent_handle + assert parent_session.state == expected_state + assert parent_session.service_session_id == "parent-service-session" + + async def test_chat_agent_as_tool_propagate_session_clears_service_session_id(client: SupportsChatGetResponse) -> None: """Test that propagate_session=True gives the child a separate session with cleared service_session_id.""" agent = Agent(client=client, name="SubAgent", description="Sub agent") diff --git a/python/packages/foundry/README.md b/python/packages/foundry/README.md index 56d13cbd872..8d1c26f9aff 100644 --- a/python/packages/foundry/README.md +++ b/python/packages/foundry/README.md @@ -82,6 +82,10 @@ entry points for existing traces, response IDs, and registered targets. ## Concurrent reuse +When a Foundry agent is used through `as_tool(propagate_session=True)`, application-owned state propagates +between the parent and child, but service-owned session handles do not. Each delegated invocation has its own +service session. This also applies to an `Agent` configured with a `FoundryAgentChatClient`. + A `FoundryChatClient` instance can be shared by concurrent asynchronous calls on the same event loop. Streaming, non-streaming, and mixed calls are supported. Keep mutable run state isolated by creating a separate `Agent` and `AgentSession` for each concurrent run and by passing separate messages and options. diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index 43374b85d29..73d59fad494 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/python/packages/foundry/agent_framework_foundry/_agent.py @@ -181,6 +181,9 @@ class MyClient(FunctionInvocationLayer, RawFoundryAgentChatClient): OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.foundry" _FEATURE_USAGE_INDEX: ClassVar[int | None] = FeatureIndex.FOUNDRY_AGENT + service_session_state_keys: ClassVar[frozenset[str]] = frozenset({FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY}) + """Service-owned state keys, including when this client is used by a generic Agent.""" + def __init__( self, *, @@ -657,7 +660,7 @@ class RawFoundryAgent( result = await agent.run("Hello!") """ - service_session_state_keys: ClassVar[frozenset[str]] = frozenset({FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY}) + service_session_state_keys: ClassVar[frozenset[str]] = RawFoundryAgentChatClient.service_session_state_keys """Session-state keys this agent owns, which untrusted input must never supply. Holds the Foundry hosted-agent session ID. Despite the attribute name, that value is not a conversation or diff --git a/python/packages/foundry/tests/foundry/test_foundry_agent.py b/python/packages/foundry/tests/foundry/test_foundry_agent.py index 3fe96541499..335cdb5756e 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_agent.py +++ b/python/packages/foundry/tests/foundry/test_foundry_agent.py @@ -1223,6 +1223,8 @@ def test_foundry_agents_declare_hosted_agent_session_id_as_server_owned() -> Non assert FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY in RawFoundryAgent.service_session_state_keys # FoundryAgent is the recommended production class, so it must inherit the same protection. assert FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY in FoundryAgent.service_session_state_keys + assert FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY in RawFoundryAgentChatClient.service_session_state_keys + assert FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY in _FoundryAgentChatClient.service_session_state_keys async def test_raw_foundry_agent_prepare_run_context_injects_agent_session_id_from_state() -> None: @@ -1263,6 +1265,65 @@ async def test_raw_foundry_agent_prepare_run_context_injects_agent_session_id_fr } +@pytest.mark.parametrize("agent_type", [RawFoundryAgent, FoundryAgent, None]) +@pytest.mark.parametrize("parent_handle", [None, "parent-agent-session"]) +async def test_foundry_agent_tools_isolate_service_state_between_children( + agent_type: type[RawFoundryAgent] | None, parent_handle: str | None +) -> None: + """Delegated Foundry calls retain application state without sharing service handles.""" + project = MagicMock() + project.get_openai_client.return_value = MagicMock() + children = [ + agent_type(project_client=project, agent_name=name) + if agent_type is not None + else Agent(client=_FoundryAgentChatClient(project_client=project, agent_name=name), name=name) + for name in ("first-child", "second-child") + ] + parent = AgentSession(service_session_id="parent-response") + parent.state["ordinary"] = "parent-value" + if parent_handle is not None: + parent.state[FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY] = parent_handle + observed_options: list[dict[str, Any]] = [] + updates: list[AgentResponseUpdate] = [] + + def respond(**kwargs: Any) -> ResponseStream[ChatResponseUpdate, ChatResponse]: + assert kwargs["stream"] is True + observed_options.append(dict(kwargs["options"])) + call_number = len(observed_options) + + async def stream() -> AsyncIterator[ChatResponseUpdate]: + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_text(text="done")], + conversation_id=f"child-response-{call_number}", + additional_properties={"agent_session_id": f"child-agent-session-{call_number}"}, + ) + + return ResponseStream(stream(), finalizer=ChatResponse.from_updates) + + with patch.object(RawFoundryAgentChatClient, "_inner_get_response", side_effect=respond): + for child in (children[0], children[1], children[0]): + delegated_tool = child.as_tool(propagate_session=True, stream_callback=updates.append) + result = await delegated_tool.invoke( + context=FunctionInvocationContext( + function=delegated_tool, + arguments={"task": "Run child"}, + session=parent, + ) + ) + assert result[0].text == "done" + + assert len(observed_options) == 3 + assert all("agent_session_id" not in options.get("extra_body", {}) for options in observed_options) + assert all(options.get("conversation_id") is None for options in observed_options) + assert len(updates) == 3 + expected_state = {"ordinary": "parent-value"} + if parent_handle is not None: + expected_state[FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY] = parent_handle + assert parent.state == expected_state + assert parent.service_session_id == "parent-response" + + def test_foundry_agent_updates_session_from_response_ids() -> None: """Test that response and hosted-agent session IDs persist independently.""" From 72f70786d8fa5cd9ad95a30b22556bf5d40318b9 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Tue, 29 Sep 2026 17:02:17 +0100 Subject: [PATCH 2/4] fix(python): preserve parent provider state Include the invoking agent and client declarations in delegated state filtering so providers with different keys retain their own state. --- python/packages/core/AGENTS.md | 5 +- .../packages/core/agent_framework/_agents.py | 26 ++++- .../packages/core/agent_framework/_tools.py | 12 +++ .../packages/core/tests/core/test_agents.py | 95 +++++++++++++++++++ python/packages/foundry/README.md | 2 +- .../foundry/agent_framework_foundry/_agent.py | 2 +- 6 files changed, 135 insertions(+), 7 deletions(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 1e4403cb727..fa2f2a9df1e 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -58,8 +58,9 @@ agent_framework/ 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 child agent and - its client to isolate those handles in both directions while ordinary application-owned state propagates. + `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. Parent declarations travel through private function-invocation metadata, not tool arguments. ### Chat Clients (`_clients.py`) diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 6f82de013e8..c4b12f844e9 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -5,7 +5,7 @@ import logging import re import sys -from collections.abc import Awaitable, Callable, Collection, Mapping, MutableMapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence from contextlib import AbstractAsyncContextManager, AsyncExitStack from copy import deepcopy from functools import partial @@ -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 @@ -742,10 +744,20 @@ 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] + from ._tools import ( + _PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage] + _PARENT_TOOL_APPROVAL_SOURCE_IDS_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage] + ) + raw_parent_service_keys = ctx.metadata.get(_PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY) + parent_service_session_state_keys = ( + cast("frozenset[str]", raw_parent_service_keys) + if isinstance(raw_parent_service_keys, frozenset) + else frozenset() + ) 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) @@ -782,6 +794,7 @@ async def _agent_wrapper(ctx: FunctionInvocationContext, **kwargs: Any) -> str: # 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, }) @@ -1601,7 +1614,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." @@ -1649,6 +1665,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) diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 4ab15a52f79..187ad79155f 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/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: @@ -2095,6 +2096,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) @@ -2104,6 +2106,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) @@ -2123,6 +2131,9 @@ async def _auto_invoke_function( tools=live_tools, ) direct_context.metadata[_PARENT_TOOL_APPROVAL_SOURCE_IDS_CONTEXT_KEY] = parent_approval_source_ids + direct_context.metadata[_PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY] = ( + parent_service_session_state_keys + ) if host_payload_budget is not None: direct_context.metadata[_FUNCTION_RESULT_PAYLOAD_BUDGET_CONTEXT_KEY] = host_payload_budget function_result = await tool.invoke( @@ -2162,6 +2173,7 @@ async def _auto_invoke_function( tools=live_tools, ) middleware_context.metadata[_PARENT_TOOL_APPROVAL_SOURCE_IDS_CONTEXT_KEY] = parent_approval_source_ids + middleware_context.metadata[_PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY] = parent_service_session_state_keys if host_payload_budget is not None: middleware_context.metadata[_FUNCTION_RESULT_PAYLOAD_BUDGET_CONTEXT_KEY] = host_payload_budget middleware_context.metadata[_AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY] = True diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index bb2d83f8be5..01c4a306405 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/python/packages/core/tests/core/test_agents.py @@ -53,6 +53,7 @@ from agent_framework._agents import _get_tool_name, _merge_options, _sanitize_agent_name from agent_framework._mcp import MCPTool, _build_prefixed_mcp_name, _normalize_mcp_name from agent_framework._middleware import FunctionInvocationContext +from agent_framework._tools import _PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY from agent_framework.exceptions import ( AgentInvalidRequestException, ChatClientInvalidResponseException, @@ -2646,6 +2647,100 @@ def capturing_run(*args: Any, _run: Callable[..., Any] = original_run, **kwargs: assert parent_session.service_session_id == "parent-service-session" +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("with_middleware", [False, True]) +@pytest.mark.parametrize("child_declares_keys", [False, True]) +async def test_chat_agent_as_tool_isolates_distinct_parent_provider_state( + stream: bool, with_middleware: bool, child_declares_keys: bool +) -> None: + """Keep parent agent and client state local even when the child uses different declarations.""" + + class ParentAgent(Agent): + service_session_state_keys = frozenset({"parent_provider", "parent_removed"}) + + class ParentClient(MockBaseChatClient): + service_session_state_keys = frozenset({"parent_client_provider", "parent_client_removed"}) + + class ChildAgent(Agent): + service_session_state_keys = frozenset({"child_provider"}) if child_declares_keys else frozenset() + + parent_state: dict[str, Any] = { + "parent_provider": {"values": ["parent"]}, + "parent_removed": "parent", + "parent_client_provider": "parent-client", + "parent_client_removed": "parent-client", + "counter": 0, + } + parent_session = AgentSession() + parent_session.state.update(parent_state) + observed_states: list[dict[str, Any]] = [] + + @tool(approval_mode="never_require") + def inspect_state(ctx: FunctionInvocationContext) -> str: + assert ctx.session is not None + observed_states.append(dict(ctx.session.state)) + assert _PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY not in ctx.kwargs + ctx.session.state.setdefault("parent_provider", {"values": []})["values"].append("child") + ctx.session.state.pop("parent_removed", None) + ctx.session.state["parent_client_provider"] = "child" + ctx.session.state.pop("parent_client_removed", None) + ctx.session.state["counter"] += 1 + if child_declares_keys: + ctx.session.state["child_provider"] = "child" + return "State inspected." + + child_client = MockBaseChatClient() + child_client.streaming_responses = [ + [ + ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="inspect", name="inspect_state", arguments={})], + ) + ], + [ChatResponseUpdate(role="assistant", contents=[Content.from_text("Child completed.")])], + ] + child = ChildAgent(client=child_client, tools=[inspect_state]) + parent_client = ParentClient() + parent_call = Content.from_function_call( + call_id="delegate", name="delegate", arguments={"task": "Inspect the delegated session"} + ) + if stream: + parent_client.streaming_responses = [ + [ChatResponseUpdate(role="assistant", contents=[parent_call])], + [ChatResponseUpdate(role="assistant", contents=[Content.from_text("Parent completed.")])], + ] + else: + parent_client.run_responses = [ + ChatResponse(messages=Message(role="assistant", contents=[parent_call])), + ChatResponse(messages=Message(role="assistant", contents=[Content.from_text("Parent completed.")])), + ] + parent = ParentAgent( + client=parent_client, + tools=[child.as_tool(name="delegate", propagate_session=True)], + middleware=[ToolApprovalMiddleware(source_id="parent_approval")] if with_middleware else [], + ) + invocation_kwargs: dict[str, frozenset[str]] = {_PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY: frozenset()} + if stream: + result = await parent.run( + "Delegate.", session=parent_session, stream=True, function_invocation_kwargs=invocation_kwargs + ).get_final_response() + else: + result = await parent.run("Delegate.", session=parent_session, function_invocation_kwargs=invocation_kwargs) + + assert result.text == "Parent completed." + assert len(observed_states) == 1 + assert all( + key not in observed_states[0] + for key in ParentAgent.service_session_state_keys | ParentClient.service_session_state_keys + ) + assert parent_session.state["parent_provider"] == {"values": ["parent"]} + assert parent_session.state["parent_removed"] == "parent" + assert parent_session.state["parent_client_provider"] == "parent-client" + assert parent_session.state["parent_client_removed"] == "parent-client" + assert "child_provider" not in parent_session.state + assert parent_session.state["counter"] == 1 + + async def test_chat_agent_as_tool_propagate_session_clears_service_session_id(client: SupportsChatGetResponse) -> None: """Test that propagate_session=True gives the child a separate session with cleared service_session_id.""" agent = Agent(client=client, name="SubAgent", description="Sub agent") diff --git a/python/packages/foundry/README.md b/python/packages/foundry/README.md index 8d1c26f9aff..8b170212cfc 100644 --- a/python/packages/foundry/README.md +++ b/python/packages/foundry/README.md @@ -84,7 +84,7 @@ entry points for existing traces, response IDs, and registered targets. When a Foundry agent is used through `as_tool(propagate_session=True)`, application-owned state propagates between the parent and child, but service-owned session handles do not. Each delegated invocation has its own -service session. This also applies to an `Agent` configured with a `FoundryAgentChatClient`. +service session. This also applies to an `Agent` configured with a `RawFoundryAgentChatClient`. A `FoundryChatClient` instance can be shared by concurrent asynchronous calls on the same event loop. Streaming, non-streaming, and mixed calls are supported. Keep mutable run state isolated by creating a separate `Agent` and diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index 73d59fad494..fbceea8cff3 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/python/packages/foundry/agent_framework_foundry/_agent.py @@ -660,7 +660,7 @@ class RawFoundryAgent( result = await agent.run("Hello!") """ - service_session_state_keys: ClassVar[frozenset[str]] = RawFoundryAgentChatClient.service_session_state_keys + service_session_state_keys: ClassVar[frozenset[str]] = frozenset({FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY}) """Session-state keys this agent owns, which untrusted input must never supply. Holds the Foundry hosted-agent session ID. Despite the attribute name, that value is not a conversation or From 72e93b0f430376148a3524a7a10baea9fa8a96f6 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Mon, 5 Oct 2026 16:07:33 +0100 Subject: [PATCH 3/4] fix(python): fail closed on unknown parent state Custom tool loops must declare parent provider-owned keys before propagating a non-empty session, preventing unknown service handles from crossing the delegation boundary. Copilot-Session: 27b76375-2d09-48e9-b487-4b343bdbeb4c --- .../specs/004-python-function-calling-loop.md | 6 +- python/packages/core/AGENTS.md | 4 +- .../packages/core/agent_framework/_agents.py | 30 +++++--- .../core/agent_framework/_middleware.py | 10 +++ .../packages/core/agent_framework/_tools.py | 6 +- .../core/tests/core/test_agent_hooks.py | 5 +- .../packages/core/tests/core/test_agents.py | 72 +++++++++++++++++++ .../tests/foundry/test_foundry_agent.py | 1 + 8 files changed, 116 insertions(+), 18 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index bdf2fc8ac23..69a0e820f3d 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -523,7 +523,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 @@ -592,6 +595,7 @@ that manually replay messages own the equivalent rule: do not resend an approval | Declaration-only call | The call is surfaced as user input and is not executed; streaming arguments appear once while finalized request metadata remains available. | `test_declaration_only_tool`, `test_streaming_declaration_only_tool_preserves_metadata_without_duplicate_arguments` | | 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` | diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index fa2f2a9df1e..fe7bcf970f5 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -60,7 +60,9 @@ agent_framework/ - 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. Parent declarations travel through private function-invocation metadata, not tool arguments. + 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`) diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index c4b12f844e9..4f99d33b047 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -694,6 +694,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 @@ -747,17 +751,21 @@ async def _agent_wrapper(ctx: FunctionInvocationContext, **kwargs: Any) -> str: parent_service_session_state_keys: frozenset[str] = frozenset() if propagate_session and parent_session is not None: - from ._tools import ( - _PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage] - _PARENT_TOOL_APPROVAL_SOURCE_IDS_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage] - ) - - raw_parent_service_keys = ctx.metadata.get(_PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY) - parent_service_session_state_keys = ( - cast("frozenset[str]", raw_parent_service_keys) - if isinstance(raw_parent_service_keys, frozenset) - else frozenset() - ) + 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) diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 8af5c885c4c..f3615eeba46 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/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,6 +525,8 @@ 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 @@ -527,6 +534,9 @@ def __init__( 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) diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 187ad79155f..5318f47177e 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -2128,12 +2128,10 @@ 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 - direct_context.metadata[_PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY] = ( - parent_service_session_state_keys - ) if host_payload_budget is not None: direct_context.metadata[_FUNCTION_RESULT_PAYLOAD_BUDGET_CONTEXT_KEY] = host_payload_budget function_result = await tool.invoke( @@ -2170,10 +2168,10 @@ 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 - middleware_context.metadata[_PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY] = parent_service_session_state_keys if host_payload_budget is not None: middleware_context.metadata[_FUNCTION_RESULT_PAYLOAD_BUDGET_CONTEXT_KEY] = host_payload_budget middleware_context.metadata[_AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY] = True diff --git a/python/packages/core/tests/core/test_agent_hooks.py b/python/packages/core/tests/core/test_agent_hooks.py index c87c424d028..e87193c8a52 100644 --- a/python/packages/core/tests/core/test_agent_hooks.py +++ b/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( diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index 01c4a306405..5bd4b4d8754 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/python/packages/core/tests/core/test_agents.py @@ -2223,6 +2223,7 @@ def capturing_run(*args: Any, **kwargs: Any) -> Any: function=tool, arguments={"task": "Hello"}, session=parent_session, + parent_service_session_state_keys=(), ) ) @@ -2507,6 +2508,7 @@ def second_write() -> str: function=delegated_tool, arguments={"task": "First delegation"}, session=parent_session, + parent_service_session_state_keys=(), ) ) @@ -2517,6 +2519,7 @@ def second_write() -> str: function=delegated_tool, arguments={"task": "Fresh delegation"}, session=parent_session, + parent_service_session_state_keys=(), ) ) @@ -2580,12 +2583,76 @@ def capturing_run(*args: Any, **kwargs: Any) -> Any: function=tool, arguments={"task": "Hello"}, session=parent_session, + parent_service_session_state_keys=(), ) ) assert parent_session.state["counter"] == 1 +async def test_chat_agent_as_tool_direct_propagation_requires_parent_provider_ownership( + client: SupportsChatGetResponse, +) -> None: + """Fail closed when a direct invocation cannot identify parent-owned provider state.""" + agent = Agent(client=client, name="SubAgent", description="Sub agent") + tool = agent.as_tool(propagate_session=True) + parent_session = AgentSession() + parent_session.state["parent_provider_session"] = "parent-handle" + child_run_called = False + original_run = agent.run + + def capturing_run(*args: Any, **kwargs: Any) -> Any: + nonlocal child_run_called + child_run_called = True + return original_run(*args, **kwargs) + + with ( + patch.object(agent, "run", side_effect=capturing_run), + raises(ToolExecutionException, match="parent provider-owned session state"), + ): + await tool.invoke( + context=FunctionInvocationContext( + function=tool, + arguments={"task": "Run child"}, + session=parent_session, + ) + ) + + assert child_run_called is False + assert parent_session.state == {"parent_provider_session": "parent-handle"} + + +async def test_chat_agent_as_tool_direct_propagation_uses_explicit_parent_provider_ownership( + client: SupportsChatGetResponse, +) -> None: + """Allow a direct host to declare parent-owned provider keys explicitly.""" + agent = Agent(client=client, name="SubAgent", description="Sub agent") + tool = agent.as_tool(propagate_session=True) + parent_session = AgentSession() + parent_session.state.update({"parent_provider_session": "parent-handle", "counter": 0}) + original_run = agent.run + + def capturing_run(*args: Any, **kwargs: Any) -> Any: + child_session = cast(AgentSession, kwargs["session"]) + assert "parent_provider_session" not in child_session.state + child_session.state["counter"] += 1 + child_session.state["parent_provider_session"] = "child-handle" + return original_run(*args, **kwargs) + + context = FunctionInvocationContext( + function=tool, + arguments={"task": "Run child"}, + session=parent_session, + parent_service_session_state_keys={"parent_provider_session"}, + ) + + with patch.object(agent, "run", side_effect=capturing_run): + await tool.invoke(context=context) + + assert parent_session.state["parent_provider_session"] == "parent-handle" + assert parent_session.state["counter"] == 1 + + @pytest.mark.parametrize("parent_handle", [None, "parent-provider-session"]) async def test_chat_agent_as_tool_propagate_session_isolates_provider_owned_state( client: SupportsChatGetResponse, parent_handle: str | None, monkeypatch: pytest.MonkeyPatch @@ -2632,6 +2699,10 @@ def capturing_run(*args: Any, _run: Callable[..., Any] = original_run, **kwargs: function=delegated_tool, arguments={"task": "Run child"}, session=parent_session, + parent_service_session_state_keys={ + *ProviderStateAgent.service_session_state_keys, + "client_provider_session", + }, ) ) @@ -2772,6 +2843,7 @@ def capturing_run(*args: Any, **kwargs: Any) -> Any: function=tool, arguments={"task": "Hello"}, session=parent_session, + parent_service_session_state_keys=(), ) ) diff --git a/python/packages/foundry/tests/foundry/test_foundry_agent.py b/python/packages/foundry/tests/foundry/test_foundry_agent.py index 335cdb5756e..a4517e60bf1 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_agent.py +++ b/python/packages/foundry/tests/foundry/test_foundry_agent.py @@ -1309,6 +1309,7 @@ async def stream() -> AsyncIterator[ChatResponseUpdate]: function=delegated_tool, arguments={"task": "Run child"}, session=parent, + parent_service_session_state_keys={FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY}, ) ) assert result[0].text == "done" From dc1be2a8f86c568620777e5d62822a08963100ce Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Tue, 6 Oct 2026 16:10:45 +0100 Subject: [PATCH 4/4] fix(python): pass Copilot parent state ownership Nested agent tools need the Copilot parent declarations after context providers populate shared session state. Copilot-Session: 27b76375-2d09-48e9-b487-4b343bdbeb4c --- python/packages/github_copilot/AGENTS.md | 7 ++ .../agent_framework_github_copilot/_agent.py | 34 +++++++- .../tests/test_github_copilot_agent.py | 85 +++++++++++++++++++ 3 files changed, 125 insertions(+), 1 deletion(-) diff --git a/python/packages/github_copilot/AGENTS.md b/python/packages/github_copilot/AGENTS.md index a4e5718d905..8f594e1c73b 100644 --- a/python/packages/github_copilot/AGENTS.md +++ b/python/packages/github_copilot/AGENTS.md @@ -38,3 +38,10 @@ When the agent applies the `enable_file_hooks` default and the working directory defines hooks, `_warn_once_about_unloaded_file_hooks` logs a warning so the change in behavior is visible. It fires at most once per agent, and never when the caller set the option explicitly either way. + +## Delegated Agent Tools + +The Copilot SDK tool adapter passes the parent agent and client +`service_session_state_keys` declarations through `FunctionInvocationContext`. +This lets nested `as_tool(propagate_session=True)` calls share application state +without inheriting or replacing provider-owned continuation state. diff --git a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py index 5ccb9032ae9..3173be60b88 100644 --- a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py +++ b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py @@ -10,7 +10,16 @@ import sys import warnings import weakref -from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Mapping, MutableMapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterable, + Awaitable, + Callable, + Collection, + Mapping, + MutableMapping, + Sequence, +) from pathlib import Path from typing import Any, ClassVar, Generic, Literal, TypedDict, cast, overload from urllib.parse import urlparse @@ -183,6 +192,17 @@ async def _resolve_function_approval( logger = logging.getLogger("agent_framework.github_copilot") + +def _provider_service_session_state_keys(agent: object, client: object | None) -> frozenset[str]: + """Return provider-owned session-state keys declared by the Copilot agent or client.""" + keys: set[str] = set() + for owner in (agent, client): + 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) + + # One lock per (CopilotClient, service session ID), shared across agents that share a client. # Entries disappear once no run holds or awaits the lock. _SESSION_RUN_LOCKS: weakref.WeakValueDictionary[tuple[int, str], asyncio.Lock] = weakref.WeakValueDictionary() @@ -1354,6 +1374,7 @@ def _prepare_tools( List of Copilot SDK tools. """ copilot_tools: list[CopilotTool] = [] + parent_service_session_state_keys = _provider_service_session_state_keys(self, self._client) for tool in tools: if isinstance(tool, CopilotTool): @@ -1364,6 +1385,7 @@ def _prepare_tools( tool, agent_session=agent_session, function_invocation_kwargs=function_invocation_kwargs, + parent_service_session_state_keys=parent_service_session_state_keys, ) ) elif isinstance(tool, MutableMapping): @@ -1378,6 +1400,7 @@ def _tool_to_copilot_tool( *, agent_session: AgentSession | None = None, function_invocation_kwargs: Mapping[str, Any] | None = None, + parent_service_session_state_keys: frozenset[str] = frozenset(), ) -> CopilotTool: """Convert an FunctionTool to a Copilot SDK tool. @@ -1392,6 +1415,14 @@ def _tool_to_copilot_tool( callback is configured instead, approval is enforced inside this handler for backward compatibility. (``on_function_approval`` and ``on_pre_tool_use`` are mutually exclusive, so only one mechanism is ever active.) + + Args: + ai_func: Agent Framework function tool to convert. + + Keyword Args: + agent_session: Parent agent session passed to the tool. + function_invocation_kwargs: Runtime keyword arguments forwarded to the tool. + parent_service_session_state_keys: Provider-owned state declared by the parent agent and client. """ approval_handler = self._function_approval_handler enforce = approval_handler is not None and ai_func.approval_mode == "always_require" @@ -1419,6 +1450,7 @@ async def handler(invocation: ToolInvocation) -> ToolResult: arguments=args, session=agent_session, kwargs=runtime_kwargs, + parent_service_session_state_keys=parent_service_session_state_keys, ) if ai_func.input_model: args_instance = ai_func.input_model(**args) diff --git a/python/packages/github_copilot/tests/test_github_copilot_agent.py b/python/packages/github_copilot/tests/test_github_copilot_agent.py index bbcf35d114a..e32d425fd3f 100644 --- a/python/packages/github_copilot/tests/test_github_copilot_agent.py +++ b/python/packages/github_copilot/tests/test_github_copilot_agent.py @@ -24,12 +24,14 @@ AgentResponse, AgentResponseUpdate, AgentSession, + BaseAgent, Content, ContextProvider, FunctionInvocationContext, HistoryProvider, MCPStdioTool, Message, + ResponseStream, tool, ) from agent_framework.exceptions import AgentException @@ -2454,6 +2456,89 @@ def lookup(query: str, ctx: FunctionInvocationContext) -> str: # The tool's JSON schema must not expose the context parameter to the model. assert "ctx" not in copilot_tool.parameters["properties"] + @pytest.mark.parametrize("parent_declares_provider_key", [False, True]) + async def test_nested_agent_tool_receives_parent_service_state_ownership( + self, + mock_client: MagicMock, + mock_session: MagicMock, + assistant_message_event: SessionEvent, + parent_declares_provider_key: bool, + ) -> None: + """Test non-empty parent state propagation through the Copilot tool adapter.""" + observed_child_states: list[dict[str, Any]] = [] + + class StateProvider(ContextProvider): + def __init__(self) -> None: + super().__init__(source_id="provider-state") + + async def before_run( + self, + *, + agent: Any, + session: AgentSession, + context: Any, + state: dict[str, Any], + ) -> None: + state["ordinary"] = "parent-value" + session.state["counter"] = 0 + if parent_declares_provider_key: + session.state["parent_provider_session"] = "parent-handle" + + class ChildAgent(BaseAgent): + def run( # type: ignore[override] + self, + messages: Any = None, + *, + stream: bool = False, + session: AgentSession | None = None, + **kwargs: Any, + ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: + assert stream is True + assert session is not None + observed_child_states.append(dict(session.state)) + session.state["counter"] += 1 + if parent_declares_provider_key: + session.state["parent_provider_session"] = "child-handle" + + async def updates() -> Any: + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text("Child completed.")]) + + return ResponseStream( + updates(), + finalizer=lambda _: AgentResponse( + messages=[Message(role="assistant", contents=[Content.from_text("Child completed.")])] + ), + ) + + mock_session.send_and_wait.return_value = assistant_message_event + child = ChildAgent(name="ChildAgent") + + class ParentAgent(GitHubCopilotAgent): + service_session_state_keys = ( + frozenset({"parent_provider_session"}) if parent_declares_provider_key else frozenset() + ) + + parent = ParentAgent( + client=mock_client, + tools=[child.as_tool(name="delegate", propagate_session=True)], + context_providers=[StateProvider()], + ) + session = parent.create_session() + + await parent.run("Hello", session=session) + copilot_tool = mock_client.create_session.call_args.kwargs["tools"][0] + result = await copilot_tool.handler(ToolInvocation(arguments={"task": "Run child"})) + + assert result.result_type == "success" + assert len(observed_child_states) == 1 + assert observed_child_states[0]["provider-state"] == {"ordinary": "parent-value"} + assert "parent_provider_session" not in observed_child_states[0] + assert session.state["counter"] == 1 + if parent_declares_provider_key: + assert session.state["parent_provider_session"] == "parent-handle" + else: + assert "parent_provider_session" not in session.state + async def test_run_streaming_forwards_function_invocation_kwargs_to_tools( self, mock_client: MagicMock,