Python: keep Foundry injected clients caller-owned on cleanup

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()
 

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,27 +132,45 @@ 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")
         if not scope:
             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 ---------------------------------------------------------

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()

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()