diff --git a/python/packages/foundry/agent_framework_foundry/_embedding_client.py b/python/packages/foundry/agent_framework_foundry/_embedding_client.py index 8050198ab30..84d9adfe906 100644 --- a/python/packages/foundry/agent_framework_foundry/_embedding_client.py +++ b/python/packages/foundry/agent_framework_foundry/_embedding_client.py @@ -215,6 +215,8 @@ def __init__( self._openai_client: AsyncOpenAI | None = None self._text_client: EmbeddingsClient | None = None self._image_client: ImageEmbeddingsClient | None = None + self._owns_text_client = False + self._owns_image_client = False self.default_headers = ( {key: value for key, value in default_headers.items() if key != USER_AGENT_KEY} if default_headers else None ) @@ -281,33 +283,49 @@ def __init__( } if IS_TELEMETRY_ENABLED: client_kwargs["user_agent"] = get_user_agent() + self._owns_text_client = text_client is None + self._owns_image_client = image_client is None self._text_client = text_client self._image_client = image_client if credential is not None: - self._text_client = text_client or EmbeddingsClient( - **client_kwargs, - per_retry_policies=[create_feature_usage_policy()], - ) - self._image_client = image_client or ImageEmbeddingsClient( - **client_kwargs, - per_retry_policies=[create_feature_usage_policy()], - ) + if text_client is None: + self._text_client = EmbeddingsClient( + **client_kwargs, + per_retry_policies=[create_feature_usage_policy()], + ) + self._owns_text_client = True + else: + self._text_client = text_client + self._owns_text_client = False + if image_client is None: + self._image_client = ImageEmbeddingsClient( + **client_kwargs, + per_retry_policies=[create_feature_usage_policy()], + ) + self._owns_image_client = True + else: + self._image_client = image_client + self._owns_image_client = False self._endpoint = resolved_models_endpoint super().__init__(additional_properties=additional_properties) async def close(self) -> None: - """Close the underlying SDK clients and release resources.""" + """Close owned SDK clients and release resources. + + Injected ``text_client``, ``image_client``, and ``project_client`` + instances remain caller-owned and are not closed here. + """ if self._openai_client is not None: with suppress(Exception): await self._openai_client.close() if self._owns_project_client and self.project_client is not None: with suppress(Exception): await self.project_client.close() - if self._text_client is not None: + if self._owns_text_client and self._text_client is not None: with suppress(Exception): await self._text_client.close() - if self._image_client is not None: + if self._owns_image_client and self._image_client is not None: with suppress(Exception): await self._image_client.close() diff --git a/python/packages/foundry/agent_framework_foundry/_memory_provider.py b/python/packages/foundry/agent_framework_foundry/_memory_provider.py index 237146c8fd9..6bdd0884e55 100644 --- a/python/packages/foundry/agent_framework_foundry/_memory_provider.py +++ b/python/packages/foundry/agent_framework_foundry/_memory_provider.py @@ -86,7 +86,8 @@ def __init__( Args: source_id: Unique identifier for this provider instance. - project_client: Azure AI Project client for memory operations. + project_client: Azure AI Project client for memory operations. Injected + clients remain caller-owned and are not entered/exited by this provider. project_endpoint: Foundry project endpoint URL. Used when project_client is not provided. credential: Azure credential for authentication. Accepts a TokenCredential, AsyncTokenCredential, or a callable token provider. @@ -109,6 +110,7 @@ def __init__( env_file_encoding=env_file_encoding, ) + owns_project_client = False if project_client is None: from azure.ai.projects.aio import AIProjectClient @@ -130,6 +132,7 @@ def __init__( if allow_preview is not None: project_client_kwargs["allow_preview"] = allow_preview project_client = AIProjectClient(**project_client_kwargs) + owns_project_client = True if not memory_store_name: raise ValueError("memory_store_name is required") @@ -137,20 +140,37 @@ def __init__( raise ValueError("scope is required") self.project_client = project_client + self._owns_project_client = owns_project_client self.memory_store_name = memory_store_name self.scope = scope self.context_prompt = context_prompt or self.DEFAULT_CONTEXT_PROMPT self.update_delay = update_delay async def __aenter__(self) -> Self: - """Async context manager entry.""" - if self.project_client and isinstance(self.project_client, AbstractAsyncContextManager): + """Async context manager entry. + + Only enters a ``project_client`` that this provider created. Injected + clients remain caller-owned. + """ + if ( + self._owns_project_client + and self.project_client + and isinstance(self.project_client, AbstractAsyncContextManager) + ): await self.project_client.__aenter__() return self async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None: - """Async context manager exit.""" - if self.project_client and isinstance(self.project_client, AbstractAsyncContextManager): + """Async context manager exit. + + Only exits a ``project_client`` that this provider created. Injected + clients remain caller-owned. + """ + if ( + self._owns_project_client + and self.project_client + and isinstance(self.project_client, AbstractAsyncContextManager) + ): await self.project_client.__aexit__(exc_type, exc_val, exc_tb) # -- Hooks pattern --------------------------------------------------------- diff --git a/python/packages/foundry/tests/foundry/test_foundry_injected_client_lifecycle.py b/python/packages/foundry/tests/foundry/test_foundry_injected_client_lifecycle.py new file mode 100644 index 00000000000..5284157f986 --- /dev/null +++ b/python/packages/foundry/tests/foundry/test_foundry_injected_client_lifecycle.py @@ -0,0 +1,56 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +from contextlib import AbstractAsyncContextManager +from typing import Any, Self +from unittest.mock import AsyncMock, patch + +from agent_framework_foundry import FoundryMemoryProvider, RawFoundryEmbeddingClient + + +class _InjectedProjectClient(AsyncMock, AbstractAsyncContextManager): # type: ignore[misc] + async def __aenter__(self) -> Self: + await AsyncMock.__aenter__(self) + return self + + async def __aexit__(self, *args: Any) -> None: + await AsyncMock.__aexit__(self, *args) + + +async def test_raw_foundry_embedding_close_skips_injected_clients() -> None: + text = AsyncMock() + image = AsyncMock() + embeddings = RawFoundryEmbeddingClient( + model="embed-model", + endpoint="https://configured.example", + text_client=text, + image_client=image, + ) + + await embeddings.close() + + text.close.assert_not_awaited() + image.close.assert_not_awaited() + + +async def test_foundry_memory_provider_skips_injected_project_client_lifecycle() -> None: + project = _InjectedProjectClient() + memory = FoundryMemoryProvider( + project_client=project, + memory_store_name="store", + scope="scope", + ) + + assert memory._owns_project_client is False + enter_mock = AsyncMock() + exit_mock = AsyncMock() + with ( + patch.object(project, "__aenter__", enter_mock), + patch.object(project, "__aexit__", exit_mock), + ): + await memory.__aenter__() + await memory.__aexit__(None, None, None) + + enter_mock.assert_not_awaited() + exit_mock.assert_not_awaited() diff --git a/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py b/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py index a175bf31d6b..1a38660f9fe 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py +++ b/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py @@ -473,23 +473,31 @@ async def test_handles_update_exception_gracefully(mock_project_client: AsyncMoc # -- Context manager tests ----------------------------------------------------- -async def test_aenter_delegates_to_client(mock_project_client: AsyncMock) -> None: - provider = FoundryMemoryProvider( - project_client=mock_project_client, - memory_store_name="test_store", - scope="user_123", - ) +async def test_aenter_delegates_to_client(mock_project_client: AsyncMock, mock_credential: Mock) -> None: + with patch("azure.ai.projects.aio.AIProjectClient") as mock_ai_project_client: + mock_ai_project_client.return_value = mock_project_client + provider = FoundryMemoryProvider( + project_endpoint="https://test.project.endpoint", + credential=mock_credential, # type: ignore[arg-type] + memory_store_name="test_store", + scope="user_123", + ) + assert provider._owns_project_client is True result = await provider.__aenter__() assert result is provider mock_project_client.__aenter__.assert_awaited_once() -async def test_aexit_delegates_to_client(mock_project_client: AsyncMock) -> None: - provider = FoundryMemoryProvider( - project_client=mock_project_client, - memory_store_name="test_store", - scope="user_123", - ) +async def test_aexit_delegates_to_client(mock_project_client: AsyncMock, mock_credential: Mock) -> None: + with patch("azure.ai.projects.aio.AIProjectClient") as mock_ai_project_client: + mock_ai_project_client.return_value = mock_project_client + provider = FoundryMemoryProvider( + project_endpoint="https://test.project.endpoint", + credential=mock_credential, # type: ignore[arg-type] + memory_store_name="test_store", + scope="user_123", + ) + assert provider._owns_project_client is True await provider.__aexit__(None, None, None) mock_project_client.__aexit__.assert_awaited_once()