22Request context middleware for automatic trace_id injection.
33"""
44
5- import json
6- import os
75import time
86
97from collections .abc import Callable
1917
2018logger = memos .log .get_logger (__name__ )
2119
22- # Maximum body size to read for logging (in bytes) - bodies larger than this will be skipped
23- MAX_BODY_LOG_SIZE = os .getenv ("MAX_BODY_LOG_SIZE" , 10 * 1024 )
24-
2520
2621def extract_trace_id_from_headers (request : Request ) -> str | None :
2722 """Extract trace_id from various possible headers with priority: g-trace-id > x-trace-id > trace-id."""
@@ -31,127 +26,6 @@ def extract_trace_id_from_headers(request: Request) -> str | None:
3126 return None
3227
3328
34- def _is_json_request (request : Request ) -> tuple [bool , str ]:
35- """
36- Check if request is a JSON request.
37-
38- Args:
39- request: The request object
40-
41- Returns:
42- Tuple of (is_json, content_type)
43- """
44- if request .method not in "POST" , "PUT" , "PATCH" , "DELETE" ):
45- return False , ""
46-
47- content_type = request .headers .get ("content-type" , "" )
48- if not content_type :
49- return False , ""
50-
51- is_json = "application/json" in content_type .lower ()
52- return is_json , content_type
53-
54-
55- def _should_read_body (content_length : str | None ) -> tuple [bool , int | None ]:
56- """
57- Check if body should be read based on content-length header.
58-
59- Args:
60- content_length: Content-Length header value
61-
62- Returns:
63- Tuple of (should_read, body_size). body_size is None if header is invalid.
64- """
65- if not content_length :
66- return True , None
67-
68- try :
69- body_size = int (content_length )
70- return body_size <= MAX_BODY_LOG_SIZE , body_size
71- except ValueError :
72- return True , None
73-
74-
75- def _create_body_info (content_type : str , body_size : int ) -> dict :
76- """Create body_info dict for large bodies that are skipped."""
77- return {
78- "content_type" : content_type ,
79- "content_length" : body_size ,
80- "note" : f"body too large ({ body_size } ,
81- }
82-
83-
84- def _parse_json_body (body_bytes : bytes ) -> dict | str :
85- """
86- Parse JSON body bytes.
87-
88- Args:
89- body_bytes: Raw body bytes
90-
91- Returns:
92- Parsed JSON dict, or error message string if parsing fails
93- """
94- try :
95- return json .loads (body_bytes )
96- except (json .JSONDecodeError , UnicodeDecodeError ) as e :
97- return f"<unable to parse JSON: { e !s}
98-
99-
100- async def get_request_params (request : Request ) -> tuple [dict , bytes | None ]:
101- """
102- Extract request parameters (query params and body) for logging.
103-
104- Only reads body for application/json requests that are within size limits.
105-
106- This function is wrapped with exception handling to ensure logging failures
107- don't affect the actual request processing.
108-
109- Args:
110- request: The incoming request object
111-
112- Returns:
113- Tuple of (params_dict, body_bytes). body_bytes is None if body was not read.
114- Returns empty dict and None on any error.
115- """
116- try :
117- params_log = {}
118-
119- # Check if this is a JSON request
120- is_json , content_type = _is_json_request (request )
121- if not is_json :
122- return params_log , None
123-
124- # Pre-check body size using content-length header
125- content_length = request .headers .get ("content-length" )
126- should_read , body_size = _should_read_body (content_length )
127-
128- if not should_read and body_size is not None :
129- params_log ["body_info" ] = _create_body_info (content_type , body_size )
130- return params_log , None
131-
132- # Read body
133- body_bytes = await request .body ()
134-
135- if not body_bytes :
136- return params_log , None
137-
138- # Post-check: verify actual size (content-length might be missing or wrong)
139- actual_size = len (body_bytes )
140- if actual_size > MAX_BODY_LOG_SIZE :
141- params_log ["body_info" ] = _create_body_info (content_type , actual_size )
142- return params_log , None
143-
144- # Parse JSON body
145- params_log ["body" ] = _parse_json_body (body_bytes )
146- return params_log , body_bytes
147-
148- except Exception as e :
149- # Catch-all for any unexpected errors
150- logger .error (f"Unexpected error in get_request_params: { e } , exc_info = True )
151- # Return empty dict to ensure request can continue
152- return {}, None
153-
154-
15529class RequestContextMiddleware (BaseHTTPMiddleware ):
15630 """
15731 Middleware to automatically inject request context for every HTTP request.
@@ -193,26 +67,9 @@ async def dispatch(self, request: Request, call_next: Callable) -> Response:
19367 )
19468 set_request_context (context )
19569
196- # Get request parameters for logging
197- # Wrap in try-catch to ensure logging failures don't break the request
198- params_log , body_bytes = await get_request_params (request )
199-
200- # Re-create the request receive function if body was read
201- # This ensures downstream handlers can still read the body
202- if body_bytes is not None :
203- try :
204-
205- async def receive ():
206- return {"type" : "http.request" , "body" : body_bytes , "more_body" : False }
207-
208- request ._receive = receive
209- except Exception as e :
210- logger .error (f"Failed to recreate request receive function: { e } )
211- # Continue without restoring body, downstream handlers will handle it
212-
21370 logger .info (
21471 f"Request started, source: { self .source } { request .method } { request .url .path }
215- f"request params: { params_log } , headers: { request .headers }
72+ f"headers: { request .headers }
21673 )
21774
21875 # Process the request
0 commit comments