Feat/add custom logger (#217)

* feat: add custom request log

* fix: format error

* fix: lint error

* feat: add request middleware

* fix: format error

* feat: support CUSTOM_LOGGER_WORKERS env

* feat: delete test_log

---------

Co-authored-by: CaralHsi <[email protected]>

Author: CarltonXiang

src/memos/api/context/dependencies.py

@@ -1,5 +1,4 @@
 import logging
-import os
 
 from fastapi import Depends, Header, Request
 
@@ -25,13 +24,6 @@ def get_trace_id_from_header(
     return g_trace_id or x_trace_id or trace_id
 
 
-def generate_trace_id() -> str:
-    """
-    Get a random trace_id.
-    """
-    return os.urandom(16).hex()
-
-
 def get_request_context(
     request: Request, trace_id: str | None = Depends(get_trace_id_from_header)
 ) -> RequestContext:
@@ -65,9 +57,6 @@ def get_g_object(trace_id: str | None = Depends(get_trace_id_from_header)) -> G:
     This creates a RequestContext and sets it globally for access
     throughout the request lifecycle.
     """
-    if trace_id is None:
-        trace_id = generate_trace_id()
-
     g = RequestContext(trace_id=trace_id)
     set_request_context(g)
     logger.info(f"Request g object created with trace_id: {g.trace_id}")

src/memos/api/middleware/request_context.py

@@ -0,0 +1,81 @@
+"""
+Request context middleware for automatic trace_id injection.
+"""
+
+import logging
+import os
+
+from collections.abc import Callable
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import Response
+
+from memos.api.context.context import RequestContext, set_request_context
+
+
+logger = logging.getLogger(__name__)
+
+
+def generate_trace_id() -> str:
+    """Generate a random trace_id."""
+    return os.urandom(16).hex()
+
+
+def extract_trace_id_from_headers(request: Request) -> str | None:
+    """Extract trace_id from various possible headers with priority: g-trace-id > x-trace-id > trace-id."""
+    trace_id = request.headers.get("g-trace-id")
+    if trace_id:
+        return trace_id
+
+    trace_id = request.headers.get("x-trace-id")
+    if trace_id:
+        return trace_id
+
+    trace_id = request.headers.get("trace-id")
+    if trace_id:
+        return trace_id
+
+    return None
+
+
+class RequestContextMiddleware(BaseHTTPMiddleware):
+    """
+    Middleware to automatically inject request context for every HTTP request.
+
+    This middleware:
+    1. Extracts trace_id from headers or generates a new one
+    2. Creates a RequestContext and sets it globally
+    3. Ensures the context is available throughout the request lifecycle
+    """
+
+    async def dispatch(self, request: Request, call_next: Callable) -> Response:
+        # Extract or generate trace_id
+        trace_id = extract_trace_id_from_headers(request)
+        if not trace_id:
+            trace_id = generate_trace_id()
+
+        # Create and set request context
+        context = RequestContext(trace_id=trace_id)
+        set_request_context(context)
+
+        # Add request metadata to context
+        context.set("method", request.method)
+        context.set("path", request.url.path)
+        context.set("client_ip", request.client.host if request.client else None)
+
+        # Log request start
+        logger.info(f"Request started: {request.method} {request.url.path} - trace_id: {trace_id}")
+
+        # Process the request
+        response = await call_next(request)
+
+        # Log request completion
+        logger.info(
+            f"Request completed: {request.method} {request.url.path} - trace_id: {trace_id} - status: {response.status_code}"
+        )
+
+        # Add trace_id to response headers for debugging
+        response.headers["x-trace-id"] = trace_id
+
+        return response

src/memos/api/product_api.py

@@ -3,6 +3,7 @@
 from fastapi import FastAPI
 
 from memos.api.exceptions import APIExceptionHandler
+from memos.api.middleware.request_context import RequestContextMiddleware
 from memos.api.routers.product_router import router as product_router
 
 
@@ -16,6 +17,9 @@
     version="1.0.0",
 )
 
+# Add request context middleware (must be added first)
+app.add_middleware(RequestContextMiddleware)
+
 # Include routers
 app.include_router(product_router)
 

src/memos/api/start_api.py

@@ -9,6 +9,7 @@
 from fastapi.responses import JSONResponse, RedirectResponse
 from pydantic import BaseModel, Field
 
+from memos.api.middleware.request_context import RequestContextMiddleware
 from memos.configs.mem_os import MOSConfig
 from memos.mem_os.main import MOS
 from memos.mem_user.user_manager import UserManager, UserRole
@@ -78,6 +79,8 @@ def get_mos_instance():
     version="1.0.0",
 )
 
+app.add_middleware(RequestContextMiddleware)
+
 
 class BaseRequest(BaseModel):
     """Base model for all requests."""

src/memos/log.py

@@ -1,12 +1,19 @@
+import atexit
 import logging
+import os
+import threading
 
+from concurrent.futures import ThreadPoolExecutor
 from logging.config import dictConfig
 from pathlib import Path
 from sys import stdout
 
+import requests
+
 from dotenv import load_dotenv
 
 from memos import settings
+from memos.api.context.context import get_current_trace_id
 
 
 # Load environment variables
@@ -26,6 +33,91 @@ def _setup_logfile() -> Path:
     return logfile
 
 
+class CustomLoggerRequestHandler(logging.Handler):
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls):
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+                    cls._instance._initialized = False
+        return cls._instance
+
+    def __init__(self):
+        """Initialize handler with minimal setup"""
+        if not self._initialized:
+            super().__init__()
+            workers = int(os.getenv("CUSTOM_LOGGER_WORKERS", "2"))
+            self._executor = ThreadPoolExecutor(
+                max_workers=workers, thread_name_prefix="log_sender"
+            )
+            self._is_shutting_down = threading.Event()
+            self._session = requests.Session()
+            self._initialized = True
+            atexit.register(self._cleanup)
+
+    def emit(self, record):
+        """Process log records of INFO or ERROR level (non-blocking)"""
+        if os.getenv("CUSTOM_LOGGER_URL") is None or self._is_shutting_down.is_set():
+            return
+
+        if record.levelno in (logging.INFO, logging.ERROR):
+            try:
+                trace_id = (
+                    get_current_trace_id()
+                )  # TODO: get trace_id from request context instead of get_current_trace_id
+                if trace_id:
+                    self._executor.submit(self._send_log_sync, record.getMessage(), trace_id)
+            except Exception as e:
+                if not self._is_shutting_down.is_set():
+                    print(f"Error sending log: {e}")
+
+    def _send_log_sync(self, message, trace_id):
+        """Send log message synchronously in a separate thread"""
+        print(f"send_log_sync: {message} {trace_id}")
+        try:
+            logger_url = os.getenv("CUSTOM_LOGGER_URL")
+            token = os.getenv("CUSTOM_LOGGER_TOKEN")
+
+            headers = {"Content-Type": "application/json"}
+            post_content = {"message": message, "trace_id": trace_id}
+
+            # Add auth token if exists
+            if token:
+                headers["Authorization"] = f"Bearer {token}"
+
+            # Add traceId to headers for consistency
+            headers["traceId"] = trace_id
+
+            # Add custom attributes from env
+            for key, value in os.environ.items():
+                if key.startswith("CUSTOM_LOGGER_ATTRIBUTE_"):
+                    attribute_key = key[len("CUSTOM_LOGGER_ATTRIBUTE_") :].lower()
+                    post_content[attribute_key] = value
+
+            self._session.post(logger_url, headers=headers, json=post_content, timeout=5)
+        except Exception:
+            # Silently ignore errors to avoid affecting main application
+            pass
+
+    def _cleanup(self):
+        """Clean up resources during program exit"""
+        if not self._initialized:
+            return
+
+        self._is_shutting_down.set()
+        try:
+            self._executor.shutdown(wait=False)
+            self._session.close()
+        except Exception as e:
+            print(f"Error during cleanup: {e}")
+
+    def close(self):
+        """Override close to prevent premature shutdown"""
+
+
 LOGGING_CONFIG = {
     "version": 1,
     "disable_existing_loggers": False,