diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index 614efcad153..f4b2b808773 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 @@ -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"): 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..e851208bde7 --- /dev/null +++ b/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()