fix: forward provider timeout and bound fallback chains - #87
Conversation
📝 WalkthroughWalkthroughThe shim adds configurable router execution deadlines and propagates optional provider-attempt timeouts through chat and Responses requests. Compact, unary, and streaming paths now enforce the deadline, returning structured timeout failures and cancelling stalled host execution. ChangesRouter timeout controls
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant FastAPIShim
participant RouterHost
Client->>FastAPIShim: Send request with timeout_ms
FastAPIShim->>RouterHost: Execute contract with timeout_ms
RouterHost-->>FastAPIShim: Complete or exceed app deadline
FastAPIShim-->>Client: Normal response or structured 504 timeout
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 |
![coderabbitai[bot]](https://avatars.github.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@shim.py`:
- Around line 264-294: Apply the existing _execute_with_deadline wrapper to the
unary stream:false host.execute_flow_async(...) call around the flow-dispatch
logic, so non-streaming requests return the structured timeout response when
request_deadline_ms expires. Leave the heartbeat-enabled streaming-flow path
unchanged and preserve the existing awaitable arguments and result handling.
- Line 950: Update the compaction result handling around _execute_with_deadline
so a timeout result is passed to _openai_error_from_router and surfaces as HTTP
504 before the empty-summary conversion. Preserve the existing 200 {"compacted":
false} soft-failure behavior for non-timeout seal failures.
🪄 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: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 2d63d567-f2ce-4f32-b47d-d7eec79d9f02
📒 Files selected for processing (5)
.env.exampleshim.pytests/test_responses_shim.pytests/test_shim.pytests/test_shim_max_tokens.py
| async def _execute_with_deadline(awaitable): | ||
| """Await one complete router run and cancel it at the outer deadline. | ||
|
|
||
| asyncio.timeout propagates cancellation into the active provider call, | ||
| so an expired request does not leave an orphan fallback chain running | ||
| after the HTTP response has finished. | ||
| """ | ||
| deadline = asyncio.timeout(request_deadline_s) | ||
| try: | ||
| async with deadline: | ||
| return await awaitable | ||
| except TimeoutError: | ||
| # Do not relabel a TimeoutError raised by the host itself: only the | ||
| # timeout context's own expiry is the outer request deadline. | ||
| if not deadline.expired(): | ||
| raise | ||
| _log.warning( | ||
| "router request deadline exceeded after %d ms", | ||
| request_deadline_ms, | ||
| ) | ||
| return { | ||
| "ok": False, | ||
| "error": "timeout", | ||
| "trace": { | ||
| "decision_path": [], | ||
| "request_deadline_exceeded": True, | ||
| "request_deadline_ms": request_deadline_ms, | ||
| "total_latency_ms": request_deadline_ms, | ||
| }, | ||
| } | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Apply the deadline to unary flow execution.
Line 264 defines the complete-execution deadline, but shim.py Line 1133 calls host.execute_flow_async(...) directly. A stream: false flow can therefore outlive request_deadline_ms and be cut off upstream instead of returning the structured 504. Keep the heartbeat-enabled streaming-flow behavior unchanged, but wrap the unary call.
Proposed fix
- result = await host.execute_flow_async(req.flow_ir, contract)
+ result = await _execute_with_deadline(
+ host.execute_flow_async(req.flow_ir, contract)
+ )🤖 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 `@shim.py` around lines 264 - 294, Apply the existing _execute_with_deadline
wrapper to the unary stream:false host.execute_flow_async(...) call around the
flow-dispatch logic, so non-streaming requests return the structured timeout
response when request_deadline_ms expires. Leave the heartbeat-enabled
streaming-flow path unchanged and preserve the existing awaitable arguments and
result handling.
| } | ||
| try: | ||
| res = await host.execute_async(contract) | ||
| res = await _execute_with_deadline(host.execute_async(contract)) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Return a 504 when compaction hits the deadline.
Line 950 can return the structured {ok: False, error: "timeout"} result, but the subsequent empty-summary path converts it into 200 {"compacted": false}. Surface deadline expiry through _openai_error_from_router before preserving the existing soft-failure behavior for other seal failures.
Proposed fix
try:
res = await _execute_with_deadline(host.execute_async(contract))
except Exception as exc:
admission = _policy_admission_error(exc)
if admission is not None:
return _invalid_policy_response(admission)
raise
+ if res.get("error") == "timeout":
+ return _openai_error_from_router(res)
def _costed(body: dict) -> dict:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| res = await _execute_with_deadline(host.execute_async(contract)) | |
| res = await _execute_with_deadline(host.execute_async(contract)) | |
| if res.get("error") == "timeout": | |
| return _openai_error_from_router(res) |
🤖 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 `@shim.py` at line 950, Update the compaction result handling around
_execute_with_deadline so a timeout result is passed to
_openai_error_from_router and surfaces as HTTP 504 before the empty-summary
conversion. Preserve the existing 200 {"compacted": false} soft-failure behavior
for non-timeout seal failures.
1 participant
Why
Production unary requests can outlive the 60-second ALB idle timeout when several provider fallbacks run sequentially. The ingress then reports an opaque gateway failure even though the router pod and target are healthy. Incoming
timeout_mswas also accepted by the permissive request model but dropped before reaching the router contract.What
timeout_msfor Chat Completions and Responses requestsROUTER_REQUEST_DEADLINE_MS, leaving margin before the ALB cutoffTesting
nix-shell --run "python -m pytest tests -q"Summary by CodeRabbit
New Features
timeout_mssupport for chat and responses requests.Bug Fixes