From 6b5b19c2ebf5b8d874d009cd5d5c9dc0e4ccf685 Mon Sep 17 00:00:00 2001 From: atty57 <99388680+atty57@users.noreply.github.com> Date: Wed, 12 Aug 2026 15:30:49 -0400 Subject: [PATCH 1/3] Python: Fix Cosmos memory provider calling renamed add_cosmos toolkit API The Agent Memory Toolkit renamed AsyncCosmosMemoryClient.add_cosmos to upsert_memory with an identical signature. The provider declares azure-cosmos-agent-memory>=0.2.0b3 with no upper bound, so a resolved install can expose either name. after_run swallows write errors and only logs a warning, so on a post-rename toolkit the agent turn still looks successful while long-term memory silently stops receiving turns. Resolve the write method once per after_run, preferring upsert_memory and falling back to add_cosmos, so both ends of the declared range keep working. Same treatment for the emulator test's direct seed call. Fixes #7633 --- .../_context_provider.py | 10 +++-- .../tests/test_context_provider.py | 40 ++++++++++++++----- .../tests/test_emulator.py | 3 +- 3 files changed, 39 insertions(+), 14 deletions(-) diff --git a/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py b/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py index 3798c3bad8..b804f3d6b0 100644 --- a/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py +++ b/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" + # ponytail: 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,7 +458,7 @@ 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), @@ -462,7 +466,7 @@ async def after_run( ) # 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 diff --git a/python/packages/azure-cosmos-memory/tests/test_context_provider.py b/python/packages/azure-cosmos-memory/tests/test_context_provider.py index ae682b237b..ebee323234 100644 --- a/python/packages/azure-cosmos-memory/tests/test_context_provider.py +++ b/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") diff --git a/python/packages/azure-cosmos-memory/tests/test_emulator.py b/python/packages/azure-cosmos-memory/tests/test_emulator.py index a456f2a19c..caad815db4 100644 --- a/python/packages/azure-cosmos-memory/tests/test_emulator.py +++ b/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", From e0e017a2f2dd7145f63b475f3480aac843fcad03 Mon Sep 17 00:00:00 2001 From: Atharva Vichare <99388680+atty57@users.noreply.github.com> Date: Wed, 12 Aug 2026 15:35:41 -0400 Subject: [PATCH 2/3] Ponytail comment erased Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../agent_framework_azure_cosmos_memory/_context_provider.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py b/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py index b804f3d6b0..5e01312cc2 100644 --- a/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py +++ b/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py @@ -435,7 +435,7 @@ async def after_run( user_id = self._resolve_user_id(state, session) thread_id = state.get("thread_id") or session.session_id or "default" - # ponytail: the toolkit renamed add_cosmos -> upsert_memory (same kwargs); accept either + # TODO: 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 From f22476c7ee65ef91962131b6704fd9aa1ba19b88 Mon Sep 17 00:00:00 2001 From: Atharva Vichare <99388680+atty57@users.noreply.github.com> Date: Wed, 12 Aug 2026 21:31:18 -0400 Subject: [PATCH 3/3] Clarify TODO comment regarding memory method rename Updated TODO comment to include author and clarify context , to resolve linting error --- .../agent_framework_azure_cosmos_memory/_context_provider.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py b/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py index 5e01312cc2..9fa93db059 100644 --- a/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py +++ b/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py @@ -435,7 +435,7 @@ 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: The toolkit renamed add_cosmos -> upsert_memory (same kwargs); accept either + # 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