diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index fae0d03641..c6d20f85ab 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -631,10 +631,22 @@ async def _agent_wrapper(ctx: FunctionInvocationContext, **kwargs: Any) -> str: ctx: the function invocation context used **kwargs: only used to dynamically load the argument that is defined for this tool. """ + session = ctx.session if propagate_session else None + + # Create a child session that shares the parent's state dict but has + # an isolated service_session_id. This avoids mutating the parent + # session in-place, which would race under concurrent asyncio.gather + # tool invocations sharing the same session. + if session is not None: + child_session = AgentSession(session_id=session.session_id) + child_session.state = session.state # shared by reference + child_session.service_session_id = None + session = child_session + stream = self.run( str(kwargs.get(arg_name, "")), stream=True, - session=ctx.session if propagate_session else None, + session=session, function_invocation_kwargs=dict(ctx.kwargs), ) if stream_callback is not None: diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index 67cddeccd1..b443b1d2fc 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/python/packages/core/tests/core/test_agents.py @@ -1737,10 +1737,14 @@ def capturing_run(*args: Any, **kwargs: Any) -> Any: ) ) - assert captured_session is parent_session + # Child receives a separate AgentSession (not the parent object) to isolate + # service_session_id, but shares the same state dict and session_id. assert captured_session is not None + assert captured_session is not parent_session assert captured_session.session_id == "parent-session-123" + assert captured_session.state is parent_session.state assert captured_session.state["shared_key"] == "shared_value" + assert captured_session.service_session_id is None async def test_chat_agent_as_tool_propagate_session_false_by_default(client: SupportsChatGetResponse) -> None: @@ -1802,6 +1806,106 @@ def capturing_run(*args: Any, **kwargs: Any) -> Any: 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") + tool = agent.as_tool(propagate_session=True) + + parent_session = AgentSession(session_id="shared-session") + parent_session.service_session_id = "resp_parent_abc123" + parent_session.state["data"] = "shared" + + original_run = agent.run + captured_session = None + + def capturing_run(*args: Any, **kwargs: Any) -> Any: + nonlocal captured_session + captured_session = kwargs.get("session") + # The child gets a different session object with isolated service_session_id + assert captured_session is not None + assert captured_session is not parent_session + assert captured_session.service_session_id is None + # But shares the same state dict by reference + assert captured_session.state is parent_session.state + assert captured_session.state["data"] == "shared" + return original_run(*args, **kwargs) + + agent.run = capturing_run # type: ignore[assignment, method-assign] # ty: ignore[invalid-assignment] + + await tool.invoke( + context=FunctionInvocationContext( + function=tool, + arguments={"task": "Hello"}, + session=parent_session, + ) + ) + + # Parent's service_session_id is never mutated + assert parent_session.service_session_id == "resp_parent_abc123" + + +async def test_chat_agent_as_tool_propagate_session_restores_service_session_id_on_error( + client: SupportsChatGetResponse, +) -> None: + """Test that parent's service_session_id is untouched even if the child agent raises.""" + agent = Agent(client=client, name="SubAgent", description="Sub agent") + tool = agent.as_tool(propagate_session=True) + + parent_session = AgentSession(session_id="shared-session") + parent_session.service_session_id = "resp_parent_xyz789" + + def failing_run(*args: Any, **kwargs: Any) -> Any: + raise RuntimeError("Child agent failed") + + agent.run = failing_run # type: ignore[assignment, method-assign] # ty: ignore[invalid-assignment] + + with raises(RuntimeError, match="Child agent failed"): + await tool.invoke( + context=FunctionInvocationContext( + function=tool, + arguments={"task": "Hello"}, + session=parent_session, + ) + ) + + # Parent's service_session_id is never mutated — child has its own session + assert parent_session.service_session_id == "resp_parent_xyz789" + + +async def test_chat_agent_as_tool_propagate_session_no_service_session_id(client: SupportsChatGetResponse) -> None: + """Test that child setting service_session_id does not leak back to the parent.""" + agent = Agent(client=client, name="SubAgent", description="Sub agent") + tool = agent.as_tool(propagate_session=True) + + parent_session = AgentSession(session_id="shared-session") + parent_session.service_session_id = None + + original_run = agent.run + captured_session = None + + def capturing_run(*args: Any, **kwargs: Any) -> Any: + nonlocal captured_session + captured_session = kwargs.get("session") + assert captured_session is not None + assert captured_session.service_session_id is None + # Simulate the child's run populating service_session_id + captured_session.service_session_id = "resp_child_leaked" + return original_run(*args, **kwargs) + + agent.run = capturing_run # type: ignore[assignment, method-assign] # ty: ignore[invalid-assignment] + + await tool.invoke( + context=FunctionInvocationContext( + function=tool, + arguments={"task": "Hello"}, + session=parent_session, + ) + ) + + # The child's service_session_id must not leak back to the parent + assert parent_session.service_session_id is None + + async def test_chat_agent_as_mcp_server_basic(client: SupportsChatGetResponse) -> None: """Test basic as_mcp_server functionality.""" agent = Agent(client=client, name="TestAgent", description="Test agent for MCP")