diff --git a/python/packages/claude/agent_framework_claude/_agent.py b/python/packages/claude/agent_framework_claude/_agent.py index 724c3575ce..e34f507c07 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,8 +389,11 @@ def __init__( self._default_options = opts self._started = False - self._current_session_id: str | None = None - 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, @@ -427,59 +431,122 @@ 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: + async with self._client_lock: + await self._connect_injected_client() 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 _connect_injected_client(self) -> None: + """Connect the injected client once. Caller must hold ``_client_lock``. - If the requested session differs from the current one, recreates the client. + 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 conversation is bound to the first session that uses it and reused + 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_id: The session ID to use, or None for a new session. - """ - needs_new_client = ( - not self._started or self._client is None or (session_id and session_id != self._current_session_id) - ) + session: The active session for this run. - 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 + 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. - # 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 + Raises: + AgentException: If the client fails to connect. + AgentInvalidRequestException: If an injected client is reused with a + 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 a different conversation reuse it. + if self._injected_session is None: + self._injected_session = session + 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 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 + + opts = self._prepare_client_options(resume_session_id=self._get_chat_conversation_id(session)) + 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 _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. - 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 + 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. @@ -638,15 +705,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: @@ -657,10 +725,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. @@ -689,16 +757,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( @@ -749,9 +819,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: @@ -764,27 +838,71 @@ 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 - # 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. + 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, run_state): + 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, + run_state: dict[str, Any], + ) -> 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. + 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)) # 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 mark_feature_used(FeatureIndex.CLAUDE) - 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 @@ -871,8 +989,8 @@ async def _get_stream( 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 4d4fe1fd03..38359f9d15 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 @@ -558,59 +559,146 @@ 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(agent.create_session()) # type: ignore[reportPrivateUsage] - # First session - await agent._ensure_session(None) # type: ignore[reportPrivateUsage] - assert agent._started # type: ignore[reportPrivateUsage] + assert client is mock_client + assert owns_client is True + mock_client.connect.assert_awaited_once() - # 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_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() + 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(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_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() - async def test_ensure_session_reuses_for_same_session(self) -> None: - """Test _ensure_session reuses client for same session ID.""" + 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 + 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_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_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() 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(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() + 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 call - await agent._ensure_session("session-123") # type: ignore[reportPrivateUsage] + 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() - # Same session should not recreate - await agent._ensure_session("session-123") # type: ignore[reportPrivateUsage] + 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") - # Only called once - assert mock_client_class.call_count == 1 + mock_client.disconnect.assert_awaited_once() # region Test ClaudeAgent Tool Conversion @@ -997,9 +1085,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: @@ -1009,9 +1096,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: @@ -1021,9 +1107,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() @@ -1034,10 +1119,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()