fix: preserve hedge model redirect audit per attempt#992
Conversation
📝 WalkthroughWalkthrough在代理会话与流式对冲逻辑中引入并传播“当前模型重定向(currentModelRedirect)”快照;在会话层增加管理接口,并将每个对冲尝试的 modelRedirect 延迟物化与记录到对冲审计链和阴影会话的克隆/合并流程中。 Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 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 |
There was a problem hiding this comment.
Code Review
This pull request introduces a mechanism to track and persist model redirects within proxy sessions, ensuring that redirect metadata is correctly associated with provider attempts and recorded in audit logs. It adds a currentModelRedirect snapshot to ProxySession and updates ProxyForwarder to handle redirect propagation during hedging and retries. Feedback was provided regarding a potential performance bottleneck in the loop used to capture redirects for all active attempts, suggesting an opportunity for optimization.
🧪 测试结果
总体结果: ✅ 所有测试通过 |
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 `@src/app/v1/_lib/proxy/forwarder.ts`:
- Around line 2998-3009: The client_abort bulk settlement path currently calls
session.addProviderToChain(...) directly on attempts and therefore can miss
snapshotting shadow attempts' modelRedirect; update the global client-abort
listener to use the existing getAttemptModelRedirect(attempt:
StreamingHedgeAttempt) helper to ensure each attempt (including shadow ones) has
its modelRedirect populated (call getAttemptModelRedirect for each attempt
before invoking session.addProviderToChain) so the audit trail includes the
redirect snapshot; also add a hedge client_abort test that verifies shadow
attempts' modelRedirects are preserved in the settlement/audit chain.
🪄 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
Run ID: 7a6bbd83-37fb-49f9-aa2c-882c5633101f
📒 Files selected for processing (4)
src/app/v1/_lib/proxy/forwarder.tssrc/app/v1/_lib/proxy/model-redirector.tssrc/app/v1/_lib/proxy/session.tstests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts
There was a problem hiding this comment.
Code Review Summary
No significant issues identified in this PR.
PR Size: L
- Lines changed: 505
- Files changed: 4
- Split suggestions: Consider splitting into (1)
ProxySession/ModelRedirectorstate tracking, (2)ProxyForwarderhedge audit plumbing, (3) hedge regression tests.
Review Coverage
- Logic and correctness - Clean
- Security (OWASP Top 10) - Clean
- Error handling - Clean
- Type safety - Clean
- Documentation accuracy - Clean
- Test coverage - Adequate
- Code clarity - Good
Automated review by Codex AI
🧪 测试结果
总体结果: ✅ 所有测试通过 |
There was a problem hiding this comment.
Code Review Summary
This PR introduces a well-designed per-attempt redirect snapshot model that correctly decouples model redirect audit from chain position. The snapshot-based approach, deep-cloning on shadow creation/sync, and lazy caching in getAttemptModelRedirect are all solid architectural choices. Five comprehensive test cases cover the critical paths.
One inconsistency found: the signal handler's client_abort path was not updated to include redirect metadata for shadow attempts, while the equivalent handleAttemptFailure path was fixed.
PR Size: L
- Lines changed: 505
- Files changed: 4
Issues Found
| Category | Critical | High | Medium | Low |
|---|---|---|---|---|
| Logic/Bugs | 0 | 0 | 1 | 0 |
| Security | 0 | 0 | 0 | 0 |
| Error Handling | 0 | 0 | 0 | 0 |
| Types | 0 | 0 | 0 | 0 |
| Comments/Docs | 0 | 0 | 0 | 0 |
| Tests | 0 | 0 | 0 | 0 |
| Simplification | 0 | 0 | 0 | 0 |
Medium Priority Issues
1. Signal handler client_abort path missing redirect metadata for shadow attempts
File: src/app/v1/_lib/proxy/forwarder.ts (~line 3539, unchanged code)
Confidence: 82
Why this is a problem: The PR adds modelRedirect: getAttemptModelRedirect(attempt) to the handleAttemptFailure CLIENT_ABORT path (line 3248), correctly capturing per-attempt redirect snapshots. However, the clientAbortSignal event handler (~line 3539) has an equivalent client_abort code path that was NOT updated. These two paths are mutually exclusive for a given attempt:
- Signal handler path (fires first when client disconnects): adds
client_abortentries for ALL unsettled attempts, then callsabortAllAttempts(undefined, "client_abort"). SinceabortAttemptsetsattempt.settled = true, subsequenthandleAttemptFailurecalls bail early. - handleAttemptFailure CLIENT_ABORT path (fires when an attempt's error is categorized as CLIENT_ABORT before the signal): adds entry with redirect metadata.
The signal handler uses session.addProviderToChain(attempt.provider, {...}) on the main session. The addProviderToChain fallback (this.getCurrentModelRedirect(provider.id)) reads from the main session's snapshot. For the initial attempt, this works (main session has the initial provider's redirect). But for shadow attempts (whose redirect lives on the shadow session, not the main session), the main session's getCurrentModelRedirect(shadowProvider.id) returns undefined because the providerId doesn't match.
This contradicts the PR's stated goal: "its chain entry would lack modelRedirect since the redirect was attached to a different chain item".
Suggested fix:
// In the clientAbortSignal handler (~line 3539 in PR version):
for (const attempt of Array.from(attempts)) {
if (!attempt.settled) {
session.addProviderToChain(attempt.provider, {
...attempt.endpointAudit,
reason: "client_abort",
attemptNumber: attempt.sequence,
errorMessage: "Client aborted request",
modelRedirect: getAttemptModelRedirect(attempt), // ADD THIS
});
}
}Review Coverage
- Logic and correctness
- Security (OWASP Top 10) - Clean
- Error handling - Clean
- Type safety - Clean
- Documentation accuracy - Clean
- Test coverage - Adequate (5 new tests cover critical paths; signal handler client_abort path not tested)
- Code clarity - Good
Automated review by Claude AI
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/app/v1/_lib/proxy/forwarder.ts (1)
96-102:⚠️ Potential issue | 🟠 Major把空快照也缓存下来。
当前 helper 只在拿到 redirect 时才写回
attempt.modelRedirect。如果这个 attempt 本来就没有 redirect,后续再次调用还是会重新读取 session 当前值,无法保持“每个 attempt 只截一次快照”的语义。这样后面的hedge_loser/retry_failed/client_abort/request_success仍然可能读到后来状态,而不是这个 attempt 当时的结果。建议修正
type StreamingHedgeAttempt = { provider: Provider; session: ProxySession; baseUrl: string; endpointAudit: { endpointId: number | null; endpointUrl: string }; modelRedirect?: ProviderChainItem["modelRedirect"]; + modelRedirectCaptured?: boolean; responseController: AbortController | null; clearResponseTimeout: (() => void) | null; firstByteTimeoutMs: number; sequence: number; requestAttemptCount: number; @@ const getAttemptModelRedirect = (attempt: StreamingHedgeAttempt) => { - if (attempt.modelRedirect !== undefined) { + if (attempt.modelRedirectCaptured) { return attempt.modelRedirect; } const redirect = attempt.session.getCurrentModelRedirect(attempt.provider.id); - if (redirect) { - attempt.modelRedirect = structuredClone(redirect); - } - + attempt.modelRedirect = redirect ? structuredClone(redirect) : undefined; + attempt.modelRedirectCaptured = true; return attempt.modelRedirect; };Also applies to: 2998-3009
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/app/v1/_lib/proxy/forwarder.ts` around lines 96 - 102, 当前 helper 只在有 redirect 时才把 snapshot 写回 attempt.modelRedirect,导致没有 redirect 的 attempt 以后再次读取会看到后续变更;请在负责设置快照的函数(修改 StreamingHedgeAttempt 的代码路径——即把 modelRedirect 写回 attempt 的 helper,另外也修改在 2998-3009 区块的同样逻辑)中无论是否有 redirect 都显式写回一个快照:例如把 attempt.modelRedirect 赋值为 session.modelRedirect 的当前值(若无则显式赋 null 或空对象),保证每个 attempt 只截取一次快照并持久化到 attempt.modelRedirect。
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@src/app/v1/_lib/proxy/forwarder.ts`:
- Around line 96-102: 当前 helper 只在有 redirect 时才把 snapshot 写回
attempt.modelRedirect,导致没有 redirect 的 attempt 以后再次读取会看到后续变更;请在负责设置快照的函数(修改
StreamingHedgeAttempt 的代码路径——即把 modelRedirect 写回 attempt 的 helper,另外也修改在
2998-3009 区块的同样逻辑)中无论是否有 redirect 都显式写回一个快照:例如把 attempt.modelRedirect 赋值为
session.modelRedirect 的当前值(若无则显式赋 null 或空对象),保证每个 attempt 只截取一次快照并持久化到
attempt.modelRedirect。
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 0d1522c2-1aaf-4840-83e7-39eb7f43cca5
📒 Files selected for processing (2)
src/app/v1/_lib/proxy/forwarder.tstests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts
Summary
Problem
When streaming requests use the hedge-based provider race (introduced in #894), each attempt may apply a different model redirect for its assigned provider. However, the
ModelRedirector.apply()was mutating the last item in the provider chain regardless of which provider attempt triggered the redirect. In a hedge scenario with shadow sessions, this caused:hedge_winner) or failed (retry_failed), its chain entry would lackmodelRedirectsince the redirect was attached to a different chain itemcurrentModelRedirectsnapshot was not deep-cloned when creating shadow sessions, allowing mutation to leak across attemptsSolution
Introduce a per-attempt redirect snapshot model that decouples the redirect audit from the chain's last item:
ProxySession.currentModelRedirect- New snapshot field that stores the redirect info keyed byproviderId, independent of chain positionsetCurrentModelRedirect()/getCurrentModelRedirect()/clearCurrentModelRedirect()- New session methods to manage the snapshot lifecycleattachCurrentModelRedirectToLastChainItem()- Opportunistically attaches the snapshot to the chain if the last item matches the provider; otherwise defers until the real chain entry is createdgetAttemptModelRedirect()in forwarder - Lazily captures and caches the redirect snapshot perStreamingHedgeAttempt, ensuring each attempt retains its own redirectcreateStreamingShadowSessionandsyncWinningAttemptSessionnowstructuredClone()the redirect snapshot so shadow and original sessions never share object referencesModelRedirector.apply()now callsclearCurrentModelRedirect()when a provider has no matching redirect rule, preventing stale data from leaking to subsequent chain entriesChanges
Core Changes
session.ts(+45): AddcurrentModelRedirectsnapshot field and accessor methods (setCurrentModelRedirect,getCurrentModelRedirect,clearCurrentModelRedirect,attachCurrentModelRedirectToLastChainItem); wiremodelRedirectintoaddProviderToChainmetadatamodel-redirector.ts(+18/-15): Replace direct chain-mutation with snapshot-based approach; callsession.setCurrentModelRedirect()andattachCurrentModelRedirectToLastChainItem()instead of mutatinglastDecision; addclearCurrentModelRedirect()calls when no redirect appliesforwarder.ts(+55/-5): AddmodelRedirecttoStreamingHedgeAttempttype; introducegetAttemptModelRedirect()helper for lazy snapshot capture; attach redirect tohedge_loser_cancelled,retry_failed,hedge_winner,request_success, and error chain entries; snapshot active attempts' redirects when a winner settles; deep-clonecurrentModelRedirectincreateStreamingShadowSessionandsyncWinningAttemptSessionTest Coverage
proxy-forwarder-hedge-first-byte.test.ts(+367): 5 new test cases:Related
Verification
bunx vitest run tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.tsbun run lint:fixbun run lintbun run typecheckbun run buildbun run testGreptile Summary
This PR fixes a redirect-audit correctness bug in the hedge streaming path: model redirect metadata was attached to whichever chain item happened to be last rather than to the attempt that actually performed the redirect, causing cross-contamination between concurrent shadow sessions.
The solution introduces a per-attempt
currentModelRedirectsnapshot onProxySessionwith providerId-scoped accessors, a lazy-capture helper (getAttemptModelRedirect) that caches the snapshot perStreamingHedgeAttempt, andstructuredCloneat shadow-creation and winner-sync boundaries to prevent reference aliasing. Test coverage is thorough, exercising winner, loser,retry_failed, snapshot isolation, and stale-clearing paths.Confidence Score: 5/5
Safe to merge; the snapshot model, pre-capture loop timing, and structuredClone boundaries are all correct across every hedge race path.
All findings are P2 style/hardening suggestions. The core correctness of the per-attempt redirect snapshot, the providerId-scoped fallback in addProviderToChain, and the pre-capture loop ordering (before syncWinningAttemptSession) is sound. No data-loss, security, or runtime-error risks identified.
No files require special attention; the two P2 comments on forwarder.ts are non-blocking.
Important Files Changed
Sequence Diagram
sequenceDiagram participant S as Main Session participant F as Forwarder (hedge) participant R as ModelRedirector participant Sh as Shadow Session F->>R: apply(session, fireworks) R->>S: setCurrentModelRedirect(fireworks.id, fw_redirect) R->>S: attachCurrentModelRedirectToLastChainItem(fireworks.id) Note over S: chain[initial_selection].modelRedirect = fw_redirect F->>Sh: createStreamingShadowSession(session, minimax) Note over Sh: currentModelRedirect = structuredClone(fw_redirect) F->>R: apply(shadow, minimax) R->>Sh: setCurrentModelRedirect(minimax.id, mm_redirect) Note over Sh: attachToLastChainItem → no-op (last item id is fireworks) F->>F: minimax first chunk arrives → commitWinner(minimaxAttempt) Note over F: Pre-capture loop: getAttemptModelRedirect(fwAttempt)<br/>reads + caches fw_redirect BEFORE session sync F->>S: syncWinningAttemptSession(session, shadow) Note over S: currentModelRedirect = structuredClone(mm_redirect) F->>S: addProviderToChain(minimax, hedge_winner, mm_redirect) F->>F: abortAllAttempts → abortAttempt(fwAttempt) F->>S: addProviderToChain(fireworks, hedge_loser_cancelled, cached fw_redirect)Prompt To Fix All With AI
Reviews (2): Last reviewed commit: "test(proxy): cover hedge redirect audit ..." | Re-trigger Greptile