Fix deduplication for "message" type long-term memories

The memory server indexes any messages from working memory into
"message" type memory records in long-term memory. However, we
were not attempting to deduplicate these memory records -- only
the higher-level long-term memory types, like semantic and episodic
memory.

This change refactors the logic we use to store messages long-term
such that they work similarly to long-term memories the client adds
to working memory. When a client adds a message to working memory,
the client gives the memory an ID, and the memory gets a blank
`persisted_at` timestamp. After setting working memory, we kick
off a background task to persist any promoted long-term memories
and messages in long-term memory. When we copy the messages into
long-term memory, we now update the `persisted_at` timestamp. If
we see that working memory again and try to persist any new memories
or messages, we'll skip any that have already been persisted.

Author: abrookins

agent-memory-client/agent_memory_client/client.py

@@ -31,6 +31,7 @@
     AckResponse,
     ClientMemoryRecord,
     HealthCheckResponse,
+    MemoryMessage,
     MemoryRecord,
     MemoryRecordResults,
     MemoryTypeEnum,
@@ -2055,7 +2056,7 @@ async def update_working_memory_data(
     async def append_messages_to_working_memory(
         self,
         session_id: str,
-        messages: list[dict[str, Any]],  # Expect proper message dicts
+        messages: list[dict[str, Any] | MemoryMessage],
         namespace: str | None = None,
         model_name: str | None = None,
         context_window_max: int | None = None,
@@ -2068,7 +2069,7 @@ async def append_messages_to_working_memory(
 
         Args:
             session_id: Target session
-            messages: List of message dictionaries with 'role' and 'content' keys
+            messages: List of message dictionaries or MemoryMessage objects
             namespace: Optional namespace
             model_name: Optional model name for token-based summarization
             context_window_max: Optional direct specification of context window max tokens
@@ -2081,19 +2082,36 @@ async def append_messages_to_working_memory(
             session_id=session_id, namespace=namespace, user_id=user_id
         )
 
-        # Validate new messages have required structure
+        # Convert messages to MemoryMessage objects
+        converted_messages = []
         for msg in messages:
-            if not isinstance(msg, dict) or "role" not in msg or "content" not in msg:
+            if isinstance(msg, MemoryMessage):
+                converted_messages.append(msg)
+            elif isinstance(msg, dict):
+                if "role" not in msg or "content" not in msg:
+                    raise ValueError("All messages must have 'role' and 'content' keys")
+                # Build message kwargs, only including non-None values
+                message_kwargs = {
+                    "role": msg["role"],
+                    "content": msg["content"],
+                }
+                if msg.get("id") is not None:
+                    message_kwargs["id"] = msg["id"]
+                if msg.get("persisted_at") is not None:
+                    message_kwargs["persisted_at"] = msg["persisted_at"]
+
+                converted_messages.append(MemoryMessage(**message_kwargs))
+            else:
                 raise ValueError(
-                    "All messages must be dictionaries with 'role' and 'content' keys"
+                    "All messages must be dictionaries or MemoryMessage objects"
                 )
 
-        # Get existing messages (already in proper dict format from get_working_memory)
+        # Get existing messages
         existing_messages = []
         if existing_memory and existing_memory.messages:
             existing_messages = existing_memory.messages
 
-        final_messages = existing_messages + messages
+        final_messages = existing_messages + converted_messages
 
         # Create updated working memory
         working_memory = WorkingMemory(

agent-memory-client/agent_memory_client/models.py

@@ -7,7 +7,7 @@
 
 from datetime import datetime, timezone
 from enum import Enum
-from typing import Any, Literal, TypedDict
+from typing import Any, Literal
 
 from pydantic import BaseModel, Field
 from ulid import ULID
@@ -48,11 +48,19 @@ class MemoryTypeEnum(str, Enum):
     MESSAGE = "message"
 
 
-class MemoryMessage(TypedDict):
+class MemoryMessage(BaseModel):
     """A message in the memory system"""
 
     role: str
     content: str
+    id: str = Field(
+        default_factory=lambda: str(ULID()),
+        description="Unique identifier for the message (auto-generated)",
+    )
+    persisted_at: datetime | None = Field(
+        default=None,
+        description="Server-assigned timestamp when message was persisted to long-term storage",
+    )
 
 
 class MemoryRecord(BaseModel):
@@ -134,9 +142,9 @@ class WorkingMemory(BaseModel):
     """Working memory for a session - contains both messages and structured memory records"""
 
     # Support both message-based memory (conversation) and structured memory records
-    messages: list[dict[str, Any]] = Field(
+    messages: list[MemoryMessage] = Field(
         default_factory=list,
-        description="Conversation messages (role/content pairs)",
+        description="Conversation messages with tracking fields",
     )
     memories: list[MemoryRecord | ClientMemoryRecord] = Field(
         default_factory=list,

agent_memory_server/api.py

@@ -2,7 +2,6 @@
 from fastapi import APIRouter, Depends, HTTPException, Query
 from mcp.server.fastmcp.prompts import base
 from mcp.types import TextContent
-from ulid import ULID
 
 from agent_memory_server import long_term_memory, working_memory
 from agent_memory_server.auth import UserInfo, get_current_user
@@ -18,7 +17,6 @@
     MemoryPromptRequest,
     MemoryPromptResponse,
     MemoryRecordResultsResponse,
-    MemoryTypeEnum,
     ModelNameLiteral,
     SearchRequest,
     SessionListResponse,
@@ -329,35 +327,13 @@ async def put_working_memory(
     )
 
     # Background tasks for long-term memory promotion and indexing (if enabled)
-    if settings.long_term_memory:
+    if settings.long_term_memory and updated_memory.memories:
         # Promote structured memories from working memory to long-term storage
-        if updated_memory.memories:
-            await background_tasks.add_task(
-                long_term_memory.promote_working_memory_to_long_term,
-                session_id,
-                updated_memory.namespace,
-            )
-
-        # Index message-based memories
-        if updated_memory.messages:
-            from agent_memory_server.models import MemoryRecord
-
-            memories = [
-                MemoryRecord(
-                    id=str(ULID()),
-                    session_id=session_id,
-                    text=f"{msg.role}: {msg.content}",
-                    namespace=updated_memory.namespace,
-                    user_id=updated_memory.user_id,
-                    memory_type=MemoryTypeEnum.MESSAGE,
-                )
-                for msg in updated_memory.messages
-            ]
-
-            await background_tasks.add_task(
-                long_term_memory.index_long_term_memories,
-                memories,
-            )
+        await background_tasks.add_task(
+            long_term_memory.promote_working_memory_to_long_term,
+            session_id,
+            updated_memory.namespace,
+        )
 
     return updated_memory
 

agent_memory_server/long_term_memory.py

@@ -1247,12 +1247,19 @@ async def promote_working_memory_to_long_term(
         if memory.persisted_at is None
     ]
 
-    if not unpersisted_memories:
-        logger.debug(f"No unpersisted memories found in session {session_id}")
+    # Find unpersisted messages (similar to unpersisted memories)
+    unpersisted_messages = [
+        msg for msg in current_working_memory.messages if msg.persisted_at is None
+    ]
+
+    if not unpersisted_memories and not unpersisted_messages:
+        logger.debug(
+            f"No unpersisted memories or messages found in session {session_id}"
+        )
         return 0
 
     logger.info(
-        f"Promoting {len(unpersisted_memories)} memories from session {session_id}"
+        f"Promoting {len(unpersisted_memories)} memories and {len(unpersisted_messages)} messages from session {session_id}"
     )
 
     promoted_count = 0
@@ -1317,10 +1324,57 @@ async def promote_working_memory_to_long_term(
         )
         updated_memories.extend(extracted_memories)
 
+    # Process unpersisted messages
+    updated_messages = []
+    for msg in current_working_memory.messages:
+        if msg.persisted_at is None:
+            # Generate ID if not present (backward compatibility)
+            if not msg.id:
+                msg.id = str(ULID())
+
+            memory_record = MemoryRecord(
+                id=msg.id,
+                session_id=session_id,
+                text=f"{msg.role}: {msg.content}",
+                namespace=namespace,
+                user_id=current_working_memory.user_id,
+                memory_type=MemoryTypeEnum.MESSAGE,
+                persisted_at=None,
+            )
+
+            # Apply same deduplication logic as structured memories
+            deduped_memory, was_overwrite = await deduplicate_by_id(
+                memory=memory_record,
+                redis_client=redis,
+            )
+
+            # Set persisted_at timestamp
+            current_memory = deduped_memory or memory_record
+            current_memory.persisted_at = datetime.now(UTC)
+
+            # Index in long-term storage
+            await index_long_term_memories(
+                [current_memory],
+                redis_client=redis,
+                deduplicate=False,  # Already deduplicated by ID
+            )
+
+            # Update message with persisted_at timestamp
+            msg.persisted_at = current_memory.persisted_at
+            promoted_count += 1
+
+            if was_overwrite:
+                logger.info(f"Overwrote existing message with id {msg.id}")
+            else:
+                logger.info(f"Promoted new message with id {msg.id}")
+
+        updated_messages.append(msg)
+
     # Update working memory with the new persisted_at timestamps and extracted memories
     if promoted_count > 0 or extracted_memories:
         updated_working_memory = current_working_memory.model_copy()
         updated_working_memory.memories = updated_memories
+        updated_working_memory.messages = updated_messages
         updated_working_memory.updated_at = datetime.now(UTC)
 
         await working_memory.set_working_memory(
@@ -1329,7 +1383,7 @@ async def promote_working_memory_to_long_term(
         )
 
         logger.info(
-            f"Successfully promoted {promoted_count} memories to long-term storage"
+            f"Successfully promoted {promoted_count} memories and messages to long-term storage"
             + (
                 f" and extracted {len(extracted_memories)} new memories"
                 if extracted_memories

agent_memory_server/mcp.py

@@ -659,7 +659,25 @@ async def set_working_memory(
     )
     ```
 
-    4. Replace entire working memory state:
+    4. Store conversation messages:
+    ```python
+    set_working_memory(
+        session_id="current_session",
+        messages=[
+            {
+                "role": "user",
+                "content": "What is the weather like?",
+                "id": "msg_001"  # Optional - auto-generated if not provided
+            },
+            {
+                "role": "assistant",
+                "content": "I'll check the weather for you."
+            }
+        ]
+    )
+    ```
+
+    5. Replace entire working memory state:
     ```python
     set_working_memory(
         session_id="current_session",
@@ -673,7 +691,7 @@ async def set_working_memory(
     Args:
         session_id: The session ID to set memory for (required)
         memories: List of structured memory records (semantic, episodic, message types)
-        messages: List of conversation messages (role/content pairs)
+        messages: List of conversation messages (role/content pairs with optional id/persisted_at)
         context: Optional summary/context text
         data: Optional dictionary for storing arbitrary JSON data
         namespace: Optional namespace for scoping
@@ -712,12 +730,35 @@ async def set_working_memory(
 
             processed_memories.append(processed_memory)
 
+    # Process messages to ensure proper format
+    processed_messages = []
+    if messages:
+        for message in messages:
+            # Handle both MemoryMessage objects and dict inputs
+            if isinstance(message, MemoryMessage):
+                # Already a MemoryMessage object, ensure persisted_at is None for new messages
+                processed_message = message.model_copy(
+                    update={
+                        "persisted_at": None,  # Mark as pending promotion
+                    }
+                )
+            else:
+                # Dictionary input, convert to MemoryMessage
+                message_dict = dict(message)
+                # Remove id=None to allow auto-generation
+                if message_dict.get("id") is None:
+                    message_dict.pop("id", None)
+                message_dict["persisted_at"] = None
+                processed_message = MemoryMessage(**message_dict)
+
+            processed_messages.append(processed_message)
+
     # Create the working memory object
     working_memory_obj = WorkingMemory(
         session_id=session_id,
         namespace=memory_namespace,
         memories=processed_memories,
-        messages=messages or [],
+        messages=processed_messages,
         context=context,
         data=data or {},
         user_id=user_id,

agent_memory_server/models.py

@@ -67,6 +67,14 @@ class MemoryMessage(BaseModel):
 
     role: str
     content: str
+    id: str = Field(
+        default_factory=lambda: str(ULID()),
+        description="Unique identifier for the message (auto-generated if not provided)",
+    )
+    persisted_at: datetime | None = Field(
+        default=None,
+        description="Server-assigned timestamp when message was persisted to long-term storage",
+    )
 
 
 class SessionListResponse(BaseModel):

tests/test_api.py

@@ -89,10 +89,22 @@ async def test_get_memory(self, client, session):
 
         data = response.json()
         response = WorkingMemoryResponse(**data)
-        assert response.messages == [
-            MemoryMessage(role="user", content="Hello"),
-            MemoryMessage(role="assistant", content="Hi there"),
-        ]
+
+        # Check that we have 2 messages with correct roles and content
+        assert len(response.messages) == 2
+
+        # Check message content and roles (IDs are auto-generated so we can't compare directly)
+        message_contents = [msg.content for msg in response.messages]
+        message_roles = [msg.role for msg in response.messages]
+        assert "Hello" in message_contents
+        assert "Hi there" in message_contents
+        assert "user" in message_roles
+        assert "assistant" in message_roles
+
+        # Check that all messages have IDs (auto-generated)
+        for msg in response.messages:
+            assert msg.id is not None
+            assert len(msg.id) > 0
 
         roles = [msg["role"] for msg in data["messages"]]
         contents = [msg["content"] for msg in data["messages"]]

tests/test_client_api.py

@@ -116,8 +116,8 @@ async def test_session_lifecycle(memory_test_client: MemoryAPIClient):
 
         # Step 1: Create new session memory
         response = await memory_test_client.put_working_memory(session_id, memory)
-        assert response.messages[0]["content"] == "Hello from the client!"
-        assert response.messages[1]["content"] == "Hi there, I'm the memory server!"
+        assert response.messages[0].content == "Hello from the client!"
+        assert response.messages[1].content == "Hi there, I'm the memory server!"
         assert response.context == "This is a test session created by the API client."
 
     # Next, mock GET response for retrieving session memory
@@ -132,8 +132,8 @@ async def test_session_lifecycle(memory_test_client: MemoryAPIClient):
         # Step 2: Retrieve the session memory
         session = await memory_test_client.get_working_memory(session_id)
         assert len(session.messages) == 2
-        assert session.messages[0]["content"] == "Hello from the client!"
-        assert session.messages[1]["content"] == "Hi there, I'm the memory server!"
+        assert session.messages[0].content == "Hello from the client!"
+        assert session.messages[1].content == "Hi there, I'm the memory server!"
         assert session.context == "This is a test session created by the API client."
 
     # Mock list sessions