Merge branch 'main' into feature/memory-duplication-fix

Author: abrookins

agent_memory_server/api.py

@@ -308,12 +308,12 @@ async def put_working_memory(
     if user_id is not None:
         memory.user_id = user_id
 
-    # Validate that all structured memories have id (if any)
-    for mem in memory.memories:
-        if not mem.id:
+    # Validate that all long-term memories have id (if any)
+    for long_term_mem in memory.memories:
+        if not long_term_mem.id:
             raise HTTPException(
                 status_code=400,
-                detail="All memory records in working memory must have an ID",
+                detail="All long-term memory records in working memory must have an ID",
             )
 
     # Validate that all messages have non-empty content

agent_memory_server/cli.py

@@ -1,19 +1,21 @@
-#!/usr/bin/env python
 """
 Command-line interface for agent-memory-server.
 """
 
 import datetime
 import importlib
-import logging
 import sys
 
 import click
 import uvicorn
 
 from agent_memory_server import __version__
 from agent_memory_server.config import settings
-from agent_memory_server.logging import configure_logging, get_logger
+from agent_memory_server.logging import (
+    configure_logging,
+    configure_mcp_logging,
+    get_logger,
+)
 from agent_memory_server.migrations import (
     migrate_add_discrete_memory_extracted_2,
     migrate_add_memory_hashes_1,
@@ -22,7 +24,6 @@
 from agent_memory_server.utils.redis import ensure_search_index_exists, get_redis_conn
 
 
-configure_logging()
 logger = get_logger(__name__)
 
 VERSION = __version__
@@ -45,6 +46,8 @@ def rebuild_index():
     """Rebuild the search index."""
     import asyncio
 
+    configure_logging()
+
     async def setup_and_run():
         redis = await get_redis_conn()
         await ensure_search_index_exists(redis, overwrite=True)
@@ -57,6 +60,7 @@ def migrate_memories():
     """Migrate memories from the old format to the new format."""
     import asyncio
 
+    configure_logging()
     click.echo("Starting memory migrations...")
 
     async def run_migrations():
@@ -82,6 +86,7 @@ def api(port: int, host: str, reload: bool):
     """Run the REST API server."""
     from agent_memory_server.main import on_start_logger
 
+    configure_logging()
     on_start_logger(port)
     uvicorn.run(
         "agent_memory_server.main:app",
@@ -103,6 +108,12 @@ def mcp(port: int, mode: str):
     """Run the MCP server."""
     import asyncio
 
+    # Configure MCP-specific logging BEFORE any imports to avoid stdout contamination
+    if mode == "stdio":
+        configure_mcp_logging()
+    else:
+        configure_logging()
+
     # Update the port in settings FIRST
     settings.mcp_port = port
 
@@ -117,14 +128,7 @@ async def setup_and_run():
             logger.info(f"Starting MCP server on port {port}\n")
             await mcp_app.run_sse_async()
         elif mode == "stdio":
-            # Try to force all logging to stderr because stdio-mode MCP servers
-            # use standard output for the protocol.
-            logging.basicConfig(
-                level=settings.log_level,
-                stream=sys.stderr,
-                force=True,  # remove any existing handlers
-                format="%(asctime)s %(name)s %(levelname)s %(message)s",
-            )
+            # Logging already configured above
             await mcp_app.run_stdio_async()
         else:
             raise ValueError(f"Invalid mode: {mode}")
@@ -154,6 +158,8 @@ def schedule_task(task_path: str, args: list[str]):
 
     from docket import Docket
 
+    configure_logging()
+
     # Parse the arguments
     task_args = {}
     for arg in args:
@@ -219,6 +225,8 @@ def task_worker(concurrency: int, redelivery_timeout: int):
 
     from docket import Worker
 
+    configure_logging()
+
     if not settings.use_docket:
         click.echo("Docket is disabled in settings. Cannot run worker.")
         sys.exit(1)

agent_memory_server/config.py

@@ -116,6 +116,8 @@ class Settings(BaseSettings):
 
     # Other Application settings
     log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "INFO"
+    default_mcp_user_id: str | None = None
+    default_mcp_namespace: str | None = None
 
     class Config:
         env_file = ".env"

agent_memory_server/logging.py

@@ -17,7 +17,7 @@ def configure_logging():
 
     # Configure standard library logging based on settings.log_level
     level = getattr(logging, settings.log_level.upper(), logging.INFO)
-    handler = logging.StreamHandler(sys.stdout)
+    handler = logging.StreamHandler(sys.stderr)  # Use stderr instead of stdout
     handler.setLevel(level)
     logging.basicConfig(level=level, handlers=[handler], format="%(message)s")
 
@@ -46,6 +46,41 @@ def configure_logging():
     _configured = True
 
 
+def configure_mcp_logging():
+    """Configure logging specifically for MCP server in stdio mode"""
+    global _configured
+
+    # Clear any existing handlers and configuration
+    root_logger = logging.getLogger()
+    root_logger.handlers.clear()
+
+    # Configure stderr-only logging for MCP stdio mode
+    level = getattr(logging, settings.log_level.upper(), logging.INFO)
+    stderr_handler = logging.StreamHandler(sys.stderr)
+    stderr_handler.setLevel(level)
+    stderr_handler.setFormatter(
+        logging.Formatter("%(asctime)s %(name)s %(levelname)s %(message)s")
+    )
+    root_logger.addHandler(stderr_handler)
+    root_logger.setLevel(level)
+
+    # Configure structlog to also use stderr
+    structlog.configure(
+        processors=[
+            structlog.stdlib.filter_by_level,
+            structlog.stdlib.add_logger_name,
+            structlog.stdlib.add_log_level,
+            structlog.processors.TimeStamper(fmt="iso"),
+            structlog.processors.format_exc_info,
+            structlog.processors.JSONRenderer(),
+        ],
+        wrapper_class=structlog.stdlib.BoundLogger,
+        logger_factory=structlog.stdlib.LoggerFactory(),
+        cache_logger_on_first_use=True,
+    )
+    _configured = True
+
+
 def get_logger(name: str | None = None) -> structlog.stdlib.BoundLogger:
     """
     Get a configured logger instance.

agent_memory_server/mcp.py

@@ -1,5 +1,4 @@
 import logging
-import os
 from typing import Any
 
 import ulid
@@ -43,9 +42,6 @@
 
 logger = logging.getLogger(__name__)
 
-# Default namespace for STDIO mode
-DEFAULT_NAMESPACE = os.getenv("MCP_NAMESPACE")
-
 
 class FastMCP(_FastMCPBase):
     """Extend FastMCP to support optional URL namespace and default STDIO namespace."""
@@ -165,7 +161,7 @@ async def run_stdio_async(self):
     "Redis Agent Memory Server",
     port=settings.mcp_port,
     instructions=INSTRUCTIONS,
-    default_namespace=DEFAULT_NAMESPACE,
+    default_namespace=settings.default_mcp_namespace,
 )
 
 
@@ -301,10 +297,11 @@ async def create_long_term_memories(
         An acknowledgement response indicating success
     """
     # Apply default namespace for STDIO if not provided in memory entries
-    if DEFAULT_NAMESPACE:
-        for mem in memories:
-            if mem.namespace is None:
-                mem.namespace = DEFAULT_NAMESPACE
+    for mem in memories:
+        if mem.namespace is None and settings.default_mcp_namespace:
+            mem.namespace = settings.default_mcp_namespace
+        if mem.user_id is None and settings.default_mcp_user_id:
+            mem.user_id = settings.default_mcp_user_id
 
     payload = CreateMemoryRecordRequest(
         memories=[MemoryRecord(**mem.model_dump()) for mem in memories]
@@ -418,6 +415,11 @@ async def search_long_term_memory(
     Returns:
         MemoryRecordResults containing matched memories sorted by relevance
     """
+    if user_id is None and settings.default_mcp_user_id:
+        user_id = UserId(eq=settings.default_mcp_user_id)
+    if namespace is None and settings.default_mcp_namespace:
+        namespace = Namespace(eq=settings.default_mcp_namespace)
+
     try:
         payload = SearchRequest(
             text=text,
@@ -545,7 +547,7 @@ async def memory_prompt(
 
     Args:
         - text: The user's query
-        - session_id: Add conversation history from a session
+        - session_id: Add conversation history from a working memory session
         - namespace: Filter session and long-term memory namespace
         - topics: Search for long-term memories matching topics
         - entities: Search for long-term memories matching entities
@@ -562,6 +564,9 @@ async def memory_prompt(
     _session_id = session_id.eq if session_id and session_id.eq else None
     session = None
 
+    if user_id is None and settings.default_mcp_user_id:
+        user_id = UserId(eq=settings.default_mcp_user_id)
+
     if _session_id is not None:
         session = WorkingMemoryRequest(
             session_id=_session_id,
@@ -602,8 +607,8 @@ async def set_working_memory(
     messages: list[MemoryMessage] | None = None,
     context: str | None = None,
     data: dict[str, Any] | None = None,
-    namespace: str | None = None,
-    user_id: str | None = None,
+    namespace: str | None = settings.default_mcp_namespace,
+    user_id: str | None = settings.default_mcp_user_id,
     ttl_seconds: int = 3600,
 ) -> WorkingMemoryResponse:
     """
@@ -701,11 +706,6 @@ async def set_working_memory(
     Returns:
         Updated working memory response (may include summarization if window exceeded)
     """
-    # Apply default namespace if configured
-    memory_namespace = namespace
-    if not memory_namespace and DEFAULT_NAMESPACE:
-        memory_namespace = DEFAULT_NAMESPACE
-
     # Auto-generate IDs for memories that don't have them
     processed_memories = []
     if memories:
@@ -717,6 +717,7 @@ async def set_working_memory(
                 processed_memory = memory.model_copy(
                     update={
                         "id": memory_id,
+                        "user_id": user_id,
                         "persisted_at": None,  # Mark as pending promotion
                     }
                 )
@@ -756,7 +757,7 @@ async def set_working_memory(
     # Create the working memory object
     working_memory_obj = WorkingMemory(
         session_id=session_id,
-        namespace=memory_namespace,
+        namespace=namespace,
         memories=processed_memories,
         messages=processed_messages,
         context=context,

agent_memory_server/vectorstore_adapter.py

@@ -829,15 +829,13 @@ async def search_memories(
             search_kwargs["score_threshold"] = score_threshold
 
         logger.debug(f"[search_memories] Search kwargs: {search_kwargs}")
-
         search_results = (
             await self.vectorstore.asimilarity_search_with_relevance_scores(
                 **search_kwargs
             )
         )
 
         logger.debug(f"[search_memories] Search results: {search_results}")
-
         # Convert results to MemoryRecordResult objects
         memory_results = []
         for i, (doc, score) in enumerate(search_results):

tests/test_cli.py

@@ -157,11 +157,11 @@ def test_mcp_command_sse_mode(self, mock_mcp_app, mock_settings):
         assert result.exit_code == 0
         mock_mcp_app.run_sse_async.assert_called_once()
 
-    @patch("agent_memory_server.cli.logging.basicConfig")
+    @patch("agent_memory_server.cli.configure_mcp_logging")
     @patch("agent_memory_server.cli.settings")
     @patch("agent_memory_server.mcp.mcp_app")
     def test_mcp_command_stdio_logging_config(
-        self, mock_mcp_app, mock_settings, mock_basic_config
+        self, mock_mcp_app, mock_settings, mock_configure_mcp_logging
     ):
         """Test that stdio mode configures logging to stderr."""
         mock_settings.mcp_port = 3001
@@ -174,7 +174,7 @@ def test_mcp_command_stdio_logging_config(
 
         assert result.exit_code == 0
         mock_mcp_app.run_stdio_async.assert_called_once()
-        mock_basic_config.assert_called_once()
+        mock_configure_mcp_logging.assert_called_once()
 
 
 class TestScheduleTask: