Skip to content

feat(gateway): add configurable storage hook request context mapping - #807

Merged
slin1237 merged 5 commits into
mainfrom
feat/storage-context-headers
Mar 24, 2026
Merged

slin1237 merged 5 commits into
mainfrom
feat/storage-context-headers

Conversation

@zhaowenzi

@zhaowenzi zhaowenzi commented Mar 18, 2026 •

Copy link
Copy Markdown
Contributor

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

  • add storage_context_headers to RouterConfig with backward-compatible default behavior
  • expose storage_context_headers through the Rust CLI and Python bindings
  • add storage context middleware in model_gateway/src/middleware.rs
  • apply the middleware only to storage-backed responses and conversations routes
  • propagate request context into spawned streaming response tasks for OpenAI and gRPC streaming paths
  • validate storage_context_headers to reject duplicate mappings to the same context key
  • add tests for configured header-to-context mapping behavior
  • document the new CLI option in docs/reference/configuration.md

Test Plan

Checklist
  • cargo +nightly fmt passes
  • cargo clippy --all-targets --all-features -- -D warnings passes
  • (Optional) Documentation updated
  • (Optional) Please join us on Slack #sig-smg to discuss, review, and merge PRs

Summary by CodeRabbit

  • New Features

    • Add CLI/config option to map incoming HTTP headers into storage hook request context and propagate it to protected endpoints, background tasks, and streaming paths.
  • Documentation

    • Add reference docs, examples, and a security note for the Storage Context Headers option.
  • Tests

    • Add/update tests for CLI parsing, propagation through startup/streaming flows, and config serialization.
  • Validation

    • New validation for storage-context mappings (reject empty names/keys and duplicate context keys).

@coderabbitai

coderabbitai Bot commented Mar 18, 2026 •

Copy link
Copy Markdown

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

Adds 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

Cohort / File(s) Summary
Python FFI & docs
bindings/python/src/lib.rs, bindings/python/src/smg/router.py, docs/reference/configuration.md
Expose storage_context_headers on Router FFI, document --storage-context-headers with example and default {}.
Python CLI args & dataclass
bindings/python/src/smg/router_args.py, bindings/python/src/smg/router_args.py
Add storage_context_headers: dict[str,str] to RouterArgs and parse --storage-context-headers into the dict via selector parsing.
Python tests
bindings/python/tests/test_arg_parser.py, bindings/python/tests/test_router_config.py, bindings/python/tests/test_startup_sequence.py, bindings/python/tests/test_validation.py
Add/adjust tests to assert CLI parsing, dataclass persistence, startup propagation, and validation behavior for storage_context_headers.
Rust config types & builder
model_gateway/src/config/types.rs, model_gateway/src/config/builder.rs
Add storage_context_headers: HashMap<String,String> to RouterConfig (serde default + Default) and builder methods storage_context_headers(...) / maybe_storage_context_headers(...).
Rust CLI wiring
model_gateway/src/main.rs
Add storage_context_headers: Vec<String> CLI arg; parse selector->map and pass via maybe_storage_context_headers into the RouterConfig builder.
Rust validation
model_gateway/src/config/validation.rs
New validator to reject empty header names or context keys and duplicate context keys; wired into config validation with tests.
Rust middleware & helpers
model_gateway/src/middleware.rs
New storage_context_middleware and helpers to extract configured headers and build a StorageRequestContext; middleware conditionally attaches context to requests.
Rust server routing
model_gateway/src/server.rs
Apply storage_context_middleware as a route layer for protected endpoints so handlers run with the storage context when present.
Request-context propagation in streaming
model_gateway/src/routers/.../responses/streaming.rs (grpc/harmony, grpc/regular, openai)
Wrap spawned/background streaming tasks with request-context propagation (current_request_context / with_request_context) so storage/request context flows into streaming/background work; notable control-flow changes in OpenAI streaming path.

Sequence Diagram

sequenceDiagram
    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
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

Suggested reviewers

  • CatherineSue
  • key4ng

Poem

🐰 I hop through headers, keys in paw,
mapping whispers that requests once saw.
I tuck context bits in a cozy nest,
so storage hooks know what to request.
A tiny hop, a helpful jest.

🚥 Pre-merge checks | ✅ 3
✅ Passed checks (3 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title accurately describes the main feature: adding configurable storage hook request context mapping via headers. It directly reflects the core change across the entire changeset.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch feat/storage-context-headers

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Summary of Changes

Hello, 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

  • Configurable Storage Hook Request Context Mapping: Introduced a new configuration option, storage_context_headers, allowing users to map incoming HTTP headers to specific keys within the storage hook request context. This enables storage hooks to access custom request-specific metadata without hardcoding header names.
  • Cross-Language Integration: The new functionality is implemented in both the Rust-based model gateway and its Python bindings, ensuring consistent behavior and configurability across different interfaces.
  • CLI Argument and Documentation: Added a new command-line argument (--storage-context-headers) for easy configuration, along with comprehensive documentation in the configuration.md file to guide users on its usage and benefits.
  • API Endpoint Integration: Integrated the storage context building and application logic into various API endpoints within the model gateway, ensuring that the configured context is available for processing requests.
  • Comprehensive Testing: New unit tests were added for the storage_context_headers configuration, validation, and context building logic, ensuring the robustness and correctness of the new feature.
Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

@github-actions github-actions Bot added documentation Improvements or additions to documentation python-bindings Python bindings changes tests Test changes model-gateway Model gateway crate changes labels Mar 18, 2026

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@slin1237 slin1237 left a comment •

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_middleware is already in middleware.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_context is well-written: graceful handling of empty/whitespace keys, clean fast-path when no headers match.
  • Nit: main.rs:1135 reuses parse_selector() for key=value parsing — 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.

@zhaowenzi
zhaowenzi force-pushed the feat/storage-context-headers branch from ad03809 to 881b0db Compare March 18, 2026 19:11
@zhaowenzi

Copy link
Copy Markdown
Contributor Author

@slin1237
I moved it to middleware in middleware.rs.
The only difference is that I applied it only to the responses and conversations routes, not globally to all routes.

@zhaowenzi
zhaowenzi marked this pull request as ready for review March 18, 2026 20:54

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread model_gateway/src/middleware.rs

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between ce6f52e and 881b0db.

📒 Files selected for processing (13)
  • bindings/python/src/lib.rs
  • bindings/python/src/smg/router.py
  • bindings/python/src/smg/router_args.py
  • bindings/python/tests/test_arg_parser.py
  • bindings/python/tests/test_router_config.py
  • bindings/python/tests/test_startup_sequence.py
  • bindings/python/tests/test_validation.py
  • docs/reference/configuration.md
  • model_gateway/src/config/builder.rs
  • model_gateway/src/config/types.rs
  • model_gateway/src/main.rs
  • model_gateway/src/middleware.rs
  • model_gateway/src/server.rs

Comment thread bindings/python/src/smg/router_args.py
Comment thread docs/reference/configuration.md
Comment thread docs/reference/configuration.md
Comment thread model_gateway/src/main.rs
Comment thread model_gateway/src/middleware.rs
@github-actions github-actions Bot added grpc gRPC client and router changes openai OpenAI router changes labels Mar 18, 2026
@zhaowenzi
zhaowenzi force-pushed the feat/storage-context-headers branch from 10259a8 to 830a3fd Compare March 18, 2026 21:36

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread model_gateway/src/middleware.rs

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 881b0db and 830a3fd.

📒 Files selected for processing (16)
  • bindings/python/src/lib.rs
  • bindings/python/src/smg/router.py
  • bindings/python/src/smg/router_args.py
  • bindings/python/tests/test_arg_parser.py
  • bindings/python/tests/test_router_config.py
  • bindings/python/tests/test_startup_sequence.py
  • bindings/python/tests/test_validation.py
  • docs/reference/configuration.md
  • model_gateway/src/config/builder.rs
  • model_gateway/src/config/types.rs
  • model_gateway/src/main.rs
  • model_gateway/src/middleware.rs
  • model_gateway/src/routers/grpc/harmony/responses/streaming.rs
  • model_gateway/src/routers/grpc/regular/responses/streaming.rs
  • model_gateway/src/routers/openai/responses/streaming.rs
  • model_gateway/src/server.rs

Comment thread docs/reference/configuration.md
Comment thread model_gateway/src/routers/openai/responses/streaming.rs Outdated

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread model_gateway/src/config/validation.rs Outdated

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread model_gateway/src/config/validation.rs
Comment thread model_gateway/src/main.rs

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 830a3fd and 2a2bfa8.

📒 Files selected for processing (1)
  • model_gateway/src/config/validation.rs

Comment thread model_gateway/src/config/validation.rs
@zhaowenzi
zhaowenzi force-pushed the feat/storage-context-headers branch from 2a2bfa8 to c397cee Compare March 18, 2026 23:59

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

♻️ Duplicate comments (2)
docs/reference/configuration.md (1)

452-471: ⚠️ Potential issue | 🟡 Minor

Clarify that --storage-context-headers is 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 | 🟡 Minor

Strengthen tests to assert the exact validation failure.

These tests only check is_err(), so they can pass on unrelated failures. Assert ConfigError::ValidationFailed with 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

📥 Commits

Reviewing files that changed from the base of the PR and between 2a2bfa8 and c397cee.

📒 Files selected for processing (17)
  • bindings/python/src/lib.rs
  • bindings/python/src/smg/router.py
  • bindings/python/src/smg/router_args.py
  • bindings/python/tests/test_arg_parser.py
  • bindings/python/tests/test_router_config.py
  • bindings/python/tests/test_startup_sequence.py
  • bindings/python/tests/test_validation.py
  • docs/reference/configuration.md
  • model_gateway/src/config/builder.rs
  • model_gateway/src/config/types.rs
  • model_gateway/src/config/validation.rs
  • model_gateway/src/main.rs
  • model_gateway/src/middleware.rs
  • model_gateway/src/routers/grpc/harmony/responses/streaming.rs
  • model_gateway/src/routers/grpc/regular/responses/streaming.rs
  • model_gateway/src/routers/openai/responses/streaming.rs
  • model_gateway/src/server.rs

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread bindings/python/src/smg/router_args.py

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread model_gateway/src/server.rs
Comment thread model_gateway/src/main.rs
@zhaowenzi
zhaowenzi force-pushed the feat/storage-context-headers branch from fa0d6e7 to f6cbbdb Compare March 20, 2026 17:57

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread bindings/python/src/smg/router_args.py
@zhaowenzi
zhaowenzi force-pushed the feat/storage-context-headers branch from f6cbbdb to 998b95b Compare March 23, 2026 06:13
@mergify

mergify Bot commented Mar 23, 2026

Copy link
Copy Markdown
Contributor

Hi @zhaowenzi, the DCO sign-off check has failed. All commits must include a Signed-off-by line.

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-lease

To sign off future commits automatically:

  • Use git commit -s every time, or
  • VSCode: enable Git: Always Sign Off in Settings
  • PyCharm: enable Sign-off commit in the Commit tool window

@zhaowenzi
zhaowenzi force-pushed the feat/storage-context-headers branch 2 times, most recently from 62df3a6 to e8ec1dc Compare March 23, 2026 20:22
@chatgpt-codex-connector

Copy link
Copy Markdown

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Repo admins can enable using credits for code reviews in their settings.

Signed-off-by: Ziwen Zhao <zzw.mose@gmail.com>
@zhaowenzi
zhaowenzi force-pushed the feat/storage-context-headers branch from e8ec1dc to a3d21ce Compare March 23, 2026 23:04
@chatgpt-codex-connector

Copy link
Copy Markdown

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Repo admins can enable using credits for code reviews in their settings.

@zhaowenzi
zhaowenzi requested a review from slin1237 March 24, 2026 04:35
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>
@chatgpt-codex-connector

Copy link
Copy Markdown

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Repo admins can enable using credits for code reviews in their settings.

@github-actions github-actions Bot removed python-bindings Python bindings changes tests Test changes labels Mar 24, 2026
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>
@chatgpt-codex-connector

Copy link
Copy Markdown

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Repo admins can enable using credits for code reviews in their settings.

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>
@chatgpt-codex-connector

Copy link
Copy Markdown

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Repo admins can enable using credits for code reviews in their settings.

@github-actions github-actions Bot added python-bindings Python bindings changes tests Test changes labels Mar 24, 2026
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>
@chatgpt-codex-connector

Copy link
Copy Markdown

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Repo admins can enable using credits for code reviews in their settings.

@slin1237
slin1237 merged commit 84b7593 into main Mar 24, 2026
30 of 35 checks passed
@slin1237
slin1237 deleted the feat/storage-context-headers branch March 24, 2026 22:35
smfirmin pushed a commit to smfirmin/smg that referenced this pull request Apr 2, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation grpc gRPC client and router changes model-gateway Model gateway crate changes openai OpenAI router changes python-bindings Python bindings changes tests Test changes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants