Merge remote-tracking branch 'upstream/dev' into dev

Author: Nyakult

src/memos/api/middleware/request_context.py

@@ -2,6 +2,8 @@
 Request context middleware for automatic trace_id injection.
 """
 
+import json
+import os
 import time
 
 from collections.abc import Callable
@@ -17,6 +19,9 @@
 
 logger = memos.log.get_logger(__name__)
 
+# Maximum body size to read for logging (in bytes) - bodies larger than this will be skipped
+MAX_BODY_LOG_SIZE = os.getenv("MAX_BODY_LOG_SIZE", 10 * 1024)
+
 
 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."""
@@ -26,6 +31,127 @@ def extract_trace_id_from_headers(request: Request) -> str | None:
     return None
 
 
+def _is_json_request(request: Request) -> tuple[bool, str]:
+    """
+    Check if request is a JSON request.
+
+    Args:
+        request: The request object
+
+    Returns:
+        Tuple of (is_json, content_type)
+    """
+    if request.method not in ("POST", "PUT", "PATCH", "DELETE"):
+        return False, ""
+
+    content_type = request.headers.get("content-type", "")
+    if not content_type:
+        return False, ""
+
+    is_json = "application/json" in content_type.lower()
+    return is_json, content_type
+
+
+def _should_read_body(content_length: str | None) -> tuple[bool, int | None]:
+    """
+    Check if body should be read based on content-length header.
+
+    Args:
+        content_length: Content-Length header value
+
+    Returns:
+        Tuple of (should_read, body_size). body_size is None if header is invalid.
+    """
+    if not content_length:
+        return True, None
+
+    try:
+        body_size = int(content_length)
+        return body_size <= MAX_BODY_LOG_SIZE, body_size
+    except ValueError:
+        return True, None
+
+
+def _create_body_info(content_type: str, body_size: int) -> dict:
+    """Create body_info dict for large bodies that are skipped."""
+    return {
+        "content_type": content_type,
+        "content_length": body_size,
+        "note": f"body too large ({body_size} bytes), skipping read",
+    }
+
+
+def _parse_json_body(body_bytes: bytes) -> dict | str:
+    """
+    Parse JSON body bytes.
+
+    Args:
+        body_bytes: Raw body bytes
+
+    Returns:
+        Parsed JSON dict, or error message string if parsing fails
+    """
+    try:
+        return json.loads(body_bytes)
+    except (json.JSONDecodeError, UnicodeDecodeError) as e:
+        return f"<unable to parse JSON: {e!s}>"
+
+
+async def get_request_params(request: Request) -> tuple[dict, bytes | None]:
+    """
+    Extract request parameters (query params and body) for logging.
+
+    Only reads body for application/json requests that are within size limits.
+
+    This function is wrapped with exception handling to ensure logging failures
+    don't affect the actual request processing.
+
+    Args:
+        request: The incoming request object
+
+    Returns:
+        Tuple of (params_dict, body_bytes). body_bytes is None if body was not read.
+        Returns empty dict and None on any error.
+    """
+    try:
+        params_log = {}
+
+        # Check if this is a JSON request
+        is_json, content_type = _is_json_request(request)
+        if not is_json:
+            return params_log, None
+
+        # Pre-check body size using content-length header
+        content_length = request.headers.get("content-length")
+        should_read, body_size = _should_read_body(content_length)
+
+        if not should_read and body_size is not None:
+            params_log["body_info"] = _create_body_info(content_type, body_size)
+            return params_log, None
+
+        # Read body
+        body_bytes = await request.body()
+
+        if not body_bytes:
+            return params_log, None
+
+        # Post-check: verify actual size (content-length might be missing or wrong)
+        actual_size = len(body_bytes)
+        if actual_size > MAX_BODY_LOG_SIZE:
+            params_log["body_info"] = _create_body_info(content_type, actual_size)
+            return params_log, None
+
+        # Parse JSON body
+        params_log["body"] = _parse_json_body(body_bytes)
+        return params_log, body_bytes
+
+    except Exception as e:
+        # Catch-all for any unexpected errors
+        logger.error(f"Unexpected error in get_request_params: {e}", exc_info=True)
+        # Return empty dict to ensure request can continue
+        return {}, None
+
+
 class RequestContextMiddleware(BaseHTTPMiddleware):
     """
     Middleware to automatically inject request context for every HTTP request.
@@ -36,6 +162,17 @@ class RequestContextMiddleware(BaseHTTPMiddleware):
     3. Ensures the context is available throughout the request lifecycle
     """
 
+    def __init__(self, app, source: str | None = None):
+        """
+        Initialize the middleware.
+
+        Args:
+            app: The ASGI application
+            source: Source identifier (e.g., 'product' or 'server') to distinguish request origin
+        """
+        super().__init__(app)
+        self.source = source or "api"
+
     async def dispatch(self, request: Request, call_next: Callable) -> Response:
         # Extract or generate trace_id
         trace_id = extract_trace_id_from_headers(request) or generate_trace_id()
@@ -52,34 +189,48 @@ async def dispatch(self, request: Request, call_next: Callable) -> Response:
             env=env,
             user_type=user_type,
             user_name=user_name,
+            source=self.source,
         )
         set_request_context(context)
 
-        # Log request start with parameters
-        params_log = {}
+        # Get request parameters for logging
+        # Wrap in try-catch to ensure logging failures don't break the request
+        params_log, body_bytes = await get_request_params(request)
+
+        # Re-create the request receive function if body was read
+        # This ensures downstream handlers can still read the body
+        if body_bytes is not None:
+            try:
 
-        # Get query parameters
-        if request.query_params:
-            params_log["query_params"] = dict(request.query_params)
+                async def receive():
+                    return {"type": "http.request", "body": body_bytes, "more_body": False}
 
-        logger.info(f"Request started, params: {params_log}, headers: {request.headers}")
+                request._receive = receive
+            except Exception as e:
+                logger.error(f"Failed to recreate request receive function: {e}")
+                # Continue without restoring body, downstream handlers will handle it
+
+        logger.info(
+            f"Request started, source: {self.source}, method: {request.method}, path: {request.url.path}, "
+            f"request params: {params_log}, headers: {request.headers}"
+        )
 
         # Process the request
         try:
             response = await call_next(request)
             end_time = time.time()
             if response.status_code == 200:
                 logger.info(
-                    f"Request completed: {request.url.path}, status: {response.status_code}, cost: {(end_time - start_time) * 1000:.2f}ms"
+                    f"Request completed: source: {self.source}, path: {request.url.path}, status: {response.status_code}, cost: {(end_time - start_time) * 1000:.2f}ms"
                 )
             else:
                 logger.error(
-                    f"Request Failed: {request.url.path}, status: {response.status_code}, cost: {(end_time - start_time) * 1000:.2f}ms"
+                    f"Request Failed: source: {self.source}, path: {request.url.path}, status: {response.status_code}, cost: {(end_time - start_time) * 1000:.2f}ms"
                 )
         except Exception as e:
             end_time = time.time()
             logger.error(
-                f"Request Exception Error: {e}, cost: {(end_time - start_time) * 1000:.2f}ms"
+                f"Request Exception Error: source: {self.source}, path: {request.url.path}, error: {e}, cost: {(end_time - start_time) * 1000:.2f}ms"
             )
             raise e
 

src/memos/api/product_api.py

@@ -17,7 +17,7 @@
     version="1.0.1",
 )
 
-app.add_middleware(RequestContextMiddleware)
+app.add_middleware(RequestContextMiddleware, source="product_api")
 # Include routers
 app.include_router(product_router)
 

src/memos/api/routers/server_router.py

@@ -412,6 +412,8 @@ def _search_pref():
         search_req.include_preference,
     )
 
+    logger.info(f"Search memories result: {memories_result}")
+
     return SearchResponse(
         message="Search completed successfully",
         data=memories_result,
@@ -618,6 +620,9 @@ def _process_pref_mem() -> list[dict[str, str]]:
         text_response_data = text_future.result()
         pref_response_data = pref_future.result()
 
+    logger.info(f"add_memories Text response data: {text_response_data}")
+    logger.info(f"add_memories Pref response data: {pref_response_data}")
+
     return MemoryResponse(
         message="Memory added successfully",
         data=text_response_data + pref_response_data,

src/memos/api/server_api.py

@@ -18,7 +18,7 @@
     version="1.0.1",
 )
 
-app.add_middleware(RequestContextMiddleware)
+app.add_middleware(RequestContextMiddleware, source="server_api")
 # Include routers
 app.include_router(server_router)
 

src/memos/context/context.py

@@ -36,12 +36,14 @@ def __init__(
         env: str | None = None,
         user_type: str | None = None,
         user_name: str | None = None,
+        source: str | None = None,
     ):
         self.trace_id = trace_id or "trace-id"
         self.api_path = api_path
         self.env = env
         self.user_type = user_type
         self.user_name = user_name
+        self.source = source
         self._data: dict[str, Any] = {}
 
     def set(self, key: str, value: Any) -> None:
@@ -59,6 +61,7 @@ def __setattr__(self, name: str, value: Any) -> None:
             "env",
             "user_type",
             "user_name",
+            "source",
         ):
             super().__setattr__(name, value)
         else:
@@ -80,6 +83,7 @@ def to_dict(self) -> dict[str, Any]:
             "env": self.env,
             "user_type": self.user_type,
             "user_name": self.user_name,
+            "source": self.source,
             "data": self._data.copy(),
         }
 
@@ -146,6 +150,16 @@ def get_current_user_name() -> str | None:
     return "memos"
 
 
+def get_current_source() -> str | None:
+    """
+    Get the current request's source (e.g., 'product_api' or 'server_api').
+    """
+    context = _request_context.get()
+    if context:
+        return context.get("source")
+    return None
+
+
 def get_current_context() -> RequestContext | None:
     """
     Get the current request context.
@@ -161,6 +175,7 @@ def get_current_context() -> RequestContext | None:
             env=context_dict.get("env"),
             user_type=context_dict.get("user_type"),
             user_name=context_dict.get("user_name"),
+            source=context_dict.get("source"),
         )
         ctx._data = context_dict.get("data", {}).copy()
         return ctx