From 84451e472dbfc4c74fa48dfabdc2aa3f4a239f1b Mon Sep 17 00:00:00 2001 From: Giles Odigwe Date: Wed, 29 Jul 2026 11:23:40 -0700 Subject: [PATCH 1/4] Python: Fix ClaudeAgent reusing one SDK client across distinct fresh sessions RawClaudeAgent kept a single mutable ClaudeSDKClient on the agent instance and reused it across distinct fresh AgentSession objects, because a fresh session passes session_id=None and the old reuse check treated that as "keep the current client". Two independent fresh sessions on one shared agent instance therefore shared a single provider conversation, so the second session continued the first session's conversation. Treat a fresh (None) continuation id as always requiring a new client, so an unbound session never inherits an existing provider conversation. Legitimate continuity is preserved: once a session runs, its service_session_id is written back, so later runs pass a real id and resume correctly. Guard client selection/creation with an asyncio.Lock so concurrent runs cannot race between the check and the client assignment. Add regression tests asserting two fresh sessions produce two clients and that an explicit continuation id still resumes the existing client. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 598a9fe1-28c5-4db1-88fd-e14acd9340af --- .../claude/agent_framework_claude/_agent.py | 32 ++++++++--- .../claude/tests/test_claude_agent.py | 54 +++++++++++++++++++ 2 files changed, 79 insertions(+), 7 deletions(-) diff --git a/python/packages/claude/agent_framework_claude/_agent.py b/python/packages/claude/agent_framework_claude/_agent.py index 7db8eb7dbfd..ca7441a0e6b 100644 --- a/python/packages/claude/agent_framework_claude/_agent.py +++ b/python/packages/claude/agent_framework_claude/_agent.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import contextlib import inspect import logging @@ -387,6 +388,7 @@ def __init__( self._started = False self._current_session_id: str | None = None self._structured_output: Any = None + self._session_lock = asyncio.Lock() def _normalize_tools( self, @@ -449,21 +451,37 @@ async def stop(self) -> None: async def _ensure_session(self, session_id: str | None = None) -> None: """Ensure the client is connected for the specified session. - If the requested session differs from the current one, recreates the client. + A ``ClaudeSDKClient`` is stateful and represents a single provider + conversation. Deciding whether to reuse it is therefore an isolation + decision, not just a connection optimization: a fresh session + (``session_id is None``) must never inherit the provider conversation of + a previously started client, otherwise independent sessions running + against the same agent instance would share conversation state. A new + client is created when there is no started client, when a fresh session + is requested, or when an explicit continuation id differs from the + currently connected one. Args: - session_id: The session ID to use, or None for a new session. + session_id: The provider continuation id to resume, or None for a + fresh session that must get its own client. """ - needs_new_client = ( - not self._started or self._client is None or (session_id and session_id != self._current_session_id) - ) + async with self._session_lock: + needs_new_client = ( + not self._started + or self._client is None + or session_id is None + or session_id != self._current_session_id + ) + + if not needs_new_client: + return - if needs_new_client: # Stop existing client if any if self._client and self._owns_client: with contextlib.suppress(Exception): await self._client.disconnect() - self._started = False + self._started = False + self._current_session_id = None # Create new client with resume option if needed opts = self._prepare_client_options(resume_session_id=session_id) diff --git a/python/packages/claude/tests/test_claude_agent.py b/python/packages/claude/tests/test_claude_agent.py index 116d4114694..0112c7bb4d1 100644 --- a/python/packages/claude/tests/test_claude_agent.py +++ b/python/packages/claude/tests/test_claude_agent.py @@ -607,6 +607,60 @@ async def test_ensure_session_reuses_for_same_session(self) -> None: # Only called once assert mock_client_class.call_count == 1 + async def test_ensure_session_creates_new_client_for_each_fresh_session(self) -> None: + """Two fresh (None) sessions must each get their own client. + + A fresh session must never inherit a previously started client's + provider conversation, otherwise independent sessions running against + the same agent instance would share conversation state. + """ + with patch("agent_framework_claude._agent.ClaudeSDKClient") as mock_client_class: + mock_client1 = MagicMock() + mock_client1.connect = AsyncMock() + mock_client1.disconnect = AsyncMock() + + mock_client2 = MagicMock() + mock_client2.connect = AsyncMock() + mock_client2.disconnect = AsyncMock() + + mock_client_class.side_effect = [mock_client1, mock_client2] + + agent = ClaudeAgent() + + # First fresh session starts a client. + await agent._ensure_session(None) # type: ignore[reportPrivateUsage] + + # A second, independent fresh session must not reuse the first client. + await agent._ensure_session(None) # type: ignore[reportPrivateUsage] + + assert mock_client_class.call_count == 2 + mock_client1.disconnect.assert_called_once() + assert agent._client is mock_client2 # type: ignore[reportPrivateUsage] + + async def test_ensure_session_resumes_bound_session(self) -> None: + """A run with a continuation id resumes rather than recreating the client.""" + with patch("agent_framework_claude._agent.ClaudeSDKClient") as mock_client_class: + mock_client1 = MagicMock() + mock_client1.connect = AsyncMock() + mock_client1.disconnect = AsyncMock() + + mock_client2 = MagicMock() + mock_client2.connect = AsyncMock() + + mock_client_class.side_effect = [mock_client1, mock_client2] + + agent = ClaudeAgent() + + # Fresh session starts a client and becomes bound to a provider id. + await agent._ensure_session(None) # type: ignore[reportPrivateUsage] + agent._current_session_id = "provider-session-1" # type: ignore[reportPrivateUsage] + + # Re-running with the same continuation id reuses the client. + await agent._ensure_session("provider-session-1") # type: ignore[reportPrivateUsage] + + assert mock_client_class.call_count == 1 + mock_client1.disconnect.assert_not_called() + # region Test ClaudeAgent Tool Conversion From 3be125e29b9c0f51eafa2530e135744783519a7b Mon Sep 17 00:00:00 2001 From: Giles Odigwe Date: Wed, 29 Jul 2026 14:49:21 -0700 Subject: [PATCH 2/4] Python: Bind Claude SDK client ownership to each run Replace the single mutable ClaudeSDKClient stored on the agent with a per-run client. Because a ClaudeSDKClient represents exactly one provider conversation, sharing one across distinct sessions collapsed them onto the same conversation and, for concurrent runs, let a fresh session disconnect a client another run was still streaming from. _acquire_client now returns a per-run client (owned) that resumes the framework session's provider conversation when one exists, and _get_stream releases it in a finally once the run completes. An injected client is reused verbatim and left to the caller. The streaming loop moves into _stream_run so the client is a local per-run value rather than shared agent state, which keeps distinct sessions isolated even under concurrency. Continuity is preserved: a session's service_session_id is written back after each run and forwarded as the resume id on subsequent runs. Replace the client-lifecycle tests with per-run ownership and end-to-end isolation tests (two fresh sessions get two separate clients, each disconnected). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 598a9fe1-28c5-4db1-88fd-e14acd9340af --- .../claude/agent_framework_claude/_agent.py | 153 ++++++++------- .../claude/tests/test_claude_agent.py | 181 ++++++++---------- 2 files changed, 166 insertions(+), 168 deletions(-) diff --git a/python/packages/claude/agent_framework_claude/_agent.py b/python/packages/claude/agent_framework_claude/_agent.py index ca7441a0e6b..89b21964743 100644 --- a/python/packages/claude/agent_framework_claude/_agent.py +++ b/python/packages/claude/agent_framework_claude/_agent.py @@ -2,7 +2,6 @@ from __future__ import annotations -import asyncio import contextlib import inspect import logging @@ -386,9 +385,7 @@ def __init__( self._default_options = opts self._started = False - self._current_session_id: str | None = None self._structured_output: Any = None - self._session_lock = asyncio.Lock() def _normalize_tools( self, @@ -426,75 +423,78 @@ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: async def start(self) -> None: """Start the Claude SDK client. - This method initializes the Claude SDK client and establishes a connection - to the Claude Code CLI. It is called automatically when using the agent - as an async context manager. + Owned clients are created per run so that distinct sessions stay isolated; + this only needs to establish a connection for a pre-configured client that + was injected at construction. It is called automatically when using the + agent as an async context manager. Raises: AgentException: If the client fails to start. """ - await self._ensure_session() + if self._client is not None and not self._owns_client and not self._started: + try: + await self._client.connect() + self._started = True + except Exception as ex: + raise AgentException(f"Failed to start Claude SDK client: {ex}") from ex async def stop(self) -> None: """Stop the Claude SDK client and clean up resources. - Stops the client if owned by this agent. Called automatically when - using the agent as an async context manager. + Per-run owned clients are disconnected at the end of each run, so this + only disconnects a long-lived client the agent still owns. A client that + was injected at construction is owned by the caller and left untouched. + Called automatically when using the agent as an async context manager. """ - if self._client and self._owns_client: + if self._client is not None and self._owns_client: with contextlib.suppress(Exception): await self._client.disconnect() self._started = False - self._current_session_id = None - async def _ensure_session(self, session_id: str | None = None) -> None: - """Ensure the client is connected for the specified session. + async def _acquire_client(self, resume_session_id: str | None = None) -> tuple[ClaudeSDKClient, bool]: + """Acquire a Claude SDK client for a single run. - A ``ClaudeSDKClient`` is stateful and represents a single provider - conversation. Deciding whether to reuse it is therefore an isolation - decision, not just a connection optimization: a fresh session - (``session_id is None``) must never inherit the provider conversation of - a previously started client, otherwise independent sessions running - against the same agent instance would share conversation state. A new - client is created when there is no started client, when a fresh session - is requested, or when an explicit continuation id differs from the - currently connected one. + A ``ClaudeSDKClient`` is stateful and represents exactly one provider + conversation, so obtaining it is an isolation decision, not just a + connection optimization. When a client was injected at construction, that + single client is reused for every run and its lifecycle is owned by the + caller. Otherwise a fresh client scoped to this run is created and + connected, resuming ``resume_session_id`` when the framework session + already carries a provider conversation id. Binding the client to the run + rather than to shared agent state keeps distinct sessions isolated even + when they run concurrently against the same agent instance. Args: - session_id: The provider continuation id to resume, or None for a - fresh session that must get its own client. - """ - async with self._session_lock: - needs_new_client = ( - not self._started - or self._client is None - or session_id is None - or session_id != self._current_session_id - ) - - if not needs_new_client: - return - - # Stop existing client if any - if self._client and self._owns_client: - with contextlib.suppress(Exception): - await self._client.disconnect() - self._started = False - self._current_session_id = None + resume_session_id: The provider continuation id to resume, or None for + a fresh conversation. - # Create new client with resume option if needed - opts = self._prepare_client_options(resume_session_id=session_id) - self._client = ClaudeSDKClient(options=opts) - self._owns_client = True + Returns: + A tuple of the client and whether the caller owns it and must + disconnect it when the run completes. An injected client is never + owned by the caller of this method. - try: - await self._client.connect() - self._started = True - self._current_session_id = session_id - except Exception as ex: - self._client = None - raise AgentException(f"Failed to start Claude SDK client: {ex}") from ex + Raises: + AgentException: If the client fails to connect. + """ + if self._client is not None and not self._owns_client: + # Injected client: a single shared conversation with a caller-managed + # lifecycle. Connect it once and reuse it verbatim. + if not self._started: + try: + await self._client.connect() + self._started = True + except Exception as ex: + raise AgentException(f"Failed to start Claude SDK client: {ex}") from ex + return self._client, False + + opts = self._prepare_client_options(resume_session_id=resume_session_id) + client = ClaudeSDKClient(options=opts) + try: + await client.connect() + except Exception as ex: + raise AgentException(f"Failed to start Claude SDK client: {ex}") from ex + return client, True def _prepare_client_options(self, resume_session_id: str | None = None) -> SDKOptions: """Prepare SDK options for client initialization. @@ -653,15 +653,16 @@ async def handler(args: dict[str, Any]) -> dict[str, Any]: handler=handler, ) - async def _apply_runtime_options(self, options: dict[str, Any] | None) -> None: + async def _apply_runtime_options(self, client: ClaudeSDKClient, options: dict[str, Any] | None) -> None: """Apply runtime options that can be changed dynamically. The Claude SDK supports changing model and permission_mode after connection. Args: + client: The per-run client to apply the options to. options: Runtime options to apply. """ - if not options or not self._client: + if not options: return if "on_function_approval" in options: @@ -672,10 +673,10 @@ async def _apply_runtime_options(self, options: dict[str, Any] | None) -> None: ) if "model" in options: - await self._client.set_model(options["model"]) + await client.set_model(options["model"]) if "permission_mode" in options: - await self._client.set_permission_mode(options["permission_mode"]) + await client.set_permission_mode(options["permission_mode"]) def _format_prompt(self, messages: list[Message] | None) -> str: """Format messages into a prompt string. @@ -783,22 +784,46 @@ async def _get_stream( """Internal streaming implementation.""" session = session or self.create_session() - # Ensure we're connected to the right session - await self._ensure_session(self._get_chat_conversation_id(session)) + # A ClaudeSDKClient represents a single provider conversation, so each run + # acquires its own client (resuming the framework session's provider + # conversation when one exists) and releases it when the run completes. + # Binding the client to the run keeps distinct sessions isolated even when + # they run concurrently against the same agent instance. + client, owns_client = await self._acquire_client(self._get_chat_conversation_id(session)) + try: + async for update in self._stream_run(client, session, messages, options): + yield update + finally: + if owns_client: + with contextlib.suppress(Exception): + await client.disconnect() - if not self._client: - raise RuntimeError("Claude SDK client not initialized.") + async def _stream_run( + self, + client: ClaudeSDKClient, + session: AgentSession, + messages: AgentRunInputs | None, + options: OptionsT | None, + ) -> AsyncIterable[AgentResponseUpdate]: + """Run a single query against ``client`` and stream response updates. + Args: + client: The per-run Claude SDK client to query. + session: The active session; its ``service_session_id`` is updated with + the provider conversation id when the run completes. + messages: The input messages for this run. + options: Runtime options (model, permission_mode) for this run. + """ prompt = self._format_prompt(normalize_messages(messages)) # Apply runtime options (model, permission_mode) - await self._apply_runtime_options(dict(options) if options else None) + await self._apply_runtime_options(client, dict(options) if options else None) session_id: str | None = None structured_output: Any = None - await self._client.query(prompt) - async for message in self._client.receive_response(): + await client.query(prompt) + async for message in client.receive_response(): if isinstance(message, StreamEvent): # Handle streaming events - extract text/thinking deltas event = message.event diff --git a/python/packages/claude/tests/test_claude_agent.py b/python/packages/claude/tests/test_claude_agent.py index 0112c7bb4d1..87d77ce482f 100644 --- a/python/packages/claude/tests/test_claude_agent.py +++ b/python/packages/claude/tests/test_claude_agent.py @@ -553,113 +553,90 @@ def test_create_session_with_service_session_id(self) -> None: session = agent.create_session(session_id="existing-session-123") assert isinstance(session, AgentSession) - async def test_ensure_session_creates_client(self) -> None: - """Test _ensure_session creates client when not started.""" - with patch("agent_framework_claude._agent.ClaudeSDKClient") as mock_client_class: - mock_client = MagicMock() - mock_client.connect = AsyncMock() - mock_client_class.return_value = mock_client - - agent = ClaudeAgent() - await agent._ensure_session(None) # type: ignore[reportPrivateUsage] - - assert agent._started # type: ignore[reportPrivateUsage] - mock_client.connect.assert_called_once() - - async def test_ensure_session_recreates_for_different_session(self) -> None: - """Test _ensure_session recreates client for different session ID.""" - with patch("agent_framework_claude._agent.ClaudeSDKClient") as mock_client_class: - mock_client1 = MagicMock() - mock_client1.connect = AsyncMock() - mock_client1.disconnect = AsyncMock() - - mock_client2 = MagicMock() - mock_client2.connect = AsyncMock() + @staticmethod + async def _create_async_generator(items: list[Any]) -> Any: + """Yield the given items as an async generator (a mock provider response).""" + for item in items: + yield item - mock_client_class.side_effect = [mock_client1, mock_client2] + def _make_mock_client(self) -> MagicMock: + """Build a mock ClaudeSDKClient exposing the methods a run exercises.""" + mock_client = MagicMock() + mock_client.connect = AsyncMock() + mock_client.disconnect = AsyncMock() + mock_client.query = AsyncMock() + mock_client.set_model = AsyncMock() + mock_client.set_permission_mode = AsyncMock() + mock_client.receive_response = MagicMock(return_value=self._create_async_generator([])) + return mock_client + async def test_acquire_client_creates_fresh_owned_client(self) -> None: + """A fresh run creates and owns a newly connected client.""" + mock_client = self._make_mock_client() + with patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client): agent = ClaudeAgent() + client, owns_client = await agent._acquire_client(None) # type: ignore[reportPrivateUsage] - # First session - await agent._ensure_session(None) # type: ignore[reportPrivateUsage] - assert agent._started # type: ignore[reportPrivateUsage] - - # Different session should recreate client - await agent._ensure_session("new-session-id") # type: ignore[reportPrivateUsage] - assert agent._current_session_id == "new-session-id" # type: ignore[reportPrivateUsage] - mock_client1.disconnect.assert_called_once() - - async def test_ensure_session_reuses_for_same_session(self) -> None: - """Test _ensure_session reuses client for same session ID.""" - with patch("agent_framework_claude._agent.ClaudeSDKClient") as mock_client_class: - mock_client = MagicMock() - mock_client.connect = AsyncMock() - mock_client_class.return_value = mock_client + assert client is mock_client + assert owns_client is True + mock_client.connect.assert_awaited_once() + async def test_acquire_client_forwards_resume_id(self) -> None: + """An existing provider conversation id is forwarded as the resume id.""" + mock_client = self._make_mock_client() + with patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client): agent = ClaudeAgent() - - # First call - await agent._ensure_session("session-123") # type: ignore[reportPrivateUsage] - - # Same session should not recreate - await agent._ensure_session("session-123") # type: ignore[reportPrivateUsage] - - # Only called once - assert mock_client_class.call_count == 1 - - async def test_ensure_session_creates_new_client_for_each_fresh_session(self) -> None: - """Two fresh (None) sessions must each get their own client. - - A fresh session must never inherit a previously started client's - provider conversation, otherwise independent sessions running against - the same agent instance would share conversation state. - """ - with patch("agent_framework_claude._agent.ClaudeSDKClient") as mock_client_class: - mock_client1 = MagicMock() - mock_client1.connect = AsyncMock() - mock_client1.disconnect = AsyncMock() - - mock_client2 = MagicMock() - mock_client2.connect = AsyncMock() - mock_client2.disconnect = AsyncMock() - - mock_client_class.side_effect = [mock_client1, mock_client2] - + with patch.object( + agent, + "_prepare_client_options", + wraps=agent._prepare_client_options, # type: ignore[reportPrivateUsage] + ) as prepare: + _, owns_client = await agent._acquire_client("provider-session-1") # type: ignore[reportPrivateUsage] + + assert owns_client is True + prepare.assert_called_once_with(resume_session_id="provider-session-1") + + async def test_acquire_client_reuses_injected_client(self) -> None: + """An injected client is reused across runs and never owned by the run.""" + injected = self._make_mock_client() + agent = ClaudeAgent(client=injected) + + client_a, owns_a = await agent._acquire_client(None) # type: ignore[reportPrivateUsage] + client_b, owns_b = await agent._acquire_client("session-123") # type: ignore[reportPrivateUsage] + + assert client_a is injected + assert client_b is injected + assert owns_a is False + assert owns_b is False + # Connected once and reused; never disconnected by the run. + injected.connect.assert_awaited_once() + injected.disconnect.assert_not_called() + + async def test_two_fresh_sessions_use_separate_clients(self) -> None: + """Two distinct fresh sessions on one agent get isolated, separately-owned clients.""" + mock_client1 = self._make_mock_client() + mock_client2 = self._make_mock_client() + with patch( + "agent_framework_claude._agent.ClaudeSDKClient", + side_effect=[mock_client1, mock_client2], + ) as mock_client_class: agent = ClaudeAgent() + await agent.run("hello", session=agent.create_session()) + await agent.run("hello", session=agent.create_session()) - # First fresh session starts a client. - await agent._ensure_session(None) # type: ignore[reportPrivateUsage] - - # A second, independent fresh session must not reuse the first client. - await agent._ensure_session(None) # type: ignore[reportPrivateUsage] - - assert mock_client_class.call_count == 2 - mock_client1.disconnect.assert_called_once() - assert agent._client is mock_client2 # type: ignore[reportPrivateUsage] - - async def test_ensure_session_resumes_bound_session(self) -> None: - """A run with a continuation id resumes rather than recreating the client.""" - with patch("agent_framework_claude._agent.ClaudeSDKClient") as mock_client_class: - mock_client1 = MagicMock() - mock_client1.connect = AsyncMock() - mock_client1.disconnect = AsyncMock() - - mock_client2 = MagicMock() - mock_client2.connect = AsyncMock() - - mock_client_class.side_effect = [mock_client1, mock_client2] + assert mock_client_class.call_count == 2 + # Each per-run client is released when its run completes. + mock_client1.disconnect.assert_awaited_once() + mock_client2.disconnect.assert_awaited_once() + async def test_owned_client_disconnected_after_run(self) -> None: + """A per-run owned client is disconnected when the run finishes.""" + mock_client = self._make_mock_client() + with patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client): agent = ClaudeAgent() + await agent.run("hello") - # Fresh session starts a client and becomes bound to a provider id. - await agent._ensure_session(None) # type: ignore[reportPrivateUsage] - agent._current_session_id = "provider-session-1" # type: ignore[reportPrivateUsage] - - # Re-running with the same continuation id reuses the client. - await agent._ensure_session("provider-session-1") # type: ignore[reportPrivateUsage] - - assert mock_client_class.call_count == 1 - mock_client1.disconnect.assert_not_called() + mock_client.disconnect.assert_awaited_once() # region Test ClaudeAgent Tool Conversion @@ -1046,9 +1023,8 @@ async def test_apply_runtime_model(self) -> None: mock_client.set_permission_mode = AsyncMock() agent = ClaudeAgent() - agent._client = mock_client # type: ignore[reportPrivateUsage] - await agent._apply_runtime_options({"model": "opus"}) # type: ignore[reportPrivateUsage] + await agent._apply_runtime_options(mock_client, {"model": "opus"}) # type: ignore[reportPrivateUsage] mock_client.set_model.assert_called_once_with("opus") async def test_apply_runtime_permission_mode(self) -> None: @@ -1058,9 +1034,8 @@ async def test_apply_runtime_permission_mode(self) -> None: mock_client.set_permission_mode = AsyncMock() agent = ClaudeAgent() - agent._client = mock_client # type: ignore[reportPrivateUsage] - await agent._apply_runtime_options({"permission_mode": "acceptEdits"}) # type: ignore[reportPrivateUsage] + await agent._apply_runtime_options(mock_client, {"permission_mode": "acceptEdits"}) # type: ignore[reportPrivateUsage] mock_client.set_permission_mode.assert_called_once_with("acceptEdits") async def test_apply_runtime_options_none(self) -> None: @@ -1070,9 +1045,8 @@ async def test_apply_runtime_options_none(self) -> None: mock_client.set_permission_mode = AsyncMock() agent = ClaudeAgent() - agent._client = mock_client # type: ignore[reportPrivateUsage] - await agent._apply_runtime_options(None) # type: ignore[reportPrivateUsage] + await agent._apply_runtime_options(mock_client, None) # type: ignore[reportPrivateUsage] mock_client.set_model.assert_not_called() mock_client.set_permission_mode.assert_not_called() @@ -1083,10 +1057,9 @@ async def test_apply_runtime_on_function_approval_rejected(self) -> None: mock_client.set_permission_mode = AsyncMock() agent = ClaudeAgent() - agent._client = mock_client # type: ignore[reportPrivateUsage] with pytest.raises(ValueError, match="on_function_approval"): - await agent._apply_runtime_options({"on_function_approval": lambda _c: True}) # type: ignore[reportPrivateUsage] + await agent._apply_runtime_options(mock_client, {"on_function_approval": lambda _c: True}) # type: ignore[reportPrivateUsage] mock_client.set_model.assert_not_called() mock_client.set_permission_mode.assert_not_called() From 02d0f2c4db72687758bd21baf102f0e879d3cce5 Mon Sep 17 00:00:00 2001 From: Giles Odigwe Date: Thu, 6 Aug 2026 12:08:42 -0700 Subject: [PATCH 3/4] Python: Close remaining Claude session-isolation gaps Address three shared-state gaps in the Claude adapter surfaced in review: - Run-scope structured output: carry the run's structured_output through a per-run state holder and a per-run finalizer instead of storing it on the agent, so a concurrent run cannot overwrite another run's value before its finalizer reads it. - Bind an injected client to one session: an injected ClaudeSDKClient is a single Claude conversation, so bind it to the first session that uses it and raise AgentInvalidRequestException if a different session tries to reuse it. A no-session run reuses the bound session so multi-turn continuity still works; multi-session callers must omit client= or use one agent per session. - Serialize the injected-client path with an asyncio.Lock so concurrent runs cannot race its connect or interleave queries on the one shared client. Owned per-run clients stay lock-free. Update and extend the tests accordingly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 598a9fe1-28c5-4db1-88fd-e14acd9340af --- .../claude/agent_framework_claude/_agent.py | 119 ++++++++++++------ .../claude/tests/test_claude_agent.py | 38 +++++- 2 files changed, 116 insertions(+), 41 deletions(-) diff --git a/python/packages/claude/agent_framework_claude/_agent.py b/python/packages/claude/agent_framework_claude/_agent.py index 6424db792f6..300755aa003 100644 --- a/python/packages/claude/agent_framework_claude/_agent.py +++ b/python/packages/claude/agent_framework_claude/_agent.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import contextlib import inspect import logging @@ -29,7 +30,7 @@ normalize_tools, ) from agent_framework._telemetry import mark_feature_used -from agent_framework.exceptions import AgentException +from agent_framework.exceptions import AgentException, AgentInvalidRequestException from agent_framework.observability import AgentTelemetryLayer from claude_agent_sdk import ( AssistantMessage, @@ -388,7 +389,11 @@ def __init__( self._default_options = opts self._started = False - self._structured_output: Any = None + # An injected client is a single Claude conversation; bind it to the first + # session that uses it and serialize access so distinct sessions cannot + # share it and concurrent runs cannot race its connection or interleave. + self._injected_session: AgentSession | None = None + self._client_lock = asyncio.Lock() def _normalize_tools( self, @@ -434,12 +439,9 @@ async def start(self) -> None: Raises: AgentException: If the client fails to start. """ - if self._client is not None and not self._owns_client and not self._started: - try: - await self._client.connect() - self._started = True - except Exception as ex: - raise AgentException(f"Failed to start Claude SDK client: {ex}") from ex + if self._client is not None and not self._owns_client: + async with self._client_lock: + await self._connect_injected_client() async def stop(self) -> None: """Stop the Claude SDK client and clean up resources. @@ -455,22 +457,37 @@ async def stop(self) -> None: self._started = False - async def _acquire_client(self, resume_session_id: str | None = None) -> tuple[ClaudeSDKClient, bool]: + async def _connect_injected_client(self) -> None: + """Connect the injected client once. Caller must hold ``_client_lock``. + + Raises: + AgentException: If the client fails to connect. + """ + if self._client is None or self._started: + return + try: + await self._client.connect() + self._started = True + except Exception as ex: + raise AgentException(f"Failed to start Claude SDK client: {ex}") from ex + + async def _acquire_client(self, session: AgentSession) -> tuple[ClaudeSDKClient, bool]: """Acquire a Claude SDK client for a single run. A ``ClaudeSDKClient`` is stateful and represents exactly one provider conversation, so obtaining it is an isolation decision, not just a connection optimization. When a client was injected at construction, that - single client is reused for every run and its lifecycle is owned by the - caller. Otherwise a fresh client scoped to this run is created and - connected, resuming ``resume_session_id`` when the framework session - already carries a provider conversation id. Binding the client to the run - rather than to shared agent state keeps distinct sessions isolated even - when they run concurrently against the same agent instance. + single conversation is bound to the first session that uses it and reused + for that session's runs; a different session is rejected because one + injected conversation cannot be shared across sessions without leaking + context. Otherwise a fresh client scoped to this run is created and + connected, resuming the session's provider conversation when it already + carries one. Binding the client to the run rather than to shared agent + state keeps distinct sessions isolated even when they run concurrently + against the same agent instance. Args: - resume_session_id: The provider continuation id to resume, or None for - a fresh conversation. + session: The active session for this run. Returns: A tuple of the client and whether the caller owns it and must @@ -479,19 +496,25 @@ async def _acquire_client(self, resume_session_id: str | None = None) -> tuple[C Raises: AgentException: If the client fails to connect. + AgentInvalidRequestException: If an injected client is reused with a + different session than the one it is bound to. """ if self._client is not None and not self._owns_client: - # Injected client: a single shared conversation with a caller-managed - # lifecycle. Connect it once and reuse it verbatim. - if not self._started: - try: - await self._client.connect() - self._started = True - except Exception as ex: - raise AgentException(f"Failed to start Claude SDK client: {ex}") from ex + # Injected client: a single caller-managed conversation. Bind it to + # the first session and refuse to let another session reuse it. + if self._injected_session is None: + self._injected_session = session + elif self._injected_session.session_id != session.session_id: + raise AgentInvalidRequestException( + "An injected ClaudeSDKClient represents a single Claude conversation and is " + "bound to one session; it cannot be reused with a different session. Omit " + "`client=` so each run gets its own isolated client, or use a separate " + "ClaudeAgent per session." + ) + await self._connect_injected_client() return self._client, False - opts = self._prepare_client_options(resume_session_id=resume_session_id) + opts = self._prepare_client_options(resume_session_id=self._get_chat_conversation_id(session)) client = ClaudeSDKClient(options=opts) try: await client.connect() @@ -708,16 +731,18 @@ def default_options(self) -> dict[str, Any]: opts["instructions"] = system_prompt return opts - def _finalize_response(self, updates: Sequence[AgentResponseUpdate]) -> AgentResponse[Any]: + def _finalize_response(self, updates: Sequence[AgentResponseUpdate], structured_output: Any) -> AgentResponse[Any]: """Build AgentResponse and propagate structured_output as value. Args: updates: The collected stream updates. + structured_output: The run-scoped structured output captured during the + run, propagated as the response value if present. Returns: An AgentResponse with structured_output set as value if present. """ - return AgentResponse.from_updates(updates, value=self._structured_output) + return AgentResponse.from_updates(updates, value=structured_output) @overload def run( @@ -768,9 +793,13 @@ def run( When stream=True: An ResponseStream for streaming updates. When stream=False: An Awaitable[AgentResponse] with the complete response. """ + # Structured output is scoped to this run so concurrent runs on a shared + # agent instance cannot overwrite each other's value before the finalizer + # reads it. + run_state: dict[str, Any] = {"structured_output": None} response = ResponseStream( - self._get_stream(messages, session=session, options=options), - finalizer=self._finalize_response, + self._get_stream(messages, session=session, options=options, run_state=run_state), + finalizer=lambda updates: self._finalize_response(updates, run_state["structured_output"]), ) if stream: @@ -783,18 +812,35 @@ async def _get_stream( *, session: AgentSession | None = None, options: OptionsT | None = None, + run_state: dict[str, Any] | None = None, ) -> AsyncIterable[AgentResponseUpdate]: """Internal streaming implementation.""" - session = session or self.create_session() + if run_state is None: + run_state = {"structured_output": None} + + # An injected client is a single caller-owned conversation: reuse the + # session it is bound to when the caller passes none (so multi-turn runs + # stay continuous), and serialize access so its connection is not raced + # and concurrent runs do not interleave on the one shared client. + if self._client is not None and not self._owns_client: + if session is None and self._injected_session is not None: + session = self._injected_session + session = session or self.create_session() + async with self._client_lock: + client, _ = await self._acquire_client(session) + async for update in self._stream_run(client, session, messages, options, run_state): + yield update + return # A ClaudeSDKClient represents a single provider conversation, so each run # acquires its own client (resuming the framework session's provider # conversation when one exists) and releases it when the run completes. # Binding the client to the run keeps distinct sessions isolated even when # they run concurrently against the same agent instance. - client, owns_client = await self._acquire_client(self._get_chat_conversation_id(session)) + session = session or self.create_session() + client, owns_client = await self._acquire_client(session) try: - async for update in self._stream_run(client, session, messages, options): + async for update in self._stream_run(client, session, messages, options, run_state): yield update finally: if owns_client: @@ -807,6 +853,7 @@ async def _stream_run( session: AgentSession, messages: AgentRunInputs | None, options: OptionsT | None, + run_state: dict[str, Any], ) -> AsyncIterable[AgentResponseUpdate]: """Run a single query against ``client`` and stream response updates. @@ -816,6 +863,8 @@ async def _stream_run( the provider conversation id when the run completes. messages: The input messages for this run. options: Runtime options (model, permission_mode) for this run. + run_state: Per-run state holder; the structured output is written here + for the run's finalizer instead of on shared agent state. """ prompt = self._format_prompt(normalize_messages(messages)) @@ -914,8 +963,8 @@ async def _stream_run( if session_id: session.service_session_id = session_id - # Store structured output for the finalizer - self._structured_output = structured_output + # Store structured output for the run's finalizer (run-scoped, not on self) + run_state["structured_output"] = structured_output class ClaudeAgent(AgentTelemetryLayer, RawClaudeAgent[OptionsT], Generic[OptionsT]): diff --git a/python/packages/claude/tests/test_claude_agent.py b/python/packages/claude/tests/test_claude_agent.py index 8d7f6f5fb94..3ed4422b0e7 100644 --- a/python/packages/claude/tests/test_claude_agent.py +++ b/python/packages/claude/tests/test_claude_agent.py @@ -6,6 +6,7 @@ import pytest from agent_framework import AgentResponseUpdate, AgentSession, Content, Message, tool from agent_framework._settings import load_settings +from agent_framework.exceptions import AgentInvalidRequestException from agent_framework_claude import ClaudeAgent, ClaudeAgentOptions, ClaudeAgentSettings from agent_framework_claude._agent import TOOLS_MCP_SERVER_NAME @@ -580,7 +581,7 @@ async def test_acquire_client_creates_fresh_owned_client(self) -> None: mock_client = self._make_mock_client() with patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client): agent = ClaudeAgent() - client, owns_client = await agent._acquire_client(None) # type: ignore[reportPrivateUsage] + client, owns_client = await agent._acquire_client(agent.create_session()) # type: ignore[reportPrivateUsage] assert client is mock_client assert owns_client is True @@ -591,23 +592,25 @@ async def test_acquire_client_forwards_resume_id(self) -> None: mock_client = self._make_mock_client() with patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client): agent = ClaudeAgent() + session = agent.get_session(service_session_id="provider-session-1") with patch.object( agent, "_prepare_client_options", wraps=agent._prepare_client_options, # type: ignore[reportPrivateUsage] ) as prepare: - _, owns_client = await agent._acquire_client("provider-session-1") # type: ignore[reportPrivateUsage] + _, owns_client = await agent._acquire_client(session) # type: ignore[reportPrivateUsage] assert owns_client is True prepare.assert_called_once_with(resume_session_id="provider-session-1") - async def test_acquire_client_reuses_injected_client(self) -> None: - """An injected client is reused across runs and never owned by the run.""" + async def test_acquire_client_reuses_injected_client_for_same_session(self) -> None: + """An injected client is reused across runs of one session and never owned by the run.""" injected = self._make_mock_client() agent = ClaudeAgent(client=injected) + session = agent.create_session() - client_a, owns_a = await agent._acquire_client(None) # type: ignore[reportPrivateUsage] - client_b, owns_b = await agent._acquire_client("session-123") # type: ignore[reportPrivateUsage] + client_a, owns_a = await agent._acquire_client(session) # type: ignore[reportPrivateUsage] + client_b, owns_b = await agent._acquire_client(session) # type: ignore[reportPrivateUsage] assert client_a is injected assert client_b is injected @@ -617,6 +620,29 @@ async def test_acquire_client_reuses_injected_client(self) -> None: injected.connect.assert_awaited_once() injected.disconnect.assert_not_called() + async def test_injected_client_rejects_second_session(self) -> None: + """An injected client is bound to one session and rejects a different one.""" + injected = self._make_mock_client() + agent = ClaudeAgent(client=injected) + + await agent._acquire_client(agent.create_session()) # type: ignore[reportPrivateUsage] + + with pytest.raises(AgentInvalidRequestException, match="single Claude conversation"): + await agent._acquire_client(agent.create_session()) # type: ignore[reportPrivateUsage] + + async def test_injected_client_reused_across_runs_without_session(self) -> None: + """No-session runs on an injected client share its one bound conversation.""" + injected = self._make_mock_client() + with patch("agent_framework_claude._agent.ClaudeSDKClient") as mock_client_class: + agent = ClaudeAgent(client=injected) + await agent.run("first") + await agent.run("second") + + # No new clients were constructed; the injected one served both runs. + mock_client_class.assert_not_called() + injected.connect.assert_awaited_once() + injected.disconnect.assert_not_called() + async def test_two_fresh_sessions_use_separate_clients(self) -> None: """Two distinct fresh sessions on one agent get isolated, separately-owned clients.""" mock_client1 = self._make_mock_client() From 6cc4bd17cf9baff78b13aa4a79ebdcec4b93a571 Mon Sep 17 00:00:00 2001 From: Giles Odigwe Date: Fri, 7 Aug 2026 13:22:17 -0700 Subject: [PATCH 4/4] Python: Bind injected Claude client on provider conversation identity Compare an injected client's binding on the session's service_session_id (the Claude conversation identity) rather than the framework-local session_id, falling back to session_id only when the incoming session has no provider id yet. A reconstructed session from get_session(service_session_id=...) carries a fresh session_id but the same provider conversation, so it now continues the bound conversation instead of raising. Sessions targeting a different conversation are still rejected. Add regression tests for reconstructed-same-conversation continuation and different-conversation rejection. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 598a9fe1-28c5-4db1-88fd-e14acd9340af --- .../claude/agent_framework_claude/_agent.py | 52 ++++++++++++++----- .../claude/tests/test_claude_agent.py | 31 +++++++++++ 2 files changed, 70 insertions(+), 13 deletions(-) diff --git a/python/packages/claude/agent_framework_claude/_agent.py b/python/packages/claude/agent_framework_claude/_agent.py index 300755aa003..e34f507c07f 100644 --- a/python/packages/claude/agent_framework_claude/_agent.py +++ b/python/packages/claude/agent_framework_claude/_agent.py @@ -478,13 +478,13 @@ async def _acquire_client(self, session: AgentSession) -> tuple[ClaudeSDKClient, conversation, so obtaining it is an isolation decision, not just a connection optimization. When a client was injected at construction, that single conversation is bound to the first session that uses it and reused - for that session's runs; a different session is rejected because one - injected conversation cannot be shared across sessions without leaking - context. Otherwise a fresh client scoped to this run is created and - connected, resuming the session's provider conversation when it already - carries one. Binding the client to the run rather than to shared agent - state keeps distinct sessions isolated even when they run concurrently - against the same agent instance. + only for that same conversation; a session that targets a different + conversation is rejected because one injected conversation cannot be + shared without leaking context. Otherwise a fresh client scoped to this + run is created and connected, resuming the session's provider conversation + when it already carries one. Binding the client to the run rather than to + shared agent state keeps distinct sessions isolated even when they run + concurrently against the same agent instance. Args: session: The active session for this run. @@ -497,19 +497,20 @@ async def _acquire_client(self, session: AgentSession) -> tuple[ClaudeSDKClient, Raises: AgentException: If the client fails to connect. AgentInvalidRequestException: If an injected client is reused with a - different session than the one it is bound to. + session that targets a different conversation than the one it is + bound to. """ if self._client is not None and not self._owns_client: # Injected client: a single caller-managed conversation. Bind it to - # the first session and refuse to let another session reuse it. + # the first session and refuse to let a different conversation reuse it. if self._injected_session is None: self._injected_session = session - elif self._injected_session.session_id != session.session_id: + elif not self._injected_session_matches(session): raise AgentInvalidRequestException( "An injected ClaudeSDKClient represents a single Claude conversation and is " - "bound to one session; it cannot be reused with a different session. Omit " - "`client=` so each run gets its own isolated client, or use a separate " - "ClaudeAgent per session." + "bound to one session; it cannot be reused with a session that targets a " + "different conversation. Omit `client=` so each run gets its own isolated " + "client, or use a separate ClaudeAgent per session." ) await self._connect_injected_client() return self._client, False @@ -522,6 +523,31 @@ async def _acquire_client(self, session: AgentSession) -> tuple[ClaudeSDKClient, raise AgentException(f"Failed to start Claude SDK client: {ex}") from ex return client, True + def _injected_session_matches(self, session: AgentSession) -> bool: + """Whether ``session`` targets the conversation the injected client is bound to. + + The isolation boundary for an injected client is the Claude conversation, + identified by ``service_session_id`` -- not the framework-local + ``session_id``. A reconstructed session (for example one from + ``get_session(service_session_id=...)``) carries a fresh ``session_id`` but + the same provider conversation id, so it legitimately continues the bound + conversation. When the incoming session has no provider id yet, fall back + to the local ``session_id`` so two distinct unbound sessions cannot share + the client. + + Args: + session: The session requesting the injected client. + + Returns: + True if the session may reuse the injected client, False otherwise. + """ + bound = self._injected_session + if bound is None or bound is session: + return True + if session.service_session_id is not None: + return bound.service_session_id == session.service_session_id + return bound.session_id == session.session_id + def _prepare_client_options(self, resume_session_id: str | None = None) -> SDKOptions: """Prepare SDK options for client initialization. diff --git a/python/packages/claude/tests/test_claude_agent.py b/python/packages/claude/tests/test_claude_agent.py index 3ed4422b0e7..38359f9d152 100644 --- a/python/packages/claude/tests/test_claude_agent.py +++ b/python/packages/claude/tests/test_claude_agent.py @@ -630,6 +630,37 @@ async def test_injected_client_rejects_second_session(self) -> None: with pytest.raises(AgentInvalidRequestException, match="single Claude conversation"): await agent._acquire_client(agent.create_session()) # type: ignore[reportPrivateUsage] + async def test_injected_client_allows_reconstructed_same_conversation(self) -> None: + """A reconstructed session with the same provider id may reuse the injected client.""" + injected = self._make_mock_client() + agent = ClaudeAgent(client=injected) + + # First run binds the injected client and the session gains a provider id. + first = agent.create_session() + await agent._acquire_client(first) # type: ignore[reportPrivateUsage] + first.service_session_id = "provider-conversation-1" + + # A new AgentSession (fresh session_id) that targets the same conversation + # continues rather than raising. + restored = agent.get_session(service_session_id="provider-conversation-1") + assert restored.session_id != first.session_id + client, owns = await agent._acquire_client(restored) # type: ignore[reportPrivateUsage] + assert client is injected + assert owns is False + + async def test_injected_client_rejects_different_conversation(self) -> None: + """A session targeting a different provider conversation is rejected.""" + injected = self._make_mock_client() + agent = ClaudeAgent(client=injected) + + first = agent.create_session() + await agent._acquire_client(first) # type: ignore[reportPrivateUsage] + first.service_session_id = "provider-conversation-1" + + other = agent.get_session(service_session_id="provider-conversation-2") + with pytest.raises(AgentInvalidRequestException, match="single Claude conversation"): + await agent._acquire_client(other) # type: ignore[reportPrivateUsage] + async def test_injected_client_reused_across_runs_without_session(self) -> None: """No-session runs on an injected client share its one bound conversation.""" injected = self._make_mock_client()