From eebb6a1a6abeda317f1bf1e3e072ef282919270e Mon Sep 17 00:00:00 2001 From: etwk <48991073+etwk@users.noreply.github.com> Date: Sun, 18 Oct 2026 14:57:53 +0800 Subject: [PATCH 1/2] feat: support native Gemini GenerateContent --- .env.example | 7 +- README.md | 53 ++++- classify.go | 12 +- config_test.go | 2 + docker-compose.yml.example | 1 + docs/ARCHITECTURE.md | 59 ++++-- docs/CONFIGURATION.md | 24 ++- docs/DEVELOPMENT.md | 6 +- docs/OPERATIONS.md | 7 +- docs/TROUBLESHOOTING.md | 19 ++ gemini.go | 248 ++++++++++++++++++++++ gemini_errors.go | 181 ++++++++++++++++ gemini_errors_test.go | 308 +++++++++++++++++++++++++++ gemini_integration_test.go | 148 +++++++++++++ gemini_test.go | 415 +++++++++++++++++++++++++++++++++++++ main.go | 54 +++-- requestlog.go | 3 + requestlog_test.go | 4 +- responses.go | 8 + responses_test.go | 6 - sse.go | 37 ++-- wire.go | 83 ++++++-- 22 files changed, 1584 insertions(+), 101 deletions(-) create mode 100644 gemini.go create mode 100644 gemini_errors.go create mode 100644 gemini_errors_test.go create mode 100644 gemini_integration_test.go create mode 100644 gemini_test.go diff --git a/.env.example b/.env.example index fc8f733..05f032a 100644 --- a/.env.example +++ b/.env.example @@ -42,9 +42,14 @@ PROXY_RESPONSES_BUFFER_MS=600000 # against later stream failures, but delays visible output while buffering. #PROXY_RESPONSES_EARLY_COMMIT=1 +# Native Gemini streamGenerateContent?alt=sse: buffer through clean completion, +# then replay unchanged. Client timeouts must exceed this window. 0 keeps full +# buffering without a timed live switch; upstream idle/request deadlines remain. +PROXY_GEMINI_BUFFER_MS=600000 + # Extra attempts for recoverable streaming failures before sending a response. # 0 retains only one limited fast retry for selected connection/HTTP failures. -# Positive values enable broader retries on both Messages and Responses. +# Positive values enable broader retries on Messages, Responses, and Gemini SSE. PROXY_TRANSACTIONAL_LOCAL_RETRIES=0 # Local retries wait for upstream Retry-After plus an exponential extra delay. diff --git a/README.md b/README.md index 8dc9b16..70ad18c 100644 --- a/README.md +++ b/README.md @@ -24,19 +24,21 @@ It provides: Agents / SDKs / applications → steady-proxy on localhost → your gateway ``` -Stream protection is built around **Anthropic Messages** (`/v1/messages`) and -**OpenAI Responses** (`/v1/responses`). Claude Code and Codex are supported clients; -other agents and SDK-based applications can use these routes too, with compatible -timeout, retry, and SSE handling. Subagents using the same API settings are covered -without terminal automation. See [client compatibility](docs/ARCHITECTURE.md#client-compatibility) -for the requirements and behavior of other routes. +Stream protection supports **Anthropic Messages** (`/v1/messages`), +**OpenAI Responses** (`/v1/responses`), and native **Gemini GenerateContent** +(`/{version}/models/{model}:streamGenerateContent?alt=sse`). Claude Code and Codex +are supported clients; other agents and SDK-based applications can use these +routes with compatible timeout, retry, and SSE handling. Subagents using the same +API settings are covered without terminal automation. See +[client compatibility](docs/ARCHITECTURE.md#client-compatibility) for requirements +and the behavior of other routes. ## What to expect -By default, Messages streams are buffered for up to 10 minutes. A response that -finishes within that window appears all at once. Responses streams go live as soon -as output arrives; you can enable full buffering to protect them from later stream -failures. +By default, Messages and Gemini SSE streams are buffered for up to 10 minutes. +A response that finishes within that window appears all at once. Responses +streams go live as soon as output arrives; you can enable full buffering to +protect them from later stream failures. Before the proxy sends a response, it can discard a failed attempt and retry. After it starts streaming to the client, recovery depends on the client's own @@ -125,9 +127,38 @@ PROXY_TRANSACTIONAL_LOCAL_RETRIES=3 ``` Full buffering delays visible output until completion or the buffering window -ends. Extra retries apply to both supported streaming APIs and must fit within +ends. Extra retries apply to all three supported streaming APIs and must fit within the active window. See [retry budgets](docs/ARCHITECTURE.md#retry-budgets). +### Gemini + +Use a gateway that accepts native Gemini requests, or set +`PROXY_UPSTREAM_URL=https://generativelanguage.googleapis.com` for Google's +Developer API. Point your SDK's HTTP base URL at `http://127.0.0.1:8789`, retaining +its API version and model settings. The proxy forwards `x-goog-api-key`, bearer +authentication, and query parameters unchanged; it does not translate protocols. + +For a direct SSE request, set `GEMINI_MODEL` to an available model and keep your +existing `GEMINI_API_KEY`: + +```bash +curl -N --max-time 720 \ + "http://127.0.0.1:8789/v1beta/models/${GEMINI_MODEL}:streamGenerateContent?alt=sse" \ + -H "x-goog-api-key: ${GEMINI_API_KEY}" \ + -H 'Content-Type: application/json' \ + -d '{"contents":[{"parts":[{"text":"Say hello."}]}]}' +``` + +Gemini selects streaming through the method and `alt=sse`, not a `stream` body +field. `generateContent` and `countTokens` remain single-shot JSON requests with +Gemini-shaped HTTP errors. Standard Vertex AI project/location/publisher model +paths are also recognized; keep the Vertex upstream host and bearer credentials. + +`PROXY_GEMINI_BUFFER_MS` controls the independent 600,000 ms buffering window. +Set client timeouts above it and configure bounded retries for HTTP `503`, honoring +`Retry-After`. After a timed switch to live output, keepalives are SSE comments; +failed streams abort the HTTP transport and require client recovery. + ### Other API clients Set your client's base URL so it sends Messages requests to diff --git a/classify.go b/classify.go index 9df8cc7..39f2ee6 100644 --- a/classify.go +++ b/classify.go @@ -333,12 +333,16 @@ func surfaceFor(path string, f failure) (status int, atype string) { return status, atype } -// Codex ignores X-Should-Retry, so exhausting/disabling the proxy retry budget -// must also produce its terminal HTTP status. Keep the original failure intact -// for diagnostics, and preserve Claude's generic error shape with its false header. +// Clients that ignore X-Should-Retry need a terminal status when retry conversion +// is disabled or exhausted. Keep the original failure intact for diagnostics. func retrySurfaceFor(path string, f failure, canRetry bool) (int, string) { - if isResponsesPath(path) && !canRetry { + if (isResponsesPath(path) || isGeminiPath(path)) && !canRetry { + wasTransient := f.transient f.transient = false + status := statusFor(f) + if isGeminiPath(path) && (wasTransient || status == http.StatusRequestTimeout || status == http.StatusTooManyRequests || status >= 500) { + f.status = http.StatusBadRequest + } } return surfaceFor(path, f) } diff --git a/config_test.go b/config_test.go index ebfa0f0..a07f1db 100644 --- a/config_test.go +++ b/config_test.go @@ -61,6 +61,7 @@ func TestParseConfigRejectsBadValues(t *testing.T) { _, err := parseConfig(envOf(map[string]string{ "PROXY_UPSTREAM_URL": "https://gw.example.com", "PROXY_KEEPALIVE_MS": "60s", + "PROXY_GEMINI_BUFFER_MS": "-1", "PROXY_UPSTREAM_BYTE_IDLE_MS": "0", "PROXY_SDK_RETRY_CAP": "-1", "PROXY_MAX_REQUEST_BYTES": "0", @@ -72,6 +73,7 @@ func TestParseConfigRejectsBadValues(t *testing.T) { } for _, want := range []string{ `PROXY_KEEPALIVE_MS="60s": not an integer`, + `PROXY_GEMINI_BUFFER_MS=-1: must be >= 0`, `PROXY_UPSTREAM_BYTE_IDLE_MS=0: must be >= 1`, `PROXY_SDK_RETRY_CAP=-1: must be >= 0`, `PROXY_MAX_REQUEST_BYTES=0: must be >= 1`, diff --git a/docker-compose.yml.example b/docker-compose.yml.example index d401f01..b65c056 100644 --- a/docker-compose.yml.example +++ b/docker-compose.yml.example @@ -26,6 +26,7 @@ services: PROXY_WORKFLOW_KEEPALIVE_MS: "${PROXY_WORKFLOW_KEEPALIVE_MS:-10000}" # Shorter window for Workflow content progress; see docs/ARCHITECTURE.md. PROXY_RESPONSES_KEEPALIVE_MS: "${PROXY_RESPONSES_KEEPALIVE_MS:-30000}" # Silent-start window when early commit is enabled. PROXY_RESPONSES_BUFFER_MS: "${PROXY_RESPONSES_BUFFER_MS:-600000}" # Full-buffer window when early commit is disabled. + PROXY_GEMINI_BUFFER_MS: "${PROXY_GEMINI_BUFFER_MS:-600000}" # Native Gemini SSE buffering window; 0 disables the timed live switch. PROXY_RESPONSES_EARLY_COMMIT: "${PROXY_RESPONSES_EARLY_COMMIT:-1}" # 0 buffers Responses; client idle timeout must exceed the buffer window. PROXY_RESPONSES_REPLAY_DELTA_BYTES: "${PROXY_RESPONSES_REPLAY_DELTA_BYTES:-65536}" # Combine compatible text deltas during replay; 0 keeps raw replay. PROXY_RESP_HEADER_TIMEOUT_MS: "${PROXY_RESP_HEADER_TIMEOUT_MS:-660000}" # Overrides the binary default of 60000 ms for slow gateway headers. diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 3741ea2..e68fb86 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -7,13 +7,14 @@ attempt possible within the configured budgets. ## Client compatibility -The proxy selects stream behavior from the HTTP method, API path, and request's -`stream` field. It does not require a Claude Code or Codex user agent. Streaming -`POST` requests on the Anthropic Messages and OpenAI Responses routes use the same -buffering and retry engine, with a separate parser for each protocol. +The proxy selects stream behavior from the HTTP method, API path, and request. +Messages and Responses use the body's `stream` field; Gemini uses the +`streamGenerateContent` method with `alt=sse`. These streaming `POST` requests +share one buffering and retry engine, with a separate interpreter per protocol. +No specific user agent is required. Other agents, SDKs, and applications can use these routes when they support the -proxy's timeout, retry, and stream behavior. The integration tests exercise both +proxy's timeout, retry, and stream behavior. Integration tests exercise all three protocols using ordinary HTTP requests. The live CLI harness uses Claude Code; it does not establish compatibility with every client. @@ -24,6 +25,7 @@ it does not establish compatibility with every client. | HTTP retries | Retry HTTP `503` with `Retry-After`. Honor `X-Should-Retry: false` where supported. Local proxy retries can recover some failures even if the client does not retry. | | Stream recovery | Handle an incomplete response after headers have been sent. The proxy cannot restart a committed response for the client. | | Responses events | Tolerate the additional `response.proxy_keepalive` event. Completion checks and content handling follow Codex's parser expectations. | +| Gemini events | Use native `GenerateContentResponse` SSE frames and tolerate SSE comment keepalives. The proxy preserves tool arguments, thought signatures, and unknown fields. | | Retry counter | Without `X-Stainless-Retry-Count`, the proxy treats the client retry count as zero; a positive `PROXY_SDK_RETRY_CAP` cannot bound those client retries. | Some compatibility choices apply to every request on a protocol's routes. @@ -35,8 +37,9 @@ system-prompt marker. Other routes, including `/v1/chat/completions`, are forwarded with one upstream attempt per client call. They do not receive stream validation, transactional -buffering, or model fallback. HTTP errors can still be normalized for retry, using -the Anthropic-style error envelope on non-Responses paths. +buffering, or model fallback. HTTP errors can still be normalized for retry. +Known Gemini model methods use Google's error envelope; Responses uses OpenAI's, +and remaining routes use the Anthropic-style envelope. ## Buffering and the point of no return @@ -51,7 +54,7 @@ cannot replace the response with an HTTP error. If a stream fails after that point, the proxy ends it and logs `DROP`; recovery depends on the client's stream retry support and remaining budget. -The two APIs choose that point differently: +The protocols choose that point differently: | Request | Default behavior | Successful completion | | --- | --- | --- | @@ -59,11 +62,30 @@ The two APIs choose that point differently: | Claude Code Workflow agent | Use a 10-second window, then switch once real content is available. | The same Messages completion. | | OpenAI Responses | Switch to live output as soon as an output event arrives; a 30-second window bounds a silent start. | A valid `response.completed`. | | Responses with `PROXY_RESPONSES_EARLY_COMMIT=0` | Buffer for up to 600 seconds, then switch to live output. | A valid `response.completed`. | +| Gemini `streamGenerateContent?alt=sse` | Buffer for up to 600 seconds, then switch to live output. | Clean transport EOF after every requested candidate has a `finishReason`, or a native prompt block. | Buffering trades immediate output for a chance to recover before the client sees anything. Choose full buffering when that protection matters more than seeing tokens arrive live. +Gemini supports native model methods under `/v1`, `/v1beta`, `/v1alpha`, and +`/v1beta1`, including standard Vertex AI +`/projects/{project}/locations/{location}/publishers/google/models/{model}` paths. +Only `streamGenerateContent?alt=sse` receives transactional protection; +`generateContent`, `countTokens`, and non-SSE streaming transports remain +single-shot. Gemini has no separate terminal SSE event, so the proxy waits for +clean EOF, retaining usage-only frames after candidate completion. Missing +candidates, unfinished SSE frames, and read faults fail the attempt. Native +`MAX_TOKENS`, safety stops, and prompt blocks are delivered unchanged, not retried +as transport failures. See Google's [GenerateContent contract](https://ai.google.dev/api/generate-content). +Candidate counts follow [ProtoJSON](https://protobuf.dev/programming-guides/json/): +camelCase and snake_case field names, quoted integers, exponent notation, and zero +fractional parts are accepted without rewriting the request. + +After a Gemini response commits, a failure aborts the downstream HTTP transport: +a normal EOF could otherwise appear successful to an SDK. No committed attempt +is replayed. Whether the client recovers depends on its stream-retry behavior. + ### One window per client request The buffering window starts when the client request reaches the proxy. Waiting @@ -73,10 +95,10 @@ that same window. None of them restart its clock. When the window ends, a pending header or error-body read is cancelled and becomes eligible for a client retry. An active SSE capture switches to live output; Workflow requests also wait for real forwardable content. In live mode, keepalives -help keep the client connection open. Messages clients receive SSE comments; -Responses clients receive a `response.proxy_keepalive` event. This event is -skippable in Codex, whose SSE reader discards comments without resetting its idle -timer; other Responses clients must also tolerate it. +help keep the client connection open. Messages and Gemini clients receive SSE +comments; Responses clients receive a `response.proxy_keepalive` event. This +event is skippable in Codex, whose SSE reader discards comments without resetting +its idle timer; other Responses clients must also tolerate it. Keepalives cannot extend an absolute request deadline or make a silent upstream produce content. The separate upstream idle timeout still applies. See @@ -88,8 +110,9 @@ Before commit, a recoverable failure has two possible paths: 1. If extra local retries are enabled and there is enough time, the proxy waits and sends the request again. The client continues waiting on its original call. -2. Otherwise, the proxy returns HTTP `503` with an `api_error`, - `X-Should-Retry: true`, and `Retry-After`. The client can then retry the call. +2. Otherwise, the proxy returns HTTP `503` with `X-Should-Retry: true` and + `Retry-After`. Messages and Responses use `api_error`; Gemini uses its native + numeric `error.code: 503` and `error.status: "UNAVAILABLE"` envelope. Using one generic error shape matters for Claude Code. Recognizable overload and rate-limit errors can enter special client handlers with smaller retry budgets. @@ -136,6 +159,12 @@ not send the Stainless counter, so any positive proxy cap leaves its retry limit to `request_max_retries` and `stream_max_retries`. Other clients without that header likewise need to enforce their own retry limit. +Gemini clients likewise own their retry budgets when they do not send the +Stainless counter. If conversion is disabled or exhausted, a transient Gemini +failure becomes terminal HTTP `400`, since its SDK may ignore `X-Should-Retry`. +Explicitly non-retryable Gemini `408`, `429`, and `5xx` errors also become `400` +so a gateway's `false` retry verdict cannot trigger an SDK's status-based retry. + The same counter determines the `Retry-After` sent to the client when the upstream provides no delay. It grows exponentially from one second to 30 seconds, plus 0–1 seconds of jitter. Without the counter, the synthesized delay remains @@ -169,7 +198,7 @@ Recognized safeguards errors can also trigger fallback when the upstream sends classification; disable fallback separately with `PROXY_REFUSAL_FALLBACK_MODEL=off` if the original safeguards error should reach the client. -Fallback does not apply to OpenAI Responses or non-streaming requests. It detects +Fallback does not apply to OpenAI Responses, Gemini, or non-streaming requests. It detects the structured refusal and recognized safeguards messages; it does not inspect ordinary assistant text for a refusal. diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index 3367988..a9e5a85 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -59,12 +59,13 @@ includes upstream waits, retry delays, and fallback attempts. See | `PROXY_WORKFLOW_KEEPALIVE_MS` | `10000` (10 s) | Workflow agent window; live output also requires real content. `0` keeps these requests buffered without a timed switch. | | `PROXY_RESPONSES_KEEPALIVE_MS` | `30000` (30 s) | Responses silent-start window when early commit is on. `0` disables this timer; output can still trigger early commit. | | `PROXY_RESPONSES_BUFFER_MS` | `600000` (10 min) | Responses window when early commit is off. `0` disables the timed switch to live output. | +| `PROXY_GEMINI_BUFFER_MS` | `600000` (10 min) | Native Gemini SSE buffering window. `0` disables the timed switch to live output; clean completion, idle limits, and request deadlines still apply. | | `PROXY_RESP_HEADER_TIMEOUT_MS` | `60000` (1 min) | Wait for upstream response headers. `0` disables this timeout, but other deadlines still apply. | | `PROXY_UPSTREAM_BYTE_IDLE_MS` | `600000` (10 min) | Maximum gap without upstream bytes during SSE capture or error-body reads. Must be at least `1`. | | `PROXY_MAX_REQUEST_DURATION_MS` | `1500000` (25 min) | Total request deadline, including local retries. Must be at least `1`. | | `PROXY_DEADLINE_MARGIN_MS` | `25000` (25 s) | Time reserved before a client's `X-Stainless-Timeout` deadline. `0` removes the margin. | -On streaming Messages and Responses requests, `X-Stainless-Timeout` is read in +On supported streaming requests, `X-Stainless-Timeout`, when present, is read in **seconds**. The effective duration is the smaller of the configured request ceiling and the client timeout minus the margin. For short client timeouts, the margin is limited to half their budget. Other routes use the configured ceiling. @@ -93,6 +94,12 @@ early-commit mode, or above `PROXY_RESPONSES_BUFFER_MS` in full-buffer mode. The README's `660000` covers either default. A zero window does not disable the client's timeouts or the proxy's request deadline. +For Gemini, keep SDK response-header, idle, and total timeouts above +`PROXY_GEMINI_BUFFER_MS`. A timeout of 660,000 ms or more covers the default +buffering window. Keepalives after the live switch are standard SSE comments; +clients that ignore comments for idle accounting must allow their actual content +gaps too. Non-SSE `streamGenerateContent` transports are forwarded without buffering. + Claude Code Workflow agents have an additional content stall timer. Keep their proxy window small and below the script's `stallMs`; see [Workflow agents](ARCHITECTURE.md#workflow-agents). @@ -107,9 +114,10 @@ small and below the script's `stallMs`; see [Workflow agents](ARCHITECTURE.md#wo Local retries need enough time left in the current request and buffering window. The SDK cap uses the Stainless counter; it does not raise client retry limits. -Clients without that header, including Codex, must enforce their own retry limits -when the proxy cap is positive. With a zero cap, Responses failures are returned -as terminal HTTP `400` because Codex does not honor `X-Should-Retry: false` alone. +Clients without that header, including Codex and Gemini SDKs, must enforce their +own retry limits when the proxy cap is positive. With a zero cap, Responses and +transient Gemini failures are returned as terminal HTTP `400` because those +clients may not honor `X-Should-Retry: false` alone. Model fallback has a separate budget and remains enabled when the SDK cap is zero. See [retry budgets](ARCHITECTURE.md#retry-budgets) and @@ -129,8 +137,8 @@ bounds grouped payloads and reduces bursts of tiny notifications when replaying buffered response. Live deltas are unchanged. Incompatible, oversized, malformed, or ambiguous frames and frames with non-empty logprobs remain raw. -Responses completion/error detection and Anthropic structured refusal detection -remain active even when `PROXY_VALIDATE_JSON=0`. +Responses completion/error detection, Gemini JSON/completion/error checks, and +Anthropic structured refusal detection remain active with `PROXY_VALIDATE_JSON=0`. ## Storage and diagnostics @@ -161,7 +169,9 @@ classification: `x-gateway-retryable` takes precedence over `x-should-retry`, which takes precedence over status and body heuristics. A terminal Responses error is rendered as HTTP -`400`; other routes retain the classified non-retryable status. +`400`. Gemini also normalizes non-retryable `408`, `429`, and `5xx` errors to `400` +because its SDKs may ignore `X-Should-Retry: false`. Other non-retryable statuses +are retained. These headers control retry classification. Recognized Anthropic safeguards errors can still trigger a separate [model fallback](ARCHITECTURE.md#model-fallback), diff --git a/docs/DEVELOPMENT.md b/docs/DEVELOPMENT.md index 81659db..305cafe 100644 --- a/docs/DEVELOPMENT.md +++ b/docs/DEVELOPMENT.md @@ -15,8 +15,8 @@ PROXY_UPSTREAM_URL=https://your-gateway.example.com ./steady-proxy The binary listens on `127.0.0.1:8789` by default. It reads environment variables directly and does not load `.env`. The README includes setup for -[Claude Code](../README.md#claude-code), [Codex](../README.md#codex), and -[other API clients](../README.md#other-api-clients). +[Claude Code](../README.md#claude-code), [Codex](../README.md#codex), +[Gemini](../README.md#gemini), and [other API clients](../README.md#other-api-clients). `make build` stamps the binary with `VERSION`, the Git commit, build date, and working-tree status. To build a versioned Docker image with explicit build @@ -87,6 +87,8 @@ Read the [architecture guide](ARCHITECTURE.md) for the request lifecycle and | [`main.go`](../main.go) | Configuration, routing, deadlines, local retries, fallback, and access logs. | | [`wire.go`](../wire.go) | Shared protocol interface and Anthropic stream behavior. | | [`responses.go`](../responses.go) | OpenAI Responses completion, errors, keepalives, and replay. | +| [`gemini.go`](../gemini.go) | Native Gemini candidate completion, usage, and unchanged SSE replay. | +| [`gemini_errors.go`](../gemini_errors.go) | Google numeric/canonical error classification and rendering. | | [`sse.go`](../sse.go) | SSE capture, buffering, validation, tool JSON, and the transition to live output. | | [`classify.go`](../classify.go) | Failure classification and client retry responses. | | [`transport.go`](../transport.go) | Upstream deadline and byte-idle handling. | diff --git a/docs/OPERATIONS.md b/docs/OPERATIONS.md index da577bc..7069326 100644 --- a/docs/OPERATIONS.md +++ b/docs/OPERATIONS.md @@ -133,7 +133,8 @@ The proxy has no authentication of its own. Keep its listening port on loopback; Compose's published port already is. Anyone with access to the port can relay requests with credentials they supply and read `/__version`. -Client authorization and Anthropic headers are forwarded to `PROXY_UPSTREAM_URL`. +Client authorization and provider-specific headers, including `x-goog-api-key`, +are forwarded to `PROXY_UPSTREAM_URL`, along with query authentication such as `key`. The proxy does not follow redirects. Its HTTP transport honors `HTTP_PROXY`/`HTTPS_PROXY` and `NO_PROXY` in the proxy process environment, so a configured forward proxy is also part of the connection path. @@ -145,8 +146,8 @@ off persistent disk. Payload archives are different: they are persistent and **unencrypted**. Known credential headers and query parameters are redacted, including `Authorization`, -`x-api-key`, `Proxy-Authorization`, `api_key`, `api_token`, `access_token`, and -`client_secret`. Redacted names are listed in `redactions`. Other headers and +`x-api-key`, `x-goog-api-key`, `Proxy-Authorization`, `key`, `api_key`, `api_token`, +`access_token`, and `client_secret`. Redacted names are listed in `redactions`. Other headers and parameters remain, and **bodies are not inspected or redacted**. Base64 is an encoding, not encryption. diff --git a/docs/TROUBLESHOOTING.md b/docs/TROUBLESHOOTING.md index 6cf945c..5835296 100644 --- a/docs/TROUBLESHOOTING.md +++ b/docs/TROUBLESHOOTING.md @@ -48,6 +48,14 @@ as success. Ordinary content frames are handled according to Codex's lenient parser rather than Anthropic's validation rules. A `response.incomplete` event also triggers a retry; it is not a successful completion. +Gemini `streamGenerateContent?alt=sse` requires a clean transport EOF after every +requested candidate finishes, or a prompt block. A `finishReason` alone does not +end capture: a later usage frame or error still matters. Missing candidates, +missing finish reasons, partial SSE frames, malformed JSON, and read faults cause +a retry before commit. After commit, the proxy aborts HTTP instead of sending a +clean EOF that could look successful. Native `MAX_TOKENS` and safety stops are +delivered unchanged; they are not stream failures. + If the log shows `DROP`, the client had already received headers. Default Responses mode reaches that point as soon as output appears. Enable full buffering with `PROXY_RESPONSES_EARLY_COMMIT=0` to protect more of the response, and set the @@ -68,6 +76,7 @@ shapes from triggering Claude Code's smaller special-case retry budgets. | `429`, `rate_limit_error`, or “Rate limited” | Retry with the upstream delay when available. | | An Anthropic SSE `error` reporting overload or rate limit | Retry before commit; end the stream after commit. | | A Responses `error` or `response.failed` reporting overload, timeout, or rate limit | Retry before commit; end the stream after commit. | +| A Gemini SSE `error` envelope reporting `UNAVAILABLE`, `RESOURCE_EXHAUSTED`, or a server failure | Retry before commit with Gemini-shaped HTTP `503`; abort HTTP after commit. | | Invalid API key, disabled organization, billing block, or unknown `4xx` | Usually retry, since gateway access or capacity may recover. | A Responses error can carry `type: "invalid_request_error"` while its more @@ -117,6 +126,7 @@ are returned with `X-Should-Retry: false`; terminal Responses errors use HTTP | “Prompt is too long” or context-length error | Compact or shorten the input. | | Explicit `model_not_found` code | Choose a model the gateway supports. | | Responses `unknown_parameter`, `invalid_type`, or another recognized validation code | Correct the named field or value. | +| Gemini `INVALID_ARGUMENT`, `NOT_FOUND`, or a bare numeric `400`/`404` | Correct the request, model, or referenced resource. The documented `API_KEY_INVALID` ErrorInfo is an authentication exception and remains retryable. | Anthropic HTTP errors use recognized message signatures. A bare HTTP `invalid_request_error` without one can still retry; an Anthropic in-stream @@ -126,6 +136,15 @@ failures remain terminal. Responses classification uses error codes, types, and recognized messages, with known transient codes taking precedence over generic invalid-request types. Gateway retry headers override HTTP heuristics. +Gemini decodes numeric `error.code`, canonical `error.status`, and `error.message` +on both HTTP and SSE errors, including the array-wrapped HTTP errors returned by +non-SSE `streamGenerateContent`. `FAILED_PRECONDITION` is not automatically permanent: +it can indicate billing or regional access. The documented Google +`google.rpc.ErrorInfo` reason `API_KEY_INVALID` remains retryable even when its +status is `INVALID_ARGUMENT`, preserving the proxy's authentication policy. +Gateway retry headers still take precedence on HTTP errors. See Google's +[GenerateContent error reference](https://ai.google.dev/gemini-api/docs/generate-content/api-errors). + ### Codex context compaction A Responses `response.failed` event whose error code is exactly diff --git a/gemini.go b/gemini.go new file mode 100644 index 0000000..119abeb --- /dev/null +++ b/gemini.go @@ -0,0 +1,248 @@ +package main + +import ( + "encoding/json" + "errors" + "io" + "math" + "strconv" + "strings" +) + +// Gemini has no separate terminal SSE event. Every requested candidate must +// finish, but a finishReason can precede a final usage-only chunk. The engine +// therefore waits for clean EOF before replaying a successful buffered response. +// Candidate stops (including token limits and safety) and prompt blocks remain +// native responses; they are not transport failures or Anthropic refusals. +type geminiWire struct { + candidateCount int + candidates map[int]geminiCandidateState + chunk uint64 + finished int + blocked bool + failed bool + sawContent bool + candidateTokens int + thoughtTokens int +} + +type geminiCandidateState struct { + chunk uint64 + finished bool +} + +func newGeminiWire(body []byte) *geminiWire { + g := &geminiWire{candidateCount: 1} + config := geminiRequestConfig{count: &g.candidateCount} + request := struct { + Config geminiRequestConfig `json:"generationConfig"` + ProtoConfig geminiRequestConfig `json:"generation_config"` + }{config, config} + if json.Unmarshal(body, &request) != nil || g.candidateCount < 1 { + g.candidateCount = 1 + } + return g +} + +// Both ProtoJSON spellings update one count in JSON field order. A later config +// replaces the earlier config, including resetting an omitted count to default. +type geminiRequestConfig struct{ count *int } + +func (c geminiRequestConfig) UnmarshalJSON(data []byte) error { + *c.count = 1 + count := geminiRequestCount{value: c.count} + fields := struct { + Count geminiRequestCount `json:"candidateCount"` + ProtoCount geminiRequestCount `json:"candidate_count"` + }{count, count} + return json.Unmarshal(data, &fields) +} + +type geminiRequestCount struct{ value *int } + +func (c geminiRequestCount) UnmarshalJSON(data []byte) error { + *c.value = 0 + text := string(data) + if text == "null" { + return nil + } + if len(data) > 0 && data[0] == '"' { + if err := json.Unmarshal(data, &text); err != nil { + return err + } + } + // ProtoJSON int32 fields accept quoted numbers, exponent notation, and + // zero fractional parts. int32's range is exactly representable in float64. + number, err := strconv.ParseFloat(text, 64) + if err != nil || math.IsNaN(number) || number < math.MinInt32 || number > math.MaxInt32 || number != math.Trunc(number) { + return errors.New("invalid Gemini candidate count") + } + *c.value = int(number) + return nil +} + +// Logging is updated by accept's single decode, not by a second scrape decode. +func (g *geminiWire) scrape(ev event, st *captureStats) {} + +func (g *geminiWire) buffered(ev event, st *captureStats) (bool, *failure) { + f := g.accept(ev, st) + g.failed = g.failed || f != nil + return false, f +} + +func (g *geminiWire) live(ev event, raw []byte, st *captureStats) ([][]byte, bool, *failure) { + if f := g.accept(ev, st); f != nil { + g.failed = true + return nil, false, f + } + return [][]byte{raw}, false, nil +} + +func (g *geminiWire) completeOnEOF() bool { + return !g.failed && (g.blocked || g.finished == g.candidateCount) +} + +func (g *geminiWire) replayBuffered(sp *spool, w io.Writer, st *captureStats) error { + return sp.replayRaw(w) +} + +func (g *geminiWire) replayPrefix(sp *spool, w io.Writer) error { + return sp.replayRaw(w) +} + +func (g *geminiWire) forwardable() bool { return g.sawContent } + +func (g *geminiWire) keepalive() []byte { return []byte(": keepalive\n\n") } + +// Only the fields needed for completion and logging are decoded. In particular, +// tool arguments, thought signatures, and all unknown metadata stay opaque and +// are always forwarded/replayed as the original SSE bytes. +type geminiChunk struct { + Candidates []*geminiCandidate `json:"candidates"` + PromptFeedback *struct { + BlockReason string `json:"blockReason"` + } `json:"promptFeedback"` + UsageMetadata *struct { + PromptTokenCount int `json:"promptTokenCount"` + CachedContentTokenCount int `json:"cachedContentTokenCount"` + CandidatesTokenCount *int `json:"candidatesTokenCount"` + ThoughtsTokenCount *int `json:"thoughtsTokenCount"` + } `json:"usageMetadata"` + ModelVersion string `json:"modelVersion"` + Error json.RawMessage `json:"error"` +} + +type geminiCandidate struct { + Index geminiCandidateIndex `json:"index"` + FinishReason string `json:"finishReason"` + Content *struct { + Parts []geminiPart `json:"parts"` + } `json:"content"` +} + +// Proto JSON may omit the default index 0. An explicit null, non-integer, or +// overflowing index is malformed, not another spelling of that default. +type geminiCandidateIndex int + +func (i *geminiCandidateIndex) UnmarshalJSON(data []byte) error { + n, err := strconv.Atoi(string(data)) + if err != nil { + return err + } + *i = geminiCandidateIndex(n) + return nil +} + +// Validate a part's object shape without materializing/copying its text, inline +// media, function arguments, or opaque thought signature. encoding/json already +// validates the complete JSON syntax before invoking this hook. +type geminiPart struct{} + +func (*geminiPart) UnmarshalJSON(data []byte) error { + if len(data) == 0 || data[0] != '{' { + return errors.New("Gemini content part is not an object") + } + return nil +} + +func (g *geminiWire) accept(ev event, st *captureStats) *failure { + data := strings.TrimSpace(ev.data) + if data == "" || data[0] != '{' { + return malformedSSE("Gemini stream event is not a JSON object") + } + var chunk geminiChunk + if err := json.Unmarshal([]byte(data), &chunk); err != nil { + return malformedSSE("invalid JSON or field shape in Gemini stream event") + } + // Both data-only Google envelopes and explicitly named error events must be + // caught even with PROXY_VALIDATE_JSON=0, including after a finishReason. + if ev.name == "error" || len(chunk.Error) != 0 { + if f := classifyGeminiSSEError(data); f != nil { + return f + } + return malformedSSE("invalid error envelope in Gemini stream event") + } + + g.chunk++ + if len(chunk.Candidates) > 0 && g.candidates == nil { + g.candidates = make(map[int]geminiCandidateState) + } + for _, candidate := range chunk.Candidates { + if candidate == nil { + return malformedSSE("Gemini stream candidate is not an object") + } + index := int(candidate.Index) + if index < 0 || index >= g.candidateCount { + return malformedSSE("Gemini stream candidate index is outside the requested range") + } + state := g.candidates[index] + if state.chunk == g.chunk { + return malformedSSE("duplicate Gemini candidate index in stream event") + } + state.chunk = g.chunk + if candidate.Content != nil && len(candidate.Content.Parts) > 0 { + g.sawContent = true + } + if reason := candidate.FinishReason; reason != "" && reason != "FINISH_REASON_UNSPECIFIED" { + if !state.finished { + state.finished = true + g.finished++ + } + if st != nil { + st.stop = reason + } + } + g.candidates[index] = state + } + if feedback := chunk.PromptFeedback; feedback != nil { + if reason := feedback.BlockReason; reason != "" && reason != "BLOCK_REASON_UNSPECIFIED" { + g.blocked = true + if st != nil { + st.stop = reason + } + } + } + if st != nil && chunk.ModelVersion != "" { + st.model = chunk.ModelVersion + } + if usage := chunk.UsageMetadata; usage != nil { + if st != nil && usage.PromptTokenCount > 0 { + // promptTokenCount already includes cachedContentTokenCount: split + // that total, never add the cached portion to it a second time. + cached := min(max(usage.CachedContentTokenCount, 0), usage.PromptTokenCount) + st.setInputUsage(usage.PromptTokenCount-cached, 0, cached) + } + // Usage values are cumulative snapshots, not per-chunk deltas. Keep + // independently supplied output components without summing repeats. + if usage.CandidatesTokenCount != nil && *usage.CandidatesTokenCount >= 0 { + g.candidateTokens = *usage.CandidatesTokenCount + } + if usage.ThoughtsTokenCount != nil && *usage.ThoughtsTokenCount >= 0 { + g.thoughtTokens = *usage.ThoughtsTokenCount + } + if st != nil { + st.outTok = g.candidateTokens + g.thoughtTokens + } + } + return nil +} diff --git a/gemini_errors.go b/gemini_errors.go new file mode 100644 index 0000000..2c59d2b --- /dev/null +++ b/gemini_errors.go @@ -0,0 +1,181 @@ +package main + +import ( + "bytes" + "encoding/json" + "net/http" + "strings" +) + +// GenerateContent uses the Google error envelope, not the Interactions API's +// string-valued error.code: code is an HTTP status and status is a canonical name. +type geminiAPIError struct { + Code int `json:"code"` + Status string `json:"status"` + Message string `json:"message"` + Details []struct { + Type string `json:"@type"` + Reason string `json:"reason"` + } `json:"details,omitempty"` +} + +func classifyGeminiHTTPErrorBytes(resp *http.Response, b []byte) failure { + var envelopes [1]struct { + Error geminiAPIError `json:"error"` + } + data := bytes.TrimSpace(b) + if len(data) > 0 && data[0] == '[' { + // The non-SSE streaming transport wraps its HTTP error in an array. + _ = json.Unmarshal(data, &envelopes) + } else { + _ = json.Unmarshal(data, &envelopes[0]) + } + return classifyGeminiErrorBytes(resp, b, envelopes[0].Error) +} + +func classifyGeminiErrorBytes(resp *http.Response, b []byte, e geminiAPIError) failure { + nativeStatus := strings.ToUpper(strings.TrimSpace(e.Status)) + f := classifyHTTPErrorBytesShaped(resp, b, func(_, _, _, _ string) bool { + // Google's documented invalid-key response is also INVALID_ARGUMENT. + // Its machine-readable ErrorInfo keeps authentication on the retry path. + for _, detail := range e.Details { + if detail.Type == "type.googleapis.com/google.rpc.ErrorInfo" && detail.Reason == "API_KEY_INVALID" { + return false + } + } + if nativeStatus == "INVALID_ARGUMENT" || nativeStatus == "NOT_FOUND" { + return true + } + // FAILED_PRECONDITION can mean disabled billing or an unsupported region. + // A named status owns the decision; bare 400/404 is only a fallback. + if nativeStatus == "" { + status := e.Code + if status == 0 { + status = resp.StatusCode + } + if status == http.StatusBadRequest || status == http.StatusNotFound { + return true + } + } + return anyContains(strings.ToLower(e.Message), requestShapeSigs) + }) + // The shared classifier expects a string error.code. Read Google's fields + // ourselves without translating or re-encoding the archived upstream body. + f.message = e.Message + if nativeStatus != "" { + f.atype = nativeStatus + } else if !f.transient { + f.atype = geminiStatusForHTTP(f.status) + } + return f +} + +// classifyGeminiSSEError is called only after the wire detects an error frame; +// a broken error payload is still a failure, never successful stream content. +func classifyGeminiSSEError(data string) *failure { + var envelope struct { + Error geminiAPIError `json:"error"` + } + b := []byte(data) + if json.Unmarshal(b, &envelope) != nil { + return &failure{transient: true, status: http.StatusBadGateway, atype: "api_error", code: "gemini_sse_error", message: "invalid Gemini error envelope"} + } + e := envelope.Error + status := e.Code + if status < 400 || status > 599 { + status = geminiHTTPStatus(strings.ToUpper(strings.TrimSpace(e.Status))) + if status == 0 { + status = http.StatusBadGateway + } + } + f := classifyGeminiErrorBytes(&http.Response{StatusCode: status}, b, e) + f.code = "gemini_sse_" + f.code + return &f +} + +func geminiHTTPStatus(status string) int { + switch status { + case "INVALID_ARGUMENT", "FAILED_PRECONDITION", "OUT_OF_RANGE": + return http.StatusBadRequest + case "UNAUTHENTICATED": + return http.StatusUnauthorized + case "PERMISSION_DENIED": + return http.StatusForbidden + case "NOT_FOUND": + return http.StatusNotFound + case "ALREADY_EXISTS", "ABORTED": + return http.StatusConflict + case "RESOURCE_EXHAUSTED": + return http.StatusTooManyRequests + case "CANCELLED": + return 499 + case "UNKNOWN", "INTERNAL", "DATA_LOSS": + return http.StatusInternalServerError + case "UNIMPLEMENTED": + return http.StatusNotImplemented + case "UNAVAILABLE": + return http.StatusServiceUnavailable + case "DEADLINE_EXCEEDED": + return http.StatusGatewayTimeout + default: + return 0 + } +} + +func geminiStatusForHTTP(status int) string { + switch status { + case http.StatusBadRequest, http.StatusRequestEntityTooLarge, http.StatusUnprocessableEntity: + return "INVALID_ARGUMENT" + case http.StatusUnauthorized: + return "UNAUTHENTICATED" + case http.StatusForbidden: + return "PERMISSION_DENIED" + case http.StatusNotFound: + return "NOT_FOUND" + case http.StatusRequestTimeout, http.StatusGatewayTimeout: + return "DEADLINE_EXCEEDED" + case http.StatusConflict: + return "ABORTED" + case http.StatusPreconditionFailed: + return "FAILED_PRECONDITION" + case http.StatusRequestedRangeNotSatisfiable: + return "OUT_OF_RANGE" + case http.StatusTooManyRequests: + return "RESOURCE_EXHAUSTED" + case 499: + return "CANCELLED" + case http.StatusInternalServerError: + return "INTERNAL" + case http.StatusNotImplemented: + return "UNIMPLEMENTED" + case http.StatusServiceUnavailable: + return "UNAVAILABLE" + default: + return "UNKNOWN" + } +} + +func writeGeminiError(w http.ResponseWriter, canRetry bool, status int, atype, message string, retryAfter int, code string) { + h := w.Header() + h.Set("Content-Type", "application/json") + if canRetry { + h.Set("X-Should-Retry", "true") + } else { + h.Set("X-Should-Retry", "false") + } + if retryAfter > 0 { + h.Set("Retry-After", itoa(retryAfter)) + } + h.Set("X-Steady-Proxy-Reason", code) + if status <= 0 { + status = http.StatusBadGateway + } + nativeStatus := strings.ToUpper(strings.TrimSpace(atype)) + if geminiHTTPStatus(nativeStatus) != status { + nativeStatus = geminiStatusForHTTP(status) + } + w.WriteHeader(status) + _ = json.NewEncoder(w).Encode(struct { + Error geminiAPIError `json:"error"` + }{Error: geminiAPIError{Code: status, Status: nativeStatus, Message: message}}) +} diff --git a/gemini_errors_test.go b/gemini_errors_test.go new file mode 100644 index 0000000..c1dfedc --- /dev/null +++ b/gemini_errors_test.go @@ -0,0 +1,308 @@ +package main + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" +) + +func TestClassifyGeminiHTTPError(t *testing.T) { + cases := []struct { + name string + http int + body string + transient bool + status int + }{ + {"invalid-argument", 400, `{"error":{"code":400,"status":"INVALID_ARGUMENT","message":"invalid contents"}}`, false, 400}, + {"not-found", 404, `{"error":{"code":404,"status":"NOT_FOUND","message":"model does not exist"}}`, false, 404}, + {"normalized-native-status", 400, `{"error":{"code":400,"status":" invalid_argument ","message":"invalid contents"}}`, false, 400}, + {"numeric-400", 400, `{"error":{"code":400,"message":"invalid contents"}}`, false, 400}, + {"numeric-404", 404, `{"error":{"code":404,"message":"model does not exist"}}`, false, 404}, + {"http-400-fallback", 400, `{"error":{"message":"bad request"}}`, false, 400}, + {"http-404-fallback", 404, `{"error":{"message":"missing resource"}}`, false, 404}, + {"billing-precondition", 400, `{"error":{"code":400,"status":"FAILED_PRECONDITION","message":"enable billing on your project"}}`, true, 502}, + {"region-precondition", 400, `{"error":{"code":400,"status":"FAILED_PRECONDITION","message":"unsupported region for this account"}}`, true, 502}, + {"auth-precondition", 400, `{"error":{"code":400,"status":"FAILED_PRECONDITION","message":"authentication is required"}}`, true, 502}, + {"unauthenticated", 401, `{"error":{"code":401,"status":"UNAUTHENTICATED","message":"invalid API key"}}`, true, 502}, + {"permission-denied", 403, `{"error":{"code":403,"status":"PERMISSION_DENIED","message":"permission denied"}}`, true, 502}, + {"invalid-key-invalid-argument", 400, `{"error":{"code":400,"status":"INVALID_ARGUMENT","message":"API key not valid. Please pass a valid API key.","details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","reason":"API_KEY_INVALID","domain":"googleapis.com"}]}}`, true, 502}, + {"quota", 429, `{"error":{"code":429,"status":"RESOURCE_EXHAUSTED","message":"quota exhausted"}}`, true, 429}, + {"internal", 500, `{"error":{"code":500,"status":"INTERNAL","message":"internal error"}}`, true, 500}, + {"unavailable", 503, `{"error":{"code":503,"status":"UNAVAILABLE","message":"busy"}}`, true, 503}, + {"unknown-named-status", 400, `{"error":{"code":400,"status":"UNKNOWN","message":"temporary gateway fault"}}`, true, 502}, + {"new-named-status", 404, `{"error":{"code":404,"status":"NEW_PROVIDER_ERROR","message":"temporary gateway fault"}}`, true, 502}, + {"unnamed-unknown", 418, `{"error":{"code":418,"message":"unknown provider fault"}}`, true, 502}, + {"named-status-beats-numeric-400", 400, `{"error":{"code":400,"status":"UNAVAILABLE","message":"temporary capacity issue"}}`, true, 502}, + {"native-request-shape-on-403", 403, `{"error":{"code":403,"status":"INVALID_ARGUMENT","message":"bad request"}}`, false, 403}, + {"numeric-envelope-owns-fallback", 403, `{"error":{"code":400,"message":"bad request"}}`, false, 403}, + {"specific-request-signature", 400, `{"error":{"code":400,"status":"FAILED_PRECONDITION","message":"prompt is too long for the context window"}}`, false, 400}, + {"specific-signature-on-unknown", 403, `{"error":{"code":403,"status":"UNKNOWN","message":"context length exceeded"}}`, false, 403}, + {"metadata-is-not-request-message", 403, `{"error":{"code":403,"status":"PERMISSION_DENIED","message":"access denied","details":[{"description":"context length"}]}}`, true, 502}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + f := classifyGeminiHTTPErrorBytes(&http.Response{StatusCode: tc.http}, []byte(tc.body)) + if f.transient != tc.transient || f.status != tc.status { + t.Fatalf("want transient=%v status=%d, got %+v", tc.transient, tc.status, f) + } + var envelope struct { + Error struct{ Message string } `json:"error"` + } + if err := json.Unmarshal([]byte(tc.body), &envelope); err != nil { + t.Fatal(err) + } + if f.message != envelope.Error.Message { + t.Fatalf("numeric error.code must not lose message: got %q, want %q", f.message, envelope.Error.Message) + } + if !f.transient && f.code != "request_shape" { + t.Fatalf("want request_shape reason, got %+v", f) + } + }) + } +} + +func TestGeminiArrayHTTPErrorRetryPolicy(t *testing.T) { + for _, tc := range []struct { + name string + status int + error string + transient bool + }{ + {"invalid-key", 400, `{"code":400,"status":"INVALID_ARGUMENT","details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","reason":"API_KEY_INVALID"}]`, true}, + {"invalid-request", 400, `{"code":400,"status":"INVALID_ARGUMENT"`, false}, + {"billing", 400, `{"code":400,"status":"FAILED_PRECONDITION"`, true}, + {"quota", 429, `{"code":429,"status":"RESOURCE_EXHAUSTED"`, true}, + } { + t.Run(tc.name, func(t *testing.T) { + body := []byte(" \n[{\"error\":" + tc.error + `,"message":"native detail"}}]`) + f := classifyGeminiHTTPErrorBytes(&http.Response{StatusCode: tc.status}, body) + if f.transient != tc.transient || f.message != "native detail" { + t.Fatalf("array-wrapped native error lost retry policy or context: %+v", f) + } + }) + } +} + +func TestGeminiHTTPStatusPrecedence(t *testing.T) { + body := []byte(`{"error":{"code":400,"status":"INVALID_ARGUMENT","message":"context length exceeded"}}`) + for _, status := range []int{408, 409, 429, 500, 502, 503, 504} { + t.Run(itoa(status), func(t *testing.T) { + f := classifyGeminiHTTPErrorBytes(&http.Response{StatusCode: status}, body) + if !f.transient || f.status != status { + t.Fatalf("shared status policy must precede body heuristics: %+v", f) + } + wantFastRetry := status == 408 || status == 409 || status == 502 || status == 503 + if f.fastRetry != wantFastRetry { + t.Fatalf("fastRetry=%v, want %v (%+v)", f.fastRetry, wantFastRetry, f) + } + }) + } +} + +func TestGeminiHTTPRetryHeaderPrecedence(t *testing.T) { + cases := []struct { + name string + status int + gateway string + upstream string + transient bool + want int + code string + delay int + }{ + {"gateway-force-retry", 400, "true", "", true, 502, "gateway_retryable_400", 7}, + {"gateway-force-permanent", 503, "false", "", false, 503, "gateway_permanent", 7}, + {"gateway-true-beats-upstream-false", 400, "TrUe", "false", true, 502, "gateway_retryable_400", 7}, + {"gateway-false-beats-upstream-true", 503, "FaLsE", "true", false, 503, "gateway_permanent", 7}, + {"upstream-force-retry", 404, "", "true", true, 502, "upstream_retry_404", 7}, + {"upstream-force-permanent", 503, "", "false", false, 503, "upstream_no_retry", 0}, + {"unknown-gateway-defers-to-upstream", 400, "unknown", "true", true, 502, "upstream_retry_400", 7}, + } + body := []byte(`{"error":{"code":400,"status":"INVALID_ARGUMENT","message":"invalid contents"}}`) + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + resp := &http.Response{StatusCode: tc.status, Header: make(http.Header)} + resp.Header.Set("X-Gateway-Retryable", tc.gateway) + resp.Header.Set("X-Should-Retry", tc.upstream) + resp.Header.Set("Retry-After", "7") + f := classifyGeminiHTTPErrorBytes(resp, body) + if f.transient != tc.transient || f.status != tc.want || f.code != tc.code || f.retryAfter != tc.delay { + t.Fatalf("header precedence mismatch: %+v", f) + } + if f.message != "invalid contents" || f.atype != "INVALID_ARGUMENT" { + t.Fatalf("native fields lost through shared classifier: %+v", f) + } + }) + } +} + +func TestGeminiHTTPRetryAfter(t *testing.T) { + for _, tc := range []struct { + name string + status int + body string + header string + want int + }{ + {"quota-backoff", 429, `{"error":{"code":429,"status":"RESOURCE_EXHAUSTED","message":"quota exhausted"}}`, "11", 11}, + {"billing-backoff", 400, `{"error":{"code":400,"status":"FAILED_PRECONDITION","message":"billing unavailable"}}`, "9", 9}, + {"server-backoff", 503, `{"error":{"code":503,"status":"UNAVAILABLE","message":"busy"}}`, "3", 3}, + {"absent-backoff", 429, `{"error":{"code":429,"status":"RESOURCE_EXHAUSTED"}}`, "", 0}, + {"malformed-backoff", 503, `{"error":{"code":503,"status":"UNAVAILABLE"}}`, "not a delay", 0}, + {"negative-backoff", 503, `{"error":{"code":503,"status":"UNAVAILABLE"}}`, "-1", 0}, + } { + t.Run(tc.name, func(t *testing.T) { + resp := &http.Response{StatusCode: tc.status, Header: make(http.Header)} + resp.Header.Set("Retry-After", tc.header) + f := classifyGeminiHTTPErrorBytes(resp, []byte(tc.body)) + if !f.transient || f.retryAfter != tc.want { + t.Fatalf("want transient backoff=%d, got %+v", tc.want, f) + } + }) + } +} + +func TestClassifyGeminiSSEError(t *testing.T) { + cases := []struct { + name string + body string + transient bool + status int + }{ + {"invalid-argument", `{"error":{"code":400,"status":"INVALID_ARGUMENT","message":"bad contents"}}`, false, 400}, + {"not-found", `{"error":{"code":404,"status":"NOT_FOUND","message":"bad model"}}`, false, 404}, + {"numeric-400", `{"error":{"code":400}}`, false, 400}, + {"numeric-404", `{"error":{"code":404}}`, false, 404}, + {"invalid-argument-without-code", `{"error":{"status":"INVALID_ARGUMENT"}}`, false, 400}, + {"not-found-without-code", `{"error":{"status":"NOT_FOUND"}}`, false, 404}, + {"billing-without-code", `{"error":{"status":"FAILED_PRECONDITION","message":"enable billing"}}`, true, 502}, + {"region", `{"error":{"code":400,"status":"FAILED_PRECONDITION","message":"unsupported region"}}`, true, 502}, + {"auth-without-code", `{"error":{"status":"UNAUTHENTICATED"}}`, true, 502}, + {"invalid-key-invalid-argument", `{"error":{"code":400,"status":"INVALID_ARGUMENT","message":"API key not valid. Please pass a valid API key.","details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","reason":"API_KEY_INVALID"}]}}`, true, 502}, + {"permission-without-code", `{"error":{"status":"PERMISSION_DENIED"}}`, true, 502}, + {"quota-without-code", `{"error":{"status":"RESOURCE_EXHAUSTED"}}`, true, 429}, + {"internal-without-code", `{"error":{"status":"INTERNAL"}}`, true, 500}, + {"unknown-canonical-without-code", `{"error":{"status":"UNKNOWN"}}`, true, 500}, + {"unavailable-without-code", `{"error":{"status":"UNAVAILABLE"}}`, true, 503}, + {"deadline-without-code", `{"error":{"status":"DEADLINE_EXCEEDED"}}`, true, 504}, + {"aborted-without-code", `{"error":{"status":"ABORTED"}}`, true, 409}, + {"already-exists-without-code", `{"error":{"status":"ALREADY_EXISTS"}}`, true, 409}, + {"cancelled-without-code", `{"error":{"status":"CANCELLED"}}`, true, 502}, + {"unimplemented-without-code", `{"error":{"status":"UNIMPLEMENTED"}}`, true, 501}, + {"data-loss-without-code", `{"error":{"status":"DATA_LOSS"}}`, true, 500}, + {"out-of-range-is-not-added-request-shape", `{"error":{"status":"OUT_OF_RANGE"}}`, true, 502}, + {"unknown-native-status", `{"error":{"status":"NEW_PROVIDER_ERROR"}}`, true, 502}, + {"normalized-native-status", `{"error":{"status":" unavailable "}}`, true, 503}, + {"numeric-status-precedes-derived-status", `{"error":{"code":503,"status":"INVALID_ARGUMENT"}}`, true, 503}, + {"invalid-numeric-code-uses-native-status", `{"error":{"code":-1,"status":"NOT_FOUND"}}`, false, 404}, + {"specific-signature", `{"error":{"status":"FAILED_PRECONDITION","message":"context length exceeded"}}`, false, 400}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + f := classifyGeminiSSEError(tc.body) + if f == nil || f.transient != tc.transient || f.status != tc.status { + t.Fatalf("want transient=%v status=%d, got %+v", tc.transient, tc.status, f) + } + if f.origStatus != 0 { + t.Fatalf("SSE error has no raw HTTP error status: %+v", f) + } + var envelope struct { + Error struct{ Message string } `json:"error"` + } + if err := json.Unmarshal([]byte(tc.body), &envelope); err != nil { + t.Fatal(err) + } + if f.message != envelope.Error.Message { + t.Fatalf("native SSE message lost: %+v", f) + } + if !f.transient && f.code != "gemini_sse_request_shape" { + t.Fatalf("want Gemini SSE request-shape reason, got %+v", f) + } + }) + } +} + +func TestGeminiMalformedErrors(t *testing.T) { + for _, body := range []string{ + "", "not JSON", `{"error":`, `[]`, `null`, `{}`, `{"error":null}`, + `{"error":"broken"}`, `{"error":{"code":"400"}}`, + `{"error":{"status":400}}`, `{"error":{"message":false}}`, + } { + t.Run(body, func(t *testing.T) { + httpFailure := classifyGeminiHTTPErrorBytes(&http.Response{StatusCode: 503}, []byte(body)) + if !httpFailure.transient || httpFailure.status != 503 { + t.Fatalf("malformed HTTP body must retain server-status policy: %+v", httpFailure) + } + f := classifyGeminiSSEError(body) + if f == nil || !f.transient || f.status != 502 { + t.Fatalf("detected malformed SSE error must remain a failure: %+v", f) + } + }) + } +} + +func TestWriteGeminiError(t *testing.T) { + cases := []struct { + name string + status int + atype string + canRetry bool + retryAfter int + wantStatus int + wantCanonical string + }{ + {"retry-surface", 503, "api_error", true, 7, 503, "UNAVAILABLE"}, + {"retry-hides-rate-limit-identity", 503, "RESOURCE_EXHAUSTED", true, 3, 503, "UNAVAILABLE"}, + {"invalid-request", 400, "invalid_request_error", false, 0, 400, "INVALID_ARGUMENT"}, + {"native-precondition", 400, "FAILED_PRECONDITION", false, 0, 400, "FAILED_PRECONDITION"}, + {"normalized-native-precondition", 400, " failed_precondition ", false, 0, 400, "FAILED_PRECONDITION"}, + {"not-found", 404, "api_error", false, 0, 404, "NOT_FOUND"}, + {"auth", 401, "authentication_error", false, 0, 401, "UNAUTHENTICATED"}, + {"permission", 403, "permission_error", false, 0, 403, "PERMISSION_DENIED"}, + {"quota", 429, "rate_limit_error", false, 0, 429, "RESOURCE_EXHAUSTED"}, + {"internal", 500, "api_error", false, 0, 500, "INTERNAL"}, + {"deadline", 504, "timeout_error", false, 0, 504, "DEADLINE_EXCEEDED"}, + {"native-status-must-agree-with-http", 404, "INVALID_ARGUMENT", false, 0, 404, "NOT_FOUND"}, + {"default-status", 0, "api_error", false, 0, 502, "UNKNOWN"}, + } + message := "steady-proxy: quoted \"failure\"\nwith detail" + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + rec := httptest.NewRecorder() + writeGeminiError(rec, tc.canRetry, tc.status, tc.atype, message, tc.retryAfter, "gemini_test_reason") + if rec.Code != tc.wantStatus || rec.Header().Get("Content-Type") != "application/json" { + t.Fatalf("wrong HTTP surface: status=%d headers=%v", rec.Code, rec.Header()) + } + wantRetry := "false" + if tc.canRetry { + wantRetry = "true" + } + if rec.Header().Get("X-Should-Retry") != wantRetry || rec.Header().Get("X-Steady-Proxy-Reason") != "gemini_test_reason" { + t.Fatalf("retry/reason headers lost: %v", rec.Header()) + } + wantDelay := "" + if tc.retryAfter > 0 { + wantDelay = itoa(tc.retryAfter) + } + if rec.Header().Get("Retry-After") != wantDelay { + t.Fatalf("Retry-After=%q, want %q", rec.Header().Get("Retry-After"), wantDelay) + } + var body map[string]json.RawMessage + if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil { + t.Fatalf("invalid JSON: %v", err) + } + if len(body) != 1 || body["error"] == nil { + t.Fatalf("want only native error envelope, got %s", rec.Body.String()) + } + var native map[string]any + if err := json.Unmarshal(body["error"], &native); err != nil { + t.Fatal(err) + } + if len(native) != 3 || native["code"] != float64(tc.wantStatus) || native["status"] != tc.wantCanonical || native["message"] != message { + t.Fatalf("want numeric code/string status/message, got %s", rec.Body.String()) + } + }) + } +} diff --git a/gemini_integration_test.go b/gemini_integration_test.go new file mode 100644 index 0000000..4ac627e --- /dev/null +++ b/gemini_integration_test.go @@ -0,0 +1,148 @@ +package main + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" +) + +const geminiStreamPath = "/v1beta/models/gemini-test:streamGenerateContent?alt=sse" + +func TestGeminiHiddenRetryDiscardsFailedOutput(t *testing.T) { + var calls atomic.Int32 + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + if calls.Add(1) == 1 { + io.WriteString(w, "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"discard me\"}]}}]}\n\n") + io.WriteString(w, "data: {\"error\":{\"code\":503,\"status\":\"UNAVAILABLE\",\"message\":\"busy\"}}\n\n") + return + } + io.WriteString(w, "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"recovered\"}]},\"finishReason\":\"STOP\"}]}\n\n") + })) + defer up.Close() + setupForTest(up.URL) + cfg.txLocalRetries = 1 + cfg.localBackoffCap = time.Millisecond + // These Anthropic-only signals must not change Gemini's window or model. + cfg.wfKeepaliveMs = time.Millisecond + body := `{"model":"not-a-gemini-body-field","system":"workflow orchestration script","contents":[{"parts":[{"text":"hi"}]}]}` + rec := httptest.NewRecorder() + handle(rec, httptest.NewRequest(http.MethodPost, geminiStreamPath, strings.NewReader(body))) + if rec.Code != http.StatusOK || calls.Load() != 2 { + t.Fatalf("status=%d calls=%d body=%s", rec.Code, calls.Load(), rec.Body) + } + if strings.Contains(rec.Body.String(), "discard me") || strings.Contains(rec.Body.String(), "busy") || !strings.Contains(rec.Body.String(), "recovered") { + t.Fatalf("failed attempt leaked or recovery missing: %s", rec.Body) + } + if rec.Header().Get("X-Steady-Proxy-Mode") != "buffered" { + t.Fatalf("Anthropic marker changed Gemini buffering: %v", rec.Header()) + } +} + +func TestGeminiInvalidArgumentDoesNotRetry(t *testing.T) { + for _, path := range []string{ + geminiStreamPath, + "/v1/models/gemini-test:generateContent", + "/v1beta1/projects/p/locations/global/publishers/google/models/gemini-test:streamGenerateContent?alt=sse", + } { + t.Run(path, func(t *testing.T) { + var calls atomic.Int32 + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.WriteHeader(http.StatusBadRequest) + io.WriteString(w, `{"error":{"code":400,"status":"INVALID_ARGUMENT","message":"bad tool schema"}}`) + })) + defer up.Close() + setupForTest(up.URL) + cfg.txLocalRetries = 3 + rec := httptest.NewRecorder() + handle(rec, httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"contents":[]}`))) + var result struct { + Error struct { + Code int `json:"code"` + Status string `json:"status"` + } `json:"error"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &result); err != nil { + t.Fatal(err) + } + if rec.Code != 400 || calls.Load() != 1 || rec.Header().Get("X-Should-Retry") != "false" || result.Error.Code != 400 || result.Error.Status != "INVALID_ARGUMENT" { + t.Fatalf("status=%d calls=%d headers=%v error=%+v", rec.Code, calls.Load(), rec.Header(), result.Error) + } + }) + } +} + +func TestGeminiDisabledRetryConversionIsTerminal(t *testing.T) { + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusTooManyRequests) + io.WriteString(w, `{"error":{"code":429,"status":"RESOURCE_EXHAUSTED","message":"quota"}}`) + })) + defer up.Close() + setupForTest(up.URL) + cfg.sdkRetryCap = 0 + rec := httptest.NewRecorder() + handle(rec, httptest.NewRequest(http.MethodPost, geminiStreamPath, strings.NewReader(`{"contents":[]}`))) + if rec.Code != 400 || rec.Header().Get("X-Should-Retry") != "false" || rec.Header().Get("Retry-After") != "" { + t.Fatalf("disabled conversion still retryable: status=%d headers=%v body=%s", rec.Code, rec.Header(), rec.Body) + } +} + +func TestGeminiRetryOverrideIsTerminalToSDKs(t *testing.T) { + for _, status := range []int{408, 429, 503} { + for _, header := range []string{"X-Gateway-Retryable", "X-Should-Retry"} { + for _, path := range []string{geminiStreamPath, "/v1/models/gemini-test:generateContent"} { + t.Run(itoa(status)+"/"+header+"/"+path, func(t *testing.T) { + var calls atomic.Int32 + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set(header, "false") + w.Header().Set("Retry-After", "7") + w.WriteHeader(status) + io.WriteString(w, `{"error":{"message":"terminal upstream failure"}}`) + })) + defer up.Close() + setupForTest(up.URL) + cfg.txLocalRetries = 3 + rec := httptest.NewRecorder() + handle(rec, httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"contents":[]}`))) + if rec.Code != 400 || calls.Load() != 1 || rec.Header().Get("X-Should-Retry") != "false" || rec.Header().Get("Retry-After") != "" { + t.Fatalf("permanent upstream error still retryable: status=%d calls=%d headers=%v", rec.Code, calls.Load(), rec.Header()) + } + }) + } + } + } +} + +func TestGeminiCommittedTruncationAbortsHTTP(t *testing.T) { + var calls atomic.Int32 + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + io.WriteString(w, "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"partial\"}]}}]}\n\n") + w.(http.Flusher).Flush() + <-r.Context().Done() + })) + defer up.Close() + setupForTest(up.URL) + cfg.geminiBufferMs = 5 * time.Millisecond + cfg.upstreamByteIdle = 40 * time.Millisecond + cfg.txLocalRetries = 3 + proxy := httptest.NewServer(http.HandlerFunc(handle)) + defer proxy.Close() + resp, err := http.Post(proxy.URL+geminiStreamPath, "application/json", strings.NewReader(`{"contents":[]}`)) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err == nil || !strings.Contains(string(body), "partial") || calls.Load() != 1 { + t.Fatalf("committed truncation must fail transport without replay: body=%q err=%v calls=%d", body, err, calls.Load()) + } +} diff --git a/gemini_test.go b/gemini_test.go new file mode 100644 index 0000000..a937263 --- /dev/null +++ b/gemini_test.go @@ -0,0 +1,415 @@ +package main + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" + "time" +) + +const goodGeminiStream = "data: {\"candidates\":[{\"index\":0,\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"Hello\"}]}}],\"modelVersion\":\"gemini-test\"}\n\n" + + "data: {\"candidates\":[{\"index\":0,\"finishReason\":\"STOP\"}]}\n\n" + + "data: {\"usageMetadata\":{\"promptTokenCount\":100,\"cachedContentTokenCount\":25,\"candidatesTokenCount\":8,\"thoughtsTokenCount\":3,\"totalTokenCount\":111}}\n\n" + +func geminiTestEvent(data string) event { + return event{data: data, raw: []byte("data: " + data + "\n\n")} +} + +func captureGemini(t *testing.T, body []byte, r io.Reader, st *captureStats, early bool) (*httptest.ResponseRecorder, bool, *failure) { + t.Helper() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + rec := httptest.NewRecorder() + window := time.Hour + wrote, f := captureSSECore(ctx, cancel, rec, http.Header{}, r, st, window, false, early, newGeminiWire(body), bufferDeadline(time.Now(), window)) + return rec, wrote, f +} + +// Each entry is a separate transport read, so tests can exercise the buffered +// prefix/live boundary without timers. The final error is independent of any +// previously observed candidate finishReason. +type geminiTestReader struct { + chunks []string + err error +} + +func (r *geminiTestReader) Read(p []byte) (int, error) { + if len(r.chunks) == 0 { + if r.err != nil { + return 0, r.err + } + return 0, io.EOF + } + n := copy(p, r.chunks[0]) + r.chunks[0] = r.chunks[0][n:] + if r.chunks[0] == "" { + r.chunks = r.chunks[1:] + } + return n, nil +} + +// Readers may return bytes and EOF in the same call. A complete native response +// must replay on that EOF, but an unfinished SSE frame must still fail. +type geminiTestFinalReader struct { + data []byte +} + +func (r *geminiTestFinalReader) Read(p []byte) (int, error) { + n := copy(p, r.data) + r.data = r.data[n:] + if len(r.data) == 0 { + return n, io.EOF + } + return n, nil +} + +func TestGeminiCaptureBufferedWaitsForEOFAndTrailingUsage(t *testing.T) { + for _, tc := range []struct { + name string + r io.Reader + }{ + {"single read", strings.NewReader(goodGeminiStream)}, + {"one byte reads", iotest1byte(goodGeminiStream)}, + {"data and EOF together", &geminiTestFinalReader{data: []byte(goodGeminiStream)}}, + } { + t.Run(tc.name, func(t *testing.T) { + var st captureStats + rec, wrote, f := captureGemini(t, nil, tc.r, &st, false) + if !wrote || f != nil { + t.Fatalf("complete stream: wrote=%v failure=%+v", wrote, f) + } + if rec.Body.String() != goodGeminiStream { + t.Fatalf("response changed or trailing usage lost: %q", rec.Body.String()) + } + if st.mode != "buffered" || rec.Header().Get("X-Steady-Proxy-Mode") != "buffered" { + t.Fatalf("expected buffered replay: mode=%q headers=%v", st.mode, rec.Header()) + } + if st.model != "gemini-test" || st.stop != "STOP" || st.inTok != 100 || st.inputTok != 75 || st.cacheReadTok != 25 || st.outTok != 11 { + t.Fatalf("usage/model/stop mismatch: %+v", st) + } + if st.bytes != int64(len(goodGeminiStream)) { + t.Fatalf("captured bytes=%d, want %d", st.bytes, len(goodGeminiStream)) + } + }) + } +} + +func TestGeminiWireCompletionTracksEveryRequestedCandidate(t *testing.T) { + for _, tc := range []struct { + name string + body string + chunks []string + complete bool + }{ + {"empty stream", "", nil, false}, + {"empty object", "", []string{`{}`}, false}, + {"no finish", "", []string{`{"candidates":[{"index":0,"content":{"parts":[{"text":"partial"}]}}]}`}, false}, + {"unspecified finish", "", []string{`{"candidates":[{"index":0,"finishReason":"FINISH_REASON_UNSPECIFIED"}]}`}, false}, + {"empty finish", "", []string{`{"candidates":[{"index":0,"finishReason":""}]}`}, false}, + {"default candidate count", "", []string{`{"candidates":[{"index":0,"finishReason":"STOP"}]}`}, true}, + {"omitted default index", `{}`, []string{`{"candidates":[{"finishReason":"STOP"}]}`}, true}, + {"explicit count one", `{"generationConfig":{"candidateCount":1}}`, []string{`{"candidates":[{"index":0,"finishReason":"STOP"}]}`}, true}, + {"two requested one finished", `{"generationConfig":{"candidateCount":2}}`, []string{`{"candidates":[{"index":0,"finishReason":"STOP"}]}`}, false}, + {"only second finished", `{"generationConfig":{"candidateCount":2}}`, []string{`{"candidates":[{"index":1,"finishReason":"STOP"}]}`}, false}, + {"repeat finish cannot replace missing candidate", `{"generationConfig":{"candidateCount":2}}`, []string{`{"candidates":[{"index":0,"finishReason":"STOP"}]}`, `{"candidates":[{"index":0,"finishReason":"STOP"}]}`}, false}, + {"second remains unfinished", `{"generationConfig":{"candidateCount":2}}`, []string{`{"candidates":[{"index":0,"finishReason":"STOP"},{"index":1,"finishReason":"FINISH_REASON_UNSPECIFIED"}]}`}, false}, + {"two finish together", `{"generationConfig":{"candidateCount":2}}`, []string{`{"candidates":[{"index":0,"finishReason":"STOP"},{"index":1,"finishReason":"MAX_TOKENS"}]}`}, true}, + {"two finish out of order", `{"generationConfig":{"candidateCount":2}}`, []string{`{"candidates":[{"index":1,"finishReason":"SAFETY"}]}`, `{"candidates":[{"finishReason":"STOP"}]}`}, true}, + {"large request does not infer unseen candidates", `{"generationConfig":{"candidateCount":2147483647}}`, []string{`{"candidates":[{"index":0,"finishReason":"STOP"}]}`}, false}, + } { + t.Run(tc.name, func(t *testing.T) { + wire := newGeminiWire([]byte(tc.body)) + for _, chunk := range tc.chunks { + done, f := wire.buffered(geminiTestEvent(chunk), nil) + if done || f != nil { + t.Fatalf("event must stay buffered until EOF: done=%v failure=%+v", done, f) + } + } + if got := wire.completeOnEOF(); got != tc.complete { + t.Fatalf("completeOnEOF=%v, want %v", got, tc.complete) + } + }) + } +} + +func TestGeminiProtoJSONCandidateCounts(t *testing.T) { + useDefaultConfig(t) + first := string(geminiTestEvent(`{"candidates":[{"index":0,"finishReason":"STOP"}]}`).raw) + second := string(geminiTestEvent(`{"candidates":[{"index":1,"finishReason":"STOP"}]}`).raw) + for _, body := range []string{ + `{"generation_config":{"candidate_count":2}}`, + `{"generation_config":{"candidateCount":2}}`, + `{"generationConfig":{"candidate_count":2}}`, + `{"generationConfig":{"candidateCount":"2"}}`, + `{"generationConfig":{"candidateCount":2.0}}`, + `{"generationConfig":{"candidateCount":2e0}}`, + `{"generationConfig":{"candidateCount":"2e0"}}`, + `{"generationConfig":{"candidateCount":"2.0"}}`, + `{"generationConfig":{"candidateCount":1,"candidate_count":2}}`, + `{"generation_config":{"candidate_count":1},"generationConfig":{"candidateCount":2}}`, + } { + t.Run(body, func(t *testing.T) { + _, wrote, f := captureGemini(t, []byte(body), strings.NewReader(first+second), nil, false) + if !wrote || f != nil { + t.Fatalf("valid two-candidate response rejected: wrote=%v failure=%+v", wrote, f) + } + rec, wrote, f := captureGemini(t, []byte(body), strings.NewReader(first), nil, false) + if wrote || f == nil || !f.transient || rec.Body.Len() != 0 { + t.Fatalf("missing second candidate accepted: wrote=%v failure=%+v body=%q", wrote, f, rec.Body.String()) + } + }) + } +} + +func TestGeminiNativeCandidateStopsAreDelivered(t *testing.T) { + for _, reason := range []string{ + "STOP", "MAX_TOKENS", "SAFETY", "RECITATION", "LANGUAGE", "OTHER", "BLOCKLIST", + "PROHIBITED_CONTENT", "SPII", "MALFORMED_FUNCTION_CALL", "UNEXPECTED_TOOL_CALL", + "TOO_MANY_TOOL_CALLS", "IMAGE_SAFETY", "IMAGE_PROHIBITED_CONTENT", "IMAGE_OTHER", + "NO_IMAGE", "CONTINUATION", "FUTURE_FINISH_REASON", + } { + t.Run(reason, func(t *testing.T) { + data := `{"candidates":[{"index":0,"finishReason":` + strconv.Quote(reason) + `}]}` + stream := string(geminiTestEvent(data).raw) + var st captureStats + rec, wrote, f := captureGemini(t, nil, strings.NewReader(stream), &st, false) + if !wrote || f != nil || rec.Body.String() != stream || st.stop != reason { + t.Fatalf("native stop was not preserved: wrote=%v failure=%+v stop=%q body=%q", wrote, f, st.stop, rec.Body.String()) + } + }) + } +} + +func TestGeminiPromptBlocksCompleteWithoutCandidates(t *testing.T) { + for _, reason := range []string{"SAFETY", "OTHER", "BLOCKLIST", "PROHIBITED_CONTENT", "IMAGE_SAFETY", "FUTURE_BLOCK_REASON"} { + t.Run(reason, func(t *testing.T) { + stream := string(geminiTestEvent(`{"promptFeedback":{"blockReason":` + strconv.Quote(reason) + `},"usageMetadata":{"promptTokenCount":12}}`).raw) + var st captureStats + rec, wrote, f := captureGemini(t, []byte(`{"generationConfig":{"candidateCount":3}}`), strings.NewReader(stream), &st, false) + if !wrote || f != nil || rec.Body.String() != stream || st.stop != reason || st.inTok != 12 { + t.Fatalf("native prompt block was not preserved: wrote=%v failure=%+v stats=%+v body=%q", wrote, f, st, rec.Body.String()) + } + }) + } + for _, feedback := range []string{`{}`, `{"blockReason":""}`, `{"blockReason":"BLOCK_REASON_UNSPECIFIED"}`} { + wire := newGeminiWire(nil) + if done, f := wire.buffered(geminiTestEvent(`{"promptFeedback":`+feedback+`}`), nil); done || f != nil || wire.completeOnEOF() { + t.Fatalf("non-blocking feedback must not complete: feedback=%s done=%v failure=%+v", feedback, done, f) + } + } +} + +func TestGeminiUsageSnapshotsDoNotDoubleCount(t *testing.T) { + wire := newGeminiWire(nil) + var st captureStats + for _, data := range []string{ + `{"modelVersion":"gemini-first","usageMetadata":{"promptTokenCount":100,"cachedContentTokenCount":25,"candidatesTokenCount":4,"thoughtsTokenCount":3,"totalTokenCount":107}}`, + `{"usageMetadata":{"promptTokenCount":100,"cachedContentTokenCount":25,"candidatesTokenCount":4,"thoughtsTokenCount":3,"totalTokenCount":107}}`, + } { + ev := geminiTestEvent(data) + wire.scrape(ev, &st) + if done, f := wire.buffered(ev, &st); done || f != nil { + t.Fatalf("usage event: done=%v failure=%+v", done, f) + } + if st.inTok != 100 || st.inputTok != 75 || st.cacheReadTok != 25 || st.cacheCreationTok != 0 || st.outTok != 7 { + t.Fatalf("repeated snapshot double-counted prompt/cache/output: %+v", st) + } + } + for _, data := range []string{ + `{"modelVersion":"gemini-resolved","usageMetadata":{"candidatesTokenCount":6,"thoughtsTokenCount":5}}`, + `{"usageMetadata":{"thoughtsTokenCount":7}}`, + } { + if _, f := wire.buffered(geminiTestEvent(data), &st); f != nil { + t.Fatalf("trailing usage: %+v", f) + } + } + if st.inTok != 100 || st.cacheReadTok != 25 || st.outTok != 13 || st.model != "gemini-resolved" { + t.Fatalf("partial trailing usage must retain independent cumulative components: %+v", st) + } + if _, f := wire.buffered(geminiTestEvent(`{"usageMetadata":{"promptTokenCount":10,"cachedContentTokenCount":20}}`), &st); f != nil { + t.Fatalf("cache clamp: %+v", f) + } + if st.inTok != 10 || st.inputTok != 0 || st.cacheReadTok != 10 || st.outTok != 13 { + t.Fatalf("cache count must not exceed the prompt total: %+v", st) + } +} + +func TestGeminiMalformedEventsFailEvenWhenValidationDisabled(t *testing.T) { + saved := cfg.validateJSON + t.Cleanup(func() { cfg.validateJSON = saved }) + cases := map[string]string{ + "empty": "", + "broken JSON": `{"candidates":[`, + "null": `null`, + "array": `[]`, + "string": `"not an object"`, + "number": `12`, + "boolean": `true`, + "trailing JSON": `{} {}`, + "candidates object": `{"candidates":{}}`, + "candidate null": `{"candidates":[null]}`, + "candidate scalar": `{"candidates":[12]}`, + "candidate array": `{"candidates":[[]]}`, + "null index": `{"candidates":[{"index":null,"finishReason":"STOP"}]}`, + "negative index": `{"candidates":[{"index":-1,"finishReason":"STOP"}]}`, + "out of range index": `{"candidates":[{"index":1,"finishReason":"STOP"}]}`, + "fractional index": `{"candidates":[{"index":0.5,"finishReason":"STOP"}]}`, + "exponent index": `{"candidates":[{"index":0e0,"finishReason":"STOP"}]}`, + "string index": `{"candidates":[{"index":"0","finishReason":"STOP"}]}`, + "boolean index": `{"candidates":[{"index":false,"finishReason":"STOP"}]}`, + "object index": `{"candidates":[{"index":{},"finishReason":"STOP"}]}`, + "overflowing index": `{"candidates":[{"index":999999999999999999999999,"finishReason":"STOP"}]}`, + "duplicate index": `{"candidates":[{"index":0,"finishReason":"STOP"},{"index":0,"finishReason":"STOP"}]}`, + "duplicate omitted index": `{"candidates":[{"finishReason":"STOP"},{"finishReason":"STOP"}]}`, + "invalid finish type": `{"candidates":[{"finishReason":12}]}`, + "content array": `{"candidates":[{"content":[]}]}`, + "parts object": `{"candidates":[{"content":{"parts":{}}}]}`, + "part null": `{"candidates":[{"content":{"parts":[null]}}]}`, + "part scalar": `{"candidates":[{"content":{"parts":[1]}}]}`, + "promptFeedback array": `{"promptFeedback":[]}`, + "usageMetadata array": `{"usageMetadata":[]}`, + } + for _, validate := range []bool{false, true} { + cfg.validateJSON = validate + for name, data := range cases { + t.Run(strconv.FormatBool(validate)+"/"+name, func(t *testing.T) { + for _, live := range []bool{false, true} { + wire := newGeminiWire(nil) + if _, f := wire.buffered(geminiTestEvent(`{"candidates":[{"finishReason":"STOP"}]}`), nil); f != nil { + t.Fatalf("initial finish: %+v", f) + } + ev := geminiTestEvent(data) + var f *failure + var done bool + if live { + var out [][]byte + out, done, f = wire.live(ev, ev.raw, nil) + if len(out) != 0 { + t.Fatalf("malformed live data must not be forwarded: %q", out) + } + } else { + done, f = wire.buffered(ev, nil) + } + if done || f == nil || !f.transient || f.code != "malformed_sse" || wire.completeOnEOF() { + t.Fatalf("malformed event accepted: live=%v done=%v failure=%+v complete=%v", live, done, f, wire.completeOnEOF()) + } + } + }) + } + } +} + +func TestGeminiInBandErrorsFailBeforeAndAfterFinish(t *testing.T) { + saved := cfg.validateJSON + cfg.validateJSON = false + t.Cleanup(func() { cfg.validateJSON = saved }) + for _, tc := range []struct { + name string + data string + eventName string + transient bool + }{ + {"native unavailable", `{"error":{"code":503,"status":"UNAVAILABLE","message":"overloaded"}}`, "", true}, + {"native invalid argument", `{"error":{"code":400,"status":"INVALID_ARGUMENT","message":"bad request"}}`, "", false}, + {"native not found", `{"error":{"code":404,"status":"NOT_FOUND","message":"model missing"}}`, "", false}, + {"named error", `{"error":{"code":503,"status":"UNAVAILABLE","message":"overloaded"}}`, "error", true}, + {"null error envelope", `{"error":null}`, "", true}, + {"missing named error envelope", `{}`, "error", true}, + } { + t.Run(tc.name, func(t *testing.T) { + for _, afterFinish := range []bool{false, true} { + for _, live := range []bool{false, true} { + wire := newGeminiWire(nil) + if afterFinish { + if _, f := wire.buffered(geminiTestEvent(`{"candidates":[{"index":0,"finishReason":"STOP"}]}`), nil); f != nil { + t.Fatalf("initial finish: %+v", f) + } + } + ev := geminiTestEvent(tc.data) + ev.name = tc.eventName + var done bool + var f *failure + if live { + var out [][]byte + out, done, f = wire.live(ev, ev.raw, nil) + if len(out) != 0 { + t.Fatalf("in-band error must not be forwarded: %q", out) + } + } else { + done, f = wire.buffered(ev, nil) + } + if done || f == nil || f.transient != tc.transient || wire.completeOnEOF() { + t.Fatalf("in-band error accepted or misclassified: live=%v afterFinish=%v done=%v failure=%+v complete=%v", live, afterFinish, done, f, wire.completeOnEOF()) + } + } + } + }) + } +} + +func TestGeminiCaptureRejectsTruncationAndNonCleanEOF(t *testing.T) { + finish := string(geminiTestEvent(`{"candidates":[{"index":0,"finishReason":"STOP"}]}`).raw) + for _, tc := range []struct { + name string + body string + stream string + err error + }{ + {"empty stream", "", "", nil}, + {"no finish", "", string(geminiTestEvent(`{"candidates":[{"index":0,"content":{"parts":[{"text":"unfinished"}]}}]}`).raw), nil}, + {"missing requested candidate", `{"generationConfig":{"candidateCount":2}}`, finish, nil}, + {"unfinished data frame", "", finish + "data: {\"usageMetadata\":{\"promptTokenCount\":9}}\n", nil}, + {"unfinished line", "", finish + "data: {\"usageMetadata\":", nil}, + {"unfinished frame with data and EOF together", "", finish + "data: {}", nil}, + {"unexpected EOF after finish", "", finish, io.ErrUnexpectedEOF}, + {"transport error after finish", "", finish, errors.New("connection reset")}, + } { + t.Run(tc.name, func(t *testing.T) { + var r io.Reader = &geminiTestReader{chunks: []string{tc.stream}, err: tc.err} + if tc.name == "unfinished frame with data and EOF together" { + r = &geminiTestFinalReader{data: []byte(tc.stream)} + } + rec, wrote, f := captureGemini(t, []byte(tc.body), r, nil, false) + if wrote || f == nil || !f.transient || rec.Body.Len() != 0 { + t.Fatalf("incomplete stream must remain uncommitted: wrote=%v failure=%+v body=%q", wrote, f, rec.Body.String()) + } + }) + } +} + +func TestGeminiCaptureLiveWaitsForCleanEOF(t *testing.T) { + chunks := strings.SplitAfter(goodGeminiStream, "\n\n") + chunks = chunks[:len(chunks)-1] + var st captureStats + rec, wrote, f := captureGemini(t, nil, &geminiTestReader{chunks: chunks}, &st, true) + if !wrote || f != nil || st.mode != "live" || rec.Body.String() != goodGeminiStream || st.outTok != 11 { + t.Fatalf("live response lost trailing usage: wrote=%v failure=%+v stats=%+v body=%q", wrote, f, st, rec.Body.String()) + } +} + +func TestGeminiCaptureLiveRejectsErrorAfterFinish(t *testing.T) { + prefix := string(geminiTestEvent(`{"candidates":[{"content":{"parts":[{"text":"prefix"}]},"finishReason":"STOP"}]}`).raw) + for _, tc := range []struct { + name string + tail string + err error + }{ + {"in-band error", string(geminiTestEvent(`{"error":{"code":503,"status":"UNAVAILABLE","message":"failed after finish"}}`).raw), nil}, + {"malformed event", "data: not-json\n\n", nil}, + {"trailing unfinished frame", "data: {}", nil}, + {"transport error", "", io.ErrUnexpectedEOF}, + } { + t.Run(tc.name, func(t *testing.T) { + var st captureStats + rec, wrote, f := captureGemini(t, nil, &geminiTestReader{chunks: []string{prefix, tc.tail}, err: tc.err}, &st, true) + if !wrote || f == nil || !f.transient || st.mode != "live" || rec.Body.String() != prefix { + t.Fatalf("post-commit failure must DROP without forwarding the bad tail: wrote=%v failure=%+v stats=%+v body=%q", wrote, f, st, rec.Body.String()) + } + }) + } +} diff --git a/main.go b/main.go index 58d8640..f47c4eb 100644 --- a/main.go +++ b/main.go @@ -1,6 +1,5 @@ // steady-proxy — a transactional, self-healing reverse proxy for API clients -// using Anthropic Messages (POST /v1/messages) or OpenAI Responses (POST -// /v1/responses), including Claude Code and Codex. +// using Anthropic Messages, OpenAI Responses, or native Gemini GenerateContent. // // Agents / SDKs / applications ──HTTP──▶ this proxy (loopback) ──HTTPS──▶ your gateway // @@ -35,6 +34,9 @@ // SSE stream (terminal `response.completed`); a start-of-stream `error` / // `response.failed` is caught pre-commit and converted to a retry, and once // output appears it commits early and streams live. See wireModel (wire.go). +// - Native Gemini streamGenerateContent?alt=sse buffers through clean EOF +// after all requested candidates finish, retaining trailing usage metadata. +// PROXY_GEMINI_BUFFER_MS bounds the hold before switching to live output. // - Any failure before that commit point (connection error, 5xx, stalled or // truncated stream, mid-stream `error` event, or a retryable status) is // turned into a *retryable* response: it stamps `x-should-retry: true` plus a @@ -91,6 +93,7 @@ type config struct { wfKeepaliveMs time.Duration responsesKeepaliveMs time.Duration responsesBufferMs time.Duration + geminiBufferMs time.Duration responsesEarlyCommit bool // Responses (Codex): commit on first output (default); false = buffer to terminal under responsesBufferMs, like /v1/messages responsesReplayBytes int // Responses buffered/prefix replay: coalesce adjacent output_text.delta events up to this many source JSON bytes; 0 = raw replay deadlineMargin time.Duration @@ -132,6 +135,7 @@ func parseConfig(lookup func(string) (string, bool)) (config, error) { wfKeepaliveMs: e.ms("PROXY_WORKFLOW_KEEPALIVE_MS", 10000, 0), // SMALL window applied ONLY to Workflow-tool agents (see isWorkflowAgent): commit early + stream live so forwarded deltas feed their per-agent stall watchdog (default 180s, no global env override). The watchdog kills on a forwarded-content-delta gap >= stallMs; while buffering the proxy forwards nothing, so the FIRST gap ~= window + upstream TTFB — keep the window small to hold that first gap under stallMs. A large window is not automatically fatal (a steady stream can survive it) but pushes the first gap toward stallMs and delays going live. It CANNOT fix a mid-turn content-silent pause >= stallMs; only a larger per-agent stallMs can. 0 = disable (falls back to full transactional; stall bug returns). responsesKeepaliveMs: e.ms("PROXY_RESPONSES_KEEPALIVE_MS", 30000, 0), // OpenAI Responses route (Codex), applies ONLY in early-commit mode: bounds a silent start (reasoning with no output) before committing + streaming keepalives. With early-commit off the route buffers under PROXY_RESPONSES_BUFFER_MS instead (see config.keepaliveWindow). Keep it below the client's stream-idle timeout. Start-of-stream errors (overload/capacity) arrive before any output, so they are always caught pre-commit regardless of this value. responsesBufferMs: e.ms("PROXY_RESPONSES_BUFFER_MS", 600000, 0), // Responses (Codex) hold used ONLY when early-commit is off: buffer to the terminal, committing live at this window as a safety valve (the /v1/messages policy applied to Responses). Responses-OWNED — independent of PROXY_KEEPALIVE_MS — so a mixed Codex+Claude proxy tunes the two horizons separately. The Codex client's stream_idle_timeout_ms must strictly exceed it (nothing is forwarded while buffering). + geminiBufferMs: e.ms("PROXY_GEMINI_BUFFER_MS", 600000, 0), // Gemini buffers through clean EOF; 0 disables the timed switch to live output. responsesEarlyCommit: e.str("PROXY_RESPONSES_EARLY_COMMIT", "1") != "0", // default on: commit as soon as the first output event is buffered. 0 -> buffer the whole response to response.completed like /v1/messages, under the PROXY_RESPONSES_BUFFER_MS hold (extends hidden-retry protection to mid-stream errors), which needs the Codex client's stream_idle_timeout_ms to exceed that hold. responsesReplayBytes: int(e.int64("PROXY_RESPONSES_REPLAY_DELTA_BYTES", 64<<10, 0)), // collapse tiny adjacent Codex text deltas during buffered/prefix replay so a completed turn cannot burst thousands of app-server notifications into the TUI; 0 preserves byte-for-byte replay. Live events remain untouched. deadlineMargin: e.ms("PROXY_DEADLINE_MARGIN_MS", 25000, 0), // finish before the client's own timeout @@ -214,8 +218,8 @@ func main() { // WriteTimeout intentionally 0: long-lived holds; ctx deadlines bound work. MaxHeaderBytes: 1 << 20, } - log.Printf(programName+" %s listening on http://%s -> %s (transactional: /v1/messages + /v1/responses, keepalive=%s, wf-keepalive=%s, responses-keepalive=%s, responses-buffer=%s, responses-early-commit=%v, responses-replay-delta=%dB, sdkRetryCap=%d, txLocalRetries=%d, refusalFallback=%s; one log line per request)", - currentVersion().token(), cfg.listenAddr, cfg.upstream, cfg.keepaliveMs, cfg.wfKeepaliveMs, cfg.responsesKeepaliveMs, cfg.responsesBufferMs, cfg.responsesEarlyCommit, cfg.responsesReplayBytes, cfg.sdkRetryCap, cfg.txLocalRetries, refusalFallbackDesc(cfg.refusalFallback)) + log.Printf(programName+" %s listening on http://%s -> %s (transactional: Messages + Responses + Gemini SSE, keepalive=%s, wf-keepalive=%s, responses-keepalive=%s, responses-buffer=%s, gemini-buffer=%s, responses-early-commit=%v, responses-replay-delta=%dB, sdkRetryCap=%d, txLocalRetries=%d, refusalFallback=%s; one log line per request)", + currentVersion().token(), cfg.listenAddr, cfg.upstream, cfg.keepaliveMs, cfg.wfKeepaliveMs, cfg.responsesKeepaliveMs, cfg.responsesBufferMs, cfg.geminiBufferMs, cfg.responsesEarlyCommit, cfg.responsesReplayBytes, cfg.sdkRetryCap, cfg.txLocalRetries, refusalFallbackDesc(cfg.refusalFallback)) log.Fatal(srv.ListenAndServe()) } @@ -251,15 +255,10 @@ func handle(w http.ResponseWriter, r *http.Request) { return } reqStart := time.Now() - // The client-facing error shape follows the request's wire (Anthropic Messages - // vs OpenAI Responses); resolve it up front so even an early reject renders in - // the shape the client expects. + // Resolve the wire before reading the body so early rejects also use its + // native error envelope. Only supported SSE generation requests are buffered; + // other routes remain single-shot with HTTP retry normalization. writeErr := errorWriterFor(r.URL.Path) - // Only streaming POST /v1/messages and POST /v1/responses get the full - // transactional treatment. Everything else (e.g. /v1/messages/count_tokens, - // model listing) is forwarded with a single attempt; transport errors AND - // upstream HTTP errors there are still converted to retryable and masked to a - // generic 503 (see proxyOnce). body, tooBig, err := readBody(r) if tooBig { // Deterministic: the same oversized request can never succeed. surfaceFor keeps @@ -281,8 +280,10 @@ func handle(w http.ResponseWriter, r *http.Request) { defer rec.finish() } - transactional := r.Method == http.MethodPost && requestWantsStream(body) && - (isMessagesPath(r.URL.Path) || isResponsesPath(r.URL.Path)) + _, geminiMethod := geminiTarget(r.URL.Path) + geminiStream := geminiMethod == "streamGenerateContent" && r.URL.Query().Get("alt") == "sse" + transactional := r.Method == http.MethodPost && (geminiStream || + (requestWantsStream(body) && (isMessagesPath(r.URL.Path) || isResponsesPath(r.URL.Path)))) if !transactional { proxyOnce(w, r, body, rec) return @@ -295,6 +296,7 @@ func handle(w http.ResponseWriter, r *http.Request) { // become hidden retries too — at the cost of a longer initial silence the // client's stream-idle timeout must cover). isResp := isResponsesPath(r.URL.Path) + isMessages := isMessagesPath(r.URL.Path) earlyCommit := isResp && cfg.responsesEarlyCommit retryCount := atoiSafe(r.Header.Get("X-Stainless-Retry-Count")) @@ -320,8 +322,11 @@ func handle(w http.ResponseWriter, r *http.Request) { // commit-early window AND gate the commit on real forwarded progress (see // captureSSEWindow); every other caller keeps the full hold. The window choice // (per wire + mode) lives in config.keepaliveWindow. - gated := !isResp && isWorkflowAgent(body) // Workflow stall gating is Anthropic-only + gated := isMessages && isWorkflowAgent(body) kaWindow := cfg.keepaliveWindow(isResp, earlyCommit, gated) + if geminiStream { + kaWindow = cfg.geminiBufferMs + } commitBy := bufferDeadline(reqStart, kaWindow) canWait := func(wait time.Duration) bool { return canWaitForLocalRetry(ctx, wait) && (commitBy.IsZero() || time.Until(commitBy) > wait) @@ -385,7 +390,7 @@ func handle(w http.ResponseWriter, r *http.Request) { // again (no recursion). The swap itself is deferred to the refusal branch so // the happy path never marshals a fallback body it won't use. reqModel := modelOf(body) - interceptRefusal := !isResp && cfg.refusalFallback != "" && reqModel != "?" && reqModel != cfg.refusalFallback + interceptRefusal := isMessages && cfg.refusalFallback != "" && reqModel != "?" && reqModel != cfg.refusalFallback rec.beginAttempt(body, reqWho, proxyRetries) resp, started, stopDeadline, rtErr := roundTripBefore(ctx, r, body, commitBy) @@ -457,7 +462,7 @@ func handle(w http.ResponseWriter, r *http.Request) { // The stream engine now owns the same absolute commit deadline; it can // switch to live output without the pre-stream timer killing the body. stopDeadline() - wrote, fail := captureSSECore(ctx, cancel, w, resp.Header, resp.Body, &st, kaWindow, gated, earlyCommit, wireFor(r.URL.Path, interceptRefusal), commitBy) + wrote, fail := captureSSECore(ctx, cancel, w, resp.Header, resp.Body, &st, kaWindow, gated, earlyCommit, wireFor(r.URL.Path, body, interceptRefusal), commitBy) resp.Body.Close() rec.noteAttemptStats(&st) if fail == nil { @@ -490,11 +495,16 @@ func handle(w http.ResponseWriter, r *http.Request) { break } if wrote { // failed AFTER committing — can't convert, response already streaming - log.Printf("DROP %s%s %s -> committed, Claude retries natively out=%s tok %s%s%s%s%s", + log.Printf("DROP %s%s %s -> committed, client recovery required out=%s tok %s%s%s%s%s", whoWithResolvedModel(reqWho, st.model), effortField(reqEffort), fail.code, htok(st.outTok), since(reqStart), att(retryCount), proxyRetryField(proxyRetries), prunField(st.maxPingRun), rec.idField()) rec.noteAttemptResult("drop") rec.noteProxyRetries(proxyRetries) rec.note("DROP", http.StatusOK, http.StatusOK, fail.code) + if geminiStream { + // Gemini readers use transport EOF, not a terminal SSE event. A clean + // HTTP end here could turn a truncated response into apparent success. + panic(http.ErrAbortHandler) + } return } last = *fail @@ -771,7 +781,11 @@ func isEventStream(ct string) bool { return strings.Contains(strings.ToLower(ct) // The method/path/gateway are static for the transactional route, so omitted. func who(r *http.Request, body []byte) string { - return modelOf(body) + "/" + agentKindFull(r, body) + model, _ := geminiTarget(r.URL.Path) + if model == "" { + model = modelOf(body) + } + return model + "/" + agentKindFull(r, body) } // agentKindFull refines agentKind with a body-derived signal: a Workflow-tool @@ -779,7 +793,7 @@ func who(r *http.Request, body []byte) string { // generic "sub"agent, and only it carries the Workflow per-agent stall). Falls // back to the header-only agentKind for everything else. func agentKindFull(r *http.Request, body []byte) string { - if isWorkflowAgent(body) { + if isMessagesPath(r.URL.Path) && isWorkflowAgent(body) { return "wf" } return agentKind(r) diff --git a/requestlog.go b/requestlog.go index 4a4da58..24ab4ff 100644 --- a/requestlog.go +++ b/requestlog.go @@ -910,6 +910,7 @@ var apiSecretHeaders = map[string]bool{ "x-stainless-api-key": true, "anthropic-api-key": true, "openai-api-key": true, + "x-goog-api-key": true, "proxy-authorization": true, } @@ -964,6 +965,8 @@ var apiSecretQueryNames = map[string]bool{ "anthropic-api-key": true, "openai_api_key": true, "openai-api-key": true, + "x-goog-api-key": true, + "x_goog_api_key": true, "access_token": true, "accesstoken": true, "access-token": true, diff --git a/requestlog_test.go b/requestlog_test.go index a5783d4..2fd1892 100644 --- a/requestlog_test.go +++ b/requestlog_test.go @@ -401,6 +401,7 @@ func TestRequestLogOnlyRedactsAPISecrets(t *testing.T) { req.Header.Set("Authorization", "Bearer sk-secret-auth") req.Header.Set("X-API-Key", "sk-secret-header") req.Header.Set("Proxy-Authorization", "Basic proxy-secret") + req.Header.Set("X-Goog-Api-Key", "goog-secret-header") req.Header.Set("Cookie", "sid=request-cookie") rec := httptest.NewRecorder() handle(rec, req) @@ -413,7 +414,7 @@ func TestRequestLogOnlyRedactsAPISecrets(t *testing.T) { t.Fatalf("want 1 archive file, got %d", len(files)) } ar, raw := readArchive(t, files[0]) - for _, leaked := range []string{"sk-secret-auth", "sk-secret-header", "proxy-secret", "sk-secret-0000", "sk-secret-key", "sk-api-token", "sk-auth-token", "sk-access", "sk-client"} { + for _, leaked := range []string{"sk-secret-auth", "sk-secret-header", "goog-secret-header", "proxy-secret", "sk-secret-0000", "sk-secret-key", "sk-api-token", "sk-auth-token", "sk-access", "sk-client"} { if bytes.Contains(raw, []byte(leaked)) { t.Fatalf("API secret %q leaked into archive:\n%s", leaked, raw) } @@ -434,6 +435,7 @@ func TestRequestLogOnlyRedactsAPISecrets(t *testing.T) { } if !containsString(ar.Redactions.HeaderNames, "Authorization") || !containsString(ar.Redactions.HeaderNames, "X-Api-Key") || + !containsString(ar.Redactions.HeaderNames, "X-Goog-Api-Key") || !containsString(ar.Redactions.HeaderNames, "Proxy-Authorization") { t.Fatalf("missing API header redactions: %+v", ar.Redactions) } diff --git a/responses.go b/responses.go index ef3c98d..17d5ae6 100644 --- a/responses.go +++ b/responses.go @@ -38,6 +38,8 @@ func newOpenAIWire() *openaiWire { return &openaiWire{replayDeltaBytes: cfg.responsesReplayBytes} } +func (o *openaiWire) completeOnEOF() bool { return false } + func (o *openaiWire) scrape(ev event, st *captureStats) { if st == nil { return @@ -1030,6 +1032,9 @@ func classifyHTTPErrorForPath(path string, resp *http.Response, b []byte) failur if isResponsesPath(path) { return classifyResponsesHTTPErrorBytes(resp, b) } + if isGeminiPath(path) { + return classifyGeminiHTTPErrorBytes(resp, b) + } return classifyHTTPErrorBytes(resp, b) } @@ -1043,6 +1048,9 @@ func errorWriterFor(path string) errorWriter { if isResponsesPath(path) { return writeOpenAIError } + if isGeminiPath(path) { + return writeGeminiError + } return writeAnthropicError } diff --git a/responses_test.go b/responses_test.go index 07a6b08..9ef4de8 100644 --- a/responses_test.go +++ b/responses_test.go @@ -1012,12 +1012,6 @@ func TestPathHelpers(t *testing.T) { if !isMessagesPath("/v1/messages/batches") || isMessagesPath("/v1/messagesfoo") { t.Fatal("isMessagesPath boundary matching wrong") } - if _, ok := wireFor("/v1/responses", false).(*openaiWire); !ok { - t.Fatal("wireFor(/v1/responses) should be openaiWire") - } - if _, ok := wireFor("/v1/messages", false).(*anthropicWire); !ok { - t.Fatal("wireFor(/v1/messages) should be anthropicWire") - } } // ───────────────────────────────────────────────────────────── e2e (handle) ── diff --git a/sse.go b/sse.go index b20c8ae..7b8b6e0 100644 --- a/sse.go +++ b/sse.go @@ -392,11 +392,17 @@ func captureSSECore(ctx context.Context, cancel context.CancelFunc, w http.Respo return wrote, fail } } - // Check the read error BEFORE any deferred commit: a read can return - // data together with io.EOF, and if that final read is a truncation - // (no message_stop) we must convert it to a clean uncommitted retry - // rather than commit the partial buffer and log a post-commit DROP. + // Gemini has no terminal SSE event: accept clean EOF only after the + // wire reports completion and no partial frame remains. Read faults and + // cancellation must still fail even after a candidate's finishReason. if r.err != nil { + if r.err == io.EOF && ctx.Err() == nil && !parser.cur.has && len(parser.pending) == 0 && model.completeOnEOF() { + if !committed { + return true, replayCaptured(sp, model, w, st, writeHead, flush) + } + flush() + return true, nil + } return onReadErr(r.err, ctx, committed, flush) } // Clean read, gated caller: the window already elapsed and real @@ -528,16 +534,7 @@ func process(data []byte, sp *spool, parser *sseParser, model wireModel, w http. return false, classifyBufferError(err), true } if done { - // Complete, valid stream: replay the whole buffer downstream. This runs - // before any bytes reached the client, so it is still a clean commit. - writeHead("buffered") - if err := model.replayBuffered(sp, w, st); err != nil { - sp.discard() - return true, &failure{transient: true, status: http.StatusBadGateway, atype: "api_error", code: "response_replay_failed", message: "failed replaying buffered response"}, true - } - flush() - sp.discard() - return true, nil, true + return true, replayCaptured(sp, model, w, st, writeHead, flush), true } } if parser.err != nil { @@ -546,6 +543,16 @@ func process(data []byte, sp *spool, parser *sseParser, model wireModel, w http. return false, nil, false } +func replayCaptured(sp *spool, model wireModel, w io.Writer, st *captureStats, writeHead func(string), flush func()) *failure { + writeHead("buffered") + if err := model.replayBuffered(sp, w, st); err != nil { + return &failure{transient: true, status: http.StatusBadGateway, atype: "api_error", code: "response_replay_failed", message: "failed replaying buffered response"} + } + flush() + sp.discard() + return nil +} + func validateSSEEventShape(ev event) *failure { switch { case ev.name != "" && ev.data == "": @@ -977,7 +984,7 @@ func onReadErr(err error, ctx context.Context, committed bool, flush func()) (bo case ctx.Err() == context.Canceled: f = &failure{transient: false, status: 499, atype: "api_error", code: "client_gone", message: "client cancelled"} case err == io.EOF || err == io.ErrUnexpectedEOF: - f = &failure{transient: true, status: 502, atype: "api_error", code: "truncated_stream", message: "EOF before message_stop"} + f = &failure{transient: true, status: 502, atype: "api_error", code: "truncated_stream", message: "stream ended before a complete response"} default: f = &failure{transient: true, status: 502, atype: "api_error", code: "stream_read_error", message: err.Error()} } diff --git a/wire.go b/wire.go index 8556ae5..6f5635e 100644 --- a/wire.go +++ b/wire.go @@ -12,12 +12,9 @@ import ( // The transactional engine (spool + keepalive→commit→live + local-retry + // idle/deadline, see captureSSECore / handle) is protocol-agnostic: its one rule // is "buffer the full response; convert any pre-commit failure into a retry; only -// give up on a request that can never succeed." What differs between the Anthropic -// Messages wire (`POST /v1/messages`) and the OpenAI Responses wire -// (`POST /v1/responses`, used by Codex) is only how the byte stream is -// *interpreted*: when it is complete, which event is an error, and how tokens are -// scraped. wireModel owns exactly that; one instance is created per upstream -// attempt. +// give up on a request that can never succeed." Each wire interprets completion, +// in-band errors, and usage for Anthropic Messages, OpenAI Responses, or Gemini +// GenerateContent. One instance is created per upstream attempt. type wireModel interface { // scrape pulls best-effort access-log facts (model/usage/stop) out of one // event into st. Called for every event; must never fail a stream. @@ -35,6 +32,10 @@ type wireModel interface { // (engine writes this event, then replays via replayBuffered). buffered(ev event, st *captureStats) (done bool, fail *failure) + // completeOnEOF accepts a clean end of transport for protocols without a + // terminal SSE event. The engine also requires a fully framed final event. + completeOnEOF() bool + // live ingests one event on the post-commit LIVE path: it returns the exact // downstream bytes to forward (empty -> engine emits keepalive()), whether the // stream terminated, and fail != nil for an in-band error (the engine ends the @@ -44,14 +45,13 @@ type wireModel interface { // keepalive is the byte frame the engine writes downstream during a silent gap // to keep the client's connection/idle timer alive. It MUST be one the client - // treats as a no-op. The two wires differ: the Anthropic SDK is fine with a - // spec SSE comment, but Codex's SSE reader only resets its idle timer on a - // yielded event and discards comments — so the Responses wire sends a skippable - // unknown-type event instead (see openaiWire.keepalive). + // treats as a no-op. Messages and Gemini use a standard SSE comment. Codex's + // reader discards comments without resetting its idle timer, so Responses + // sends a skippable unknown-type event (see openaiWire.keepalive). keepalive() []byte // replayBuffered writes the whole buffered stream downstream once, at the - // terminal-in-buffer commit. + // terminal-event or clean-EOF commit. replayBuffered(sp *spool, w io.Writer, st *captureStats) error // replayPrefix writes the buffered prefix at a keepalive/early commit, before @@ -63,14 +63,14 @@ type wireModel interface { forwardable() bool } -// wireFor selects the wire model for a request path. Only the two transactional -// routes reach here; everything else uses proxyOnce and never needs a model. The -// Anthropic-only refusal intercept is baked into the returned wire, so the shared -// engine never carries it (interceptRefusal is meaningless for Responses). -func wireFor(path string, interceptRefusal bool) wireModel { +// wireFor selects a stream interpreter. Refusal interception is Anthropic-only. +func wireFor(path string, body []byte, interceptRefusal bool) wireModel { if isResponsesPath(path) { return newOpenAIWire() } + if isGeminiPath(path) { + return newGeminiWire(body) + } return newAnthropicWire(interceptRefusal) } @@ -89,6 +89,55 @@ func isResponsesPath(path string) bool { return pathIs(path, "/v1/responses") } +// geminiTarget recognizes native Developer API and Vertex AI model methods. +// Exact resource boundaries keep look-alike and unrelated routes single-shot. +func geminiTarget(path string) (model, method string) { + version, resource, ok := strings.Cut(strings.TrimPrefix(path, "/"), "/") + if !ok { + return "", "" + } + switch version { + case "v1", "v1beta", "v1alpha", "v1beta1": + default: + return "", "" + } + var name string + if rest, vertex := strings.CutPrefix(resource, "projects/"); vertex { + project, rest, ok := strings.Cut(rest, "/locations/") + if !ok || project == "" || strings.Contains(project, "/") { + return "", "" + } + location, rest, ok := strings.Cut(rest, "/publishers/google/models/") + if !ok || location == "" || strings.Contains(location, "/") { + return "", "" + } + name = rest + } else { + name, ok = strings.CutPrefix(resource, "models/") + if !ok { + return "", "" + } + } + if strings.Contains(name, "/") { + return "", "" + } + model, method, ok = strings.Cut(name, ":") + if !ok || model == "" { + return "", "" + } + switch method { + case "generateContent", "streamGenerateContent", "countTokens": + return model, method + default: + return "", "" + } +} + +func isGeminiPath(path string) bool { + _, method := geminiTarget(path) + return method != "" +} + // ─────────────────────────────────────────────────────── anthropic wire model ── // // anthropicWire is a thin adapter over the existing, unchanged Anthropic stream @@ -116,6 +165,8 @@ func newAnthropicWire(interceptRefusal bool) *anthropicWire { func (a *anthropicWire) scrape(ev event, st *captureStats) { scrapeUsage(ev, st) } +func (a *anthropicWire) completeOnEOF() bool { return false } + func (a *anthropicWire) buffered(ev event, st *captureStats) (bool, *failure) { if cfg.validateJSON { if f := validateSSEEventShape(ev); f != nil { From deb09016a3bb2d615fd4ff1bf024cdd40a84817c Mon Sep 17 00:00:00 2001 From: etwk <48991073+etwk@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:49:23 +0800 Subject: [PATCH 2/2] fix: report Gemini thinking levels and preserve retry signals --- classify.go | 34 ++++++++---- docs/ARCHITECTURE.md | 9 ++-- docs/OPERATIONS.md | 8 ++- docs/TROUBLESHOOTING.md | 12 +++++ gemini.go | 78 ++++++++++++++++++++++----- gemini_errors.go | 62 +++++++++++++++++++-- gemini_errors_test.go | 52 +++++++++++++++++- gemini_integration_test.go | 108 +++++++++++++++++++++++++++++++++++++ gemini_test.go | 49 ++++++++++++++++- main.go | 43 +++++++++------ reliability_test.go | 33 ++++++++++++ requestlog.go | 2 +- responses.go | 18 ++++++- transport.go | 13 ++++- 14 files changed, 467 insertions(+), 54 deletions(-) diff --git a/classify.go b/classify.go index 39f2ee6..a060a59 100644 --- a/classify.go +++ b/classify.go @@ -111,8 +111,7 @@ func anyContains(s string, subs []string) bool { return false } -// classifyTransport handles a failure that occurred before/around the response -// headers (dial/TLS/reset/timeout). +// classifyTransport handles failed upstream I/O and request cancellation. func classifyTransport(err error, ctx context.Context) failure { switch { case ctx.Err() == context.DeadlineExceeded || errors.Is(err, context.DeadlineExceeded): @@ -156,6 +155,19 @@ func classifyHTTPErrorBytes(resp *http.Response, b []byte) failure { return classifyHTTPErrorBytesShaped(resp, b, anthropicRequestShaped) } +// httpRetryHint returns the first recognized authoritative retry verdict. +func httpRetryHint(h http.Header) (source string, retry bool) { + for _, name := range [...]string{"x-gateway-retryable", "x-should-retry"} { + switch strings.ToLower(h.Get(name)) { + case "true": + return name, true + case "false": + return name, false + } + } + return "", false +} + // classifyHTTPErrorBytesShaped is classifyHTTPErrorBytes with a pluggable // request-shape decision so each wire keeps its own "can never succeed" rule // while sharing the status/header handling. @@ -178,17 +190,17 @@ func classifyHTTPErrorBytesShaped(resp *http.Response, b []byte, reqShaped reque // beats every heuristic below. x-should-retry is the Anthropic API's own near- // equivalent. Both are absent from a plain provider, which falls through // to status/body classification. - switch strings.ToLower(resp.Header.Get("x-gateway-retryable")) { - case "true": - return failure{transient: true, status: mapTransientStatus(st), atype: atype, code: "gateway_retryable_" + itoa(st), message: ae.Error.Message, retryAfter: ra} - case "false": + switch source, retry := httpRetryHint(resp.Header); source { + case "x-gateway-retryable": + if retry { + return failure{transient: true, status: mapTransientStatus(st), atype: atype, code: "gateway_retryable_" + itoa(st), message: ae.Error.Message, retryAfter: ra} + } return failure{transient: false, status: st, atype: atype, code: "gateway_permanent", message: ae.Error.Message, retryAfter: ra} - } - switch strings.ToLower(resp.Header.Get("x-should-retry")) { - case "false": + case "x-should-retry": + if retry { + return failure{transient: true, status: mapTransientStatus(st), atype: atype, code: "upstream_retry_" + itoa(st), message: ae.Error.Message, retryAfter: ra} + } return failure{transient: false, status: st, atype: atype, code: "upstream_no_retry", message: ae.Error.Message} - case "true": - return failure{transient: true, status: mapTransientStatus(st), atype: atype, code: "upstream_retry_" + itoa(st), message: ae.Error.Message, retryAfter: ra} } switch { diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index e68fb86..7fe2cec 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -75,9 +75,12 @@ Only `streamGenerateContent?alt=sse` receives transactional protection; `generateContent`, `countTokens`, and non-SSE streaming transports remain single-shot. Gemini has no separate terminal SSE event, so the proxy waits for clean EOF, retaining usage-only frames after candidate completion. Missing -candidates, unfinished SSE frames, and read faults fail the attempt. Native -`MAX_TOKENS`, safety stops, and prompt blocks are delivered unchanged, not retried -as transport failures. See Google's [GenerateContent contract](https://ai.google.dev/api/generate-content). +candidates, unfinished SSE frames, and read faults fail the attempt. SSE control +frames and empty data messages do not advance candidate completion. Native +`MAX_TOKENS`, safety stops, and prompt blocks without candidates are delivered +unchanged, not retried as transport failures. A prompt block accompanied by +candidates is contradictory and cannot satisfy completion. See Google's +[GenerateContent contract](https://ai.google.dev/api/generate-content). Candidate counts follow [ProtoJSON](https://protobuf.dev/programming-guides/json/): camelCase and snake_case field names, quoted integers, exponent notation, and zero fractional parts are accepted without rewriting the request. diff --git a/docs/OPERATIONS.md b/docs/OPERATIONS.md index 7069326..a8c363b 100644 --- a/docs/OPERATIONS.md +++ b/docs/OPERATIONS.md @@ -74,7 +74,7 @@ The fields after the outcome explain the request: | Field | Meaning | | --- | --- | | `model/main`, `model/sub`, `model/wf` | Requested model and detected caller kind. `requested->resolved` shows a different model reported by the upstream. | -| `high` or another effort value | Reasoning effort requested through Messages `output_config.effort` or Responses `reasoning.effort`. | +| `high` or another effort value | Requested Messages `output_config.effort`, Responses `reasoning.effort`, or Gemini `generationConfig.thinkingConfig.thinkingLevel` (lowercased). | | `in=… out=… tok` | Input and output token counts. Input includes cache reads and creation. | | `end_turn`, `completed`, etc. | The reported stop reason. | | `truncated_stream`, `sse_overloaded`, etc. | The classified failure code; use it to identify the actual cause. | @@ -89,6 +89,12 @@ Caller labels use Claude Code headers and the Workflow system-prompt marker. Requests without those signals appear as `main`, including requests from other applications; the label does not identify or restrict the client software. +Gemini thinking-level logging also accepts ProtoJSON snake_case field names. +Only an explicit named level is reported; the proxy does not infer one from +`thinkingBudget` or a model's default. The original request is forwarded unchanged. +Gemini model labels and archived attempt models come from the model URL, not a +body-level `model` field. + Verbose `[local-retry]` lines include the cause, wait, and retry index. `skip=not-enough-time` explains why another attempt could not fit the time budget. The response headers also expose `X-Steady-Proxy-Mode` on streaming delivery and diff --git a/docs/TROUBLESHOOTING.md b/docs/TROUBLESHOOTING.md index 5835296..8651899 100644 --- a/docs/TROUBLESHOOTING.md +++ b/docs/TROUBLESHOOTING.md @@ -145,6 +145,18 @@ status is `INVALID_ARGUMENT`, preserving the proxy's authentication policy. Gateway retry headers still take precedence on HTTP errors. See Google's [GenerateContent error reference](https://ai.google.dev/gemini-api/docs/generate-content/api-errors). +Gemini's native [`google.rpc.RetryInfo`](https://github.com/googleapis/googleapis/blob/master/google/rpc/error_details.proto) +cooldown is honored for HTTP and in-band SSE errors when no valid positive HTTP +`Retry-After` is present. Fractional seconds round up. A delay that cannot fit the +remaining local retry window is returned to the client without an early retry. + +Interrupted HTTP error-body reads are transport failures, not proof of an invalid +request. Incomplete or malformed Gemini error JSON likewise cannot establish a +validation error. Explicit retry headers remain authoritative unless the request +context itself has ended. +Proxy buffering deadlines remain timeout failures during HTTP/2 body reads; they +are not treated as client cancellations. + ### Codex context compaction A Responses `response.failed` event whose error code is exactly diff --git a/gemini.go b/gemini.go index 119abeb..4ab4a24 100644 --- a/gemini.go +++ b/gemini.go @@ -32,32 +32,75 @@ type geminiCandidateState struct { } func newGeminiWire(body []byte) *geminiWire { - g := &geminiWire{candidateCount: 1} - config := geminiRequestConfig{count: &g.candidateCount} + return &geminiWire{candidateCount: geminiSettingsOf(body).candidateCount} +} + +type geminiRequestSettings struct { + candidateCount int + thinkingLevel string +} + +func geminiSettingsOf(body []byte) geminiRequestSettings { + settings := geminiRequestSettings{candidateCount: 1} + config := geminiRequestConfig{settings: &settings} request := struct { Config geminiRequestConfig `json:"generationConfig"` ProtoConfig geminiRequestConfig `json:"generation_config"` }{config, config} - if json.Unmarshal(body, &request) != nil || g.candidateCount < 1 { - g.candidateCount = 1 + if json.Unmarshal(body, &request) != nil { + return geminiRequestSettings{candidateCount: 1} + } + if settings.candidateCount < 1 { + settings.candidateCount = 1 } - return g + return settings } -// Both ProtoJSON spellings update one count in JSON field order. A later config -// replaces the earlier config, including resetting an omitted count to default. -type geminiRequestConfig struct{ count *int } +// Both ProtoJSON spellings update one configuration in JSON field order. A later +// config replaces the earlier config, including resetting omitted settings. +type geminiRequestConfig struct{ settings *geminiRequestSettings } func (c geminiRequestConfig) UnmarshalJSON(data []byte) error { - *c.count = 1 - count := geminiRequestCount{value: c.count} + *c.settings = geminiRequestSettings{candidateCount: 1} + count := geminiRequestCount{value: &c.settings.candidateCount} + thinking := geminiThinkingConfig{level: &c.settings.thinkingLevel} + fields := struct { + Count geminiRequestCount `json:"candidateCount"` + ProtoCount geminiRequestCount `json:"candidate_count"` + Thinking geminiThinkingConfig `json:"thinkingConfig"` + ProtoThinking geminiThinkingConfig `json:"thinking_config"` + }{count, count, thinking, thinking} + return json.Unmarshal(data, &fields) +} + +type geminiThinkingConfig struct{ level *string } + +func (c geminiThinkingConfig) UnmarshalJSON(data []byte) error { + *c.level = "" + level := geminiThinkingLevel{value: c.level} fields := struct { - Count geminiRequestCount `json:"candidateCount"` - ProtoCount geminiRequestCount `json:"candidate_count"` - }{count, count} + Level geminiThinkingLevel `json:"thinkingLevel"` + ProtoLevel geminiThinkingLevel `json:"thinking_level"` + }{level, level} return json.Unmarshal(data, &fields) } +type geminiThinkingLevel struct{ value *string } + +func (l geminiThinkingLevel) UnmarshalJSON(data []byte) error { + // Diagnostic metadata must not invalidate otherwise usable candidate counts. + // Do not infer a level from a token budget or a model's implicit defaults. + var level string + _ = json.Unmarshal(data, &level) + switch level = strings.ToLower(strings.TrimSpace(level)); level { + case "minimal", "low", "medium", "high": + *l.value = level + default: + *l.value = "" + } + return nil +} + type geminiRequestCount struct{ value *int } func (c geminiRequestCount) UnmarshalJSON(data []byte) error { @@ -99,7 +142,10 @@ func (g *geminiWire) live(ev event, raw []byte, st *captureStats) ([][]byte, boo } func (g *geminiWire) completeOnEOF() bool { - return !g.failed && (g.blocked || g.finished == g.candidateCount) + if g.blocked { + return !g.failed && len(g.candidates) == 0 + } + return !g.failed && g.finished == g.candidateCount } func (g *geminiWire) replayBuffered(sp *spool, w io.Writer, st *captureStats) error { @@ -167,6 +213,10 @@ func (*geminiPart) UnmarshalJSON(data []byte) error { func (g *geminiWire) accept(ev event, st *captureStats) *failure { data := strings.TrimSpace(ev.data) + // SSE control frames and empty data messages carry no model response. + if ev.data == "" && ev.name != "error" { + return nil + } if data == "" || data[0] != '{' { return malformedSSE("Gemini stream event is not a JSON object") } diff --git a/gemini_errors.go b/gemini_errors.go index 2c59d2b..6c736df 100644 --- a/gemini_errors.go +++ b/gemini_errors.go @@ -4,6 +4,7 @@ import ( "bytes" "encoding/json" "net/http" + "strconv" "strings" ) @@ -14,8 +15,9 @@ type geminiAPIError struct { Status string `json:"status"` Message string `json:"message"` Details []struct { - Type string `json:"@type"` - Reason string `json:"reason"` + Type string `json:"@type"` + Reason string `json:"reason"` + RetryDelay json.RawMessage `json:"retryDelay"` } `json:"details,omitempty"` } @@ -24,11 +26,17 @@ func classifyGeminiHTTPErrorBytes(resp *http.Response, b []byte) failure { Error geminiAPIError `json:"error"` } data := bytes.TrimSpace(b) + var decodeErr error if len(data) > 0 && data[0] == '[' { // The non-SSE streaming transport wraps its HTTP error in an array. - _ = json.Unmarshal(data, &envelopes) + decodeErr = json.Unmarshal(data, &envelopes) } else { - _ = json.Unmarshal(data, &envelopes[0]) + decodeErr = json.Unmarshal(data, &envelopes[0]) + } + if decodeErr != nil { + // A malformed or classification-size-limited body cannot establish a + // deterministic validation error. Header retry verdicts still win. + return classifyHTTPErrorBytesShaped(resp, b, func(_, _, _, _ string) bool { return false }) } return classifyGeminiErrorBytes(resp, b, envelopes[0].Error) } @@ -67,9 +75,55 @@ func classifyGeminiErrorBytes(resp *http.Response, b []byte, e geminiAPIError) f } else if !f.transient { f.atype = geminiStatusForHTTP(f.status) } + if f.transient && f.retryAfter == 0 { + f.retryAfter = retryAfterSeconds(resp.Header.Get("Retry-After")) + if f.retryAfter == 0 { + for _, detail := range e.Details { + if detail.Type != "type.googleapis.com/google.rpc.RetryInfo" { + continue + } + var delay string + if json.Unmarshal(detail.RetryDelay, &delay) == nil { + f.retryAfter = max(f.retryAfter, geminiRetryDelaySeconds(delay)) + } + } + } + } return f } +// ProtoJSON Duration uses seconds with up to nine fractional digits. Round up +// to the HTTP Retry-After granularity so the client never retries early. +func geminiRetryDelaySeconds(value string) int { + text, ok := strings.CutSuffix(value, "s") + if !ok || text == "" { + return 0 + } + whole, fraction, fractional := strings.Cut(text, ".") + if whole == "" || (fractional && (fraction == "" || len(fraction) > 9)) { + return 0 + } + for _, part := range [...]string{whole, fraction} { + for i := range len(part) { + if part[i] < '0' || part[i] > '9' { + return 0 + } + } + } + seconds, err := strconv.ParseInt(whole, 10, 64) + if err != nil || seconds > 315576000000 { + return 0 + } + if strings.Trim(fraction, "0") != "" { + seconds++ + } + maxInt := int(^uint(0) >> 1) + if seconds > int64(maxInt) { + return maxInt + } + return int(seconds) +} + // classifyGeminiSSEError is called only after the wire detects an error frame; // a broken error payload is still a failure, never successful stream content. func classifyGeminiSSEError(data string) *failure { diff --git a/gemini_errors_test.go b/gemini_errors_test.go index c1dfedc..f16c1a8 100644 --- a/gemini_errors_test.go +++ b/gemini_errors_test.go @@ -120,7 +120,7 @@ func TestGeminiHTTPRetryHeaderPrecedence(t *testing.T) { {"upstream-force-permanent", 503, "", "false", false, 503, "upstream_no_retry", 0}, {"unknown-gateway-defers-to-upstream", 400, "unknown", "true", true, 502, "upstream_retry_400", 7}, } - body := []byte(`{"error":{"code":400,"status":"INVALID_ARGUMENT","message":"invalid contents"}}`) + body := []byte(`{"error":{"code":400,"status":"INVALID_ARGUMENT","message":"invalid contents","details":[{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"23s"}]}}`) for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { resp := &http.Response{StatusCode: tc.status, Header: make(http.Header)} @@ -164,6 +164,47 @@ func TestGeminiHTTPRetryAfter(t *testing.T) { } } +func TestGeminiRetryInfoBackoff(t *testing.T) { + for _, tc := range []struct { + name string + delay string + header string + want int64 + }{ + {"whole seconds", `"23s"`, "", 23}, + {"fractional seconds round up", `"23.000000001s"`, "", 24}, + {"subsecond cooldown", `"0.000000001s"`, "", 1}, + {"whole fractional form", `"23.000s"`, "", 23}, + {"protobuf maximum seconds", `"315576000000s"`, "", 315576000000}, + {"HTTP header takes precedence", `"23s"`, "7", 7}, + {"invalid header uses native delay", `"23s"`, "invalid", 23}, + {"zero cooldown", `"0s"`, "", 0}, + {"negative cooldown", `"-1s"`, "", 0}, + {"out of range duration", `"315576000001s"`, "", 0}, + {"excess fractional precision", `"0.0000000001s"`, "", 0}, + {"wrong units", `"1m"`, "", 0}, + {"compound duration", `"1s2s"`, "", 0}, + {"invalid field type", `23`, "", 0}, + } { + t.Run(tc.name, func(t *testing.T) { + body := `{"error":{"code":429,"status":"RESOURCE_EXHAUSTED","message":"quota exhausted","details":[{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":` + tc.delay + `}]}}` + resp := &http.Response{StatusCode: 429, Header: make(http.Header)} + resp.Header.Set("Retry-After", tc.header) + f := classifyGeminiHTTPErrorBytes(resp, []byte(body)) + want := min(tc.want, int64(int(^uint(0)>>1))) + if !f.transient || int64(f.retryAfter) != want { + t.Fatalf("HTTP cooldown=%d, want %d: %+v", f.retryAfter, want, f) + } + if tc.header == "" { + f := classifyGeminiSSEError(body) + if f == nil || !f.transient || int64(f.retryAfter) != want { + t.Fatalf("SSE cooldown must preserve native RetryInfo: %+v, want %d", f, want) + } + } + }) + } +} + func TestClassifyGeminiSSEError(t *testing.T) { cases := []struct { name string @@ -243,6 +284,15 @@ func TestGeminiMalformedErrors(t *testing.T) { } } +func TestGeminiIncompleteHTTPErrorCannotEstablishValidationFailure(t *testing.T) { + for _, body := range []string{`{"error":{"code":400,`, `{"error":{"status":"INVALID_ARGUMENT","message":`, "not JSON"} { + f := classifyGeminiHTTPErrorBytes(&http.Response{StatusCode: 400}, []byte(body)) + if !f.transient || f.status != 502 { + t.Fatalf("incomplete error metadata became terminal: %+v", f) + } + } +} + func TestWriteGeminiError(t *testing.T) { cases := []struct { name string diff --git a/gemini_integration_test.go b/gemini_integration_test.go index 4ac627e..f29c4ce 100644 --- a/gemini_integration_test.go +++ b/gemini_integration_test.go @@ -5,6 +5,7 @@ import ( "io" "net/http" "net/http/httptest" + "strconv" "strings" "sync/atomic" "testing" @@ -26,6 +27,9 @@ func TestGeminiHiddenRetryDiscardsFailedOutput(t *testing.T) { })) defer up.Close() setupForTest(up.URL) + dir := t.TempDir() + cfg.requestLogDir = dir + t.Cleanup(func() { cfg.requestLogDir = "" }) cfg.txLocalRetries = 1 cfg.localBackoffCap = time.Millisecond // These Anthropic-only signals must not change Gemini's window or model. @@ -42,6 +46,19 @@ func TestGeminiHiddenRetryDiscardsFailedOutput(t *testing.T) { if rec.Header().Get("X-Steady-Proxy-Mode") != "buffered" { t.Fatalf("Anthropic marker changed Gemini buffering: %v", rec.Header()) } + files := archiveFiles(t, dir) + if len(files) != 1 { + t.Fatalf("want one request archive, got %d", len(files)) + } + archive, _ := readArchive(t, files[0]) + if len(archive.Attempts) != 2 { + t.Fatalf("archive lost retry history: %+v", archive.Attempts) + } + for _, attempt := range archive.Attempts { + if attempt.Model != "gemini-test" { + t.Fatalf("attempt model must come from URL, not JSON body: %+v", attempt) + } + } } func TestGeminiInvalidArgumentDoesNotRetry(t *testing.T) { @@ -146,3 +163,94 @@ func TestGeminiCommittedTruncationAbortsHTTP(t *testing.T) { t.Fatalf("committed truncation must fail transport without replay: body=%q err=%v calls=%d", body, err, calls.Load()) } } + +func TestGeminiTruncatedHTTPErrorCanRecover(t *testing.T) { + for _, archived := range []bool{false, true} { + t.Run(strconv.FormatBool(archived), func(t *testing.T) { + var calls atomic.Int32 + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if calls.Add(1) == 1 { + w.Header().Set("Content-Length", "1024") + w.WriteHeader(400) + io.WriteString(w, `{"error":{"code":400,`) + return + } + w.Header().Set("Content-Type", "text/event-stream") + io.WriteString(w, goodGeminiStream) + })) + defer up.Close() + setupForTest(up.URL) + cfg.txLocalRetries = 1 + cfg.localBackoffCap = 0 + if archived { + cfg.requestLogDir = t.TempDir() + t.Cleanup(func() { cfg.requestLogDir = "" }) + } + rec := httptest.NewRecorder() + handle(rec, httptest.NewRequest(http.MethodPost, geminiStreamPath, strings.NewReader(`{}`))) + if rec.Code != 200 || calls.Load() != 2 || !strings.Contains(rec.Body.String(), `"text":"Hello"`) || strings.Contains(rec.Body.String(), `"error"`) { + t.Fatalf("truncated error suppressed recovery: status=%d calls=%d body=%s", rec.Code, calls.Load(), rec.Body) + } + }) + } +} + +func TestGeminiTruncatedUnaryErrorHonorsRetryHints(t *testing.T) { + for _, hint := range []string{"", "false", "true"} { + t.Run("hint="+hint, func(t *testing.T) { + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("X-Gateway-Retryable", hint) + w.Header().Set("Retry-After", "17") + w.Header().Set("Content-Length", "1024") + w.WriteHeader(400) + io.WriteString(w, `{"error":{"code":400,`) + })) + defer up.Close() + setupForTest(up.URL) + rec := httptest.NewRecorder() + handle(rec, httptest.NewRequest(http.MethodPost, "/v1/models/gemini-test:generateContent", strings.NewReader(`{}`))) + if hint == "false" { + if rec.Code != 400 || rec.Header().Get("X-Should-Retry") != "false" || rec.Header().Get("Retry-After") != "" { + t.Fatalf("explicit permanent hint lost: status=%d headers=%v", rec.Code, rec.Header()) + } + } else if rec.Code != 503 || rec.Header().Get("X-Should-Retry") != "true" || rec.Header().Get("Retry-After") != "17" { + t.Fatalf("read failure lost retry eligibility or cooldown: status=%d headers=%v", rec.Code, rec.Header()) + } + }) + } +} + +func TestGeminiRetryInfoPreventsPrematureLocalRetry(t *testing.T) { + for _, sse := range []bool{false, true} { + for _, delay := range []string{"23", "315576000000"} { + t.Run(strconv.FormatBool(sse)+"/"+delay, func(t *testing.T) { + var calls atomic.Int32 + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + body := `{"error":{"code":429,"status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"` + delay + `s"}]}}` + if sse { + w.Header().Set("Content-Type", "text/event-stream") + io.WriteString(w, "data: "+body+"\n\n") + } else { + w.WriteHeader(429) + io.WriteString(w, body) + } + })) + defer up.Close() + setupForTest(up.URL) + cfg.txLocalRetries = 1 + cfg.localBackoffCap = 0 + cfg.geminiBufferMs = time.Second + rec := httptest.NewRecorder() + handle(rec, httptest.NewRequest(http.MethodPost, geminiStreamPath, strings.NewReader(`{}`))) + wantDelay := delay + if strconv.IntSize == 32 && delay == "315576000000" { + wantDelay = "2147483647" + } + if rec.Code != 503 || calls.Load() != 1 || rec.Header().Get("Retry-After") != wantDelay { + t.Fatalf("native cooldown was shortened or bypassed: status=%d calls=%d headers=%v", rec.Code, calls.Load(), rec.Header()) + } + }) + } + } +} diff --git a/gemini_test.go b/gemini_test.go index a937263..6096db6 100644 --- a/gemini_test.go +++ b/gemini_test.go @@ -99,6 +99,29 @@ func TestGeminiCaptureBufferedWaitsForEOFAndTrailingUsage(t *testing.T) { } } +func TestGeminiControlFramesDoNotFailOrCompleteResponse(t *testing.T) { + controls := ": keepalive\n\nid: 123\nretry: 1000\n\nevent: ping\n\ndata:\n\n" + for _, early := range []bool{false, true} { + t.Run(strconv.FormatBool(early), func(t *testing.T) { + for _, complete := range []bool{false, true} { + stream := controls + if complete { + stream += goodGeminiStream + controls + } + var st captureStats + rec, wrote, f := captureGemini(t, nil, iotest1byte(stream), &st, early) + if complete { + if !wrote || f != nil || st.stop != "STOP" || st.outTok != 11 || !strings.Contains(rec.Body.String(), `"text":"Hello"`) { + t.Fatalf("control frames broke a complete response: wrote=%v failure=%+v stats=%+v body=%q", wrote, f, st, rec.Body.String()) + } + } else if wrote || f == nil || !f.transient || rec.Body.Len() != 0 { + t.Fatalf("control frames completed an empty response: wrote=%v failure=%+v body=%q", wrote, f, rec.Body.String()) + } + } + }) + } +} + func TestGeminiWireCompletionTracksEveryRequestedCandidate(t *testing.T) { for _, tc := range []struct { name string @@ -115,6 +138,7 @@ func TestGeminiWireCompletionTracksEveryRequestedCandidate(t *testing.T) { {"omitted default index", `{}`, []string{`{"candidates":[{"finishReason":"STOP"}]}`}, true}, {"explicit count one", `{"generationConfig":{"candidateCount":1}}`, []string{`{"candidates":[{"index":0,"finishReason":"STOP"}]}`}, true}, {"two requested one finished", `{"generationConfig":{"candidateCount":2}}`, []string{`{"candidates":[{"index":0,"finishReason":"STOP"}]}`}, false}, + {"thinking enum encoding cannot mask missing candidate", `{"generationConfig":{"candidateCount":2,"thinkingConfig":{"thinkingLevel":2}}}`, []string{`{"candidates":[{"index":0,"finishReason":"STOP"}]}`}, false}, {"only second finished", `{"generationConfig":{"candidateCount":2}}`, []string{`{"candidates":[{"index":1,"finishReason":"STOP"}]}`}, false}, {"repeat finish cannot replace missing candidate", `{"generationConfig":{"candidateCount":2}}`, []string{`{"candidates":[{"index":0,"finishReason":"STOP"}]}`, `{"candidates":[{"index":0,"finishReason":"STOP"}]}`}, false}, {"second remains unfinished", `{"generationConfig":{"candidateCount":2}}`, []string{`{"candidates":[{"index":0,"finishReason":"STOP"},{"index":1,"finishReason":"FINISH_REASON_UNSPECIFIED"}]}`}, false}, @@ -204,6 +228,19 @@ func TestGeminiPromptBlocksCompleteWithoutCandidates(t *testing.T) { } } +func TestGeminiPromptBlocksCannotHideCandidates(t *testing.T) { + block := string(geminiTestEvent(`{"promptFeedback":{"blockReason":"SAFETY"}}`).raw) + for _, finish := range []string{"", "STOP"} { + candidate := string(geminiTestEvent(`{"candidates":[{"content":{"parts":[{"text":"candidate"}]},"finishReason":` + strconv.Quote(finish) + `}]}`).raw) + for _, stream := range []string{block + candidate, candidate + block} { + rec, wrote, f := captureGemini(t, nil, strings.NewReader(stream), nil, false) + if wrote || f == nil || !f.transient || rec.Body.Len() != 0 { + t.Fatalf("contradictory prompt block accepted: wrote=%v failure=%+v body=%q", wrote, f, rec.Body.String()) + } + } + } +} + func TestGeminiUsageSnapshotsDoNotDoubleCount(t *testing.T) { wire := newGeminiWire(nil) var st captureStats @@ -243,7 +280,6 @@ func TestGeminiMalformedEventsFailEvenWhenValidationDisabled(t *testing.T) { saved := cfg.validateJSON t.Cleanup(func() { cfg.validateJSON = saved }) cases := map[string]string{ - "empty": "", "broken JSON": `{"candidates":[`, "null": `null`, "array": `[]`, @@ -320,6 +356,7 @@ func TestGeminiInBandErrorsFailBeforeAndAfterFinish(t *testing.T) { {"named error", `{"error":{"code":503,"status":"UNAVAILABLE","message":"overloaded"}}`, "error", true}, {"null error envelope", `{"error":null}`, "", true}, {"missing named error envelope", `{}`, "error", true}, + {"empty named error", ``, "error", true}, } { t.Run(tc.name, func(t *testing.T) { for _, afterFinish := range []bool{false, true} { @@ -413,3 +450,13 @@ func TestGeminiCaptureLiveRejectsErrorAfterFinish(t *testing.T) { }) } } + +func TestGeminiLongRetryInfoCannotFitRequestDeadline(t *testing.T) { + useDefaultConfig(t) + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + f := classifyGeminiSSEError(`{"error":{"code":429,"details":[{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"315576000000s"}]}}`) + if f == nil || canWaitForLocalRetry(ctx, localRetryDelay(f.retryAfter, 0)) { + t.Fatalf("native cooldown exceeded the request deadline but was accepted: %+v", f) + } +} diff --git a/main.go b/main.go index f47c4eb..8694582 100644 --- a/main.go +++ b/main.go @@ -414,8 +414,8 @@ func handle(w http.ResponseWriter, r *http.Request) { } if resp.StatusCode < 200 || resp.StatusCode >= 300 { - respBody := readHTTPErrorBody(resp, rec) - last = classifyHTTPErrorForPath(r.URL.Path, resp, respBody) + respBody, readErr := readHTTPErrorBody(resp, rec) + last = classifyHTTPErrorForPath(ctx, r.URL.Path, resp, respBody, readErr) last.origStatus = resp.StatusCode rec.noteAttemptFailure(last, "http_error") resp.Body.Close() @@ -570,19 +570,18 @@ func (p *prefixCapture) Write(b []byte) (int, error) { return len(b), nil } -func readHTTPErrorBody(resp *http.Response, rec *reqRecorder) []byte { +func readHTTPErrorBody(resp *http.Response, rec *reqRecorder) ([]byte, error) { body := withByteIdle(resp.Body, cfg.upstreamByteIdle) defer body.stop() if rec == nil { - b, _ := io.ReadAll(io.LimitReader(body, httpErrorClassifyBytes)) - return b + return io.ReadAll(io.LimitReader(body, httpErrorClassifyBytes)) } rec.noteAttemptResponseHeaders(resp.StatusCode, resp.Header) prefix := &prefixCapture{max: httpErrorClassifyBytes} dst := io.MultiWriter(prefix, rec.respWriter()) _, err := io.Copy(dst, body) rec.noteResponseReadError(err) - return prefix.buf.Bytes() + return prefix.buf.Bytes(), err } func copyResponseBodyToArchive(resp *http.Response, rec *reqRecorder) { @@ -621,8 +620,8 @@ func proxyOnce(w http.ResponseWriter, r *http.Request, body []byte, rec *reqReco f = classifyTransport(err, ctx) rec.noteAttemptFailure(f, "transport_error") case resp.StatusCode < 200 || resp.StatusCode >= 300: - respBody := readHTTPErrorBody(resp, rec) - f = classifyHTTPErrorForPath(r.URL.Path, resp, respBody) + respBody, readErr := readHTTPErrorBody(resp, rec) + f = classifyHTTPErrorForPath(ctx, r.URL.Path, resp, respBody, readErr) f.origStatus = resp.StatusCode rec.noteAttemptFailure(f, "http_error") resp.Body.Close() @@ -781,11 +780,14 @@ func isEventStream(ct string) bool { return strings.Contains(strings.ToLower(ct) // The method/path/gateway are static for the transactional route, so omitted. func who(r *http.Request, body []byte) string { - model, _ := geminiTarget(r.URL.Path) - if model == "" { - model = modelOf(body) + return requestModelOf(r.URL.Path, body) + "/" + agentKindFull(r, body) +} + +func requestModelOf(path string, body []byte) string { + if model, _ := geminiTarget(path); model != "" { + return model } - return model + "/" + agentKindFull(r, body) + return modelOf(body) } // agentKindFull refines agentKind with a body-derived signal: a Workflow-tool @@ -889,7 +891,12 @@ func localRetryDelay(upstreamRetryAfter int, retryIndex int) time.Duration { if upstreamRetryAfter < 0 { upstreamRetryAfter = 0 } - return time.Duration(upstreamRetryAfter)*time.Second + localRetryExtraBackoff(retryIndex) + extra := localRetryExtraBackoff(retryIndex) + maxDelay := time.Duration(math.MaxInt64) + if int64(upstreamRetryAfter) > int64((maxDelay-extra)/time.Second) { + return maxDelay + } + return time.Duration(upstreamRetryAfter)*time.Second + extra } func localRetryExtraBackoff(retryIndex int) time.Duration { @@ -917,8 +924,11 @@ func canWaitForLocalRetry(ctx context.Context, wait time.Duration) bool { if wait <= 0 { return true } - if deadline, ok := ctx.Deadline(); ok && time.Until(deadline) <= wait+time.Second { - return false + if deadline, ok := ctx.Deadline(); ok { + remaining := time.Until(deadline) + if remaining <= time.Second || wait >= remaining-time.Second { + return false + } } return true } @@ -949,6 +959,9 @@ func modelOf(body []byte) string { } func reasoningEffortOf(path string, body []byte) string { + if isGeminiPath(path) { + return geminiSettingsOf(body).thinkingLevel + } var b struct { Reasoning struct { Effort string `json:"effort"` diff --git a/reliability_test.go b/reliability_test.go index fb0810b..f63c44c 100644 --- a/reliability_test.go +++ b/reliability_test.go @@ -381,3 +381,36 @@ func TestErrorArchiveObeysByteIdle(t *testing.T) { }) } } + +func TestHTTP2ErrorBodyDeadlineRemainsRetryable(t *testing.T) { + for _, path := range []string{geminiStreamPath, "/v1/responses", "/v1/messages"} { + t.Run(path, func(t *testing.T) { + up := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.ProtoMajor != 2 { + t.Errorf("expected HTTP/2 upstream, got %s", r.Proto) + } + w.WriteHeader(503) + io.WriteString(w, `{"error":{"message":"unavailable",`) + w.(http.Flusher).Flush() + <-r.Context().Done() + })) + up.EnableHTTP2 = true + up.StartTLS() + defer up.Close() + setupForTest(up.URL) + client = up.Client() + client.CheckRedirect = stopRedirect + cfg.geminiBufferMs = 100 * time.Millisecond + cfg.responsesEarlyCommit = false + cfg.responsesBufferMs = 100 * time.Millisecond + cfg.keepaliveMs = 100 * time.Millisecond + cfg.upstreamByteIdle = 2 * time.Second + cfg.maxRequestDur = time.Second + rec := httptest.NewRecorder() + handle(rec, httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"stream":true}`))) + if rec.Code != 503 || rec.Header().Get("X-Should-Retry") != "true" || rec.Header().Get("X-Steady-Proxy-Reason") != "deadline" { + t.Fatalf("attempt timeout became terminal client cancellation: status=%d headers=%v", rec.Code, rec.Header()) + } + }) + } +} diff --git a/requestlog.go b/requestlog.go index 24ab4ff..f87bbab 100644 --- a/requestlog.go +++ b/requestlog.go @@ -296,7 +296,7 @@ func (rec *reqRecorder) beginAttempt(body []byte, reqWho string, proxyRetries in StartedAt: rec.whenString(started), started: started, Who: reqWho, - Model: modelOf(body), + Model: requestModelOf(rec.path, body), ProxyRetry: proxyRetries, RequestBodyRef: ref, RequestBodySize: len(body), diff --git a/responses.go b/responses.go index 17d5ae6..26cf5fa 100644 --- a/responses.go +++ b/responses.go @@ -2,6 +2,7 @@ package main import ( "bufio" + "context" "encoding/json" "io" "net/http" @@ -1027,8 +1028,21 @@ func responsesRequestShaped(atype, code, msg, _ string) bool { return !classifyResponsesError(atype, code, msg).transient } -// classifyHTTPErrorForPath dispatches HTTP-error classification by wire. -func classifyHTTPErrorForPath(path string, resp *http.Response, b []byte) failure { +// classifyHTTPErrorForPath applies read failures before body heuristics, retaining +// explicit retry hints unless the request context itself has ended. +func classifyHTTPErrorForPath(ctx context.Context, path string, resp *http.Response, b []byte, readErr error) failure { + if readErr != nil { + if source, _ := httpRetryHint(resp.Header); source == "" || ctx.Err() != nil { + f := classifyTransport(readErr, ctx) + if f.transient { + f.retryAfter = retryAfterSeconds(resp.Header.Get("Retry-After")) + } + return f + } + // An authoritative header owns the verdict; the incomplete body cannot + // supply trustworthy error metadata. + b = nil + } if isResponsesPath(path) { return classifyResponsesHTTPErrorBytes(resp, b) } diff --git a/transport.go b/transport.go index 9e4f2e9..2a09399 100644 --- a/transport.go +++ b/transport.go @@ -46,15 +46,26 @@ func roundTripBefore(ctx context.Context, r *http.Request, body []byte, deadline cleanup() return nil, started, stop, err } - resp.Body = &cancelBody{ReadCloser: resp.Body, cancel: cleanup} + resp.Body = &cancelBody{ReadCloser: resp.Body, ctx: attemptCtx, cancel: cleanup} return resp, started, stop, nil } type cancelBody struct { io.ReadCloser + ctx context.Context cancel func() } +func (b *cancelBody) Read(p []byte) (int, error) { + n, err := b.ReadCloser.Read(p) + if err != nil && err != io.EOF { + if cause := context.Cause(b.ctx); cause != nil { + err = errors.Join(err, cause) + } + } + return n, err +} + func (b *cancelBody) Close() error { b.cancel() return b.ReadCloser.Close()