Repository navigation
feat(gateway): add configurable storage hook request context mapping - #807
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds end-to-end support for storage_context_headers: CLI flag, docs, Python bindings and Router field, RouterConfig field and builder methods, config validation, middleware that builds StorageRequestContext from request headers, route-layer integration, and propagation into streaming/background tasks. Changes
Sequence DiagramsequenceDiagram
participant Client
participant Router as "Router / Server"
participant Middleware as "storage_context_middleware"
participant Builder as "build_storage_request_context"
participant Handler as "Route Handler / Streaming Task"
participant Storage as "Storage Hook / Downstream"
Client->>Router: HTTP request with headers
Router->>Middleware: enter protected layer
Middleware->>Builder: extract configured headers -> StorageRequestContext
Builder-->>Middleware: StorageRequestContext
Middleware->>Handler: invoke handler (with_request_context(ctx, next))
Handler->>Storage: perform storage ops using context
Storage-->>Handler: response/side-effects
Handler-->>Client: HTTP/streaming response
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enhances the gateway's request handling by introducing a flexible mechanism to pass custom HTTP header information into the storage hook's request context. This allows for more dynamic and context-aware operations within storage hooks, improving their adaptability to various application requirements. The changes involve extending configuration options, updating API call flows to incorporate this context, and ensuring proper validation and testing across the system's components. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
There was a problem hiding this comment.
Code Review
The pull request introduces configurable storage hook request context mapping, allowing HTTP headers to be mapped to context keys for storage hooks. This feature is implemented across the Rust core and Python bindings, including CLI argument parsing, configuration, and integration into various API endpoints. The changes are well-structured, with appropriate additions to data structures, builder patterns, and new utility functions for context building. Comprehensive unit tests have been added to validate the new functionality, including edge cases like empty headers and missing values. The documentation has also been updated to reflect the new configuration options. Overall, the implementation is robust and correctly integrates the new feature.
There was a problem hiding this comment.
Architecture Concern: Use middleware instead of per-handler wrapping
The current approach wraps 20+ endpoint handlers in server.rs with with_storage_context(), adding ~200 lines of repetitive boilerplate and requiring 8 conversation handlers to add headers: http::HeaderMap to their signatures.
This should be a single axum middleware in middleware.rs (where auth_middleware and RequestIdLayer already live):
// middleware.rs
pub async fn storage_context_layer(
State(state): State<Arc<AppState>>,
request: Request,
next: Next,
) -> Response {
if state.context.router_config.storage_context_headers.is_empty() {
return next.run(request).await;
}
match build_storage_request_context(&state.context.router_config, request.headers()) {
Some(ctx) => with_request_context(ctx, next.run(request)).await,
None => next.run(request).await,
}
}Then in server.rs, add it once when building the router:
.layer(axum::middleware::from_fn_with_state(state.clone(), storage_context_layer))Benefits:
- 0 endpoint handler changes (vs 20+)
- 0 conversation handler signature changes (vs 8)
- ~15 lines vs ~200 lines in server.rs
- Same performance — identical code runs, just in one place
- Follows existing patterns (
auth_middlewareis already inmiddleware.rs)
Other notes
- Config plumbing is correct:
types.rs,builder.rs,main.rs, Python bindings all updated properly with#[serde(default)]for backward compat. build_storage_request_contextis well-written: graceful handling of empty/whitespace keys, clean fast-path when no headers match.- Nit:
main.rs:1135reusesparse_selector()forkey=valueparsing — works fine but the function name implies k8s label selectors. A comment would help. - Nit:
server.rs:82—header_value()is a very generic name for a module-level function.extract_header_str()would be clearer.
Summary
The feature design (task-local RequestContext for WASM storage hooks) is sound. The implementation just needs to move from per-handler wrapping to middleware. This is a straightforward refactor that cuts ~185 lines of diff from server.rs.
ad03809 to
881b0db
Compare
|
@slin1237 |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 881b0dba65
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@bindings/python/src/smg/router_args.py`:
- Around line 1135-1137: The current assignment to
args_dict["storage_context_headers"] reuses cls._parse_selector which silently
drops malformed HEADER=CONTEXT_KEY tokens and skips the unprefixed fallback;
replace this with a dedicated parser method (e.g., a static method named
_parse_storage_context_headers) and call it here with
cli_args_dict.get(f"{prefix}storage_context_headers", None) so it: returns {}
for empty input, splits a single space-delimited string into entries, validates
each entry contains an '=' (raising ValueError on malformed tokens), and builds
the header->context_key mapping; update the code to call
cls._parse_storage_context_headers(...) instead of cls._parse_selector(...) for
storage_context_headers.
In `@docs/reference/configuration.md`:
- Around line 466-467: Update the configuration docs to explicitly call out the
trust boundary for mapped headers used by storage hooks: state that mappings
like tenant_id and user_id (used by storage hooks/read values from request
context) are only safe when the headers are injected or sanitized by a trusted
upstream (e.g., API gateway or reverse proxy); warn that if the gateway accepts
client-provided headers without sanitization those values can be spoofed and
suggest best practices (strip/overwrite untrusted headers, validate tokens, or
only accept headers from internal network sources).
- Around line 461-464: The fenced code block containing the example
"--storage-context-headers x-tenant-id=tenant_id x-user-id=user_id" needs blank
lines before and after the triple-backtick fence to satisfy markdownlint MD031;
update the surrounding text in docs/reference/configuration.md so there is an
empty line immediately above the opening ```bash and an empty line immediately
below the closing ``` to resolve the lint error.
In `@model_gateway/src/main.rs`:
- Around line 1136-1139: Replace the current use of parse_selector for
storage_context_headers with a strict parser that rejects malformed entries:
implement a new function parse_storage_context_headers(values: &[String]) ->
ConfigResult<HashMap<String,String>> which iterates values, uses split_once('=')
and returns ConfigError::InvalidValue (field "storage_context_headers", include
the bad value and reason "expected HEADER=CONTEXT_KEY") when any entry lacks
'='; then wire maybe_storage_context_headers to call
parse_storage_context_headers(&self.storage_context_headers) instead of
parse_selector so startup fails fast on typos.
In `@model_gateway/src/middleware.rs`:
- Around line 61-81: The header lookup in build_storage_request_context uses raw
header names from RouterConfig.storage_context_headers which may be mixed-case
and can fail because HeaderMap matches lowercased names; normalize header_name
before lookup by trimming and converting to lowercase (e.g.,
header_name.to_ascii_lowercase()) and then pass that normalized name into
extract_header_str (or update extract_header_str to perform case-insensitive
lookup), keeping the rest of the logic the same and leaving a note in
RouterConfig docs that header entries are normalized to lowercase.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: bec9714e-fdd8-4482-9b62-a9d4c82a6b63
📒 Files selected for processing (13)
bindings/python/src/lib.rsbindings/python/src/smg/router.pybindings/python/src/smg/router_args.pybindings/python/tests/test_arg_parser.pybindings/python/tests/test_router_config.pybindings/python/tests/test_startup_sequence.pybindings/python/tests/test_validation.pydocs/reference/configuration.mdmodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/main.rsmodel_gateway/src/middleware.rsmodel_gateway/src/server.rs
10259a8 to
830a3fd
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 830a3fd706
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@docs/reference/configuration.md`:
- Around line 452-471: Update the Storage Context Headers docs to explicitly
state the scope: clarify that the `--storage-context-headers` mapping only
applies to the storage-backed Responses/Conversations surface and is not applied
gateway-wide (it does not populate context for other endpoints such as
`/generate`, chat/completions, or other non-storage routes). Mention the exact
flag `--storage-context-headers` and the surfaces affected
(Responses/Conversations) so operators understand the limited scope and don’t
expect these headers on non-storage endpoints.
In `@model_gateway/src/routers/openai/responses/streaming.rs`:
- Around line 821-855: The current early returns when tx send/forwarding fails
(e.g., inside forward_streaming_event, send_mcp_list_tools_events and the
surrounding branches handling is_in_progress/mcp_list_tools_sent) drop remaining
processing and thus skip the persistence phase; instead, detect tx/receiver
disconnection (mirror the receiver_connected handling in
handle_simple_streaming_passthrough) by turning off forwarding: mark a local
receiver_connected=false or forwarding_disabled flag when a send to tx fails,
stop calling forward_streaming_event/send_mcp_list_tools_events thereafter, but
continue draining/parsing events, executing tool outputs and running the final
persistence logic (the code that writes final responses/conversation-backed
storage) so that store=true and conversation-backed requests still persist even
if the client disconnected; keep function names to locate edits:
forward_streaming_event, SseEventData, send_mcp_list_tools_events,
receiver_connected, handle_simple_streaming_passthrough, mcp_list_tools_sent,
and sequence_number.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: b7b77f0f-2063-49b0-bd6e-644249f61979
📒 Files selected for processing (16)
bindings/python/src/lib.rsbindings/python/src/smg/router.pybindings/python/src/smg/router_args.pybindings/python/tests/test_arg_parser.pybindings/python/tests/test_router_config.pybindings/python/tests/test_startup_sequence.pybindings/python/tests/test_validation.pydocs/reference/configuration.mdmodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/main.rsmodel_gateway/src/middleware.rsmodel_gateway/src/routers/grpc/harmony/responses/streaming.rsmodel_gateway/src/routers/grpc/regular/responses/streaming.rsmodel_gateway/src/routers/openai/responses/streaming.rsmodel_gateway/src/server.rs
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 480c2bbdb8
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2a2bfa88dc
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/config/validation.rs`:
- Around line 1088-1136: Update the three tests to assert the specific
validation failure rather than only is_err(); call
ConfigValidator::validate(&config), unwrap the Err, and assert it matches
ConfigError::ValidationFailed (or that its Display/Debug string contains the
specific substring for the storage_context_headers error) so the failure is
coming from the storage context validation; reference the test functions
(test_reject_duplicate_storage_context_keys,
test_reject_empty_storage_context_key,
test_reject_empty_storage_context_header_name), RouterConfig construction, and
ConfigValidator::validate to locate where to replace assert!(result.is_err())
with an assertion that inspects the Err variant or message text indicating the
expected validation failure.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 7e36b4e0-3c8b-48fc-8d1b-438081c40804
📒 Files selected for processing (1)
model_gateway/src/config/validation.rs
2a2bfa8 to
c397cee
Compare
There was a problem hiding this comment.
♻️ Duplicate comments (2)
docs/reference/configuration.md (1)
452-471:⚠️ Potential issue | 🟡 MinorClarify that
--storage-context-headersis not gateway-wide.This section still reads as global behavior. Please explicitly state it applies only to storage-backed Responses/Conversations routes, not
/generate,/v1/chat/completions, or other non-storage endpoints.✏️ Suggested doc patch
This lets storage hooks read values such as `tenant_id` and `user_id` from the request context without hard-coding specific headers in the gateway. + +This mapping is applied only on the storage-backed Responses/Conversations +surface (for example, `/v1/responses*` and `/v1/conversations*`). It is not +applied gateway-wide to endpoints such as `/generate`, `/v1/chat/completions`, +or other non-storage routes. Only map headers that are injected or sanitized by a trusted upstream. Client-supplied headers can otherwise spoof storage hook request context values.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@docs/reference/configuration.md` around lines 452 - 471, The section for Storage Context Headers currently implies gateway-wide behavior for the `--storage-context-headers` option; update the text under the "Storage Context Headers" heading (and the example) to explicitly state that `--storage-context-headers` applies only to storage-backed Responses and Conversations routes (e.g., storage-backed /responses or /conversations handlers) and does NOT apply to non-storage endpoints such as `/generate`, `/v1/chat/completions`, or other non-storage routes; keep the example but add one sentence clarifying scope and a short warning to avoid assuming global header mapping.model_gateway/src/config/validation.rs (1)
1088-1136:⚠️ Potential issue | 🟡 MinorStrengthen tests to assert the exact validation failure.
These tests only check
is_err(), so they can pass on unrelated failures. AssertConfigError::ValidationFailedwith the expected reason substring per case.✅ Suggested test hardening
- let result = ConfigValidator::validate(&config); - assert!(result.is_err()); + match ConfigValidator::validate(&config) { + Err(ConfigError::ValidationFailed { reason }) => { + assert!(reason.contains("same context key")); + } + other => panic!("expected ValidationFailed for duplicate context key, got: {other:?}"), + }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/config/validation.rs` around lines 1088 - 1136, Tests test_reject_duplicate_storage_context_keys, test_reject_empty_storage_context_key, and test_reject_empty_storage_context_header_name only assert result.is_err(); update each to assert the specific ConfigError::ValidationFailed variant and that its message contains the expected reason substring. Locate the tests and replace the loose is_err() checks with pattern matching or equality checks against ConfigError::ValidationFailed produced by ConfigValidator::validate(&config), and assert the error string contains the relevant text (e.g., "duplicate storage context key", "empty storage context key", "empty storage context header name") to ensure failures are for the intended validation in RouterConfig.storage_context_headers.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@docs/reference/configuration.md`:
- Around line 452-471: The section for Storage Context Headers currently implies
gateway-wide behavior for the `--storage-context-headers` option; update the
text under the "Storage Context Headers" heading (and the example) to explicitly
state that `--storage-context-headers` applies only to storage-backed Responses
and Conversations routes (e.g., storage-backed /responses or /conversations
handlers) and does NOT apply to non-storage endpoints such as `/generate`,
`/v1/chat/completions`, or other non-storage routes; keep the example but add
one sentence clarifying scope and a short warning to avoid assuming global
header mapping.
In `@model_gateway/src/config/validation.rs`:
- Around line 1088-1136: Tests test_reject_duplicate_storage_context_keys,
test_reject_empty_storage_context_key, and
test_reject_empty_storage_context_header_name only assert result.is_err();
update each to assert the specific ConfigError::ValidationFailed variant and
that its message contains the expected reason substring. Locate the tests and
replace the loose is_err() checks with pattern matching or equality checks
against ConfigError::ValidationFailed produced by
ConfigValidator::validate(&config), and assert the error string contains the
relevant text (e.g., "duplicate storage context key", "empty storage context
key", "empty storage context header name") to ensure failures are for the
intended validation in RouterConfig.storage_context_headers.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 2338d80c-f793-4c1b-9fa1-98a984b3a37b
📒 Files selected for processing (17)
bindings/python/src/lib.rsbindings/python/src/smg/router.pybindings/python/src/smg/router_args.pybindings/python/tests/test_arg_parser.pybindings/python/tests/test_router_config.pybindings/python/tests/test_startup_sequence.pybindings/python/tests/test_validation.pydocs/reference/configuration.mdmodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/config/validation.rsmodel_gateway/src/main.rsmodel_gateway/src/middleware.rsmodel_gateway/src/routers/grpc/harmony/responses/streaming.rsmodel_gateway/src/routers/grpc/regular/responses/streaming.rsmodel_gateway/src/routers/openai/responses/streaming.rsmodel_gateway/src/server.rs
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c397cee2d5
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fa0d6e7cd8
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
fa0d6e7 to
f6cbbdb
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f6cbbdb4e4
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
f6cbbdb to
998b95b
Compare
|
Hi @zhaowenzi, the DCO sign-off check has failed. All commits must include a To fix existing commits: # Sign off the last N commits (replace N with the number of unsigned commits)
git rebase HEAD~N --signoff
git push --force-with-leaseTo sign off future commits automatically:
|
62df3a6 to
e8ec1dc
Compare
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
Signed-off-by: Ziwen Zhao <zzw.mose@gmail.com>
e8ec1dc to
a3d21ce
Compare
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
Replace the tokio::task_local propagation pattern (wrapping entire spawned streaming task bodies in with_request_context) with a struct- carried approach where StorageHandles carries the RequestContext and persist functions apply the scope internally. What changed: - context.rs: add request_context field to StorageHandles and storage_request_context to openai RequestContext - route.rs: capture current_request_context() into ctx when building the openai RequestContext (still within middleware task-local scope) - persistence_utils.rs: persist_conversation_items now accepts Option<StorageRequestContext> and wraps the inner logic in with_request_context when present - grpc/common/responses/utils.rs: persist_response_if_needed passes request_context through to persist_conversation_items - streaming files: pass storage.request_context.clone() or None to persist calls — no spawn-body wrapping, no indentation changes What removed: - CLI arg --storage-context-headers (config-file-only now) - Python bindings for storage_context_headers (7 files reverted) - All current_request_context/with_request_context usage in streaming spawned tasks (3 files reverted to main) Why: The original approach required wrapping entire tokio::spawn bodies in with_request_context, causing ~500 lines of indentation churn across 3 streaming files. The refactored approach threads context through structs and applies the task-local scope only at the persist call site, reducing router changes to +1 line per persist call. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
Add back the CLI flag since there's no config file mechanism yet — without it the storage_context_headers field in RouterConfig would be unreachable. Format: --storage-context-headers header=context_key Example: --storage-context-headers x-tenant-id=tenant_id x-user-id=user_id Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
Restore the Python RouterArgs field and CLI arg parser for storage_context_headers, matching the Rust CLI --storage-context-headers. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
gRPC routes run through the same HTTP middleware as OpenAI routes, so they have access to storage request context. Thread it through ResponsesContext the same way OpenAI uses StorageHandles — capture current_request_context() in the handler before any spawn, carry it on the struct, and pass it to persist calls. What changed: - grpc/common/responses/context.rs: add request_context field to ResponsesContext - grpc/router.rs: capture current_request_context() when building ResponsesContext (both regular and harmony paths) - grpc/regular/responses/streaming.rs: thread request_context through process_and_transform_sse_stream to persist_response_if_needed - All gRPC persist call sites: pass ctx.request_context.clone() instead of None Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
…mg-project#807) Signed-off-by: Ziwen Zhao <zzw.mose@gmail.com> Signed-off-by: Simo Lin <linsimo.mark@gmail.com> Co-authored-by: Simo Lin <linsimo.mark@gmail.com>
Description
Problem
Storage hooks already support per-request RequestContext, but the gateway did not provide a general way to populate that context from incoming HTTP headers. Downstream integrations had to add custom request handling in server.rs to pass request metadata into storage-backed responses and conversations flows.
In addition, streaming response persistence runs from spawned Tokio tasks, so task-local request context needed to be propagated into those background tasks for storage hooks to see the same request metadata.
Solution
Add a configurable storage_context_headers mapping and use route-scoped middleware to populate storage hook request context from selected HTTP headers. Also propagate the current request context into spawned streaming response tasks so storage-backed persistence continues to see the mapped context outside the original request task.
This keeps the gateway behavior general while avoiding per-handler boilerplate and preserving a narrow endpoint surface.
Changes
Test Plan
Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
New Features
Documentation
Tests
Validation