fix(gateway): harden AI provider HTTP error handling and context propagation#258
Conversation
…agation - Cap error response body reads to 4KB in both OpenRouter and Ollama providers. Previously, io.ReadAll(resp.Body) on non-2xx responses could read unlimited data from a malicious/misconfigured upstream, causing OOM. Now uses io.LimitReader to cap at 4096 bytes, matching the existing pattern in verifyPayment(). - Handle context.Canceled alongside context.DeadlineExceeded in AI provider error paths. During graceful shutdown, Go cancels contexts (not deadline-exceeds them). Without this fix, shutdown cancellations were misclassified and wrapped in connection-error messages instead of being returned as clean context errors. - Apply the same context.Canceled fix to handleSummarize() in main.go so the gateway returns 504 (not 502) during shutdown-induced cancellations, giving clients actionable retry signals. - Add unit tests covering both the body-limit cap and context.Canceled behavior for both providers.
|
@tlnarayana005 is attempting to deploy a commit to the ankanmisra's projects Team on Vercel. A member of the Team first needs to authorize it. |
|
Hi @tlnarayana005, thanks for opening this PR. Every contribution helps MicroAI-Paygate grow. If you find the project useful, consider starring the repository — it helps others discover it. Star MicroAI-Paygate on GitHub Looking forward to reviewing this PR. |
github-actions
Bot
added
go
📝 WalkthroughWalkthroughThis PR adds byte limits on error response bodies read from OpenRouter and Ollama providers, refines context deadline/cancellation detection in both providers' ChangesAI provider error handling improvements
Estimated code review effort: 2 (Simple) | ~12 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.12.2)level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies" 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.
🧹 Nitpick comments (3)
gateway/main.go (1)
484-494: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winVerifier-call error classification wasn't updated for
context.Canceled, unlike the AI-provider path just below.Line 488 still only checks
errors.Is(err, context.DeadlineExceeded); a canceled context (e.g., graceful shutdown) duringverifyPaymentwould fall into the502 verification_unavailablebranch instead of the504 verifier_timeoutbranch that the AI path now gets via the update at Lines 526-528. As per coding guidelines, "Every outbound verifier/AI/Redis call in the gateway should respect context deadlines and handle cancellation" — worth aligning this branch for consistency.♻️ Suggested alignment
- if errors.Is(err, context.DeadlineExceeded) { + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { respondError(c, 504, "verifier_timeout", err) } else { respondError(c, 502, "verification_unavailable", err) }🤖 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 `@gateway/main.go` around lines 484 - 494, In the verifyPayment error handling inside the gateway main flow, the classification only treats context.DeadlineExceeded as a timeout, so context.Canceled still falls through to verification_unavailable. Update the error branch around verifyPayment to treat both context.DeadlineExceeded and context.Canceled as verifier_timeout, matching the AI-provider handling below and keeping cancellation semantics consistent.Source: Coding guidelines
gateway/internal/ai/openrouter.go (2)
68-73: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRedundant condition — simplify to
ctx.Err() != nil.
context.Context.Err()only ever returnsnil,context.Canceled, orcontext.DeadlineExceeded, and once non-nil it stays non-nil. Sincectxhere is the exact context attached viahttp.NewRequestWithContext(ctx, ...), anyerrfromDo()wrapping either sentinel impliesctx.Err()is already set. Theerrors.Is(...)checks add no coverage beyondctx.Err() != nil.♻️ Simplification
- if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) || - ctx.Err() == context.DeadlineExceeded || ctx.Err() == context.Canceled { + if ctx.Err() != nil { return "", ctx.Err() } return "", err🤖 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 `@gateway/internal/ai/openrouter.go` around lines 68 - 73, The error handling in the OpenRouter request path is overly defensive and duplicates cancellation checks. In the function that calls http.NewRequestWithContext and Do in openrouter.go, simplify the branch by relying on ctx.Err() != nil instead of checking both errors.Is against context.DeadlineExceeded/context.Canceled and ctx.Err() values; keep the behavior of returning ctx.Err() when the request was canceled or timed out, otherwise return err unchanged.
74-80: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueBounded error-body read looks correct; consider draining remainder for connection reuse.
The
io.LimitReader(resp.Body, maxErrorResponseBytes)cap correctly prevents unbounded reads. Since only up to 4KB is consumed beforeClose(), the underlyingnet/httptransport will typically close rather than reuse the TCP connection when the upstream error body exceeds the cap.♻️ Optional: drain remaining body before close
if resp.StatusCode < 200 || resp.StatusCode >= 300 { body, _ := io.ReadAll(io.LimitReader(resp.Body, maxErrorResponseBytes)) + io.Copy(io.Discard, resp.Body) // drain to allow connection reuse return "", fmt.Errorf("openrouter returned status %d: %s", resp.StatusCode, string(body)) }🤖 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 `@gateway/internal/ai/openrouter.go` around lines 74 - 80, The error path in openrouter.go only reads a capped portion of resp.Body before closing it, which can prevent HTTP connection reuse on oversized upstream error responses. Update the status-check block in the OpenRouter request handling to explicitly drain any remaining response body after reading the limited error bytes, while keeping the maxErrorResponseBytes cap and the existing error message in the same resp.StatusCode handling path.
🤖 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 `@gateway/internal/ai/openrouter.go`:
- Around line 68-73: The error handling in the OpenRouter request path is overly
defensive and duplicates cancellation checks. In the function that calls
http.NewRequestWithContext and Do in openrouter.go, simplify the branch by
relying on ctx.Err() != nil instead of checking both errors.Is against
context.DeadlineExceeded/context.Canceled and ctx.Err() values; keep the
behavior of returning ctx.Err() when the request was canceled or timed out,
otherwise return err unchanged.
- Around line 74-80: The error path in openrouter.go only reads a capped portion
of resp.Body before closing it, which can prevent HTTP connection reuse on
oversized upstream error responses. Update the status-check block in the
OpenRouter request handling to explicitly drain any remaining response body
after reading the limited error bytes, while keeping the maxErrorResponseBytes
cap and the existing error message in the same resp.StatusCode handling path.
In `@gateway/main.go`:
- Around line 484-494: In the verifyPayment error handling inside the gateway
main flow, the classification only treats context.DeadlineExceeded as a timeout,
so context.Canceled still falls through to verification_unavailable. Update the
error branch around verifyPayment to treat both context.DeadlineExceeded and
context.Canceled as verifier_timeout, matching the AI-provider handling below
and keeping cancellation semantics consistent.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 2411a38d-7728-42ca-a3ec-6937828f15dc
📒 Files selected for processing (4)
gateway/internal/ai/error_handling_test.gogateway/internal/ai/ollama.gogateway/internal/ai/openrouter.gogateway/main.go
|
@codex review |
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e6ab8f128c
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) || | ||
| ctx.Err() == context.DeadlineExceeded || ctx.Err() == context.Canceled { | ||
| return "", ctx.Err() |
There was a problem hiding this comment.
When http.DefaultClient.Do returns a deadline/cancellation error from the client or transport rather than from the passed request context, ctx.Err() is still nil; for example, setting http.DefaultClient = &http.Client{Timeout: ...} makes this branch return ("", nil). That makes Generate look successful with an empty summary, so a paid summarize request can get a 200/receipt instead of an upstream error. Return ctx.Err() only when it is non-nil, otherwise return the matched err or the appropriate context sentinel; the identical Ollama branch needs the same treatment.
Useful? React with 👍 / 👎.
AnkanMisra
added
bug
ChatGPT Codex Connector
2 participants
fix(gateway): harden AI provider HTTP error handling and context propagation
Cap error response body reads to 4KB in both OpenRouter and Ollama providers. Previously, io.ReadAll(resp.Body) on non-2xx responses could read unlimited data from a malicious/misconfigured upstream, causing OOM. Now uses io.LimitReader to cap at 4096 bytes, matching the existing pattern in verifyPayment().
Handle context.Canceled alongside context.DeadlineExceeded in AI provider error paths. During graceful shutdown, Go cancels contexts (not deadline-exceeds them). Without this fix, shutdown cancellations were misclassified and wrapped in connection-error messages instead of being returned as clean context errors.
Apply the same context.Canceled fix to handleSummarize() in main.go so the gateway returns 504 (not 502) during shutdown-induced cancellations, giving clients actionable retry signals.
Add unit tests covering both the body-limit cap and context.Canceled behavior for both providers.
Summary
io.LimitReader) to prevent potential OOM DOS vulnerabilities from malicious or misconfigured upstreams.context.Canceledhandling inopenrouter.go,ollama.go, andmain.goto properly return 504 Gateway Timeout instead of 502 Bad Gateway during graceful shutdowns.Type Of Change
Affected Areas
gateway/)verifier/)web/)tests/,run_e2e.sh)bench/)deploy/, Docker, env, workflows)Contributor Checklist
Verification