diff --git a/.cursor/rules/deepseek.mdc b/.cursor/rules/deepseek.mdc index 318989c..db0154b 100644 --- a/.cursor/rules/deepseek.mdc +++ b/.cursor/rules/deepseek.mdc @@ -53,4 +53,15 @@ Official first call: https://api-docs.deepseek.com/ | ----- | ---- | ----- | | Upstream 401 after alias maps to DeepSeek | Only Moonshot key saved | Require `deepseekKeyEncrypted` when resolved provider is `deepseek` (doctor / set-key) | - +### `Downgraded flash turns are nearly free via prompt cache` + +In a multi-turn agent flow, the conversation context grows monotonically, so +every successive tool-result request re-sends nearly the whole prior context. +DeepSeek auto-caches this prefix, making iterated cheap turns extraordinarily +cheap. + +| Symptom | Cause | Fix | +| ----- | ---- | ----- | +| `deepseek-v4-flash` tool-result turns log `est_usd ≈ 0.000` with ~99% `cache_hit_tokens` (e.g. 40,576 hit / 46 miss) while `tokens_in` climbs each turn | DeepSeek prompt caching keys on a matching prefix; because each turn appends one tool result to unchanged prior messages, the prefix (and thus the cache) is reused | Expected and desirable — do not treat near-zero flash-cost turns in a stable conversation as under-counting. It means the per-turn downgrade is already paying for itself; run medium write-turns on `pro` and small/read-only on `flash` without fear of cost blow-up | + +`` diff --git a/.cursor/rules/gateway.mdc b/.cursor/rules/gateway.mdc index 4d2bff7..39b70c6 100644 --- a/.cursor/rules/gateway.mdc +++ b/.cursor/rules/gateway.mdc @@ -164,9 +164,35 @@ when the model has a mapped cheaper variant: | Symptom | Cause | Fix | | ------- | ----- | --- | | "fix linter errors in some files" stays on the expensive model instead of flash | `isStrongEditing` matches "fix" and blocks the `automation` downgrade | Let lint tasks classify as automation even with a strong-editing word present (`isLintTask` → bypass the `!isStrongEditing` guard) | -| `glm-5.2` / `kimi-k3` never downgrade on automation tasks | Model has no mapped cheaper variant and falls back to `deepseek-v4-flash` (`defaultSubagentModel`). Override also changes provider (re-resolved via `ResolveModel`), so cross-provider routing works: e.g. `glm-5.2` → `deepseek-v4-flash` routes to `provider=deepseek`, not Z.AI. To add a same-provider cheap variant (e.g. `glm-4.7` as flash for `glm-5.2`), uncomment the entry in `modelOverrideMap` — it takes priority over the universal fallback. | +| `glm-5.2` / `kimi-k3` never downgrade on automation tasks | Model has no mapped cheaper variant and falls back to `deepseek-v4-flash` (`defaultFlashModel`). Override also changes provider (re-resolved via `ResolveModel`), so cross-provider routing works: e.g. `glm-5.2` → `deepseek-v4-flash` routes to `provider=deepseek`, not Z.AI. To add a same-provider cheap variant (e.g. `glm-4.7` as flash for `glm-5.2`), uncomment the entry in `modelOverrideMap` — it takes priority over the universal fallback. | - +#### Per-turn routing tiers (`flash` / `pro` / `keep`) + +Within one multi-step agent flow, each tool-result request is routed by an +`OverrideTier` decided from the **tool result size + the tool name** (not the +content class). Real 71s GLM 5.2 refactor session: 7/10 turns went `flash`, +2 decision-heavy `pro`, initial prompt + one large result stayed `glm-5.2`. + +| Signal | Tier | Model | Rationale | +| ------ | ---- | ----- | --------- | +| tool result, size `small` | flash | `deepseek-v4-flash` | short output → cheap interpretation | +| tool result, size `medium`, read-only tool (`Read`/`Grep`/`Glob`/search) | flash | `deepseek-v4-flash` | just reading → no pro needed | +| tool result, size `medium`, write/decision tool (`Shell`/`StrReplace`/tests/unknown) | pro | `deepseek-v4-pro` | may decide state changes → needs a reasoning model | +| tool result, size `large` | keep | original | big output needs pro interpretation | +| user prompt | keep | original | content classifier decides (see above) | + +Route on the resolved `tool_name` (matched from the last `tool_call_id` back to +the assistant `tool_calls[].function.name`). Treat any non-read-only tool name — +**including empty/unknown** — as write/decision → pro. A cheap tier is only safe +when the result can be interpreted without deep reasoning. + +#### Flash verbosity cap (output token ceiling on cheap turns) + +| Symptom | Cause | Fix | +| ------- | ----- | --- | +| Agent flow "takes too long / overthought a lot" for a simple change | On `flash` turn the model is over-confident and emits 4–16k output tokens (~20–106s) rather than a terse next step | Cap cheap tiers' `max_tokens` to `toolVerbosityCap` (1500) so flash turns stay terse. Apply the cap even when the model override is a no-op (flow already on flash). Cheap turns should land at 78–600 output tokens, not thousands | + +`` ### `DiagnosticDump: gated message-array dumps for content-extraction analysis` diff --git a/.cursor/rules/release.mdc b/.cursor/rules/release.mdc index b3bb872..311eb40 100644 --- a/.cursor/rules/release.mdc +++ b/.cursor/rules/release.mdc @@ -9,8 +9,8 @@ alwaysApply: false ## Release Process - Version is tracked in `VERSION`. -- `make release-patch` / `make release-minor` / `make release-major` bumps the - version in `VERSION`, commits, tags, and pushes. +- `make release-patch` / `make release-minor` bumps the version in `VERSION`, + commits, tags, and pushes. - **Tag push triggers CI:** the `.github/workflows/release-binaries.yaml` workflow runs on push to `main`, pull requests to `main`, and pushes of tags matching `v*`. diff --git a/.cursor/skills/open-pr/SKILL.md b/.cursor/skills/open-pr/SKILL.md new file mode 100644 index 0000000..46c3bcb --- /dev/null +++ b/.cursor/skills/open-pr/SKILL.md @@ -0,0 +1,159 @@ +--- +name: open-pr +description: > + Create a GitHub PR from the current branch in this Discursive Go repo. + Generates PR description with functional line count, key files to review, + and non-technical summary from the git diff. Supports draft mode. +disable-model-invocation: true +allowed-tools: Bash, Read +--- + +# /open-pr — Create a GitHub Pull Request + +## Call budget + +This skill MUST complete in **≤8 total tool calls**: +- Step 0: 1 Bash — batch state gathering +- Step 1: 1 AskQuestion +- Step 2: 1 Bash — dump diff + stats +- Step 3: 1 Read — read the dumped diff once (100ms, tiny) +- Step 4: 2 Bash — `git push` then `gh pr create` (dependent; can't chain a + heredoc body onto a `&&` push safely) +- Step 5: 1 Bash — cleanup `tmp/diff.diff` +- Step 6: inline report + +Every call beyond 8 (and any Bash call that isn't one of the above) is a bug. +The prior session made ~15 redundant Shell calls (multiple `git status`/`git +log`/`git branch`, per-file `git diff`, `git fetch`, `gh pr view`). None of +those are allowed in this flow. + +## Flow (strict) + +### Step 0: gather state (1 batch Bash call) + +With `working_directory` set to the repo root: + +```bash +git branch -vv && git branch -r && git log -15 --oneline +``` + +From this output, identify: +- Current branch name +- Remote base branch options (pick `main`, plus any other active branches) +- Recent commits (the branch's story) + +**Do NOT run `git status`, `git log`, or `git branch` separately.** + +### Step 1: ask base branch (AskQuestion) + +**Always** present a choice of base branches derived from step 0. Include `main` +and any other visible active branches plus an "Other" option. Never default. + +### Step 2: dump functional diff + stats (1 Bash call) + +Dump to gitignored `./tmp/diff.diff`, plus `--shortstat` and `--stat`: + +```bash +git diff ...HEAD -- ':(exclude)*.md' ':(exclude)*.mdc' ':(exclude)**/test/**' ':(exclude)**/*_test.go' ':(exclude)**/*.txt' ':(exclude).cursor/rules/**' ':(exclude).cursor/skills/**' ':(exclude)examples/**' ':(exclude)VERSION' > ./tmp/diff.diff && echo "---SHORTSTAT---" && git diff ...HEAD --shortstat -- ':(exclude)*.md' ':(exclude)*.mdc' ':(exclude)**/test/**' ':(exclude)**/*_test.go' ':(exclude)**/*.txt' ':(exclude).cursor/rules/**' ':(exclude).cursor/skills/**' ':(exclude)examples/**' ':(exclude)VERSION' && echo "---STAT---" && git diff ...HEAD --stat -- ':(exclude)*.md' ':(exclude)*.mdc' ':(exclude)**/test/**' ':(exclude)**/*_test.go' ':(exclude)**/*.txt' ':(exclude).cursor/rules/**' ':(exclude).cursor/skills/**' ':(exclude)examples/**' ':(exclude)VERSION' && echo "---NONFUNCSTAT---" && git diff ...HEAD --shortstat -- '*.md' '*.mdc' '.cursor/skills/**' '*.txt' +``` + +If the functional diff is empty (zero lines, or only binary `usage.db` type +junk), fall back to the full diff (no pathspec) and treat as a rules/docs-only +PR per that section below. + +### Step 3: read the diff ONCE (1 Read call) + +Read `./tmp/diff.diff` with the Read tool. Derive from it: + +- **Summary bullets (4–6):** what changed and why. Each 1–2 lines. Do not + enumerate markdown/rules/test file changes here. +- **Key files to review:** pick 3–8 from the `--stat` in step 2 (largest first). + +Prefer this single Read over many per-file `git diff` calls. + +### Step 4: push then create PR (2 Bash calls) + +Push the branch: + +```bash +git push -u origin HEAD +``` + +If the branch is already pushed (no-op), that's fine — `gh pr create` still +works. If push fails for a non-obvious reason (not "Everything up-to-date"), +surface the error and stop. + +Then create the PR (body assembled from step 3's read): + +```bash +gh pr create --base --title "" --body "$(cat <<'EOF' +**Functional lines changed:** files, + − + +## Key files to review +- `…` — … +- `…` — … + +## Summary +- **…**: … +- **…**: … + +## Non-functional changes # only if substantial rules/skills/README changes +- `…` — … +EOF +)" +``` + +For draft PRs, add `--draft`. Capture the PR URL from the output. The +`## Non-functional changes` section is included only when step 2's +`---NONFUNCSTAT---` shows substantial changed files (e.g. README + rules large +enough to matter to reviewers). Keep it brief. + +### Step 5: cleanup (1 Bash call) + +```bash +rm -f ./tmp/diff.diff +``` + +### Step 6: report (inline, no call) + +Report the PR URL captured from step 4's `gh pr create` output. Done. Do not +`gh pr view` after creation unless the PR appeared to fail. + +## PR body order (always) + +1. `**Functional lines changed:** …` +2. `## Key files to review` +3. Optional `## TODO_IN_THIS_PR` +4. `## Summary` +5. Optional `## Non-functional changes` + +## PR title format + +``` + +``` + +No JIRA key, no `[WIP]` prefix. For draft PRs, use `--draft`. + +Examples: +- `Add smart router with content-based model downgrade` +- `Fix usage query time-window handling` + +## Rules-only / docs-only PRs (zero functional files) + +When the functional diff is empty: + +- `**Functional lines changed:**` → `0` explicitly. +- `## Key files to review` → 3–8 paths from the **full** diff (largest first). +- `## Summary` → describe the actual changed files. + +## Anti-patterns (do NOT do these) + +- Do NOT read individual file diffs with `git diff -- ` — use the single + temp-file dump instead. +- Do NOT run `git status` — the branch state is from step 0. +- Do NOT run `git fetch` as a separate call — if needed, prepend to step 0. +- Do NOT run `gh pr view` after creation unless step 4 output is suspicious. +- Do NOT run more than one Bash call for independent commands — chain with `&&`. +- Do NOT forget `working_directory` — every Bash call that touches the repo + needs it set explicitly. diff --git a/README.md b/README.md index b75c847..0960619 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,7 @@ ### Table of Contents - [📦 Quickstart](#-quickstart) +- [⚡ Smart Routing](#-smart-routing) - [☁️ Setting up Cloudflare](#️-setting-up-cloudflare) - [📊 Usage Dashboard](#-usage-dashboard) - [🪐 Providers](#-providers) @@ -89,6 +90,14 @@ discursive status --show-key | jq Gateway keys are masked by default. Pass `--show-key` to print the full `gateway_key` for Cursor setup. +> **💡 Smart routing is on by default.** The gateway's *smart router* inspects +> every request and, when safe, routes simpler work (short lookups, code +> search, structured extraction, and cheap per-turn tool results) to a cheaper +> model — typically `deepseek-v4-flash` — to cut cost and latency. More complex +> work (editing/refactoring, reasoning, large outputs) keeps the original model. +> See [Smart Routing](#-smart-routing) below, or disable it with +> `discursive start --smart-router=false`. + ### 3. Configure Cursor Open **Cursor Settings → Models** and enter: @@ -112,16 +121,16 @@ Reload Cursor: **Cmd+Shift+P → Reload Window**. You should see Change the model alias in Cursor's model picker — no restart needed: -| Cursor alias | Provider | Real model | Use | -| ------------- | -------- | ------------------- | -------------------------------------- | -| `gpt-4o` | Moonshot | `kimi-k3` | Planning / flagship | -| `gpt-4o-mini` | Moonshot | `kimi-k2.7-code` | Coding; always thinks | -| `o1` | DeepSeek | `deepseek-v4-pro` | Harder execution | -| `o3-mini` | DeepSeek | `deepseek-v4-flash` | Cheap execution | -| `gpt-5-nano` | Thaura | `thaura` | Ethical AI; optional provider | -| `gpt-4.1-turbo` | Z.AI | `glm-5.2` | Planning; cheaper than K3 | -| `gpt-4.1` | Z.AI | `glm-4.7` | Cheap execution | -| `gpt-4-turbo` | Z.AI | `glm-5.2` | Compat alias (Cursor may rewrite `gpt-4.1-turbo` to this) | +| Cursor alias | Provider | Real model | Use | +| --------------- | -------- | ------------------- | --------------------------------------------------------- | +| `gpt-4o` | Moonshot | `kimi-k3` | Planning / flagship | +| `gpt-4o-mini` | Moonshot | `kimi-k2.7-code` | Coding; always thinks | +| `o1` | DeepSeek | `deepseek-v4-pro` | Harder execution | +| `o3-mini` | DeepSeek | `deepseek-v4-flash` | Cheap execution | +| `gpt-5-nano` | Thaura | `thaura` | Ethical AI; optional provider | +| `gpt-4.1-turbo` | Z.AI | `glm-5.2` | Planning; cheaper than K3 | +| `gpt-4.1` | Z.AI | `glm-4.7` | Cheap execution | +| `gpt-4-turbo` | Z.AI | `glm-5.2` | Compat alias (Cursor may rewrite `gpt-4.1-turbo` to this) | @@ -133,6 +142,81 @@ In Cursor Settings → Models: turn off "Override OpenAI API Key" and --- +## ⚡ Smart Routing + +The gateway can **automatically downgrade individual requests** to a cheaper, +faster model when the work is simple enough — cutting token cost and latency +without changing what you pick in Cursor. Smart routing is **on by default** and +requires no configuration, but every behavior is configurable via +`discursive start` flags. + +> The router runs entirely **inside the gateway**. Cursor still sends every +> request to the gateway under whatever model alias you chose; the gateway +> inspects each request, may route it to a cheaper model, and proxies upstream. +> Cursor's model picker is unaware of the routing. + +### What gets downgraded + +Each incoming request is classified and routed to one of three tiers +(`keep` / `pro` / `flash`), based on the **content class** and, for per-turn +tool results, the **result size + tool name**: + +| Request type | Tier | Model | +| --------------------------------------------------------------- | ----- | --------------------- | +| Simple lookup / explanation | flash | `deepseek-v4-flash` | +| Code search / exploration | flash | `deepseek-v4-flash` | +| Structured extraction (`json_object` / `json_schema`) | flash | `deepseek-v4-flash` | +| Automation / mechanical work (lint, git, scripts, PR) | flash | `deepseek-v4-flash` | +| Tool result — small | flash | `deepseek-v4-flash` | +| Tool result — medium, read-only (`Read`/`Grep`/`Glob`) | flash | `deepseek-v4-flash` | +| Tool result — medium, write/decision (`Shell`/`StrReplace`) | pro | `deepseek-v4-pro` | +| Editing / refactoring | keep | original model | +| Complex reasoning / architecture | keep | original model | +| Tool result — large | keep | original model | + +Full details and rationale live in the agent rules (`.cursor/rules/gateway.mdc`, +*Per-turn routing tiers*). In practice a real GLM 5.2 refactor session routed +7 of 10 per-turn requests to `flash` and 2 decision-heavy ones to `pro`, keeping +the original model for the initial prompt and the one large tool result. + +### Flash output cap + +Flash-tier requests are limited to **1,500 output tokens** (`max_tokens = +1500`), so cheap turns stay terse instead of "overthinking" and emitting +4–16k tokens. This is the main lever that keeps agent flows fast. The cap is +applied even when the flow is already on `flash` (model override is a no-op). + +### `discursive start` flags + +| Flag | Default | Purpose | +| -------------------- | ------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `--smart-router` | `true` | Enable the smart router (classification + downgrade + per-turn routing + flash cap). Set `--smart-router=false` to run the gateway with no automatic model changes. | +| `--diagnostic-dump` | `false` | Debugging aid: write the raw `messages` array Cursor sent to `/tmp/discursive-msgdump-.json` whenever content extraction from the last user message returns nothing (e.g. tool-result rounds). Independent of `--smart-router` — dumps can fire even when routing is off, or stay silent while it's on. | +| `--log-level` | `info` | Log verbosity: `debug`, `info`, `warn`, `error`. Use `debug` to see the per-request `request_class`, `turn_type`, `tool_result_size`, `tool_name`, and `override_tier` lines from the router. Overrides `DISCURSIVE_LOG_LEVEL`. | +| `--background` | `false` | Detach and run in the background. Logs to `{dataRoot}/gateway.log`. | +| `--tunnel` | (config) | Tunnel mode: `named`, `none`, or `quick` (persists to config). | +| `--public-url` | (config) | Public HTTPS base URL ending in `/v1` (persists to config). | + +Examples: + +```bash +# Routing on (default) + debug logging +discursive start --smart-router --log-level debug + +# Routing on + debug logging + message dumps for a debugging session +discursive start --smart-router --log-level debug --diagnostic-dump + +# Disable routing entirely +discursive start --smart-router=false +``` + +> **💡 Tip:** At `--log-level debug`, the router logs one line per request with +> `request_class`, `turn_type`, `tool_result_size`, `tool_name`, and +> `override_tier`. This is the easiest way to see exactly what the router is +> doing and tune your expectations. + +--- + ## ☁️ Setting up Cloudflare @@ -190,7 +274,7 @@ logs include an `effort` field on request/response/usage lines. | --------------------------------------- | -------------------- | ---------------------------------------------------------------------------------------- | | `kimi-k3` | `low`, `high`, `max` | `low` (API default is `max`; we default lower for cost) | | `deepseek-v4-pro` / `deepseek-v4-flash` | `off`, `high`, `max` | `off` (`off` → `thinking: disabled`; otherwise `thinking: enabled` + `reasoning_effort`) | -| `glm-5.2` | `off`, `high`, `max` | `off` (`off` → `thinking: disabled`; otherwise `thinking: enabled` + `reasoning_effort`) | +| `glm-5.2` | `off`, `high`, `max` | `off` (`off` → `thinking: disabled`; otherwise `thinking: enabled` + `reasoning_effort`) | - Lower effort usually means fewer thinking tokens and lower cost. `thaura` does not expose this control. @@ -204,10 +288,10 @@ selector: [Kimi K2.7 Code](https://www.kimi.com/resources/kimi-k2-7-code) windows and native reasoning capabilities. -| API model ID | Cache hit / MTok | Input / MTok | Output / MTok | Role | -| ------------ | ---------------- | ------------ | ------------- | ----------------------------------------------- | -| `kimi-k3` | $0.30 | $3.00 | $15.00 | Flagship; 1M-token context, always thinks | -| `kimi-k2.7-code` | $0.19 | $0.95 | $4.00 | Coding model; always thinks | +| API model ID | Cache hit / MTok | Input / MTok | Output / MTok | Role | +| ---------------- | ---------------- | ------------ | ------------- | ----------------------------------------- | +| `kimi-k3` | $0.30 | $3.00 | $15.00 | Flagship; 1M-token context, always thinks | +| `kimi-k2.7-code` | $0.19 | $0.95 | $4.00 | Coding model; always thinks | - Pricing: [https://platform.kimi.ai/docs/pricing/chat](https://platform.kimi.ai/docs/pricing/chat) @@ -236,6 +320,37 @@ models at a fraction of the cost per token. --- +### 🪻 Z.AI + +[Z.AI](https://docs.z.ai/) provides GLM-series models with +thinking support and prompt caching. Z.AI is used via the **GLM Coding Plan** +(subscription, credits quota), which exposes the OpenAI-compatible base URL +`https://api.z.ai/api/coding/paas/v4`. + +| API model ID | Cache hit / MTok | Input / MTok | Output / MTok | Role | +| ------------ | ---------------- | ------------ | ------------- | ------------------------------------------------------------------------ | +| `glm-5.2` | $0.26 | $1.40 | $4.40 | Planning model; reasoning_effort + cache | +| `glm-4.7` | $0.11 | $0.60 | $2.20 | Budget execution; thinking on/off | +| `glm-4.6v` | $0.05 | $0.30 | $0.90 | Vision worker — describes images for ALL providers (not user-selectable) | + +> **Image routing:** any request (any provider) that contains image content is +> intercepted by the gateway and each image is described by Z.AI `glm-4.6v` +> (coding-plan endpoint) before the selected text model is called. A Z.AI API +> key is therefore required to send images. If it is missing or the vision +> model rejects the image, the request **fails fast** with a clear `vision_error` +> rather than silently dropping the image. + +- Pricing: [https://docs.z.ai/guides/overview/pricing](https://docs.z.ai/guides/overview/pricing) +- API docs: [https://docs.z.ai/api-reference/introduction](https://docs.z.ai/api-reference/introduction) +- API key: [https://z.ai/manage-apikey/apikey-list](https://z.ai/manage-apikey/apikey-list) (GLM Coding Plan key) + +| Parameter | `glm-5.2` | `glm-4.7` | +| ------------------ | ------------------------------------------------------------- | ----------------------- | +| `thinking` | `{type: "enabled"}` when reasoning; else `{type: "disabled"}` | `{type: "enabled" | "disabled"}` | +| `reasoning_effort` | Normalized → `off`/`high`/`max` | Deleted (not supported) | + +--- + ### 🐪 Thaura @@ -245,9 +360,9 @@ excellence with ethical principles, designed to support Palestinian liberation and mission-aligned technology development. -| API model ID | Input / MTok | Output / MTok | Role | -| ------------ | ------------ | ------------- | -------------------------------------------- | -| `thaura` | $0.50 | $2.00 | OpenAI-compatible chat and tool use | +| API model ID | Input / MTok | Output / MTok | Role | +| ------------ | ------------ | ------------- | ----------------------------------- | +| `thaura` | $0.50 | $2.00 | OpenAI-compatible chat and tool use | - Pricing: [https://thaura.ai/api-platform](https://thaura.ai/api-platform) @@ -286,37 +401,9 @@ and mission-aligned technology development. > +--- -### 🪻 Z.AI - -[Z.AI](https://docs.z.ai/) provides GLM-series models with -thinking support and prompt caching. Z.AI is used via the **GLM Coding Plan** -(subscription, credits quota), which exposes the OpenAI-compatible base URL -`https://api.z.ai/api/coding/paas/v4`. - -| API model ID | Cache hit / MTok | Input / MTok | Output / MTok | Role | -| ------------ | ---------------- | ------------ | ------------- | ----------------------------------------- | -| `glm-5.2` | $0.26 | $1.40 | $4.40 | Planning model; reasoning_effort + cache | -| `glm-4.7` | $0.11 | $0.60 | $2.20 | Budget execution; thinking on/off | -| `glm-4.6v` | $0.05 | $0.30 | $0.90 | Vision worker — describes images for ALL providers (not user-selectable) | - -> **Image routing:** any request (any provider) that contains image content is -> intercepted by the gateway and each image is described by Z.AI `glm-4.6v` -> (coding-plan endpoint) before the selected text model is called. A Z.AI API -> key is therefore required to send images. If it is missing or the vision -> model rejects the image, the request **fails fast** with a clear `vision_error` -> rather than silently dropping the image. - -- Pricing: [https://docs.z.ai/guides/overview/pricing](https://docs.z.ai/guides/overview/pricing) -- API docs: [https://docs.z.ai/api-reference/introduction](https://docs.z.ai/api-reference/introduction) -- API key: [https://z.ai/manage-apikey/apikey-list](https://z.ai/manage-apikey/apikey-list) (GLM Coding Plan key) - -| Parameter | `glm-5.2` | `glm-4.7` | -| ------------------ | ------------------------------------------------------------- | -------------------------------------------- | -| `thinking` | `{type: "enabled"}` when reasoning; else `{type: "disabled"}` | `{type: "enabled"|"disabled"}` | -| `reasoning_effort` | Normalized → `off`/`high`/`max` | Deleted (not supported) | - ## 🛠 Tech Stack @@ -327,11 +414,9 @@ thinking support and prompt caching. Z.AI is used via the **GLM Coding Plan** | Tunnel | [cloudflared](https://developers.cloudflare.com/cloudflare-one/connections/connect-networks/) named tunnel | | Upstream APIs | OpenAI-compatible chat completions (Moonshot + DeepSeek + Thaura + Z.AI) | - --- - ## 📁 File Structure ``` @@ -363,18 +448,18 @@ planning/ # MVP task sequence (T01–T10) All output is JSON on stdout. Pipe through `jq` for readability. -| Command | Description | -| ---------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `discursive start` | Start gateway on `localhost:4001`. `--background` forks to daemon. `--log-level` (debug/info/warn/error). `--tunnel` (named/none/quick), `--public-url`. Auto-invokes `init` if config is incomplete on first run. | -| `discursive stop` | Send SIGTERM via PID file. No-op if not running. | -| `discursive status` | Config dump + runtime state: PID alive? uptime? log file path/size, tunnel mode, model mapping. Gateway key masked by default; `--show-key` prints the full key. | -| `discursive logs` | Pretty-print `gateway.log` with colored level prefixes. `--follow` (`-f`) for live tail (uses fsnotify — no polling). `-n N` for last N lines. File auto-rotates at ~2 MB, keeps 2 backups. | -| `discursive log-level [debug | info | warn | error]` | Show or set log verbosity. Set persists per-process; hints how to export `DISCURSIVE_LOG_LEVEL` for persistence. | -| `discursive doctor` | Health checks: keys present, port available, local/public HTTP health, tunnel mode, cloudflared binary, logs writable. | -| `discursive usage` | Token + cost estimates per session/model. | +| Command | Description | +| ---------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| `discursive start` | Start gateway on `localhost:4001`. `--background` forks to daemon. `--log-level` (debug/info/warn/error). `--tunnel` (named/none/quick), `--public-url`. `--smart-router` (on by default), `--diagnostic-dump` (off by default). Auto-invokes `init` if config is incomplete on first run. See [Smart Routing](#-smart-routing). | +| `discursive stop` | Send SIGTERM via PID file. No-op if not running. | +| `discursive status` | Config dump + runtime state: PID alive? uptime? log file path/size, tunnel mode, model mapping. Gateway key masked by default; `--show-key` prints the full key. | +| `discursive logs` | Pretty-print `gateway.log` with colored level prefixes. `--follow` (`-f`) for live tail (uses fsnotify — no polling). `-n N` for last N lines. File auto-rotates at ~2 MB, keeps 2 backups. | +| `discursive log-level [debug | info | warn | error]` | Show or set log verbosity. Set persists per-process; hints how to export `DISCURSIVE_LOG_LEVEL` for persistence. | +| `discursive doctor` | Health checks: keys present, port available, local/public HTTP health, tunnel mode, cloudflared binary, logs writable. | +| `discursive usage` | Token + cost estimates per session/model. | | `discursive set` | Configure settings via flags. `--moonshot-key`, `--deepseek-key`, `--thaura-key`, `--zai-key`, `--tunnel-token`, `--public-url`, `--rotate-gateway-key`, `--model`. Combine several in one call. `--show-key` prints the full gateway key. | -| `discursive completion [bash | zsh | fish | powershell]` | Generate a shell completion script (see [Shell Completion](#️-shell-completion)). | -| `discursive version` | Print version. | +| `discursive completion [bash | zsh | fish | powershell]` | Generate a shell completion script (see [Shell Completion](#️-shell-completion)). | +| `discursive version` | Print version. | JSON slog on **stdout**, interactive prompts on **stderr** — pipe-friendly. diff --git a/internal/gateway/auth.go b/internal/gateway/auth.go index e188fc6..1ad4cc9 100644 --- a/internal/gateway/auth.go +++ b/internal/gateway/auth.go @@ -1,22 +1,32 @@ package gateway import ( + "crypto/subtle" "encoding/json" "net/http" "strings" ) +// bearerScheme is the OpenAI-compatible Authorization scheme used by Cursor. +// It is matched case-insensitively so operators can send "Bearer", "bearer", +// or any mixed-case variant and still authenticate. +const bearerScheme = "Bearer" + +// apiKeyHeaders lists the fallback headers (after Authorization) that may +// carry the gateway key, in priority order. +var apiKeyHeaders = []string{"api-key", "x-api-key", "x-openai-api-key"} + // ExtractAPIKey reads the gateway API key from common OpenAI-style headers. +// The Authorization scheme is matched case-insensitively. If Authorization is +// present with a non-Bearer scheme (or empty), control falls through to the +// api-key fallback headers. func ExtractAPIKey(r *http.Request) string { if auth := r.Header.Get("Authorization"); auth != "" { - if rest, ok := strings.CutPrefix(auth, "Bearer "); ok { - return strings.TrimSpace(rest) - } - if rest, ok := strings.CutPrefix(auth, "bearer "); ok { - return strings.TrimSpace(rest) + if scheme, token, ok := strings.Cut(auth, " "); ok && strings.EqualFold(scheme, bearerScheme) { + return strings.TrimSpace(token) } } - for _, name := range []string{"api-key", "x-api-key", "x-openai-api-key"} { + for _, name := range apiKeyHeaders { if v := strings.TrimSpace(r.Header.Get(name)); v != "" { return v } @@ -24,12 +34,13 @@ func ExtractAPIKey(r *http.Request) string { return "" } -// GatewayKeyMatches reports whether provided equals expected (constant-ish compare). +// GatewayKeyMatches reports whether provided equals expected using a +// constant-time comparison. func GatewayKeyMatches(provided, expected string) bool { if provided == "" || expected == "" { return false } - return provided == expected + return subtle.ConstantTimeCompare([]byte(provided), []byte(expected)) == 1 } // WriteUnauthorized writes an OpenAI-shaped 401 JSON body. @@ -46,6 +57,7 @@ func WriteUnauthorized(w http.ResponseWriter) { }) } +// writeJSONError writes a JSON error response. func writeJSONError(w http.ResponseWriter, status int, message, typ string) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(status) @@ -57,10 +69,18 @@ func writeJSONError(w http.ResponseWriter, status int, message, typ string) { }) } -func requireGatewayKey(w http.ResponseWriter, r *http.Request, expected string) bool { - if !GatewayKeyMatches(ExtractAPIKey(r), expected) { - WriteUnauthorized(w) - return false +// AuthMiddleware returns an http.Handler that wraps next with gateway key +// authentication. Requests that fail auth receive an OpenAI-shaped 401. +// If expected is empty, auth is bypassed. +func AuthMiddleware(next http.Handler, expected string) http.Handler { + if expected == "" { + return next } - return true + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !GatewayKeyMatches(ExtractAPIKey(r), expected) { + WriteUnauthorized(w) + return + } + next.ServeHTTP(w, r) + }) } diff --git a/internal/gateway/auth_test.go b/internal/gateway/auth_test.go index 1f5d8d9..1c47501 100644 --- a/internal/gateway/auth_test.go +++ b/internal/gateway/auth_test.go @@ -1,82 +1,107 @@ package gateway import ( + "encoding/json" + "net/http" "net/http/httptest" "strings" "testing" -) - -func TestExtractAPIKey_Bearer(t *testing.T) { - req := httptest.NewRequest("GET", "/", nil) - req.Header.Set("Authorization", "Bearer sk-test-key-12345") - if got := ExtractAPIKey(req); got != "sk-test-key-12345" { - t.Fatalf("got %q want %q", got, "sk-test-key-12345") - } -} - -func TestExtractAPIKey_BearerLowercase(t *testing.T) { - req := httptest.NewRequest("GET", "/", nil) - req.Header.Set("Authorization", "bearer sk-test-key-12345") - if got := ExtractAPIKey(req); got != "sk-test-key-12345" { - t.Fatalf("got %q want %q", got, "sk-test-key-12345") - } -} - -func TestExtractAPIKey_BearerWithSpaces(t *testing.T) { - // Authorization header is checked as-is (no leading-space trim on the header value itself). - // The key value IS trimmed after extraction. - req := httptest.NewRequest("GET", "/", nil) - req.Header.Set("Authorization", "Bearer sk-test-key-12345 ") - if got := ExtractAPIKey(req); got != "sk-test-key-12345" { - t.Fatalf("got %q want %q", got, "sk-test-key-12345") - } -} - -func TestExtractAPIKey_XApiKey(t *testing.T) { - req := httptest.NewRequest("GET", "/", nil) - req.Header.Set("x-api-key", "sk-x-api-key") - if got := ExtractAPIKey(req); got != "sk-x-api-key" { - t.Fatalf("got %q want %q", got, "sk-x-api-key") - } -} - -func TestExtractAPIKey_ApiKeyHeader(t *testing.T) { - req := httptest.NewRequest("GET", "/", nil) - req.Header.Set("api-key", "sk-api-key") - if got := ExtractAPIKey(req); got != "sk-api-key" { - t.Fatalf("got %q want %q", got, "sk-api-key") - } -} - -func TestExtractAPIKey_OpenAIKeyHeader(t *testing.T) { - req := httptest.NewRequest("GET", "/", nil) - req.Header.Set("x-openai-api-key", "sk-openai-key") - if got := ExtractAPIKey(req); got != "sk-openai-key" { - t.Fatalf("got %q want %q", got, "sk-openai-key") - } -} -func TestExtractAPIKey_NoHeader(t *testing.T) { - req := httptest.NewRequest("GET", "/", nil) - if got := ExtractAPIKey(req); got != "" { - t.Fatalf("got %q want empty", got) - } -} + "github.com/commoddity/discursive/internal/config" +) -func TestExtractAPIKey_BearerTakesPriority(t *testing.T) { - req := httptest.NewRequest("GET", "/", nil) - req.Header.Set("Authorization", "Bearer sk-bearer-key") - req.Header.Set("x-api-key", "sk-x-key") - if got := ExtractAPIKey(req); got != "sk-bearer-key" { - t.Fatalf("got %q want %q", got, "sk-bearer-key") +func TestExtractAPIKey(t *testing.T) { + tests := []struct { + name string + headers map[string]string + want string + }{ + { + name: "Bearer", + headers: map[string]string{"Authorization": "Bearer sk-test-key-12345"}, + want: "sk-test-key-12345", + }, + { + name: "bearer lowercase", + headers: map[string]string{"Authorization": "bearer sk-test-key-12345"}, + want: "sk-test-key-12345", + }, + { + name: "bearer uppercase scheme", + headers: map[string]string{"Authorization": "BEARER sk-test-key-12345"}, + want: "sk-test-key-12345", + }, + { + name: "bearer mixed-case scheme", + headers: map[string]string{"Authorization": "bEaReR sk-test-key-12345"}, + want: "sk-test-key-12345", + }, + { + name: "bearer separators between scheme and token", + headers: map[string]string{"Authorization": "Bearer sk-test-key-12345 "}, + want: "sk-test-key-12345", + }, + { + name: "Bearer with extra whitespace", + headers: map[string]string{"Authorization": "Bearer sk-test-key-12345 "}, + want: "sk-test-key-12345", + }, + { + name: "x-api-key header", + headers: map[string]string{"x-api-key": "sk-x-api-key"}, + want: "sk-x-api-key", + }, + { + name: "api-key header", + headers: map[string]string{"api-key": "sk-api-key"}, + want: "sk-api-key", + }, + { + name: "x-openai-api-key header", + headers: map[string]string{"x-openai-api-key": "sk-openai-key"}, + want: "sk-openai-key", + }, + { + name: "no header", + headers: map[string]string{}, + want: "", + }, + { + name: "Bearer takes priority over x-api-key", + headers: map[string]string{"Authorization": "Bearer sk-bearer-key", "x-api-key": "sk-x-key"}, + want: "sk-bearer-key", + }, + { + name: "Authorization without Bearer prefix", + headers: map[string]string{"Authorization": "Basic dXNlcjpwYXNz"}, + want: "", + }, + { + name: "non-Bearer scheme falls through to api-key header", + headers: map[string]string{"Authorization": "Basic dXNlcjpwYXNz", "x-api-key": "sk-fallback"}, + want: "sk-fallback", + }, + { + name: "empty Authorization falls through to x-api-key", + headers: map[string]string{"Authorization": "", "x-api-key": "sk-fallback"}, + want: "sk-fallback", + }, + { + name: "Bearer with empty value returns empty", + headers: map[string]string{"Authorization": "Bearer "}, + want: "", + }, } -} - -func TestExtractAPIKey_NoBearerPrefix(t *testing.T) { - req := httptest.NewRequest("GET", "/", nil) - req.Header.Set("Authorization", "Basic dXNlcjpwYXNz") - if got := ExtractAPIKey(req); got != "" { - t.Fatalf("got %q want empty", got) + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req := httptest.NewRequest("GET", "/", nil) + for k, v := range tt.headers { + req.Header.Set(k, v) + } + if got := ExtractAPIKey(req); got != tt.want { + t.Fatalf("ExtractAPIKey() = %q, want %q", got, tt.want) + } + }) } } @@ -89,6 +114,7 @@ func TestGatewayKeyMatches(t *testing.T) { }{ {"exact match", "sk-abc123", "sk-abc123", true}, {"mismatch", "sk-abc123", "sk-xyz789", false}, + {"mismatch same length", "sk-aaaaaaaaaa", "sk-bbbbbbbbbb", false}, {"empty provided", "", "sk-abc123", false}, {"empty expected", "sk-abc123", "", false}, {"both empty", "", "", false}, @@ -106,17 +132,199 @@ func TestWriteUnauthorized(t *testing.T) { w := httptest.NewRecorder() WriteUnauthorized(w) - if w.Code != 401 { - t.Fatalf("status = %d, want 401", w.Code) + if w.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want %d", w.Code, http.StatusUnauthorized) } if ct := w.Header().Get("Content-Type"); ct != "application/json" { t.Fatalf("Content-Type = %q, want application/json", ct) } - body := w.Body.String() - if !strings.Contains(body, "invalid_api_key") { - t.Fatalf("missing invalid_api_key in body: %s", body) + + var resp struct { + Error struct { + Message string `json:"message"` + Type string `json:"type"` + Param any `json:"param"` + Code string `json:"code"` + } `json:"error"` + } + if err := json.NewDecoder(w.Body).Decode(&resp); err != nil { + t.Fatalf("invalid JSON body: %v", err) } - if !strings.Contains(body, "invalid_request_error") { - t.Fatalf("missing invalid_request_error in body: %s", body) + if resp.Error.Code != "invalid_api_key" { + t.Fatalf("error.code = %q, want invalid_api_key", resp.Error.Code) } + if resp.Error.Type != "invalid_request_error" { + t.Fatalf("error.type = %q, want invalid_request_error", resp.Error.Type) + } + if resp.Error.Param != nil { + t.Fatalf("error.param = %v, want nil", resp.Error.Param) + } + if !strings.Contains(resp.Error.Message, "invalid_api_key") { + t.Fatalf("error.message missing invalid_api_key: %q", resp.Error.Message) + } +} + +func TestWriteJSONError(t *testing.T) { + tests := []struct { + name string + status int + message string + typ string + }{ + {"not_found", http.StatusNotFound, "resource not found", "not_found_error"}, + {"bad_request", http.StatusBadRequest, "invalid input", "invalid_request_error"}, + {"internal", http.StatusInternalServerError, "something went wrong", "internal_error"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + w := httptest.NewRecorder() + writeJSONError(w, tt.status, tt.message, tt.typ) + + if w.Code != tt.status { + t.Fatalf("status = %d, want %d", w.Code, tt.status) + } + if ct := w.Header().Get("Content-Type"); ct != "application/json" { + t.Fatalf("Content-Type = %q, want application/json", ct) + } + + var resp struct { + Error struct { + Message string `json:"message"` + Type string `json:"type"` + } `json:"error"` + } + if err := json.NewDecoder(w.Body).Decode(&resp); err != nil { + t.Fatalf("invalid JSON body: %v", err) + } + if resp.Error.Message != tt.message { + t.Fatalf("error.message = %q, want %q", resp.Error.Message, tt.message) + } + if resp.Error.Type != tt.typ { + t.Fatalf("error.type = %q, want %q", resp.Error.Type, tt.typ) + } + }) + } +} + +func TestAuthMiddleware(t *testing.T) { + okHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + }) + + t.Run("valid key passes through", func(t *testing.T) { + mw := AuthMiddleware(okHandler, "sk-correct") + req := httptest.NewRequest("GET", "/", nil) + req.Header.Set("Authorization", "Bearer sk-correct") + w := httptest.NewRecorder() + mw.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want %d", w.Code, http.StatusOK) + } + }) + + t.Run("invalid key returns 401", func(t *testing.T) { + mw := AuthMiddleware(okHandler, "sk-correct") + req := httptest.NewRequest("GET", "/", nil) + req.Header.Set("Authorization", "Bearer sk-wrong") + w := httptest.NewRecorder() + mw.ServeHTTP(w, req) + + if w.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want %d", w.Code, http.StatusUnauthorized) + } + }) + + t.Run("missing key returns 401", func(t *testing.T) { + mw := AuthMiddleware(okHandler, "sk-correct") + req := httptest.NewRequest("GET", "/", nil) + w := httptest.NewRecorder() + mw.ServeHTTP(w, req) + + if w.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want %d", w.Code, http.StatusUnauthorized) + } + }) + + t.Run("empty expected bypasses auth", func(t *testing.T) { + mw := AuthMiddleware(okHandler, "") + req := httptest.NewRequest("GET", "/", nil) + w := httptest.NewRecorder() + mw.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want %d (bypass when key is empty)", w.Code, http.StatusOK) + } + }) + + t.Run("wrong key format returns 401", func(t *testing.T) { + mw := AuthMiddleware(okHandler, "sk-correct") + req := httptest.NewRequest("GET", "/", nil) + req.Header.Set("Authorization", "Basic dXNlcjpwYXNz") + w := httptest.NewRecorder() + mw.ServeHTTP(w, req) + + if w.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want %d", w.Code, http.StatusUnauthorized) + } + }) + + t.Run("key via x-api-key header passes through", func(t *testing.T) { + mw := AuthMiddleware(okHandler, "sk-correct") + req := httptest.NewRequest("GET", "/", nil) + req.Header.Set("x-api-key", "sk-correct") + w := httptest.NewRecorder() + mw.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want %d", w.Code, http.StatusOK) + } + }) +} + +func TestAuthMiddleware_IntegrationWithServer(t *testing.T) { + srv, err := NewServer(ServerConfig{ + ListenAddr: "127.0.0.1:0", + GatewayKey: "sk-test-gateway", + Settings: &config.AppSettings{}, + }) + if err != nil { + t.Fatal(err) + } + ts := httptest.NewServer(srv.Handler()) + t.Cleanup(ts.Close) + t.Cleanup(func() { _ = srv.Shutdown(t.Context()) }) + + t.Run("GET /health bypasses auth", func(t *testing.T) { + res, err := http.Get(ts.URL + "/health") + if err != nil { + t.Fatal(err) + } + defer func() { _ = res.Body.Close() }() + if res.StatusCode != http.StatusOK { + t.Fatalf("health status = %d, want %d", res.StatusCode, http.StatusOK) + } + }) + + t.Run("GET /v1/models without auth returns 401", func(t *testing.T) { + res, err := http.Get(ts.URL + "/v1/models") + if err != nil { + t.Fatal(err) + } + defer func() { _ = res.Body.Close() }() + if res.StatusCode != http.StatusUnauthorized { + t.Fatalf("status = %d, want %d", res.StatusCode, http.StatusUnauthorized) + } + }) + + t.Run("POST /v1/chat/completions without auth returns 401", func(t *testing.T) { + res, err := http.Post(ts.URL+"/v1/chat/completions", "application/json", strings.NewReader(`{}`)) + if err != nil { + t.Fatal(err) + } + defer func() { _ = res.Body.Close() }() + if res.StatusCode != http.StatusUnauthorized { + t.Fatalf("status = %d, want %d", res.StatusCode, http.StatusUnauthorized) + } + }) } diff --git a/internal/gateway/health.go b/internal/gateway/health.go index db3ee6b..8804e56 100644 --- a/internal/gateway/health.go +++ b/internal/gateway/health.go @@ -15,9 +15,6 @@ func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) { } func (s *Server) handleModels(w http.ResponseWriter, r *http.Request) { - if !requireGatewayKey(w, r, s.cfg.GatewayKey) { - return - } listed := ListAdvertisedModels() data := make([]map[string]any, 0, len(listed)) for _, m := range listed { diff --git a/internal/gateway/proxy.go b/internal/gateway/proxy.go index bbfbdad..04908ea 100644 --- a/internal/gateway/proxy.go +++ b/internal/gateway/proxy.go @@ -13,9 +13,6 @@ import ( ) func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) { - if !requireGatewayKey(w, r, s.cfg.GatewayKey) { - return - } started := time.Now() requestID := newRequestID() @@ -55,6 +52,19 @@ func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) { // Classify request for subagent detection and downgrade to cheaper model. if result := s.router.ClassifyAndOverride(sanitized.Body, requestID); result.OverrideApplied { sanitized.Model = result.OverrideModel + // The override may map to a different provider (e.g. kimi-k3 → + // deepseek-v4-flash via the default fallback). Re-resolve the + // provider from the overridden model so the request is sent to the + // correct upstream. If the override model is unknown, fall back to a + // best-effort client error — we must not send a model to the wrong + // provider's endpoint. + if route, err := ResolveModel(result.OverrideModel); err != nil { + logRequest(requestID, "status", http.StatusBadGateway, "error", fmt.Sprintf("override model %q not resolvable: %v", result.OverrideModel, err), "provider", string(sanitized.Provider), "model", result.OverrideModel) + writeJSONError(w, http.StatusBadGateway, fmt.Sprintf("override model %q not resolvable: %v", result.OverrideModel, err), "upstream_error") + return + } else { + sanitized.Provider = route.Provider + } } upstreamKey, err := s.upstreamKey(sanitized.Provider) @@ -194,31 +204,28 @@ func (s *Server) writeBufferedResponse(w http.ResponseWriter, status int, respBo w.Header().Set("Content-Type", "application/json") w.WriteHeader(status) + + // Always log the full upstream error body at ERROR level. + slog.Error("upstream_error", + "request_id", requestID, + "status", status, + "provider", string(provider), + "model", model, + "effort", effort, + "body", string(respBody), + ) + + // Surface provider errors verbatim. If the upstream body is valid JSON, + // pass it through untouched (preserves the provider's error shape). If + // it's not JSON, wrap the raw body in an OpenAI-shaped envelope so the + // actual provider message reaches Cursor instead of a generic placeholder. var errObj map[string]any if json.Unmarshal(respBody, &errObj) == nil { _ = json.NewEncoder(w).Encode(errObj) - - // Always log the full upstream error body at ERROR level. - slog.Error("upstream_error", - "request_id", requestID, - "status", status, - "provider", string(provider), - "model", model, - "effort", effort, - "body", string(respBody), - ) } else { - slog.Error("upstream_error", - "request_id", requestID, - "status", status, - "provider", string(provider), - "model", model, - "effort", effort, - "body", string(respBody), - ) _ = json.NewEncoder(w).Encode(map[string]any{ "error": map[string]any{ - "message": fmt.Sprintf("upstream status %d", status), + "message": fmt.Sprintf("upstream status %d: %s", status, string(respBody)), "type": "upstream_error", }, }) diff --git a/internal/gateway/proxy_test.go b/internal/gateway/proxy_test.go index cf4bbb2..f81353e 100644 --- a/internal/gateway/proxy_test.go +++ b/internal/gateway/proxy_test.go @@ -341,3 +341,90 @@ func TestProxyImageWithoutVisionKeyFailsFast(t *testing.T) { t.Fatalf("text model called %d times despite fail-fast", textCalled.Load()) } } + +// TestProxy_SmartRouterProviderSwitch verifies that when the smart router +// downgrades a model to a model from a DIFFERENT provider, the gateway +// re-resolves the upstream provider so the request is sent to the correct +// endpoint. Regression for the "modelCode: does not exist" cross-provider +// bug (e.g. gpt-4o→kimi-k3 downgraded to deepseek-v4-flash must reach the +// DeepSeek endpoint, not the Moonshot endpoint). +func TestProxy_SmartRouterProviderSwitch(t *testing.T) { + var moonshotCalled atomic.Int32 + moonshotUp := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + moonshotCalled.Add(1) + t.Log("moonshot upstream hit (should not happen for downgraded request)") + _ = json.NewEncoder(w).Encode(mockCompletion("kimi-k3")) + })) + t.Cleanup(moonshotUp.Close) + + var deepseekModel string + var deepseekCalled atomic.Int32 + deepseekUp := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + deepseekCalled.Add(1) + var body map[string]any + _ = json.NewDecoder(r.Body).Decode(&body) + deepseekModel, _ = body["model"].(string) + _ = json.NewEncoder(w).Encode(mockCompletion(deepseekModel)) + })) + t.Cleanup(deepseekUp.Close) + + dataRoot := t.TempDir() + settings := config.DefaultSettings() + if err := settings.EnsureGatewayKey(); err != nil { + t.Fatal(err) + } + if err := settings.SetMoonshotKey(dataRoot, "sk-ms"); err != nil { + t.Fatal(err) + } + if err := settings.SetDeepSeekKey(dataRoot, "sk-ds"); err != nil { + t.Fatal(err) + } + if err := config.Save(dataRoot, settings); err != nil { + t.Fatal(err) + } + + srv, err := gateway.NewServer(gateway.ServerConfig{ + ListenAddr: "127.0.0.1:0", + GatewayKey: settings.GatewayKey, + DataRoot: dataRoot, + Settings: &settings, + HTTPClient: deepseekUp.Client(), + SmartRouterEnabled: true, + ChatURLOverride: map[config.Provider]string{ + config.ProviderMoonshot: moonshotUp.URL + "/moonshot/chat/completions", + config.ProviderDeepSeek: deepseekUp.URL + "/deepseek/chat/completions", + }, + }) + if err != nil { + t.Fatal(err) + } + ts := httptest.NewServer(srv.Handler()) + t.Cleanup(ts.Close) + t.Cleanup(func() { _ = srv.Shutdown(t.Context()) }) + + env := &testEnv{srv: srv, ts: ts, gatewayKey: settings.GatewayKey, dataRoot: dataRoot} + + // gpt-4o resolves to kimi-k3 (Moonshot). It's a short, simple lookup — + // classified SimpleLookup → downgraded to defaultFlashModel + // (deepseek-v4-flash), which lives on DeepSeek. The request must reach the + // DeepSeek upstream with the overridden model. + res, body := env.doJSON(t, http.MethodPost, "/v1/chat/completions", true, map[string]any{ + "model": "gpt-4o", + "messages": []any{ + map[string]any{"role": "user", "content": "what is a goroutine?"}, + }, + }) + if res.StatusCode != 200 { + t.Fatalf("status %d body %s", res.StatusCode, body) + } + + if deepseekCalled.Load() != 1 { + t.Fatalf("expected DeepSeek upstream to be called exactly once, got %d", deepseekCalled.Load()) + } + if deepseekModel != "deepseek-v4-flash" { + t.Fatalf("expected overridden model deepseek-v4-flash sent to DeepSeek, got %q", deepseekModel) + } + if moonshotCalled.Load() != 0 { + t.Fatalf("Moonshot upstream was called %d times for a downgraded request, it must not be", moonshotCalled.Load()) + } +} diff --git a/internal/gateway/router.go b/internal/gateway/router.go index 6da3b4a..c322f6c 100644 --- a/internal/gateway/router.go +++ b/internal/gateway/router.go @@ -2,6 +2,7 @@ package gateway import ( "encoding/json" + "hash/fnv" "log/slog" "os" "path/filepath" @@ -53,6 +54,63 @@ type RouterConfig struct { DiagnosticDump bool } +// TurnType classifies what kind of agent-loop turn a request represents, +// based on the role of the last message in the messages array. +type TurnType string + +const ( + // TurnToolResult means the last message role is "tool" — the model just + // received a tool result and must decide the next step. These are the + // candidate turns for downgrade to a cheaper model. + TurnToolResult TurnType = "tool_result" + // TurnUserPrompt means the last message role is "user" or "developer" — + // a fresh human prompt. The content classifier (classifyRequest) handles + // these. + TurnUserPrompt TurnType = "user_prompt" + // TurnAgentContinue means the last message role is "assistant" — + // continuation or prefill scenario. Conservative: keep model. + TurnAgentContinue TurnType = "agent_continue" + // TurnUnknown means empty messages, system-only, or unparseable + // structure. Conservative: keep model. + TurnUnknown TurnType = "unknown" +) + +// OverrideTier states how aggressively a request should be downgraded. +// It replaces the old boolean "should downgrade" with a 3-way decision so the +// router can route cheap turns to flash, decision-heavy cheap turns to pro, +// and everything else to the original (kept) model. +type OverrideTier string + +const ( + // TierKeep means keep the original model — no override. + TierKeep OverrideTier = "keep" + // TierPro means downgrade to a mid-tier model that still reasons well + // (e.g. deepseek-v4-pro) — used for decision-heavy cheap tool results. + TierPro OverrideTier = "pro" + // TierFlash means downgrade to the cheapest model (e.g. deepseek-v4-flash) + // — used for read-only / low-risk tool results and simple content classes. + TierFlash OverrideTier = "flash" +) + +// toolVerbosityCap is the max completion tokens we inject into cheap +// (flash) tool-result turns. It exists so cheap turns cannot balloon into +// multi-thousand-token streams (4–16k tokens / 20–106s latencies were observed +// in real runs), which is the real cause of "the flow took too long". Keeping +// cheap turns terse makes the whole agent loop faster at negligible cost. +const toolVerbosityCap = 1500 + +// Read-only tools never change state; their results only need cheap +// interpretation. Write / decision tools (StrReplace, Shell, notebooks, etc.) +// may warrant a mid-tier (pro) model when their results are non-trivial. +var readOnlyToolNames = map[string]bool{ + "Read": true, + "Grep": true, + "Glob": true, + "WebSearch": true, + "WebFetch": true, + "FetchMcpResource": true, +} + // ClassifierResult contains the classification outcome for logging/shadowing. type ClassifierResult struct { IsSubagent bool `json:"is_subagent"` @@ -65,19 +123,64 @@ type ClassifierResult struct { OverrideApplied bool `json:"override_applied"` OverrideReason string `json:"override_reason,omitempty"` ClassificationAge string `json:"classification_age,omitempty"` + TurnType TurnType `json:"turn_type"` + ToolResultSize string `json:"tool_result_size,omitempty"` + ToolName string `json:"tool_name,omitempty"` + OverrideTier OverrideTier `json:"override_tier,omitempty"` + MaxTokens int `json:"max_tokens,omitempty"` } -// modelOverrideMap maps expensive models to cheaper equivalents for classification-based downgrades. -// Entries here take priority. Models not in this map fall back to defaultSubagentModel. +// modelOverrideMap maps expensive models to cheaper equivalents for flash-tier +// downgrades. Entries take priority; models not in the map fall back to +// defaultFlashModel. var modelOverrideMap = map[string]string{ "deepseek-v4-pro": "deepseek-v4-flash", // "glm-5.2": "glm-4.7", // Uncomment to use glm-4.7 as flash for glm-5.2 } -// defaultSubagentModel is the fallback model for any traffic that triggers a -// downgrade but whose original model isn't in modelOverrideMap. Anything not -// in the map (including models that aren't DeepSeek at all) falls back to this. -const defaultSubagentModel = "deepseek-v4-flash" +// proModelOverrideMap maps expensive models to a mid-tier pro equivalent for +// decision-heavy downgrades (medium write/error tool results). Entries take +// priority; models not in the map fall back to defaultProModel. Models that +// are already pro-or-below (e.g. deepseek-v4-pro itself) have no entry, so a +// pro-tier override is a no-op for them. +var proModelOverrideMap = map[string]string{ + "glm-5.2": "deepseek-v4-pro", + "kimi-k3": "deepseek-v4-pro", + "kimi-k2.7-code": "deepseek-v4-pro", + "deepseek-v4": "deepseek-v4-pro", + "deepseek-v4-pro": "deepseek-v4-pro", +} + +// defaultFlashModel is the fallback model for any traffic that triggers a +// cheap downgrade but whose original model isn't in modelOverrideMap. Anything +// not in the map (including models that aren't DeepSeek at all) falls back to +// this. +const defaultFlashModel = "deepseek-v4-flash" + +// defaultProModel is the fallback model for pro-tier downgrades of originals +// not present in proModelOverrideMap. +const defaultProModel = "deepseek-v4-pro" + +// overrideModelForTier resolves the model to send for a downgrade of the given +// tier. flash uses modelOverrideMap (fallback defaultFlashModel); pro uses +// proModelOverrideMap (fallback defaultProModel). keep returns the original +// model unchanged. +func overrideModelForTier(original string, tier OverrideTier) string { + switch tier { + case TierPro: + if m, ok := proModelOverrideMap[original]; ok { + return m + } + return defaultProModel + case TierFlash: + if m, ok := modelOverrideMap[original]; ok { + return m + } + return defaultFlashModel + default: + return original + } +} // SmartRouter performs request classification and optional model override. type SmartRouter struct { @@ -123,15 +226,41 @@ func (r *SmartRouter) ClassifyAndOverride(body map[string]any, requestID string) // Instead, content-based classification handles all traffic uniformly. result.IsSubagent = false - // Step 2: content-based classification. + // Step 2: turn type detection for per-turn routing instrumentation. + result.TurnType = detectTurnType(body) + result.ToolResultSize = toolResultSize(body) + result.ToolName = toolResultName(body) + + // Step 3: content-based classification. result.RequestClass = classifyRequest(body) + // Determine this turn's routing tier. Two independent signals contribute: + // (a) content class — simple lookup, code search, structured extraction, ... + // (b) per-turn — tool-result rounds within a multi-step agent flow. + // The per-turn signal is decision-aware: cheap (small / medium read-only) + // results route to flash, non-trivial write/error results route to pro, + // and large results keep the original model (big outputs may need + // pro-level interpretation). + tier := overrideTier(result, body) + result.OverrideTier = tier + + // Verbosity cap: cheap (flash) tiers get a tight max_tokens so they cannot + // balloon into multi-thousand-token streams that stall the agent loop. + if tier == TierFlash { + result.MaxTokens = toolVerbosityCap + } + // Always log classification at DEBUG so operators can tune thresholds. lastMsg := lastUserMessage(body) lastMsgStripped := stripCursorNoise(lastMsg) slog.Debug("router: classify", "request_id", requestID, "request_class", result.RequestClass, + "turn_type", result.TurnType, + "tool_result_size", result.ToolResultSize, + "tool_name", result.ToolName, + "override_tier", result.OverrideTier, + "max_tokens", result.MaxTokens, "sys_prompt_len", result.SysPromptLen, "msg_count", result.MsgCount, "has_tools", result.HasTools, @@ -143,9 +272,20 @@ func (r *SmartRouter) ClassifyAndOverride(body map[string]any, requestID string) ) // Dump the full messages array to a temp file when content extraction fails, - // so we can analyze Cursor's message structure and refine our extraction. - // Gated behind the separate --diagnostic-dump flag. - if r.cfg.DiagnosticDump && len(lastMsgStripped) == 0 { + // so we can analyze Cursor's message structure and refine our extraction, + // and — during debugging sessions (--diagnostic-dump is always operator-gated + // and off by default) — when a classify event is otherwise not visible from + // the log lines alone. These are: + // * stripped_len == 0 → extraction failed (original reason) + // * tool-result turns → validate size bucketing + tool schema + // * overrides actually applied → confirm the downgrade (model+provider) + // * no-op extraction → stripped_len == last_msg_len warns that + // stripCursorNoise stripped nothing, so the + // classifier may be missing structure + // * 1% random sample → catch unforeseen shapes current + // heuristics can't predict + // Dump happens before routing so the raw message array is captured verbatim. + if r.cfg.DiagnosticDump && shouldDumpForDiagnostics(result.TurnType, result.ToolResultSize, tier != TierKeep, len(lastMsg), len(lastMsgStripped), requestID) { dumpMessages(body, requestID) } @@ -153,30 +293,47 @@ func (r *SmartRouter) ClassifyAndOverride(body map[string]any, requestID string) return result } - // Determine whether to downgrade based on request class. - if !shouldDowngrade(result.RequestClass) { + // No override when the tier says to keep the original model. + if tier == TierKeep { return result } - // Apply model override: prefer an explicit map entry; fall back to the - // universal default (deepseek-v4-flash) for models not in the map. - override, ok := modelOverrideMap[result.OriginalModel] - if !ok { - override = defaultSubagentModel + // Apply the verbosity cap first so it also applies when the effective tier + // is flash but the model override below is a no-op (e.g. the flow already + // runs on deepseek-v4-flash). Flash-tier turns are always kept terse to + // avoid multi-thousand-token streams that stall the agent loop. Applied + // before the model override so a no-op override still gets the cap. + if tier == TierFlash { + body["max_tokens"] = toolVerbosityCap } - // No-op if the resolved override equals the current model. + // Resolve the override model for the chosen tier. Both maps fall back to a + // provider-agnostic default so models not explicitly mapped still downgrade. + override := overrideModelForTier(result.OriginalModel, tier) + + // No-op if the resolved override equals the current model (e.g. a pro-tier + // request that is already on deepseek-v4-pro). if override == result.OriginalModel { return result } result.OverrideModel = override result.OverrideApplied = true - result.OverrideReason = string(result.RequestClass) + " downgrade (" + result.OriginalModel + " → " + override + ")" + + // Build a reason that identifies which signal triggered the downgrade. + result.OverrideReason = overrideReason(result, tier) + " downgrade (" + result.OriginalModel + " → " + override + ")" + + // Apply the model override to the body (proxy re-resolves the provider). body["model"] = override + slog.Info("router: model_overridden", "request_id", requestID, "request_class", result.RequestClass, + "turn_type", result.TurnType, + "tool_result_size", result.ToolResultSize, + "tool_name", result.ToolName, + "override_tier", result.OverrideTier, + "max_tokens", result.MaxTokens, "from", result.OriginalModel, "to", override, "sys_prompt_len", result.SysPromptLen, @@ -201,6 +358,82 @@ func shouldDowngrade(c RequestClass) bool { } } +// overrideTier decides the routing tier for a request. It combines the content +// class (a) with the per-turn tool-result decision (b): +// +// Content-driven downgrade triggered → TierFlash (cheap is fine for simple +// lookup / code search / structured extraction / automation). +// +// Tool-result turn (last role == tool): +// small → TierFlash (short outputs, low risk) +// medium + read-only tool → TierFlash (just reading: grep/read/search) +// medium + write/error tool → TierPro (decision-heavy: shell, edits, +// tests, failures — needs pro) +// large → TierKeep (big outputs need pro reasoning) +// unknown / empty size → TierKeep (can't gauge risk → conservative) +// +// Everything else → content-class result (shouldDowngrade) or TierKeep. +func overrideTier(result ClassifierResult, body map[string]any) OverrideTier { + // Per-turn tool-result signal first — it is the primary driver of cheap + // turns within a multi-step agent flow. + if result.TurnType == TurnToolResult { + switch result.ToolResultSize { + case "small": + return TierFlash + case "medium": + if isWriteTool(result.ToolName) { + return TierPro + } + return TierFlash + case "large": + return TierKeep + default: + // Unknown/empty size: we can't gauge risk, so keep the original + // model (conservative). + return TierKeep + } + } + + // Non-tool turns follow the content classifier. + if shouldDowngrade(result.RequestClass) { + return TierFlash + } + return TierKeep +} + +// overrideReason builds a human-readable reason for the override log, naming +// whichever signal(s) triggered the tier. +func overrideReason(result ClassifierResult, tier OverrideTier) string { + var parts []string + if result.TurnType == TurnToolResult && result.ToolResultSize != "" { + parts = append(parts, string(result.TurnType)+"/"+result.ToolResultSize) + if result.ToolName != "" { + parts = append(parts, "tool="+result.ToolName) + } + } else if result.RequestClass != "" && result.RequestClass != ClassUnknown { + parts = append(parts, string(result.RequestClass)) + } + if tier == TierPro { + parts = append(parts, "decision-heavy→pro") + } + if len(parts) == 0 { + return string(tier) + } + return strings.Join(parts, "+") +} + +// isWriteTool returns true when a tool name represents a state-changing or +// decision-heavy operation whose (non-small) result warrants a pro model. +// Read-only tools (Read/Grep/Glob/search) → false. Everything unknown is +// treated as a write/decision tool so we default to the more conservative +// (pro) tier for medium results. +func isWriteTool(name string) bool { + if name == "" { + return true + } + return !readOnlyToolNames[name] +} + // classifyRequest inspects the last user message and request structure to // infer the task type. Returns ClassUnknown when no clear signal is found. // @@ -497,6 +730,137 @@ func messageRoles(body map[string]any) string { return strings.Join(roles, ", ") } +// detectTurnType classifies the request by the role of its last message. +// Tool-result rounds (last role == "tool") are the primary targets for +// per-turn model downgrade within a multi-step agent flow. +func detectTurnType(body map[string]any) TurnType { + msgs, ok := body["messages"].([]any) + if !ok || len(msgs) == 0 { + return TurnUnknown + } + last, ok := msgs[len(msgs)-1].(map[string]any) + if !ok { + return TurnUnknown + } + role, _ := last["role"].(string) + switch role { + case "tool": + return TurnToolResult + case "user", "developer": + return TurnUserPrompt + case "assistant": + return TurnAgentContinue + default: + return TurnUnknown + } +} + +// toolResultSize buckets the content size of the last tool-result message. +// Zero means the messages array is empty or the last message is not a tool. +// Buckets help identify where token spend concentrates: +// +// <=512 chars → "small" (quick shell command outputs) +// <=4096 chars → "medium" (diff outputs, file reads) +// >4096 chars → "large" (full logs, stack traces) +func toolResultSize(body map[string]any) string { + msgs, ok := body["messages"].([]any) + if !ok || len(msgs) == 0 { + return "" + } + last, ok := msgs[len(msgs)-1].(map[string]any) + if !ok { + return "" + } + role, _ := last["role"].(string) + if role != "tool" { + return "" + } + content, _ := last["content"].(string) + n := len(content) + switch { + case n <= 512: + return "small" + case n <= 4096: + return "medium" + default: + return "large" + } +} + +// toolResultName returns the tool's function name for the last tool-result +// message, or "" when it can't be determined. It works by matching the last +// tool message's tool_call_id against the preceding assistant message's +// tool_calls[].id (or function call), returning the associated function name. +// This lets the router distinguish read-only tools (Read/Grep/Glob/search) +// from write/decision tools (Shell/StrReplace/notebook) to thread the +// decision-aware pro tier. +func toolResultName(body map[string]any) string { + msgs, ok := body["messages"].([]any) + if !ok || len(msgs) == 0 { + return "" + } + + // Find the last tool message and its tool_call_id. + var callID string + lastIdx := -1 + for i := len(msgs) - 1; i >= 0; i-- { + msg, ok := msgs[i].(map[string]any) + if !ok { + continue + } + role, _ := msg["role"].(string) + if role != "tool" { + continue + } + lastIdx = i + callID, _ = msg["tool_call_id"].(string) + break + } + if lastIdx < 0 { + return "" + } + + // Walk backward from the tool message for the assistant message that + // issued the tool call with the matching id. + for i := lastIdx - 1; i >= 0; i-- { + msg, ok := msgs[i].(map[string]any) + if !ok { + continue + } + role, _ := msg["role"].(string) + if role != "assistant" { + continue + } + // Legacy shape: single function_call object. + if fc, ok := msg["function_call"].(map[string]any); ok { + name, _ := fc["name"].(string) + if name != "" { + return name + } + } + // Standard shape: tool_calls array. + tcs, ok := msg["tool_calls"].([]any) + if !ok { + continue + } + for _, tc := range tcs { + entry, ok := tc.(map[string]any) + if !ok { + continue + } + if id, _ := entry["id"].(string); id != "" && id != callID { + continue + } + if fn, ok := entry["function"].(map[string]any); ok { + if name, _ := fn["name"].(string); name != "" { + return name + } + } + } + } + return "" +} + // systemPromptLength returns the length (in chars) of the first system or // developer message content, or 0 if none found. func systemPromptLength(body map[string]any) int { @@ -649,3 +1013,48 @@ func dumpMessages(body map[string]any, requestID string) { } slog.Info("router: dump_messages written", "request_id", requestID, "file", filename, "size_bytes", len(b)) } + +// dumpSampleRate is the denominator for the deterministic sample of classify +// events we dump during a debugging session (1 in 100). It exists so a +// long-running --diagnostic-dump session surfaces shapes the targeted triggers +// can't predict, without writing a file per request. +const dumpSampleRate uint32 = 100 + +// shouldDumpForDiagnostics decides whether a classify event warrants a raw +// message dump during a debugging session. Dumping is always operator-gated by +// the --diagnostic-dump flag; this function only narrows which requests within +// that session produce a file so we don't write one per request even when the +// flag is on. It returns true when the event is otherwise invisible from the +// log lines alone: +// - stripped_len == 0 extraction failed (always, original trigger) +// - tool-result turn validate size bucketing + tool content schema +// - would_downgrade confirm the model+provider downgrade +// - no-op extraction stripCursorNoise removed nothing — the +// classifier may be missing structure +// - 1% deterministic sample catch unforeseen shapes heuristics miss +func shouldDumpForDiagnostics(turnType TurnType, toolResultSize string, wouldDowngrade bool, lastMsgLen, strippedLen int, requestID string) bool { + // Extraction failed — original diagnostic trigger, always dump. + if strippedLen == 0 { + return true + } + // Tool-result turns: validate size bucketing and tool content schema. + if turnType == TurnToolResult { + return true + } + // Overrides actually applied: confirm the downgrade sent the expected model + // and provider, and nothing was lost in provider re-resolution. + if wouldDowngrade { + return true + } + // No-op extraction: if stripping removed nothing, the classifier may be + // missing structure it should have seen. + if lastMsgLen > 0 && strippedLen == lastMsgLen { + return true + } + // Deterministic sample so long-running sessions surface unforeseen shapes. + // Hash the requestID (a random per-request string) and gate on a modulus; + // hashing avoids mutable per-process state and is reproducible. + h := fnv.New32a() + _, _ = h.Write([]byte(requestID)) + return h.Sum32()%dumpSampleRate == 0 +} diff --git a/internal/gateway/router_test.go b/internal/gateway/router_test.go index 5390d20..44a180a 100644 --- a/internal/gateway/router_test.go +++ b/internal/gateway/router_test.go @@ -1,6 +1,9 @@ package gateway import ( + "encoding/json" + "fmt" + "strings" "testing" ) @@ -550,9 +553,9 @@ func TestSmartRouter_UnknownProviderDefaultFallback(t *testing.T) { t.Errorf("%q should be overridden (not in map → fallback default), class=%q", model, result.RequestClass) } - if result.OverrideModel != defaultSubagentModel { + if result.OverrideModel != defaultFlashModel { t.Errorf("%q override model = %q, want default %q", - model, result.OverrideModel, defaultSubagentModel) + model, result.OverrideModel, defaultFlashModel) } }) } @@ -573,8 +576,8 @@ func TestSmartRouter_Glm52ToFlashViaDefault(t *testing.T) { if !result.OverrideApplied { t.Errorf("glm-5.2 should be overridden (not in map → default), got class=%q", result.RequestClass) } - if stringField(body, "model") != defaultSubagentModel { - t.Errorf("expected default %q, got %q", defaultSubagentModel, stringField(body, "model")) + if stringField(body, "model") != defaultFlashModel { + t.Errorf("expected default %q, got %q", defaultFlashModel, stringField(body, "model")) } } @@ -622,6 +625,217 @@ func TestSmartRouter_DisabledPreservesModelAndClassifies(t *testing.T) { } } +// TestPerTurnDowngrade_ToolResultSmall verifies that a tool-result turn with a +// small result gets downgraded to flash EVEN when the content class is +// "editing" (which would normally keep the model). This is the core per-turn +// routing behavior: the model just received a tool result and needs to decide +// the next step — a task flash can handle. +func TestPerTurnDowngrade_ToolResultSmall(t *testing.T) { + body := map[string]any{ + "model": "deepseek-v4-pro", + "messages": []any{ + map[string]any{"role": "system", "content": strings.Repeat("x", 15000)}, + map[string]any{"role": "user", "content": "refactor the auth module and add tests"}, + map[string]any{"role": "assistant", "content": "", "tool_calls": []any{}}, + // Small tool result: "go test" output, ~30 chars + map[string]any{"role": "tool", "content": "ok\tinternal/gateway\t0.123s"}, + }, + } + r := NewSmartRouter(RouterConfig{Enabled: true}) + result := r.ClassifyAndOverride(body, "req_test") + + if !result.OverrideApplied { + t.Errorf("expected override applied for small tool-result turn, got none (class=%q, turn=%q, size=%q)", + result.RequestClass, result.TurnType, result.ToolResultSize) + } + if result.OverrideModel != "deepseek-v4-flash" { + t.Errorf("expected deepseek-v4-flash, got %q", result.OverrideModel) + } + if result.TurnType != TurnToolResult { + t.Errorf("expected turn_type=tool_result, got %q", result.TurnType) + } + if result.ToolResultSize != "small" { + t.Errorf("expected tool_result_size=small, got %q", result.ToolResultSize) + } +} + +// TestPerTurnDowngrade_ToolResultLargeKeptModel verifies that a tool-result +// turn with a LARGE result keeps the original model, even though it's a +// tool-result turn. This is the conservative policy: large results may need +// pro-level interpretation. +func TestPerTurnDowngrade_ToolResultLargeKeptModel(t *testing.T) { + body := map[string]any{ + "model": "deepseek-v4-pro", + "messages": []any{ + map[string]any{"role": "system", "content": strings.Repeat("x", 15000)}, + map[string]any{"role": "user", "content": "refactor the auth module and add tests"}, + map[string]any{"role": "assistant", "content": "", "tool_calls": []any{}}, + // Large tool result: >4096 chars + map[string]any{"role": "tool", "content": strings.Repeat("x", 5000)}, + }, + } + r := NewSmartRouter(RouterConfig{Enabled: true}) + result := r.ClassifyAndOverride(body, "req_test") + + if result.OverrideApplied { + t.Errorf("expected NO override for large tool-result turn, got override to %q (class=%q, turn=%q, size=%q)", + result.OverrideModel, result.RequestClass, result.TurnType, result.ToolResultSize) + } + if result.ToolResultSize != "large" { + t.Errorf("expected tool_result_size=large, got %q", result.ToolResultSize) + } +} + +// TestPerTurnDowngrade_ToolResultMediumWriteToPro verifies the decision-aware +// tier: a MEDIUM tool result from a write/decision tool (StrReplace) on an +// expensive original model (glm-5.2) downgrades to deepseek-v4-pro (not flash), +// and does NOT apply the flash verbosity cap. +func TestPerTurnDowngrade_ToolResultMediumWriteToPro(t *testing.T) { + body := map[string]any{ + "model": "glm-5.2", + "messages": []any{ + map[string]any{"role": "system", "content": strings.Repeat("x", 15000)}, + map[string]any{"role": "user", "content": "refactor the auth module"}, + map[string]any{ + "role": "assistant", + "content": "", + "tool_calls": []any{ + map[string]any{ + "id": "call_edit", + "type": "function", + "function": map[string]any{"name": "StrReplace", "arguments": "{}"}, + }, + }, + }, + // Medium tool result: <4096 chars, write result + map[string]any{"role": "tool", "tool_call_id": "call_edit", "content": strings.Repeat("y", 2000)}, + }, + } + r := NewSmartRouter(RouterConfig{Enabled: true}) + result := r.ClassifyAndOverride(body, "req_test") + + if !result.OverrideApplied { + t.Fatalf("expected override for medium write tool-result turn (class=%q, turn=%q, size=%q, tool=%q)", + result.RequestClass, result.TurnType, result.ToolResultSize, result.ToolName) + } + if result.OverrideTier != TierPro { + t.Errorf("expected override_tier=pro, got %q", result.OverrideTier) + } + if result.OverrideModel != "deepseek-v4-pro" { + t.Errorf("expected deepseek-v4-pro for medium write tool result, got %q", result.OverrideModel) + } + if result.ToolName != "StrReplace" { + t.Errorf("expected tool_name=StrReplace, got %q", result.ToolName) + } + // Pro tier must NOT receive the flash verbosity cap. + switch v := body["max_tokens"].(type) { + case json.Number: + if n, _ := v.Int64(); n == toolVerbosityCap { + t.Errorf("pro tier must not cap max_tokens to %d, got %d", toolVerbosityCap, n) + } + case float64: + if int(v) == toolVerbosityCap { + t.Errorf("pro tier must not cap max_tokens to %d, got %v", toolVerbosityCap, v) + } + } +} + +// TestPerTurnDowngrade_ToolResultMediumReadToFlash verifies a MEDIUM read-only +// tool result (Read) downgrades to flash and gets the verbosity cap. +func TestPerTurnDowngrade_ToolResultMediumReadToFlash(t *testing.T) { + body := map[string]any{ + "model": "glm-5.2", + "messages": []any{ + map[string]any{"role": "system", "content": strings.Repeat("x", 15000)}, + map[string]any{"role": "user", "content": "refactor the auth module"}, + map[string]any{ + "role": "assistant", + "content": "", + "tool_calls": []any{ + map[string]any{ + "id": "call_read", + "type": "function", + "function": map[string]any{"name": "Read", "arguments": "{}"}, + }, + }, + }, + map[string]any{"role": "tool", "tool_call_id": "call_read", "content": strings.Repeat("z", 2000)}, + }, + } + r := NewSmartRouter(RouterConfig{Enabled: true}) + result := r.ClassifyAndOverride(body, "req_test") + + if !result.OverrideApplied { + t.Fatalf("expected override for medium read tool-result turn (turn=%q, size=%q, tool=%q)", + result.TurnType, result.ToolResultSize, result.ToolName) + } + if result.OverrideTier != TierFlash { + t.Errorf("expected override_tier=flash, got %q", result.OverrideTier) + } + if result.OverrideModel != "deepseek-v4-flash" { + t.Errorf("expected deepseek-v4-flash for medium read tool result, got %q", result.OverrideModel) + } + // Flash tier must cap max_tokens on the body. + var capVal int + switch v := body["max_tokens"].(type) { + case int: + capVal = v + case json.Number: + n, _ := v.Int64() + capVal = int(n) + case float64: + capVal = int(v) + } + if capVal != toolVerbosityCap { + t.Errorf("expected body max_tokens capped to %d, got %d", toolVerbosityCap, capVal) + } +} + +// TestPerTurnVerbosityCap_AppliedEvenWhenModelOverrideNoop verifies that when +// the flow already runs on flash (so the model override is a no-op), a small +// tool-result turn still receives the flash verbosity cap. +func TestPerTurnVerbosityCap_AppliedEvenWhenModelOverrideNoop(t *testing.T) { + body := map[string]any{ + "model": "deepseek-v4-flash", + "messages": []any{ + map[string]any{"role": "system", "content": strings.Repeat("x", 15000)}, + map[string]any{"role": "user", "content": "refactor the auth module"}, + map[string]any{ + "role": "assistant", + "content": "", + "tool_calls": []any{ + map[string]any{ + "id": "call_shell", + "type": "function", + "function": map[string]any{"name": "Shell", "arguments": "{}"}, + }, + }, + }, + map[string]any{"role": "tool", "tool_call_id": "call_shell", "content": "ok"}, + }, + } + r := NewSmartRouter(RouterConfig{Enabled: true}) + result := r.ClassifyAndOverride(body, "req_test") + + // Override is a no-op (flash → flash) but the cap should still be applied. + if result.OverrideApplied { + t.Errorf("expected NO override (flash → flash no-op), got override to %q", result.OverrideModel) + } + var capVal int + switch v := body["max_tokens"].(type) { + case int: + capVal = v + case json.Number: + n, _ := v.Int64() + capVal = int(n) + case float64: + capVal = int(v) + } + if capVal != toolVerbosityCap { + t.Errorf("expected body max_tokens capped to %d on flash no-op turn, got %d", toolVerbosityCap, capVal) + } +} + func TestShouldDowngrade(t *testing.T) { tests := []struct { class RequestClass @@ -647,6 +861,200 @@ func TestShouldDowngrade(t *testing.T) { } } +// TestOverrideTier covers the decision-aware per-turn and content routing tier +// logic. Tool-result turns route cheap (small / medium read-only) results to +// flash, decision-heavy (medium write/error) results to pro, and large results +// to keep. Non-tool turns follow the content classifier. +func TestOverrideTier(t *testing.T) { + tests := []struct { + name string + result ClassifierResult + body map[string]any + want OverrideTier + }{ + {"tool small → flash", ClassifierResult{TurnType: TurnToolResult, ToolResultSize: "small", ToolName: "Shell"}, nil, TierFlash}, + {"tool medium read → flash", ClassifierResult{TurnType: TurnToolResult, ToolResultSize: "medium", ToolName: "Read"}, nil, TierFlash}, + {"tool medium grep → flash", ClassifierResult{TurnType: TurnToolResult, ToolResultSize: "medium", ToolName: "Grep"}, nil, TierFlash}, + {"tool medium glob → flash", ClassifierResult{TurnType: TurnToolResult, ToolResultSize: "medium", ToolName: "Glob"}, nil, TierFlash}, + {"tool medium write → pro", ClassifierResult{TurnType: TurnToolResult, ToolResultSize: "medium", ToolName: "StrReplace"}, nil, TierPro}, + {"tool medium shell → pro", ClassifierResult{TurnType: TurnToolResult, ToolResultSize: "medium", ToolName: "Shell"}, nil, TierPro}, + {"tool medium unknown tool → pro (write default)", ClassifierResult{TurnType: TurnToolResult, ToolResultSize: "medium", ToolName: "CustomThing"}, nil, TierPro}, + {"tool medium empty tool name → pro (conservative default)", ClassifierResult{TurnType: TurnToolResult, ToolResultSize: "medium", ToolName: ""}, nil, TierPro}, + {"tool large → keep", ClassifierResult{TurnType: TurnToolResult, ToolResultSize: "large", ToolName: "Shell"}, nil, TierKeep}, + {"tool empty size → keep", ClassifierResult{TurnType: TurnToolResult, ToolResultSize: "", ToolName: "Read"}, nil, TierKeep}, + {"content simple lookup → flash", ClassifierResult{TurnType: TurnUserPrompt, RequestClass: ClassSimpleLookup}, nil, TierFlash}, + {"content code search → flash", ClassifierResult{TurnType: TurnUserPrompt, RequestClass: ClassCodeSearch}, nil, TierFlash}, + {"content editing → keep", ClassifierResult{TurnType: TurnUserPrompt, RequestClass: ClassEditing}, nil, TierKeep}, + {"content complex reasoning → keep", ClassifierResult{TurnType: TurnUserPrompt, RequestClass: ClassComplexReasoning}, nil, TierKeep}, + {"content unknown → keep", ClassifierResult{TurnType: TurnUserPrompt, RequestClass: ClassUnknown}, nil, TierKeep}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // overrideTier uses result fields only; body is unused today but kept + // in the signature for future signals. + if got := overrideTier(tt.result, tt.body); got != tt.want { + t.Errorf("overrideTier(%+v) = %q, want %q", tt.result, got, tt.want) + } + }) + } +} + +// TestIsWriteTool verifies the read-only vs write/decision tool classification. +func TestIsWriteTool(t *testing.T) { + readOnly := []string{"Read", "Grep", "Glob", "WebSearch", "WebFetch", "FetchMcpResource"} + for _, name := range readOnly { + if isWriteTool(name) { + t.Errorf("isWriteTool(%q) = true, want false (read-only)", name) + } + } + + writeTools := []string{"Shell", "StrReplace", "EditNotebook", "Write", "Delete", "", "CustomTool"} + for _, name := range writeTools { + if !isWriteTool(name) { + t.Errorf("isWriteTool(%q) = false, want true (write/decision/unknown)", name) + } + } +} + +// TestToolResultName verifies we can resolve the function name for the last +// tool-result message from the preceding assistant tool_calls entry. +func TestToolResultName(t *testing.T) { + body := map[string]any{ + "messages": []any{ + map[string]any{"role": "system", "content": "sys"}, + map[string]any{"role": "user", "content": "do the thing"}, + map[string]any{ + "role": "assistant", + "content": "", + "tool_calls": []any{ + map[string]any{ + "id": "call_read", + "type": "function", + "function": map[string]any{ + "name": "Read", + "arguments": "{}", + }, + }, + map[string]any{ + "id": "call_shell", + "type": "function", + "function": map[string]any{ + "name": "Shell", + "arguments": "{}", + }, + }, + }, + }, + map[string]any{"role": "tool", "tool_call_id": "call_shell", "content": "exit 0"}, + }, + } + if got := toolResultName(body); got != "Shell" { + t.Errorf("toolResultName() = %q, want Shell", got) + } + + // Empty messages → "". + if got := toolResultName(map[string]any{"messages": []any{}}); got != "" { + t.Errorf("toolResultName(empty) = %q, want empty", got) + } + + // Not a tool-result last message → "". + noTool := map[string]any{ + "messages": []any{map[string]any{"role": "user", "content": "hi"}}, + } + if got := toolResultName(noTool); got != "" { + t.Errorf("toolResultName(no tool) = %q, want empty", got) + } +} + +// TestOverrideModelForTier verifies tier→model resolution across the maps and +// default fallbacks, including the pro/no-op cases. +func TestOverrideModelForTier(t *testing.T) { + tests := []struct { + name string + original string + tier OverrideTier + want string + }{ + {"flash deepseek-pro → flash", "deepseek-v4-pro", TierFlash, "deepseek-v4-flash"}, + {"flash unknown → default flash", "kimi-k3", TierFlash, "deepseek-v4-flash"}, + {"pro glm-5.2 → deepseek-pro", "glm-5.2", TierPro, "deepseek-v4-pro"}, + {"pro kimi-k3 → deepseek-pro", "kimi-k3", TierPro, "deepseek-v4-pro"}, + {"pro deepseek-pro → deepseek-pro (no-op)", "deepseek-v4-pro", TierPro, "deepseek-v4-pro"}, + {"pro unknown → default pro", "thaura", TierPro, "deepseek-v4-pro"}, + {"keep deepseek-pro → unchanged", "deepseek-v4-pro", TierKeep, "deepseek-v4-pro"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := overrideModelForTier(tt.original, tt.tier); got != tt.want { + t.Errorf("overrideModelForTier(%q, %q) = %q, want %q", tt.original, tt.tier, got, tt.want) + } + }) + } +} + +// TestShouldDumpForDiagnostics validates the targeted diagnostic-dump triggers. +// The full-message dump is always operator-gated by --diagnostic-dump; this +// function only narrows which requests within that session produce a file. +func TestShouldDumpForDiagnostics(t *testing.T) { + tests := []struct { + name string + turnType TurnType + toolResultSize string + wouldDowngrade bool + lastMsgLen int + strippedLen int + requestID string + want bool + }{ + {"extraction failed (stripped_len 0) → dump", TurnUserPrompt, "", false, 100, 0, "req_a", true}, + {"tool-result turn small → dump (schema validation)", TurnToolResult, "small", true, 100, 40, "req_b", true}, + {"tool-result turn large → dump", TurnToolResult, "large", false, 100, 40, "req_c", true}, + {"would downgrade → dump", TurnUserPrompt, "", true, 100, 40, "req_d", true}, + {"no-op extraction (stripped==last, non-zero) → dump", TurnUserPrompt, "", false, 40, 40, "req_e", true}, + // A user_prompt that extracted fine with no override hits only the 1% + // deterministic sample, which is not guaranteed true for a fixed ID, so + // we do not assert it here; see TestShouldDumpForDiagnostics_Sample. + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := shouldDumpForDiagnostics(tt.turnType, tt.toolResultSize, tt.wouldDowngrade, tt.lastMsgLen, tt.strippedLen, tt.requestID); got != tt.want { + t.Errorf("shouldDumpForDiagnostics(%q, %q, %v, %d, %d, %q) = %v, want %v", + tt.turnType, tt.toolResultSize, tt.wouldDowngrade, tt.lastMsgLen, tt.strippedLen, tt.requestID, got, tt.want) + } + }) + } +} + +// TestShouldDumpForDiagnostics_Sample verifies the deterministic 1%-sample +// path: the same requestID always yields the same decision (tokenized by FNV, +// not by a clock or counter), and at most 1-in-100 of distinct IDs dump on the +// sample branch alone. +func TestShouldDumpForDiagnostics_Sample(t *testing.T) { + // Same ID, twice → same result (deterministic, no per-process state). + id := "req_deterministic_sample" + a := shouldDumpForDiagnostics(TurnUserPrompt, "", false, 100, 40, id) + b := shouldDumpForDiagnostics(TurnUserPrompt, "", false, 100, 40, id) + if a != b { + t.Fatalf("sample decision not deterministic for same requestID: %v vs %v", a, b) + } + + // Distinct IDs: count how many of 10k hit the sample branch only. + var sampled int + for i := 0; i < 10000; i++ { + if shouldDumpForDiagnostics(TurnUserPrompt, "", false, 100, 40, fmt.Sprintf("req_sample_%d", i)) { + sampled++ + } + } + // Allow slack for the exact 1% boundary across the modulus; require it to be + // on the order of 1% (roughly 100 +/- 30), not 0 or 50%+. + if sampled < 70 || sampled > 130 { + t.Fatalf("deterministic sample rate off: got %d/10000 (~%0.2f%%), want ~1%%", sampled, float64(sampled)*0.01) + } +} + func TestLastUserMessage(t *testing.T) { tests := []struct { name string @@ -914,3 +1322,187 @@ func TestContentClassification_WithCursorXML(t *testing.T) { }) } } + +func TestDetectTurnType(t *testing.T) { + tests := []struct { + name string + body map[string]any + want TurnType + }{ + { + name: "tool result turn", + body: map[string]any{ + "messages": []any{ + map[string]any{"role": "system", "content": "You are a coding agent."}, + map[string]any{"role": "user", "content": "Run the tests"}, + map[string]any{"role": "assistant", "content": "", "tool_calls": []any{}}, + map[string]any{"role": "tool", "content": "ok\tinternal/gateway\t0.123s"}, + }, + }, + want: TurnToolResult, + }, + { + name: "user prompt turn", + body: map[string]any{ + "messages": []any{ + map[string]any{"role": "system", "content": "You are a coding agent."}, + map[string]any{"role": "user", "content": "What is a goroutine?"}, + }, + }, + want: TurnUserPrompt, + }, + { + name: "developer prompt turn", + body: map[string]any{ + "messages": []any{ + map[string]any{"role": "developer", "content": "You are a coding agent."}, + map[string]any{"role": "user", "content": "Hello"}, + }, + }, + want: TurnUserPrompt, + }, + { + name: "assistant continuation turn", + body: map[string]any{ + "messages": []any{ + map[string]any{"role": "system", "content": "You are a coding agent."}, + map[string]any{"role": "assistant", "content": "Let me think about the next step..."}, + }, + }, + want: TurnAgentContinue, + }, + { + name: "empty messages → unknown", + body: map[string]any{"messages": []any{}}, + want: TurnUnknown, + }, + { + name: "no messages key → unknown", + body: map[string]any{}, + want: TurnUnknown, + }, + { + name: "nil body → unknown", + body: nil, + want: TurnUnknown, + }, + { + name: "unknown role → unknown", + body: map[string]any{ + "messages": []any{ + map[string]any{"role": "system", "content": "Hello"}, + map[string]any{"role": "bogus", "content": "?"}, + }, + }, + want: TurnUnknown, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := detectTurnType(tt.body) + if got != tt.want { + t.Errorf("detectTurnType() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestToolResultSize(t *testing.T) { + tests := []struct { + name string + body map[string]any + want string + }{ + { + name: "small tool result (≤512)", + body: map[string]any{ + "messages": []any{ + map[string]any{"role": "system", "content": "Agent"}, + map[string]any{"role": "tool", "content": "ok"}, + }, + }, + want: "small", + }, + { + name: "medium tool result (≤4096)", + body: map[string]any{ + "messages": []any{ + map[string]any{"role": "system", "content": "Agent"}, + map[string]any{ + "role": "tool", + "content": strings.Repeat("x", 1024), + }, + }, + }, + want: "medium", + }, + { + name: "large tool result (>4096)", + body: map[string]any{ + "messages": []any{ + map[string]any{"role": "system", "content": "Agent"}, + map[string]any{ + "role": "tool", + "content": strings.Repeat("x", 5000), + }, + }, + }, + want: "large", + }, + { + name: "boundary: exactly 512 → small", + body: map[string]any{ + "messages": []any{ + map[string]any{"role": "system", "content": "Agent"}, + map[string]any{ + "role": "tool", + "content": strings.Repeat("x", 512), + }, + }, + }, + want: "small", + }, + { + name: "boundary: exactly 4096 → medium", + body: map[string]any{ + "messages": []any{ + map[string]any{"role": "system", "content": "Agent"}, + map[string]any{ + "role": "tool", + "content": strings.Repeat("x", 4096), + }, + }, + }, + want: "medium", + }, + { + name: "not a tool turn → empty", + body: map[string]any{ + "messages": []any{ + map[string]any{"role": "user", "content": "Hello"}, + }, + }, + want: "", + }, + { + name: "empty messages → empty", + body: map[string]any{"messages": []any{}}, + want: "", + }, + { + name: "nil body → empty", + body: nil, + want: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := toolResultSize(tt.body) + if got != tt.want { + t.Errorf("toolResultSize() = %q, want %q", got, tt.want) + } + }) + } +} diff --git a/internal/gateway/server.go b/internal/gateway/server.go index c6b7309..9d3f53b 100644 --- a/internal/gateway/server.go +++ b/internal/gateway/server.go @@ -106,11 +106,15 @@ func NewServer(cfg ServerConfig) (*Server, error) { } func (s *Server) routes() { + auth := func(h http.HandlerFunc) http.Handler { + return AuthMiddleware(h, s.cfg.GatewayKey) + } + s.mux.HandleFunc("GET /health", s.handleHealth) - s.mux.HandleFunc("GET /v1/models", s.handleModels) - s.mux.HandleFunc("POST /v1/models", s.handleModels) - s.mux.HandleFunc("POST /v1/chat/completions", s.handleChatCompletions) - s.mux.HandleFunc("POST /v1/responses", s.handleChatCompletions) + s.mux.Handle("GET /v1/models", auth(s.handleModels)) + s.mux.Handle("POST /v1/models", auth(s.handleModels)) + s.mux.Handle("POST /v1/chat/completions", auth(s.handleChatCompletions)) + s.mux.Handle("POST /v1/responses", auth(s.handleChatCompletions)) } // Handler returns the HTTP handler for httptest tests. diff --git a/internal/gateway/server_test.go b/internal/gateway/server_test.go index ccb2199..c0fa98d 100644 --- a/internal/gateway/server_test.go +++ b/internal/gateway/server_test.go @@ -456,3 +456,76 @@ func TestToolCallIDRetry(t *testing.T) { t.Fatalf("expected retry, calls=%d", calls.Load()) } } + +func TestUpstreamError_JSONErrorPassedVerbatim(t *testing.T) { + upstreamErr := `{"error":{"message":"rate limit exceeded","type":"rate_limit_error","code":"rate_limit_reached"}}` + env := setupEnv(t, "sk-moon", "", "", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(upstreamErr)) + }) + + res, body := env.doJSON(t, http.MethodPost, "/v1/chat/completions", true, map[string]any{ + "model": "gpt-4o", + "messages": []any{map[string]any{"role": "user", "content": "hi"}}, + }) + if res.StatusCode != http.StatusTooManyRequests { + t.Fatalf("status = %d, want %d", res.StatusCode, http.StatusTooManyRequests) + } + + var resp struct { + Error struct { + Message string `json:"message"` + Type string `json:"type"` + Code string `json:"code"` + } `json:"error"` + } + if err := json.Unmarshal(body, &resp); err != nil { + t.Fatalf("invalid JSON body: %v\nbody: %s", err, body) + } + if resp.Error.Message != "rate limit exceeded" { + t.Fatalf("error.message = %q, want %q", resp.Error.Message, "rate limit exceeded") + } + if resp.Error.Type != "rate_limit_error" { + t.Fatalf("error.type = %q, want %q", resp.Error.Type, "rate_limit_error") + } + if resp.Error.Code != "rate_limit_reached" { + t.Fatalf("error.code = %q, want %q", resp.Error.Code, "rate_limit_reached") + } +} + +func TestUpstreamError_NonJSONErrorWrappedWithRawBody(t *testing.T) { + rawErr := `503 Service Unavailable` + env := setupEnv(t, "sk-moon", "", "", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html") + w.WriteHeader(http.StatusServiceUnavailable) + _, _ = w.Write([]byte(rawErr)) + }) + + res, body := env.doJSON(t, http.MethodPost, "/v1/chat/completions", true, map[string]any{ + "model": "gpt-4o", + "messages": []any{map[string]any{"role": "user", "content": "hi"}}, + }) + if res.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want %d", res.StatusCode, http.StatusServiceUnavailable) + } + + var resp struct { + Error struct { + Message string `json:"message"` + Type string `json:"type"` + } `json:"error"` + } + if err := json.Unmarshal(body, &resp); err != nil { + t.Fatalf("invalid JSON body: %v\nbody: %s", err, body) + } + if !strings.Contains(resp.Error.Message, rawErr) { + t.Fatalf("error.message should contain raw upstream body:\ngot: %q\nwant substring: %q", resp.Error.Message, rawErr) + } + if !strings.Contains(resp.Error.Message, "503") { + t.Fatalf("error.message should contain status code 503: %q", resp.Error.Message) + } + if resp.Error.Type != "upstream_error" { + t.Fatalf("error.type = %q, want %q", resp.Error.Type, "upstream_error") + } +} diff --git a/internal/gateway/usage/usage.db b/internal/gateway/usage/usage.db new file mode 100644 index 0000000..7efe6dc Binary files /dev/null and b/internal/gateway/usage/usage.db differ diff --git a/internal/usage/query.go b/internal/usage/query.go index 8232199..9564386 100644 --- a/internal/usage/query.go +++ b/internal/usage/query.go @@ -54,96 +54,125 @@ type BucketModelBreakdown struct { EstUSD float64 `json:"est_usd"` } +// Shared SQL prefixes. Each query method appends its WHERE / GROUP BY / +// ORDER BY clause. The full SQL emitted (column selection, GROUP BY, ORDER BY) +// is the contract for upstream callers and must not change. +const ( + // eventDetailQuery selects the raw event columns for single-summary queries. + eventDetailQuery = `SELECT id, session_id, timestamp, provider, model, + prompt_tokens, completion_tokens, cache_hit_tokens, cache_miss_tokens, + est_usd, request_id, latency_ms + FROM events` + // dayAggregateQuery groups rows by calendar day. + dayAggregateQuery = `SELECT date(timestamp) as day, + COUNT(*) as reqs, + COALESCE(SUM(prompt_tokens),0), + COALESCE(SUM(completion_tokens),0), + COALESCE(SUM(cache_hit_tokens),0), + COALESCE(SUM(cache_miss_tokens),0), + COALESCE(SUM(est_usd),0) + FROM events` + // modelBreakdownQuery groups rows by provider and model. + modelBreakdownQuery = `SELECT provider, model, + COUNT(*) as reqs, + COALESCE(SUM(prompt_tokens),0), + COALESCE(SUM(completion_tokens),0), + COALESCE(SUM(cache_hit_tokens),0), + COALESCE(SUM(cache_miss_tokens),0), + COALESCE(SUM(est_usd),0) + FROM events` + // providerBreakdownQuery groups rows by provider. + providerBreakdownQuery = `SELECT provider, + COUNT(*) as reqs, + COALESCE(SUM(prompt_tokens),0), + COALESCE(SUM(completion_tokens),0), + COALESCE(SUM(cache_hit_tokens),0), + COALESCE(SUM(cache_miss_tokens),0), + COALESCE(SUM(est_usd),0) + FROM events` + // sessionsQueryBase aggregates per-session totals; QuerySessionsSince appends + // an optional WHERE clause and the GROUP BY / ORDER BY. + sessionsQueryBase = `SELECT session_id, + COUNT(*) as reqs, + COALESCE(SUM(prompt_tokens),0), + COALESCE(SUM(completion_tokens),0), + COALESCE(SUM(est_usd),0), + MIN(timestamp) as first_seen, + MAX(timestamp) as last_seen + FROM events` +) + // QueryDailyTotals returns a DailySummary for a specific date (YYYY-MM-DD). func (s *Store) QueryDailyTotals(date string) (DailySummary, error) { - rows, err := s.db.Query( - `SELECT id, session_id, timestamp, provider, model, - prompt_tokens, completion_tokens, cache_hit_tokens, cache_miss_tokens, - est_usd, request_id, latency_ms - FROM events WHERE date(timestamp) = ? ORDER BY timestamp ASC`, date) - if err != nil { - return DailySummary{}, fmt.Errorf("query daily totals: %w", err) - } - defer func() { _ = rows.Close() }() - - return buildDailySummary(rows, date, "") + return queryAndScan(s, "query daily totals", + eventDetailQuery+` WHERE date(timestamp) = ? ORDER BY timestamp ASC`, []any{date}, + func(rows *sql.Rows) (DailySummary, error) { + return buildDailySummary(rows, date, "") + }) } // QueryLastNDays returns DailySummary entries for the last N calendar days. func (s *Store) QueryLastNDays(n int) ([]DailySummary, error) { - rows, err := s.db.Query( - `SELECT date(timestamp) as day, - COUNT(*) as reqs, - COALESCE(SUM(prompt_tokens),0), - COALESCE(SUM(completion_tokens),0), - COALESCE(SUM(cache_hit_tokens),0), - COALESCE(SUM(cache_miss_tokens),0), - COALESCE(SUM(est_usd),0) - FROM events - GROUP BY day - ORDER BY day DESC - LIMIT ?`, n) - if err != nil { - return nil, fmt.Errorf("query last n days: %w", err) - } - defer func() { _ = rows.Close() }() - return scanDaySummaries(rows) + return queryAndScan(s, "query last n days", + dayAggregateQuery+` GROUP BY day ORDER BY day DESC LIMIT ?`, []any{n}, + func(rows *sql.Rows) ([]DailySummary, error) { + return scanRows(rows, scanDaySummaryRow) + }) } // QueryByDaySince returns DailySummary entries grouped by day since a given time. // When window is sub-day (e.g. 1h/3h/12h), groups by a configurable bucket. // bucketMinutes: 0 means group by day (date). >0 means group by floor(timestamp / bucket). func (s *Store) QueryByDaySince(since time.Time, bucketMinutes int) ([]DailySummary, error) { - var rows *sql.Rows - var err error if bucketMinutes > 0 { bucketSecs := bucketMinutes * 60 - rows, err = s.db.Query( - `SELECT strftime('%Y-%m-%dT%H:%M:00', - datetime((CAST(strftime('%s', timestamp) AS INTEGER) / ?) * ?, 'unixepoch')) as bucket, - COUNT(*) as reqs, - COALESCE(SUM(prompt_tokens),0), - COALESCE(SUM(completion_tokens),0), - COALESCE(SUM(cache_hit_tokens),0), - COALESCE(SUM(cache_miss_tokens),0), - COALESCE(SUM(est_usd),0) - FROM events WHERE timestamp >= ? - GROUP BY bucket - ORDER BY bucket ASC`, bucketSecs, bucketSecs, since.UTC().Format(time.RFC3339)) - } else { - rows, err = s.db.Query( - `SELECT date(timestamp) as day, - COUNT(*) as reqs, - COALESCE(SUM(prompt_tokens),0), - COALESCE(SUM(completion_tokens),0), - COALESCE(SUM(cache_hit_tokens),0), - COALESCE(SUM(cache_miss_tokens),0), - COALESCE(SUM(est_usd),0) - FROM events WHERE timestamp >= ? - GROUP BY day - ORDER BY day ASC`, since.UTC().Format(time.RFC3339)) + q := `SELECT strftime('%Y-%m-%dT%H:%M:00', + datetime((CAST(strftime('%s', timestamp) AS INTEGER) / ?) * ?, 'unixepoch')) as bucket, + COUNT(*) as reqs, + COALESCE(SUM(prompt_tokens),0), + COALESCE(SUM(completion_tokens),0), + COALESCE(SUM(cache_hit_tokens),0), + COALESCE(SUM(cache_miss_tokens),0), + COALESCE(SUM(est_usd),0) + FROM events WHERE timestamp >= ? + GROUP BY bucket + ORDER BY bucket ASC` + return queryAndScan(s, "query by day since", q, + []any{bucketSecs, bucketSecs, since.UTC().Format(time.RFC3339)}, + func(rows *sql.Rows) ([]DailySummary, error) { + return scanRows(rows, scanDaySummaryRow) + }) } + return queryAndScan(s, "query by day since", + dayAggregateQuery+` WHERE timestamp >= ? GROUP BY day ORDER BY day ASC`, + []any{since.UTC().Format(time.RFC3339)}, + func(rows *sql.Rows) ([]DailySummary, error) { + return scanRows(rows, scanDaySummaryRow) + }) +} + +// queryAndScan runs q with args, closes the result rows, and hands them to build. +// build owns scanning/consuming rows; any failure there is returned as-is. +func queryAndScan[T any](s *Store, label, q string, args []any, build func(*sql.Rows) (T, error)) (T, error) { + rows, err := s.db.Query(q, args...) if err != nil { - return nil, fmt.Errorf("query by day since: %w", err) + var zero T + return zero, fmt.Errorf("%s: %w", label, err) } defer func() { _ = rows.Close() }() - return scanDaySummaries(rows) + return build(rows) } -// scanDaySummaries scans rows grouped by a time bucket into DailySummary slices. -func scanDaySummaries(rows *sql.Rows) ([]DailySummary, error) { - var out []DailySummary +// scanRows walks each row, using perRow to scan a single typed value, and returns +// the collected slice. It surfaces rows.Err() after iteration. +func scanRows[T any](rows *sql.Rows, perRow func(*sql.Rows) (T, error)) ([]T, error) { + var out []T for rows.Next() { - var bucket string - var ds DailySummary - if err := rows.Scan(&bucket, &ds.RequestCount, &ds.TokensIn, - &ds.TokensOut, &ds.CacheHitTokens, &ds.CacheMissTokens, &ds.EstUSD); err != nil { - return nil, fmt.Errorf("scan day/bucket: %w", err) + item, err := perRow(rows) + if err != nil { + return nil, err } - ds.Date = bucket - ds.EstUSD = RoundUSD(ds.EstUSD) - ds.CursorReference = cursorRef() - out = append(out, ds) + out = append(out, item) } if err := rows.Err(); err != nil { return nil, fmt.Errorf("rows: %w", err) @@ -151,6 +180,18 @@ func scanDaySummaries(rows *sql.Rows) ([]DailySummary, error) { return out, nil } +// scanDaySummaryRow scans one row of dayAggregateQuery (or the bucket variant). +func scanDaySummaryRow(rows *sql.Rows) (DailySummary, error) { + var ds DailySummary + if err := rows.Scan(&ds.Date, &ds.RequestCount, &ds.TokensIn, + &ds.TokensOut, &ds.CacheHitTokens, &ds.CacheMissTokens, &ds.EstUSD); err != nil { + return DailySummary{}, fmt.Errorf("scan day/bucket: %w", err) + } + ds.EstUSD = RoundUSD(ds.EstUSD) + ds.CursorReference = cursorRef() + return ds, nil +} + // QueryByDayModelSince returns per-model breakdown per bucket since a given time. // This powers the Spend by Period chart split by model instead of cache hit/miss. func (s *Store) QueryByDayModelSince(since time.Time, bucketMinutes int) ([]BucketModelBreakdown, error) { @@ -200,17 +241,60 @@ func (s *Store) QueryByDayModelSince(since time.Time, bucketMinutes int) ([]Buck // QuerySessionDetail returns a DailySummary for a specific session ID. func (s *Store) QuerySessionDetail(sessionID string) (DailySummary, error) { - rows, err := s.db.Query( - `SELECT id, session_id, timestamp, provider, model, - prompt_tokens, completion_tokens, cache_hit_tokens, cache_miss_tokens, - est_usd, request_id, latency_ms - FROM events WHERE session_id = ? ORDER BY timestamp ASC`, sessionID) - if err != nil { - return DailySummary{}, fmt.Errorf("query session: %w", err) + return queryAndScan(s, "query session", + eventDetailQuery+` WHERE session_id = ? ORDER BY timestamp ASC`, []any{sessionID}, + func(rows *sql.Rows) (DailySummary, error) { + return buildDailySummary(rows, "", sessionID) + }) +} + +// accumulateDailyEvent folds a single event into ds and its byModel breakdown. +// It is pure: it reads only ev and mutates only ds/byModel, never the DB. +func accumulateDailyEvent(ds *DailySummary, byModel map[string]*ModelBreakdown, ev Event) { + ds.RequestCount++ + ds.TokensIn += ev.PromptTokens + ds.TokensOut += ev.CompletionTokens + ds.CacheHitTokens += ev.CacheHitTokens + ds.CacheMissTokens += ev.CacheMissTokens + ds.EstUSD += ev.EstUSD + + mb, ok := byModel[ev.Model] + if !ok { + mb = &ModelBreakdown{Model: ev.Model, Provider: string(ev.Provider)} + byModel[ev.Model] = mb } - defer func() { _ = rows.Close() }() + mb.RequestCount++ + mb.TokensIn += ev.PromptTokens + mb.TokensOut += ev.CompletionTokens + mb.CacheHitTokens += ev.CacheHitTokens + mb.CacheMissTokens += ev.CacheMissTokens + mb.EstUSD += ev.EstUSD +} - return buildDailySummary(rows, "", sessionID) +// finalizeDailySummary rounds estimates and folds the model map into ds.ByModel. +func finalizeDailySummary(ds DailySummary, byModel map[string]*ModelBreakdown) DailySummary { + ds.EstUSD = RoundUSD(ds.EstUSD) + for _, mb := range byModel { + mb.EstUSD = RoundUSD(mb.EstUSD) + ds.ByModel = append(ds.ByModel, *mb) + } + return ds +} + +// scanEventRow scans one raw event row into an Event, parsing its provider/timestamp. +func scanEventRow(rows *sql.Rows) (Event, error) { + var ev Event + var tsStr, provStr string + if err := rows.Scan( + &ev.ID, &ev.SessionID, &tsStr, &provStr, &ev.Model, + &ev.PromptTokens, &ev.CompletionTokens, &ev.CacheHitTokens, &ev.CacheMissTokens, + &ev.EstUSD, &ev.RequestID, &ev.LatencyMS, + ); err != nil { + return Event{}, fmt.Errorf("scan event: %w", err) + } + ev.Timestamp, _ = time.Parse(time.RFC3339Nano, tsStr) + ev.Provider = config.Provider(provStr) + return ev, nil } // buildDailySummary scans event rows into a DailySummary. @@ -223,56 +307,21 @@ func buildDailySummary(rows *sql.Rows, date, sessionID string) (DailySummary, er byModel := make(map[string]*ModelBreakdown) for rows.Next() { - var ev Event - var tsStr, provStr string - if err := rows.Scan( - &ev.ID, &ev.SessionID, &tsStr, &provStr, &ev.Model, - &ev.PromptTokens, &ev.CompletionTokens, &ev.CacheHitTokens, &ev.CacheMissTokens, - &ev.EstUSD, &ev.RequestID, &ev.LatencyMS, - ); err != nil { - return DailySummary{}, fmt.Errorf("scan event: %w", err) + ev, err := scanEventRow(rows) + if err != nil { + return DailySummary{}, err } - ev.Timestamp, _ = time.Parse(time.RFC3339Nano, tsStr) - ev.Provider = config.Provider(provStr) - if sessionID != "" && ev.SessionID != sessionID { continue } - - ds.RequestCount++ - ds.TokensIn += ev.PromptTokens - ds.TokensOut += ev.CompletionTokens - ds.CacheHitTokens += ev.CacheHitTokens - ds.CacheMissTokens += ev.CacheMissTokens - ds.EstUSD += ev.EstUSD - - mb, ok := byModel[ev.Model] - if !ok { - mb = &ModelBreakdown{ - Model: ev.Model, - Provider: string(ev.Provider), - } - byModel[ev.Model] = mb - } - mb.RequestCount++ - mb.TokensIn += ev.PromptTokens - mb.TokensOut += ev.CompletionTokens - mb.CacheHitTokens += ev.CacheHitTokens - mb.CacheMissTokens += ev.CacheMissTokens - mb.EstUSD += ev.EstUSD + accumulateDailyEvent(&ds, byModel, ev) } if err := rows.Err(); err != nil { return DailySummary{}, fmt.Errorf("rows: %w", err) } - ds.EstUSD = RoundUSD(ds.EstUSD) - for _, mb := range byModel { - mb.EstUSD = RoundUSD(mb.EstUSD) - ds.ByModel = append(ds.ByModel, *mb) - } - - return ds, nil + return finalizeDailySummary(ds, byModel), nil } func cursorRef() CursorReference { @@ -310,132 +359,73 @@ type SessionInfo struct { func (s *Store) QueryMonthToDate() (DailySummary, error) { now := time.Now().UTC() start := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC).Format("2006-01-02") - rows, err := s.db.Query( - `SELECT id, session_id, timestamp, provider, model, - prompt_tokens, completion_tokens, cache_hit_tokens, cache_miss_tokens, - est_usd, request_id, latency_ms - FROM events WHERE date(timestamp) >= ? ORDER BY timestamp ASC`, start) - if err != nil { - return DailySummary{}, fmt.Errorf("query mtd: %w", err) - } - defer func() { _ = rows.Close() }() - return buildDailySummary(rows, start, "") + return queryAndScan(s, "query mtd", + eventDetailQuery+` WHERE date(timestamp) >= ? ORDER BY timestamp ASC`, []any{start}, + func(rows *sql.Rows) (DailySummary, error) { + return buildDailySummary(rows, start, "") + }) } // QueryByModelSince returns usage breakdown by model since a given time. func (s *Store) QueryByModelSince(since time.Time) ([]ModelBreakdown, error) { - rows, err := s.db.Query( - `SELECT provider, model, - COUNT(*) as reqs, - COALESCE(SUM(prompt_tokens),0), - COALESCE(SUM(completion_tokens),0), - COALESCE(SUM(cache_hit_tokens),0), - COALESCE(SUM(cache_miss_tokens),0), - COALESCE(SUM(est_usd),0) - FROM events WHERE timestamp >= ? - GROUP BY provider, model - ORDER BY SUM(est_usd) DESC`, since.UTC().Format(time.RFC3339)) - if err != nil { - return nil, fmt.Errorf("query by model since: %w", err) - } - defer func() { _ = rows.Close() }() - return scanModelBreakdowns(rows) + return queryAndScan(s, "query by model since", + modelBreakdownQuery+` WHERE timestamp >= ? GROUP BY provider, model ORDER BY SUM(est_usd) DESC`, + []any{since.UTC().Format(time.RFC3339)}, + func(rows *sql.Rows) ([]ModelBreakdown, error) { + return scanRows(rows, scanModelBreakdownRow) + }) } // QueryByModel returns all-time usage breakdown by model. func (s *Store) QueryByModel() ([]ModelBreakdown, error) { - rows, err := s.db.Query( - `SELECT provider, model, - COUNT(*) as reqs, - COALESCE(SUM(prompt_tokens),0), - COALESCE(SUM(completion_tokens),0), - COALESCE(SUM(cache_hit_tokens),0), - COALESCE(SUM(cache_miss_tokens),0), - COALESCE(SUM(est_usd),0) - FROM events - GROUP BY provider, model - ORDER BY SUM(est_usd) DESC`) - if err != nil { - return nil, fmt.Errorf("query by model: %w", err) - } - defer func() { _ = rows.Close() }() - return scanModelBreakdowns(rows) + return queryAndScan(s, "query by model", + modelBreakdownQuery+` GROUP BY provider, model ORDER BY SUM(est_usd) DESC`, nil, + func(rows *sql.Rows) ([]ModelBreakdown, error) { + return scanRows(rows, scanModelBreakdownRow) + }) } -func scanModelBreakdowns(rows *sql.Rows) ([]ModelBreakdown, error) { - var out []ModelBreakdown - for rows.Next() { - var mb ModelBreakdown - if err := rows.Scan(&mb.Provider, &mb.Model, - &mb.RequestCount, &mb.TokensIn, &mb.TokensOut, - &mb.CacheHitTokens, &mb.CacheMissTokens, &mb.EstUSD); err != nil { - return nil, fmt.Errorf("scan model: %w", err) - } - mb.EstUSD = RoundUSD(mb.EstUSD) - out = append(out, mb) - } - if err := rows.Err(); err != nil { - return nil, fmt.Errorf("rows: %w", err) +// scanModelBreakdownRow scans one row of modelBreakdownQuery. +func scanModelBreakdownRow(rows *sql.Rows) (ModelBreakdown, error) { + var mb ModelBreakdown + if err := rows.Scan(&mb.Provider, &mb.Model, + &mb.RequestCount, &mb.TokensIn, &mb.TokensOut, + &mb.CacheHitTokens, &mb.CacheMissTokens, &mb.EstUSD); err != nil { + return ModelBreakdown{}, fmt.Errorf("scan model: %w", err) } - return out, nil + mb.EstUSD = RoundUSD(mb.EstUSD) + return mb, nil } // QueryByProviderSince returns usage breakdown by provider since a given time. func (s *Store) QueryByProviderSince(since time.Time) ([]ProviderBreakdown, error) { - rows, err := s.db.Query( - `SELECT provider, - COUNT(*) as reqs, - COALESCE(SUM(prompt_tokens),0), - COALESCE(SUM(completion_tokens),0), - COALESCE(SUM(cache_hit_tokens),0), - COALESCE(SUM(cache_miss_tokens),0), - COALESCE(SUM(est_usd),0) - FROM events WHERE timestamp >= ? - GROUP BY provider - ORDER BY SUM(est_usd) DESC`, since.UTC().Format(time.RFC3339)) - if err != nil { - return nil, fmt.Errorf("query by provider since: %w", err) - } - defer func() { _ = rows.Close() }() - return scanProviderBreakdowns(rows) + return queryAndScan(s, "query by provider since", + providerBreakdownQuery+` WHERE timestamp >= ? GROUP BY provider ORDER BY SUM(est_usd) DESC`, + []any{since.UTC().Format(time.RFC3339)}, + func(rows *sql.Rows) ([]ProviderBreakdown, error) { + return scanRows(rows, scanProviderBreakdownRow) + }) } // QueryByProvider returns all-time usage breakdown by provider. func (s *Store) QueryByProvider() ([]ProviderBreakdown, error) { - rows, err := s.db.Query( - `SELECT provider, - COUNT(*) as reqs, - COALESCE(SUM(prompt_tokens),0), - COALESCE(SUM(completion_tokens),0), - COALESCE(SUM(cache_hit_tokens),0), - COALESCE(SUM(cache_miss_tokens),0), - COALESCE(SUM(est_usd),0) - FROM events - GROUP BY provider - ORDER BY SUM(est_usd) DESC`) - if err != nil { - return nil, fmt.Errorf("query by provider: %w", err) - } - defer func() { _ = rows.Close() }() - return scanProviderBreakdowns(rows) + return queryAndScan(s, "query by provider", + providerBreakdownQuery+` GROUP BY provider ORDER BY SUM(est_usd) DESC`, nil, + func(rows *sql.Rows) ([]ProviderBreakdown, error) { + return scanRows(rows, scanProviderBreakdownRow) + }) } -func scanProviderBreakdowns(rows *sql.Rows) ([]ProviderBreakdown, error) { - var out []ProviderBreakdown - for rows.Next() { - var pb ProviderBreakdown - if err := rows.Scan(&pb.Provider, - &pb.RequestCount, &pb.TokensIn, &pb.TokensOut, - &pb.CacheHitTokens, &pb.CacheMissTokens, &pb.EstUSD); err != nil { - return nil, fmt.Errorf("scan provider: %w", err) - } - pb.EstUSD = RoundUSD(pb.EstUSD) - out = append(out, pb) +// scanProviderBreakdownRow scans one row of providerBreakdownQuery. +func scanProviderBreakdownRow(rows *sql.Rows) (ProviderBreakdown, error) { + var pb ProviderBreakdown + if err := rows.Scan(&pb.Provider, + &pb.RequestCount, &pb.TokensIn, &pb.TokensOut, + &pb.CacheHitTokens, &pb.CacheMissTokens, &pb.EstUSD); err != nil { + return ProviderBreakdown{}, fmt.Errorf("scan provider: %w", err) } - if err := rows.Err(); err != nil { - return nil, fmt.Errorf("rows: %w", err) - } - return out, nil + pb.EstUSD = RoundUSD(pb.EstUSD) + return pb, nil } // QuerySessions returns a list of all unique sessions with summary info. @@ -445,53 +435,33 @@ func (s *Store) QuerySessions() ([]SessionInfo, error) { // QuerySessionsSince returns sessions whose last event is since a given time. func (s *Store) QuerySessionsSince(since time.Time) ([]SessionInfo, error) { - var rows *sql.Rows - var err error - if since.IsZero() { - rows, err = s.db.Query( - `SELECT session_id, - COUNT(*) as reqs, - COALESCE(SUM(prompt_tokens),0), - COALESCE(SUM(completion_tokens),0), - COALESCE(SUM(est_usd),0), - MIN(timestamp) as first_seen, - MAX(timestamp) as last_seen - FROM events - GROUP BY session_id - ORDER BY MAX(timestamp) DESC`) - } else { - rows, err = s.db.Query( - `SELECT session_id, - COUNT(*) as reqs, - COALESCE(SUM(prompt_tokens),0), - COALESCE(SUM(completion_tokens),0), - COALESCE(SUM(est_usd),0), - MIN(timestamp) as first_seen, - MAX(timestamp) as last_seen - FROM events WHERE timestamp >= ? - GROUP BY session_id - ORDER BY MAX(timestamp) DESC`, since.UTC().Format(time.RFC3339)) - } - if err != nil { - return nil, fmt.Errorf("query sessions: %w", err) + q := sessionsQueryBase + var args []any + if !since.IsZero() { + q += " WHERE timestamp >= ?" + args = append(args, since.UTC().Format(time.RFC3339)) } - defer func() { _ = rows.Close() }() + q += " GROUP BY session_id ORDER BY MAX(timestamp) DESC" - var out = make([]SessionInfo, 0) - for rows.Next() { - var si SessionInfo - if err := rows.Scan(&si.SessionID, - &si.RequestCount, &si.TokensIn, &si.TokensOut, &si.EstUSD, - &si.FirstSeen, &si.LastSeen); err != nil { - return nil, fmt.Errorf("scan session: %w", err) - } - si.EstUSD = RoundUSD(si.EstUSD) - out = append(out, si) + sessions, err := queryAndScan(s, "query sessions", q, args, func(rows *sql.Rows) ([]SessionInfo, error) { + return scanRows(rows, scanSessionInfoRow) + }) + if sessions == nil { + sessions = []SessionInfo{} } - if err := rows.Err(); err != nil { - return nil, fmt.Errorf("rows: %w", err) + return sessions, err +} + +// scanSessionInfoRow scans one row of sessionsQueryBase (+ WHERE/GROUP BY/ORDER BY). +func scanSessionInfoRow(rows *sql.Rows) (SessionInfo, error) { + var si SessionInfo + if err := rows.Scan(&si.SessionID, + &si.RequestCount, &si.TokensIn, &si.TokensOut, &si.EstUSD, + &si.FirstSeen, &si.LastSeen); err != nil { + return SessionInfo{}, fmt.Errorf("scan session: %w", err) } - return out, nil + si.EstUSD = RoundUSD(si.EstUSD) + return si, nil } // DBStats holds database-level statistics. diff --git a/internal/usage/query_test.go b/internal/usage/query_test.go index 58df72a..2077d7c 100644 --- a/internal/usage/query_test.go +++ b/internal/usage/query_test.go @@ -1,6 +1,8 @@ package usage import ( + "database/sql" + "reflect" "testing" "time" @@ -193,3 +195,406 @@ func TestQueryEmptyReturns(t *testing.T) { } }) } + +func TestAccumulateDailyEvent(t *testing.T) { + existing := &ModelBreakdown{ + Model: "kimi-k3", Provider: string(config.ProviderMoonshot), + RequestCount: 1, TokensIn: 100, TokensOut: 50, + CacheHitTokens: 10, CacheMissTokens: 90, EstUSD: 1.25, + } + + tests := []struct { + name string + ev Event + seedDS *DailySummary + seedModels map[string]ModelBreakdown + wantDS DailySummary + wantModels map[string]ModelBreakdown + }{ + { + name: "first event into empty summary", + ev: Event{Model: "kimi-k3", Provider: config.ProviderMoonshot, + PromptTokens: 100, CompletionTokens: 50, + CacheHitTokens: 10, CacheMissTokens: 90, EstUSD: 1.25}, + wantDS: DailySummary{RequestCount: 1, TokensIn: 100, TokensOut: 50, + CacheHitTokens: 10, CacheMissTokens: 90, EstUSD: 1.25}, + wantModels: map[string]ModelBreakdown{ + "kimi-k3": {Model: "kimi-k3", Provider: string(config.ProviderMoonshot), + RequestCount: 1, TokensIn: 100, TokensOut: 50, + CacheHitTokens: 10, CacheMissTokens: 90, EstUSD: 1.25}, + }, + }, + { + name: "second event into existing model", + ev: Event{Model: "kimi-k3", Provider: config.ProviderMoonshot, + PromptTokens: 50, CompletionTokens: 25, EstUSD: 0.75}, + seedDS: &DailySummary{RequestCount: 1, TokensIn: 100, TokensOut: 50, + CacheHitTokens: 10, CacheMissTokens: 90, EstUSD: 1.25}, + seedModels: map[string]ModelBreakdown{"kimi-k3": *existing}, + wantDS: DailySummary{RequestCount: 2, TokensIn: 150, TokensOut: 75, + CacheHitTokens: 10, CacheMissTokens: 90, EstUSD: 2.0}, + wantModels: map[string]ModelBreakdown{ + "kimi-k3": {Model: "kimi-k3", Provider: string(config.ProviderMoonshot), + RequestCount: 2, TokensIn: 150, TokensOut: 75, + CacheHitTokens: 10, CacheMissTokens: 90, EstUSD: 2.0}, + }, + }, + { + name: "event starts a distinct model breakdown", + ev: Event{Model: "deepseek-v4-flash", Provider: config.ProviderDeepSeek, + PromptTokens: 5, EstUSD: 0.5}, + seedDS: &DailySummary{RequestCount: 1, TokensIn: 100, TokensOut: 50, + CacheHitTokens: 10, CacheMissTokens: 90, EstUSD: 1.25}, + seedModels: map[string]ModelBreakdown{"kimi-k3": *existing}, + wantDS: DailySummary{RequestCount: 2, TokensIn: 105, TokensOut: 50, + CacheHitTokens: 10, CacheMissTokens: 90, EstUSD: 1.75}, + wantModels: map[string]ModelBreakdown{ + "kimi-k3": *existing, + "deepseek-v4-flash": {Model: "deepseek-v4-flash", Provider: string(config.ProviderDeepSeek), + RequestCount: 1, TokensIn: 5, TokensOut: 0, EstUSD: 0.5}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ds := DailySummary{} + if tt.seedDS != nil { + ds = *tt.seedDS + } + byModel := make(map[string]*ModelBreakdown) + for k, v := range tt.seedModels { + mb := v + byModel[k] = &mb + } + + accumulateDailyEvent(&ds, byModel, tt.ev) + + // ByModel is never set by the accumulator; normalize to nil and + // compare deeply. + ds.ByModel = nil + if !reflect.DeepEqual(ds, tt.wantDS) { + t.Fatalf("summary mismatch:\n got %+v\nwant %+v", ds, tt.wantDS) + } + got := make(map[string]ModelBreakdown, len(byModel)) + for k, mb := range byModel { + got[k] = *mb + } + if !reflect.DeepEqual(got, tt.wantModels) { + t.Fatalf("model breakdowns mismatch:\n got %+v\nwant %+v", got, tt.wantModels) + } + }) + } +} + +func TestFinalizeDailySummary(t *testing.T) { + tests := []struct { + name string + ds DailySummary + models map[string]ModelBreakdown + want DailySummary + }{ + { + name: "rounds totals and folds model map", + ds: DailySummary{EstUSD: 0.123456}, + models: map[string]ModelBreakdown{ + "a": {Model: "a", EstUSD: 0.9876}, + }, + want: DailySummary{EstUSD: RoundUSD(0.123456), + ByModel: []ModelBreakdown{{Model: "a", EstUSD: 0.988}}}, + }, + { + name: "rounds floats to three decimals", + ds: DailySummary{EstUSD: 1.00005}, + models: map[string]ModelBreakdown{ + "b": {Model: "b", EstUSD: 2.34567}, + "c": {Model: "c", EstUSD: 0.0001}, + }, + want: DailySummary{EstUSD: RoundUSD(1.00005), + ByModel: []ModelBreakdown{ + {Model: "b", EstUSD: RoundUSD(2.34567)}, + {Model: "c", EstUSD: 0.0}, + }}, + }, + { + name: "no models yields empty ByModel", + ds: DailySummary{EstUSD: 2.5}, + want: DailySummary{EstUSD: 2.5, ByModel: nil}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + byModel := make(map[string]*ModelBreakdown, len(tt.models)) + for k, v := range tt.models { + mb := v + byModel[k] = &mb + } + got := finalizeDailySummary(tt.ds, byModel) + + if got.EstUSD != tt.want.EstUSD { + t.Fatalf("EstUSD mismatch: got %v, want %v", got.EstUSD, tt.want.EstUSD) + } + gotModels := make(map[string]ModelBreakdown, len(got.ByModel)) + for _, mb := range got.ByModel { + gotModels[mb.Model] = mb + } + wantModels := make(map[string]ModelBreakdown, len(tt.want.ByModel)) + for _, mb := range tt.want.ByModel { + wantModels[mb.Model] = mb + } + if !reflect.DeepEqual(gotModels, wantModels) { + t.Fatalf("ByModel mismatch:\n got %+v\nwant %+v", gotModels, wantModels) + } + }) + } +} + +func TestScanRows(t *testing.T) { + root := t.TempDir() + store, err := NewStore(root) + if err != nil { + t.Fatal(err) + } + + _, _ = store.Record(Event{ + SessionID: "s1", Provider: config.ProviderMoonshot, Model: "kimi-k3", + PromptTokens: 1000, CompletionTokens: 500, CacheHitTokens: 100, + Timestamp: time.Date(2026, 7, 15, 12, 0, 0, 0, time.UTC), + }) + _, _ = store.Record(Event{ + SessionID: "s1", Provider: config.ProviderMoonshot, Model: "kimi-k3", + PromptTokens: 2000, CompletionTokens: 1000, CacheHitTokens: 200, + Timestamp: time.Date(2026, 7, 15, 13, 0, 0, 0, time.UTC), + }) + _, _ = store.Record(Event{ + SessionID: "s2", Provider: config.ProviderDeepSeek, Model: "deepseek-v4-flash", + PromptTokens: 3000, Timestamp: time.Date(2026, 7, 16, 9, 0, 0, 0, time.UTC), + }) + + t.Run("day summaries grouped by day", func(t *testing.T) { + want := map[string]DailySummary{ + "2026-07-15": {Date: "2026-07-15", RequestCount: 2, TokensIn: 3000, TokensOut: 1500, CacheHitTokens: 300}, + "2026-07-16": {Date: "2026-07-16", RequestCount: 1, TokensIn: 3000, TokensOut: 0, CacheHitTokens: 0}, + } + rows := queryRows(t, store, dayAggregateQuery+" GROUP BY day ORDER BY day ASC", nil) + got, err := scanRows(rows, scanDaySummaryRow) + if err != nil { + t.Fatal(err) + } + assertDailyMap(t, got, want) + }) + + t.Run("model breakdowns grouped by provider,model", func(t *testing.T) { + want := map[string]ModelBreakdown{ + "kimi-k3": {Model: "kimi-k3", Provider: string(config.ProviderMoonshot), RequestCount: 2, TokensIn: 3000, TokensOut: 1500, CacheHitTokens: 300}, + "deepseek-v4-flash": {Model: "deepseek-v4-flash", Provider: string(config.ProviderDeepSeek), RequestCount: 1, TokensIn: 3000}, + } + rows := queryRows(t, store, modelBreakdownQuery+" GROUP BY provider, model ORDER BY SUM(est_usd) DESC", nil) + got, err := scanRows(rows, scanModelBreakdownRow) + if err != nil { + t.Fatal(err) + } + assertModelMap(t, got, want) + }) + + t.Run("provider breakdowns grouped by provider", func(t *testing.T) { + want := map[string]ProviderBreakdown{ + string(config.ProviderMoonshot): {Provider: string(config.ProviderMoonshot), RequestCount: 2, TokensIn: 3000, TokensOut: 1500, CacheHitTokens: 300}, + string(config.ProviderDeepSeek): {Provider: string(config.ProviderDeepSeek), RequestCount: 1, TokensIn: 3000}, + } + rows := queryRows(t, store, providerBreakdownQuery+" GROUP BY provider ORDER BY SUM(est_usd) DESC", nil) + got, err := scanRows(rows, scanProviderBreakdownRow) + if err != nil { + t.Fatal(err) + } + assertProviderMap(t, got, want) + }) + + t.Run("session info grouped by session", func(t *testing.T) { + want := map[string]SessionInfo{ + "s1": {SessionID: "s1", RequestCount: 2, TokensIn: 3000, TokensOut: 1500, FirstSeen: "2026-07-15T12:00:00Z", LastSeen: "2026-07-15T13:00:00Z"}, + "s2": {SessionID: "s2", RequestCount: 1, TokensIn: 3000, FirstSeen: "2026-07-16T09:00:00Z", LastSeen: "2026-07-16T09:00:00Z"}, + } + rows := queryRows(t, store, sessionsQueryBase+" GROUP BY session_id ORDER BY MAX(timestamp) DESC", nil) + got, err := scanRows(rows, scanSessionInfoRow) + if err != nil { + t.Fatal(err) + } + assertSessionMap(t, got, want) + }) + + t.Run("empty result yields nil slice", func(t *testing.T) { + rows := queryRows(t, store, dayAggregateQuery+" WHERE date(timestamp) >= ? GROUP BY day ORDER BY day ASC", []any{"2099-01-01"}) + got, err := scanRows(rows, scanDaySummaryRow) + if err != nil { + t.Fatal(err) + } + if got != nil { + t.Fatalf("expected nil slice, got %+v", got) + } + }) +} + +func TestQueryAndScan(t *testing.T) { + root := t.TempDir() + store, err := NewStore(root) + if err != nil { + t.Fatal(err) + } + _, _ = store.Record(Event{ + SessionID: "s1", Provider: config.ProviderMoonshot, Model: "kimi-k3", + PromptTokens: 1000, Timestamp: time.Date(2026, 7, 15, 12, 0, 0, 0, time.UTC), + }) + + tests := []struct { + name string + probe func() (int, error) + want int + }{ + { + name: "queries and scans rows", + probe: func() (int, error) { + return queryAndScan(store, "test query", modelBreakdownQuery, nil, func(rows *sql.Rows) (int, error) { + got, err := scanRows(rows, scanModelBreakdownRow) + return len(got), err + }) + }, + want: 1, + }, + { + name: "returns empty count for no rows", + probe: func() (int, error) { + return queryAndScan(store, "test query", modelBreakdownQuery+ + " WHERE timestamp >= ? GROUP BY provider, model", []any{"2099-01-01"}, + func(rows *sql.Rows) (int, error) { + got, err := scanRows(rows, scanModelBreakdownRow) + return len(got), err + }) + }, + want: 0, + }, + { + name: "propagates query errors", + probe: func() (int, error) { + return queryAndScan(store, "test query", "SELECT * FROM no_such_table", nil, + func(rows *sql.Rows) (int, error) { return 0, nil }) + }, + want: -1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := tt.probe() + if tt.want == -1 { + if err == nil { + t.Fatal("expected an error, got nil") + } + return + } + if err != nil { + t.Fatal(err) + } + if got != tt.want { + t.Fatalf("got %d, want %d", got, tt.want) + } + }) + } +} + +// queryRows runs q on the store and returns the rows for scanning. +func queryRows(t *testing.T, store *Store, q string, args []any) *sql.Rows { + t.Helper() + rows, err := store.db.Query(q, args...) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = rows.Close() }) + return rows +} + +func assertDailyMap(t *testing.T, got []DailySummary, want map[string]DailySummary) { + t.Helper() + gotMap := make(map[string]DailySummary, len(got)) + for _, d := range got { + gotMap[d.Date] = d + } + if len(gotMap) != len(want) { + t.Fatalf("day count mismatch: got %v, want %v", keysOf(gotMap), keysOf(want)) + } + for k, w := range want { + g, ok := gotMap[k] + if !ok { + t.Fatalf("missing day %q", k) + } + g.EstUSD, w.EstUSD = 0, 0 + g.CursorReference, w.CursorReference = CursorReference{}, CursorReference{} + g.ByModel, w.ByModel = nil, nil + if !reflect.DeepEqual(g, w) { + t.Errorf("day %q mismatch:\n got %+v\nwant %+v", k, g, w) + } + } +} + +func assertModelMap(t *testing.T, got []ModelBreakdown, want map[string]ModelBreakdown) { + t.Helper() + gotMap := make(map[string]ModelBreakdown, len(got)) + for _, m := range got { + gotMap[m.Model] = m + } + for k, w := range want { + g, ok := gotMap[k] + if !ok { + t.Fatalf("missing model %q", k) + } + g.EstUSD, w.EstUSD = 0, 0 + if g != w { + t.Errorf("model %q mismatch:\n got %+v\nwant %+v", k, g, w) + } + } +} + +func assertProviderMap(t *testing.T, got []ProviderBreakdown, want map[string]ProviderBreakdown) { + t.Helper() + gotMap := make(map[string]ProviderBreakdown, len(got)) + for _, p := range got { + gotMap[p.Provider] = p + } + for k, w := range want { + g, ok := gotMap[k] + if !ok { + t.Fatalf("missing provider %q", k) + } + g.EstUSD, w.EstUSD = 0, 0 + if g != w { + t.Errorf("provider %q mismatch:\n got %+v\nwant %+v", k, g, w) + } + } +} + +func assertSessionMap(t *testing.T, got []SessionInfo, want map[string]SessionInfo) { + t.Helper() + gotMap := make(map[string]SessionInfo, len(got)) + for _, si := range got { + gotMap[si.SessionID] = si + } + for k, w := range want { + g, ok := gotMap[k] + if !ok { + t.Fatalf("missing session %q", k) + } + g.EstUSD, w.EstUSD = 0, 0 + if g != w { + t.Errorf("session %q mismatch:\n got %+v\nwant %+v", k, g, w) + } + } +} + +func keysOf[K comparable](m map[K]DailySummary) []K { + out := make([]K, 0, len(m)) + for k := range m { + out = append(out, k) + } + return out +} diff --git a/internal/usageui/server_test.go b/internal/usageui/server_test.go index 45a9570..9ebb97d 100644 --- a/internal/usageui/server_test.go +++ b/internal/usageui/server_test.go @@ -111,193 +111,215 @@ func TestIndexPage(t *testing.T) { } } -func TestAPISummary(t *testing.T) { - srv := newTestServer(t) - w := doJSON(t, srv, "/api/summary") - if w.Code != http.StatusOK { - t.Fatalf("status %d", w.Code) - } - var ds usage.DailySummary - if err := json.Unmarshal(w.Body.Bytes(), &ds); err != nil { - t.Fatal(err) - } - if ds.RequestCount < 1 { - t.Fatalf("expected at least 1 request, got %d", ds.RequestCount) - } +// apiTest describes a single table-driven API endpoint test. On success +// (expected status code match), the response body is unmarshaled into a fresh +// T and check is invoked with the decoded value. When check is nil only the +// status code is asserted. +type apiTest[T any] struct { + name string + path string + expectedStatus int + check func(t *testing.T, got T) } -func TestAPIByDay(t *testing.T) { +// runAPITest executes one apiTest entry against a fresh test server. +func runAPITest[T any](t *testing.T, tc apiTest[T]) { + t.Helper() srv := newTestServer(t) - w := doJSON(t, srv, "/api/by-day") - if w.Code != http.StatusOK { - t.Fatalf("status %d", w.Code) - } - var days []usage.DailySummary - if err := json.Unmarshal(w.Body.Bytes(), &days); err != nil { - t.Fatal(err) + w := doJSON(t, srv, tc.path) + if w.Code != tc.expectedStatus { + t.Fatalf("status %d, want %d", w.Code, tc.expectedStatus) } - if len(days) < 1 { - t.Fatalf("expected at least 1 day, got %d", len(days)) - } -} - -func TestAPIByModel(t *testing.T) { - srv := newTestServer(t) - w := doJSON(t, srv, "/api/by-model") - if w.Code != http.StatusOK { - t.Fatalf("status %d", w.Code) + if tc.check == nil { + return } - var models []usage.ModelBreakdown - if err := json.Unmarshal(w.Body.Bytes(), &models); err != nil { + var got T + if err := json.Unmarshal(w.Body.Bytes(), &got); err != nil { t.Fatal(err) } - if len(models) < 1 { - t.Fatalf("expected at least 1 model, got %d", len(models)) - } + tc.check(t, got) } -func TestAPIByProvider(t *testing.T) { - srv := newTestServer(t) - w := doJSON(t, srv, "/api/by-provider") - if w.Code != http.StatusOK { - t.Fatalf("status %d", w.Code) - } - var provs []usage.ProviderBreakdown - if err := json.Unmarshal(w.Body.Bytes(), &provs); err != nil { - t.Fatal(err) - } - if len(provs) < 1 { - t.Fatalf("expected at least 1 provider, got %d", len(provs)) - } -} +func TestAPIEndpoints(t *testing.T) { + // --- /api/summary --- + t.Run("TestAPISummary", func(t *testing.T) { + runAPITest(t, apiTest[usage.DailySummary]{ + name: "TestAPISummary", + path: "/api/summary", + expectedStatus: http.StatusOK, + check: func(t *testing.T, ds usage.DailySummary) { + if ds.RequestCount < 1 { + t.Fatalf("expected at least 1 request, got %d", ds.RequestCount) + } + }, + }) + }) -func TestAPISessions(t *testing.T) { - srv := newTestServer(t) - w := doJSON(t, srv, "/api/sessions") - if w.Code != http.StatusOK { - t.Fatalf("status %d", w.Code) - } - var sessions []usage.SessionInfo - if err := json.Unmarshal(w.Body.Bytes(), &sessions); err != nil { - t.Fatal(err) - } - if len(sessions) < 1 { - t.Fatalf("expected at least 1 session, got %d", len(sessions)) - } -} + // --- /api/by-day --- + t.Run("TestAPIByDay", func(t *testing.T) { + runAPITest(t, apiTest[[]usage.DailySummary]{ + name: "TestAPIByDay", + path: "/api/by-day", + expectedStatus: http.StatusOK, + check: func(t *testing.T, days []usage.DailySummary) { + if len(days) < 1 { + t.Fatalf("expected at least 1 day, got %d", len(days)) + } + }, + }) + }) -func TestAPIHealth(t *testing.T) { - srv := newTestServer(t) - w := doJSON(t, srv, "/api/health") - if w.Code != http.StatusOK { - t.Fatalf("status %d", w.Code) - } - var h HealthInfo - if err := json.Unmarshal(w.Body.Bytes(), &h); err != nil { - t.Fatal(err) - } - if h.Version != "0.0.0-test" { - t.Fatalf("version: %q", h.Version) - } - if h.TunnelMode != "quick" { - t.Fatalf("tunnel_mode: %q", h.TunnelMode) - } - if !h.HasMoonshotKey { - t.Fatal("expected has_moonshot_key") - } - if !h.HasThauraKey { - t.Fatal("expected has_thaura_key") - } - if !h.HasZaiKey { - t.Fatal("expected has_zai_key") - } -} + // --- /api/by-model --- + t.Run("TestAPIByModel", func(t *testing.T) { + runAPITest(t, apiTest[[]usage.ModelBreakdown]{ + name: "TestAPIByModel", + path: "/api/by-model", + expectedStatus: http.StatusOK, + check: func(t *testing.T, models []usage.ModelBreakdown) { + if len(models) < 1 { + t.Fatalf("expected at least 1 model, got %d", len(models)) + } + }, + }) + }) -func TestAPISessionDetail(t *testing.T) { - srv := newTestServer(t) - w := doJSON(t, srv, "/api/sessions?session_id=sess-test") - if w.Code != http.StatusOK { - t.Fatalf("status %d", w.Code) - } - var ds usage.DailySummary - if err := json.Unmarshal(w.Body.Bytes(), &ds); err != nil { - t.Fatal(err) - } - if ds.RequestCount < 1 { - t.Fatalf("expected at least 1 request, got %d", ds.RequestCount) - } - if len(ds.ByModel) < 1 { - t.Fatalf("expected by_model breakdown, got %d", len(ds.ByModel)) - } -} + // --- /api/by-provider --- + t.Run("TestAPIByProvider", func(t *testing.T) { + runAPITest(t, apiTest[[]usage.ProviderBreakdown]{ + name: "TestAPIByProvider", + path: "/api/by-provider", + expectedStatus: http.StatusOK, + check: func(t *testing.T, provs []usage.ProviderBreakdown) { + if len(provs) < 1 { + t.Fatalf("expected at least 1 provider, got %d", len(provs)) + } + }, + }) + }) -func TestAPIByDaySince(t *testing.T) { - srv := newTestServer(t) - w := doJSON(t, srv, "/api/by-day?since=2025-01-01T00:00:00Z") - if w.Code != http.StatusOK { - t.Fatalf("status %d", w.Code) - } - var days []usage.DailySummary - if err := json.Unmarshal(w.Body.Bytes(), &days); err != nil { - t.Fatal(err) - } - // All seeded events are after 2025, so should produce at least 1 day. - if len(days) < 1 { - t.Fatalf("expected at least 1 day with since filter, got %d", len(days)) - } -} + // --- /api/sessions --- + t.Run("TestAPISessions", func(t *testing.T) { + runAPITest(t, apiTest[[]usage.SessionInfo]{ + name: "TestAPISessions", + path: "/api/sessions", + expectedStatus: http.StatusOK, + check: func(t *testing.T, sessions []usage.SessionInfo) { + if len(sessions) < 1 { + t.Fatalf("expected at least 1 session, got %d", len(sessions)) + } + }, + }) + }) -func TestAPIByModelSince(t *testing.T) { - srv := newTestServer(t) - w := doJSON(t, srv, "/api/by-model?since=2025-01-01T00:00:00Z") - if w.Code != http.StatusOK { - t.Fatalf("status %d", w.Code) - } - var models []usage.ModelBreakdown - if err := json.Unmarshal(w.Body.Bytes(), &models); err != nil { - t.Fatal(err) - } - if len(models) < 1 { - t.Fatalf("expected at least 1 model, got %d", len(models)) - } -} + // --- /api/health --- + t.Run("TestAPIHealth", func(t *testing.T) { + runAPITest(t, apiTest[HealthInfo]{ + name: "TestAPIHealth", + path: "/api/health", + expectedStatus: http.StatusOK, + check: func(t *testing.T, h HealthInfo) { + if h.Version != "0.0.0-test" { + t.Fatalf("version: %q", h.Version) + } + if h.TunnelMode != "quick" { + t.Fatalf("tunnel_mode: %q", h.TunnelMode) + } + if !h.HasMoonshotKey { + t.Fatal("expected has_moonshot_key") + } + if !h.HasThauraKey { + t.Fatal("expected has_thaura_key") + } + if !h.HasZaiKey { + t.Fatal("expected has_zai_key") + } + }, + }) + }) -func TestAPIByProviderSince(t *testing.T) { - srv := newTestServer(t) - w := doJSON(t, srv, "/api/by-provider?since=2025-01-01T00:00:00Z") - if w.Code != http.StatusOK { - t.Fatalf("status %d", w.Code) - } - var provs []usage.ProviderBreakdown - if err := json.Unmarshal(w.Body.Bytes(), &provs); err != nil { - t.Fatal(err) - } - if len(provs) < 1 { - t.Fatalf("expected at least 1 provider, got %d", len(provs)) - } -} + // --- /api/sessions?session_id=sess-test --- + t.Run("TestAPISessionDetail", func(t *testing.T) { + runAPITest(t, apiTest[usage.DailySummary]{ + name: "TestAPISessionDetail", + path: "/api/sessions?session_id=sess-test", + expectedStatus: http.StatusOK, + check: func(t *testing.T, ds usage.DailySummary) { + if ds.RequestCount < 1 { + t.Fatalf("expected at least 1 request, got %d", ds.RequestCount) + } + if len(ds.ByModel) < 1 { + t.Fatalf("expected by_model breakdown, got %d", len(ds.ByModel)) + } + }, + }) + }) -func TestAPISessionsSince(t *testing.T) { - srv := newTestServer(t) - w := doJSON(t, srv, "/api/sessions?since=2025-01-01T00:00:00Z") - if w.Code != http.StatusOK { - t.Fatalf("status %d", w.Code) - } - var sessions []usage.SessionInfo - if err := json.Unmarshal(w.Body.Bytes(), &sessions); err != nil { - t.Fatal(err) - } - if len(sessions) < 1 { - t.Fatalf("expected at least 1 session, got %d", len(sessions)) - } -} + // --- /api/by-day?since=2025-01-01T00:00:00Z --- + t.Run("TestAPIByDaySince", func(t *testing.T) { + runAPITest(t, apiTest[[]usage.DailySummary]{ + name: "TestAPIByDaySince", + path: "/api/by-day?since=2025-01-01T00:00:00Z", + expectedStatus: http.StatusOK, + check: func(t *testing.T, days []usage.DailySummary) { + // All seeded events are after 2025, so should produce at least 1 day. + if len(days) < 1 { + t.Fatalf("expected at least 1 day with since filter, got %d", len(days)) + } + }, + }) + }) -func TestAPIBadSince(t *testing.T) { - srv := newTestServer(t) - w := doJSON(t, srv, "/api/by-day?since=not-a-date") - if w.Code != http.StatusBadRequest { - t.Fatalf("expected 400 for bad since, got %d", w.Code) - } + // --- /api/by-model?since=2025-01-01T00:00:00Z --- + t.Run("TestAPIByModelSince", func(t *testing.T) { + runAPITest(t, apiTest[[]usage.ModelBreakdown]{ + name: "TestAPIByModelSince", + path: "/api/by-model?since=2025-01-01T00:00:00Z", + expectedStatus: http.StatusOK, + check: func(t *testing.T, models []usage.ModelBreakdown) { + if len(models) < 1 { + t.Fatalf("expected at least 1 model, got %d", len(models)) + } + }, + }) + }) + + // --- /api/by-provider?since=2025-01-01T00:00:00Z --- + t.Run("TestAPIByProviderSince", func(t *testing.T) { + runAPITest(t, apiTest[[]usage.ProviderBreakdown]{ + name: "TestAPIByProviderSince", + path: "/api/by-provider?since=2025-01-01T00:00:00Z", + expectedStatus: http.StatusOK, + check: func(t *testing.T, provs []usage.ProviderBreakdown) { + if len(provs) < 1 { + t.Fatalf("expected at least 1 provider, got %d", len(provs)) + } + }, + }) + }) + + // --- /api/sessions?since=2025-01-01T00:00:00Z --- + t.Run("TestAPISessionsSince", func(t *testing.T) { + runAPITest(t, apiTest[[]usage.SessionInfo]{ + name: "TestAPISessionsSince", + path: "/api/sessions?since=2025-01-01T00:00:00Z", + expectedStatus: http.StatusOK, + check: func(t *testing.T, sessions []usage.SessionInfo) { + if len(sessions) < 1 { + t.Fatalf("expected at least 1 session, got %d", len(sessions)) + } + }, + }) + }) + + // --- /api/by-day?since=not-a-date (bad input → 400, no body check) --- + t.Run("TestAPIBadSince", func(t *testing.T) { + runAPITest(t, apiTest[json.RawMessage]{ + name: "TestAPIBadSince", + path: "/api/by-day?since=not-a-date", + expectedStatus: http.StatusBadRequest, + check: nil, + }) + }) } func TestAPIByDayModelEmptyPadsHourBuckets(t *testing.T) {