From f282f1515b6aaa7abccee6cee0cedd16de0fae85 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+universeplayer@users.noreply.github.com> Date: Fri, 24 Apr 2026 13:37:17 +0800 Subject: [PATCH 1/4] Python: make FoundryChatClient an async context manager MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #5428. `built_in_chat_clients` and similar samples wrap the client in `async with client:`, which works for OpenAIChatClient and the Azure variants but fails on FoundryChatClient with `TypeError: object does not support the asynchronous context manager protocol`. The chat client holds an AIProjectClient internally (self- created when the caller passes project_endpoint + credential), so it also needs a lifecycle hook — otherwise the project client leaks. Mirror the pattern already used in `FoundryChatAgent`: - Track `_should_close_client` in `__init__`, flipped on only when we construct the project client ourselves. - Add `async def close()` that awaits `project_client.close()` only when owned. - Implement `__aenter__` (no-op, returns self) and `__aexit__` (awaits close). `FoundryChatClient` inherits both via the existing class hierarchy, so the sample starts working without any call-site changes. --- .../agent_framework_foundry/_chat_client.py | 24 +++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index 614efcad153..040445b207b 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -46,9 +46,9 @@ else: from typing_extensions import override # type: ignore # pragma: no cover if sys.version_info >= (3, 11): - from typing import TypedDict # type: ignore # pragma: no cover + from typing import Self, TypedDict # type: ignore # pragma: no cover else: - from typing_extensions import TypedDict # type: ignore # pragma: no cover + from typing_extensions import Self, TypedDict # type: ignore # pragma: no cover if TYPE_CHECKING: from agent_framework import ChatAndFunctionMiddlewareTypes, ToolTypes @@ -187,6 +187,7 @@ def __init__( "Either 'project_endpoint' or 'project_client' is required. " "Set project_endpoint via parameter or 'FOUNDRY_PROJECT_ENDPOINT' environment variable." ) + self._should_close_client = False if not project_client: if not project_endpoint: raise ValueError( @@ -204,6 +205,7 @@ def __init__( if allow_preview is not None: project_client_kwargs["allow_preview"] = allow_preview project_client = AIProjectClient(**project_client_kwargs) + self._should_close_client = True openai_kwargs: dict[str, Any] = {} if default_headers: @@ -220,6 +222,24 @@ def __init__( ) self.project_client = project_client + async def close(self) -> None: + """Close the project client if we created it.""" + if self._should_close_client: + await self.project_client.close() + + async def __aenter__(self) -> Self: + """Enter the async context manager.""" + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: Any, + ) -> None: + """Close the project client on exit (only when owned).""" + await self.close() + @override def _check_model_presence(self, options: dict[str, Any]) -> None: if not options.get("model"): From bf11675801ec841932ae1652006f706e6f97bfaf Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Wed, 6 May 2026 23:22:15 +0800 Subject: [PATCH 2/4] test: cover FoundryChatClient context ownership --- ...est_foundry_chat_client_context_manager.py | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) create mode 100644 python/packages/foundry/tests/foundry/test_foundry_chat_client_context_manager.py diff --git a/python/packages/foundry/tests/foundry/test_foundry_chat_client_context_manager.py b/python/packages/foundry/tests/foundry/test_foundry_chat_client_context_manager.py new file mode 100644 index 00000000000..98081d38222 --- /dev/null +++ b/python/packages/foundry/tests/foundry/test_foundry_chat_client_context_manager.py @@ -0,0 +1,47 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +from agent_framework_foundry import FoundryChatClient + +_TEST_FOUNDRY_PROJECT_ENDPOINT = "https://test-project.services.ai.azure.com/" +_TEST_FOUNDRY_MODEL = "test-gpt-4o" + + +def _make_mock_openai_client() -> MagicMock: + client = MagicMock() + client.default_headers = {} + client.responses = MagicMock() + client.responses.create = AsyncMock() + client.responses.parse = AsyncMock() + return client + + +async def test_context_manager_closes_owned_project_client() -> None: + credential = MagicMock() + project_client = MagicMock() + project_client.get_openai_client.return_value = _make_mock_openai_client() + project_client.close = AsyncMock() + + with patch("agent_framework_foundry._chat_client.AIProjectClient", return_value=project_client): + async with FoundryChatClient( + project_endpoint=_TEST_FOUNDRY_PROJECT_ENDPOINT, + model=_TEST_FOUNDRY_MODEL, + credential=credential, + ) as client: + assert client.project_client is project_client + + project_client.close.assert_awaited_once_with() + + +async def test_context_manager_does_not_close_injected_project_client() -> None: + project_client = MagicMock() + project_client.get_openai_client.return_value = _make_mock_openai_client() + project_client.close = AsyncMock() + + async with FoundryChatClient(project_client=project_client, model=_TEST_FOUNDRY_MODEL) as client: + assert client.project_client is project_client + + project_client.close.assert_not_awaited() From f423efa08858544ea79d33346f0b97740dbf38a3 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Thu, 7 May 2026 11:39:07 +0800 Subject: [PATCH 3/4] fix: close owned Foundry client on init failure --- .../agent_framework_foundry/_chat_client.py | 42 ++++++++++++++----- ...est_foundry_chat_client_context_manager.py | 23 ++++++++++ 2 files changed, 54 insertions(+), 11 deletions(-) diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index 040445b207b..35987786ee6 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import logging import sys from collections.abc import Awaitable, Callable, Mapping, Sequence @@ -131,6 +132,15 @@ class RawFoundryChatClient( # type: ignore[misc] OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.foundry" # type: ignore[reportIncompatibleVariableOverride, misc] SUPPORTS_RICH_FUNCTION_OUTPUT: ClassVar[bool] = False # type: ignore[reportIncompatibleVariableOverride, misc] + @staticmethod + def _close_project_client_after_init_error(project_client: AIProjectClient) -> None: + try: + loop = asyncio.get_running_loop() + except RuntimeError: + asyncio.run(project_client.close()) + else: + loop.create_task(project_client.close()) + def __init__( self, *, @@ -188,7 +198,8 @@ def __init__( "Set project_endpoint via parameter or 'FOUNDRY_PROJECT_ENDPOINT' environment variable." ) self._should_close_client = False - if not project_client: + project_client_created = False + if project_client is None: if not project_endpoint: raise ValueError( "Azure AI project endpoint is required. Set via 'project_endpoint' parameter " @@ -205,21 +216,30 @@ def __init__( if allow_preview is not None: project_client_kwargs["allow_preview"] = allow_preview project_client = AIProjectClient(**project_client_kwargs) - self._should_close_client = True + project_client_created = True openai_kwargs: dict[str, Any] = {} if default_headers: openai_kwargs["default_headers"] = default_headers - super().__init__( - model=resolved_model, - async_client=project_client.get_openai_client(**openai_kwargs), - default_headers=default_headers, - instruction_role=instruction_role, - compaction_strategy=compaction_strategy, - tokenizer=tokenizer, - additional_properties=additional_properties, - ) + try: + super().__init__( + model=resolved_model, + async_client=project_client.get_openai_client(**openai_kwargs), + default_headers=default_headers, + instruction_role=instruction_role, + compaction_strategy=compaction_strategy, + tokenizer=tokenizer, + additional_properties=additional_properties, + ) + except Exception: + if project_client_created: + try: + self._close_project_client_after_init_error(project_client) + except Exception: + logger.warning("Failed to close Foundry project client after initialization error.", exc_info=True) + raise + self._should_close_client = project_client_created self.project_client = project_client async def close(self) -> None: diff --git a/python/packages/foundry/tests/foundry/test_foundry_chat_client_context_manager.py b/python/packages/foundry/tests/foundry/test_foundry_chat_client_context_manager.py index 98081d38222..e851208bde7 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_chat_client_context_manager.py +++ b/python/packages/foundry/tests/foundry/test_foundry_chat_client_context_manager.py @@ -2,8 +2,11 @@ from __future__ import annotations +import asyncio from unittest.mock import AsyncMock, MagicMock, patch +import pytest + from agent_framework_foundry import FoundryChatClient _TEST_FOUNDRY_PROJECT_ENDPOINT = "https://test-project.services.ai.azure.com/" @@ -45,3 +48,23 @@ async def test_context_manager_does_not_close_injected_project_client() -> None: assert client.project_client is project_client project_client.close.assert_not_awaited() + + +async def test_constructor_closes_owned_project_client_when_openai_client_creation_fails() -> None: + credential = MagicMock() + project_client = MagicMock() + project_client.get_openai_client.side_effect = RuntimeError("boom") + project_client.close = AsyncMock() + + with ( + patch("agent_framework_foundry._chat_client.AIProjectClient", return_value=project_client), + pytest.raises(RuntimeError, match="boom"), + ): + FoundryChatClient( + project_endpoint=_TEST_FOUNDRY_PROJECT_ENDPOINT, + model=_TEST_FOUNDRY_MODEL, + credential=credential, + ) + + await asyncio.sleep(0) + project_client.close.assert_awaited_once_with() From 2037ee3ca5448b9a83cb4cdfd5527c7cf2d5f45b Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Thu, 7 May 2026 11:44:38 +0800 Subject: [PATCH 4/4] fix: keep async cleanup task referenced --- .../foundry/agent_framework_foundry/_chat_client.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index 35987786ee6..f4b2b808773 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -139,7 +139,15 @@ def _close_project_client_after_init_error(project_client: AIProjectClient) -> N except RuntimeError: asyncio.run(project_client.close()) else: - loop.create_task(project_client.close()) + close_task = loop.create_task(project_client.close()) + close_task.add_done_callback(RawFoundryChatClient._log_project_client_close_failure) + + @staticmethod + def _log_project_client_close_failure(close_task: asyncio.Task[None]) -> None: + try: + close_task.result() + except Exception: + logger.warning("Failed to close Foundry project client after initialization error.", exc_info=True) def __init__( self,