From de44e9878e18dc7ca10c96385f3646a16c0932f5 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 27 Oct 2026 10:21:33 +0100 Subject: [PATCH 1/3] Python: Add async context managers to chat clients --- python/README.md | 39 ++++++++++++ .../_bedrock_client.py | 2 + .../agent_framework_anthropic/_chat_client.py | 10 +++ .../_foundry_client.py | 2 + .../_vertex_client.py | 2 + .../tests/test_anthropic_provider_clients.py | 36 +++++++++++ .../agent_framework_bedrock/_chat_client.py | 11 +++- .../bedrock/tests/test_bedrock_client.py | 14 +++++ python/packages/core/AGENTS.md | 5 +- .../packages/core/agent_framework/_clients.py | 29 +++++++++ .../packages/core/tests/core/test_agents.py | 29 +++++++++ .../packages/core/tests/core/test_clients.py | 50 ++++++++++++++- python/packages/foundry/README.md | 8 +++ .../foundry/agent_framework_foundry/_agent.py | 13 +++- .../agent_framework_foundry/_chat_client.py | 14 +++++ .../tests/foundry/test_foundry_agent.py | 62 ++++++++++++++++++- .../_foundry_local_client.py | 1 + .../tests/test_foundry_local_client.py | 19 ++++++ .../agent_framework_gemini/_chat_client.py | 12 ++++ .../gemini/tests/test_gemini_client.py | 31 ++++++++++ python/packages/mistral/AGENTS.md | 4 +- .../agent_framework_mistral/_chat_client.py | 11 +++- .../tests/mistral/test_mistral_chat_client.py | 44 ++++++++++++- .../agent_framework_ollama/_chat_client.py | 9 +++ .../ollama/tests/test_ollama_chat_client.py | 17 +++++ .../agent_framework_openai/_chat_client.py | 9 +++ .../_chat_completion_client.py | 9 +++ .../tests/openai/test_openai_chat_client.py | 42 ++++++++++++- .../references/integrations.md | 10 +++ 29 files changed, 527 insertions(+), 17 deletions(-) diff --git a/python/README.md b/python/README.md index 080c4d45bb8..97d059bc00a 100644 --- a/python/README.md +++ b/python/README.md @@ -158,6 +158,45 @@ async def main(): asyncio.run(main()) ``` +### Connection and credential lifetimes + +Chat clients support async context managers for deterministic cleanup: + +```python +from agent_framework import Agent, Message +from agent_framework.openai import OpenAIChatClient + +async with OpenAIChatClient() as client: + response = await client.get_response([Message("user", ["Hello!"])]) + +# Alternatively, let the agent enter and exit its chat client's context. +async with Agent(client=OpenAIChatClient(), instructions="Be helpful.") as agent: + response = await agent.run("Hello!") +``` + +Calling clients without a context manager remains supported; use `await client.close()` when finished. +Context exit releases connections created by the connector, including on exceptions. Caller-supplied SDK clients, +HTTP transports, and credentials remain caller-owned. In particular, closing an Azure SDK connection does not +close the credential supplied to it. Scope credentials outside the agent or client that uses them: + +```python +from azure.identity.aio import AzureCliCredential + +async with AzureCliCredential() as credential: + async with OpenAIChatClient( + azure_endpoint="https://YOUR-RESOURCE.openai.azure.com", + model="YOUR-DEPLOYMENT", + credential=credential, + ) as client: + response = await client.get_response([Message("user", ["Hello!"])]) +``` + +Treat closing a resource-owning client as the end of its lifetime; reopening is SDK-dependent. Avoid overlapping +agent contexts sharing a connector-owned connection, since exiting either agent closes that connection. To share +an SDK client across agents, inject it into their chat clients and manage its lifetime outside those agent contexts. +`BaseChatClient` provides default context methods and a no-op `close()`, so custom subclasses do not need to +implement lifecycle methods unless they own resources. Protocol-only chat clients do not need context methods. + ## 4. Build an Agent with Tools and Functions Enhance your agent with custom tools and function calling: diff --git a/python/packages/anthropic/agent_framework_anthropic/_bedrock_client.py b/python/packages/anthropic/agent_framework_anthropic/_bedrock_client.py index 3af207e20be..c46098f0bb9 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_bedrock_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_bedrock_client.py @@ -113,6 +113,7 @@ def __init__( secret_key_secret = settings.get("aws_secret_access_key") session_token_secret = settings.get("aws_session_token") + owns_client = anthropic_client is None if anthropic_client is None: anthropic_client = AsyncAnthropicBedrock( aws_secret_key=secret_key_secret.get_secret_value() if secret_key_secret is not None else None, @@ -130,6 +131,7 @@ def __init__( additional_beta_flags=additional_beta_flags, additional_properties=additional_properties, ) + self._owns_client = owns_client class AnthropicBedrockClient( diff --git a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py index dd894913609..9e4e1949369 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py @@ -418,6 +418,7 @@ class MyOptions(AnthropicChatOptions, total=False): model_setting = anthropic_settings.get("chat_model") base_url_setting = anthropic_settings.get("base_url") + owns_client = anthropic_client is None if anthropic_client is None: if api_key_secret is None: raise ValueError( @@ -438,9 +439,18 @@ class MyOptions(AnthropicChatOptions, total=False): # Initialize instance variables self.anthropic_client = anthropic_client + self._owns_client = owns_client + self._client_closed = False self.additional_beta_flags = additional_beta_flags or [] self.model = model_setting + @override + async def close(self) -> None: + """Close the internally created SDK client, leaving supplied credentials open.""" + if self._owns_client and not self._client_closed: + await self.anthropic_client.close() + self._client_closed = True + # region Static factory methods for hosted tools @staticmethod diff --git a/python/packages/anthropic/agent_framework_anthropic/_foundry_client.py b/python/packages/anthropic/agent_framework_anthropic/_foundry_client.py index 2b1ee0cd405..6660f574f46 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_foundry_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_foundry_client.py @@ -79,6 +79,7 @@ def __init__( base_url_setting = settings.get("anthropic_foundry_base_url") api_key_value = api_key_secret.get_secret_value() if api_key_secret is not None else None + owns_client = anthropic_client is None if anthropic_client is None: if base_url_setting is None and resource_setting is None: message = ( @@ -107,6 +108,7 @@ def __init__( additional_beta_flags=additional_beta_flags, additional_properties=additional_properties, ) + self._owns_client = owns_client class AnthropicFoundryClient( diff --git a/python/packages/anthropic/agent_framework_anthropic/_vertex_client.py b/python/packages/anthropic/agent_framework_anthropic/_vertex_client.py index 71e9b8ef69c..79e7a76da6f 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_vertex_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_vertex_client.py @@ -81,6 +81,7 @@ def __init__( region_setting = settings.get("cloud_ml_region") project_id_setting = settings.get("anthropic_vertex_project_id") + owns_client = anthropic_client is None if anthropic_client is None: resolved_region = region_setting if region_setting is not None else NOT_GIVEN resolved_project_id = project_id_setting if project_id_setting is not None else NOT_GIVEN @@ -102,6 +103,7 @@ def __init__( additional_beta_flags=additional_beta_flags, additional_properties=additional_properties, ) + self._owns_client = owns_client class AnthropicVertexClient( diff --git a/python/packages/anthropic/tests/test_anthropic_provider_clients.py b/python/packages/anthropic/tests/test_anthropic_provider_clients.py index 182091c8c9b..5c9477ea2e2 100644 --- a/python/packages/anthropic/tests/test_anthropic_provider_clients.py +++ b/python/packages/anthropic/tests/test_anthropic_provider_clients.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import logging from unittest.mock import AsyncMock, MagicMock, patch @@ -11,9 +13,11 @@ from agent_framework_anthropic import ( AnthropicBedrockClient, + AnthropicClient, AnthropicFoundryClient, AnthropicVertexClient, RawAnthropicBedrockClient, + RawAnthropicClient, RawAnthropicFoundryClient, RawAnthropicVertexClient, ) @@ -28,6 +32,38 @@ def _create_mock_transport(base_url: str) -> MagicMock: return transport +@pytest.mark.parametrize( + ("client_type", "factory", "kwargs"), + [ + (RawAnthropicClient, "_chat_client.AsyncAnthropic", {"api_key": "test-key"}), + (AnthropicClient, "_chat_client.AsyncAnthropic", {"api_key": "test-key"}), + (RawAnthropicFoundryClient, "_foundry_client.AsyncAnthropicFoundry", {"resource": "test", "api_key": "test"}), + (AnthropicFoundryClient, "_foundry_client.AsyncAnthropicFoundry", {"resource": "test", "api_key": "test"}), + (RawAnthropicBedrockClient, "_bedrock_client.AsyncAnthropicBedrock", {"aws_region": "us-east-1"}), + (AnthropicBedrockClient, "_bedrock_client.AsyncAnthropicBedrock", {"aws_region": "us-east-1"}), + ( + RawAnthropicVertexClient, + "_vertex_client.AsyncAnthropicVertex", + {"project_id": "test", "region": "us-central1"}, + ), + (AnthropicVertexClient, "_vertex_client.AsyncAnthropicVertex", {"project_id": "test", "region": "us-central1"}), + ], +) +@pytest.mark.parametrize("injected", [False, True]) +async def test_anthropic_client_context_ownership(client_type, factory: str, kwargs, injected: bool) -> None: + sdk_client = _create_mock_transport("https://test.example.com") + sdk_client.close = AsyncMock() + if injected: + kwargs = {**kwargs, "anthropic_client": sdk_client} + with patch(f"agent_framework_anthropic.{factory}", return_value=sdk_client): + client = client_type(model="test-model", **kwargs) + async with client as entered: + assert entered is client + sdk_client.close.assert_not_awaited() + await client.close() + assert sdk_client.close.await_count == (0 if injected else 1) + + @pytest.mark.parametrize( ("public_client", "raw_client"), [ diff --git a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py index 455c271ee54..b9f125f7bc8 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py +++ b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py @@ -326,7 +326,9 @@ class MyOptions(BedrockChatOptions, total=False): region = settings.get("region") or session_region or DEFAULT_REGION chat_model = settings.get("chat_model") - if client: + self._owns_client = client is None + self._client_closed = False + if client is not None: self._bedrock_client = client else: session = boto3_session or self._create_session(settings) @@ -344,6 +346,13 @@ class MyOptions(BedrockChatOptions, total=False): self.model = chat_model self.region = region + @override + async def close(self) -> None: + """Close the owned runtime connection without closing caller-owned sessions.""" + if self._owns_client and not self._client_closed: + await asyncio.to_thread(self._bedrock_client.close) + self._client_closed = True + @staticmethod def _create_session(settings: BedrockSettings) -> Boto3Session: session_kwargs: dict[str, Any] = {"region_name": settings.get("region") or DEFAULT_REGION} diff --git a/python/packages/bedrock/tests/test_bedrock_client.py b/python/packages/bedrock/tests/test_bedrock_client.py index 4cdb6bf89ea..fbe23f4847d 100644 --- a/python/packages/bedrock/tests/test_bedrock_client.py +++ b/python/packages/bedrock/tests/test_bedrock_client.py @@ -78,6 +78,20 @@ def test_agent_accepts_bedrock_chat_client() -> None: assert agent.client is client +@pytest.mark.parametrize("injected", [False, True]) +async def test_bedrock_client_context_ownership(injected: bool) -> None: + runtime = MagicMock(spec=BaseClient) + session = MagicMock(spec=Boto3Session) + session.region_name = "us-west-2" + session.client.return_value = runtime + client = BedrockChatClient(model="test-model", boto3_session=session, client=runtime if injected else None) + async with client as entered: + assert entered is client + runtime.close.assert_not_called() + await client.close() + assert runtime.close.call_count == (0 if injected else 1) + + async def test_get_response_invokes_bedrock_runtime() -> None: stub = _StubBedrockRuntime() client = BedrockChatClient( diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 895764eb93d..99746590285 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -67,7 +67,10 @@ agent_framework/ ### Chat Clients (`_clients.py`) - **`SupportsChatGetResponse`** - Protocol for chat client implementations -- **`BaseChatClient`** - Abstract base class with middleware support; subclasses implement `_inner_get_response()` and `_inner_get_streaming_response()` +- **`BaseChatClient`** - Abstract base class without middleware; subclasses implement `_inner_get_response()` + for streaming and non-streaming requests. Provides `__aenter__()` returning `self`, `__aexit__()` delegating to + `close()`, and a default no-op async `close()`. Connectors override cleanup for owned connections and credentials; + injected resources remain caller-owned. `SupportsChatGetResponse` does not require lifecycle methods. ### Types (`_types.py`) diff --git a/python/packages/core/agent_framework/_clients.py b/python/packages/core/agent_framework/_clients.py index 564d78aa660..bef2485eba8 100644 --- a/python/packages/core/agent_framework/_clients.py +++ b/python/packages/core/agent_framework/_clients.py @@ -12,6 +12,7 @@ Mapping, Sequence, ) +from types import TracebackType from typing import ( TYPE_CHECKING, Any, @@ -44,6 +45,11 @@ else: from typing_extensions import TypeVar # pragma: no cover +if sys.version_info >= (3, 11): + from typing import Self # pragma: no cover +else: + from typing_extensions import Self # pragma: no cover + if TYPE_CHECKING: from pydantic import BaseModel @@ -309,6 +315,29 @@ def __init__( self.tokenizer = tokenizer super().__init__() + async def __aenter__(self) -> Self: + """Enter the client context without requiring it for ordinary requests.""" + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + /, + ) -> None: + """Release resources owned by the client without suppressing exceptions.""" + await self.close() + + async def close(self) -> None: + """Release client-owned connections and credentials. + + The default implementation does nothing. Resource-owning implementations + override this method and must leave caller-supplied resources open. + Clients remain usable without entering an async context; callers can use + ``await client.close()`` to release their resources explicitly. + """ + def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: """Convert the instance to a dictionary. diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index 1481ce8a77b..de6b2c36793 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/python/packages/core/tests/core/test_agents.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import asyncio import contextlib import inspect @@ -220,6 +222,33 @@ async def test_agent_explicit_open_close(agent_type: type[RawAgent]) -> None: assert events == ["client.enter", "mcp.enter", "mcp.exit", "client.exit"] +@pytest.mark.parametrize("agent_type", [Agent, RawAgent]) +@pytest.mark.parametrize("fail_on_enter", [False, True]) +async def test_agent_manages_inherited_client_context( + agent_type: type[RawAgent], fail_on_enter: bool, chat_client_base: MockBaseChatClient +) -> None: + events: list[str] = [] + mcp_tool = _LifecycleMCPTool("mcp", events, fail_on_enter=fail_on_enter) + agent = agent_type(client=chat_client_base, tools=[mcp_tool]) + with patch.object(chat_client_base, "close", new_callable=AsyncMock) as close: + if fail_on_enter: + with pytest.raises(RuntimeError, match="MCP entry failed"): + await agent.open() + else: + async with agent: + close.assert_not_awaited() + close.assert_awaited_once_with() + await agent.close() + close.assert_awaited_once_with() + + +async def test_agent_still_accepts_client_without_lifecycle(client: SupportsChatGetResponse) -> None: + assert not hasattr(client, "__aenter__") + async with RawAgent(client=client) as agent: + response = await agent.run("Hello") + assert response.text == "test response" + + @pytest.mark.parametrize("agent_type", [Agent, RawAgent]) async def test_agent_async_with_delegates_to_open_close(agent_type: type[RawAgent]) -> None: events: list[str] = [] diff --git a/python/packages/core/tests/core/test_clients.py b/python/packages/core/tests/core/test_clients.py index b122b9c5ebf..8b59181aa28 100644 --- a/python/packages/core/tests/core/test_clients.py +++ b/python/packages/core/tests/core/test_clients.py @@ -1,9 +1,11 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations +import asyncio import logging from typing import Any -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest @@ -115,6 +117,52 @@ def test_base_client(chat_client_base: SupportsChatGetResponse): assert isinstance(chat_client_base, SupportsChatGetResponse) +async def test_base_client_default_lifecycle(chat_client_base: BaseChatClient) -> None: + await chat_client_base.close() + async with chat_client_base as entered: + assert entered is chat_client_base + response = await entered.get_response([Message(role="user", contents=["Hello"])]) + assert response.text == "test response - Hello" + updates = [update async for update in entered.get_response([Message("user", ["Hello"])], stream=True)] + assert updates + await chat_client_base.close() + assert (await chat_client_base.get_response([Message("user", ["Again"])])).text == "test response - Again" + + +@pytest.mark.parametrize("error", [None, RuntimeError("body failed"), asyncio.CancelledError()]) +async def test_base_client_context_calls_close(chat_client_base: BaseChatClient, error: BaseException | None) -> None: + with patch.object(chat_client_base, "close", new_callable=AsyncMock) as close: + if error is None: + async with chat_client_base: + close.assert_not_awaited() + else: + with pytest.raises(type(error)) as raised: + async with chat_client_base: + raise error + assert raised.value is error + close.assert_awaited_once_with() + + +async def test_base_client_context_propagates_cleanup_failure(chat_client_base: BaseChatClient) -> None: + with ( + patch.object(chat_client_base, "close", new_callable=AsyncMock, side_effect=RuntimeError("close failed")), + pytest.raises(RuntimeError, match="close failed"), + ): + async with chat_client_base: + pass + + +async def test_base_client_context_closes_after_early_stream_exit(chat_client_base: BaseChatClient) -> None: + with patch.object(chat_client_base, "close", new_callable=AsyncMock) as close: + async with chat_client_base: + async with chat_client_base.get_response([Message("user", ["Hello"])], stream=True) as stream: + async for update in stream: + assert update.text + break + close.assert_not_awaited() + close.assert_awaited_once_with() + + def test_base_client_rejects_direct_additional_properties(chat_client_base: SupportsChatGetResponse) -> None: with pytest.raises(TypeError): type(chat_client_base)(legacy_key="legacy-value") # type: ignore[call-arg] # pyrefly: ignore[bad-instantiation, unexpected-keyword] # ty: ignore[unknown-argument] diff --git a/python/packages/foundry/README.md b/python/packages/foundry/README.md index 7100d475182..f6be0f52062 100644 --- a/python/packages/foundry/README.md +++ b/python/packages/foundry/README.md @@ -8,6 +8,14 @@ This package supports `azure-ai-projects>=2.2.0,<2.8.0`. Projects 2.5 and later `openai>=3.0.0`, so `agent-framework-foundry` requires `agent-framework-openai>=1.14.2`, which supports both OpenAI 2.x and 3.x. +## Chat client lifecycle + +Use `async with FoundryChatClient(...)` or `async with FoundryAgent(...)` to release owned connections, or call +`await client.close()` explicitly when using a chat client without a context manager. Chat clients close their +derived OpenAI connection, including when a timeout-configured copy shares its transport. They also close the +project connection if they created it. A supplied `project_client` and credential remain caller-owned; manage the +credential's context outside the chat client or agent. Cleanup does not delete remote agents or conversations. + ## Tool ownership for preconfigured Foundry agents `FoundryAgent` connects to a PromptAgent or HostedAgent whose instructions, tool declarations, and tool-selection diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index 5b0c2e4bfa3..e6636d38355 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/python/packages/foundry/agent_framework_foundry/_agent.py @@ -292,6 +292,8 @@ def __init__( tokenizer=tokenizer, additional_properties=additional_properties, ) + self._owns_client = True + self._project_client_closed = False self.model = "" # Foundry agents resolve model server-side; ignore env vars (issue #7272). @override @@ -541,10 +543,15 @@ def _transform_input_for_azure_ai(self, input_items: list[dict[str, Any]]) -> li return transformed + @override async def close(self) -> None: - """Close the project client if we created it.""" - if self._should_close_client: - await self.project_client.close() + """Close owned OpenAI and project connections, not the supplied credential.""" + try: + await super().close() + finally: + if self._should_close_client and not self._project_client_closed: + await self.project_client.close() + self._project_client_closed = True class _FoundryAgentChatClient( diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index e41a700c1d2..f134729b7df 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -272,6 +272,20 @@ def __init__( additional_properties=additional_properties, ) self.project_client = project_client + self._owns_project_client = owns_project_client + self._project_client_closed = False + # get_openai_client creates an independent connection, even for an injected project. + self._owns_client = True + + @override + async def close(self) -> None: + """Close owned OpenAI and project connections, not the supplied credential.""" + try: + await super().close() + finally: + if self._owns_project_client and not self._project_client_closed: + await self.project_client.close() + self._project_client_closed = True @override async def _prepare_options( diff --git a/python/packages/foundry/tests/foundry/test_foundry_agent.py b/python/packages/foundry/tests/foundry/test_foundry_agent.py index 5bba9eb78c0..d7cb0913d8f 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_agent.py +++ b/python/packages/foundry/tests/foundry/test_foundry_agent.py @@ -49,7 +49,7 @@ RawFoundryAgentChatClient, _FoundryAgentChatClient, ) -from agent_framework_foundry._chat_client import FoundryChatClient +from agent_framework_foundry._chat_client import FoundryChatClient, RawFoundryChatClient from agent_framework_foundry._feature_usage import FeatureIndex, FeatureUsagePolicy _OPENAI_HTTPX = cast(Any, import_module(DefaultAsyncHttpxClient.__mro__[1].__module__.partition(".")[0])) @@ -74,6 +74,66 @@ def test_raw_foundry_agent_chat_client_does_not_mark_openai_feature() -> None: assert RawFoundryAgentChatClient._FEATURE_USAGE_INDEX is FeatureIndex.FOUNDRY_AGENT +@pytest.mark.parametrize( + "client_type", [RawFoundryChatClient, FoundryChatClient, RawFoundryAgentChatClient, _FoundryAgentChatClient] +) +@pytest.mark.parametrize("injected_project", [False, True]) +@pytest.mark.parametrize("timeout", [None, 60.0]) +async def test_foundry_chat_client_context_ownership( + client_type, injected_project: bool, timeout: float | None +) -> None: + sdk_client = AsyncOpenAI(api_key="test-key") + project = MagicMock() + project.get_openai_client.return_value = sdk_client + project.close = AsyncMock() + credential = MagicMock() + kwargs: dict[str, Any] = ( + {"agent_name": "test-agent"} if issubclass(client_type, RawFoundryAgentChatClient) else {"model": "test-model"} + ) + module = "_agent" if "agent_name" in kwargs else "_chat_client" + with patch(f"agent_framework_foundry.{module}.AIProjectClient", return_value=project): + client = client_type( + project_endpoint="https://test.services.ai.azure.com/api/projects/test", + project_client=project if injected_project else None, + credential=credential, + timeout=timeout, + **kwargs, + ) + with patch.object(client.client, "close", wraps=client.client.close) as close: + async with client as entered: + assert entered is client + assert not sdk_client.is_closed() + await client.close() + close.assert_awaited_once() + assert sdk_client.is_closed() + assert project.close.await_count == (0 if injected_project else 1) + credential.close.assert_not_called() + + +@pytest.mark.parametrize("client_type", [RawFoundryChatClient, RawFoundryAgentChatClient]) +async def test_foundry_chat_client_closes_project_when_openai_close_fails(client_type) -> None: + sdk_client = MagicMock(spec=AsyncOpenAI) + sdk_client.close = AsyncMock(side_effect=RuntimeError("close failed")) + project = MagicMock() + project.get_openai_client.return_value = sdk_client + project.close = AsyncMock() + kwargs: dict[str, Any] = ( + {"agent_name": "test-agent"} if issubclass(client_type, RawFoundryAgentChatClient) else {"model": "test-model"} + ) + module = "_agent" if "agent_name" in kwargs else "_chat_client" + with patch(f"agent_framework_foundry.{module}.AIProjectClient", return_value=project): + client = client_type( + project_endpoint="https://test.services.ai.azure.com/api/projects/test", + credential=MagicMock(), + **kwargs, + ) + with pytest.raises(RuntimeError, match="close failed"): + async with client: + pass + sdk_client.close.assert_awaited_once() + project.close.assert_awaited_once() + + def _get_foundry_azure_ai_search_model() -> str | None: """Return the model/deployment to use for local Azure AI Search integration validation.""" return next((os.environ[key] for key in _FOUNDRY_AZURE_AI_SEARCH_MODEL_ENV_VARS if os.getenv(key)), None) diff --git a/python/packages/foundry_local/agent_framework_foundry_local/_foundry_local_client.py b/python/packages/foundry_local/agent_framework_foundry_local/_foundry_local_client.py index a36647b8e35..e7a2018e427 100644 --- a/python/packages/foundry_local/agent_framework_foundry_local/_foundry_local_client.py +++ b/python/packages/foundry_local/agent_framework_foundry_local/_foundry_local_client.py @@ -349,4 +349,5 @@ class MyOptions(FoundryLocalChatOptions, total=False): middleware=middleware, function_invocation_configuration=function_invocation_configuration, ) + self._owns_client = True self.manager = manager diff --git a/python/packages/foundry_local/tests/test_foundry_local_client.py b/python/packages/foundry_local/tests/test_foundry_local_client.py index 2ba5e9d89c0..18bc70fe160 100644 --- a/python/packages/foundry_local/tests/test_foundry_local_client.py +++ b/python/packages/foundry_local/tests/test_foundry_local_client.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import inspect from unittest.mock import MagicMock, patch @@ -83,6 +85,23 @@ def test_agent_accepts_foundry_local_client(mock_foundry_local_manager: MagicMoc assert agent.client is client +async def test_foundry_local_context_closes_connection_not_service(mock_foundry_local_manager: MagicMock) -> None: + with patch( + "agent_framework_foundry_local._foundry_local_client.FoundryLocalManager", + return_value=mock_foundry_local_manager, + ): + client = FoundryLocalClient(model="test-model-id") + with patch.object(client.client, "close", wraps=client.client.close) as close: + async with client as entered: + assert entered is client + assert not client.client.is_closed() + await client.close() + close.assert_awaited_once() + assert client.client.is_closed() + mock_foundry_local_manager.unload_model.assert_not_called() + mock_foundry_local_manager.stop_service.assert_not_called() + + def test_foundry_local_client_get_response_uses_explicit_runtime_buckets() -> None: """Foundry Local should expose explicit runtime buckets instead of raw kwargs.""" signature = inspect.signature(FoundryLocalClient.get_response) diff --git a/python/packages/gemini/agent_framework_gemini/_chat_client.py b/python/packages/gemini/agent_framework_gemini/_chat_client.py index b64c9e58cf7..ebfcf566238 100644 --- a/python/packages/gemini/agent_framework_gemini/_chat_client.py +++ b/python/packages/gemini/agent_framework_gemini/_chat_client.py @@ -341,6 +341,8 @@ def __init__( location=google_settings.get("cloud_location"), credentials=credentials, ) + self._owns_client = client is None + self._client_closed = False self.model = google_settings.get("model") self.max_tracked_thought_signatures = max_tracked_thought_signatures @@ -348,6 +350,16 @@ def __init__( super().__init__(additional_properties=additional_properties) + @override + async def close(self) -> None: + """Close both owned SDK transports, leaving caller-supplied credentials open.""" + if self._owns_client and not self._client_closed: + try: + await self._genai_client.aio.aclose() + finally: + self._genai_client.close() + self._client_closed = True + @staticmethod def get_code_interpreter_tool() -> types.Tool: """Create a code execution tool. diff --git a/python/packages/gemini/tests/test_gemini_client.py b/python/packages/gemini/tests/test_gemini_client.py index 873654676de..5c83d199683 100644 --- a/python/packages/gemini/tests/test_gemini_client.py +++ b/python/packages/gemini/tests/test_gemini_client.py @@ -76,6 +76,37 @@ def test_enterprise_integration_gate_accepts_either_mode(monkeypatch: pytest.Mon assert _has_gemini_integration_credentials() +@pytest.mark.parametrize("client_type", [RawGeminiChatClient, GeminiChatClient]) +@pytest.mark.parametrize("injected", [False, True]) +@pytest.mark.parametrize("fail_async_close", [False, True]) +async def test_gemini_client_context_ownership(client_type, injected: bool, fail_async_close: bool) -> None: + sdk_client = MagicMock() + sdk_client.aio.aclose = AsyncMock(side_effect=RuntimeError("close failed") if fail_async_close else None) + credential = MagicMock() + with patch("agent_framework_gemini._sdk_client.genai.Client", return_value=sdk_client): + client = client_type( + model="test-model", + enterprise=True, + project="test-project", + location="us-central1", + credentials=credential, + client=sdk_client if injected else None, + ) + if fail_async_close and not injected: + with pytest.raises(RuntimeError, match="close failed"): + async with client: + pass + sdk_client.aio.aclose.assert_awaited_once() + sdk_client.close.assert_called_once() + else: + async with client as entered: + assert entered is client + await client.close() + assert sdk_client.aio.aclose.await_count == (0 if injected else 1) + assert sdk_client.close.call_count == (0 if injected else 1) + credential.close.assert_not_called() + + @pytest.mark.integration @pytest.mark.skipif( os.getenv("GOOGLE_GENAI_USE_ENTERPRISE", "").lower() not in {"true", "1", "yes", "on"}, diff --git a/python/packages/mistral/AGENTS.md b/python/packages/mistral/AGENTS.md index 1304b6267e0..bd9b0ec623b 100644 --- a/python/packages/mistral/AGENTS.md +++ b/python/packages/mistral/AGENTS.md @@ -25,11 +25,9 @@ from agent_framework.mistral import MistralChatClient # Requires MISTRAL_API_KEY environment variable (or pass api_key= directly) client = MistralChatClient(model="mistral-large-latest") -try: +async with client: agent = Agent(client=client) result = await agent.run("Hello!") -finally: - await client.close() ``` ```python diff --git a/python/packages/mistral/agent_framework_mistral/_chat_client.py b/python/packages/mistral/agent_framework_mistral/_chat_client.py index 15ebd072158..1cd384e34ab 100644 --- a/python/packages/mistral/agent_framework_mistral/_chat_client.py +++ b/python/packages/mistral/agent_framework_mistral/_chat_client.py @@ -401,6 +401,7 @@ def __init__( self.model = mistral_settings.get("chat_model") self.server_url = mistral_settings.get("server_url") self._owns_client = not isinstance(client, Mistral) + self._client_closed = False if isinstance(client, Mistral): self.client = client @@ -421,11 +422,15 @@ def __init__( super().__init__(additional_properties=additional_properties) + @override async def close(self) -> None: """Close the internally created Mistral SDK client.""" - if self._owns_client: - await self.client.__aexit__(None, None, None) # type: ignore[no-untyped-call] - self.client.__exit__(None, None, None) # type: ignore[no-untyped-call] + if self._owns_client and not self._client_closed: + try: + await self.client.__aexit__(None, None, None) # type: ignore[no-untyped-call] + finally: + self.client.__exit__(None, None, None) # type: ignore[no-untyped-call] + self._client_closed = True @override def service_url(self) -> str: diff --git a/python/packages/mistral/tests/mistral/test_mistral_chat_client.py b/python/packages/mistral/tests/mistral/test_mistral_chat_client.py index ec65d53d396..9430f61bb01 100644 --- a/python/packages/mistral/tests/mistral/test_mistral_chat_client.py +++ b/python/packages/mistral/tests/mistral/test_mistral_chat_client.py @@ -1,10 +1,13 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import json import logging import os from collections.abc import AsyncIterator, Sequence from typing import Any, cast +from unittest.mock import AsyncMock, patch import httpx import pytest @@ -194,23 +197,58 @@ def test_mistral_chat_service_url_default() -> None: assert client.service_url() == "https://api.mistral.ai" -async def test_mistral_chat_close_only_closes_owned_client() -> None: - owned = MistralChatClient(model="mistral-large-latest", api_key="test-key") +@pytest.mark.parametrize("client_type", [RawMistralChatClient, MistralChatClient]) +async def test_mistral_chat_close_only_closes_owned_client(client_type) -> None: + owned = client_type(model="mistral-large-latest", api_key="test-key") owned_http_client = owned.client.sdk_configuration.async_client + async with owned as entered: + assert entered is owned await owned.close() assert isinstance(owned_http_client, httpx.AsyncClient) assert owned_http_client.is_closed http_client = httpx.AsyncClient(base_url="https://custom.mistral.ai") - injected = MistralChatClient(model="mistral-large-latest", http_client=http_client) + injected = client_type(model="mistral-large-latest", http_client=http_client) assert injected.service_url() == "https://custom.mistral.ai" + async with injected: + pass await injected.close() assert not http_client.is_closed await http_client.aclose() +async def test_mistral_chat_context_leaves_injected_sdk_client_open() -> None: + sdk_client = Mistral(api_key="test-key") + async_transport = sdk_client.sdk_configuration.async_client + sync_transport = sdk_client.sdk_configuration.client + client = MistralChatClient(model="test-model", client=sdk_client) + async with client: + pass + await client.close() + assert isinstance(async_transport, httpx.AsyncClient) + assert isinstance(sync_transport, httpx.Client) + assert not async_transport.is_closed + assert not sync_transport.is_closed + await async_transport.aclose() + sync_transport.close() + + +async def test_mistral_chat_closes_sync_transport_when_async_cleanup_fails() -> None: + client = MistralChatClient(model="test-model", api_key="test-key") + sync_transport = client.client.sdk_configuration.client + assert isinstance(sync_transport, httpx.Client) + with ( + patch.object(client.client, "__aexit__", new_callable=AsyncMock, side_effect=RuntimeError("close failed")), + pytest.raises(RuntimeError, match="close failed"), + ): + async with client: + pass + assert sync_transport.is_closed + await client.close() + + async def test_mistral_chat_missing_model_raises_at_request(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("MISTRAL_CHAT_MODEL", raising=False) http_client = httpx.AsyncClient(base_url="https://api.mistral.ai") diff --git a/python/packages/ollama/agent_framework_ollama/_chat_client.py b/python/packages/ollama/agent_framework_ollama/_chat_client.py index 970840ed6e4..a5f06ba9c5b 100644 --- a/python/packages/ollama/agent_framework_ollama/_chat_client.py +++ b/python/packages/ollama/agent_framework_ollama/_chat_client.py @@ -340,6 +340,8 @@ def __init__( self.model = ollama_settings["model"] # type: ignore[assignment, reportTypedDictNotRequiredAccess] # we can just pass in None for the host, the default is set by the Ollama package. self.client = client or AsyncClient(host=ollama_settings.get("host")) + self._owns_client = client is None + self._client_closed = False # Save Host URL for serialization with to_dict() self.host = str(self.client._client.base_url) # type: ignore[reportUnknownMemberType,reportPrivateUsage,reportUnknownArgumentType] @@ -350,6 +352,13 @@ def __init__( ) self.middleware = list(self.chat_middleware) + @override + async def close(self) -> None: + """Close the owned HTTP transport; Ollama exposes no public close method.""" + if self._owns_client and not self._client_closed: + await self.client._client.aclose() # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] + self._client_closed = True + @override def service_url(self) -> str: """Get the URL of the Ollama service. diff --git a/python/packages/ollama/tests/test_ollama_chat_client.py b/python/packages/ollama/tests/test_ollama_chat_client.py index b404bf2701c..706497f0ad9 100644 --- a/python/packages/ollama/tests/test_ollama_chat_client.py +++ b/python/packages/ollama/tests/test_ollama_chat_client.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import base64 import os import uuid @@ -35,6 +37,21 @@ from agent_framework_ollama import OllamaChatClient + +@pytest.mark.parametrize("injected", [False, True]) +async def test_ollama_client_context_ownership(injected: bool) -> None: + sdk_client = AsyncClient(host="http://localhost:11434") + with patch("agent_framework_ollama._chat_client.AsyncClient", return_value=sdk_client): + client = OllamaChatClient(model="test-model", client=sdk_client if injected else None) + async with client as entered: + assert entered is client + assert not sdk_client._client.is_closed + await client.close() + assert sdk_client._client.is_closed is not injected + if injected: + await sdk_client._client.aclose() + + # region Service Setup skip_if_azure_integration_tests_disabled = pytest.mark.skipif( diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index fdb8c3d8eae..65fafa4584d 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -817,6 +817,8 @@ def __init__( ) self.client = client + self._owns_client = async_client is None + self._client_closed = False self.model: str = settings.get("model") or "" # Store configuration for serialization @@ -876,6 +878,13 @@ def _handle_request_error(self, ex: Exception) -> NoReturn: inner_exception=ex, ) from ex + @override + async def close(self) -> None: + """Close the connector-owned SDK connection, not caller-supplied credentials.""" + if self._owns_client and not self._client_closed: + await self.client.close() + self._client_closed = True + @override def service_url(self) -> str: """Get the URL of the service. diff --git a/python/packages/openai/agent_framework_openai/_chat_completion_client.py b/python/packages/openai/agent_framework_openai/_chat_completion_client.py index 92be0f53933..27da4fea330 100644 --- a/python/packages/openai/agent_framework_openai/_chat_completion_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_completion_client.py @@ -488,6 +488,8 @@ def __init__( ) self.client = client + self._owns_client = async_client is None + self._client_closed = False self.model: str = settings.get("model") or "" # Store configuration for serialization @@ -1415,6 +1417,13 @@ def _prepare_content_for_openai(self, content: Content) -> dict[str, Any]: logger.debug("Unsupported content type passed (type: %s)", content.type) return {} + @override + async def close(self) -> None: + """Close the connector-owned SDK connection, not caller-supplied credentials.""" + if self._owns_client and not self._client_closed: + await self.client.close() + self._client_closed = True + @override def service_url(self) -> str: """Get the URL of the service. diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index ca9870e1ae3..40c507a1aa6 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import asyncio import base64 import copy @@ -80,7 +82,13 @@ from pydantic import BaseModel from pytest import param -from agent_framework_openai import OpenAIChatClient, OpenAIChatOptions, RawOpenAIChatClient +from agent_framework_openai import ( + OpenAIChatClient, + OpenAIChatCompletionClient, + OpenAIChatOptions, + RawOpenAIChatClient, + RawOpenAIChatCompletionClient, +) from agent_framework_openai._chat_client import OPENAI_LOCAL_SHELL_CALL_ITEM_ID_KEY from agent_framework_openai._exceptions import OpenAIContentFilterException @@ -99,6 +107,38 @@ class OutputStruct(BaseModel): weather: str | None = None +@pytest.mark.parametrize( + "client_type", [RawOpenAIChatClient, OpenAIChatClient, RawOpenAIChatCompletionClient, OpenAIChatCompletionClient] +) +@pytest.mark.parametrize("azure", [False, True]) +@pytest.mark.parametrize("injected", [False, True]) +async def test_openai_chat_client_context_ownership(client_type, azure: bool, injected: bool) -> None: + from azure.core.credentials_async import AsyncTokenCredential + + sdk_client = AsyncOpenAI(api_key="test-key") if injected else None + credential = MagicMock(spec=AsyncTokenCredential) + credential.close = AsyncMock() + kwargs: dict[str, Any] = {"async_client": sdk_client} if injected else {} + if azure: + kwargs.update(azure_endpoint="https://test.openai.azure.com", credential=credential) + else: + kwargs["api_key"] = "test-key" + + client = client_type(model="test-model", **kwargs) + with patch.object(client.client, "close", wraps=client.client.close) as close: + async with client as entered: + assert entered is client + close.assert_not_awaited() + await client.close() + assert close.await_count == (0 if injected else 1) + assert client.client.is_closed() is not injected + credential.close.assert_not_awaited() + assert "_owns_client" not in client.to_dict() + assert "_client_closed" not in client.to_dict() + if sdk_client is not None: + await sdk_client.close() + + class _FakeAsyncEventStream: def __init__(self, events: Sequence[object], headers: dict[str, str] | None = None) -> None: self._events = events diff --git a/skills/agent-framework-builder-python/references/integrations.md b/skills/agent-framework-builder-python/references/integrations.md index 2b298e0f103..689008db810 100644 --- a/skills/agent-framework-builder-python/references/integrations.md +++ b/skills/agent-framework-builder-python/references/integrations.md @@ -21,6 +21,16 @@ Verify authentication, streaming, structured output, hosted tools, history owner support instead of assuming provider parity. If no provider matches, use [Extensibility](extensibility.md#custom-providers). +## Connection and credential ownership + +Use `async with client` for standalone chat clients, or `async with Agent(client=client, ...)` to have the agent +enter and exit the client's context. Direct requests without a context remain supported; call +`await client.close()` when finished. Connectors release their own connections, not caller-injected SDK clients, +HTTP transports, or credentials. Scope supplied credentials outside the clients and agents that use them. +For Foundry, the derived OpenAI connection is owned by the connector even when the project client is supplied. +Do not share a connector-owned connection across overlapping agent contexts; inject an externally managed SDK +client when its lifetime must extend beyond any one agent context. Do not assume a closed SDK connection can reopen. + ## Capability categories - Use function tools for application-owned code and MCP for external tool services. From 8b2a429bcd7b14a6106c673039df55c0efc66bf5 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 27 Oct 2026 15:22:16 +0100 Subject: [PATCH 2/3] Python: Apply pyupgrade annotation fixes --- python/packages/core/tests/core/test_agents.py | 4 ++-- .../openai/tests/openai/test_openai_chat_client.py | 14 +++++++------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index de6b2c36793..b2d0dd90033 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/python/packages/core/tests/core/test_agents.py @@ -103,7 +103,7 @@ def __init__(self, events: list[str]) -> None: self.events = events self.exit_exception_type: type[BaseException] | None = None - async def __aenter__(self) -> "_LifecycleClient": + async def __aenter__(self) -> _LifecycleClient: self.events.append("client.enter") return self @@ -119,7 +119,7 @@ def __init__(self, name: str, events: list[str], *, fail_on_enter: bool = False) self.fail_on_enter = fail_on_enter self.exit_exception_type: type[BaseException] | None = None - async def __aenter__(self) -> "_LifecycleMCPTool": + async def __aenter__(self) -> _LifecycleMCPTool: self.events.append(f"{self.name}.enter") if self.fail_on_enter: raise RuntimeError("MCP entry failed") diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index 40c507a1aa6..f558003c725 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -145,7 +145,7 @@ def __init__(self, events: Sequence[object], headers: dict[str, str] | None = No self._iterator: Iterator[object] = iter(()) self._headers = headers or {} - def __aiter__(self) -> "_FakeAsyncEventStream": + def __aiter__(self) -> _FakeAsyncEventStream: self._iterator = iter(self._events) return self @@ -160,14 +160,14 @@ async def __anext__(self) -> object: # ``.headers``. The chat client then ``async with``-s the parsed stream so the # underlying socket is closed deterministically. Mimic both interfaces here so # test mocks remain a single object. - def parse(self) -> "_FakeAsyncEventStream": + def parse(self) -> _FakeAsyncEventStream: return self @property def headers(self) -> dict[str, str]: return self._headers - async def __aenter__(self) -> "_FakeAsyncEventStream": + async def __aenter__(self) -> _FakeAsyncEventStream: return self async def __aexit__( @@ -206,7 +206,7 @@ def _as_raw(mock_response: MagicMock, *, headers: dict[str, str] | None = None) class _FakeAsyncEventStreamContext(_FakeAsyncEventStream): - async def __aenter__(self) -> "_FakeAsyncEventStreamContext": + async def __aenter__(self) -> _FakeAsyncEventStreamContext: return self async def __aexit__( @@ -889,7 +889,7 @@ def __init__(self, items: Sequence[object]) -> None: self._items = list(items) self._iterator: Iterator[object] = iter(()) - def __aiter__(self) -> "_BareEventStream": + def __aiter__(self) -> _BareEventStream: self._iterator = iter(self._items) return self @@ -920,7 +920,7 @@ class _FakeTelemetryStreamWrapper: def __init__(self, stream_async_iter: Any) -> None: self.stream_async_iter: Any = stream_async_iter - def __aiter__(self) -> "_FakeTelemetryStreamWrapper": + def __aiter__(self) -> _FakeTelemetryStreamWrapper: self.stream_async_iter = self.stream_async_iter.__aiter__() return self @@ -7652,7 +7652,7 @@ async def test_streaming_azure_ai_search_output_enriches_mapped_agent_response() ), ] - async def _stream() -> AsyncGenerator[ChatResponseUpdate, None]: + async def _stream() -> AsyncGenerator[ChatResponseUpdate]: for update in updates: yield update From d49dca0b363bc7564527e3b58b3ee7ea2d5a59b6 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 27 Oct 2026 15:26:10 +0100 Subject: [PATCH 3/3] Python: Preserve synchronous chat client cleanup compatibility --- python/README.md | 1 + .../packages/core/agent_framework/_clients.py | 6 ++- .../packages/core/tests/core/test_agents.py | 18 +++++++++ .../packages/core/tests/core/test_clients.py | 38 ++++++++++++++++++- 4 files changed, 60 insertions(+), 3 deletions(-) diff --git a/python/README.md b/python/README.md index 97d059bc00a..d6fbd768d2c 100644 --- a/python/README.md +++ b/python/README.md @@ -196,6 +196,7 @@ agent contexts sharing a connector-owned connection, since exiting either agent an SDK client across agents, inject it into their chat clients and manage its lifetime outside those agent contexts. `BaseChatClient` provides default context methods and a no-op `close()`, so custom subclasses do not need to implement lifecycle methods unless they own resources. Protocol-only chat clients do not need context methods. +Inherited context exit also supports existing subclasses whose `close()` is synchronous. ## 4. Build an Agent with Tools and Functions diff --git a/python/packages/core/agent_framework/_clients.py b/python/packages/core/agent_framework/_clients.py index bef2485eba8..447cb730b0b 100644 --- a/python/packages/core/agent_framework/_clients.py +++ b/python/packages/core/agent_framework/_clients.py @@ -12,6 +12,7 @@ Mapping, Sequence, ) +from inspect import isawaitable from types import TracebackType from typing import ( TYPE_CHECKING, @@ -327,13 +328,16 @@ async def __aexit__( /, ) -> None: """Release resources owned by the client without suppressing exceptions.""" - await self.close() + result = self.close() + if isawaitable(result): + await result async def close(self) -> None: """Release client-owned connections and credentials. The default implementation does nothing. Resource-owning implementations override this method and must leave caller-supplied resources open. + Context exit also accepts legacy synchronous ``close()`` overrides. Clients remain usable without entering an async context; callers can use ``await client.close()`` to release their resources explicitly. """ diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index b2d0dd90033..b28da937cce 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/python/packages/core/tests/core/test_agents.py @@ -249,6 +249,24 @@ async def test_agent_still_accepts_client_without_lifecycle(client: SupportsChat assert response.text == "test response" +@pytest.mark.parametrize("agent_type", [Agent, RawAgent]) +@pytest.mark.parametrize("error", [None, RuntimeError("body failed")]) +async def test_agent_context_accepts_client_with_sync_close( + agent_type: type[RawAgent], chat_client_base: MockBaseChatClient, error: RuntimeError | None +) -> None: + with patch.object(chat_client_base, "close", new_callable=MagicMock, return_value=None) as close: + agent = agent_type(client=chat_client_base) + if error is None: + async with agent: + close.assert_not_called() + else: + with pytest.raises(RuntimeError) as raised: + async with agent: + raise error + assert raised.value is error + close.assert_called_once_with() + + @pytest.mark.parametrize("agent_type", [Agent, RawAgent]) async def test_agent_async_with_delegates_to_open_close(agent_type: type[RawAgent]) -> None: events: list[str] = [] diff --git a/python/packages/core/tests/core/test_clients.py b/python/packages/core/tests/core/test_clients.py index 8b59181aa28..55a4c470450 100644 --- a/python/packages/core/tests/core/test_clients.py +++ b/python/packages/core/tests/core/test_clients.py @@ -5,7 +5,7 @@ import asyncio import logging from typing import Any -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -129,7 +129,7 @@ async def test_base_client_default_lifecycle(chat_client_base: BaseChatClient) - assert (await chat_client_base.get_response([Message("user", ["Again"])])).text == "test response - Again" -@pytest.mark.parametrize("error", [None, RuntimeError("body failed"), asyncio.CancelledError()]) +@pytest.mark.parametrize("error", [None, RuntimeError("body failed")]) async def test_base_client_context_calls_close(chat_client_base: BaseChatClient, error: BaseException | None) -> None: with patch.object(chat_client_base, "close", new_callable=AsyncMock) as close: if error is None: @@ -143,6 +143,40 @@ async def test_base_client_context_calls_close(chat_client_base: BaseChatClient, close.assert_awaited_once_with() +async def test_base_client_context_closes_on_cancellation(chat_client_base: BaseChatClient) -> None: + entered = asyncio.Event() + + async def run_client() -> None: + async with chat_client_base: + entered.set() + await asyncio.Event().wait() + + with patch.object(chat_client_base, "close", new_callable=AsyncMock) as close: + task = asyncio.create_task(run_client()) + await entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert task.cancelled() + close.assert_awaited_once_with() + + +@pytest.mark.parametrize("error", [None, RuntimeError("body failed")]) +async def test_base_client_context_accepts_sync_close( + chat_client_base: BaseChatClient, error: RuntimeError | None +) -> None: + with patch.object(chat_client_base, "close", new_callable=MagicMock, return_value=None) as close: + if error is None: + async with chat_client_base: + close.assert_not_called() + else: + with pytest.raises(RuntimeError) as raised: + async with chat_client_base: + raise error + assert raised.value is error + close.assert_called_once_with() + + async def test_base_client_context_propagates_cleanup_failure(chat_client_base: BaseChatClient) -> None: with ( patch.object(chat_client_base, "close", new_callable=AsyncMock, side_effect=RuntimeError("close failed")),