Stop logging to sdtio with mcp server in stdio mode

Author: abrookins

agent_memory_server/cli.py

@@ -5,14 +5,17 @@
 
 import datetime
 import importlib
-import logging
 import sys
 
 import click
 import uvicorn
 
 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,
@@ -21,7 +24,7 @@
 from agent_memory_server.utils.redis import ensure_search_index_exists, get_redis_conn
 
 
-configure_logging()
+# Don't configure logging at module level - let each command handle it
 logger = get_logger(__name__)
 
 VERSION = "0.2.0"
@@ -44,6 +47,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)
@@ -56,6 +61,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():
@@ -81,6 +87,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",
@@ -102,6 +109,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
 
@@ -116,14 +129,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}")
@@ -153,6 +159,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:
@@ -218,6 +226,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/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")
 
@@ -38,6 +38,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/vectorstore_factory.py

@@ -38,9 +38,10 @@ def patched_create_ulid() -> str:
 
     # Replace the broken function with our working one
     redisvl.utils.utils.create_ulid = patched_create_ulid
-    logging.info("Successfully patched RedisVL ULID function")
-except Exception as e:
-    logging.warning(f"Could not patch RedisVL ULID function: {e}")
+    # Note: Successfully patched RedisVL ULID function
+except Exception:
+    # Note: Could not patch RedisVL ULID function
+    pass
 
 from agent_memory_server.config import settings
 from agent_memory_server.vectorstore_adapter import (

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: