Python: make FoundryChatClient an async context manager

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
@@ -46,9 +47,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
@@ -131,6 +132,23 @@ 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:
+            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,
         *,
@@ -187,7 +205,9 @@ def __init__(
                 "Either 'project_endpoint' or 'project_client' is required. "
                 "Set project_endpoint via parameter or 'FOUNDRY_PROJECT_ENDPOINT' environment variable."
             )
-        if not project_client:
+        self._should_close_client = False
+        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 "
@@ -204,22 +224,50 @@ def __init__(
             if allow_preview is not None:
                 project_client_kwargs["allow_preview"] = allow_preview
             project_client = AIProjectClient(**project_client_kwargs)
+            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:
+        """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"):

python/packages/foundry/tests/foundry/test_foundry_chat_client_context_manager.py

@@ -0,0 +1,70 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+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/"
+_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()
+
+
+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()