feat: support session IDs in request bodies#1137
Conversation
Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>
Try out this PRQuick install: pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@63ee21b6ecf485df6de23b788c9aeca68c6b4b97Recommended with virtual environment (using uv): uv venv --python 3.12 && source .venv/bin/activate
uv pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@63ee21b6ecf485df6de23b788c9aeca68c6b4b97Last updated for commit: |
WalkthroughAdds a new optional ChangesSession Body Field Feature
Estimated code review effort: 2 (Simple) | ~15 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/unit/config/test_converter_endpoint_session.py (1)
8-17: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueTest name doesn't follow the
test_<function>_<scenario>_<expected>convention.
test_session_affinity_fields_are_added_to_endpointis descriptive but doesn't name the function under test (build_endpoint) or a clear expected outcome. Considertest_build_endpoint_session_fields_set_included_in_output.As per path instructions,
tests/**/*.{py,pyi}: "Name tests astest_<function>_<scenario>_<expected>."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/config/test_converter_endpoint_session.py` around lines 8 - 17, Rename the test to follow the required `test_<function>_<scenario>_<expected>` pattern by including the function under test and the outcome; update `test_session_affinity_fields_are_added_to_endpoint` in the `build_endpoint` test to a name like `test_build_endpoint_session_fields_included_in_output`, keeping the assertions unchanged.Source: Path instructions
tests/unit/transports/test_aiohttp_transport.py (1)
302-321: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider adding a negative-path test for a missing session header.
Once
headers[session_header]inaiohttp_transport.pyis guarded (see companion comment on that file), add a test wheresession_body_fieldis set but the configuredsession_headeris absent from the outgoing headers, to lock in the non-crashing/skip behavior.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/transports/test_aiohttp_transport.py` around lines 302 - 321, Add a negative-path test alongside test_send_request_copies_resolved_session_header_to_body in test_aiohttp_transport.py to cover AioHttpTransport.send_request when session_body_field is configured but the session_header key is missing from the outgoing headers. Verify the request still succeeds without crashing and that the body does not get a copied session field when the header is absent, matching the guarded behavior in aiohttp_transport.py.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/aiperf/transports/aiohttp_transport.py`:
- Around line 322-328: The session-body injection in
aiohttp_transport.send_request currently does a direct headers[session_header]
lookup, which can raise KeyError and get swallowed by the broad exception
handling. Update the session_body_field block to use a safe header lookup with
an explicit fallback/diagnostic when
request_info.model_endpoint.endpoint.session_header (or default
X-Correlation-ID) is missing, and surface a clear error before payload
assignment. Keep the fix localized around send_request and the
session_header/session_body_field handling so misconfigured headers fail with a
useful message instead of silently breaking every request.
---
Nitpick comments:
In `@tests/unit/config/test_converter_endpoint_session.py`:
- Around line 8-17: Rename the test to follow the required
`test_<function>_<scenario>_<expected>` pattern by including the function under
test and the outcome; update
`test_session_affinity_fields_are_added_to_endpoint` in the `build_endpoint`
test to a name like `test_build_endpoint_session_fields_included_in_output`,
keeping the assertions unchanged.
In `@tests/unit/transports/test_aiohttp_transport.py`:
- Around line 302-321: Add a negative-path test alongside
test_send_request_copies_resolved_session_header_to_body in
test_aiohttp_transport.py to cover AioHttpTransport.send_request when
session_body_field is configured but the session_header key is missing from the
outgoing headers. Verify the request still succeeds without crashing and that
the body does not get a copied session field when the header is absent, matching
the guarded behavior in aiohttp_transport.py.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f7a1c7ad-556a-4f0e-bee5-058fba2f83e6
📒 Files selected for processing (11)
docs/benchmark-datasets.mddocs/cli-options.mdsrc/aiperf/common/models/model_endpoint_info.pysrc/aiperf/config/endpoint.pysrc/aiperf/config/flags/_converter_endpoint.pysrc/aiperf/config/flags/_section_fields.pysrc/aiperf/config/flags/cli_config.pysrc/aiperf/config/schema/aiperf-config.schema.jsonsrc/aiperf/transports/aiohttp_transport.pytests/unit/config/test_converter_endpoint_session.pytests/unit/transports/test_aiohttp_transport.py
| session_body_field = request_info.model_endpoint.endpoint.session_body_field | ||
| if session_body_field: | ||
| session_header = ( | ||
| request_info.model_endpoint.endpoint.session_header | ||
| or "X-Correlation-ID" | ||
| ) | ||
| payload[session_body_field] = headers[session_header] |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Guard the header lookup to avoid silent per-request failures.
headers[session_header] will raise KeyError whenever the configured (or default X-Correlation-ID) header isn't present in the built headers for a given request/turn. Because send_request's broad except Exception swallows this, a misconfigured --session-header/--session-body-field pair — or a turn missing the expected header — silently fails every request with an opaque error instead of a clear diagnostic.
🐛 Proposed fix using a safe lookup
- session_body_field = request_info.model_endpoint.endpoint.session_body_field
- if session_body_field:
- session_header = (
- request_info.model_endpoint.endpoint.session_header
- or "X-Correlation-ID"
- )
- payload[session_body_field] = headers[session_header]
+ session_body_field = request_info.model_endpoint.endpoint.session_body_field
+ if session_body_field:
+ session_header = (
+ request_info.model_endpoint.endpoint.session_header
+ or "X-Correlation-ID"
+ )
+ session_value = headers.get(session_header)
+ if session_value is not None:
+ payload[session_body_field] = session_value
+ else:
+ self.warning(
+ f"session_body_field={session_body_field!r} configured but "
+ f"header {session_header!r} not found in request headers; "
+ "skipping body-field passthrough for this request."
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| session_body_field = request_info.model_endpoint.endpoint.session_body_field | |
| if session_body_field: | |
| session_header = ( | |
| request_info.model_endpoint.endpoint.session_header | |
| or "X-Correlation-ID" | |
| ) | |
| payload[session_body_field] = headers[session_header] | |
| session_body_field = request_info.model_endpoint.endpoint.session_body_field | |
| if session_body_field: | |
| session_header = ( | |
| request_info.model_endpoint.endpoint.session_header | |
| or "X-Correlation-ID" | |
| ) | |
| session_value = headers.get(session_header) | |
| if session_value is not None: | |
| payload[session_body_field] = session_value | |
| else: | |
| self.warning( | |
| f"session_body_field={session_body_field!r} configured but " | |
| f"header {session_header!r} not found in request headers; " | |
| "skipping body-field passthrough for this request." | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/aiperf/transports/aiohttp_transport.py` around lines 322 - 328, The
session-body injection in aiohttp_transport.send_request currently does a direct
headers[session_header] lookup, which can raise KeyError and get swallowed by
the broad exception handling. Update the session_body_field block to use a safe
header lookup with an explicit fallback/diagnostic when
request_info.model_endpoint.endpoint.session_header (or default
X-Correlation-ID) is missing, and surface a clear error before payload
assignment. Keep the fix localized around send_request and the
session_header/session_body_field handling so misconfigured headers fail with a
useful message instead of silently breaking every request.
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
1 participant
Exgentic replay added in #1064 carries each recorded session ID in
x-dynamo-session-id, while SGLang's session-aware cache expects the same ID in the request body's top-levelsession_idfield. This adds an opt-in--session-body-fieldthat copies the resolved per-request session header value into a configured body field, avoiding a dataset-specific adapter.What changed
--session-body-fieldand threaded it through endpoint configuration and generated docs/schema.flowchart LR A["Exgentic source session"] --> B["Resolved session header"] B --> C["x-dynamo-session-id header"] B --> D["session_id request field"]Usage
Validation
pre-commit run --all-files: passed.Summary by CodeRabbit
New Features
--session-body-fieldoption in the CLI and corresponding configuration support.Documentation
Bug Fixes