Python: Fix Cosmos memory provider calling renamed add_cosmos toolkit API

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

@@ -435,13 +435,17 @@ async def after_run(
         user_id = self._resolve_user_id(state, session)
         thread_id = state.get("thread_id") or session.session_id or "default"
 
+        # TODO(atty57): The toolkit renamed add_cosmos -> upsert_memory (same kwargs); accept either
+        # until the declared azure-cosmos-agent-memory floor is past the rename, then inline it.
+        write_turn = getattr(self.memory_client, "upsert_memory", None) or self.memory_client.add_cosmos
+
         try:
             # Store input messages (skip empty/whitespace-only content to avoid junk turns)
             for msg in context.input_messages:
                 if hasattr(msg, "role") and hasattr(msg, "text") and msg.text and msg.text.strip():
                     role_value = getattr(msg.role, "value", None) or str(msg.role)
                     if role_value in {"user", "assistant", "system"}:
-                        await self.memory_client.add_cosmos(
+                        await write_turn(
                             user_id=user_id,
                             thread_id=thread_id,
                             role=self._ROLE_MAP.get(role_value, role_value),
@@ -454,15 +458,15 @@ async def after_run(
                     if hasattr(msg, "role") and hasattr(msg, "text") and msg.text and msg.text.strip():
                         role_value = getattr(msg.role, "value", None) or str(msg.role)
                         if role_value in {"user", "assistant", "system"}:
-                            await self.memory_client.add_cosmos(
+                            await write_turn(
                                 user_id=user_id,
                                 thread_id=thread_id,
                                 role=self._ROLE_MAP.get(role_value, role_value),
                                 content=msg.text.strip(),
                             )
 
             # Auto-extraction and processing:
-            # When auto_extract is True (default), add_cosmos() schedules cadence-aware background
+            # 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

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

@@ -52,7 +52,7 @@ def mock_memory_client() -> AsyncMock:
     mock_client = AsyncMock()
     mock_client.search_cosmos = AsyncMock(return_value=[])
     mock_client.get_user_summary = AsyncMock(return_value=None)
-    mock_client.add_cosmos = AsyncMock()
+    mock_client.upsert_memory = AsyncMock()
     mock_client.create_memory_store = AsyncMock()
     mock_client.__aenter__ = AsyncMock(return_value=mock_client)
     mock_client.__aexit__ = AsyncMock()
@@ -468,8 +468,8 @@ async def test_stores_input_and_response_messages(self, mock_memory_client: Asyn
             agent=_STUB_AGENT, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
         )
 
-        assert mock_memory_client.add_cosmos.await_count == 2
-        calls = mock_memory_client.add_cosmos.await_args_list
+        assert mock_memory_client.upsert_memory.await_count == 2
+        calls = mock_memory_client.upsert_memory.await_args_list
 
         # Check input message stored
         assert calls[0].kwargs["role"] == "user"
@@ -499,7 +499,7 @@ async def test_assistant_role_mapped_to_agent(self, mock_memory_client: AsyncMoc
             agent=_STUB_AGENT, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
         )
 
-        stored_roles = [c.kwargs["role"] for c in mock_memory_client.add_cosmos.await_args_list]
+        stored_roles = [c.kwargs["role"] for c in mock_memory_client.upsert_memory.await_args_list]
         assert stored_roles == ["user", "agent"]
         # No raw "assistant" role should ever be sent to the toolkit.
         assert "assistant" not in stored_roles
@@ -520,7 +520,7 @@ async def test_uses_custom_user_and_thread_ids(self, mock_memory_client: AsyncMo
             agent=_STUB_AGENT, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
         )
 
-        call_kwargs = mock_memory_client.add_cosmos.await_args_list[0].kwargs
+        call_kwargs = mock_memory_client.upsert_memory.await_args_list[0].kwargs
         assert call_kwargs["user_id"] == "user-456"
         assert call_kwargs["thread_id"] == "thread-789"
 
@@ -541,8 +541,8 @@ async def test_skips_empty_messages(self, mock_memory_client: AsyncMock) -> None
         )
 
         # Only one message should be stored
-        assert mock_memory_client.add_cosmos.await_count == 1
-        call_kwargs = mock_memory_client.add_cosmos.await_args_list[0].kwargs
+        assert mock_memory_client.upsert_memory.await_count == 1
+        call_kwargs = mock_memory_client.upsert_memory.await_args_list[0].kwargs
         assert call_kwargs["content"] == "Valid message"
 
     async def test_skips_whitespace_only_messages(self, mock_memory_client: AsyncMock) -> None:
@@ -563,15 +563,35 @@ async def test_skips_whitespace_only_messages(self, mock_memory_client: AsyncMoc
         )
 
         # Whitespace-only input and the whitespace-only response are both skipped.
-        assert mock_memory_client.add_cosmos.await_count == 1
-        call_kwargs = mock_memory_client.add_cosmos.await_args_list[0].kwargs
+        assert mock_memory_client.upsert_memory.await_count == 1
+        call_kwargs = mock_memory_client.upsert_memory.await_args_list[0].kwargs
         assert call_kwargs["content"] == "Trimmed message"
 
+    async def test_falls_back_to_add_cosmos_on_older_toolkit(self) -> None:
+        """Toolkit versions predating the upsert_memory rename still receive turns.
+
+        The declared azure-cosmos-agent-memory range spans both names, so a resolved
+        install can expose either one; picking neither would silently drop every turn.
+        """
+        legacy_client = AsyncMock(spec=["add_cosmos", "search_cosmos", "get_user_summary"])
+        legacy_client.add_cosmos = AsyncMock()
+
+        provider = CosmosMemoryContextProvider(memory_client=legacy_client)
+        session = AgentSession(session_id="test-session")
+        ctx = SessionContext(input_messages=[Message(role="user", contents=["Hello"])], session_id="s1")
+
+        await provider.after_run(
+            agent=_STUB_AGENT, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
+        )
+
+        assert legacy_client.add_cosmos.await_count == 1
+        assert legacy_client.add_cosmos.await_args_list[0].kwargs["content"] == "Hello"
+
     async def test_storage_failure_logs_warning(
         self, mock_memory_client: AsyncMock, caplog: pytest.LogCaptureFixture
     ) -> None:
         """Storage failures are logged but don't raise."""
-        mock_memory_client.add_cosmos.side_effect = Exception("Storage failed")
+        mock_memory_client.upsert_memory.side_effect = Exception("Storage failed")
 
         provider = CosmosMemoryContextProvider(memory_client=mock_memory_client)
         session = AgentSession(session_id="test-session")

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

@@ -183,7 +183,8 @@ async def test_before_run_retrieves_seeded_fact(self, emulator_provider: CosmosM
         # embeddings client). This lands in the memories container under the quantizedFlat
         # vector index, without needing LLM extraction.
         assert provider.memory_client is not None
-        await provider.memory_client.add_cosmos(
+        seed = getattr(provider.memory_client, "upsert_memory", None) or provider.memory_client.add_cosmos
+        await seed(
             user_id=user_id,
             thread_id=thread_id,
             role="user",