fix: resolve working memory key via search index when user_id/namespace omitted (#239)

* fix: resolve working memory key via search index when user_id/namespace omitted (#235)

PUT /v1/working-memory/{session_id} stores user_id and namespace from
the request body into the Redis key, but GET/DELETE construct the key
from query params. When callers omit those params, a different key is
produced and the session appears missing (404).

Add an index-based fallback: when a direct key lookup returns nothing,
query the working memory search index by session_id to resolve the
actual Redis key. This preserves multi-tenant isolation (namespace +
user_id remain in keys) while making GET/DELETE work without requiring
callers to repeat scoping parameters.

Changes:
- Add _resolve_working_memory_key_via_index() helper in working_memory.py
- Use fallback in get_working_memory() and delete_working_memory()
- Return stored namespace/user_id from JSON data instead of caller params
- Add user_id/namespace params to MCP get_working_memory tool
- Add comprehensive tests covering the fix and multi-tenant isolation

* Update agent_memory_server/working_memory.py

Co-authored-by: Vishal Bala <vishal-bala@users.noreply.github.com>

* Fix linter issues

---------

Co-authored-by: Tyler Hutcherson <tyler.hutcherson@redis.com>
Co-authored-by: Vishal Bala <vishal-bala@users.noreply.github.com>
Co-authored-by: Vishal Bala <vishalbala.1994@gmail.com>

Author: 3 people

agent_memory_server/mcp.py

@@ -979,20 +979,27 @@ async def set_working_memory(
 @mcp_app.tool()
 async def get_working_memory(
     session_id: str,
+    user_id: str | None = None,
+    namespace: str | None = None,
     recent_messages_limit: int | None = None,
 ) -> WorkingMemory:
     """
     Get working memory for a session. This works like the GET /sessions/{id}/memory API endpoint.
 
     Args:
         session_id: The session ID to retrieve working memory for
+        user_id: Optional user ID to scope the session lookup
+        namespace: Optional namespace to scope the session lookup
         recent_messages_limit: Optional limit on number of recent messages to return (most recent first)
 
     Returns:
         Working memory containing messages, context, and structured memory records
     """
     result = await working_memory_core.get_working_memory(
-        session_id=session_id, recent_messages_limit=recent_messages_limit
+        session_id=session_id,
+        user_id=user_id,
+        namespace=namespace,
+        recent_messages_limit=recent_messages_limit,
     )
     if result is None:
         return WorkingMemory(session_id=session_id, messages=[], memories=[])

agent_memory_server/working_memory.py

@@ -356,6 +356,84 @@ async def list_sessions(
         return 0, []
 
 
+async def _resolve_working_memory_key_via_index(
+    redis_client: Redis,
+    session_id: str,
+    user_id: str | None = None,
+    namespace: str | None = None,
+) -> str | None:
+    """
+    Resolve the actual Redis key for a working memory session using the search index.
+
+    When a direct key lookup fails (e.g. because user_id/namespace were provided
+    during PUT but omitted during GET), this function queries the search index
+    by session_id to find the document and derive the correct Redis key.
+
+    If multiple sessions share the same session_id (different namespace/user_id)
+    and the caller did not supply enough filters to disambiguate, the function
+    logs a warning and returns ``None`` to avoid silently returning the wrong
+    session.
+
+    Args:
+        redis_client: Redis client
+        session_id: The session ID to look up
+        user_id: Optional user_id filter (narrows results if multiple sessions
+            share an ID)
+        namespace: Optional namespace filter
+
+    Returns:
+        The Redis key string if exactly one match is found, None otherwise
+    """
+
+    from agent_memory_server.working_memory_index import get_working_memory_index
+
+    try:
+        index = await get_working_memory_index(redis_client)
+
+        filter_expression = Tag("session_id") == session_id
+        if namespace:
+            filter_expression &= Tag("namespace") == namespace
+        if user_id:
+            filter_expression &= Tag("user_id") == user_id
+
+        # Request up to 2 results so we can detect ambiguity.
+        filter_query = FilterQuery(
+            filter_expression=filter_expression,
+            return_fields=["session_id", "namespace", "user_id"],
+            num_results=2,
+        )
+
+        raw_results = await index.search(filter_query)
+        docs = getattr(raw_results, "docs", raw_results) or []
+
+        if not docs:
+            return None
+
+        total = getattr(raw_results, "total", len(docs))
+        if total > 1:
+            logger.warning(
+                "Ambiguous working-memory lookup for session_id=%s: "
+                "%d sessions matched. Provide namespace/user_id to disambiguate.",
+                session_id,
+                total,
+            )
+            return None
+
+        doc = docs[0]
+        # RedisVL returns doc.id as the full Redis key
+        doc_key = getattr(doc, "id", None)
+        if doc_key:
+            if isinstance(doc_key, bytes):
+                doc_key = doc_key.decode("utf-8")
+            return doc_key
+
+        return None
+
+    except Exception as e:
+        logger.debug(f"Index-based key resolution failed for session {session_id}: {e}")
+        return None
+
+
 async def get_working_memory(
     session_id: str,
     user_id: str | None = None,
@@ -418,6 +496,21 @@ async def get_working_memory(
                     )
             # If key_type is "none", the key doesn't exist - working_memory_data stays None
 
+        # Fallback: if direct key lookup failed, try resolving via the search
+        # index.  This handles the case where PUT stored with user_id/namespace
+        # but GET was called without them (issue #235).
+        if not working_memory_data:
+            resolved_key = await _resolve_working_memory_key_via_index(
+                redis_client, session_id, user_id, namespace
+            )
+            if resolved_key and resolved_key != key:
+                logger.debug(
+                    f"Resolved working memory key via index: {resolved_key} "
+                    f"(original key: {key})"
+                )
+                key = resolved_key
+                working_memory_data = await redis_client.json().get(key)
+
         if not working_memory_data:
             logger.debug(
                 f"No working memory found for parameters: {session_id}, {user_id}, {namespace}"
@@ -467,14 +560,19 @@ async def get_working_memory(
                 MemoryStrategyConfig()
             )  # Default to discrete strategy
 
+        # Use stored values for namespace/user_id — the caller may not have
+        # provided them (index-fallback path, issue #235).
+        stored_namespace = working_memory_data.get("namespace") or namespace
+        stored_user_id = working_memory_data.get("user_id") or user_id
+
         return WorkingMemory(
             messages=messages,
             memories=memories,
             context=working_memory_data.get("context"),
-            user_id=working_memory_data.get("user_id"),
+            user_id=stored_user_id,
             tokens=working_memory_data.get("tokens", 0),
             session_id=session_id,
-            namespace=namespace,
+            namespace=stored_namespace,
             ttl_seconds=working_memory_data.get("ttl_seconds", None),
             data=working_memory_data.get("data") or {},
             long_term_memory_strategy=long_term_memory_strategy,
@@ -589,6 +687,16 @@ async def delete_working_memory(
     )
 
     try:
+        # Check if the key exists; if not, try resolving via the search index
+        # (same fallback as get_working_memory for issue #235).
+        exists = await redis_client.exists(key)
+        if not exists:
+            resolved_key = await _resolve_working_memory_key_via_index(
+                redis_client, session_id, user_id, namespace
+            )
+            if resolved_key:
+                key = resolved_key
+
         # Delete the JSON key - the working memory search index automatically
         # removes the document from the index when the key is deleted
         await redis_client.delete(key)