fix: restore legacy X-Ray-Serve-Request-Id header for sticky routing

src/twinkle/server/gateway/proxy.py

@@ -14,7 +14,7 @@
 
 from twinkle.server.telemetry.tracing import inject_context
 from twinkle.utils.logger import get_logger
-from twinkle_client.http.headers import H_MULTIPLEX, H_MULTIPLEX_LEGACY, H_REQUEST_ID
+from twinkle_client.http.headers import H_MULTIPLEX, H_MULTIPLEX_LEGACY, H_REQUEST_ID, H_REQUEST_ID_LEGACY
 
 logger = get_logger()
 
@@ -72,7 +72,7 @@ def _prepare_headers(self, request_headers) -> dict[str, str]:
         headers = dict(request_headers)
         headers.pop('host', None)
         headers.pop('content-length', None)
-        request_id = request_headers.get(H_REQUEST_ID)
+        request_id = request_headers.get(H_REQUEST_ID) or request_headers.get(H_REQUEST_ID_LEGACY)
         if request_id is not None and not request_headers.get(H_MULTIPLEX_LEGACY):
             headers[H_MULTIPLEX_LEGACY] = request_id
         if request_id is not None:

src/twinkle/server/utils/ray_serve_patch.py

@@ -66,7 +66,7 @@ def _patched_setup_request_context_and_handle(
             logger.debug(f'[Ray Serve Patch] Matched multiplexed_model_id: {multiplexed_model_id}')
 
         # Original logic for other headers (unchanged)
-        if decoded_key == 'x-request-id':
+        if decoded_key in ('x-request-id', 'x-ray-serve-request-id'):
             request_context_info['request_id'] = value.decode()
 
     import ray.serve.context as serve_context

src/twinkle_client/http/headers.py

@@ -13,6 +13,7 @@
 
 # -- request id --
 H_REQUEST_ID = 'x-request-id'
+H_REQUEST_ID_LEGACY = 'X-Ray-Serve-Request-Id'
 
 # -- multiplexed model id (sticky routing) --
 H_MULTIPLEX = 'serve_multiplexed_model_id'
@@ -22,7 +23,7 @@
 H_AUTH = 'Authorization'
 H_AUTH_TWINKLE = 'Twinkle-Authorization'
 
-_ROUTING_HEADERS = (H_REQUEST_ID, H_MULTIPLEX, H_MULTIPLEX_LEGACY)
+_ROUTING_HEADERS = (H_REQUEST_ID, H_REQUEST_ID_LEGACY, H_MULTIPLEX, H_MULTIPLEX_LEGACY)
 _AUTH_HEADERS = (H_AUTH, H_AUTH_TWINKLE)