MCP memory_prompt: support session-only and long_term_search modes

agent_memory_server/mcp.py

@@ -678,13 +678,14 @@ async def memory_prompt(
     recency_half_life_last_access_days: float | None = None,
     recency_half_life_created_days: float | None = None,
     server_side_recency: bool | None = None,
+    include_long_term_search: bool = True,
 ) -> dict[str, Any]:
     """
     Hydrate a query with relevant session history and long-term memories.
 
     This tool enriches the query by retrieving:
-    1. Context from the current conversation session
-    2. Relevant long-term memories related to the query
+    1. Context from the current conversation session (when session_id is provided)
+    2. Relevant long-term memories related to the query (when include_long_term_search is True)
 
     The tool returns both the relevant memories AND the user's query in a format ready for
     generating comprehensive responses.
@@ -703,41 +704,37 @@ async def memory_prompt(
     - NEVER invent or guess a session ID - if you don't know it, omit this filter
     - Session IDs from examples will NOT work with real data
 
+    MODES:
+    - Session + long-term search (default): Provide session_id and leave include_long_term_search=True
+    - Session-only: Provide session_id and set include_long_term_search=False
+    - Long-term search only: Omit session_id, leave include_long_term_search=True
+
+    At least one of session_id or include_long_term_search=True must be active.
+
     COMMON USAGE PATTERNS:
     ```python
     1. Hydrate a user prompt with long-term memory search:
     memory_prompt(query="What was my favorite color?")
-    ```
 
-    2. Answer "what do you remember about me?" type questions:
+    2. Session-only (no long-term search):
+    memory_prompt(
+        query="Continue our conversation",
+        session_id={"eq": "session_12345"},
+        include_long_term_search=False
+    )
+
+    3. Answer "what do you remember about me?" type questions:
     memory_prompt(
         query="What do you remember about me?",
         user_id={"eq": "user_123"},
         limit=50
     )
-    ```
 
-    3. Hydrate a user prompt with long-term memory search and session filter:
+    4. Hydrate with session context and filtered long-term search:
     memory_prompt(
         query="What is my favorite color?",
-        session_id={
-            "eq": "session_12345"
-        },
-        namespace={
-            "eq": "user_preferences"
-        }
-    )
-
-    4. Hydrate a user prompt with long-term memory search and complex filters:
-    memory_prompt(
-        query="What was my favorite color?",
-        topics={
-            "any": ["preferences", "settings"]
-        },
-        created_at={
-            "gt": "2023-01-01T00:00:00Z"
-        },
-        limit=5
+        session_id={"eq": "session_12345"},
+        namespace={"eq": "user_preferences"}
     )
 
     5. Search with datetime range filters:
@@ -746,9 +743,6 @@ async def memory_prompt(
         created_at={
             "gte": "2024-01-01T00:00:00Z",
             "lt": "2024-02-01T00:00:00Z"
-        },
-        last_accessed={
-            "gt": "2024-01-15T12:00:00Z"
         }
     )
     ```
@@ -777,15 +771,11 @@ async def memory_prompt(
         - recency_half_life_last_access_days: Half-life (days) for last_accessed decay
         - recency_half_life_created_days: Half-life (days) for created_at decay
         - server_side_recency: If true, attempt server-side recency-aware re-ranking
+        - include_long_term_search: Whether to include long-term memory search (default: True). Set to False for session-only prompts.
 
     Returns:
         JSON-serializable memory prompt payload including memory context and the user's query
     """
-    if distance_threshold is not None and search_mode != SearchModeEnum.SEMANTIC:
-        raise ValueError(
-            "distance_threshold is only supported for semantic search mode"
-        )
-
     _session_id = session_id.eq if session_id and session_id.eq else None
     session = None
 
@@ -801,38 +791,49 @@ async def memory_prompt(
             context_window_max=context_window_max,
         )
 
-    # Do NOT pass session_id to the long-term search — it scopes working
-    # memory retrieval, not the long-term memory search.  The REST API keeps
-    # these separate (session vs long_term_search); the MCP flat parameter
-    # space must not conflate them or long-term memories from other sessions
-    # will be excluded.
-    search_payload = SearchRequest(
-        text=query,
-        namespace=namespace,
-        topics=topics,
-        entities=entities,
-        created_at=created_at,
-        last_accessed=last_accessed,
-        user_id=user_id,
-        distance_threshold=distance_threshold,
-        memory_type=memory_type,
-        event_date=event_date,
-        search_mode=search_mode,
-        limit=limit,
-        offset=offset,
-        recency_boost=recency_boost,
-        recency_semantic_weight=recency_semantic_weight,
-        recency_recency_weight=recency_recency_weight,
-        recency_freshness_weight=recency_freshness_weight,
-        recency_novelty_weight=recency_novelty_weight,
-        recency_half_life_last_access_days=recency_half_life_last_access_days,
-        recency_half_life_created_days=recency_half_life_created_days,
-        server_side_recency=server_side_recency,
-    )
-    _params = {}
+    if not session and not include_long_term_search:
+        raise ValueError(
+            "Either session_id or include_long_term_search=True must be provided"
+        )
+
+    _params: dict[str, Any] = {}
     if session is not None:
         _params["session"] = session
-    if search_payload is not None:
+
+    if include_long_term_search:
+        if distance_threshold is not None and search_mode != SearchModeEnum.SEMANTIC:
+            raise ValueError(
+                "distance_threshold is only supported for semantic search mode"
+            )
+
+        # Do NOT pass session_id to the long-term search -- it scopes working
+        # memory retrieval, not the long-term memory search.  The REST API keeps
+        # these separate (session vs long_term_search); the MCP flat parameter
+        # space must not conflate them or long-term memories from other sessions
+        # will be excluded.
+        search_payload = SearchRequest(
+            text=query,
+            namespace=namespace,
+            topics=topics,
+            entities=entities,
+            created_at=created_at,
+            last_accessed=last_accessed,
+            user_id=user_id,
+            distance_threshold=distance_threshold,
+            memory_type=memory_type,
+            event_date=event_date,
+            search_mode=search_mode,
+            limit=limit,
+            offset=offset,
+            recency_boost=recency_boost,
+            recency_semantic_weight=recency_semantic_weight,
+            recency_recency_weight=recency_recency_weight,
+            recency_freshness_weight=recency_freshness_weight,
+            recency_novelty_weight=recency_novelty_weight,
+            recency_half_life_last_access_days=recency_half_life_last_access_days,
+            recency_half_life_created_days=recency_half_life_created_days,
+            server_side_recency=server_side_recency,
+        )
         _params["long_term_search"] = search_payload
 
     # Create a background tasks instance for the MCP call

tests/test_mcp.py

@@ -293,6 +293,98 @@ async def mock_core_memory_prompt(
                 captured_params["long_term_search"].search_mode == SearchModeEnum.HYBRID
             )
 
+    @pytest.mark.asyncio
+    async def test_memory_prompt_session_only_mode(self, session, monkeypatch):
+        """Test memory_prompt with include_long_term_search=False (session-only mode)."""
+        captured_params = {}
+
+        async def mock_core_memory_prompt(
+            params: MemoryPromptRequest, background_tasks, optimize_query: bool = False
+        ):
+            captured_params["query"] = params.query
+            captured_params["session"] = params.session
+            captured_params["long_term_search"] = params.long_term_search
+
+            return MemoryPromptResponse(
+                messages=[
+                    SystemMessage(
+                        content=TextContent(type="text", text="Session only response")
+                    )
+                ]
+            )
+
+        monkeypatch.setattr(
+            "agent_memory_server.mcp.core_memory_prompt", mock_core_memory_prompt
+        )
+
+        async with client_session(mcp_app._mcp_server) as client:
+            prompt = await client.call_tool(
+                "memory_prompt",
+                {
+                    "query": "Continue our conversation",
+                    "session_id": {"eq": session},
+                    "include_long_term_search": False,
+                },
+            )
+
+            assert isinstance(prompt, CallToolResult)
+            assert captured_params["session"] is not None
+            assert captured_params["session"].session_id == session
+            assert captured_params["long_term_search"] is None
+
+    @pytest.mark.asyncio
+    async def test_memory_prompt_session_only_ignores_search_params(
+        self, session, monkeypatch
+    ):
+        """Test that session-only mode does not validate search-only params."""
+        captured_params = {}
+
+        async def mock_core_memory_prompt(
+            params: MemoryPromptRequest, background_tasks, optimize_query: bool = False
+        ):
+            captured_params["long_term_search"] = params.long_term_search
+            return MemoryPromptResponse(
+                messages=[SystemMessage(content=TextContent(type="text", text="OK"))]
+            )
+
+        monkeypatch.setattr(
+            "agent_memory_server.mcp.core_memory_prompt", mock_core_memory_prompt
+        )
+
+        async with client_session(mcp_app._mcp_server) as client:
+            # This should NOT raise even though distance_threshold + keyword is invalid
+            prompt = await client.call_tool(
+                "memory_prompt",
+                {
+                    "query": "test",
+                    "session_id": {"eq": session},
+                    "include_long_term_search": False,
+                    "distance_threshold": 0.8,
+                    "search_mode": "keyword",
+                },
+            )
+
+            assert isinstance(prompt, CallToolResult)
+            assert captured_params["long_term_search"] is None
+
+    @pytest.mark.asyncio
+    async def test_memory_prompt_no_session_no_search_raises(self, mcp_test_setup):
+        """Test that omitting both session_id and include_long_term_search raises."""
+        async with client_session(mcp_app._mcp_server) as client:
+            result = await client.call_tool(
+                "memory_prompt",
+                {
+                    "query": "test",
+                    "include_long_term_search": False,
+                },
+            )
+            assert result.isError
+            assert any(
+                "session_id" in str(c.text).lower()
+                or "include_long_term_search" in str(c.text).lower()
+                for c in result.content
+            )
+
     @pytest.mark.asyncio
     async def test_set_working_memory_tool(self, mcp_test_setup):
         """Test the set_working_memory tool function"""