Python: Preserve injected Cosmos memory client configuration ownership

python/packages/azure-cosmos-memory/README.md

@@ -420,8 +420,10 @@ For fine-grained control over memory processing:
 from azure.cosmos.agent_memory.aio import AsyncCosmosMemoryClient
 
 # Create a custom memory client. To disable automatic extraction, zero the cadence thresholds
-# on the client you build - the provider cannot reconfigure a client you pass in, so supplying
-# a memory_client together with auto_extract=False or processor_config raises ValueError.
+# on the client you build. When memory_client is supplied, the provider ignores its connection,
+# credential, auto_extract, processor_config and prompts_dir options and preserves the client's
+# configuration. Retrieval options (top_k, min_confidence, memory_types and context_prompt) still apply.
+# The caller is responsible for closing the supplied client.
 memory_client = AsyncCosmosMemoryClient(
     cosmos_endpoint=cosmos_endpoint,
     cosmos_database="ai_memory",

python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py

@@ -119,26 +119,26 @@ def __init__(
             credential: Azure credential for authentication. When provided it is used for both
                 Cosmos DB and AI Foundry; when ``None`` the toolkit builds (and owns) a
                 ``DefaultAzureCredential``.
-            memory_client: Pre-created AsyncCosmosMemoryClient.
+            memory_client: Pre-created AsyncCosmosMemoryClient, used as configured by the caller.
+                Connection settings, credential, auto_extract, processor_config and prompts_dir
+                are ignored when supplied. The caller remains responsible for closing the client.
             top_k: Number of memories to retrieve in search.
             min_confidence: Minimum confidence score (0.0-1.0) for retrieved memories.
             memory_types: Types of memories to retrieve. Default: ["fact", "procedural"].
             context_prompt: Prompt to prepend to retrieved memories.
             auto_extract: Enable automatic background memory extraction/summarization after
                 turn writes. When ``False`` the cadence thresholds are zeroed so nothing runs
                 automatically and callers drive processing via ``memory_client.process_now()``.
-                Only applied when the provider builds the client; supplying ``memory_client``
-                together with ``auto_extract=False`` raises ``ValueError``.
+                Only applied when the provider builds the client; ignored with ``memory_client``.
             processor_config: Optional processor cadence configuration, forwarded to the toolkit
                 client via ``cadence_thresholds``. Only applied when the provider builds the
-                client; supplying ``memory_client`` together with ``processor_config`` raises
-                ``ValueError`` (configure cadence on your own client instead).
+                client; ignored with ``memory_client`` (configure cadence on your own client instead).
             prompts_dir: Optional directory of Prompty templates for the memory pipeline. When
                 set, the extraction and summarization steps read their templates (including
                 ``extract_memories.prompty``) from this directory instead of the toolkit's
                 bundled defaults, letting you customize what the extraction LLM produces. The
-                directory must contain the full template set. Applies whether the client is built
-                by the provider or supplied via ``memory_client``.
+                directory must contain the full template set. Only applied when the provider builds
+                the client; ignored with ``memory_client``.
 
         Raises:
             SettingNotFoundError: If ``cosmos_endpoint``, ``foundry_endpoint``, ``embedding_model``,
@@ -163,34 +163,24 @@ def __init__(
         self.memory_types: list[MemoryType] = list(memory_types) if memory_types else ["fact", "procedural"]
         self.context_prompt = context_prompt
         self.auto_extract = auto_extract
-        self._prompts_dir = prompts_dir
-
-        # Build the per-instance cadence override for the toolkit client. The Agent Memory Toolkit
-        # accepts these thresholds directly via ``cadence_thresholds=`` (v0.2.0b3+), so the provider
-        # configures the processor without mutating global ``os.environ``. ``auto_extract=False``
-        # zeroes the extraction/summary steps so the toolkit's background auto-trigger never runs on
-        # turn writes; callers then drive processing explicitly via ``memory_client.process_now(...)``.
-        # Keys not present fall back to the toolkit's environment/defaults.
-        cadence_thresholds: dict[str, int] = {
-            str(k): int(v) for k, v in cast("Mapping[str, int]", processor_config or {}).items()
-        }
-        if not auto_extract:
-            cadence_thresholds["FACT_EXTRACTION_EVERY_N"] = 0
-            cadence_thresholds["THREAD_SUMMARY_EVERY_N"] = 0
-            cadence_thresholds["USER_SUMMARY_EVERY_N"] = 0
-
-        # A caller-supplied client owns its own cadence configuration; the provider cannot apply
-        # ``cadence_thresholds`` to an already-constructed client. Reject the combination instead of
-        # silently ignoring the requested configuration.
-        if memory_client is not None and cadence_thresholds:
-            raise ValueError(
-                "processor_config and auto_extract=False only take effect when the provider builds "
-                "the memory client. When supplying your own memory_client, configure cadence via "
-                "AsyncCosmosMemoryClient(cadence_thresholds=...) directly."
-            )
+        self._prompts_dir = prompts_dir if memory_client is None else None
 
-        # Initialize memory client if not provided
+        # Caller-supplied clients retain their own connection, cadence and prompt configuration.
         if memory_client is None:
+            # Build the per-instance cadence override for the toolkit client. The Agent Memory Toolkit
+            # accepts these thresholds directly via ``cadence_thresholds=`` (v0.2.0b3+), so the provider
+            # configures the processor without mutating global ``os.environ``. ``auto_extract=False``
+            # zeroes the extraction/summary steps so the toolkit's background auto-trigger never runs on
+            # turn writes; callers then drive processing explicitly via ``memory_client.process_now(...)``.
+            # Keys not present fall back to the toolkit's environment/defaults.
+            cadence_thresholds: dict[str, int] = {
+                str(k): int(v) for k, v in cast("Mapping[str, int]", processor_config or {}).items()
+            }
+            if not auto_extract:
+                cadence_thresholds["FACT_EXTRACTION_EVERY_N"] = 0
+                cadence_thresholds["THREAD_SUMMARY_EVERY_N"] = 0
+                cadence_thresholds["USER_SUMMARY_EVERY_N"] = 0
+
             # Resolve connection settings from explicit args, then the environment. ``load_settings``
             # validates that the required endpoints are present (raising if not), replacing manual
             # ``os.getenv`` + ``if not ...: raise`` blocks.
@@ -512,13 +502,13 @@ async def after_run(
                                 content=msg.text.strip(),
                             )
 
-            # Auto-extraction and processing:
+            # Auto-extraction and processing for provider-created clients:
             # When auto_extract is True (default), the turn write schedules cadence-aware background
             # processing (fact extraction, summaries, reconciliation) based on the configured
             # thresholds (FACT_EXTRACTION_EVERY_N, DEDUP_EVERY_N, etc.), so no explicit
             # process_now() call is needed. When auto_extract is False, those thresholds were
             # zeroed in __init__ so nothing runs automatically; call memory_client.process_now()
-            # to drive extraction manually.
+            # to drive extraction manually. Injected clients retain their own cadence settings.
 
         except Exception as e:
             logger.warning("Failed to store conversation turns: %s", e, exc_info=True)

python/packages/azure-cosmos-memory/tests/test_context_provider.py

@@ -349,17 +349,50 @@ def test_default_cadence_thresholds_is_none(self) -> None:
             _, kwargs = mock_client_class.call_args
             assert kwargs["cadence_thresholds"] is None
 
-    def test_processor_config_with_supplied_client_raises(self, mock_memory_client: AsyncMock) -> None:
-        """Cadence config cannot apply to a caller-supplied client, so combining them raises."""
-        with pytest.raises(ValueError, match="processor_config"):
-            CosmosMemoryContextProvider(
-                memory_client=mock_memory_client, processor_config={"FACT_EXTRACTION_EVERY_N": 10}
+    @pytest.mark.parametrize("auto_extract", [True, False])
+    def test_supplied_client_ignores_construction_options(
+        self, mock_memory_client: AsyncMock, auto_extract: bool, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Injected clients bypass connection resolution and cadence configuration."""
+        for name in (
+            "COSMOS_ENDPOINT",
+            "FOUNDRY_PROJECT_ENDPOINT",
+            "FOUNDRY_ENDPOINT",
+            "EMBEDDING_MODEL",
+            "CHAT_MODEL",
+        ):
+            monkeypatch.delenv(name, raising=False)
+        with (
+            patch("agent_framework_azure_cosmos_memory._context_provider.load_settings") as mock_load_settings,
+            patch("agent_framework_azure_cosmos_memory._context_provider.AsyncCosmosMemoryClient") as mock_client_class,
+        ):
+            provider = CosmosMemoryContextProvider(
+                memory_client=mock_memory_client,
+                processor_config={"FACT_EXTRACTION_EVERY_N": 2},
+                auto_extract=auto_extract,
+                top_k=7,
+                min_confidence=0.8,
+                memory_types=["episodic"],
+                context_prompt="Selected memories:",
             )
 
-    def test_auto_extract_false_with_supplied_client_raises(self, mock_memory_client: AsyncMock) -> None:
-        """auto_extract=False cannot apply to a caller-supplied client, so combining them raises."""
-        with pytest.raises(ValueError, match="processor_config"):
-            CosmosMemoryContextProvider(memory_client=mock_memory_client, auto_extract=False)
+        assert provider.memory_client is mock_memory_client
+        assert provider.top_k == 7
+        assert provider.min_confidence == 0.8
+        assert provider.memory_types == ["episodic"]
+        assert provider.context_prompt == "Selected memories:"
+        mock_load_settings.assert_not_called()
+        mock_client_class.assert_not_called()
+
+    def test_supplied_client_does_not_parse_cadence(self, mock_memory_client: AsyncMock) -> None:
+        """Construction-only cadence values are not evaluated for injected clients."""
+        cadence = MagicMock()
+        cadence.items.side_effect = AssertionError("cadence should not be inspected")
+
+        provider = CosmosMemoryContextProvider(memory_client=mock_memory_client, processor_config=cadence)
+
+        assert provider.memory_client is mock_memory_client
+        cadence.items.assert_not_called()
 
 
 # -- before_run tests ----------------------------------------------------------
@@ -955,9 +988,12 @@ async def test_provided_client_not_closed(self, mock_memory_client: AsyncMock) -
         """When client is provided externally, provider should not close it."""
         provider = CosmosMemoryContextProvider(memory_client=mock_memory_client)
 
-        async with provider:
-            pass
+        with patch.object(provider, "flush", new_callable=AsyncMock) as mock_flush:
+            async with provider:
+                pass
 
+        mock_flush.assert_awaited_once()
+        mock_memory_client.close.assert_not_awaited()
         # Should still enter the client
         mock_memory_client.__aenter__.assert_awaited_once()
         # But should NOT exit it (caller owns it)
@@ -1052,12 +1088,19 @@ async def test_prompts_dir_redirects_pipeline_loader(self, mock_memory_client: A
         # _get_pipeline is synchronous on the toolkit client; return our stand-in pipeline.
         mock_memory_client._get_pipeline = MagicMock(return_value=mock_pipeline)
 
-        provider = CosmosMemoryContextProvider(
-            memory_client=mock_memory_client,
-            prompts_dir="/custom/prompts",
-        )
-        async with provider:
-            pass
+        with patch(
+            "agent_framework_azure_cosmos_memory._context_provider.AsyncCosmosMemoryClient",
+            return_value=mock_memory_client,
+        ):
+            provider = CosmosMemoryContextProvider(
+                cosmos_endpoint="https://test.documents.azure.com:443/",
+                foundry_endpoint="https://test.ai.azure.com",
+                embedding_model="text-embedding-3-large",
+                chat_model="gpt-4o-mini",
+                prompts_dir="/custom/prompts",
+            )
+            async with provider:
+                pass
 
         from azure.cosmos.agent_memory.services._pipeline_helpers import PromptyLoader
 
@@ -1074,3 +1117,15 @@ async def test_no_prompts_dir_leaves_pipeline_untouched(self, mock_memory_client
             pass
 
         mock_memory_client._get_pipeline.assert_not_called()
+
+    async def test_supplied_client_ignores_prompts_dir(self, mock_memory_client: AsyncMock) -> None:
+        """Custom templates must not replace the caller's existing pipeline loader."""
+        mock_memory_client._get_pipeline = MagicMock()
+        provider = CosmosMemoryContextProvider(memory_client=mock_memory_client, prompts_dir="/custom/prompts")
+
+        async with provider:
+            pass
+
+        mock_memory_client._get_pipeline.assert_not_called()
+        mock_memory_client.__aexit__.assert_not_awaited()
+        mock_memory_client.close.assert_not_awaited()

python/packages/azure-cosmos-memory/tests/test_emulator.py

@@ -335,9 +335,22 @@ async def test_prompts_dir_overrides_extraction_prompt(
         custom_dir = _make_custom_prompts_dir(tmp_path / "prompts", marker)
 
         chat = _FakeChat()
-        client = _build_emulator_client(monkeypatch, chat)
+
+        def _create_emulator_client(**kwargs: Any) -> AsyncCosmosMemoryClient:
+            return _build_emulator_client(monkeypatch, chat)
+
+        # Keep the real emulator-backed toolkit client, but let the provider create and own it
+        # so its prompts_dir option applies. Foundry calls still use the deterministic fakes.
+        monkeypatch.setattr(
+            "agent_framework_azure_cosmos_memory._context_provider.AsyncCosmosMemoryClient",
+            _create_emulator_client,
+        )
         provider = CosmosMemoryContextProvider(
-            memory_client=client,
+            cosmos_endpoint=_EMULATOR_ENDPOINT,
+            cosmos_database="test_af_mem",
+            foundry_endpoint="https://test.ai.azure.com",
+            embedding_model="test-embedding",
+            chat_model="test-chat",
             top_k=5,
             min_confidence=0.0,
             memory_types=["fact"],
@@ -346,7 +359,7 @@ async def test_prompts_dir_overrides_extraction_prompt(
         try:
             await provider.__aenter__()
         except Exception as exc:  # noqa: BLE001 - clear skip on any connectivity/setup failure
-            await client.close()
+            await provider.memory_client.close()
             pytest.skip(f"Cosmos DB emulator not reachable or vector search unavailable at {_EMULATOR_ENDPOINT}: {exc}")
 
         try:
@@ -372,4 +385,3 @@ async def test_prompts_dir_overrides_extraction_prompt(
         # in the overridden template, so its presence proves prompts_dir was honored end to end.
         sent = "\n".join(str(msg.get("content", "")) for call in chat.calls for msg in call)
         assert marker in sent, "custom extract_memories.prompty was not used by the extraction pipeline"
-        await client.close()