From 6bf32709707ee6fce18fe1b192ed29b4256529d7 Mon Sep 17 00:00:00 2001 From: ninglinLiu Date: Tue, 3 Nov 2026 08:44:50 +0800 Subject: [PATCH 1/3] Python: keep Foundry injected clients caller-owned on cleanup --- .../_embedding_client.py | 40 ++++++++++---- .../_memory_provider.py | 38 ++++++++----- .../test_foundry_injected_client_lifecycle.py | 54 +++++++++++++++++++ 3 files changed, 109 insertions(+), 23 deletions(-) create mode 100644 python/packages/foundry/tests/foundry/test_foundry_injected_client_lifecycle.py 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..ed725496212 100644 --- a/python/packages/foundry/agent_framework_foundry/_memory_provider.py +++ b/python/packages/foundry/agent_framework_foundry/_memory_provider.py @@ -9,9 +9,8 @@ from __future__ import annotations import logging -import sys from contextlib import AbstractAsyncContextManager -from typing import TYPE_CHECKING, Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar, Self, TypedDict from agent_framework import ( AgentSession, @@ -26,11 +25,6 @@ from ._feature_usage import FeatureIndex, create_feature_usage_policy -if sys.version_info >= (3, 11): - from typing import Self, TypedDict # pragma: no cover -else: - from typing_extensions import Self, TypedDict # pragma: no cover - if TYPE_CHECKING: from agent_framework import SupportsAgentRun from azure.ai.projects.aio import AIProjectClient @@ -86,7 +80,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 +104,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 +126,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 +134,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..c12d357c6aa --- /dev/null +++ b/python/packages/foundry/tests/foundry/test_foundry_injected_client_lifecycle.py @@ -0,0 +1,54 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +from contextlib import AbstractAsyncContextManager +from typing import Any, Self +from unittest.mock import AsyncMock + +import pytest + +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) + + +@pytest.mark.asyncio +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() + + +@pytest.mark.asyncio +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 + await memory.__aenter__() + await memory.__aexit__(None, None, None) + + project.__aenter__.assert_not_awaited() + project.__aexit__.assert_not_awaited() From 367e0044b7f52678e6df98122b0f9db42417b9af Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 3 Nov 2026 06:54:05 +0000 Subject: [PATCH 2/3] Python: restore Foundry 3.10 Self fallback and owned-client lifecycle tests Keep the version-gated typing_extensions Self and TypedDict import so FoundryMemoryProvider loads on Python 3.10, import Self from typing_extensions in the injected-client lifecycle tests, and construct the aenter/aexit delegation tests through a patched AIProjectClient. Co-authored-by: NingLin Liu --- .../_memory_provider.py | 8 ++++- .../test_foundry_injected_client_lifecycle.py | 3 +- .../foundry/test_foundry_memory_provider.py | 32 ++++++++++++------- 3 files changed, 29 insertions(+), 14 deletions(-) diff --git a/python/packages/foundry/agent_framework_foundry/_memory_provider.py b/python/packages/foundry/agent_framework_foundry/_memory_provider.py index ed725496212..6bdd0884e55 100644 --- a/python/packages/foundry/agent_framework_foundry/_memory_provider.py +++ b/python/packages/foundry/agent_framework_foundry/_memory_provider.py @@ -9,8 +9,9 @@ from __future__ import annotations import logging +import sys from contextlib import AbstractAsyncContextManager -from typing import TYPE_CHECKING, Any, ClassVar, Self, TypedDict +from typing import TYPE_CHECKING, Any, ClassVar from agent_framework import ( AgentSession, @@ -25,6 +26,11 @@ from ._feature_usage import FeatureIndex, create_feature_usage_policy +if sys.version_info >= (3, 11): + from typing import Self, TypedDict # pragma: no cover +else: + from typing_extensions import Self, TypedDict # pragma: no cover + if TYPE_CHECKING: from agent_framework import SupportsAgentRun from azure.ai.projects.aio import AIProjectClient 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 index c12d357c6aa..4f7ce21d1b4 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_injected_client_lifecycle.py +++ b/python/packages/foundry/tests/foundry/test_foundry_injected_client_lifecycle.py @@ -3,10 +3,11 @@ from __future__ import annotations from contextlib import AbstractAsyncContextManager -from typing import Any, Self +from typing import Any from unittest.mock import AsyncMock import pytest +from typing_extensions import Self from agent_framework_foundry import FoundryMemoryProvider, RawFoundryEmbeddingClient 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() From edb6609f3d655eb0245ee7bd45d9f8c40867e9be Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 3 Nov 2026 08:30:24 +0000 Subject: [PATCH 3/3] Python: satisfy py311 pyupgrade and typing on injected-client tests Import Self from typing so pyupgrade --py311-plus leaves the lifecycle test unchanged, and assert injected project-client entry and exit against patched AsyncMocks. Co-authored-by: NingLin Liu --- .../test_foundry_injected_client_lifecycle.py | 25 ++++++++++--------- 1 file changed, 13 insertions(+), 12 deletions(-) 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 index 4f7ce21d1b4..5284157f986 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_injected_client_lifecycle.py +++ b/python/packages/foundry/tests/foundry/test_foundry_injected_client_lifecycle.py @@ -3,11 +3,8 @@ from __future__ import annotations from contextlib import AbstractAsyncContextManager -from typing import Any -from unittest.mock import AsyncMock - -import pytest -from typing_extensions import Self +from typing import Any, Self +from unittest.mock import AsyncMock, patch from agent_framework_foundry import FoundryMemoryProvider, RawFoundryEmbeddingClient @@ -21,7 +18,6 @@ async def __aexit__(self, *args: Any) -> None: await AsyncMock.__aexit__(self, *args) -@pytest.mark.asyncio async def test_raw_foundry_embedding_close_skips_injected_clients() -> None: text = AsyncMock() image = AsyncMock() @@ -38,7 +34,6 @@ async def test_raw_foundry_embedding_close_skips_injected_clients() -> None: image.close.assert_not_awaited() -@pytest.mark.asyncio async def test_foundry_memory_provider_skips_injected_project_client_lifecycle() -> None: project = _InjectedProjectClient() memory = FoundryMemoryProvider( @@ -48,8 +43,14 @@ async def test_foundry_memory_provider_skips_injected_project_client_lifecycle() ) assert memory._owns_project_client is False - await memory.__aenter__() - await memory.__aexit__(None, None, None) - - project.__aenter__.assert_not_awaited() - project.__aexit__.assert_not_awaited() + 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()