From c1f71574a792095a87e06dfb906d9078500dc907 Mon Sep 17 00:00:00 2001 From: "Shangxiang Fan (from Dev Box)" Date: Mon, 27 Jul 2026 19:09:25 +0800 Subject: [PATCH] =?UTF-8?q?release:=20v1.3.0=20=E2=80=94=20streaming=20int?= =?UTF-8?q?egrity,=20Claude=20Code=20compatibility,=20health/model=20APIs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes auto-upgrade, which panicked on startup, and several defects that could silently corrupt or truncate model output. Streaming integrity - Decode SSE by buffering raw bytes and splitting complete lines only. Chunks do not respect UTF-8 boundaries, so per-chunk from_utf8_lossy turned any multi-byte character split across chunks into U+FFFD; the damaged payload then failed to parse as JSON and was dropped, deleting text mid-response and replaying the corruption into the next request as history. - Detect interrupted streams and emit a protocol-appropriate terminator (OpenAI error chunk + [DONE], Anthropic/Responses `event: error`, Gemini finishReason OTHER), recording the request as 502 instead of 200. - Close Anthropic streams that end without a finish_reason so clients neither block on a missing message_stop nor keep a partial turn as complete. - Flush a trailing event that arrives without a newline; forward unparsable data payloads instead of dropping them. - Emit a `: keepalive` comment after 15s of silence, at event boundaries only, so extended thinking does not trip the ~60s upstream idle timeout. Auto-upgrade - Run self_update on a blocking thread. It owns a private Tokio runtime, and dropping that from the async main panicked the process, which made auto_upgrade: true unusable. - Fail the release workflow when the tag does not match Cargo.toml. v1.2.3 shipped as 1.2.2, so its binaries re-downloaded the same release forever. Claude Code / MCP compatibility - Forward context_management plus the derived context-management beta; it was stripped, silently disabling context editing. - Merge the client's anthropic-beta header instead of overwriting it. - Normalize tool_result content arrays; MCP tools returning blocks rather than a bare string were rejected upstream. - Detect images nested in tool_result for the vision header, and translate url-source images. - Fill missing max_tokens from the model catalog; retry once with max_completion_tokens for models that require it. - Count per-message and tool-schema overhead in the local token estimate. - Serve /v1/messages/count_tokens from upstream when available, else a local estimate, instead of always failing with 400. Other fixes - Stop scraping the AUR for the VS Code version; use Microsoft's release API. - Price requests by the served model, and unshadow the gpt-4o rate. - Clamp pagination and use saturating arithmetic; aggregate metrics and audit queries in place rather than cloning the whole request store. - Add a connect timeout, rotate error.log, and drop a duplicate startup model fetch. Added - GET /health, GET /v1/models/{model}, graceful shutdown, ghc_proxy_uptime_seconds, and a model= filter on /api/audit. --- .github/workflows/release.yml | 26 + CHANGELOG.md | 149 ++++ Cargo.lock | 2 +- Cargo.toml | 2 +- README.md | 4 +- docs/api.md | 107 ++- src/anthropic.rs | 539 +++++++++++++- src/config.rs | 2 +- src/main.rs | 70 +- src/server.rs | 1307 ++++++++++++++++++++++++--------- src/state.rs | 69 ++ src/store.rs | 120 +++ src/util.rs | 212 +++++- 13 files changed, 2193 insertions(+), 416 deletions(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index c3bc025..f03798b 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -19,8 +19,34 @@ env: CARGO_TERM_COLOR: always jobs: + verify: + name: Verify tag matches Cargo.toml + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v7 + with: + ref: ${{ github.event.inputs.tag || github.ref }} + + # A release whose binary reports a different version than its tag makes + # the built-in auto-upgrade re-download the same release on every start, + # because `current_version` never catches up with the tag. Fail loudly + # instead of publishing that. + - name: Compare tag with package version + shell: bash + run: | + tag="${{ github.event.inputs.tag || github.ref_name }}" + crate_version=$(grep -m1 '^version = ' Cargo.toml | cut -d'"' -f2) + echo "tag=$tag crate_version=$crate_version" + if [ "$tag" != "v$crate_version" ]; then + echo "::error::Tag $tag does not match Cargo.toml version $crate_version." \ + "Bump Cargo.toml (and Cargo.lock) before tagging." + exit 1 + fi + build: name: Build ${{ matrix.target }} + needs: verify runs-on: ${{ matrix.os }} strategy: fail-fast: false diff --git a/CHANGELOG.md b/CHANGELOG.md index d00ac0b..b45d401 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,155 @@ All notable changes to this project will be documented in this file. ## [Unreleased] +## [1.3.0] - 2026-07-27 + +### Added +- `GET /health` liveness/readiness probe reporting version, uptime, Copilot + token status and remaining lifetime, loaded model count, requests served, and + whether API-key auth is enabled. Answers without contacting upstream and is + never guarded by the API key; `?strict=true` returns `503` when not ready +- `GET /v1/models/{model}` (and `/models/{model}`) OpenAI-compatible single + model retrieval, including `capabilities` and `supported_endpoints`. Model + aliases from `model_mappings` are resolved; unknown ids return `404` +- Graceful shutdown on Ctrl-C and `SIGTERM`, draining in-flight requests and SSE + streams instead of dropping them +- `ghc_proxy_uptime_seconds` gauge on `/metrics` +- `model=` filter on `GET /api/audit`, matching the requested or translated model +- `/health` and `/v1/models/{model}` documented in `openapi.json` +- **SSE keepalive.** Streams emit a `: keepalive` comment after 15s of silence + so an idle connection is not dropped at the ~60s timeout GitHub's upstream + load balancer and most intermediaries enforce. Extended thinking can easily + exceed that with no tokens emitted, which previously surfaced as + `user_request_timeout` or a stream that died mid-answer. The comment is only + injected at an event boundary, so the verbatim passthrough paths are never + spliced mid-event +- **`anthropic-beta` request headers are forwarded.** The client's flags are + preserved and merged with the ones the proxy derives, instead of being + overwritten by the 1M-context flag +- Missing `max_tokens` on `/v1/chat/completions` is filled from the model + catalog's `capabilities.limits.max_output_tokens`, which several Copilot + models reject the request without +- Requesting a `/responses`-only model on `/v1/chat/completions` now returns an + actionable `unsupported_api_for_model` error naming the right endpoint, + instead of an opaque upstream 400 + +### Fixed +- **`context_management` was silently stripped, disabling Claude Code's context + editing.** The field was not in the Anthropic passthrough allowlist, and the + upstream rejects it with a misleading `Extra inputs are not permitted` 400 + unless the `context-management-2025-06-27` beta is requested. It is now + forwarded and the matching beta flag is derived automatically +- **`tool_result` content arrays were forwarded raw, breaking every MCP tool + that returns anything but a bare string.** Anthropic allows `content` to be an + array of blocks, which MCP servers routinely use; passing it through made the + upstream reject the whole request with + `type has to be either 'image_url' or 'text'`. Text blocks are now joined and + image blocks rewritten as `image_url` data URLs +- **Images nested inside a `tool_result` did not set `Copilot-Vision-Request`.** + `has_image` only inspected top-level content blocks, so screenshots returned + by MCP tools were sent as non-vision requests +- `image` blocks with a `url` source (rather than `base64`) are now translated + instead of being turned into an empty data URL +- Models that replaced `max_tokens` with `max_completion_tokens` (`gpt-5.3-codex` + and friends) previously failed outright; the parameter is renamed and the + request retried once, on both the streaming and non-streaming paths +- **Stopped scraping the Arch User Repository for the latest VS Code version.** + The AUR maintainers asked proxies of this kind to stop, having become the + single most-requested endpoint on their service. `dynamic_vscode_version` now + uses Microsoft's own `update.code.visualstudio.com` release API and ignores + non-`major.minor.patch` builds +- The local `/v1/messages/count_tokens` estimate ignored per-message framing and + tool schemas, under-reporting badly enough that Claude Code compacted too late + and then hit a hard `prompt token count exceeds the limit` failure. Tool + definitions are now counted as serialized JSON, plus per-message and + per-request overhead +- **`/v1/messages/count_tokens` always failed with HTTP 400.** The handler only + tried the upstream when the catalog advertised `/v1/messages/count_tokens`, + which Copilot never does, and then returned an error. It now forwards to the + upstream for any model exposing the native `/v1/messages` surface (returning + exact counts) and falls back to a local tiktoken estimate — marked + `"estimated": true` — instead of erroring. This unblocks Claude Code, which + calls the endpoint before every request +- **Streaming `/v1/responses` returned upstream errors as a `200` SSE body.** + A 401/429/5xx from upstream was wrapped in a success stream, so clients saw a + broken response instead of the real failure. Non-2xx upstream responses are + now surfaced with their status code, matching the other streaming paths +- **Cost estimates were wrong for mapped models and for `gpt-4o`.** Costs were + computed from the client-supplied model name rather than the model actually + served, so an alias such as `opus` priced at the fallback rate (~38x too low). + The `gpt-4o` rate was also unreachable because the broader `gpt-4` arm matched + first. Rates now key off the translated model, order specific families before + general ones, strip `publisher/` prefixes for GitHub Models ids, and cover the + gpt-4.1/gpt-5/o-series/Gemini families +- Dashboard pagination could overflow `usize` (panicking the handler) and an + unbounded `per_page` forced a full clone of the request store; `page`/`per_page` + are now parsed with saturating arithmetic and clamped to 500 +- `error.log` grew without limit; it is now rotated to `error.log.1` past 8 MB +- The upstream HTTP client had no connect timeout, so a dead upstream could wedge + a request indefinitely (30s connect timeout, 90s pool idle timeout; no overall + request timeout so SSE streams are unaffected) +- The model catalog was fetched twice at startup because the periodic refresh + timer fired immediately on its first tick +- `/metrics` and `/api/audit*` cloned every retained record (including captured + request/response bodies in debug mode) on each call; aggregation and filtering + now happen in place under the store lock +- Removed a dead `is_github_models` binding in the translated Anthropic path + +#### Streaming integrity +- **Streaming responses could be silently corrupted mid-text.** Every SSE chunk + was decoded with `String::from_utf8_lossy` as it arrived, but upstream chunks + do not respect UTF-8 character boundaries. Any multi-byte character (CJK, + emoji, `—`, box drawing) split across two chunks became `U+FFFD`, and because + the damaged payload then failed to parse as JSON it was **dropped entirely**, + deleting text from the middle of a response with no error. The corrupted or + missing text was then replayed into the next request as conversation history. + SSE parsing now buffers raw bytes and only decodes complete lines + (`util::SseLineBuffer`), matching the `TextDecoderStream` + `TextLineStream` + pipeline the TypeScript `copilot-api` proxy gets from its platform +- **Truncated streams were delivered as if they were complete.** A transport + error mid-stream only did `break`, so the client received a partial answer + with no terminator and the request was still recorded as `200`. All five + streaming paths now detect the interruption, emit a protocol-appropriate + terminator, and record the request as `502`: + - OpenAI: `data: {"error": …, "code": "stream_interrupted"}` followed by `[DONE]` + - Anthropic: `event: error` with an `api_error` payload + - Responses: `event: error` with `code: "stream_interrupted"` + - Gemini: final chunk with `finishReason: "OTHER"` +- **Anthropic streams that ended without a `finish_reason` never sent + `message_stop`.** Claude Code either blocks waiting for it or keeps the + partial text as a finished assistant turn. `AnthropicStreamState::finish()` + now closes any open content block and emits `message_delta` + `message_stop` +- A final SSE event arriving without a trailing newline was discarded; the line + buffer is now flushed when the stream ends +- `/v1/chat/completions` streaming dropped any `data:` payload that failed to + parse as JSON instead of forwarding it +- `/v1/responses` streaming now also flags a stream that ends before + `response.completed` +- Direct-Anthropic streaming extracted `stop_reason` and `tools_called` by + re-scanning the captured response body, so those audit fields were only ever + populated in debug mode; they are now collected while the stream is parsed +- SSE `data:` parsing accepts `data:{…}` (no space) per the SSE spec, not only + `data: {…}` + +### Changed +- `/api/audit/summary` now rounds cost figures to 4 decimals so small per-request + costs are no longer reported as `0.0` +- Default `vscode_version` bumped to `1.130.0` +- The release workflow now fails if the tag does not match the `Cargo.toml` + version. `v1.2.3` shipped with `version = "1.2.2"`, so binaries from that + release report `1.2.2` and `auto_upgrade` re-downloaded it on every start + +## [1.2.3] - 2026-07-20 + +### Fixed +- `--setup --claudecode` no longer overwrites unrelated Claude Code settings + ([#24](https://github.com/MartinForReal/ghc-proxy/pull/24)) + +### Known issue +- Released with `Cargo.toml` still at `1.2.2`, so binaries from this release + identify themselves as `1.2.2`. With `auto_upgrade` enabled they re-download + this release on every start. Upgrade to 1.3.0 to clear it. + ## [1.2.2] - 2026-07-14 ### Added diff --git a/Cargo.lock b/Cargo.lock index bfe33b8..5b59bdb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -743,7 +743,7 @@ dependencies = [ [[package]] name = "ghc-proxy" -version = "1.2.2" +version = "1.3.0" dependencies = [ "async-stream", "axum", diff --git a/Cargo.toml b/Cargo.toml index 0dc1762..a3ece90 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "ghc-proxy" -version = "1.2.2" +version = "1.3.0" edition = "2021" description = "GitHub Copilot API Proxy - Provides OpenAI and Anthropic compatible endpoints via GitHub Copilot (Rust port of ghc-tunnel)" license = "MIT" diff --git a/README.md b/README.md index 0e6e9b2..f100440 100644 --- a/README.md +++ b/README.md @@ -240,11 +240,13 @@ the dashboard and model listings. | `POST /v1/chat/completions` | OpenAI chat completions | | `POST /v1/responses` | OpenAI responses API (Codex) | | `GET /v1/models` | List available models | +| `GET /v1/models/{model}` | Retrieve a single model (aliases resolved) | | `POST /v1/messages` | Anthropic messages API | -| `POST /v1/messages/count_tokens` | Anthropic token counting | +| `POST /v1/messages/count_tokens` | Anthropic token counting (local estimate fallback) | | `POST /v1beta/models/{model}:generateContent` | Gemini generate content | | `POST /v1beta/models/{model}:streamGenerateContent` | Gemini streaming (SSE) | | `POST /v1beta/models/{model}:countTokens` | Gemini token counting | +| `GET /health` | Liveness/readiness probe (`?strict=true` for 503 when not ready) | | `GET /openapi.json` | OpenAPI v3 specification | | `GET /` | Web dashboard | | `GET /metrics/dashboard` | Metrics dashboard UI | diff --git a/docs/api.md b/docs/api.md index 93c6136..b89baea 100644 --- a/docs/api.md +++ b/docs/api.md @@ -22,14 +22,16 @@ require a key on the LLM endpoints; see [Authentication](#authentication) below. | `POST /v1/chat/completions` | OpenAI chat completions (also `/chat/completions`) | | `POST /v1/responses` | OpenAI Responses API for Codex (also `/responses`) | | `POST /v1/messages` | Anthropic Messages API | -| `POST /v1/messages/count_tokens` | Anthropic token counting (real BPE) | +| `POST /v1/messages/count_tokens` | Anthropic token counting (real BPE, local estimate fallback) | | `POST /v1beta/models/{model}:generateContent` | Gemini generate content | | `POST /v1beta/models/{model}:streamGenerateContent` | Gemini streaming (SSE) | | `POST /v1beta/models/{model}:countTokens` | Gemini token counting | | `POST /v1/embeddings` | Embeddings (also `/embeddings`) | | `GET /v1/models` | List available models (also `/models`, `/api/models`) | +| `GET /v1/models/{model}` | Retrieve a single model (also `/models/{model}`) | | `GET /v1/models/full/` | Raw upstream model catalog with capabilities | | `GET /usage` | Copilot plan and quota usage | +| `GET /health` | Liveness/readiness probe | | `GET /` | Web analytics dashboard | | `GET /metrics/dashboard` | Metrics dashboard UI | | `GET /metrics` | OpenMetrics exposition endpoint | @@ -45,6 +47,58 @@ Streaming (SSE) is supported on the chat, responses, and messages endpoints by setting `"stream": true` in the request body. The Gemini surface streams via the dedicated `:streamGenerateContent` action. +## Health check + +`GET /health` answers without contacting the upstream, so it is cheap enough for +a service supervisor or container probe to poll frequently. It is never guarded +by the optional API key. + +```bash +curl http://127.0.0.1:8314/health +``` + +```json +{ + "status": "ok", + "ready": true, + "version": "1.3.0", + "uptime_seconds": 128, + "copilot_token": { "present": true, "expires_in_seconds": 1487 }, + "models_loaded": 77, + "requests_served": 42, + "auth_required": false +} +``` + +`ready` is true once a Copilot token has been obtained **and** the model catalog +has loaded. A degraded proxy still answers `200` with `ready: false` so probes +can distinguish "process alive" from "able to serve traffic". Add `?strict=true` +to get `503 Service Unavailable` instead when the proxy is not ready. + +## Retrieve a model + +`GET /v1/models/{model}` returns a single catalog entry in the OpenAI shape, +including the raw `capabilities` and `supported_endpoints` reported upstream. +Model aliases from `model_mappings` are resolved, so `/v1/models/opus` returns +the mapped Copilot model. Unknown ids return `404` with an OpenAI-style error +body. + +```bash +curl http://127.0.0.1:8314/v1/models/claude-opus-4.8 +``` + +## Token counting + +`POST /v1/messages/count_tokens` forwards to the upstream Anthropic +`count_tokens` endpoint for models that expose the native `/v1/messages` +surface, returning exact counts. For every other model (and whenever the +upstream call fails) the proxy falls back to a local tiktoken estimate using the +tokenizer advertised in the model catalog. Estimated responses are marked: + +```json +{ "input_tokens": 812, "estimated": true } +``` + ## OpenAI SDK ```python @@ -131,6 +185,23 @@ curl http://127.0.0.1:8314/v1/models/full/ This is the authoritative source for which models support a 1M-token context window (those advertising `max_context_window_tokens` greater than 200,000). +## Audit API + +`GET /api/audit` returns recent request records with their extracted audit +fields. All filters are optional and combine with AND: + +| Parameter | Effect | +|-----------|--------| +| `endpoint` | Substring match on the endpoint path | +| `status` | Exact HTTP status code | +| `tool_name` | Keeps records whose request offered a matching tool | +| `agent` | `true`/`false` — agent- vs user-initiated requests | +| `model` | Substring match on the requested or translated model | +| `page`, `per_page` | Pagination (`per_page` is clamped to 500) | + +`GET /api/audit/summary` aggregates the same records into top tools, stop-reason +counts, estimated cost, and prompt-cache hit rate. + ## Notable behaviors - **GitHub Models routing** — when enabled (default), requests whose translated @@ -152,5 +223,39 @@ window (those advertising `max_context_window_tokens` greater than 200,000). - **Adaptive-thinking migration** — when an upstream model rejects `thinking.type = "enabled"`, the proxy automatically retries using the adaptive format. +- **Upstream errors are never disguised as streams** — a non-2xx upstream + response on a `"stream": true` request is returned as a normal error response + with the upstream status code, not as a `200` SSE body. +- **Interrupted streams are reported, not silently truncated** — if the upstream + connection drops mid-response, the proxy emits a protocol-appropriate + terminator (`data: {"error": …}` + `[DONE]` for OpenAI, `event: error` for + Anthropic and Responses, `finishReason: "OTHER"` for Gemini) and records the + request as `502`. Anthropic streams that end without a `finish_reason` are + also closed with `message_stop`, so clients never block on a half-open + message or keep a partial answer as a completed turn. +- **Byte-exact streaming** — SSE parsing buffers raw bytes and decodes only + complete lines, so multi-byte characters split across network chunks are + never mangled into `U+FFFD` or dropped. +- **SSE keepalive** — a `: keepalive` comment is emitted after 15 seconds of + silence so extended thinking does not trip the ~60 second idle timeout + enforced by the upstream load balancer. Comments are ignored by every + spec-compliant SSE client. +- **`anthropic-beta` passthrough** — the client's beta flags are forwarded and + merged with the ones the proxy derives (`context-1m-2025-08-07` for extended + context models, `context-management-2025-06-27` when the request uses + `context_management`). +- **MCP tool results** — a `tool_result` whose `content` is an array of blocks + is normalized before translation: text blocks are joined and image blocks + become `image_url` data URLs. Images nested there also enable the vision + header. +- **Parameter migration** — a model that rejects `max_tokens` in favour of + `max_completion_tokens` is retried once with the renamed parameter, and a + missing `max_tokens` is filled from the model catalog. +- **Cost estimates use the served model** — the `estimated_cost_usd` field and + the `ghc_proxy_estimated_cost_usd_total` metric price requests using the model + actually sent upstream (after translation), not the alias the client asked + for. +- **Graceful shutdown** — on Ctrl-C (or `SIGTERM` on Unix) the proxy stops + accepting connections and lets in-flight requests and SSE streams finish. - **Content filtering** — system-prompt add/remove and tool-result suffix removal are applied per your configuration. diff --git a/src/anthropic.rs b/src/anthropic.rs index 62ea2bd..2b41a53 100644 --- a/src/anthropic.rs +++ b/src/anthropic.rs @@ -69,16 +69,10 @@ pub fn anthropic_to_openai(req: &Value, cfg: &Config) -> Value { .cloned() .collect(); for tr in tool_results { - let mut c = tr - .get("content") - .cloned() - .unwrap_or(Value::String(String::new())); - if let Value::String(s) = &c { - c = Value::String(strip_tool_result_suffix( - s, - &cfg.tool_result_suffix_remove, - )); - } + let c = normalize_tool_result_content( + tr.get("content"), + &cfg.tool_result_suffix_remove, + ); messages.push(json!({ "role": "tool", "tool_call_id": tr.get("tool_use_id").cloned().unwrap_or(Value::Null), @@ -214,6 +208,85 @@ pub fn anthropic_to_openai(req: &Value, cfg: &Config) -> Value { Value::Object(out) } +/// Normalizes the `content` of an Anthropic `tool_result` block into something +/// the OpenAI chat-completions API accepts on a `tool` message. +/// +/// Anthropic allows `content` to be either a plain string or an array of +/// content blocks, and MCP servers routinely return the array form (a +/// screenshot, or text plus an image). Forwarding that array unchanged makes +/// the upstream reject the whole request with +/// `type has to be either 'image_url' or 'text'`, so every MCP tool that +/// returns anything but a bare string fails. Text blocks are joined and image +/// blocks are rewritten as `image_url` data URLs; a pure-text result collapses +/// back to a plain string, which is what the API expects most of the time. +fn normalize_tool_result_content(content: Option<&Value>, suffixes: &[String]) -> Value { + match content { + None | Some(Value::Null) => Value::String(String::new()), + Some(Value::String(s)) => Value::String(strip_tool_result_suffix(s, suffixes)), + Some(Value::Array(blocks)) => { + let has_image = blocks.iter().any(|b| type_of(b) == "image"); + if !has_image { + let text = blocks + .iter() + .filter_map(|b| match type_of(b) { + "text" => b.get("text").and_then(|t| t.as_str()), + _ => None, + }) + .collect::>() + .join("\n"); + return Value::String(strip_tool_result_suffix(&text, suffixes)); + } + let mut out: Vec = Vec::new(); + for b in blocks { + match type_of(b) { + "text" => { + let text = b.get("text").and_then(|t| t.as_str()).unwrap_or(""); + out.push(json!({ + "type": "text", + "text": strip_tool_result_suffix(text, suffixes) + })); + } + "image" => out.push(image_block_to_image_url(b)), + _ => {} + } + } + Value::Array(out) + } + // Anything else (a number, an object) is forwarded as a JSON string so + // the upstream still receives valid `tool` content. + Some(other) => Value::String(other.to_string()), + } +} + +/// Converts an Anthropic `image` content block into an OpenAI `image_url` part. +/// Both the `base64` and `url` source forms are supported. +fn image_block_to_image_url(block: &Value) -> Value { + let src = block.get("source"); + let source_type = src + .and_then(|s| s.get("type")) + .and_then(|t| t.as_str()) + .unwrap_or("base64"); + if source_type == "url" { + let url = src + .and_then(|s| s.get("url")) + .and_then(|u| u.as_str()) + .unwrap_or(""); + return json!({"type": "image_url", "image_url": {"url": url}}); + } + let media = src + .and_then(|s| s.get("media_type")) + .and_then(|m| m.as_str()) + .unwrap_or("image/png"); + let data = src + .and_then(|s| s.get("data")) + .and_then(|d| d.as_str()) + .unwrap_or(""); + json!({ + "type": "image_url", + "image_url": {"url": format!("data:{media};base64,{data}")} + }) +} + /// Extracts the OpenAI `content` value (string or multimodal array) from a set /// of Anthropic user content blocks. fn extract_user_content(blocks: &[Value]) -> Option { @@ -248,21 +321,7 @@ fn extract_user_content(blocks: &[Value]) -> Option { match type_of(b) { "text" => out.push(json!({"type": "text", "text": b.get("text")})), "thinking" => out.push(json!({"type": "text", "text": b.get("thinking")})), - "image" => { - let src = b.get("source"); - let media = src - .and_then(|s| s.get("media_type")) - .and_then(|m| m.as_str()) - .unwrap_or(""); - let data = src - .and_then(|s| s.get("data")) - .and_then(|d| d.as_str()) - .unwrap_or(""); - out.push(json!({ - "type": "image_url", - "image_url": {"url": format!("data:{media};base64,{data}")} - })); - } + "image" => out.push(image_block_to_image_url(b)), _ => {} } } @@ -450,6 +509,8 @@ pub struct AnthropicStreamState { message_start_sent: bool, content_block_index: i64, content_block_open: bool, + /// Set once a terminating `message_stop` has been emitted. + finished: bool, /// OpenAI tool-call index -> anthropic content block index. tool_calls: std::collections::HashMap, } @@ -607,10 +668,49 @@ impl AnthropicStreamState { "usage": usage_out })); events.push(json!({"type": "message_stop"})); + self.finished = true; } events } + + /// Whether the upstream already delivered a `finish_reason`, i.e. the + /// Anthropic event sequence was properly terminated. + pub fn is_finished(&self) -> bool { + self.finished + } + + /// Closes an unterminated stream. + /// + /// When the upstream connection drops (or simply ends) before sending a + /// `finish_reason`, no `message_stop` was emitted. Anthropic clients block + /// waiting for one, and a partially received assistant turn can end up + /// recorded as if it were complete. This emits the missing + /// `content_block_stop` / `message_delta` / `message_stop` sequence with an + /// explicit `stop_reason`, so the client sees a terminated — and, when + /// `stop_reason` is `"error"`, visibly incomplete — message. + /// + /// Returns no events when the stream was already terminated or never + /// started. + pub fn finish(&mut self, stop_reason: &str) -> Vec { + if self.finished || !self.message_start_sent { + self.finished = true; + return Vec::new(); + } + let mut events = Vec::new(); + if self.content_block_open { + events.push(json!({"type": "content_block_stop", "index": self.content_block_index})); + self.content_block_open = false; + } + events.push(json!({ + "type": "message_delta", + "delta": {"stop_reason": stop_reason, "stop_sequence": Value::Null}, + "usage": {"output_tokens": 0} + })); + events.push(json!({"type": "message_stop"})); + self.finished = true; + events + } } /// Anthropic request keys forwarded to the upstream `/v1/messages` endpoint. @@ -630,8 +730,49 @@ const ALLOWED_ANTHROPIC_KEYS: &[&str] = &[ "thinking", "output_config", "service_tier", + // Claude Code's automatic context editing. Copilot accepts this once the + // matching beta is requested; silently dropping it disables compaction on + // the client without any diagnostic. + "context_management", ]; +/// Anthropic beta flag that unlocks the `context_management` request field. +pub const CONTEXT_MANAGEMENT_BETA: &str = "context-management-2025-06-27"; +/// Anthropic beta flag that unlocks the 1M-token context window. +pub const CONTEXT_1M_BETA: &str = "context-1m-2025-08-07"; + +/// Builds the `anthropic-beta` header value for an upstream request. +/// +/// The client's own `anthropic-beta` header is preserved — Claude Code sends +/// flags there that the proxy has no business dropping — and the flags this +/// proxy derives are appended when missing. Returns `None` when there is +/// nothing to send. +pub fn merge_anthropic_beta(client_value: Option<&str>, derived: &[&str]) -> Option { + let mut flags: Vec = Vec::new(); + if let Some(value) = client_value { + for flag in value.split(',') { + let flag = flag.trim(); + if !flag.is_empty() && !flags.iter().any(|f| f == flag) { + flags.push(flag.to_string()); + } + } + } + for flag in derived { + if !flags.iter().any(|f| f == flag) { + flags.push((*flag).to_string()); + } + } + (!flags.is_empty()).then(|| flags.join(",")) +} + +/// Whether the request asks for Claude Code's context-editing feature, which +/// requires the `context-management-2025-06-27` beta to be requested. +pub fn uses_context_management(req: &Value) -> bool { + req.get("context_management") + .map(|v| !v.is_null()) + .unwrap_or(false) +} + fn clean_cache_control(block: &mut Value) { if let Some(cc) = block.get_mut("cache_control") { if cc.get("type").and_then(|t| t.as_str()) == Some("ephemeral") { @@ -897,15 +1038,136 @@ pub fn apply_tool_result_suffix(req: &Value, cfg: &Config) -> Value { } /// Whether the message list contains an image content block. +/// +/// Images can appear either directly in a message's content or nested inside a +/// `tool_result` block — the shape every MCP screenshot tool produces. Missing +/// the nested case means the request goes upstream without +/// `Copilot-Vision-Request`, so the image is silently ignored or rejected. pub fn has_image(req: &Value) -> bool { arr(req, "messages").iter().any(|m| { m.get("content") .and_then(|c| c.as_array()) - .map(|blocks| blocks.iter().any(|b| type_of(b) == "image")) + .map(|blocks| blocks.iter().any(block_has_image)) .unwrap_or(false) }) } +/// Whether a content block is an image, or carries one in its nested +/// `tool_result` content. +fn block_has_image(block: &Value) -> bool { + match type_of(block) { + "image" => true, + "tool_result" => block + .get("content") + .and_then(|c| c.as_array()) + .map(|inner| inner.iter().any(|b| type_of(b) == "image")) + .unwrap_or(false), + _ => false, + } +} + +/// Object keys skipped when flattening a request for token estimation: opaque +/// identifiers and, most importantly, base64 image payloads which would +/// otherwise dominate the count. +const NON_COUNTABLE_KEYS: &[&str] = &[ + "data", + "type", + "id", + "tool_use_id", + "media_type", + "cache_control", +]; + +/// Recursively appends every string leaf of `value` to `out`. +fn push_strings(value: &Value, out: &mut String) { + match value { + Value::String(s) => { + out.push_str(s); + out.push('\n'); + } + Value::Array(items) => { + for item in items { + push_strings(item, out); + } + } + Value::Object(map) => { + for (k, v) in map { + if NON_COUNTABLE_KEYS.contains(&k.as_str()) { + continue; + } + push_strings(v, out); + } + } + _ => {} + } +} + +/// Flattens the countable text of an Anthropic Messages request: the system +/// prompt, every message content block, and the tool definitions. Used for the +/// local token estimate served by `/v1/messages/count_tokens` when the upstream +/// native counting endpoint is unavailable. +pub fn collect_countable_text(req: &Value) -> String { + let mut out = String::new(); + if let Some(system) = req.get("system") { + push_strings(system, &mut out); + } + for msg in arr(req, "messages") { + if let Some(content) = msg.get("content") { + push_strings(content, &mut out); + } + } + for tool in arr(req, "tools") { + push_strings(&tool, &mut out); + } + out +} + +/// Fixed token overhead the API charges for each message's role/framing. +const PER_MESSAGE_OVERHEAD: u64 = 4; +/// Fixed token overhead for each tool definition beyond its serialized schema. +const PER_TOOL_OVERHEAD: u64 = 8; +/// Fixed overhead charged once per request. +const REQUEST_OVERHEAD: u64 = 8; + +/// Estimates the input token count of an Anthropic Messages request when the +/// upstream native counting endpoint is unavailable. +/// +/// Counting only the visible text under-reports badly: the API also charges for +/// per-message framing and for the full JSON schema of every tool definition. +/// A client that trusts a low number (Claude Code decides when to compact from +/// it) keeps growing the conversation until the request is hard-rejected, so +/// this deliberately errs on the side of the structural overhead being present. +pub fn estimate_input_tokens(req: &Value, tokenizer: &str) -> u64 { + use crate::filters::count_tokens; + + let mut total = REQUEST_OVERHEAD; + + if let Some(system) = req.get("system") { + let mut text = String::new(); + push_strings(system, &mut text); + total += count_tokens(&text, tokenizer); + } + + for msg in arr(req, "messages") { + total += PER_MESSAGE_OVERHEAD; + if let Some(content) = msg.get("content") { + let mut text = String::new(); + push_strings(content, &mut text); + total += count_tokens(&text, tokenizer); + } + } + + // Tool definitions are sent as JSON, so the schema punctuation and keys are + // billed too — counting only the description would miss most of it. + for tool in arr(req, "tools") { + total += PER_TOOL_OVERHEAD; + let serialized = serde_json::to_string(&tool).unwrap_or_default(); + total += count_tokens(&serialized, tokenizer); + } + + total +} + #[cfg(test)] mod tests { use super::*; @@ -929,6 +1191,41 @@ mod tests { assert_eq!(out["max_tokens"], 100); } + #[test] + fn countable_text_covers_system_messages_and_tools() { + let req = json!({ + "system": [{"type": "text", "text": "system rules"}], + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "hello there"}]}, + {"role": "assistant", "content": "sure thing"} + ], + "tools": [{"name": "read_file", "description": "reads a file"}] + }); + let text = collect_countable_text(&req); + assert!(text.contains("system rules")); + assert!(text.contains("hello there")); + assert!(text.contains("sure thing")); + assert!(text.contains("read_file")); + assert!(text.contains("reads a file")); + } + + #[test] + fn countable_text_skips_base64_image_payloads() { + let blob = "A".repeat(5000); + let req = json!({ + "messages": [{ + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": blob}} + ] + }] + }); + let text = collect_countable_text(&req); + assert!(text.contains("describe")); + assert!(!text.contains("AAAA")); + } + #[test] fn tool_result_becomes_tool_message() { let cfg = Config::default(); @@ -945,6 +1242,153 @@ mod tests { assert_eq!(out["messages"][0]["content"], "ok"); } + #[test] + fn tool_result_block_array_is_flattened_to_text() { + // MCP servers commonly return `content` as an array of blocks. Passing + // the array through unchanged makes the upstream reject the request. + let cfg = Config::default(); + let req = json!({ + "model": "claude-3", + "messages": [{ + "role": "user", + "content": [{ + "type": "tool_result", + "tool_use_id": "abc", + "content": [ + {"type": "text", "text": "line one"}, + {"type": "text", "text": "line two"} + ] + }] + }] + }); + let out = anthropic_to_openai(&req, &cfg); + assert_eq!(out["messages"][0]["role"], "tool"); + assert_eq!(out["messages"][0]["content"], "line one\nline two"); + } + + #[test] + fn tool_result_image_block_becomes_image_url() { + let cfg = Config::default(); + let req = json!({ + "model": "claude-3", + "messages": [{ + "role": "user", + "content": [{ + "type": "tool_result", + "tool_use_id": "abc", + "content": [ + {"type": "text", "text": "screenshot:"}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "QUJD"}} + ] + }] + }] + }); + let out = anthropic_to_openai(&req, &cfg); + let content = &out["messages"][0]["content"]; + assert_eq!(content[0]["type"], "text"); + assert_eq!(content[1]["type"], "image_url"); + assert_eq!(content[1]["image_url"]["url"], "data:image/png;base64,QUJD"); + } + + #[test] + fn url_source_images_are_passed_through() { + let block = + json!({"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}); + let out = image_block_to_image_url(&block); + assert_eq!(out["image_url"]["url"], "https://example.com/a.png"); + } + + #[test] + fn vision_is_detected_inside_tool_results() { + let img = json!({"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "x"}}); + // Directly in the message content. + let direct = json!({"messages": [{"role": "user", "content": [img.clone()]}]}); + assert!(has_image(&direct)); + // Nested in a tool_result — the shape MCP screenshot tools produce. + let nested = json!({"messages": [{ + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "t", "content": [ + {"type": "text", "text": "here"}, img + ]}] + }]}); + assert!(has_image(&nested)); + // No image anywhere. + let none = + json!({"messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]}); + assert!(!has_image(&none)); + // A string-content tool_result must not panic or false-positive. + let plain = json!({"messages": [{ + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "t", "content": "ok"}] + }]}); + assert!(!has_image(&plain)); + } + + #[test] + fn tool_result_suffix_is_stripped_inside_blocks() { + let suffixes = vec!["\n[end]".to_string()]; + let content = json!([{"type": "text", "text": "result\n[end]"}]); + let out = normalize_tool_result_content(Some(&content), &suffixes); + assert_eq!(out, json!("result")); + } + + #[test] + fn beta_header_merges_client_and_derived_flags() { + // The client's flags are preserved and derived ones appended. + assert_eq!( + merge_anthropic_beta(Some("claude-code-20250219"), &[CONTEXT_1M_BETA]), + Some(format!("claude-code-20250219,{CONTEXT_1M_BETA}")) + ); + // Duplicates are collapsed. + assert_eq!( + merge_anthropic_beta(Some(CONTEXT_1M_BETA), &[CONTEXT_1M_BETA]), + Some(CONTEXT_1M_BETA.to_string()) + ); + // Whitespace around client flags is tolerated. + assert_eq!( + merge_anthropic_beta(Some("a , b"), &[]), + Some("a,b".to_string()) + ); + // Nothing to send. + assert_eq!(merge_anthropic_beta(None, &[]), None); + assert_eq!(merge_anthropic_beta(Some(""), &[]), None); + } + + #[test] + fn context_management_survives_sanitization() { + let req = json!({ + "model": "m", + "messages": [], + "context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]} + }); + assert!(uses_context_management(&req)); + let out = sanitize_anthropic_request(&req); + assert!(out.get("context_management").is_some()); + } + + #[test] + fn token_estimate_includes_message_and_tool_overhead() { + let bare = json!({"messages": [{"role": "user", "content": "hello"}]}); + let with_tools = json!({ + "messages": [{"role": "user", "content": "hello"}], + "tools": [{ + "name": "read_file", + "description": "Read a file from disk", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}} + }] + }); + let bare_count = estimate_input_tokens(&bare, "cl100k_base"); + let tool_count = estimate_input_tokens(&with_tools, "cl100k_base"); + // The tool schema must be billed, otherwise clients compact too late. + assert!( + tool_count > bare_count + 10, + "tools added only {} tokens", + tool_count - bare_count + ); + // Even a bare request carries per-message framing overhead. + assert!(bare_count > crate::filters::count_tokens("hello", "cl100k_base")); + } + #[test] fn openai_response_to_anthropic() { let resp = json!({ @@ -980,6 +1424,49 @@ mod tests { assert_eq!(start[0]["type"], "message_start"); let end = st.process(&json!({"choices": [{"delta": {}, "finish_reason": "stop"}]})); assert!(end.iter().any(|e| e["type"] == "message_stop")); + assert!(st.is_finished()); + // A properly terminated stream needs no synthetic closing events. + assert!(st.finish("error").is_empty()); + } + + #[test] + fn finish_closes_a_stream_cut_before_finish_reason() { + let mut st = AnthropicStreamState::new(); + st.process(&json!({"id": "1", "model": "m", "choices": [{"delta": {"content": "par"}}]})); + assert!(!st.is_finished()); + // Upstream died here: no finish_reason, so no message_stop was sent and + // an Anthropic client would block forever. + let events = st.finish("error"); + assert_eq!(events[0]["type"], "content_block_stop"); + assert_eq!(events[1]["type"], "message_delta"); + assert_eq!(events[1]["delta"]["stop_reason"], "error"); + assert_eq!(events[2]["type"], "message_stop"); + assert!(st.is_finished()); + // Idempotent. + assert!(st.finish("error").is_empty()); + } + + #[test] + fn finish_is_noop_before_any_chunk_arrived() { + // Nothing was ever emitted, so there is no half-open message to close. + let mut st = AnthropicStreamState::new(); + assert!(st.finish("error").is_empty()); + } + + #[test] + fn finish_closes_an_open_tool_use_block() { + let mut st = AnthropicStreamState::new(); + st.process(&json!({ + "id": "1", + "model": "m", + "choices": [{"delta": {"tool_calls": [ + {"index": 0, "id": "toolu_1", "function": {"name": "read_file", "arguments": "{\"p"}} + ]}}] + })); + // The tool_use block is open with truncated JSON arguments. + let events = st.finish("error"); + assert_eq!(events[0]["type"], "content_block_stop"); + assert!(events.iter().any(|e| e["type"] == "message_stop")); } #[test] diff --git a/src/config.rs b/src/config.rs index 137e4a8..d1f4174 100644 --- a/src/config.rs +++ b/src/config.rs @@ -10,7 +10,7 @@ use std::path::PathBuf; /// Kept in sync with the `engines.vscode` baseline of the latest /// `microsoft/vscode-copilot-chat` release (see "Mimicking the Copilot client" /// in the README for how to refresh these values). -pub const VSCODE_VERSION: &str = "1.123.0"; +pub const VSCODE_VERSION: &str = "1.130.0"; /// Default GitHub Copilot API version header value (`X-GitHub-Api-Version`), /// matching the `X-GitHub-Api-Version` constant in the Copilot Chat client /// source (`src/platform/networking/common/networking.ts`). diff --git a/src/main.rs b/src/main.rs index a3642d8..affc0f9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -602,6 +602,14 @@ fn print_setup_guide( println!("{bar}"); } +/// Checks GitHub releases and replaces this binary when a newer version is +/// available. +/// +/// This is fully blocking: `self_update` builds its own blocking HTTP client, +/// which stands up a private Tokio runtime. Dropping that runtime from inside +/// an async context panics with "Cannot drop a runtime in a context where +/// blocking is not allowed", so [`run_auto_upgrade`] must keep this on a +/// blocking thread. fn maybe_auto_upgrade(enabled: bool) { if !enabled { return; @@ -640,6 +648,21 @@ fn maybe_auto_upgrade(enabled: bool) { } } +/// Runs the auto-upgrade check on a blocking thread. +/// +/// `self_update` is synchronous and owns a private Tokio runtime internally. +/// Calling it directly from the async `main` panicked the process on startup +/// ("Cannot drop a runtime in a context where blocking is not allowed"), which +/// made `auto_upgrade: true` unusable. +async fn run_auto_upgrade(enabled: bool) { + if !enabled { + return; + } + if let Err(e) = tokio::task::spawn_blocking(move || maybe_auto_upgrade(enabled)).await { + tracing::warn!("Auto-upgrade task failed: {e}"); + } +} + #[tokio::main] async fn main() { tracing_subscriber::fmt() @@ -838,7 +861,7 @@ async fn main() { } // Optionally self-update from GitHub releases before serving traffic. - maybe_auto_upgrade(cfg.auto_upgrade); + run_auto_upgrade(cfg.auto_upgrade).await; // Optionally refresh the VS Code version used in upstream headers. if cfg.dynamic_vscode_version { @@ -895,7 +918,9 @@ async fn main() { { let state = app_state.clone(); tokio::spawn(async move { - let mut interval = tokio::time::interval(std::time::Duration::from_secs(30 * 60)); + let period = std::time::Duration::from_secs(30 * 60); + let mut interval = + tokio::time::interval_at(tokio::time::Instant::now() + period, period); loop { interval.tick().await; if let Err(e) = state.load_models().await { @@ -917,6 +942,7 @@ async fn main() { println!("\nStarting GitHub Copilot API Proxy on {host}:{port}"); println!("Dashboard: http://{host}:{port}/"); + println!("Health check: http://{host}:{port}/health"); println!("Metrics UI: http://{host}:{port}/metrics/dashboard"); println!("OpenMetrics: http://{host}:{port}/metrics"); println!("Reload config: POST http://{host}:{port}/api/config/reload"); @@ -933,10 +959,48 @@ async fn main() { std::process::exit(1); } }; - if let Err(e) = axum::serve(listener, app).await { + if let Err(e) = axum::serve(listener, app) + .with_graceful_shutdown(shutdown_signal()) + .await + { eprintln!("Server error: {e}"); std::process::exit(1); } + tracing::info!("Server stopped."); +} + +/// Resolves when the process receives Ctrl-C (all platforms) or SIGTERM (Unix), +/// letting in-flight requests and SSE streams finish before the socket closes. +async fn shutdown_signal() { + let ctrl_c = async { + if let Err(e) = tokio::signal::ctrl_c().await { + tracing::warn!("Failed to install Ctrl-C handler: {e}"); + // Never resolve, so shutdown is only driven by the other signal. + std::future::pending::<()>().await; + } + }; + + #[cfg(unix)] + let terminate = async { + match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) { + Ok(mut sig) => { + sig.recv().await; + } + Err(e) => { + tracing::warn!("Failed to install SIGTERM handler: {e}"); + std::future::pending::<()>().await; + } + } + }; + + #[cfg(not(unix))] + let terminate = std::future::pending::<()>(); + + tokio::select! { + _ = ctrl_c => {} + _ = terminate => {} + } + tracing::info!("Shutdown signal received; draining in-flight requests..."); } #[cfg(test)] diff --git a/src/server.rs b/src/server.rs index 9cc87bb..fc90242 100644 --- a/src/server.rs +++ b/src/server.rs @@ -28,6 +28,8 @@ pub fn router(state: SharedState) -> Router { .route("/models", get(get_models)) .route("/v1/models/full/", get(get_models_full)) .route("/models/full/", get(get_models_full)) + .route("/v1/models/{model_id}", get(get_model)) + .route("/models/{model_id}", get(get_model)) .route("/v1/chat/completions", post(chat_completions)) .route("/chat/completions", post(chat_completions)) .route("/v1/responses", post(responses)) @@ -37,6 +39,7 @@ pub fn router(state: SharedState) -> Router { .route("/v1/embeddings", post(embeddings)) .route("/embeddings", post(embeddings)) .route("/usage", get(usage)) + .route("/health", get(health)) .route("/metrics", get(metrics_openmetrics)) .route("/", get(dashboard)) .route("/requests", get(requests_page)) @@ -174,7 +177,12 @@ fn set_initiator(headers: &mut HeaderMap, agent: bool) { headers.insert("X-Initiator", HeaderValue::from_static(v)); } -/// Logs an upstream error to `error.log` in the config directory. +/// Maximum size of `error.log` before it is rotated to `error.log.1`. +const ERROR_LOG_MAX_BYTES: u64 = 8 * 1024 * 1024; + +/// Logs an upstream error to `error.log` in the config directory. The file is +/// rotated once it exceeds `ERROR_LOG_MAX_BYTES` so a persistently failing +/// upstream cannot fill the disk. fn log_error(endpoint: &str, request: &Value, response: &str, status: u16) { let dir = crate::config::config_dir(); let _ = std::fs::create_dir_all(&dir); @@ -186,6 +194,12 @@ fn log_error(endpoint: &str, request: &Value, response: &str, status: u16) { "response": response, }); let path = dir.join("error.log"); + if std::fs::metadata(&path) + .map(|m| m.len() > ERROR_LOG_MAX_BYTES) + .unwrap_or(false) + { + let _ = std::fs::rename(&path, dir.join("error.log.1")); + } use std::io::Write; if let Ok(mut f) = std::fs::OpenOptions::new() .create(true) @@ -243,6 +257,26 @@ fn error_response(status: StatusCode, msg: String) -> Response { (status, Json(json!({"error": msg}))).into_response() } +/// Renames `max_tokens` to `max_completion_tokens`, the parameter newer +/// OpenAI-family models require. Returns false when there is nothing to rename, +/// so callers can avoid a pointless retry. +fn rewrite_max_tokens_param(req: &mut Value) -> bool { + let Some(obj) = req.as_object_mut() else { + return false; + }; + if obj.contains_key("max_completion_tokens") { + // Already migrated; just drop the rejected alias. + return obj.remove("max_tokens").is_some(); + } + match obj.remove("max_tokens") { + Some(v) => { + obj.insert("max_completion_tokens".to_string(), v); + true + } + None => false, + } +} + // --------------------------------------------------------------------------- // Audit extraction helpers (Phase 1: Foundation for analytics) // --------------------------------------------------------------------------- @@ -272,60 +306,40 @@ fn extract_message_count(body: &Value) -> usize { .unwrap_or(0) } -/// Extract stop reason from SSE response body (may contain multiple events). -fn extract_stop_reason_from_sse(body: &str) -> Option { - for line in body.lines() { - if let Some(data) = line.strip_prefix("data: ") { - if let Ok(event) = serde_json::from_str::(data) { - if let Some(sr) = event - .get("delta") - .and_then(|d| d.get("stop_reason")) - .and_then(|s| s.as_str()) - { - return Some(sr.to_string()); - } - } - } - } - None -} - -/// Extract tool calls from SSE response (streaming events). -fn extract_tools_called_from_sse(body: &str) -> Vec { - let mut tools = Vec::new(); - - for line in body.lines() { - if let Some(data) = line.strip_prefix("data: ") { - if let Ok(event) = serde_json::from_str::(data) { - // Check for tool_use in content_block - if let Some(block) = event.get("content_block") { - if block.get("type").and_then(|t| t.as_str()) == Some("tool_use") { - if let Some(name) = block.get("name").and_then(|n| n.as_str()) { - if !tools.contains(&name.to_string()) { - tools.push(name.to_string()); - } - } - } - } - } - } +/// Per-1K-token (input, output) rates used for cost estimation. +/// +/// Arms are evaluated in order, so more specific model families must come +/// before the prefixes that would otherwise swallow them (`gpt-4o` before +/// `gpt-4`, `opus`/`sonnet`/`haiku` before the generic `claude` fallback). +fn model_rates(model: &str) -> (f64, f64) { + let m = model.to_ascii_lowercase(); + // Strip a `publisher/` prefix so GitHub Models ids price like their base + // model (e.g. `openai/gpt-4o` -> `gpt-4o`). + let m = m.rsplit('/').next().unwrap_or(&m); + match m { + m if m.contains("opus") => (0.015, 0.075), + m if m.contains("sonnet") => (0.003, 0.015), + m if m.contains("haiku") => (0.0008, 0.004), + m if m.contains("gpt-4o-mini") => (0.00015, 0.0006), + m if m.contains("gpt-4o") => (0.005, 0.015), + m if m.contains("gpt-4.1-mini") => (0.0004, 0.0016), + m if m.contains("gpt-4.1") => (0.002, 0.008), + m if m.contains("gpt-4") => (0.03, 0.06), + m if m.contains("gpt-5-mini") || m.contains("gpt-5.5-mini") => (0.00025, 0.002), + m if m.contains("gpt-5") => (0.00125, 0.01), + m if m.contains("o3-mini") || m.contains("o4-mini") => (0.0011, 0.0044), + m if m.contains("gemini") && m.contains("flash") => (0.0003, 0.0025), + m if m.contains("gemini") => (0.00125, 0.01), + m if m.contains("embedding") => (0.00002, 0.0), + _ => (0.0005, 0.0015), } - - tools } /// Calculate estimated cost in USD based on token counts and model. -/// Uses simplified rates: Claude $0.003/$0.015 (input/output), GPT-4 $0.03/$0.06, etc. +/// Rates are approximate public list prices per 1K tokens and are only used for +/// the dashboard/metrics estimate, never for billing. fn calculate_cost(model: &str, input_tokens: u64, output_tokens: u64) -> f64 { - let (input_rate, output_rate) = match model { - m if m.contains("opus-4") => (0.015, 0.075), // claude-opus - m if m.contains("sonnet") => (0.003, 0.015), // claude-sonnet - m if m.contains("haiku") => (0.0008, 0.004), // claude-haiku - m if m.contains("gpt-4") => (0.03, 0.06), // gpt-4 - m if m.contains("gpt-4o") => (0.005, 0.015), // gpt-4o - _ => (0.0005, 0.0015), // fallback - }; - + let (input_rate, output_rate) = model_rates(model); (input_tokens as f64 * input_rate + output_tokens as f64 * output_rate) / 1000.0 } @@ -452,6 +466,54 @@ async fn get_models_full(State(state): State) -> Response { Json(models.clone().unwrap_or(Value::Null)).into_response() } +/// OpenAI-compatible single model retrieval (`GET /v1/models/{model}`). +/// +/// Model ids are matched after translation so aliases configured in +/// `model_mappings` (e.g. `sonnet`) resolve the same way they do on the +/// inference endpoints. Unknown ids return a 404 in the OpenAI error shape. +async fn get_model(State(state): State, Path(model_id): Path) -> Response { + if let Err(e) = state.ensure_copilot_token().await { + return error_response(StatusCode::INTERNAL_SERVER_ERROR, e); + } + if let Err(e) = state + .ensure_models_fresh(Duration::from_secs(30 * 60)) + .await + { + tracing::warn!("model refresh failed: {e}"); + } + let translated = translate::translate(&state.model_mappings(), &model_id); + let found = match state.find_model(&model_id).await { + Some(m) => Some(m), + None => state.find_model(&translated).await, + }; + let Some(entry) = found else { + return ( + StatusCode::NOT_FOUND, + Json(json!({ + "error": { + "message": format!("The model '{model_id}' does not exist."), + "type": "invalid_request_error", + "code": "model_not_found" + } + })), + ) + .into_response(); + }; + let id = entry.get("id").cloned().unwrap_or(Value::Null); + Json(json!({ + "id": id, + "object": "model", + "type": "model", + "created": 0, + "created_at": "1970-01-01T00:00:00.000Z", + "owned_by": entry.get("vendor").cloned().unwrap_or(Value::String("unknown".into())), + "display_name": entry.get("name").cloned().or_else(|| entry.get("id").cloned()).unwrap_or(Value::Null), + "capabilities": entry.get("capabilities").cloned().unwrap_or(Value::Null), + "supported_endpoints": entry.get("supported_endpoints").cloned().unwrap_or(Value::Null), + })) + .into_response() +} + // --------------------------------------------------------------------------- // Chat completions // --------------------------------------------------------------------------- @@ -484,6 +546,42 @@ async fn chat_completions(State(state): State, body: Bytes) -> Resp return error_response(StatusCode::TOO_MANY_REQUESTS, e); } + // Some Copilot models are only reachable through `/responses` and answer a + // chat-completions call with an opaque `unsupported_api_for_model` 400. + // Turn that into an actionable message before spending the round trip. + if !to_github_models + && !state + .model_supports_endpoint(&translated, "/chat/completions") + .await + && state + .model_supports_endpoint(&translated, "/responses") + .await + { + return ( + StatusCode::BAD_REQUEST, + Json(json!({ + "error": { + "message": format!( + "Model '{original_model}' is not available on /v1/chat/completions. \ + Use /v1/responses with '{translated}' instead." + ), + "type": "invalid_request_error", + "code": "unsupported_api_for_model" + } + })), + ) + .into_response(); + } + + // Copilot rejects a chat-completions request that omits `max_tokens` for + // some models. Fill it from the model catalog rather than surfacing an + // avoidable 400. + if req.get("max_tokens").is_none() && req.get("max_completion_tokens").is_none() { + if let Some(limit) = state.model_max_output_tokens(&translated).await { + req["max_tokens"] = json!(limit); + } + } + let messages = req .get("messages") .and_then(|m| m.as_array()) @@ -527,7 +625,7 @@ async fn chat_completions(State(state): State, body: Bytes) -> Resp state.clone(), &url, headers, - payload, + req, "/v1/chat/completions", original_model, translated, @@ -537,7 +635,14 @@ async fn chat_completions(State(state): State, body: Bytes) -> Resp .await; } - let resp = util::post_with_retry(&state, &url, headers, payload, "/v1/chat/completions").await; + let resp = util::post_with_retry( + &state, + &url, + headers.clone(), + payload, + "/v1/chat/completions", + ) + .await; let Some(resp) = resp else { return error_response( StatusCode::GATEWAY_TIMEOUT, @@ -547,8 +652,24 @@ async fn chat_completions(State(state): State, body: Bytes) -> Resp ), ); }; - let status = resp.status(); - let text = resp.text().await.unwrap_or_default(); + let mut status = resp.status(); + let mut text = resp.text().await.unwrap_or_default(); + // Newer OpenAI-family models reject `max_tokens` and demand + // `max_completion_tokens`. Migrate and retry once instead of surfacing a + // parameter-naming error the caller cannot act on. + if util::is_max_tokens_unsupported_error(status.as_u16(), &text) + && rewrite_max_tokens_param(&mut req) + { + tracing::info!("[/v1/chat/completions] retrying with max_completion_tokens"); + let retry_payload = serde_json::to_vec(&req).unwrap_or_default(); + if let Some(retry) = + util::post_with_retry(&state, &url, headers, retry_payload, "/v1/chat/completions") + .await + { + status = retry.status(); + text = retry.text().await.unwrap_or_default(); + } + } let resp_size = text.len(); log_debug_response(&state, "/v1/chat/completions", &text); if status.is_success() { @@ -563,6 +684,7 @@ async fn chat_completions(State(state): State, body: Bytes) -> Resp .and_then(|t| t.as_u64()) .unwrap_or(0); let (tool_count, tool_names) = extract_tools_from_request(&req); + let cost = calculate_cost(&translated, input_tokens, output_tokens); state.store.add(RequestRecord { id: uuid::Uuid::new_v4().to_string(), timestamp: now_iso(), @@ -584,7 +706,7 @@ async fn chat_completions(State(state): State, body: Bytes) -> Resp tools_called: None, is_agent_initiated: Some(agent), prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost(&original_model, input_tokens, output_tokens)), + estimated_cost_usd: Some(cost), }); Json(parsed).into_response() } else { @@ -709,6 +831,7 @@ async fn responses(State(state): State, body: Bytes) -> Response { .and_then(|t| t.as_u64()) .unwrap_or(0); let (tool_count, tool_names) = extract_tools_from_request(&req); + let cost = calculate_cost(&translated, input_tokens, output_tokens); state.store.add(RequestRecord { id: uuid::Uuid::new_v4().to_string(), timestamp: now_iso(), @@ -730,7 +853,7 @@ async fn responses(State(state): State, body: Bytes) -> Response { tools_called: None, is_agent_initiated: Some(agent), prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost(&original_model, input_tokens, output_tokens)), + estimated_cost_usd: Some(cost), }); Json(parsed).into_response() } else { @@ -743,7 +866,11 @@ async fn responses(State(state): State, body: Bytes) -> Response { // Anthropic messages // --------------------------------------------------------------------------- -async fn messages(State(state): State, body: Bytes) -> Response { +async fn messages( + State(state): State, + client_headers: HeaderMap, + body: Bytes, +) -> Response { let start = Instant::now(); let mut req = match parse_body(&body) { Ok(v) => v, @@ -772,18 +899,56 @@ async fn messages(State(state): State, body: Bytes) -> Response { req = anthropic::apply_system_prompt(&req, &cfg); req = anthropic::apply_tool_result_suffix(&req, &cfg); + let client_beta = client_beta_header(&client_headers); if state.use_direct_anthropic(&translated).await { - messages_direct(state, req, original_model, translated, start).await + messages_direct(state, req, original_model, translated, client_beta, start).await } else { messages_translated(state, req, original_model, translated, start).await } } +/// Reads the client's `anthropic-beta` header, if any. +fn client_beta_header(headers: &HeaderMap) -> Option { + headers + .get("anthropic-beta") + .and_then(|v| v.to_str().ok()) + .map(|v| v.to_string()) +} + +/// Applies the `anthropic-beta` header for an upstream Anthropic-native call, +/// preserving whatever the client asked for and adding the flags this request +/// needs. +async fn apply_anthropic_beta( + state: &SharedState, + headers: &mut HeaderMap, + model: &str, + req: &Value, + client_beta: Option<&str>, +) { + let mut derived: Vec<&str> = Vec::new(); + // Mirror the official Anthropic API pattern for unlocking the 1M-token + // context window; only for models whose catalog advertises it. + if state.model_supports_1m(model).await { + derived.push(anthropic::CONTEXT_1M_BETA); + } + // `context_management` is rejected with a misleading "Extra inputs are not + // permitted" 400 unless the matching beta is requested. + if anthropic::uses_context_management(req) { + derived.push(anthropic::CONTEXT_MANAGEMENT_BETA); + } + if let Some(value) = anthropic::merge_anthropic_beta(client_beta, &derived) { + if let Ok(v) = HeaderValue::from_str(&value) { + headers.insert("anthropic-beta", v); + } + } +} + async fn messages_direct( state: SharedState, req: Value, original_model: String, translated: String, + client_beta: Option, start: Instant, ) -> Response { let vision = anthropic::has_image(&req); @@ -797,15 +962,14 @@ async fn messages_direct( .unwrap_or(false); let mut headers = state.copilot_headers(vision).await; headers.insert("anthropic-version", HeaderValue::from_static("2023-06-01")); - // Mirror the official Anthropic API pattern for unlocking the 1M-token - // context window. Copilot accepts this header harmlessly, so only send it - // for models whose catalog actually advertises the extended window. - if state.model_supports_1m(&translated).await { - headers.insert( - "anthropic-beta", - HeaderValue::from_static("context-1m-2025-08-07"), - ); - } + apply_anthropic_beta( + &state, + &mut headers, + &translated, + &req, + client_beta.as_deref(), + ) + .await; set_initiator(&mut headers, agent); let url = format!("{}/v1/messages", state.copilot_base_url()); @@ -928,11 +1092,7 @@ async fn messages_direct( tools_called: (!tools_called.is_empty()).then_some(tools_called), is_agent_initiated: Some(agent), prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost( - &original_model, - input_tokens, - output_tokens, - )), + estimated_cost_usd: Some(calculate_cost(&translated, input_tokens, output_tokens)), }); return Json(parsed).into_response(); } @@ -994,9 +1154,7 @@ async fn messages_translated( // at the /v1/chat/completions level only. let url = format!("{}/chat/completions", state.copilot_base_url()); let mut headers = state.copilot_headers(vision).await; - let is_github_models = false; set_initiator(&mut headers, agent); - let _ = is_github_models; // used below for store record let req_size = serde_json::to_vec(¤t).map(|v| v.len()).unwrap_or(0); let payload = serde_json::to_vec(&openai_req).unwrap_or_default(); log_debug_request(&state, "/v1/messages", &openai_req); @@ -1066,11 +1224,7 @@ async fn messages_translated( tools_called: None, is_agent_initiated: Some(agent), prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost( - &original_model, - input_tokens, - output_tokens, - )), + estimated_cost_usd: Some(calculate_cost(&translated, input_tokens, output_tokens)), }); return Json(anthropic_resp).into_response(); } @@ -1092,7 +1246,11 @@ async fn messages_translated( anthropic_error(StatusCode::BAD_GATEWAY, "Exhausted retries".into()) } -async fn count_tokens(State(state): State, body: Bytes) -> Response { +async fn count_tokens( + State(state): State, + client_headers: HeaderMap, + body: Bytes, +) -> Response { if state.ensure_copilot_token().await.is_err() { return Json(json!({"input_tokens": 1})).into_response(); } @@ -1114,39 +1272,65 @@ async fn count_tokens(State(state): State, body: Bytes) -> Response .ensure_models_fresh(Duration::from_secs(30 * 60)) .await; - // Prefer real token counting from upstream responses whenever the model - // supports the Anthropic native count-tokens endpoint. - if state + // Prefer real token counting from upstream whenever the model can plausibly + // serve it. Copilot's catalog does not advertise + // `/v1/messages/count_tokens` separately, so any model exposing the native + // Anthropic `/v1/messages` surface is worth trying. + let native_count = state .model_supports_endpoint(&translated, "/v1/messages/count_tokens") .await - { + || state + .model_supports_endpoint(&translated, "/v1/messages") + .await; + if native_count { let vision = anthropic::has_image(&req); let mut headers = state.copilot_headers(vision).await; headers.insert("anthropic-version", HeaderValue::from_static("2023-06-01")); - if state.model_supports_1m(&translated).await { - headers.insert( - "anthropic-beta", - HeaderValue::from_static("context-1m-2025-08-07"), - ); - } + apply_anthropic_beta( + &state, + &mut headers, + &translated, + &req, + client_beta_header(&client_headers).as_deref(), + ) + .await; let url = format!("{}/v1/messages/count_tokens", state.copilot_base_url()); - let payload = serde_json::to_vec(&req).unwrap_or_default(); + let payload = + serde_json::to_vec(&anthropic::sanitize_anthropic_request(&req)).unwrap_or_default(); if let Some(resp) = util::post_with_retry(&state, &url, headers, payload, "/v1/messages/count_tokens").await { - if resp.status().is_success() { + let status = resp.status(); + if status.is_success() { let parsed: Value = resp.json().await.unwrap_or(json!({"input_tokens": 1})); - return Json(parsed).into_response(); + if parsed.get("input_tokens").is_some() { + return Json(parsed).into_response(); + } + } else { + tracing::debug!( + "[count_tokens] upstream returned {status} for '{translated}'; \ + falling back to a local estimate" + ); } } } - error_response( - StatusCode::BAD_REQUEST, - format!( - "Real token counting is unavailable for model '{original_model}'. The upstream endpoint /v1/messages/count_tokens is not supported or failed." - ), - ) + // Fall back to a local tiktoken estimate. Returning an error here would + // break clients such as Claude Code, which call this endpoint before every + // request purely to decide when to compact the conversation. + // + // A bare text count is systematically low: it ignores the per-message + // framing tokens and the JSON schema of every tool definition. Under- + // reporting makes Claude Code compact too late and then hit a hard + // `prompt token count exceeds the limit` failure, so the estimate adds + // both back. + let tokenizer = state.model_tokenizer(&translated).await; + let total = anthropic::estimate_input_tokens(&req, &tokenizer); + tracing::debug!( + "[count_tokens] upstream counting unavailable for '{original_model}'; \ + returning local {tokenizer} estimate of {total} tokens" + ); + Json(json!({"input_tokens": total, "estimated": true})).into_response() } // --------------------------------------------------------------------------- @@ -1263,6 +1447,7 @@ async fn gemini_generate( .get("completion_tokens") .and_then(|t| t.as_u64()) .unwrap_or(0); + let cost = calculate_cost(&translated, input_tokens, output_tokens); state.store.add(RequestRecord { id: uuid::Uuid::new_v4().to_string(), timestamp: now_iso(), @@ -1284,7 +1469,7 @@ async fn gemini_generate( tools_called: None, is_agent_initiated: Some(agent), prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost(&raw_model, input_tokens, output_tokens)), + estimated_cost_usd: Some(cost), }); Json(gemini_resp).into_response() } else { @@ -1368,24 +1553,23 @@ async fn stream_gemini( let stream = async_stream::stream! { use futures_util::StreamExt; let mut byte_stream = upstream.bytes_stream(); - let mut buf = String::new(); + let mut lines = util::SseLineBuffer::new(); let mut input_tokens = 0u64; let mut output_tokens = 0u64; let mut resp_size = 0usize; let mut finish: Option = None; - let mut debug_resp = String::new(); - while let Some(chunk) = byte_stream.next().await { - let Ok(chunk) = chunk else { break; }; - let chunk_str = String::from_utf8_lossy(&chunk); - if state.is_debug() { debug_resp.push_str(&chunk_str); } - buf.push_str(&chunk_str); - let mut lines: Vec<&str> = buf.split('\n').collect(); - let remainder = lines.pop().unwrap_or("").to_string(); - for line in lines { - let line = line.trim_end_matches('\r'); - if !line.starts_with("data: ") { continue; } - let data = &line[6..]; - if data == "[DONE]" { continue; } + let mut debug_raw: Vec = Vec::new(); + let mut interrupted: Option = None; + loop { + let chunk = match byte_stream.next().await { + Some(Ok(chunk)) => chunk, + Some(Err(e)) => { interrupted = Some(e.to_string()); break; } + None => break, + }; + if state.is_debug() { debug_raw.extend_from_slice(&chunk); } + for line in lines.push(&chunk) { + let Some(data) = util::sse_data(&line) else { continue }; + if data == "[DONE]" || data.is_empty() { continue; } if let Ok(v) = serde_json::from_str::(data) { if let Some(u) = v.get("usage") { if !u.is_null() { @@ -1408,14 +1592,20 @@ async fn stream_gemini( } } } - buf = remainder; } - // Final chunk with finish reason + usage. + // Final chunk with finish reason + usage. A stream cut short reports + // `OTHER` so the client can tell the candidate is incomplete rather + // than treating a partial answer as a finished one. + if let Some(ref reason) = interrupted { + tracing::warn!("[/v1beta/models] upstream stream interrupted: {reason}"); + finish = Some("OTHER".to_string()); + } let usage = json!({"prompt_tokens": input_tokens, "completion_tokens": output_tokens}); let ev = gemini::gemini_stream_final_chunk(finish.as_deref(), &usage, &model_json); let payload = format!("data: {}\n\n", serde_json::to_string(&ev).unwrap_or_default()); resp_size += payload.len(); yield Ok(Bytes::from(payload)); + let debug_resp = String::from_utf8_lossy(&debug_raw).into_owned(); log_debug_response(&state, "/v1beta/models", &debug_resp); state.store.add(RequestRecord { id: uuid::Uuid::new_v4().to_string(), @@ -1423,7 +1613,7 @@ async fn stream_gemini( endpoint: "/v1beta/models".to_string(), model: original_model.clone(), translated_model: (translated != original_model).then_some(translated.clone()), - status_code: status, + status_code: if interrupted.is_some() { 502 } else { status }, request_size: req_size, response_size: resp_size, input_tokens, @@ -1434,11 +1624,11 @@ async fn stream_gemini( message_count: None, tool_count: None, tool_names: None, - stop_reason: None, + stop_reason: interrupted.is_some().then(|| "stream_interrupted".to_string()), tools_called: None, is_agent_initiated: None, prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost(&original_model, input_tokens, output_tokens)), + estimated_cost_usd: Some(calculate_cost(&translated, input_tokens, output_tokens)), }); }; build_sse_response(stream) @@ -1536,6 +1726,49 @@ async fn usage(State(state): State) -> Response { } } +// --------------------------------------------------------------------------- +// Health +// --------------------------------------------------------------------------- + +/// Liveness/readiness probe. Always answers without contacting upstream so it +/// stays cheap enough for a service supervisor to poll frequently. +/// +/// `ready` is true once a Copilot token has been obtained and the model catalog +/// has been loaded. A degraded proxy answers `200` with `ready: false` so +/// monitoring can distinguish "process alive" from "able to serve traffic"; +/// pass `?strict=true` to get a `503` instead when not ready. +async fn health( + State(state): State, + Query(params): Query>, +) -> Response { + let (token_present, token_expires_in) = state.copilot_token_status().await; + let model_count = state.model_count().await; + let stats = state.store.stats(); + let ready = token_present && model_count > 0; + let body = json!({ + "status": if ready { "ok" } else { "degraded" }, + "ready": ready, + "version": env!("CARGO_PKG_VERSION"), + "uptime_seconds": state.uptime_secs(), + "copilot_token": { + "present": token_present, + "expires_in_seconds": token_expires_in, + }, + "models_loaded": model_count, + "requests_served": stats.request_count, + "auth_required": state.api_key().is_some(), + }); + let strict = params + .get("strict") + .map(|v| v == "1" || v.eq_ignore_ascii_case("true")) + .unwrap_or(false); + if strict && !ready { + (StatusCode::SERVICE_UNAVAILABLE, Json(body)).into_response() + } else { + Json(body).into_response() + } +} + // --------------------------------------------------------------------------- // Streaming helpers // --------------------------------------------------------------------------- @@ -1562,76 +1795,138 @@ async fn stream_openai( state: SharedState, url: &str, headers: HeaderMap, - payload: Vec, + mut req: Value, endpoint: &'static str, original_model: String, translated: String, req_size: usize, start: Instant, ) -> Response { - let req_body = state - .is_debug() - .then(|| String::from_utf8_lossy(&payload).into_owned()); - let upstream = state - .http - .post(url) - .headers(headers) - .body(payload) - .send() - .await; - let upstream = match upstream { - Ok(r) => r, - Err(e) => return error_response(StatusCode::GATEWAY_TIMEOUT, e.to_string()), + let req_body = capture_json(&state, &req); + let mut payload = serde_json::to_vec(&req).unwrap_or_default(); + let mut retried_param = false; + let upstream = loop { + let upstream = state + .http + .post(url) + .headers(headers.clone()) + .body(payload.clone()) + .send() + .await; + let upstream = match upstream { + Ok(r) => r, + Err(e) => return error_response(StatusCode::GATEWAY_TIMEOUT, e.to_string()), + }; + let status = upstream.status().as_u16(); + // A non-2xx upstream (e.g. GitHub Models returning 401/403 as JSON when + // the token lacks the `models: read` permission) is not an SSE stream — + // surface it as a normal error instead of forwarding a broken "stream". + if !(200..300).contains(&status) { + let text = upstream.text().await.unwrap_or_default(); + log_debug_response(&state, endpoint, &text); + // Newer OpenAI-family models reject `max_tokens` in favour of + // `max_completion_tokens`; migrate and retry once. + if !retried_param + && util::is_max_tokens_unsupported_error(status, &text) + && rewrite_max_tokens_param(&mut req) + { + tracing::info!("[{endpoint}] retrying stream with max_completion_tokens"); + payload = serde_json::to_vec(&req).unwrap_or_default(); + retried_param = true; + continue; + } + log_error(endpoint, &json!({"model": &translated}), &text, status); + return passthrough_error( + StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_GATEWAY), + text, + ); + } + break upstream; }; let status = upstream.status().as_u16(); - // A non-2xx upstream (e.g. GitHub Models returning 401/403 as JSON when the - // token lacks the `models: read` permission) is not an SSE stream — surface it as a - // normal error instead of forwarding a broken "stream". - if !(200..300).contains(&status) { - let text = upstream.text().await.unwrap_or_default(); - log_debug_response(&state, endpoint, &text); - log_error(endpoint, &json!({"model": &translated}), &text, status); - return passthrough_error( - StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_GATEWAY), - text, - ); - } let model = translated.clone(); let stream = async_stream::stream! { use futures_util::StreamExt; let mut byte_stream = upstream.bytes_stream(); - let mut buf = String::new(); + let mut lines = util::SseLineBuffer::new(); let mut input_tokens = 0u64; let mut output_tokens = 0u64; let mut resp_size = 0usize; - let mut debug_resp = String::new(); - while let Some(chunk) = byte_stream.next().await { - let Ok(chunk) = chunk else { break; }; - let chunk_str = String::from_utf8_lossy(&chunk); - if state.is_debug() { debug_resp.push_str(&chunk_str); } - buf.push_str(&chunk_str); - let mut lines: Vec<&str> = buf.split('\n').collect(); - let remainder = lines.pop().unwrap_or("").to_string(); - for line in lines { - let line = line.trim_end_matches('\r'); - if !line.starts_with("data: ") { continue; } - let data = &line[6..]; + let mut debug_raw: Vec = Vec::new(); + let mut interrupted: Option = None; + let mut saw_done = false; + loop { + let chunk = match byte_stream.next().await { + Some(Ok(chunk)) => chunk, + // The upstream connection dropped mid-response. Everything + // emitted so far is a partial answer; record why and tell the + // client below instead of ending the stream as if it were + // complete. + Some(Err(e)) => { interrupted = Some(e.to_string()); break; } + None => break, + }; + if state.is_debug() { debug_raw.extend_from_slice(&chunk); } + for line in lines.push(&chunk) { + let Some(data) = util::sse_data(&line) else { continue }; if data == "[DONE]" { + saw_done = true; yield Ok::(Bytes::from_static(b"data: [DONE]\n\n")); continue; } - if let Ok(v) = serde_json::from_str::(data) { - if let Some(u) = v.get("usage") { - input_tokens = u.get("prompt_tokens").and_then(|t| t.as_u64()).unwrap_or(input_tokens); - output_tokens = u.get("completion_tokens").and_then(|t| t.as_u64()).unwrap_or(output_tokens); + if data.is_empty() { continue; } + match serde_json::from_str::(data) { + Ok(v) => { + if let Some(u) = v.get("usage") { + input_tokens = u.get("prompt_tokens").and_then(|t| t.as_u64()).unwrap_or(input_tokens); + output_tokens = u.get("completion_tokens").and_then(|t| t.as_u64()).unwrap_or(output_tokens); + } + resp_size += data.len(); + yield Ok(Bytes::from(format!("data: {data}\n\n"))); + } + // Forward payloads we cannot parse rather than dropping + // them: silently deleting a delta removes text from the + // middle of the answer with no trace. + Err(e) => { + tracing::warn!("[{endpoint}] forwarding unparsable SSE payload: {e}"); + resp_size += data.len(); + yield Ok(Bytes::from(format!("data: {data}\n\n"))); + } + } + } + } + // A final event that arrived without its trailing newline is still a + // real event; do not drop it. + if let Some(line) = lines.flush() { + if let Some(data) = util::sse_data(&line) { + if data == "[DONE]" { + saw_done = true; + yield Ok(Bytes::from_static(b"data: [DONE]\n\n")); + } else if !data.is_empty() { + if let Ok(v) = serde_json::from_str::(data) { + if let Some(u) = v.get("usage") { + input_tokens = u.get("prompt_tokens").and_then(|t| t.as_u64()).unwrap_or(input_tokens); + output_tokens = u.get("completion_tokens").and_then(|t| t.as_u64()).unwrap_or(output_tokens); + } } resp_size += data.len(); yield Ok(Bytes::from(format!("data: {data}\n\n"))); } } - buf = remainder; + } + if let Some(ref reason) = interrupted { + tracing::warn!("[{endpoint}] upstream stream interrupted: {reason}"); + let ev = json!({"error": { + "message": format!("Upstream stream ended prematurely: {reason}. The response above is incomplete."), + "type": "upstream_error", + "code": "stream_interrupted" + }}); + yield Ok(Bytes::from(format!("data: {ev}\n\n"))); + } + if !saw_done { + yield Ok(Bytes::from_static(b"data: [DONE]\n\n")); } let _ = model; + let debug_resp = String::from_utf8_lossy(&debug_raw).into_owned(); log_debug_response(&state, endpoint, &debug_resp); state.store.add(RequestRecord { id: uuid::Uuid::new_v4().to_string(), @@ -1639,7 +1934,9 @@ async fn stream_openai( endpoint: endpoint.to_string(), model: original_model.clone(), translated_model: (translated != original_model).then_some(translated.clone()), - status_code: status, + // An interrupted stream is not a successful request, even though + // the response headers were already a 200. + status_code: if interrupted.is_some() { 502 } else { status }, request_size: req_size, response_size: resp_size, input_tokens, @@ -1650,11 +1947,11 @@ async fn stream_openai( message_count: None, // Streaming doesn't have access to parsed req tool_count: None, tool_names: None, - stop_reason: None, + stop_reason: interrupted.is_some().then(|| "stream_interrupted".to_string()), tools_called: None, is_agent_initiated: None, prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost(&original_model, input_tokens, output_tokens)), + estimated_cost_usd: Some(calculate_cost(&translated, input_tokens, output_tokens)), }); }; build_sse_response(stream) @@ -1687,39 +1984,72 @@ async fn stream_responses( Err(e) => return error_response(StatusCode::GATEWAY_TIMEOUT, e.to_string()), }; let status = upstream.status().as_u16(); + // A non-2xx upstream returns a JSON error body, not an SSE stream. Forward + // it as a normal error response instead of wrapping it in a 200 "stream", + // which would make clients treat an auth/quota failure as a valid answer. + if !(200..300).contains(&status) { + let text = upstream.text().await.unwrap_or_default(); + log_debug_response(&state, "/v1/responses", &text); + log_error("/v1/responses", &req, &text, status); + return passthrough_error( + StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_GATEWAY), + text, + ); + } let stream = async_stream::stream! { use futures_util::StreamExt; let mut byte_stream = upstream.bytes_stream(); - let mut buf = String::new(); + let mut lines = util::SseLineBuffer::new(); let mut input_tokens = 0u64; let mut output_tokens = 0u64; let mut resp_size = 0usize; - let mut debug_resp = String::new(); - while let Some(chunk) = byte_stream.next().await { - let Ok(chunk) = chunk else { break; }; + let mut debug_raw: Vec = Vec::new(); + let mut interrupted: Option = None; + let mut completed = false; + loop { + let chunk = match byte_stream.next().await { + Some(Ok(chunk)) => chunk, + Some(Err(e)) => { interrupted = Some(e.to_string()); break; } + None => break, + }; resp_size += chunk.len(); - // Verbatim passthrough of raw bytes. + // Verbatim passthrough of raw bytes. Bytes are never re-decoded on + // this path, so a multi-byte character split across chunks reaches + // the client intact. yield Ok::(Bytes::copy_from_slice(&chunk)); - let chunk_str = String::from_utf8_lossy(&chunk); - if state.is_debug() { debug_resp.push_str(&chunk_str); } - buf.push_str(&chunk_str); - let mut lines: Vec<&str> = buf.split('\n').collect(); - let remainder = lines.pop().unwrap_or("").to_string(); - for line in lines { - let line = line.trim_end_matches('\r'); - if !line.starts_with("data: ") { continue; } - let data = &line[6..]; - if data == "[DONE]" { continue; } + if state.is_debug() { debug_raw.extend_from_slice(&chunk); } + for line in lines.push(&chunk) { + let Some(data) = util::sse_data(&line) else { continue }; + if data == "[DONE]" || data.is_empty() { continue; } if let Ok(v) = serde_json::from_str::(data) { if v.get("type").and_then(|t| t.as_str()) == Some("response.completed") { + completed = true; let usage = v.get("response").and_then(|r| r.get("usage")).cloned().unwrap_or(json!({})); input_tokens = usage.get("input_tokens").and_then(|t| t.as_u64()).unwrap_or(0); output_tokens = usage.get("output_tokens").and_then(|t| t.as_u64()).unwrap_or(0); } } } - buf = remainder; } + // The Responses protocol terminates with `response.completed`. If it + // never arrived the answer is partial, so emit an explicit `error` + // event rather than letting the client treat the truncated output as a + // finished turn. + if interrupted.is_some() || !completed { + let reason = interrupted + .clone() + .unwrap_or_else(|| "upstream closed the stream before response.completed".to_string()); + tracing::warn!("[/v1/responses] incomplete stream: {reason}"); + let ev = json!({ + "type": "error", + "code": "stream_interrupted", + "message": format!("Upstream stream ended prematurely: {reason}. The response above is incomplete."), + "param": Value::Null, + "sequence_number": Value::Null + }); + yield Ok(Bytes::from(format!("event: error\ndata: {ev}\n\n"))); + } + let debug_resp = String::from_utf8_lossy(&debug_raw).into_owned(); log_debug_response(&state, "/v1/responses", &debug_resp); state.store.add(RequestRecord { id: uuid::Uuid::new_v4().to_string(), @@ -1727,7 +2057,7 @@ async fn stream_responses( endpoint: "/v1/responses".to_string(), model: original_model.clone(), translated_model: (translated != original_model).then_some(translated.clone()), - status_code: status, + status_code: if interrupted.is_some() { 502 } else { status }, request_size: req_size, response_size: resp_size, input_tokens, @@ -1738,16 +2068,16 @@ async fn stream_responses( message_count: None, tool_count: None, tool_names: None, - stop_reason: None, + stop_reason: (interrupted.is_some() || !completed) + .then(|| "stream_interrupted".to_string()), tools_called: None, is_agent_initiated: None, prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost(&original_model, input_tokens, output_tokens)), + estimated_cost_usd: Some(calculate_cost(&translated, input_tokens, output_tokens)), }); }; build_sse_response(stream) } - /// Streams a direct Anthropic SSE response back to the client verbatim. async fn stream_anthropic_direct( state: SharedState, @@ -1765,35 +2095,72 @@ async fn stream_anthropic_direct( let mut input_tokens = 0u64; let mut output_tokens = 0u64; let mut resp_size = 0usize; - let mut buf = String::new(); - let mut debug_resp = String::new(); - while let Some(chunk) = byte_stream.next().await { - let Ok(chunk) = chunk else { break; }; + let mut lines = util::SseLineBuffer::new(); + let mut debug_raw: Vec = Vec::new(); + let mut interrupted: Option = None; + let mut saw_message_stop = false; + let mut stop_reason: Option = None; + let mut tools_called: Vec = Vec::new(); + loop { + let chunk = match byte_stream.next().await { + Some(Ok(chunk)) => chunk, + Some(Err(e)) => { interrupted = Some(e.to_string()); break; } + None => break, + }; resp_size += chunk.len(); + // Verbatim passthrough: the client receives the exact upstream + // bytes, so characters split across chunks are never mangled. yield Ok::(Bytes::copy_from_slice(&chunk)); - let chunk_str = String::from_utf8_lossy(&chunk); - if state.is_debug() { debug_resp.push_str(&chunk_str); } - buf.push_str(&chunk_str); - let mut lines: Vec<&str> = buf.split('\n').collect(); - let remainder = lines.pop().unwrap_or("").to_string(); - for line in lines { - let line = line.trim_end_matches('\r'); - if !line.starts_with("data: ") { continue; } - let data = &line[6..]; - if let Ok(v) = serde_json::from_str::(data) { - match v.get("type").and_then(|t| t.as_str()) { - Some("message_start") => { - input_tokens = v.get("message").and_then(|m| m.get("usage")).and_then(|u| u.get("input_tokens")).and_then(|t| t.as_u64()).unwrap_or(0); + if state.is_debug() { debug_raw.extend_from_slice(&chunk); } + for line in lines.push(&chunk) { + let Some(data) = util::sse_data(&line) else { continue }; + if data.is_empty() { continue; } + let Ok(v) = serde_json::from_str::(data) else { continue }; + match v.get("type").and_then(|t| t.as_str()) { + Some("message_start") => { + input_tokens = v.get("message").and_then(|m| m.get("usage")).and_then(|u| u.get("input_tokens")).and_then(|t| t.as_u64()).unwrap_or(0); + } + Some("message_delta") => { + output_tokens = v.get("usage").and_then(|u| u.get("output_tokens")).and_then(|t| t.as_u64()).unwrap_or(output_tokens); + if let Some(sr) = v.get("delta").and_then(|d| d.get("stop_reason")).and_then(|s| s.as_str()) { + stop_reason = Some(sr.to_string()); } - Some("message_delta") => { - output_tokens = v.get("usage").and_then(|u| u.get("output_tokens")).and_then(|t| t.as_u64()).unwrap_or(output_tokens); + } + Some("content_block_start") => { + if let Some(block) = v.get("content_block") { + if block.get("type").and_then(|t| t.as_str()) == Some("tool_use") { + if let Some(name) = block.get("name").and_then(|n| n.as_str()) { + if !tools_called.iter().any(|t| t == name) { + tools_called.push(name.to_string()); + } + } + } } - _ => {} } + Some("message_stop") => saw_message_stop = true, + _ => {} } } - buf = remainder; } + // Anthropic clients wait for `message_stop`. Without it they either + // hang or record the partial text as a completed assistant turn, which + // then poisons every following request in the conversation. + if interrupted.is_some() || !saw_message_stop { + let reason = interrupted + .clone() + .unwrap_or_else(|| "upstream closed the stream before message_stop".to_string()); + tracing::warn!("[/v1/messages] incomplete stream: {reason}"); + let err = json!({ + "type": "error", + "error": { + "type": "api_error", + "message": format!("Upstream stream ended prematurely: {reason}. The response above is incomplete.") + } + }); + yield Ok(Bytes::from(format!("event: error\ndata: {err}\n\n"))); + stop_reason = Some("stream_interrupted".to_string()); + } + let debug_resp = String::from_utf8_lossy(&debug_raw).into_owned(); log_debug_response(&state, "/v1/messages", &debug_resp); state.store.add(RequestRecord { id: uuid::Uuid::new_v4().to_string(), @@ -1801,22 +2168,22 @@ async fn stream_anthropic_direct( endpoint: "/v1/messages".to_string(), model: original_model.clone(), translated_model: (translated != original_model).then_some(translated.clone()), - status_code: status, + status_code: if interrupted.is_some() { 502 } else { status }, request_size: req_size, response_size: resp_size, input_tokens, output_tokens, duration: elapsed_secs(start), request_body: req_body, - response_body: if state.is_debug() { Some(debug_resp.clone()) } else { None }, + response_body: if state.is_debug() { Some(debug_resp) } else { None }, message_count: None, tool_count: None, tool_names: None, - stop_reason: extract_stop_reason_from_sse(&debug_resp), - tools_called: (!extract_tools_called_from_sse(&debug_resp).is_empty()).then(|| extract_tools_called_from_sse(&debug_resp)), + stop_reason, + tools_called: (!tools_called.is_empty()).then_some(tools_called), is_agent_initiated: None, prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost(&original_model, input_tokens, output_tokens)), + estimated_cost_usd: Some(calculate_cost(&translated, input_tokens, output_tokens)), }); }; build_sse_response(stream) @@ -1868,40 +2235,66 @@ async fn stream_anthropic_translated( let stream = async_stream::stream! { use futures_util::StreamExt; let mut byte_stream = upstream.bytes_stream(); - let mut buf = String::new(); + let mut lines = util::SseLineBuffer::new(); let mut conv = AnthropicStreamState::new(); let mut chunks: Vec = Vec::new(); let mut input_tokens = 0u64; let mut output_tokens = 0u64; let mut resp_size = 0usize; - let mut debug_resp = String::new(); - while let Some(chunk) = byte_stream.next().await { - let Ok(chunk) = chunk else { break; }; - let chunk_str = String::from_utf8_lossy(&chunk); - if state.is_debug() { debug_resp.push_str(&chunk_str); } - buf.push_str(&chunk_str); - let mut lines: Vec<&str> = buf.split('\n').collect(); - let remainder = lines.pop().unwrap_or("").to_string(); - for line in lines { - let line = line.trim_end_matches('\r'); - if !line.starts_with("data: ") { continue; } - let data = &line[6..]; - if data == "[DONE]" { continue; } - if let Ok(v) = serde_json::from_str::(data) { - if let Some(u) = v.get("usage") { - input_tokens = u.get("prompt_tokens").and_then(|t| t.as_u64()).unwrap_or(input_tokens); - output_tokens = u.get("completion_tokens").and_then(|t| t.as_u64()).unwrap_or(output_tokens); - } - chunks.push(v.clone()); - for event in conv.process(&v) { - let ev_type = event.get("type").and_then(|t| t.as_str()).unwrap_or("message"); - let payload = format!("event: {ev_type}\ndata: {}\n\n", serde_json::to_string(&event).unwrap_or_default()); - resp_size += payload.len(); - yield Ok::(Bytes::from(payload)); - } + let mut debug_raw: Vec = Vec::new(); + let mut interrupted: Option = None; + loop { + let chunk = match byte_stream.next().await { + Some(Ok(chunk)) => chunk, + Some(Err(e)) => { interrupted = Some(e.to_string()); break; } + None => break, + }; + if state.is_debug() { debug_raw.extend_from_slice(&chunk); } + for line in lines.push(&chunk) { + let Some(data) = util::sse_data(&line) else { continue }; + if data == "[DONE]" || data.is_empty() { continue; } + let Ok(v) = serde_json::from_str::(data) else { + tracing::warn!("[/v1/messages] skipping unparsable upstream SSE payload"); + continue; + }; + if let Some(u) = v.get("usage") { + input_tokens = u.get("prompt_tokens").and_then(|t| t.as_u64()).unwrap_or(input_tokens); + output_tokens = u.get("completion_tokens").and_then(|t| t.as_u64()).unwrap_or(output_tokens); + } + chunks.push(v.clone()); + for event in conv.process(&v) { + let ev_type = event.get("type").and_then(|t| t.as_str()).unwrap_or("message"); + let payload = format!("event: {ev_type}\ndata: {}\n\n", serde_json::to_string(&event).unwrap_or_default()); + resp_size += payload.len(); + yield Ok::(Bytes::from(payload)); } } - buf = remainder; + } + // The upstream may stop without ever sending a `finish_reason`, in + // which case no `message_stop` was emitted and an Anthropic client + // would wait forever — or, worse, keep the partial text as a finished + // assistant turn. Always close the event sequence explicitly. + if let Some(ref reason) = interrupted { + tracing::warn!("[/v1/messages] upstream stream interrupted: {reason}"); + } + let stop_reason = if interrupted.is_some() { "error" } else { "end_turn" }; + let mut synthesized_stop = false; + for event in conv.finish(stop_reason) { + synthesized_stop = true; + let ev_type = event.get("type").and_then(|t| t.as_str()).unwrap_or("message"); + let payload = format!("event: {ev_type}\ndata: {}\n\n", serde_json::to_string(&event).unwrap_or_default()); + resp_size += payload.len(); + yield Ok(Bytes::from(payload)); + } + if let Some(ref reason) = interrupted { + let err = json!({ + "type": "error", + "error": { + "type": "api_error", + "message": format!("Upstream stream ended prematurely: {reason}. The response above is incomplete.") + } + }); + yield Ok(Bytes::from(format!("event: error\ndata: {err}\n\n"))); } // Fall back to merged usage if streaming chunks did not carry usage. if input_tokens == 0 && output_tokens == 0 { @@ -1911,6 +2304,7 @@ async fn stream_anthropic_translated( output_tokens = usage.get("completion_tokens").and_then(|t| t.as_u64()).unwrap_or(0); } } + let debug_resp = String::from_utf8_lossy(&debug_raw).into_owned(); log_debug_response(&state, "/v1/messages", &debug_resp); state.store.add(RequestRecord { id: uuid::Uuid::new_v4().to_string(), @@ -1918,7 +2312,7 @@ async fn stream_anthropic_translated( endpoint: "/v1/messages".to_string(), model: original_model.clone(), translated_model: (translated != original_model).then_some(translated.clone()), - status_code: status, + status_code: if interrupted.is_some() { 502 } else { status }, request_size: req_size, response_size: resp_size, input_tokens, @@ -1929,21 +2323,77 @@ async fn stream_anthropic_translated( message_count: None, tool_count: None, tool_names: None, - stop_reason: None, + stop_reason: synthesized_stop.then(|| stop_reason.to_string()), tools_called: None, is_agent_initiated: None, prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost(&original_model, input_tokens, output_tokens)), + estimated_cost_usd: Some(calculate_cost(&translated, input_tokens, output_tokens)), }); }; build_sse_response(stream) } +/// How long a stream may stay silent before a keepalive comment is emitted. +/// +/// GitHub's upstream load balancer, and most intermediaries, drop an idle +/// connection at around 60 seconds. A model doing extended thinking can easily +/// produce no tokens for longer than that, which surfaces to the user as +/// `user_request_timeout` or a stream that simply dies mid-answer. +const SSE_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15); + +/// Wraps an SSE stream so a comment line is emitted whenever it goes quiet. +/// +/// `: keepalive` is an SSE comment: the spec requires clients to ignore it, so +/// it holds the connection open without being visible to the consumer. It is +/// only injected at an event boundary — after bytes ending in a blank line, or +/// before anything has been sent — because the verbatim passthrough paths yield +/// raw upstream chunks and splicing a comment into a half-written event would +/// corrupt it. +fn keepalive( + stream: S, +) -> impl futures_util::Stream> +where + S: futures_util::Stream> + Send + 'static, +{ + async_stream::stream! { + use futures_util::StreamExt; + let mut inner = Box::pin(stream); + // True while the last bytes written ended an SSE event. + let mut at_boundary = true; + loop { + tokio::select! { + biased; + item = inner.next() => { + match item { + Some(chunk) => { + if let Ok(bytes) = &chunk { + at_boundary = ends_at_event_boundary(bytes); + } + yield chunk; + } + None => break, + } + } + _ = tokio::time::sleep(SSE_KEEPALIVE_INTERVAL) => { + if at_boundary { + yield Ok(Bytes::from_static(b": keepalive\n\n")); + } + } + } + } + } +} + +/// Whether a chunk ends on an SSE event boundary (a blank line). +fn ends_at_event_boundary(bytes: &[u8]) -> bool { + bytes.ends_with(b"\n\n") || bytes.ends_with(b"\r\n\r\n") +} + fn build_sse_response(stream: S) -> Response where S: futures_util::Stream> + Send + 'static, { - let body = Body::from_stream(stream); + let body = Body::from_stream(keepalive(stream)); let mut resp = Response::new(body); *resp.headers_mut() = sse_headers(); resp @@ -2011,6 +2461,26 @@ async fn openapi_spec() -> Response { "responses": { "200": { "description": "Model list" } } } }, + "/v1/models/{model}": { + "get": { + "summary": "Retrieve a single model", + "parameters": [{ "name": "model", "in": "path", "required": true, "schema": { "type": "string" } }], + "responses": { + "200": { "description": "Model description" }, + "404": { "description": "Unknown model" } + } + } + }, + "/health": { + "get": { + "summary": "Liveness/readiness probe", + "parameters": [{ "name": "strict", "in": "query", "required": false, "schema": { "type": "boolean" }, "description": "Return 503 instead of 200 when the proxy is not ready." }], + "responses": { + "200": { "description": "Health report" }, + "503": { "description": "Not ready (only with strict=true)" } + } + } + }, "/v1beta/models/{model}:generateContent": { "post": { "summary": "Gemini generateContent", @@ -2069,7 +2539,6 @@ fn metrics_label_escape(v: &str) -> String { async fn metrics_openmetrics(State(state): State) -> Response { let stats = state.store.stats(); - let (records, _) = state.store.recent(usize::MAX, 0); let mut req_total_by_labels: HashMap<(String, u16, String), u64> = HashMap::new(); let mut in_tokens_by_model: HashMap = HashMap::new(); @@ -2077,26 +2546,30 @@ async fn metrics_openmetrics(State(state): State) -> Response { let mut duration_sum_by_endpoint: HashMap = HashMap::new(); let mut duration_count_by_endpoint: HashMap = HashMap::new(); let mut cost_total = 0.0_f64; - - for rec in &records { - let model = rec - .translated_model - .as_ref() - .unwrap_or(&rec.model) - .to_string(); - *req_total_by_labels - .entry((rec.endpoint.clone(), rec.status_code, model.clone())) - .or_insert(0) += 1; - *in_tokens_by_model.entry(model.clone()).or_insert(0) += rec.input_tokens; - *out_tokens_by_model.entry(model).or_insert(0) += rec.output_tokens; - *duration_sum_by_endpoint - .entry(rec.endpoint.clone()) - .or_insert(0.0) += rec.duration; - *duration_count_by_endpoint - .entry(rec.endpoint.clone()) - .or_insert(0) += 1; - cost_total += rec.estimated_cost_usd.unwrap_or(0.0); - } + let mut record_count = 0usize; + + state.store.with_records(|records| { + for rec in records { + record_count += 1; + let model = rec + .translated_model + .as_ref() + .unwrap_or(&rec.model) + .to_string(); + *req_total_by_labels + .entry((rec.endpoint.clone(), rec.status_code, model.clone())) + .or_insert(0) += 1; + *in_tokens_by_model.entry(model.clone()).or_insert(0) += rec.input_tokens; + *out_tokens_by_model.entry(model).or_insert(0) += rec.output_tokens; + *duration_sum_by_endpoint + .entry(rec.endpoint.clone()) + .or_insert(0.0) += rec.duration; + *duration_count_by_endpoint + .entry(rec.endpoint.clone()) + .or_insert(0) += 1; + cost_total += rec.estimated_cost_usd.unwrap_or(0.0); + } + }); let mut out = String::new(); out.push_str( @@ -2165,7 +2638,14 @@ async fn metrics_openmetrics(State(state): State) -> Response { "# HELP ghc_proxy_store_records Number of request records currently retained in memory.\n", ); out.push_str("# TYPE ghc_proxy_store_records gauge\n"); - out.push_str(&format!("ghc_proxy_store_records {}\n", records.len())); + out.push_str(&format!("ghc_proxy_store_records {}\n", record_count)); + + out.push_str("# HELP ghc_proxy_uptime_seconds Seconds since the proxy started.\n"); + out.push_str("# TYPE ghc_proxy_uptime_seconds gauge\n"); + out.push_str(&format!( + "ghc_proxy_uptime_seconds {}\n", + state.uptime_secs() + )); out.push_str( "# HELP ghc_proxy_estimated_cost_usd_total Total estimated request cost in USD.\n", @@ -2219,26 +2699,35 @@ async fn api_stats(State(state): State) -> Response { Json(state.store.stats()).into_response() } -async fn api_requests( - State(state): State, - Query(params): Query>, -) -> Response { - let page: usize = params +/// Largest page size a dashboard API will honour. Keeps a hostile `per_page` +/// from forcing an unbounded clone of the request store. +const MAX_PAGE_SIZE: usize = 500; + +/// Parses `page`/`per_page` query parameters into a `(page, per_page, offset)` +/// triple. `per_page` is clamped to `MAX_PAGE_SIZE` and the offset uses +/// saturating arithmetic so extreme values cannot overflow. +fn parse_pagination(params: &HashMap) -> (usize, usize, usize) { + let page = params .get("page") - .and_then(|p| p.parse().ok()) + .and_then(|p| p.parse::().ok()) .unwrap_or(1) .max(1); - let per_page: usize = params + let per_page = params .get("per_page") - .and_then(|p| p.parse().ok()) - .unwrap_or(50); - let offset = (page - 1) * per_page; + .and_then(|p| p.parse::().ok()) + .unwrap_or(50) + .clamp(1, MAX_PAGE_SIZE); + let offset = page.saturating_sub(1).saturating_mul(per_page); + (page, per_page, offset) +} + +async fn api_requests( + State(state): State, + Query(params): Query>, +) -> Response { + let (page, per_page, offset) = parse_pagination(¶ms); let (items, total) = state.store.recent(per_page, offset); - let total_pages = if per_page > 0 { - total.div_ceil(per_page) - } else { - 0 - }; + let total_pages = total.div_ceil(per_page); Json(json!({ "items": items, "total": total, @@ -2250,75 +2739,51 @@ async fn api_requests( } /// Audit API: Returns filtered request records with audit fields. -/// Query params: endpoint=, status=, tool_name=, agent=true|false, page=, per_page= +/// Query params: endpoint=, status=, tool_name=, agent=true|false, model=, page=, per_page= async fn api_audit( State(state): State, Query(params): Query>, ) -> Response { - let page: usize = params - .get("page") - .and_then(|p| p.parse().ok()) - .unwrap_or(1) - .max(1); - let per_page: usize = params - .get("per_page") - .and_then(|p| p.parse().ok()) - .unwrap_or(50); + let (page, per_page, offset) = parse_pagination(¶ms); let endpoint_filter = params.get("endpoint").map(|s| s.as_str()); let status_filter = params.get("status").and_then(|s| s.parse::().ok()); let tool_filter = params.get("tool_name").map(|s| s.as_str()); let agent_filter = params.get("agent").and_then(|s| s.parse::().ok()); + let model_filter = params.get("model").map(|s| s.as_str()); - let (records, _total) = state.store.recent(usize::MAX, 0); - - // Apply filters - let filtered: Vec = records - .into_iter() - .filter(|rec| { - // Endpoint filter - if let Some(ep) = endpoint_filter { - if !rec.endpoint.contains(ep) { - return false; - } + // Filter under the store lock so only the returned page is cloned. + let (items, filtered_total) = state.store.filtered_page(per_page, offset, |rec| { + if let Some(ep) = endpoint_filter { + if !rec.endpoint.contains(ep) { + return false; } - // Status code filter - if let Some(st) = status_filter { - if rec.status_code != st { - return false; - } + } + if let Some(st) = status_filter { + if rec.status_code != st { + return false; } - // Tool name filter - if let Some(tool) = tool_filter { - if let Some(ref tools) = rec.tool_names { - if !tools.iter().any(|t| t.contains(tool)) { - return false; - } - } else { - return false; - } + } + if let Some(tool) = tool_filter { + match rec.tool_names { + Some(ref tools) if tools.iter().any(|t| t.contains(tool)) => {} + _ => return false, } - // Agent filter - if let Some(is_agent) = agent_filter { - if rec.is_agent_initiated != Some(is_agent) { - return false; - } + } + if let Some(is_agent) = agent_filter { + if rec.is_agent_initiated != Some(is_agent) { + return false; } - true - }) - .collect(); + } + if let Some(model) = model_filter { + let served = rec.translated_model.as_deref().unwrap_or(&rec.model); + if !rec.model.contains(model) && !served.contains(model) { + return false; + } + } + true + }); - let filtered_total = filtered.len(); - let offset = (page - 1) * per_page; - let items = filtered - .into_iter() - .skip(offset) - .take(per_page) - .collect::>(); - let total_pages = if per_page > 0 { - filtered_total.div_ceil(per_page) - } else { - 0 - }; + let total_pages = filtered_total.div_ceil(per_page); Json(json!({ "items": items, @@ -2332,45 +2797,47 @@ async fn api_audit( /// Audit summary API: Returns aggregated statistics about tools, stop reasons, and costs. async fn api_audit_summary(State(state): State) -> Response { - let (records, _) = state.store.recent(usize::MAX, 0); - let mut tool_usage: HashMap = HashMap::new(); let mut stop_reason_counts: HashMap = HashMap::new(); let mut total_cost = 0.0; let mut agent_count = 0usize; let mut cache_hit_count = 0usize; let mut cache_write_count = 0usize; - - for rec in &records { - // Tool usage aggregation - if let Some(ref tools) = rec.tool_names { - for tool in tools { - *tool_usage.entry(tool.clone()).or_insert(0) += 1; + let mut record_count = 0usize; + + state.store.with_records(|records| { + for rec in records { + record_count += 1; + // Tool usage aggregation + if let Some(ref tools) = rec.tool_names { + for tool in tools { + *tool_usage.entry(tool.clone()).or_insert(0) += 1; + } } - } - // Stop reason aggregation - if let Some(ref sr) = rec.stop_reason { - *stop_reason_counts.entry(sr.clone()).or_insert(0) += 1; - } + // Stop reason aggregation + if let Some(ref sr) = rec.stop_reason { + *stop_reason_counts.entry(sr.clone()).or_insert(0) += 1; + } - // Cost aggregation - if let Some(cost) = rec.estimated_cost_usd { - total_cost += cost; - } + // Cost aggregation + if let Some(cost) = rec.estimated_cost_usd { + total_cost += cost; + } - // Agent tracking - if rec.is_agent_initiated == Some(true) { - agent_count += 1; - } + // Agent tracking + if rec.is_agent_initiated == Some(true) { + agent_count += 1; + } - // Cache tracking - if rec.prompt_cache_hit == Some(true) { - cache_hit_count += 1; - } else if rec.prompt_cache_hit == Some(false) { - cache_write_count += 1; + // Cache tracking + if rec.prompt_cache_hit == Some(true) { + cache_hit_count += 1; + } else if rec.prompt_cache_hit == Some(false) { + cache_write_count += 1; + } } - } + }); // Sort tools by usage let mut tools_sorted: Vec<_> = tool_usage.into_iter().collect(); @@ -2382,10 +2849,10 @@ async fn api_audit_summary(State(state): State) -> Response { stop_reasons_sorted.sort_by_key(|b| std::cmp::Reverse(b.1)); Json(json!({ - "total_requests": records.len(), + "total_requests": record_count, "agent_initiated": agent_count, - "total_cost_usd": (total_cost * 100.0).round() / 100.0, - "avg_cost_usd": if !records.is_empty() { (total_cost / records.len() as f64 * 100.0).round() / 100.0 } else { 0.0 }, + "total_cost_usd": (total_cost * 10_000.0).round() / 10_000.0, + "avg_cost_usd": if record_count > 0 { (total_cost / record_count as f64 * 10_000.0).round() / 10_000.0 } else { 0.0 }, "top_tools": top_tools, "stop_reasons": stop_reasons_sorted, "cache_hits": cache_hit_count, @@ -2454,4 +2921,114 @@ mod tests { let h = HeaderMap::new(); assert_eq!(presented_api_key(&h), None); } + + #[test] + fn health_path_is_not_protected() { + assert!(!is_protected_path("/health")); + assert!(!is_protected_path("/openapi.json")); + // Single-model retrieval is an LLM endpoint and must stay guarded. + assert!(is_protected_path("/v1/models/claude-opus-4.8")); + } + + #[test] + fn model_rates_prefer_the_most_specific_family() { + // `gpt-4o` must not be swallowed by the broader `gpt-4` arm. + assert_eq!(model_rates("gpt-4o"), (0.005, 0.015)); + assert_eq!(model_rates("gpt-4-turbo"), (0.03, 0.06)); + assert_eq!(model_rates("gpt-4o-mini"), (0.00015, 0.0006)); + // Publisher-qualified GitHub Models ids price like their base model. + assert_eq!(model_rates("openai/gpt-4o"), model_rates("gpt-4o")); + // Claude tiers are distinct. + assert_eq!(model_rates("claude-opus-4.8"), (0.015, 0.075)); + assert_eq!(model_rates("claude-haiku-4.5"), (0.0008, 0.004)); + // Unknown models fall back rather than panicking. + assert_eq!(model_rates("something-new"), (0.0005, 0.0015)); + } + + #[test] + fn cost_scales_with_tokens() { + let cost = calculate_cost("claude-opus-4.8", 1000, 1000); + assert!((cost - (0.015 + 0.075)).abs() < 1e-9); + assert_eq!(calculate_cost("gpt-4o", 0, 0), 0.0); + } + + #[test] + fn pagination_is_clamped_and_overflow_safe() { + let mut p = HashMap::new(); + // Defaults. + assert_eq!(parse_pagination(&p), (1, 50, 0)); + // Page 3 of 10. + p.insert("page".into(), "3".into()); + p.insert("per_page".into(), "10".into()); + assert_eq!(parse_pagination(&p), (3, 10, 20)); + // Hostile values must clamp instead of overflowing usize. + p.insert("page".into(), usize::MAX.to_string()); + p.insert("per_page".into(), usize::MAX.to_string()); + let (page, per_page, offset) = parse_pagination(&p); + assert_eq!(page, usize::MAX); + assert_eq!(per_page, MAX_PAGE_SIZE); + assert_eq!(offset, usize::MAX); + // Zero/garbage page sizes fall back to a usable window. + p.insert("page".into(), "0".into()); + p.insert("per_page".into(), "0".into()); + assert_eq!(parse_pagination(&p), (1, 1, 0)); + p.insert("per_page".into(), "abc".into()); + assert_eq!(parse_pagination(&p).1, 50); + } + + #[test] + fn model_routes_do_not_conflict() { + // `/v1/models/full/` and `/v1/models/{model_id}` are registered on the + // same prefix; building the routes must not panic on a route conflict. + async fn noop() {} + let _: Router = Router::new() + .route("/v1/models", get(noop)) + .route("/v1/models/full/", get(noop)) + .route("/v1/models/{model_id}", get(noop)) + .route("/models", get(noop)) + .route("/models/full/", get(noop)) + .route("/models/{model_id}", get(noop)); + } + + #[test] + fn max_tokens_is_renamed_for_newer_models() { + let mut req = json!({"model": "gpt-5.3-codex", "max_tokens": 4096}); + assert!(rewrite_max_tokens_param(&mut req)); + assert!(req.get("max_tokens").is_none()); + assert_eq!(req["max_completion_tokens"], 4096); + // Nothing left to rename, so no pointless retry. + assert!(!rewrite_max_tokens_param(&mut req)); + } + + #[test] + fn max_tokens_rename_drops_the_alias_when_both_are_present() { + let mut req = json!({"max_tokens": 10, "max_completion_tokens": 20}); + assert!(rewrite_max_tokens_param(&mut req)); + assert!(req.get("max_tokens").is_none()); + assert_eq!(req["max_completion_tokens"], 20); + } + + #[test] + fn keepalive_only_injects_at_event_boundaries() { + assert!(ends_at_event_boundary(b"data: {}\n\n")); + assert!(ends_at_event_boundary(b"data: {}\r\n\r\n")); + // Mid-event: injecting here would corrupt the passthrough stream. + assert!(!ends_at_event_boundary(b"data: {\"partial\":")); + assert!(!ends_at_event_boundary(b"data: {}\n")); + assert!(!ends_at_event_boundary(b"")); + } + + #[test] + fn client_beta_header_is_read() { + let mut h = HeaderMap::new(); + assert_eq!(client_beta_header(&h), None); + h.insert( + "anthropic-beta", + HeaderValue::from_static("claude-code-20250219"), + ); + assert_eq!( + client_beta_header(&h).as_deref(), + Some("claude-code-20250219") + ); + } } diff --git a/src/state.rs b/src/state.rs index e8a2346..2319329 100644 --- a/src/state.rs +++ b/src/state.rs @@ -35,6 +35,8 @@ pub struct AppState { /// Per-process session id (`vscode-sessionid` header): a UUID followed by a /// 13-digit millisecond timestamp, matching the real Copilot client format. pub session_id: String, + /// Instant the process started serving, used to report uptime on `/health`. + pub started_at: Instant, } pub type SharedState = Arc; @@ -56,7 +58,12 @@ fn now_millis() -> u128 { impl AppState { pub fn new(config: Config, github_token: String) -> Self { + // No global request timeout: SSE responses are long-lived streams. Only + // the connect phase and idle pooled connections are bounded so a dead + // upstream cannot wedge a request forever. let http = reqwest::Client::builder() + .connect_timeout(Duration::from_secs(30)) + .pool_idle_timeout(Duration::from_secs(90)) .build() .expect("failed to build HTTP client"); AppState { @@ -73,9 +80,51 @@ impl AppState { last_request: Mutex::new(None), machine_id: auth::load_or_create_machine_id(), session_id: format!("{}{}", uuid::Uuid::new_v4(), now_millis()), + started_at: Instant::now(), } } + /// Seconds this process has been running. + pub fn uptime_secs(&self) -> u64 { + self.started_at.elapsed().as_secs() + } + + /// Copilot token status as `(present, seconds_until_expiry)`. The remaining + /// lifetime is zero when the token is missing or already expired. + pub async fn copilot_token_status(&self) -> (bool, u64) { + let tokens = self.tokens.lock().await; + let present = tokens.copilot_token.is_some(); + let remaining = tokens.expires_at.saturating_sub(now_secs()); + (present, if present { remaining } else { 0 }) + } + + /// Number of models currently held in the catalog cache. + pub async fn model_count(&self) -> usize { + self.models + .read() + .await + .as_ref() + .and_then(|m| m.get("data")) + .and_then(|d| d.as_array()) + .map(|a| a.len()) + .unwrap_or(0) + } + + /// Looks up a single entry from the cached model catalog by id. + pub async fn find_model(&self, id: &str) -> Option { + self.models + .read() + .await + .as_ref() + .and_then(|m| m.get("data")) + .and_then(|d| d.as_array()) + .and_then(|arr| { + arr.iter() + .find(|m| m.get("id").and_then(|i| i.as_str()) == Some(id)) + .cloned() + }) + } + pub fn config_snapshot(&self) -> Config { self.config.read().unwrap().clone() } @@ -428,6 +477,26 @@ impl AppState { self.model_supports_endpoint(model, "/v1/messages").await } + /// Maximum output tokens advertised for a model + /// (`capabilities.limits.max_output_tokens`). Used to fill in a missing + /// `max_tokens`, which some Copilot models reject the request without. + pub async fn model_max_output_tokens(&self, model: &str) -> Option { + self.models + .read() + .await + .as_ref() + .and_then(|m| m.get("data")) + .and_then(|d| d.as_array()) + .and_then(|arr| { + arr.iter() + .find(|m| m.get("id").and_then(|i| i.as_str()) == Some(model)) + .and_then(|m| m.get("capabilities")) + .and_then(|c| c.get("limits")) + .and_then(|l| l.get("max_output_tokens")) + .and_then(|t| t.as_u64()) + }) + } + /// Returns the tokenizer name advertised by the model's catalog entry /// (`capabilities.tokenizer`), used to pick a tiktoken encoder for local /// token estimation. Falls back to `cl100k_base` when unknown. diff --git a/src/store.rs b/src/store.rs index 643a192..3b15934 100644 --- a/src/store.rs +++ b/src/store.rs @@ -115,4 +115,124 @@ impl RequestStore { .collect(); (items, total) } + + /// Runs `f` over the retained records (newest first) while holding the + /// store lock. Used by aggregation endpoints (`/metrics`, the audit APIs) + /// so they never clone the entire ring buffer — records may carry captured + /// request/response bodies when debug mode is enabled. + pub fn with_records( + &self, + f: impl FnOnce(&mut dyn Iterator) -> R, + ) -> R { + let inner = self.inner.lock().unwrap(); + let mut iter = inner.records.iter(); + f(&mut iter) + } + + /// Returns the records matching `predicate` for a page, plus the total + /// number of matches. Filtering happens under the lock so only the records + /// actually returned are cloned. + pub fn filtered_page( + &self, + per_page: usize, + offset: usize, + predicate: impl Fn(&RequestRecord) -> bool, + ) -> (Vec, usize) { + let inner = self.inner.lock().unwrap(); + let total = inner.records.iter().filter(|r| predicate(r)).count(); + let items = inner + .records + .iter() + .filter(|r| predicate(r)) + .skip(offset) + .take(per_page) + .cloned() + .collect(); + (items, total) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn record(endpoint: &str, status: u16) -> RequestRecord { + RequestRecord { + id: endpoint.to_string(), + timestamp: String::new(), + endpoint: endpoint.to_string(), + model: "m".into(), + translated_model: None, + status_code: status, + request_size: 1, + response_size: 2, + input_tokens: 3, + output_tokens: 4, + duration: 0.5, + request_body: None, + response_body: None, + message_count: None, + tool_count: None, + tool_names: None, + stop_reason: None, + tools_called: None, + is_agent_initiated: None, + estimated_cost_usd: None, + prompt_cache_hit: None, + } + } + + #[test] + fn evicts_oldest_beyond_capacity() { + let store = RequestStore::new(2); + store.add(record("a", 200)); + store.add(record("b", 200)); + store.add(record("c", 200)); + let (items, total) = store.recent(10, 0); + assert_eq!(total, 2); + // Newest first; the oldest record was evicted. + assert_eq!(items[0].endpoint, "c"); + assert_eq!(items[1].endpoint, "b"); + // Aggregate stats keep counting evicted requests. + assert_eq!(store.stats().request_count, 3); + assert_eq!(store.stats().total_input_tokens, 9); + } + + #[test] + fn filtered_page_paginates_matches_only() { + let store = RequestStore::new(10); + store.add(record("/v1/messages", 200)); + store.add(record("/v1/chat/completions", 500)); + store.add(record("/v1/messages", 429)); + + let (items, total) = store.filtered_page(10, 0, |r| r.endpoint == "/v1/messages"); + assert_eq!(total, 2); + assert_eq!(items.len(), 2); + + // Offset applies to the filtered set, not the raw store. + let (items, total) = store.filtered_page(1, 1, |r| r.endpoint == "/v1/messages"); + assert_eq!(total, 2); + assert_eq!(items.len(), 1); + assert_eq!(items[0].status_code, 200); + } + + #[test] + fn with_records_visits_every_record() { + let store = RequestStore::new(10); + store.add(record("a", 200)); + store.add(record("b", 404)); + let (count, errors) = store.with_records(|records| { + let mut count = 0; + let mut errors = 0; + for r in records { + count += 1; + if r.status_code >= 400 { + errors += 1; + } + } + (count, errors) + }); + assert_eq!(count, 2); + assert_eq!(errors, 1); + } } diff --git a/src/util.rs b/src/util.rs index 9eb57ad..c13f027 100644 --- a/src/util.rs +++ b/src/util.rs @@ -1,10 +1,84 @@ -//! Miscellaneous helpers: orphaned tool-result detection/removal and an HTTP -//! request retry wrapper with exponential backoff. +//! Miscellaneous helpers: orphaned tool-result detection/removal, SSE line +//! buffering, and an HTTP request retry wrapper with exponential backoff. use crate::state::AppState; use serde_json::Value; use std::time::Duration; +/// Incremental, chunk-boundary-safe splitter for SSE byte streams. +/// +/// Upstream responses arrive as arbitrary byte chunks that do **not** respect +/// either line or UTF-8 character boundaries. Decoding each chunk with +/// `String::from_utf8_lossy` as it arrives corrupts any multi-byte character +/// (CJK, emoji, box-drawing, …) that straddles two chunks, replacing it with +/// `U+FFFD` — and if that character sits inside a `data:` payload, the whole +/// event stops parsing as JSON and is dropped, silently deleting text from the +/// middle of a response. +/// +/// This buffer keeps raw bytes until a `\n` is seen and only then decodes, so a +/// partial character simply waits for the rest of its bytes. Anything left over +/// when the stream ends is returned by [`SseLineBuffer::flush`], so a final +/// event that arrives without a trailing newline is not lost either. +#[derive(Default)] +pub struct SseLineBuffer { + buf: Vec, +} + +impl SseLineBuffer { + pub fn new() -> Self { + Self::default() + } + + /// Appends a chunk and returns every complete line it completed, with the + /// trailing `\r` (CRLF streams) already stripped. + pub fn push(&mut self, chunk: &[u8]) -> Vec { + self.buf.extend_from_slice(chunk); + let mut lines = Vec::new(); + let mut start = 0usize; + while let Some(pos) = self.buf[start..].iter().position(|&b| b == b'\n') { + let end = start + pos; + lines.push(decode_line(&self.buf[start..end])); + start = end + 1; + } + if start > 0 { + self.buf.drain(..start); + } + lines + } + + /// Returns any buffered bytes not terminated by a newline, consuming them. + /// Called once the upstream stream ends so a final unterminated event is + /// still delivered. + pub fn flush(&mut self) -> Option { + if self.buf.is_empty() { + return None; + } + let line = decode_line(&self.buf); + self.buf.clear(); + if line.is_empty() { + None + } else { + Some(line) + } + } +} + +/// Decodes one complete SSE line. A complete line is always valid UTF-8 in +/// practice, so the lossy conversion never actually substitutes anything here. +fn decode_line(bytes: &[u8]) -> String { + String::from_utf8_lossy(bytes) + .trim_end_matches('\r') + .to_string() +} + +/// Returns the payload of an SSE `data:` line, or `None` for any other field +/// (`event:`, `id:`, `retry:`, comments, blank separators). +pub fn sse_data(line: &str) -> Option<&str> { + let rest = line.strip_prefix("data:")?; + // The spec allows exactly one optional space after the colon. + Some(rest.strip_prefix(' ').unwrap_or(rest)) +} + /// Detects whether an upstream error indicates an orphaned `tool_use_id` in a /// `tool_result` block. pub fn is_orphaned_tool_error(status: u16, body: &str) -> bool { @@ -18,15 +92,25 @@ pub fn is_thinking_enabled_unsupported_error(status: u16, body: &str) -> bool { status == 400 && body.contains("thinking.type.enabled") && body.contains("adaptive") } -/// Fetches the latest stable VS Code version from the Arch User Repository -/// `visual-studio-code-bin` PKGBUILD, parsing the `pkgver=` line. Returns -/// `None` on any network or parse error so callers can fall back to a -/// configured default. +/// Detects the upstream 400 returned by newer OpenAI-family models that +/// replaced `max_tokens` with `max_completion_tokens`. +pub fn is_max_tokens_unsupported_error(status: u16, body: &str) -> bool { + status == 400 && body.contains("max_tokens") && body.contains("max_completion_tokens") +} + +/// Fetches the latest stable VS Code version from Microsoft's own update +/// service (`update.code.visualstudio.com`), the same endpoint the editor's +/// updater uses. Returns `None` on any network or parse error so callers can +/// fall back to a configured default. +/// +/// This deliberately does **not** scrape the Arch User Repository PKGBUILD: +/// the AUR maintainers asked proxies of this kind to stop, having become the +/// single most-requested endpoint on their service. pub async fn fetch_latest_vscode_version(client: &reqwest::Client) -> Option { - const URL: &str = - "https://aur.archlinux.org/cgit/aur.git/plain/PKGBUILD?h=visual-studio-code-bin"; + const URL: &str = "https://update.code.visualstudio.com/api/releases/stable"; let resp = client .get(URL) + .header("Accept", "application/json") .timeout(Duration::from_secs(5)) .send() .await @@ -34,17 +118,22 @@ pub async fn fetch_latest_vscode_version(client: &reqwest::Client) -> Option = resp.json().await.ok()?; + // Newest first; keep only plain `major.minor.patch` entries so insider or + // recovery builds never leak into the `Editor-Version` header. + releases.into_iter().find(|v| is_release_version(v)) +} + +/// Whether a string is a plain `major.minor.patch` version. +fn is_release_version(v: &str) -> bool { + let mut parts = 0; + for part in v.split('.') { + if part.is_empty() || !part.bytes().all(|b| b.is_ascii_digit()) { + return false; } + parts += 1; } - None + parts == 3 } /// Extracts orphaned tool-use ids referenced in an error message. @@ -219,4 +308,93 @@ mod tests { "some other validation error" )); } + + #[test] + fn sse_buffer_splits_on_newlines_only() { + let mut b = SseLineBuffer::new(); + assert_eq!( + b.push(b"data: one\ndata: two\n"), + vec!["data: one", "data: two"] + ); + // A partial line is held back until its newline arrives. + assert!(b.push(b"data: thr").is_empty()); + assert_eq!(b.push(b"ee\n"), vec!["data: three"]); + assert!(b.flush().is_none()); + } + + #[test] + fn sse_buffer_strips_crlf() { + let mut b = SseLineBuffer::new(); + assert_eq!(b.push(b"data: x\r\n\r\n"), vec!["data: x", ""]); + } + + #[test] + fn sse_buffer_survives_split_multibyte_characters() { + // "你好" split in the middle of the first character's 3 bytes. Decoding + // each chunk independently would yield U+FFFD and destroy the payload. + let full = "data: {\"t\":\"你好🎉\"}\n".as_bytes().to_vec(); + let cut = 9; // inside the first multi-byte character + let mut b = SseLineBuffer::new(); + assert!(b.push(&full[..cut]).is_empty()); + let lines = b.push(&full[cut..]); + assert_eq!(lines, vec!["data: {\"t\":\"你好🎉\"}"]); + assert!(!lines[0].contains('\u{FFFD}')); + // The payload still parses as JSON, which is what keeps the delta alive. + let data = sse_data(&lines[0]).unwrap(); + let v: Value = serde_json::from_str(data).unwrap(); + assert_eq!(v["t"], "你好🎉"); + } + + #[test] + fn sse_buffer_survives_byte_at_a_time_delivery() { + let payload = "data: {\"t\":\"日本語テキスト\"}\n"; + let mut b = SseLineBuffer::new(); + let mut lines = Vec::new(); + for byte in payload.as_bytes() { + lines.extend(b.push(&[*byte])); + } + assert_eq!(lines.len(), 1); + let v: Value = serde_json::from_str(sse_data(&lines[0]).unwrap()).unwrap(); + assert_eq!(v["t"], "日本語テキスト"); + } + + #[test] + fn sse_buffer_flush_returns_unterminated_tail() { + let mut b = SseLineBuffer::new(); + assert!(b.push(b"data: last-event-no-newline").is_empty()); + assert_eq!(b.flush().as_deref(), Some("data: last-event-no-newline")); + // Flushing twice is a no-op. + assert!(b.flush().is_none()); + } + + #[test] + fn sse_data_extracts_payload_only() { + assert_eq!(sse_data("data: {}"), Some("{}")); + // No space after the colon is also valid. + assert_eq!(sse_data("data:{}"), Some("{}")); + // Only a single leading space is consumed. + assert_eq!(sse_data("data: x"), Some(" x")); + assert_eq!(sse_data("event: message"), None); + assert_eq!(sse_data(": comment"), None); + assert_eq!(sse_data(""), None); + } + + #[test] + fn release_versions_exclude_insider_builds() { + assert!(is_release_version("1.130.0")); + assert!(is_release_version("1.99.12")); + assert!(!is_release_version("1.130.0-insider")); + assert!(!is_release_version("1.130")); + assert!(!is_release_version("1.130.0.1")); + assert!(!is_release_version("")); + } + + #[test] + fn detects_max_tokens_parameter_migration() { + let body = "Unsupported parameter: 'max_tokens' is not supported with this model. \ + Use 'max_completion_tokens' instead."; + assert!(is_max_tokens_unsupported_error(400, body)); + assert!(!is_max_tokens_unsupported_error(200, body)); + assert!(!is_max_tokens_unsupported_error(400, "some other error")); + } }