diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 3e1be99b38..068caceb15 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -137,6 +137,7 @@ HTTP client socket and connection configuration. Controls low-level socket optio | `AIPERF_HTTP_TRUST_ENV` | `False` | — | Trust environment variables for HTTP client configuration. When enabled, aiohttp will read proxy settings from HTTP_PROXY, HTTPS_PROXY, and NO_PROXY environment variables. | | `AIPERF_HTTP_X_SESSION_ID_FROM_CORRELATION_ID` | `False` | — | Also send X-Session-ID with the stable X-Correlation-ID value. Use this when an external router requires a session-affinity header. | | `AIPERF_HTTP_X_SMG_ROUTING_KEY_FROM_CORRELATION_ID` | `False` | — | Also send X-SMG-Routing-Key with the stable X-Correlation-ID value. Use this with the SGLang Model Gateway manual routing policy. | +| `AIPERF_HTTP_X_DYNAMO_SESSION_ID_FROM_CORRELATION_ID` | `False` | — | Also send X-Dynamo-Session-ID with the stable X-Correlation-ID value, plus X-Dynamo-Parent-Session-ID on subagent children. Use this with a Dynamo frontend running --router-session-affinity-ttl-secs to pin every turn of a session to the replica holding its KV prefix. | | `AIPERF_HTTP_VIDEO_POLL_INTERVAL` | `0.1` | ≥ 0.001, ≤ 10.0 | Interval in seconds between status polls for async video generation jobs. Lower values provide faster completion detection but increase server load. Applies to the aiohttp transport. | ## LOGGING diff --git a/src/aiperf/common/environment.py b/src/aiperf/common/environment.py index a7c6510188..7271ca955b 100644 --- a/src/aiperf/common/environment.py +++ b/src/aiperf/common/environment.py @@ -685,6 +685,13 @@ class _HTTPSettings(BaseSettings): description="Also send X-SMG-Routing-Key with the stable X-Correlation-ID value. " "Use this with the SGLang Model Gateway manual routing policy.", ) + X_DYNAMO_SESSION_ID_FROM_CORRELATION_ID: bool = Field( + default=False, + description="Also send X-Dynamo-Session-ID with the stable X-Correlation-ID value, " + "plus X-Dynamo-Parent-Session-ID on subagent children. Use this with a Dynamo " + "frontend running --router-session-affinity-ttl-secs to pin every turn of a " + "session to the replica holding its KV prefix.", + ) VIDEO_POLL_INTERVAL: float = Field( ge=0.001, le=10.0, diff --git a/src/aiperf/transports/base_transports.py b/src/aiperf/transports/base_transports.py index 1b57f453ce..b34da70c98 100644 --- a/src/aiperf/transports/base_transports.py +++ b/src/aiperf/transports/base_transports.py @@ -121,6 +121,12 @@ def build_headers(self, request_info: RequestInfo) -> dict[str, str]: headers["X-Session-ID"] = request_info.x_correlation_id if Environment.HTTP.X_SMG_ROUTING_KEY_FROM_CORRELATION_ID: headers["X-SMG-Routing-Key"] = request_info.x_correlation_id + if Environment.HTTP.X_DYNAMO_SESSION_ID_FROM_CORRELATION_ID: + headers["X-Dynamo-Session-ID"] = request_info.x_correlation_id + if request_info.parent_correlation_id: + headers["X-Dynamo-Parent-Session-ID"] = ( + request_info.parent_correlation_id + ) headers.update(request_info.endpoint_headers) headers.update(self.get_transport_headers(request_info)) diff --git a/tests/unit/transports/test_base_transport.py b/tests/unit/transports/test_base_transport.py index 1f0aa690ca..f4034083ff 100644 --- a/tests/unit/transports/test_base_transport.py +++ b/tests/unit/transports/test_base_transport.py @@ -176,6 +176,33 @@ def test_build_headers_can_alias_correlation_id_as_smg_routing_key( assert headers["X-SMG-Routing-Key"] == "test-correlation-id" + def test_build_headers_dynamo_session_header_root_has_no_parent( + self, transport, request_info, monkeypatch + ): + """Root session emits X-Dynamo-Session-ID but no parent header.""" + monkeypatch.setattr( + Environment.HTTP, "X_DYNAMO_SESSION_ID_FROM_CORRELATION_ID", True + ) + + headers = transport.build_headers(request_info) + + assert headers["X-Dynamo-Session-ID"] == "test-correlation-id" + assert "X-Dynamo-Parent-Session-ID" not in headers + + def test_build_headers_dynamo_session_header_child_adds_parent( + self, transport, request_info, monkeypatch + ): + """Subagent child also emits X-Dynamo-Parent-Session-ID.""" + monkeypatch.setattr( + Environment.HTTP, "X_DYNAMO_SESSION_ID_FROM_CORRELATION_ID", True + ) + request_info.parent_correlation_id = "parent-correlation-id" + + headers = transport.build_headers(request_info) + + assert headers["X-Dynamo-Session-ID"] == "test-correlation-id" + assert headers["X-Dynamo-Parent-Session-ID"] == "parent-correlation-id" + def test_build_headers_transport_headers_override(self, request_info): """Test that transport headers can override endpoint headers."""