Propagate stable agent session identity from the OpenAI boundary - #1152
Propagate stable agent session identity from the OpenAI boundary#1152i386 wants to merge 0 commit into
Conversation
📝 WalkthroughWalkthroughThe change adds validated agent-session identities from trusted headers or Responses conversations. It attaches identity metadata to frontend requests, includes the ID in generation identifiers, and propagates it into generation receipts. ChangesAgent session propagation
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant OpenAIClient
participant OpenAiFrontendRouter
participant AgentSessionResolver
participant OpenAiGenerationIds
participant GenerationReceipt
OpenAIClient->>OpenAiFrontendRouter: Send session header or conversation
OpenAiFrontendRouter->>AgentSessionResolver: Resolve agent-session identity
AgentSessionResolver->>OpenAiGenerationIds: Attach agent_session_id
OpenAiGenerationIds->>GenerationReceipt: Propagate agent_session_id
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
crates/openai-frontend/src/router.rs (1)
1148-1175: 🔒 Security & Privacy | 🔵 Trivial | ⚡ Quick winAdd a test that a client cannot inject the internal agent-session keys through the request body.
extrais#[serde(flatten)], so a client can placemesh_internal_agent_session_iddirectly in the JSON body.set_agent_sessionremoves both internal keys before it inserts, so the current code strips such input. No test locks that behavior. A future change that skipsset_agent_sessionon any path would silently open a spoofing hole.Do you want me to generate this test?
💚 Sketch of the anti-spoofing test
#[tokio::test] async fn body_supplied_internal_agent_session_keys_are_stripped() { let backend = Arc::new(SessionCaptureBackend::default()); let app = router_for(backend.clone()); let response = app .oneshot( Request::builder() .method("POST") .uri("/v1/chat/completions") .header("content-type", "application/json") .body(Body::from( json!({ "model": "capture-model", "messages": [{"role": "user", "content": "hello"}], "mesh_internal_agent_session_id": "spoofed", "mesh_internal_agent_session_source": "spoofed-source" }) .to_string(), )) .unwrap(), ) .await .unwrap(); assert_eq!(response.status(), StatusCode::OK); let requests = backend.requests.lock().unwrap(); assert!(requests[0].agent_session().is_none()); assert!(requests[0].agent_session_source().is_none()); }🤖 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 `@crates/openai-frontend/src/router.rs` around lines 1148 - 1175, Add a regression test alongside unconfigured_session_header_is_ignored that sends mesh_internal_agent_session_id and mesh_internal_agent_session_source in the request JSON body, then verifies the request captured by SessionCaptureBackend has neither agent session value nor source. Preserve the successful response assertion and use the existing router_for and backend capture flow.crates/skippy-server/src/frontend/tests/generation.rs (1)
189-199: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAssert that a shared agent session still yields distinct generation IDs.
The central constraint of this PR is that the agent-session identity must not be conflated with the runtime session ID.
OpenAiGenerationIds::newcurrently derivessession_idandrequest_idfromsession_labelonly, so the constraint holds. This test passesNoneand therefore does not cover it. Add a case that reuses one agent session across constructions.💚 Proposed additional test
#[test] fn shared_agent_session_still_yields_distinct_generation_ids() { let ids = (0..64) .map(|_| OpenAiGenerationIds::new(OpenAiCacheHints::default(), Some("agent-thread-1"))) .collect::<Vec<_>>(); let mut sessions = std::collections::BTreeSet::new(); for id in &ids { assert_eq!(id.agent_session_id.as_deref(), Some("agent-thread-1")); assert!(sessions.insert(id.session_id)); } }🤖 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 `@crates/skippy-server/src/frontend/tests/generation.rs` around lines 189 - 199, Add a test alongside generation_ids_are_unique_under_fast_creation that repeatedly constructs OpenAiGenerationIds with the same agent session label, such as “agent-thread-1”. Verify each result preserves that agent_session_id and that the generated session_id values remain distinct using a BTreeSet.crates/openai-frontend/src/completions.rs (1)
47-74: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated agent-session metadata code in
crates/openai-frontend/src/completions.rsandcrates/openai-frontend/src/chat.rs. Both files define the same two internal key constants and the sameset_agent_session,agent_session, andagent_session_sourcebodies. The shared root cause is one piece of logic copied into two request types instead of living in the module that already ownsAgentSessionIdentity. If one copy changes, the two endpoints disagree on the key names, and the router's strip-then-set anti-spoofing guarantee breaks for one of them.
crates/openai-frontend/src/completions.rs#L47-L74: replace the three method bodies with calls to shared helpers incommon, and delete the local key constants at Lines 14-15.crates/openai-frontend/src/chat.rs#L14-L15: delete the duplicate key constants and route the three methods at Lines 51-78 through the same shared helpers incommon.🤖 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 `@crates/openai-frontend/src/completions.rs` around lines 47 - 74, Move the shared agent-session metadata logic into helpers in common, then update the methods in crates/openai-frontend/src/completions.rs:47-74 and crates/openai-frontend/src/chat.rs:51-78 to delegate to those helpers; delete the duplicate key constants from crates/openai-frontend/src/completions.rs:14-15 and crates/openai-frontend/src/chat.rs:14-15. Preserve the existing set, get, and anti-spoofing behavior through the centralized helpers.
🤖 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 `@crates/openai-frontend/src/router.rs`:
- Around line 462-495: The agent-session extraction and resolution
responsibility should move out of router.rs into a dedicated agent_session
module. Relocate agent_session_from_header, resolve_agent_session, their
required imports/types, and related tests into the new module, expose the
functions as needed, and update router handlers to call the module APIs while
preserving existing validation and conflict behavior.
- Around line 469-471: Update the trusted-header extraction around
headers.get(name) to inspect all occurrences of the configured header. Return
the existing missing-header result when absent, accept repeated values only when
they are identical, and fail closed with the same conflict behavior used for
header-versus-conversation mismatches when duplicate values differ.
---
Nitpick comments:
In `@crates/openai-frontend/src/completions.rs`:
- Around line 47-74: Move the shared agent-session metadata logic into helpers
in common, then update the methods in
crates/openai-frontend/src/completions.rs:47-74 and
crates/openai-frontend/src/chat.rs:51-78 to delegate to those helpers; delete
the duplicate key constants from crates/openai-frontend/src/completions.rs:14-15
and crates/openai-frontend/src/chat.rs:14-15. Preserve the existing set, get,
and anti-spoofing behavior through the centralized helpers.
In `@crates/openai-frontend/src/router.rs`:
- Around line 1148-1175: Add a regression test alongside
unconfigured_session_header_is_ignored that sends mesh_internal_agent_session_id
and mesh_internal_agent_session_source in the request JSON body, then verifies
the request captured by SessionCaptureBackend has neither agent session value
nor source. Preserve the successful response assertion and use the existing
router_for and backend capture flow.
In `@crates/skippy-server/src/frontend/tests/generation.rs`:
- Around line 189-199: Add a test alongside
generation_ids_are_unique_under_fast_creation that repeatedly constructs
OpenAiGenerationIds with the same agent session label, such as “agent-thread-1”.
Verify each result preserves that agent_session_id and that the generated
session_id values remain distinct using a BTreeSet.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 3dd53144-a96d-4709-b03b-16967c9cfcac
📒 Files selected for processing (12)
crates/openai-frontend/src/chat.rscrates/openai-frontend/src/common.rscrates/openai-frontend/src/completions.rscrates/openai-frontend/src/lib.rscrates/openai-frontend/src/responses.rscrates/openai-frontend/src/router.rscrates/skippy-server/src/frontend/backend.rscrates/skippy-server/src/frontend/generation/cache_hints.rscrates/skippy-server/src/frontend/generation_flow.rscrates/skippy-server/src/frontend/generation_receipt.rscrates/skippy-server/src/frontend/local_generation.rscrates/skippy-server/src/frontend/tests/generation.rs
|
This pull request is currently a draft. Reviews will not take place until the PR is marked as ready for review. |
|
Oh this is interesting but not sure if the "why" - is it if they are from same session from routing or kv optimisations? |
|
Yeah, it's basically so that proxy sitting in front of the open AI gateway can't tell whether or not to request has been processed by the same agent or person. I need it for stuff. |
michaelneale
left a comment
michaelneale
left a comment
There was a problem hiding this comment.
Low risk: opt-in (defaults off), identity is opaque/length-bounded/control-char-rejected and explicitly kept out of prompt-cache keys and runtime session IDs, so no cache-poisoning or KV-collision risk. Client spoofing blocked via extra.remove() before set. GenerationReceipt field is not serialized on any wire. Clean, well-isolated. LGTM.
There was a problem hiding this comment.
🧹 Nitpick comments (1)
crates/openai-frontend/src/router.rs (1)
205-218: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd direct coverage for the new Responses and completion paths.
The current tests prove trusted-header propagation on
/v1/chat/completionsand conversation propagation on/v1/responses. They do not prove header-only propagation on/v1/responsesor/v1/completions. If identical normalized values from both sources are valid, add that acceptance case as well.Also applies to: 429-433
🤖 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 `@crates/openai-frontend/src/router.rs` around lines 205 - 218, Add direct tests for the `/v1/responses` and `/v1/completions` handlers covering header-only agent-session propagation, alongside the existing chat-completions and Responses conversation tests. Include an acceptance case where normalized header and conversation session values are identical and verify the resolved session is accepted.
🤖 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.
Nitpick comments:
In `@crates/openai-frontend/src/router.rs`:
- Around line 205-218: Add direct tests for the `/v1/responses` and
`/v1/completions` handlers covering header-only agent-session propagation,
alongside the existing chat-completions and Responses conversation tests.
Include an acceptance case where normalized header and conversation session
values are identical and verify the resolved session is accepted.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: df3b4ff5-edee-4b7f-9a1b-5f8759467099
📒 Files selected for processing (3)
crates/openai-frontend/src/agent_session.rscrates/openai-frontend/src/lib.rscrates/openai-frontend/src/router.rs
🚧 Files skipped from review as they are similar to previous changes (2)
- crates/openai-frontend/src/lib.rs
- crates/openai-frontend/src/agent_session.rs
27809cd to
4bd1453
Compare
2 participants
Summary
conversation.idand one operator-configured trusted headerWhy
Long-lived consumers need a neutral way to correlate multiple OpenAI requests from the same agent session. This adds the transport contract only; downstream products remain responsible for authentication, persistence, and policy.
Validation
cargo fmt --all -- --checkcargo check --workspace --all-targetscargo clippy -p openai-frontend -p skippy-server --all-targets -- -D warningscargo test -p openai-frontend -p skippy-server --quietgit diff --checkSummary by CodeRabbit
New Features
Bug Fixes
Tests