diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..4289c5a --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,68 @@ +name: CI + +on: + push: + branches: [main, develop] + pull_request: + +permissions: + contents: read + +jobs: + lint-test: + name: lint & test (Go ${{ matrix.go-version }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + go-version: ["1.26", "1.27"] + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-go@v5 + with: + go-version: ${{ matrix.go-version }} + + - name: Install golangci-lint + run: go install github.com/golangci/golangci-lint/v2/cmd/golangci-lint@latest + + - name: make lint + run: make lint + + - name: make test + run: make test + + # make test doesn't pass -race (CLAUDE.md: offline, no network, no + # keys is the only hard requirement on the default target); CI runs + # the race detector as a separate, still fully offline step. + - name: go test -race + run: go test ./... -race + + build: + name: build (${{ matrix.goos }}/${{ matrix.goarch }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + include: + - goos: linux + goarch: amd64 + - goos: linux + goarch: arm64 + - goos: darwin + goarch: amd64 + - goos: darwin + goarch: arm64 + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-go@v5 + with: + go-version: "1.27" + + - name: go build + env: + GOOS: ${{ matrix.goos }} + GOARCH: ${{ matrix.goarch }} + CGO_ENABLED: "0" + run: go build -o bin/pico ./cmd/pico diff --git a/README.md b/README.md new file mode 100644 index 0000000..232168b --- /dev/null +++ b/README.md @@ -0,0 +1,211 @@ +# pico code + +A small CLI AI agent written in Go. One agent loop, two LLM backends — +[Anthropic](https://docs.anthropic.com/) (cloud) and [Ollama](https://ollama.com/) +(local) — and a handful of sandboxed tools (read a file, list a directory, run +an allowlisted command, write a file behind a flag). + +This project exists to work through the real engineering problems of +building an agent — tool-call normalization across providers, streaming, +context management, cancellation, safety — not to ship a product. If you're +reading the source alongside this README, `.claude/CLAUDE.md` is the +project's own working spec and is worth a read. + +## Quickstart + +Requires Go 1.26+. + +```bash +git clone https://github.com/reno/pico-code.git +cd pico-code +make build # -> bin/pico +``` + +Pick a backend and talk to it: + +```bash +# Local, no account needed — see "Ollama" below for the one-time setup. +./bin/pico chat --provider=ollama --model=qwen3:8b + +# Cloud — needs ANTHROPIC_API_KEY set first. +./bin/pico chat --provider=anthropic --model=claude-sonnet-4-5 +``` + +Either command drops you into a `> ` prompt. Ctrl+D (EOF) exits cleanly; +Ctrl+C interrupts the process the way it would any CLI tool — one to stop, +a second if the first doesn't. `--tui` behaves differently on purpose: see +[While chatting](#while-chatting). + +## The two backends + +### Anthropic + +```bash +export ANTHROPIC_API_KEY=sk-ant-... +./bin/pico chat --provider=anthropic --model= +``` + +Nothing else to install. Retries on 429/5xx happen inside the SDK; a bad key +or an unreachable network surfaces as a clear error rather than a hang. + +### Ollama + +```bash +# once: install Ollama (https://ollama.com/download), make sure it's +# running (the desktop app does this for you; on Linux, `ollama serve`), +# and pull a model +ollama pull qwen3:8b + +./bin/pico chat --provider=ollama --model=qwen3:8b +``` + +`OLLAMA_HOST` overrides the default `http://localhost:11434`. Ollama +silently truncates context if its `num_ctx` option is left unset — pico +code always sets it explicitly (`--num-ctx`, default 4096) and logs the +effective value at startup. + +Not every local model supports native tool calling. pico code probes +`/api/show` once per process and fails fast with a message pointing at +`--tools=prompted` if the model doesn't advertise tool support. +`--tools=prompted` injects tool schemas into the system prompt and parses a +fenced JSON block back out of the reply instead of relying on the model's +native tool-calling — it can't stream, so it also forces `--stream=false` +for that turn (a log line says so), and it isn't available at all under +`--tui`, which always streams. + +## Flags + +| Flag | Default | Meaning | +| ----------------- | ------------ | -------------------------------------------------------------- | +| `--provider` | `anthropic` | `anthropic` or `ollama` | +| `--model` | (none) | model ID for the selected provider | +| `--max-turns` | `25` | agent loop turns before it stops on its own | +| `--token-budget` | `100000` | cumulative tokens before it stops on its own | +| `--workspace` | `.` | root directory filesystem tools are confined to | +| `--yes` | `false` | skip interactive approval prompts | +| `--tools` | `native` | `native` or `prompted`; prompted mode ignores `--tui` and forces non-streaming | +| `--stream` | `true` | stream a reply as it arrives, plain mode only (the TUI always streams) | +| `--tui` | `false` | bubbletea TUI instead of the plain read-eval-print loop | +| `--log-level` | `info` | `debug`, `info`, `warn`, or `error` | +| `--num-ctx` | `4096` | Ollama's context window (`num_ctx`); ignored by Anthropic | +| `--allow-write` | `false` | register the `write_file` tool | +| `--session` | (none) | name a session to resume or start; saved after every turn | +| `--allow-commands`| (none) | comma-separated binary allowlist; registers `run_command` only if non-empty | +| `--context-window`| `200000` | context window compaction measures usage against; ignored by Ollama, which uses `--num-ctx` | + +`PICO_CODE_PROVIDER` is an environment fallback for `--provider`, checked +only when the flag isn't explicitly set. `ANTHROPIC_API_KEY` and +`OLLAMA_HOST` are read from the environment only — there's no flag for +either, so a credential never ends up in shell history. + +## While chatting + +Both the plain REPL and the TUI understand the same commands: + +- `/usage` — cumulative and last-turn token counts +- `/new` — start a fresh, unsaved conversation +- `/save [name]` — save the current conversation as a session +- `/load ` — replace the current conversation with a saved session + +`--session ` does the same save/resume automatically, after every +turn, in both — kill the process mid-conversation and `--session ` +again picks up where it left off. + +`--tui` swaps the REPL for a full-screen interface (scrollback, a spinner +while the model thinks, tool calls shown with a running/✓/✗ status, +markdown rendering for finished replies, an approval modal for tools that +need sign-off). A command's output appears in the scrollback prefixed with +the command line itself, since the TUI otherwise never echoes what you +typed. Its Ctrl+C/Ctrl+D also differ from the plain REPL's on purpose: +Ctrl+C there cancels only the in-flight turn and returns to the prompt, +and Ctrl+D is the one that exits — a per-turn cancellation the plain REPL +doesn't have, since its context is shared process-wide rather than derived +fresh per turn. + +## Safety + +- Filesystem tools are confined to `--workspace`, resolved through + `EvalSymlinks` so a symlink can't point a read or write outside it. Reads + of `.env`, `.git/config`, `*.pem`, and `id_*` are denied outright. +- `run_command` (registered only when `--allow-commands` names at least one + binary) executes only that allowlisted binary, directly (`exec.Command`, + never a shell — no `|`, `&&`, `$()`, or similar can do anything), with a + timeout and truncated output. +- `write_file` and `run_command` both need approval unless `--yes` is + passed; the plain REPL prompts on the terminal, the TUI shows a modal. +- Tool output is truncated (head + tail, with an elision marker) before it + ever reaches the model's context, so one large file can't blow the + context window. + +## How the agent loop works + +```mermaid +flowchart TD + U[user turn] --> C[call the provider] + C --> T{reply has tool calls} + T -- no --> Done[return the reply] + T -- yes --> R[run every tool call concurrently] + R --> A[append results in call order] + A --> G{guard tripped} + G -- yes --> Stop[stop with an explanation appended] + G -- no --> C +``` + +A guard trips on max turns, the token budget, a wall-clock timeout, or the +same tool call repeating three times in a row. + +Every tool call the model issues is answered by exactly one tool result, in +the same message, in the same order — even if the tool panics, times out, +or the run is cancelled mid-call. That invariant is asserted in the history +package itself (`history.Validate`), not just hoped for, and it's what makes +compaction and session persistence safe to build on top of: a saved or +compacted history is always in a state the model can resume from cleanly. + +## Development + +```bash +make build # bin/pico +make test # go test ./... — offline, no network, no API keys +make lint # gofmt -l . && go vet ./... && golangci-lint run +make run # go run ./cmd/pico chat --provider=ollama +``` + +`make test` never touches the network. Provider adapters replay recorded +HTTP exchanges from `testdata/golden/` through an `httptest.Server`; running +a test with `RECORD=1` (and real credentials, for the adapters that need +them) re-records against the live API instead, scrubbing any credential +before the fixture is written back to disk. + +## What I learned / what breaks with small models + +Small local models are the reason most of the guard rails in this codebase +exist: + +- **They invent tool names.** The registry rejects an unknown name with an + error result — handed back to the model as data, not a crash — and the + model usually corrects itself on the next turn. +- **They repeat themselves.** A model stuck in a loop calling the same tool + with the same arguments gets cut off after three identical calls in a + row, with an explanation appended to the transcript instead of spinning + forever. +- **They emit arguments that don't match the schema**, or send them as a + JSON-encoded *string* instead of an object (Ollama does this often enough + that the adapter normalizes both shapes before decoding). Validation + failures go back to the model as an error result with the actual message, + which is usually enough for a capable-but-small model to fix on retry. +- **Some don't call tools at all — they narrate the call in prose.** That's + what `--tools=prompted` is for: schemas injected into the system prompt, + a fenced JSON block parsed back out of the reply instead of relying on + native tool-calling. +- **Streaming exposed a real gap in `encoding/json`.** A JSON string chunk + containing a truncated multi-byte UTF-8 sequence gets silently replaced + with U+FFFD by Go's own decoder — which mattered for testing "a tool + argument split mid-character" honestly: splitting raw UTF-8 bytes across + SSE chunks is untestable for this exact reason, so the fixture instead + splits a `\uXXXX` escape sequence, which is always well-formed at the + outer JSON-encoding level regardless of where it's cut. + +More generally: almost everything CLAUDE.md calls a "provider gotcha" was +discovered by actually running a small model against this loop, not +anticipated up front. The cloud model mostly behaves; the local one is what +taught this codebase to be defensive. diff --git a/cmd/pico/chat.go b/cmd/pico/chat.go index 490bbca..b9218af 100644 --- a/cmd/pico/chat.go +++ b/cmd/pico/chat.go @@ -1,11 +1,33 @@ package main import ( + "bufio" + "bytes" + "context" + "errors" "fmt" + "io" + "log/slog" + "os" + "strings" + "time" + tea "github.com/charmbracelet/bubbletea" "github.com/spf13/cobra" + "github.com/reno/pico-code/internal/agent" "github.com/reno/pico-code/internal/config" + "github.com/reno/pico-code/internal/history" + "github.com/reno/pico-code/internal/llm" + "github.com/reno/pico-code/internal/llm/prompted" + "github.com/reno/pico-code/internal/tools" + "github.com/reno/pico-code/internal/ui" +) + +const ( + systemPrompt = "You are pico code, a terminal coding agent. Be direct and concise." + defaultMaxTokens = 4096 + defaultToolTimeout = 2 * time.Minute ) // newChatCmd builds the chat subcommand. Flag values win over the @@ -14,17 +36,21 @@ import ( // config.Load since they are credentials, not flags. func newChatCmd(getenv func(string) string) *cobra.Command { var flags struct { - provider string - model string - maxTurns int - tokenBudget int - workspace string - yes bool - tools string - stream bool - tui bool - logLevel string - numCtx int + provider string + model string + maxTurns int + tokenBudget int + workspace string + yes bool + tools string + stream bool + tui bool + logLevel string + numCtx int + allowWrite bool + session string + allowCommands []string + contextWindow int } cmd := &cobra.Command{ @@ -39,17 +65,21 @@ func newChatCmd(getenv func(string) string) *cobra.Command { } cfg, err := config.Load(config.Flags{ - Provider: provider, - Model: flags.model, - MaxTurns: flags.maxTurns, - TokenBudget: flags.tokenBudget, - Workspace: flags.workspace, - Yes: flags.yes, - Tools: flags.tools, - Stream: flags.stream, - TUI: flags.tui, - LogLevel: flags.logLevel, - NumCtx: flags.numCtx, + Provider: provider, + Model: flags.model, + MaxTurns: flags.maxTurns, + TokenBudget: flags.tokenBudget, + Workspace: flags.workspace, + Yes: flags.yes, + Tools: flags.tools, + Stream: flags.stream, + TUI: flags.tui, + LogLevel: flags.logLevel, + NumCtx: flags.numCtx, + AllowWrite: flags.allowWrite, + Session: flags.session, + AllowCommands: flags.allowCommands, + ContextWindow: flags.contextWindow, }, getenv) if err != nil { return fmt.Errorf("resolving config: %w", err) @@ -71,12 +101,298 @@ func newChatCmd(getenv func(string) string) *cobra.Command { f.BoolVar(&flags.tui, "tui", false, "use the bubbletea TUI instead of plain output") f.StringVar(&flags.logLevel, "log-level", "info", "log level (debug|info|warn|error)") f.IntVar(&flags.numCtx, "num-ctx", 4096, "context window size passed to Ollama's num_ctx (ignored by other providers)") + f.BoolVar(&flags.allowWrite, "allow-write", false, "register the write_file tool (off by default)") + f.StringVar(&flags.session, "session", "", "name a session to resume or start; saved after every turn") + f.StringSliceVar(&flags.allowCommands, "allow-commands", nil, "comma-separated binary allowlist for run_command; registers the tool only if non-empty") + f.IntVar(&flags.contextWindow, "context-window", defaultContextWindow, "context window size compaction measures usage against (ignored by Ollama, which uses --num-ctx instead)") return cmd } -// runChat is a placeholder until the agent loop (phase 4) exists. -func runChat(cmd *cobra.Command, cfg *config.Config) error { - _, err := fmt.Fprintf(cmd.OutOrStdout(), "pico chat: provider=%s model=%s (agent loop not yet implemented)\n", cfg.Provider, cfg.Model) +// runChat is a package var so tests can substitute a stub instead of +// driving a real provider/agent turn end to end. +var runChat = func(cmd *cobra.Command, cfg *config.Config) error { + if err := setupLogging(cmd.ErrOrStderr(), cfg.LogLevel); err != nil { + return err + } + + provider, err := llm.New(cfg) + if err != nil { + return fmt.Errorf("resolving provider: %w", err) + } + provider, err = resolveProvider(cfg, provider) + if err != nil { + return err + } + + registry, err := buildRegistry(cfg) + if err != nil { + return err + } + + sessDir, err := defaultSessionsDir() + if err != nil { + return err + } + sess, err := newSession(sessDir, cfg.Session) + if err != nil { + return err + } + h, err := sess.loadOrCreateHistory() + if err != nil { + return err + } + + guards := agent.Guards{MaxTurns: cfg.MaxTurns, TokenBudget: cfg.TokenBudget} + + if cfg.TUI { + return runTUIChat(cmd, cfg, provider, registry, h, guards, sess) + } + return runPlainChat(cmd, cfg, provider, registry, h, guards, sess) +} + +// defaultContextWindow is --context-window's default: a conservative +// stand-in for a real model's actual context size, used as-is unless the +// flag overrides it. Ollama ignores it entirely and uses NumCtx instead — +// that's already an explicit, required value (CLAUDE.md), so compaction +// reuses it rather than tracking a second number that could drift from it. +const ( + defaultContextWindow = 200_000 + compactionTriggerFraction = 0.75 + compactionKeepTurns = 6 +) + +func compactionPolicy(cfg *config.Config) agent.CompactionPolicy { + window := cfg.ContextWindow + if cfg.Provider == config.ProviderOllama { + window = cfg.NumCtx + } + return agent.CompactionPolicy{ContextWindow: window, TriggerFraction: compactionTriggerFraction, KeepTurns: compactionKeepTurns} +} + +// resolveProvider wraps provider with prompted-tools support when cfg.Tools +// asks for it, rejecting that combined with --tui: the TUI always drives +// the agent through RunStream, and prompted.Provider.Stream always errors +// (it needs the full reply before it can look for a fenced tool call), so +// the combination would fail confusingly on the very first turn instead of +// with a clear message at startup. +func resolveProvider(cfg *config.Config, provider llm.Provider) (llm.Provider, error) { + if cfg.Tools != config.ToolsPrompted { + return provider, nil + } + if cfg.TUI { + return nil, errors.New("--tui does not support --tools=prompted: the TUI requires streaming, and prompted mode can't stream") + } + return prompted.Wrap(provider), nil +} + +// buildRegistry wires the built-in tools this session offers. run_command +// is only registered when --allow-commands names at least one binary — an +// empty allowlist would make the tool present but unconditionally useless, +// worse than leaving it out. +func buildRegistry(cfg *config.Config) (*tools.Registry, error) { + sandbox, err := tools.NewSandbox(cfg.Workspace, nil) + if err != nil { + return nil, fmt.Errorf("resolving workspace sandbox: %w", err) + } + + registry := tools.NewRegistry() + readTool, err := tools.NewReadFileTool(sandbox) + if err != nil { + return nil, err + } + if err := registry.Register(readTool); err != nil { + return nil, err + } + listTool, err := tools.NewListDirTool(sandbox) + if err != nil { + return nil, err + } + if err := registry.Register(listTool); err != nil { + return nil, err + } + if cfg.AllowWrite { + writeTool, err := tools.NewWriteFileTool(sandbox) + if err != nil { + return nil, err + } + if err := registry.Register(writeTool); err != nil { + return nil, err + } + } + if len(cfg.AllowCommands) > 0 { + runTool, err := tools.NewRunCommandTool(cfg.AllowCommands, defaultToolTimeout) + if err != nil { + return nil, err + } + if err := registry.Register(runTool); err != nil { + return nil, err + } + } + return registry, nil +} + +// runPlainChat drives the agent through ui.PlainRenderer. Piped/non-TTY +// stdin (the common case for scripting and for 7.1's AC) is read in full as +// one message and answered once; a TTY gets a minimal read-eval-print loop. +func runPlainChat(cmd *cobra.Command, cfg *config.Config, provider llm.Provider, registry *tools.Registry, h *history.History, guards agent.Guards, sess *session) error { + approver := agent.AutoApprove + if !cfg.Yes { + approver = agent.ConsoleApprover{In: cmd.InOrStdin(), Out: cmd.OutOrStdout()} + } + ag := agent.New(provider, registry, h, systemPrompt, defaultMaxTokens, guards, defaultToolTimeout, approver) + ag.SetCompactionPolicy(compactionPolicy(cfg)) + out := cmd.OutOrStdout() + ctx := cmd.Context() + in := cmd.InOrStdin() + + runTurn := newTurnRunner(cfg, ag, out) + + if !isInteractive(in) { + data, err := io.ReadAll(in) + if err != nil { + return fmt.Errorf("reading stdin: %w", err) + } + input := strings.TrimSpace(string(data)) + if input == "" { + return nil + } + if err := runTurn(ctx, input); err != nil { + return err + } + return sess.saveIfActive(h) + } + + return runREPL(ctx, in, out, ag, h, sess, runTurn) +} + +// newTurnRunner picks Run or RunStream per cfg.Stream, printing Run's +// final text itself since PlainRenderer only ever sees a stream. +// tools=prompted forces non-streaming regardless of --stream: the +// prompted.Provider decorator needs a full reply before it can look for a +// fenced tool call, so its Stream always errors. +func newTurnRunner(cfg *config.Config, ag *agent.Agent, out io.Writer) func(context.Context, string) error { + stream := cfg.Stream + if cfg.Tools == config.ToolsPrompted && cfg.Stream { + slog.Info("chat: streaming disabled", "reason", "tools=prompted can't stream") + stream = false + } + if !stream { + return func(ctx context.Context, input string) error { + text, err := ag.Run(ctx, input) + if err != nil { + return err + } + _, err = fmt.Fprintln(out, text) + return err + } + } + renderer := ui.PlainRenderer{Out: out} + return func(ctx context.Context, input string) error { + _, err := ag.RunStream(ctx, input, renderer) + return err + } +} + +// runREPL reads one line at a time from in, dispatching a leading slash +// command to handleCommand and everything else to runTurn (saving to +// sess's file afterward, if one is active), until in hits EOF or ctx is +// cancelled. Split out from runPlainChat so a test can drive it directly +// without needing a real terminal for isInteractive. +func runREPL(ctx context.Context, in io.Reader, out io.Writer, ag *agent.Agent, h *history.History, sess *session, runTurn func(context.Context, string) error) error { + prompt := func() { _, _ = fmt.Fprint(out, "> ") } + + scanner := bufio.NewScanner(in) + prompt() + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + prompt() + continue + } + if name, arg, ok := slashCommand(line); ok { + if _, err := handleCommand(out, ag, h, sess, name, arg); err != nil { + return err + } + prompt() + continue + } + if err := runTurn(ctx, line); err != nil { + return err + } + if ctx.Err() != nil { + return ctx.Err() + } + if err := sess.saveIfActive(h); err != nil { + return err + } + prompt() + } + return scanner.Err() +} + +// isInteractive reports whether in is a terminal. A non-*os.File reader +// (e.g. a test's bytes.Buffer) is treated as non-interactive. +func isInteractive(in io.Reader) bool { + f, ok := in.(*os.File) + if !ok { + return false + } + fi, err := f.Stat() + if err != nil { + return false + } + return fi.Mode()&os.ModeCharDevice != 0 +} + +// runTUIChat drives the agent through the bubbletea Model: a background +// goroutine reads submitted input and calls RunStream, sending its +// progress and result back into the program as messages, while +// program.Run() owns the terminal until Ctrl+D quits it. +func runTUIChat(cmd *cobra.Command, cfg *config.Config, provider llm.Provider, registry *tools.Registry, h *history.History, guards agent.Guards, sess *session) error { + ctx := cmd.Context() + submit := make(chan string, 1) + program := tea.NewProgram(ui.NewModel(submit), tea.WithAltScreen(), tea.WithContext(ctx)) + + renderer := &ui.TUIRenderer{Program: program} + approver := &ui.TUIApprover{Program: program} + ag := agent.New(provider, registry, h, systemPrompt, defaultMaxTokens, guards, defaultToolTimeout, approver) + ag.SetCompactionPolicy(compactionPolicy(cfg)) + + go func() { + for input := range submit { + if name, arg, ok := slashCommand(input); ok { + ui.CommandOutput(program, runTUICommand(ag, h, sess, input, name, arg)) + continue + } + + turnCtx, cancel := context.WithCancel(ctx) + ui.TurnStarted(program, cancel) + text, err := ag.RunStream(turnCtx, input, renderer) + cancel() + if err == nil { + err = sess.saveIfActive(h) + } + ui.TurnDone(program, text, err) + } + }() + + _, err := program.Run() + close(submit) return err } + +// runTUICommand runs a parsed slash command and formats its output for the +// TUI's transcript, prefixed with the command line itself (the plain +// REPL's prompt naturally shows what was typed; the TUI doesn't echo user +// input into its transcript at all, so a command needs to say what it is). +// Split out from runTUIChat's driver goroutine so it's testable without a +// running *tea.Program. +func runTUICommand(ag *agent.Agent, h *history.History, sess *session, input, name, arg string) string { + var buf bytes.Buffer + _, _ = fmt.Fprintf(&buf, "> %s\n", input) + // bytes.Buffer.Write never errors, so handleCommand's error return + // (only ever a write failure) is unreachable here. + _, _ = handleCommand(&buf, ag, h, sess, name, arg) + return buf.String() +} diff --git a/cmd/pico/chat_test.go b/cmd/pico/chat_test.go new file mode 100644 index 0000000..17c3a7b --- /dev/null +++ b/cmd/pico/chat_test.go @@ -0,0 +1,387 @@ +package main + +import ( + "bytes" + "context" + "errors" + "io" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/spf13/cobra" + + "github.com/reno/pico-code/internal/agent" + "github.com/reno/pico-code/internal/config" + "github.com/reno/pico-code/internal/history" + "github.com/reno/pico-code/internal/llm" + "github.com/reno/pico-code/internal/tools" +) + +// fakeChatProvider returns one canned reply, enough to drive +// runPlainChat's single-shot piped-input path without a real backend. +// disallowStream, when set, makes Stream fail instead of succeed — used to +// prove a --stream=false turn never calls it at all. +type fakeChatProvider struct { + reply string + disallowStream bool +} + +func (f *fakeChatProvider) Name() string { return "fake" } + +func (f *fakeChatProvider) Chat(context.Context, llm.Request) (*llm.Response, error) { + return &llm.Response{ + Message: llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.Text{Text: f.reply}}}, + StopReason: "end_turn", + }, nil +} + +func (f *fakeChatProvider) Stream(ctx context.Context, _ llm.Request) (<-chan llm.Event, error) { + if f.disallowStream { + return nil, errors.New("fakeChatProvider: Stream must not be called when streaming is disabled") + } + return llm.StreamEvents(ctx, func(send func(llm.Event) bool) { + if !send(llm.TextDelta{Text: f.reply}) { + return + } + send(llm.MessageDone{StopReason: "end_turn"}) + }), nil +} + +// TestPipedChatProducesCleanANSIFreeOutput is 7.1's AC: piping input to a +// non-interactive chat session produces clean, ANSI-free output. +func TestPipedChatProducesCleanANSIFreeOutput(t *testing.T) { + cmd := &cobra.Command{} + cmd.SetContext(context.Background()) + cmd.SetIn(strings.NewReader("hi\n")) + var out bytes.Buffer + cmd.SetOut(&out) + + cfg := &config.Config{Yes: true, Stream: true} + provider := &fakeChatProvider{reply: "hello there"} + h := history.New() + + err := runPlainChat(cmd, cfg, provider, tools.NewRegistry(), h, agent.Guards{}, noSession(t)) + if err != nil { + t.Fatalf("runPlainChat() error = %v", err) + } + + got := out.String() + if !strings.Contains(got, "hello there") { + t.Errorf("output = %q, want it to contain the reply", got) + } + if strings.ContainsRune(got, '\x1b') { + t.Errorf("output contains an ANSI escape byte: %q", got) + } + if err := h.Validate(); err != nil { + t.Errorf("history.Validate() error = %v", err) + } +} + +// TestUsageCommandReportsCumulativeAcrossTurns is 8.1's AC ("expose /usage +// in the chat"), driven through runREPL so it never needs a real terminal. +func TestUsageCommandReportsCumulativeAcrossTurns(t *testing.T) { + provider := &fakeChatProvider{reply: "ok"} + h := history.New() + ag := agent.New(provider, tools.NewRegistry(), h, "", 1024, agent.Guards{}, 0, agent.AutoApprove) + + var out bytes.Buffer + in := strings.NewReader("hello\n/usage\n") + runTurn := newTurnRunner(&config.Config{Stream: true}, ag, &out) + if err := runREPL(context.Background(), in, &out, ag, h, noSession(t), runTurn); err != nil { + t.Fatalf("runREPL() error = %v", err) + } + + got := out.String() + if !strings.Contains(got, "1 turn") { + t.Errorf("output = %q, want it to report 1 completed turn", got) + } + if strings.ContainsRune(got, '\x1b') { + t.Errorf("output contains an ANSI escape byte: %q", got) + } +} + +func TestSlashCommandParsesNameAndArg(t *testing.T) { + name, arg, ok := slashCommand("/save my-session") + if !ok || name != "save" || arg != "my-session" { + t.Errorf("slashCommand(%q) = (%q, %q, %v), want (\"save\", \"my-session\", true)", "/save my-session", name, arg, ok) + } + + if _, _, ok := slashCommand("not a command"); ok { + t.Error("slashCommand on ordinary text should return ok = false") + } +} + +// noSession returns an inactive session (no name) rooted at a temp +// directory, for tests that need a *session parameter but don't care about +// persistence. +func noSession(t *testing.T) *session { + t.Helper() + sess, err := newSession(t.TempDir(), "") + if err != nil { + t.Fatalf("newSession() error = %v", err) + } + return sess +} + +// TestSessionAutoSavesAfterEachTurnAndResumes is 8.3's AC in the plain +// REPL: killing the process mid-session and resuming has the model see the +// full prior context. "Killing the process" is simulated by driving runREPL +// once for the first half of the conversation, then constructing a brand +// new History via the session's on-disk file (as a fresh process would) +// and driving runREPL again. +func TestSessionAutoSavesAfterEachTurnAndResumes(t *testing.T) { + dir := t.TempDir() + + sess1, err := newSession(dir, "work") + if err != nil { + t.Fatalf("newSession() error = %v", err) + } + h1, err := sess1.loadOrCreateHistory() + if err != nil { + t.Fatalf("loadOrCreateHistory() error = %v", err) + } + provider := &fakeChatProvider{reply: "first reply"} + ag1 := agent.New(provider, tools.NewRegistry(), h1, "", 1024, agent.Guards{}, 0, agent.AutoApprove) + runTurn1 := newTurnRunner(&config.Config{Stream: true}, ag1, io.Discard) + if err := runREPL(context.Background(), strings.NewReader("first message\n"), io.Discard, ag1, h1, sess1, runTurn1); err != nil { + t.Fatalf("runREPL() (process 1) error = %v", err) + } + + // A brand new process: nothing but the session name and the file it + // wrote to disk. + sess2, err := newSession(dir, "work") + if err != nil { + t.Fatalf("newSession() error = %v", err) + } + h2, err := sess2.loadOrCreateHistory() + if err != nil { + t.Fatalf("loadOrCreateHistory() (process 2) error = %v", err) + } + + if err := h2.Validate(); err != nil { + t.Fatalf("Validate() on resumed history error = %v", err) + } + if got, want := len(h2.Snapshot()), len(h1.Snapshot()); got != want { + t.Fatalf("resumed history has %d messages, want %d (the full prior context)", got, want) + } +} + +func TestSessionCommandsNewSaveLoad(t *testing.T) { + dir := t.TempDir() + sess, err := newSession(dir, "") + if err != nil { + t.Fatalf("newSession() error = %v", err) + } + h := history.New() + provider := &fakeChatProvider{reply: "reply"} + ag := agent.New(provider, tools.NewRegistry(), h, "", 1024, agent.Guards{}, 0, agent.AutoApprove) + + var out bytes.Buffer + in := strings.NewReader("hello there\n/save first\n/new\n/load first\n") + runTurn := newTurnRunner(&config.Config{Stream: true}, ag, &out) + if err := runREPL(context.Background(), in, &out, ag, h, sess, runTurn); err != nil { + t.Fatalf("runREPL() error = %v", err) + } + + got := out.String() + if !strings.Contains(got, `saved session "first"`) { + t.Errorf("output missing save confirmation: %q", got) + } + if !strings.Contains(got, "started a new, unsaved session") { + t.Errorf("output missing /new confirmation: %q", got) + } + if !strings.Contains(got, `loaded session "first"`) { + t.Errorf("output missing load confirmation: %q", got) + } + // /load restored the turn /new had cleared. + if len(h.Snapshot()) == 0 { + t.Error("history is empty after /load, want the saved turn restored") + } +} + +// TestSessionLoadOfCorruptFileIsReadableErrorNotPanic is 8.3's other AC: a +// corrupt session file fails with a readable error rather than a panic. +func TestSessionLoadOfCorruptFileIsReadableErrorNotPanic(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "broken.json"), []byte("not json"), 0o600); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + + sess, err := newSession(dir, "broken") + if err != nil { + t.Fatalf("newSession() error = %v", err) + } + if _, err := sess.loadOrCreateHistory(); err == nil { + t.Fatal("loadOrCreateHistory() = nil error, want a readable error for a corrupt session file") + } +} + +// TestStreamFalseUsesNonStreamingPath proves --stream=false actually has +// an effect: the fake provider's Stream errors, so a passing turn here +// means newTurnRunner drove it through Chat/ag.Run instead. +func TestStreamFalseUsesNonStreamingPath(t *testing.T) { + cmd := &cobra.Command{} + cmd.SetContext(context.Background()) + cmd.SetIn(strings.NewReader("hi\n")) + var out bytes.Buffer + cmd.SetOut(&out) + + cfg := &config.Config{Yes: true, Stream: false} + provider := &fakeChatProvider{reply: "non-streamed reply", disallowStream: true} + h := history.New() + + if err := runPlainChat(cmd, cfg, provider, tools.NewRegistry(), h, agent.Guards{}, noSession(t)); err != nil { + t.Fatalf("runPlainChat() error = %v", err) + } + if got := out.String(); !strings.Contains(got, "non-streamed reply") { + t.Errorf("output = %q, want it to contain the reply", got) + } +} + +// TestStreamTrueWithPromptedToolsFallsBackToNonStreaming is the other half +// of newTurnRunner's --tools=prompted handling: even with --stream=true, +// prompted mode can't stream, so the turn must still go through Chat. +func TestStreamTrueWithPromptedToolsFallsBackToNonStreaming(t *testing.T) { + cmd := &cobra.Command{} + cmd.SetContext(context.Background()) + cmd.SetIn(strings.NewReader("hi\n")) + var out bytes.Buffer + cmd.SetOut(&out) + + cfg := &config.Config{Yes: true, Stream: true, Tools: config.ToolsPrompted} + provider := &fakeChatProvider{reply: "fallback reply", disallowStream: true} + h := history.New() + + if err := runPlainChat(cmd, cfg, provider, tools.NewRegistry(), h, agent.Guards{}, noSession(t)); err != nil { + t.Fatalf("runPlainChat() error = %v", err) + } + if got := out.String(); !strings.Contains(got, "fallback reply") { + t.Errorf("output = %q, want it to contain the reply", got) + } +} + +func TestResolveProviderWrapsForPromptedTools(t *testing.T) { + provider := &fakeChatProvider{reply: "unused"} + got, err := resolveProvider(&config.Config{Tools: config.ToolsPrompted}, provider) + if err != nil { + t.Fatalf("resolveProvider() error = %v", err) + } + if got == llm.Provider(provider) { + t.Error("resolveProvider() returned the provider unwrapped, want it wrapped for prompted-tools mode") + } +} + +func TestResolveProviderLeavesNativeToolsUnwrapped(t *testing.T) { + provider := &fakeChatProvider{reply: "unused"} + got, err := resolveProvider(&config.Config{Tools: config.ToolsNative}, provider) + if err != nil { + t.Fatalf("resolveProvider() error = %v", err) + } + if got != llm.Provider(provider) { + t.Error("resolveProvider() should return native-mode providers unwrapped") + } +} + +func TestResolveProviderRejectsTUIWithPromptedTools(t *testing.T) { + provider := &fakeChatProvider{reply: "unused"} + _, err := resolveProvider(&config.Config{Tools: config.ToolsPrompted, TUI: true}, provider) + if err == nil { + t.Fatal("resolveProvider() error = nil, want an error for --tui with --tools=prompted") + } + if !strings.Contains(err.Error(), "tui") || !strings.Contains(err.Error(), "prompted") { + t.Errorf("resolveProvider() error = %q, want it to name both --tui and --tools=prompted", err) + } +} + +func TestBuildRegistryRegistersRunCommandOnlyWhenAllowlisted(t *testing.T) { + dir := t.TempDir() + + withCommands, err := buildRegistry(&config.Config{Workspace: dir, AllowCommands: []string{"echo"}}) + if err != nil { + t.Fatalf("buildRegistry() error = %v", err) + } + if _, err := withCommands.Get("run_command"); err != nil { + t.Errorf("Get(\"run_command\") error = %v, want it registered when --allow-commands is non-empty", err) + } + + without, err := buildRegistry(&config.Config{Workspace: dir}) + if err != nil { + t.Fatalf("buildRegistry() error = %v", err) + } + if _, err := without.Get("run_command"); !errors.Is(err, tools.ErrToolNotFound) { + t.Errorf("Get(\"run_command\") error = %v, want ErrToolNotFound when --allow-commands is empty", err) + } +} + +// TestRunTUICommandUsage is the TUI counterpart of +// TestUsageCommandReportsCumulativeAcrossTurns: /usage must work the same +// way whether it's typed into the plain REPL or the TUI's textarea. +func TestRunTUICommandUsage(t *testing.T) { + provider := &fakeChatProvider{reply: "ok"} + h := history.New() + ag := agent.New(provider, tools.NewRegistry(), h, "", 1024, agent.Guards{}, 0, agent.AutoApprove) + runTurn := newTurnRunner(&config.Config{Stream: true}, ag, io.Discard) + if err := runTurn(context.Background(), "hello"); err != nil { + t.Fatalf("runTurn() error = %v", err) + } + + out := runTUICommand(ag, h, noSession(t), "/usage", "usage", "") + + if !strings.Contains(out, "> /usage") { + t.Errorf("output = %q, want it to echo the command line", out) + } + if !strings.Contains(out, "1 turn") { + t.Errorf("output = %q, want it to report 1 completed turn", out) + } +} + +// TestRunTUICommandNewSaveLoad is the TUI counterpart of +// TestSessionCommandsNewSaveLoad. +func TestRunTUICommandNewSaveLoad(t *testing.T) { + dir := t.TempDir() + sess, err := newSession(dir, "") + if err != nil { + t.Fatalf("newSession() error = %v", err) + } + h := history.New() + h.Append(llm.Message{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "hello there"}}}) + provider := &fakeChatProvider{reply: "reply"} + ag := agent.New(provider, tools.NewRegistry(), h, "", 1024, agent.Guards{}, 0, agent.AutoApprove) + + saveOut := runTUICommand(ag, h, sess, "/save first", "save", "first") + if !strings.Contains(saveOut, `saved session "first"`) { + t.Errorf("save output = %q, want a save confirmation", saveOut) + } + + newOut := runTUICommand(ag, h, sess, "/new", "new", "") + if !strings.Contains(newOut, "started a new, unsaved session") { + t.Errorf("new output = %q, want a /new confirmation", newOut) + } + if len(h.Snapshot()) != 0 { + t.Errorf("history has %d messages after /new, want 0", len(h.Snapshot())) + } + + loadOut := runTUICommand(ag, h, sess, "/load first", "load", "first") + if !strings.Contains(loadOut, `loaded session "first"`) { + t.Errorf("load output = %q, want a load confirmation", loadOut) + } + if len(h.Snapshot()) == 0 { + t.Error("history is empty after /load, want the saved turn restored") + } +} + +func TestCompactionPolicyUsesContextWindowForNonOllamaProviders(t *testing.T) { + got := compactionPolicy(&config.Config{Provider: config.ProviderAnthropic, ContextWindow: 123_456, NumCtx: 4096}) + if got.ContextWindow != 123_456 { + t.Errorf("ContextWindow = %d, want the configured --context-window value (4096 is Ollama's NumCtx and must not leak in here)", got.ContextWindow) + } +} + +func TestCompactionPolicyUsesNumCtxForOllama(t *testing.T) { + got := compactionPolicy(&config.Config{Provider: config.ProviderOllama, ContextWindow: 123_456, NumCtx: 8192}) + if got.ContextWindow != 8192 { + t.Errorf("ContextWindow = %d, want NumCtx (8192), not --context-window (which Ollama ignores)", got.ContextWindow) + } +} diff --git a/cmd/pico/commands.go b/cmd/pico/commands.go new file mode 100644 index 0000000..bbaa5ca --- /dev/null +++ b/cmd/pico/commands.go @@ -0,0 +1,89 @@ +package main + +import ( + "fmt" + "io" + "strings" + + "github.com/reno/pico-code/internal/agent" + "github.com/reno/pico-code/internal/history" +) + +// slashCommand parses a leading "/name arg" out of line. ok is false for +// ordinary chat input. +func slashCommand(line string) (name, arg string, ok bool) { + line = strings.TrimSpace(line) + if !strings.HasPrefix(line, "/") { + return "", "", false + } + fields := strings.SplitN(strings.TrimPrefix(line, "/"), " ", 2) + name = fields[0] + if len(fields) > 1 { + arg = strings.TrimSpace(fields[1]) + } + return name, arg, name != "" +} + +// handleCommand runs a parsed slash command, writing its output to out. It +// always returns handled == true, even for an unknown command (reported as +// an error line rather than sent to the model as chat input). +func handleCommand(out io.Writer, ag *agent.Agent, h *history.History, sess *session, name, arg string) (handled bool, err error) { + switch name { + case "usage": + printUsage(out, ag) + case "new": + h.Reset() + sess.name = "" + _, err = fmt.Fprintln(out, "started a new, unsaved session") + case "save": + err = runSaveCommand(out, h, sess, arg) + case "load": + err = runLoadCommand(out, h, sess, arg) + default: + _, err = fmt.Fprintf(out, "unknown command /%s\n", name) + } + return true, err +} + +func runSaveCommand(out io.Writer, h *history.History, sess *session, arg string) error { + name := arg + if name == "" { + name = sess.name + } + if name == "" { + _, err := fmt.Fprintln(out, "usage: /save (no active session to save without one)") + return err + } + if err := h.Save(sess.path(name)); err != nil { + _, werr := fmt.Fprintf(out, "save failed: %v\n", err) + return werr + } + sess.name = name + _, err := fmt.Fprintf(out, "saved session %q\n", name) + return err +} + +func runLoadCommand(out io.Writer, h *history.History, sess *session, arg string) error { + if arg == "" { + _, err := fmt.Fprintln(out, "usage: /load ") + return err + } + if err := h.LoadInto(sess.path(arg)); err != nil { + _, werr := fmt.Fprintf(out, "load failed: %v\n", err) + return werr + } + sess.name = arg + _, err := fmt.Fprintf(out, "loaded session %q (%d message(s))\n", arg, len(h.Snapshot())) + return err +} + +func printUsage(out io.Writer, ag *agent.Agent) { + turns := ag.TurnUsages() + cumulative := ag.CumulativeUsage() + _, _ = fmt.Fprintf(out, "cumulative: %d input, %d output token(s) across %d turn(s)\n", + cumulative.InputTokens, cumulative.OutputTokens, len(turns)) + if len(turns) > 0 { + last := turns[len(turns)-1] + _, _ = fmt.Fprintf(out, "last turn: %d input, %d output token(s)\n", last.InputTokens, last.OutputTokens) + } +} diff --git a/cmd/pico/logging.go b/cmd/pico/logging.go new file mode 100644 index 0000000..c7d3926 --- /dev/null +++ b/cmd/pico/logging.go @@ -0,0 +1,40 @@ +package main + +import ( + "fmt" + "io" + "log/slog" + + "github.com/reno/pico-code/internal/config" +) + +// setupLogging points the default slog logger at w with cfg.LogLevel's +// severity. Called once at startup; only --log-level=debug ever dumps +// provider request/response bodies (9.2), and everything logs to w +// (stderr in production) rather than stdout, so piped/plain chat output +// stays clean regardless of level. +func setupLogging(w io.Writer, level config.LogLevel) error { + slogLevel, err := parseLogLevel(level) + if err != nil { + return err + } + slog.SetDefault(slog.New(slog.NewTextHandler(w, &slog.HandlerOptions{Level: slogLevel}))) + return nil +} + +func parseLogLevel(level config.LogLevel) (slog.Level, error) { + switch level { + case config.LogLevelDebug: + return slog.LevelDebug, nil + case config.LogLevelInfo: + return slog.LevelInfo, nil + case config.LogLevelWarn: + return slog.LevelWarn, nil + case config.LogLevelError: + return slog.LevelError, nil + default: + // config.Load already validates this; reaching here means a caller + // built a Config by hand rather than through Load. + return 0, fmt.Errorf("logging: unknown level %q", level) + } +} diff --git a/cmd/pico/logging_test.go b/cmd/pico/logging_test.go new file mode 100644 index 0000000..dfd7191 --- /dev/null +++ b/cmd/pico/logging_test.go @@ -0,0 +1,34 @@ +package main + +import ( + "bytes" + "log/slog" + "testing" + + "github.com/reno/pico-code/internal/config" +) + +func TestSetupLoggingAppliesConfiguredLevel(t *testing.T) { + var buf bytes.Buffer + if err := setupLogging(&buf, config.LogLevelWarn); err != nil { + t.Fatalf("setupLogging() error = %v", err) + } + t.Cleanup(func() { slog.SetDefault(slog.New(slog.NewTextHandler(&bytes.Buffer{}, nil))) }) + + slog.Info("should not appear") + slog.Warn("should appear") + + out := buf.String() + if bytes.Contains([]byte(out), []byte("should not appear")) { + t.Errorf("Info-level line logged despite --log-level=warn: %q", out) + } + if !bytes.Contains([]byte(out), []byte("should appear")) { + t.Errorf("Warn-level line missing: %q", out) + } +} + +func TestParseLogLevelRejectsUnknownValue(t *testing.T) { + if _, err := parseLogLevel(config.LogLevel("verbose")); err == nil { + t.Fatal("parseLogLevel() = nil error, want one for an unknown level") + } +} diff --git a/cmd/pico/main.go b/cmd/pico/main.go index fe40e1d..9d45f88 100644 --- a/cmd/pico/main.go +++ b/cmd/pico/main.go @@ -6,6 +6,9 @@ import ( "os" "os/signal" "syscall" + + _ "github.com/reno/pico-code/internal/llm/anthropic" + _ "github.com/reno/pico-code/internal/llm/ollama" ) func main() { diff --git a/cmd/pico/root_test.go b/cmd/pico/root_test.go index 5dfde86..2d0cfd5 100644 --- a/cmd/pico/root_test.go +++ b/cmd/pico/root_test.go @@ -4,8 +4,30 @@ import ( "bytes" "strings" "testing" + + "github.com/spf13/cobra" + + "github.com/reno/pico-code/internal/config" ) +// stubRunChat swaps the package-level runChat var for a stub that records +// the resolved cfg instead of driving a real provider/agent turn, returning +// a getter for it (valid only after Execute runs) and restoring the +// original on cleanup — tests that only care about how flags/env resolve +// into a Config use this so they never touch the network or block on +// stdin. +func stubRunChat(t *testing.T) func() *config.Config { + t.Helper() + var got *config.Config + orig := runChat + runChat = func(_ *cobra.Command, cfg *config.Config) error { + got = cfg + return nil + } + t.Cleanup(func() { runChat = orig }) + return func() *config.Config { return got } +} + func TestChatHelpListsEveryFlag(t *testing.T) { root := newRootCmd(func(string) string { return "" }) buf := &bytes.Buffer{} @@ -21,7 +43,7 @@ func TestChatHelpListsEveryFlag(t *testing.T) { wantFlags := []string{ "--provider", "--model", "--max-turns", "--token-budget", "--workspace", "--yes", "--tools", "--stream", "--tui", "--log-level", - "--num-ctx", + "--num-ctx", "--allow-write", "--session", "--allow-commands", "--context-window", } for _, flag := range wantFlags { if !strings.Contains(out, flag) { @@ -31,36 +53,38 @@ func TestChatHelpListsEveryFlag(t *testing.T) { } func TestChatProviderEnvFallback(t *testing.T) { + getCfg := stubRunChat(t) env := map[string]string{"PICO_CODE_PROVIDER": "ollama"} getenv := func(k string) string { return env[k] } root := newRootCmd(getenv) buf := &bytes.Buffer{} root.SetOut(buf) - root.SetArgs([]string{"chat"}) + root.SetArgs([]string{"chat", "--model=test-model"}) if err := root.Execute(); err != nil { t.Fatalf("chat returned error: %v", err) } - if !strings.Contains(buf.String(), "provider=ollama") { - t.Errorf("expected env fallback to select ollama provider, got: %s", buf.String()) + if got := getCfg().Provider; got != config.ProviderOllama { + t.Errorf("Provider = %q, want env fallback to select %q", got, config.ProviderOllama) } } func TestChatFlagOverridesEnv(t *testing.T) { + getCfg := stubRunChat(t) env := map[string]string{"PICO_CODE_PROVIDER": "ollama"} getenv := func(k string) string { return env[k] } root := newRootCmd(getenv) buf := &bytes.Buffer{} root.SetOut(buf) - root.SetArgs([]string{"chat", "--provider=anthropic"}) + root.SetArgs([]string{"chat", "--provider=anthropic", "--model=test-model"}) if err := root.Execute(); err != nil { t.Fatalf("chat returned error: %v", err) } - if !strings.Contains(buf.String(), "provider=anthropic") { - t.Errorf("expected explicit flag to override env, got: %s", buf.String()) + if got := getCfg().Provider; got != config.ProviderAnthropic { + t.Errorf("Provider = %q, want the explicit flag to override env and select %q", got, config.ProviderAnthropic) } } @@ -79,3 +103,19 @@ func TestChatUnknownProviderIsClearError(t *testing.T) { t.Errorf("expected error to name the offending value, got: %v", err) } } + +func TestChatMissingModelIsClearError(t *testing.T) { + root := newRootCmd(func(string) string { return "" }) + buf := &bytes.Buffer{} + root.SetOut(buf) + root.SetErr(buf) + root.SetArgs([]string{"chat"}) + + err := root.Execute() + if err == nil { + t.Fatal("expected an error when --model is not passed") + } + if !strings.Contains(err.Error(), "--model") { + t.Errorf("expected error to mention --model, got: %v", err) + } +} diff --git a/cmd/pico/sessions.go b/cmd/pico/sessions.go new file mode 100644 index 0000000..e0ccf0a --- /dev/null +++ b/cmd/pico/sessions.go @@ -0,0 +1,79 @@ +package main + +import ( + "errors" + "fmt" + "os" + "path/filepath" + + "github.com/reno/pico-code/internal/history" +) + +// session tracks which on-disk file (if any) the current conversation +// should be saved to after each turn — set at startup by --session, and +// changeable at runtime by /save and /load. +type session struct { + dir string + name string +} + +// newSession returns a session rooted at dir (created if missing), named +// initially by name ("" for no active session). +func newSession(dir, name string) (*session, error) { + if err := os.MkdirAll(dir, 0o700); err != nil { + return nil, fmt.Errorf("creating sessions directory %s: %w", dir, err) + } + return &session{dir: dir, name: name}, nil +} + +// defaultSessionsDir is where sessions live absent any override; a real +// user's home directory in production, a test's t.TempDir() in tests. +func defaultSessionsDir() (string, error) { + home, err := os.UserHomeDir() + if err != nil { + return "", fmt.Errorf("resolving home directory for sessions: %w", err) + } + return filepath.Join(home, ".pico-code", "sessions"), nil +} + +func (s *session) path(name string) string { + return filepath.Join(s.dir, name+".json") +} + +// active reports whether a session is currently named (and so gets +// auto-saved after each turn). +func (s *session) active() bool { + return s.name != "" +} + +// loadOrCreateHistory resolves cfg's --session into a History: a fresh one +// if unnamed or the file doesn't exist yet, or the resumed one if it does. +// A corrupt session file surfaces as a wrapped error here rather than a +// panic (8.3's AC) — history.Load already guarantees that, this just +// preserves it up the call chain. +func (s *session) loadOrCreateHistory() (*history.History, error) { + if !s.active() { + return history.New(), nil + } + path := s.path(s.name) + if _, err := os.Stat(path); errors.Is(err, os.ErrNotExist) { + return history.New(), nil + } + h, err := history.Load(path) + if err != nil { + return nil, fmt.Errorf("resuming session %q: %w", s.name, err) + } + return h, nil +} + +// saveIfActive persists h to the current session's file, a no-op if no +// session is active. +func (s *session) saveIfActive(h *history.History) error { + if !s.active() { + return nil + } + if err := h.Save(s.path(s.name)); err != nil { + return fmt.Errorf("saving session %q: %w", s.name, err) + } + return nil +} diff --git a/docs/tui-check.md b/docs/tui-check.md new file mode 100644 index 0000000..3873864 --- /dev/null +++ b/docs/tui-check.md @@ -0,0 +1,29 @@ +# TUI manual check + +Run `./bin/pico chat --tui --provider=ollama --model=` (or +`--provider=anthropic`) against a live backend and walk this list. Tick each +box only after actually observing the behavior — this is 7.2's AC, and it is +explicitly manual: nothing here is exercised by `make test`. + +- [ ] Scrollback viewport shows prior turns; PgUp/PgDn scroll it. +- [ ] Textarea accepts multi-line input; Enter (not shift+Enter) submits. +- [ ] Spinner appears while waiting on the provider and disappears once the + reply starts. +- [ ] Assistant text appears incrementally as it streams, not all at once. +- [ ] A tool call shows a running indicator, then flips to ✓ or ✗ once it + finishes. +- [ ] Markdown in a finished reply (headings, code fences, lists) renders + formatted, not as raw `#`/`` ``` `` characters. +- [ ] A tool that needs approval (`write_file` with `--allow-write`, or + `run_command`) shows a modal; `y` approves, `n`/Esc denies, and the + turn continues either way. +- [ ] Two tools needing approval in the same round show their modals one at + a time, not stacked. +- [ ] Ctrl+C while a turn is in flight cancels that turn and returns to the + prompt; the TUI keeps running. +- [ ] Ctrl+D exits the program cleanly, including while a turn is in flight. +- [ ] Resizing the terminal reflows the viewport and textarea without + corrupting prior output. +- [ ] `/usage`, `/new`, `/save `, and `/load ` each show their + output in the scrollback, prefixed with the command line typed; `/new` + clears history, `/load` restores it. diff --git a/go.mod b/go.mod index 422f502..f3f2b82 100644 --- a/go.mod +++ b/go.mod @@ -4,20 +4,50 @@ go 1.26.0 require ( github.com/anthropics/anthropic-sdk-go v1.67.0 + github.com/charmbracelet/bubbles v1.0.0 + github.com/charmbracelet/bubbletea v1.3.10 + github.com/charmbracelet/glamour v1.0.0 github.com/google/go-cmp v0.7.0 github.com/invopop/jsonschema v0.14.0 github.com/ollama/ollama v0.33.1 github.com/spf13/cobra v1.10.2 go.uber.org/goleak v1.3.0 + golang.org/x/sync v0.17.0 ) require ( + github.com/alecthomas/chroma/v2 v2.20.0 // indirect + github.com/atotto/clipboard v0.1.4 // indirect + github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect + github.com/aymerick/douceur v0.2.0 // indirect github.com/bahlo/generic-list-go v0.2.0 // indirect github.com/buger/jsonparser v1.1.2 // indirect + github.com/charmbracelet/colorprofile v0.4.1 // indirect + github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834 // indirect + github.com/charmbracelet/x/ansi v0.11.6 // indirect + github.com/charmbracelet/x/cellbuf v0.0.15 // indirect + github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf // indirect + github.com/charmbracelet/x/term v0.2.2 // indirect + github.com/clipperhouse/displaywidth v0.9.0 // indirect + github.com/clipperhouse/stringish v0.1.1 // indirect + github.com/clipperhouse/uax29/v2 v2.5.0 // indirect + github.com/dlclark/regexp2 v1.11.5 // indirect + github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect github.com/google/uuid v1.6.0 // indirect + github.com/gorilla/css v1.0.1 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect + github.com/lucasb-eyer/go-colorful v1.3.0 // indirect github.com/mailru/easyjson v0.7.7 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/mattn/go-localereader v0.0.1 // indirect + github.com/mattn/go-runewidth v0.0.19 // indirect + github.com/microcosm-cc/bluemonday v1.0.27 // indirect + github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect + github.com/muesli/cancelreader v0.2.2 // indirect + github.com/muesli/reflow v0.3.0 // indirect + github.com/muesli/termenv v0.16.0 // indirect github.com/pb33f/ordered-map/v2 v2.3.1 // indirect + github.com/rivo/uniseg v0.4.7 // indirect github.com/spf13/pflag v1.0.9 // indirect github.com/standard-webhooks/standard-webhooks/libraries v0.0.1 // indirect github.com/tidwall/gjson v1.18.0 // indirect @@ -25,9 +55,14 @@ require ( github.com/tidwall/pretty v1.2.1 // indirect github.com/tidwall/sjson v1.2.5 // indirect github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect + github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect + github.com/yuin/goldmark v1.7.13 // indirect + github.com/yuin/goldmark-emoji v1.0.6 // indirect go.yaml.in/yaml/v4 v4.0.0-rc.2 // indirect golang.org/x/crypto v0.43.0 // indirect - golang.org/x/sync v0.17.0 // indirect - golang.org/x/sys v0.37.0 // indirect + golang.org/x/net v0.46.0 // indirect + golang.org/x/sys v0.38.0 // indirect + golang.org/x/term v0.36.0 // indirect + golang.org/x/text v0.30.0 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index 126955e..38a82b5 100644 --- a/go.sum +++ b/go.sum @@ -1,18 +1,68 @@ +github.com/MakeNowJust/heredoc v1.0.0 h1:cXCdzVdstXyiTqTvfqk9SDHpKNjxuom+DOlyEeQ4pzQ= +github.com/MakeNowJust/heredoc v1.0.0/go.mod h1:mG5amYoWBHf8vpLOuehzbGGw0EHxpZZ6lCpQ4fNJ8LE= +github.com/alecthomas/assert/v2 v2.11.0 h1:2Q9r3ki8+JYXvGsDyBXwH3LcJ+WK5D0gc5E8vS6K3D0= +github.com/alecthomas/assert/v2 v2.11.0/go.mod h1:Bze95FyfUr7x34QZrjL+XP+0qgp/zg8yS+TtBj1WA3k= +github.com/alecthomas/chroma/v2 v2.20.0 h1:sfIHpxPyR07/Oylvmcai3X/exDlE8+FA820NTz+9sGw= +github.com/alecthomas/chroma/v2 v2.20.0/go.mod h1:e7tViK0xh/Nf4BYHl00ycY6rV7b8iXBksI9E359yNmA= +github.com/alecthomas/repr v0.5.1 h1:E3G4t2QbHTSNpPKBgMTln5KLkZHLOcU7r37J4pXBuIg= +github.com/alecthomas/repr v0.5.1/go.mod h1:Fr0507jx4eOXV7AlPV6AVZLYrLIuIeSOWtW57eE/O/4= github.com/anthropics/anthropic-sdk-go v1.67.0 h1:tTsVSR+YD6R+1EPoMyZ9im7YKEjJpYGZk5klcDi9HRg= github.com/anthropics/anthropic-sdk-go v1.67.0/go.mod h1:3EfIfmFqxH6rbiLcIP4tPFyXL/IHakx2wDG4OU+TIEI= +github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4= +github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI= +github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= +github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= +github.com/aymanbagabas/go-udiff v0.3.1 h1:LV+qyBQ2pqe0u42ZsUEtPiCaUoqgA9gYRDs3vj1nolY= +github.com/aymanbagabas/go-udiff v0.3.1/go.mod h1:G0fsKmG+P6ylD0r6N/KgQD/nWzgfnl8ZBcNLgcbrw8E= +github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk= +github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4= github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= github.com/buger/jsonparser v1.1.2 h1:frqHqw7otoVbk5M8LlE/L7HTnIq2v9RX6EJ48i9AxJk= github.com/buger/jsonparser v1.1.2/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= +github.com/charmbracelet/bubbles v1.0.0 h1:12J8/ak/uCZEMQ6KU7pcfwceyjLlWsDLAxB5fXonfvc= +github.com/charmbracelet/bubbles v1.0.0/go.mod h1:9d/Zd5GdnauMI5ivUIVisuEm3ave1XwXtD1ckyV6r3E= +github.com/charmbracelet/bubbletea v1.3.10 h1:otUDHWMMzQSB0Pkc87rm691KZ3SWa4KUlvF9nRvCICw= +github.com/charmbracelet/bubbletea v1.3.10/go.mod h1:ORQfo0fk8U+po9VaNvnV95UPWA1BitP1E0N6xJPlHr4= +github.com/charmbracelet/colorprofile v0.4.1 h1:a1lO03qTrSIRaK8c3JRxJDZOvhvIeSco3ej+ngLk1kk= +github.com/charmbracelet/colorprofile v0.4.1/go.mod h1:U1d9Dljmdf9DLegaJ0nGZNJvoXAhayhmidOdcBwAvKk= +github.com/charmbracelet/glamour v1.0.0 h1:AWMLOVFHTsysl4WV8T8QgkQ0s/ZNZo7CiE4WKhk8l08= +github.com/charmbracelet/glamour v1.0.0/go.mod h1:DSdohgOBkMr2ZQNhw4LZxSGpx3SvpeujNoXrQyH2hxo= +github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834 h1:ZR7e0ro+SZZiIZD7msJyA+NjkCNNavuiPBLgerbOziE= +github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834/go.mod h1:aKC/t2arECF6rNOnaKaVU6y4t4ZeHQzqfxedE/VkVhA= +github.com/charmbracelet/x/ansi v0.11.6 h1:GhV21SiDz/45W9AnV2R61xZMRri5NlLnl6CVF7ihZW8= +github.com/charmbracelet/x/ansi v0.11.6/go.mod h1:2JNYLgQUsyqaiLovhU2Rv/pb8r6ydXKS3NIttu3VGZQ= +github.com/charmbracelet/x/cellbuf v0.0.15 h1:ur3pZy0o6z/R7EylET877CBxaiE1Sp1GMxoFPAIztPI= +github.com/charmbracelet/x/cellbuf v0.0.15/go.mod h1:J1YVbR7MUuEGIFPCaaZ96KDl5NoS0DAWkskup+mOY+Q= +github.com/charmbracelet/x/exp/golden v0.0.0-20241011142426-46044092ad91 h1:payRxjMjKgx2PaCWLZ4p3ro9y97+TVLZNaRZgJwSVDQ= +github.com/charmbracelet/x/exp/golden v0.0.0-20241011142426-46044092ad91/go.mod h1:wDlXFlCrmJ8J+swcL/MnGUuYnqgQdW9rhSD61oNMb6U= +github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf h1:rLG0Yb6MQSDKdB52aGX55JT1oi0P0Kuaj7wi1bLUpnI= +github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf/go.mod h1:B3UgsnsBZS/eX42BlaNiJkD1pPOUa+oF1IYC6Yd2CEU= +github.com/charmbracelet/x/term v0.2.2 h1:xVRT/S2ZcKdhhOuSP4t5cLi5o+JxklsoEObBSgfgZRk= +github.com/charmbracelet/x/term v0.2.2/go.mod h1:kF8CY5RddLWrsgVwpw4kAa6TESp6EB5y3uxGLeCqzAI= +github.com/clipperhouse/displaywidth v0.9.0 h1:Qb4KOhYwRiN3viMv1v/3cTBlz3AcAZX3+y9OLhMtAtA= +github.com/clipperhouse/displaywidth v0.9.0/go.mod h1:aCAAqTlh4GIVkhQnJpbL0T/WfcrJXHcj8C0yjYcjOZA= +github.com/clipperhouse/stringish v0.1.1 h1:+NSqMOr3GR6k1FdRhhnXrLfztGzuG+VuFDfatpWHKCs= +github.com/clipperhouse/stringish v0.1.1/go.mod h1:v/WhFtE1q0ovMta2+m+UbpZ+2/HEXNWYXQgCt4hdOzA= +github.com/clipperhouse/uax29/v2 v2.5.0 h1:x7T0T4eTHDONxFJsL94uKNKPHrclyFI0lm7+w94cO8U= +github.com/clipperhouse/uax29/v2 v2.5.0/go.mod h1:Wn1g7MK6OoeDT0vL+Q0SQLDz/KpfsVRgg6W7ihQeh4g= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dlclark/regexp2 v1.11.5 h1:Q/sSnsKerHeCkc/jSTNq1oCm7KiVgUMZRDUoRu0JQZQ= +github.com/dlclark/regexp2 v1.11.5/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= github.com/dnaeon/go-vcr v1.2.0 h1:zHCHvJYTMh1N7xnV7zf1m1GPBF9Ad0Jk/whtQ1663qI= github.com/dnaeon/go-vcr v1.2.0/go.mod h1:R4UdLID7HZT3taECzJs4YgbbH6PIGXB6W/sc5OLb6RQ= +github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4= +github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/gorilla/css v1.0.1 h1:ntNaBIghp6JmvWnxbZKANoLyuXTPZ4cAMlo6RyhlbO8= +github.com/gorilla/css v1.0.1/go.mod h1:BvnYkspnSzMmwRK+b8/xgNPLiIuNZr6vbZBTPQ2A3b0= +github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM= +github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= github.com/invopop/jsonschema v0.14.0 h1:MHQqLhvpNUZfw+hM3AZDYK7jxO8FZoQeQM77g8iyZjg= @@ -22,14 +72,37 @@ github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI= github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag= +github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4= +github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88= +github.com/mattn/go-runewidth v0.0.12/go.mod h1:RAqKPSqVFrSLVXbA8x7dzmKdmGzieGRCM46jaSJTDAk= +github.com/mattn/go-runewidth v0.0.19 h1:v++JhqYnZuu5jSKrk9RbgF5v4CGUjqRfBm05byFGLdw= +github.com/mattn/go-runewidth v0.0.19/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs= +github.com/microcosm-cc/bluemonday v1.0.27 h1:MpEUotklkwCSLeH+Qdx1VJgNqLlpY2KXwXFM08ygZfk= +github.com/microcosm-cc/bluemonday v1.0.27/go.mod h1:jFi9vgW+H7c3V0lb6nR74Ib/DIB5OBs92Dimizgw2cA= +github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI= +github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo= +github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA= +github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo= +github.com/muesli/reflow v0.3.0 h1:IFsN6K9NfGtjeggFP+68I4chLZV2yIKsXJFNZ+eWh6s= +github.com/muesli/reflow v0.3.0/go.mod h1:pbwTDkVPibjO2kyvBQRBxTWEEGDGq0FlB1BIKtnHY/8= +github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc= +github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk= github.com/ollama/ollama v0.33.1 h1:nEeqiZzqlO+RfjJ7AvfUgrYlV/wbrxPl2szRg9mtxPU= github.com/ollama/ollama v0.33.1/go.mod h1:Kekx/+OtFZHmqbkVH/QUUDVcMQS+1pg1dcz4Qy7TGn4= github.com/pb33f/ordered-map/v2 v2.3.1 h1:5319HDO0aw4DA4gzi+zv4FXU9UlSs3xGZ40wcP1nBjY= github.com/pb33f/ordered-map/v2 v2.3.1/go.mod h1:qxFQgd0PkVUtOMCkTapqotNgzRhMPL7VvaHKbd1HnmQ= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rivo/uniseg v0.1.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= +github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= +github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= +github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU= github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4= @@ -51,6 +124,12 @@ github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= +github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= +github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= +github.com/yuin/goldmark v1.7.13 h1:GPddIs617DnBLFFVJFgpo1aBfe/4xcvMc3SB5t/D0pA= +github.com/yuin/goldmark v1.7.13/go.mod h1:ip/1k0VRfGynBgxOz0yCqHrbZXhcjxyuS66Brc7iBKg= +github.com/yuin/goldmark-emoji v1.0.6 h1:QWfF2FYaXwL74tfGOW5izeiZepUDroDJfWubQI9HTHs= +github.com/yuin/goldmark-emoji v1.0.6/go.mod h1:ukxJDKFpdFb5x0a5HqbdlcKtebh086iJpI31LTKmWuA= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= @@ -58,12 +137,20 @@ go.yaml.in/yaml/v4 v4.0.0-rc.2 h1:/FrI8D64VSr4HtGIlUtlFMGsm7H7pWTbj6vOLVZcA6s= go.yaml.in/yaml/v4 v4.0.0-rc.2/go.mod h1:aZqd9kCMsGL7AuUv/m/PvWLdg5sjJsZ4oHDEnfPPfY0= golang.org/x/crypto v0.43.0 h1:dduJYIi3A3KOfdGOHX8AVZ/jGiyPa3IbBozJ5kNuE04= golang.org/x/crypto v0.43.0/go.mod h1:BFbav4mRNlXJL4wNeejLpWxB7wMbc79PdRGhWKncxR0= +golang.org/x/exp v0.0.0-20250218142911-aa4b98e5adaa h1:t2QcU6V556bFjYgu4L6C+6VrCPyJZ+eyRsABUPs1mz4= +golang.org/x/exp v0.0.0-20250218142911-aa4b98e5adaa/go.mod h1:BHOTPb3L19zxehTsLoJXVaTktb06DFgmdW6Wb9s8jqk= +golang.org/x/net v0.46.0 h1:giFlY12I07fugqwPuWJi68oOnpfqFnJIJzaIIm2JVV4= +golang.org/x/net v0.46.0/go.mod h1:Q9BGdFy1y4nkUwiLvT5qtyhAnEHgnQ/zd8PfU6nc210= golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug= golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= -golang.org/x/sys v0.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ= -golang.org/x/sys v0.37.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= +golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc= +golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= golang.org/x/term v0.36.0 h1:zMPR+aF8gfksFprF/Nc/rd1wRS1EI6nDBGyWAvDzx2Q= golang.org/x/term v0.36.0/go.mod h1:Qu394IJq6V6dCBRgwqshf3mPF85AqzYEzofzRdZkWss= +golang.org/x/text v0.30.0 h1:yznKA/E9zq54KzlzBEAWn1NXSQ8DIp/NYMy88xJjl4k= +golang.org/x/text v0.30.0/go.mod h1:yDdHFIX9t+tORqspjENWgzaCVXgk0yYnYuSZ8UzzBVM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/internal/agent/approval.go b/internal/agent/approval.go new file mode 100644 index 0000000..b50f676 --- /dev/null +++ b/internal/agent/approval.go @@ -0,0 +1,56 @@ +package agent + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "io" + "strings" +) + +// Approver decides whether a tool call needing approval (see +// tools.ApprovalRequired) proceeds. CLAUDE.md invariant 5 makes the agent +// loop the sole owner of this decision — a Tool itself never prompts. +type Approver interface { + Approve(ctx context.Context, toolName string, input json.RawMessage, preview string) (bool, error) +} + +type autoApprove struct{} + +func (autoApprove) Approve(context.Context, string, json.RawMessage, string) (bool, error) { + return true, nil +} + +// AutoApprove approves every call without prompting — used for --yes. +var AutoApprove Approver = autoApprove{} + +// ConsoleApprover prompts on Out and reads a y/N answer from In. +type ConsoleApprover struct { + In io.Reader + Out io.Writer +} + +// Approve implements Approver. +func (c ConsoleApprover) Approve(_ context.Context, toolName string, input json.RawMessage, preview string) (bool, error) { + if _, err := fmt.Fprintf(c.Out, "Approve call to %q with input %s?\n", toolName, input); err != nil { + return false, fmt.Errorf("agent: write approval prompt: %w", err) + } + if preview != "" { + if _, err := fmt.Fprintln(c.Out, preview); err != nil { + return false, fmt.Errorf("agent: write approval preview: %w", err) + } + } + if _, err := fmt.Fprint(c.Out, "[y/N]: "); err != nil { + return false, fmt.Errorf("agent: write approval prompt: %w", err) + } + + line, err := bufio.NewReader(c.In).ReadString('\n') + if err != nil && err != io.EOF { + return false, fmt.Errorf("agent: read approval answer: %w", err) + } + answer := strings.ToLower(strings.TrimSpace(line)) + return answer == "y" || answer == "yes", nil +} + +var _ Approver = ConsoleApprover{} diff --git a/internal/agent/approval_test.go b/internal/agent/approval_test.go new file mode 100644 index 0000000..46d11e2 --- /dev/null +++ b/internal/agent/approval_test.go @@ -0,0 +1,51 @@ +package agent + +import ( + "bytes" + "context" + "encoding/json" + "strings" + "testing" +) + +func TestAutoApproveAlwaysApproves(t *testing.T) { + ok, err := AutoApprove.Approve(context.Background(), "any_tool", json.RawMessage(`{}`), "") + if err != nil { + t.Fatalf("Approve() error = %v", err) + } + if !ok { + t.Error("AutoApprove.Approve() = false, want true") + } +} + +func TestConsoleApprover(t *testing.T) { + tests := []struct { + name string + input string + want bool + }{ + {"yes", "y\n", true}, + {"full yes", "yes\n", true}, + {"no", "n\n", false}, + {"empty defaults to no", "\n", false}, + {"garbage defaults to no", "sure\n", false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + out := &bytes.Buffer{} + c := ConsoleApprover{In: strings.NewReader(tt.input), Out: out} + + got, err := c.Approve(context.Background(), "run_command", json.RawMessage(`{"command":"echo"}`), "preview text") + if err != nil { + t.Fatalf("Approve() error = %v", err) + } + if got != tt.want { + t.Errorf("Approve() = %v, want %v", got, tt.want) + } + if !strings.Contains(out.String(), "run_command") || !strings.Contains(out.String(), "preview text") { + t.Errorf("prompt output = %q, want it to mention the tool name and preview", out.String()) + } + }) + } +} diff --git a/internal/agent/compact.go b/internal/agent/compact.go new file mode 100644 index 0000000..edd1c61 --- /dev/null +++ b/internal/agent/compact.go @@ -0,0 +1,110 @@ +package agent + +import ( + "context" + "fmt" + "log/slog" + "strings" + + "github.com/reno/pico-code/internal/history" + "github.com/reno/pico-code/internal/llm" +) + +// CompactionPolicy configures when Run/RunStream summarize the oldest +// turns to keep context small. The zero value (ContextWindow == 0) +// disables it. +type CompactionPolicy struct { + // ContextWindow is the provider's context window, in tokens. + ContextWindow int + // TriggerFraction is the fraction of ContextWindow that, once + // estimated usage exceeds it, triggers a compaction. + TriggerFraction float64 + // KeepTurns is how many of the most recent user turns are kept + // verbatim; everything older is summarized into one message. + KeepTurns int +} + +// SetCompactionPolicy configures p for future turns; the zero value +// disables compaction. It is not a New parameter since most callers +// (every existing test, and any caller not tracking context budget) don't +// need it. +func (a *Agent) SetCompactionPolicy(p CompactionPolicy) { + a.compaction = p +} + +// maybeCompact summarizes the oldest turns into a single synthetic message +// once history's estimated size crosses the configured threshold. It is +// best-effort: a failed summarization call (or a disabled/inapplicable +// policy) leaves history untouched rather than failing the turn — the +// worst case is simply not compacting yet. +func (a *Agent) maybeCompact(ctx context.Context) { + p := a.compaction + if p.ContextWindow <= 0 || p.TriggerFraction <= 0 || p.KeepTurns <= 0 { + return + } + + messages := a.history.Snapshot() + threshold := int(float64(p.ContextWindow) * p.TriggerFraction) + if history.EstimateTokens(messages) <= threshold { + return + } + + boundary := history.CompactionBoundary(messages, p.KeepTurns) + if boundary <= 0 { + return + } + + summary, err := a.summarize(ctx, messages[:boundary]) + if err != nil { + slog.Warn("agent: skipping compaction, summarization call failed", "error", err) + return + } + if err := a.history.Compact(p.KeepTurns, summary); err != nil { + slog.Warn("agent: skipping compaction", "error", err) + } +} + +func (a *Agent) summarize(ctx context.Context, elided []llm.Message) (string, error) { + prompt := buildSummarizationPrompt(elided) + resp, err := a.provider.Chat(ctx, llm.Request{ + System: a.system, + Messages: []llm.Message{{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: prompt}}}}, + MaxTokens: a.maxTokens, + }) + if err != nil { + return "", fmt.Errorf("agent: summarize: %w", err) + } + return textOf(resp.Message), nil +} + +// summarizationPromptTemplate is pinned by a golden test: changing its +// wording is a deliberate choice, not an incidental refactor. +const summarizationPromptTemplate = `Summarize the earlier part of this conversation between you (the assistant) and a user, so the summary can replace the original messages in your context. Preserve any facts, decisions, file paths, commands, or conclusions that later turns might depend on. Do not include pleasantries or filler. Write the summary as plain prose, not a transcript. + +Conversation to summarize: +%s` + +func buildSummarizationPrompt(messages []llm.Message) string { + return fmt.Sprintf(summarizationPromptTemplate, renderTranscript(messages)) +} + +func renderTranscript(messages []llm.Message) string { + var b strings.Builder + for _, m := range messages { + for _, blk := range m.Blocks { + switch v := blk.(type) { + case llm.Text: + fmt.Fprintf(&b, "%s: %s\n", m.Role, v.Text) + case llm.ToolUse: + fmt.Fprintf(&b, "%s: [called tool %s with %s]\n", m.Role, v.Name, v.Input) + case llm.ToolResult: + status := "ok" + if v.IsError { + status = "error" + } + fmt.Fprintf(&b, "%s: [tool result (%s): %s]\n", m.Role, status, v.Content) + } + } + } + return strings.TrimRight(b.String(), "\n") +} diff --git a/internal/agent/compact_test.go b/internal/agent/compact_test.go new file mode 100644 index 0000000..1274b8c --- /dev/null +++ b/internal/agent/compact_test.go @@ -0,0 +1,178 @@ +package agent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "testing" + + "github.com/reno/pico-code/internal/history" + "github.com/reno/pico-code/internal/llm" + "github.com/reno/pico-code/internal/tools" +) + +const summarizationPromptGoldenPath = "testdata/golden/summarization_prompt.txt" + +// TestSummarizationPromptGolden is 8.2's AC: a golden test pins the +// summarization prompt. Run with RECORD=1 to (re)write the golden file, +// mirroring the RECORD=1 convention CLAUDE.md's testing section already +// uses for provider fixtures. +func TestSummarizationPromptGolden(t *testing.T) { + elided := []llm.Message{ + {Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "read config.go and tell me the default port"}}}, + {Role: llm.RoleAssistant, Blocks: []llm.Block{llm.ToolUse{ID: "call_1", Name: "read_file", Input: json.RawMessage(`{"path":"config.go"}`)}}}, + {Role: llm.RoleUser, Blocks: []llm.Block{llm.ToolResult{ToolUseID: "call_1", Content: "const defaultPort = 8080"}}}, + {Role: llm.RoleAssistant, Blocks: []llm.Block{llm.Text{Text: "The default port is 8080."}}}, + } + got := buildSummarizationPrompt(elided) + + if os.Getenv("RECORD") == "1" { + if err := os.WriteFile(summarizationPromptGoldenPath, []byte(got), 0o600); err != nil { + t.Fatalf("WriteFile(%s) error = %v", summarizationPromptGoldenPath, err) + } + } + + want, err := os.ReadFile(summarizationPromptGoldenPath) + if err != nil { + t.Fatalf("ReadFile(%s) error = %v (run with RECORD=1 to generate it)", summarizationPromptGoldenPath, err) + } + if got != string(want) { + t.Errorf("buildSummarizationPrompt() mismatch against %s\ngot:\n%s\nwant:\n%s", summarizationPromptGoldenPath, got, want) + } +} + +// compactionScriptedProvider answers the first Chat call — maybeCompact's +// summarization request, which always happens before the turn's own — with +// summary, then falls through to the wrapped scriptedProvider's script for +// the turn itself. +type compactionScriptedProvider struct { + *scriptedProvider + summary string + summarized bool +} + +func (p *compactionScriptedProvider) Chat(ctx context.Context, req llm.Request) (*llm.Response, error) { + if !p.summarized { + p.summarized = true + return textResponse(p.summary), nil + } + return p.scriptedProvider.Chat(ctx, req) +} + +// TestMaybeCompactTriggersAboveThresholdAndKeepsHistoryValid drives a +// history well past a tiny configured context window and checks that +// compaction fires, the result still passes Validate() (CLAUDE.md +// invariant 3 — the one this task is most likely to break, since +// compaction discards history wholesale), and the synthetic summary +// message is what the (extra) summarization call returned. +func TestMaybeCompactTriggersAboveThresholdAndKeepsHistoryValid(t *testing.T) { + h := history.New() + for i := 0; i < 10; i++ { + h.Append(llm.Message{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "a fairly long question to pad out the estimated token count here"}}}) + h.Append(llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.Text{Text: "a fairly long answer to pad out the estimated token count here too"}}}) + } + + provider := &compactionScriptedProvider{ + scriptedProvider: &scriptedProvider{responses: []*llm.Response{textResponse("final reply")}}, + summary: "everything before this was summarized", + } + a := New(provider, tools.NewRegistry(), h, "", 1024, Guards{}, 0, AutoApprove) + a.SetCompactionPolicy(CompactionPolicy{ContextWindow: 100, TriggerFraction: 0.5, KeepTurns: 2}) + + if _, err := a.Run(context.Background(), "one more question"); err != nil { + t.Fatalf("Run() error = %v", err) + } + + if err := h.Validate(); err != nil { + t.Fatalf("Validate() after compaction error = %v", err) + } + + msgs := h.Snapshot() + if len(msgs) == 0 { + t.Fatal("history is empty after Run()") + } + first := msgs[0] + text, ok := first.Blocks[0].(llm.Text) + if !ok || text.Text != "everything before this was summarized" || first.Role != llm.RoleAssistant { + t.Errorf("messages[0] = %+v, want the synthetic assistant summary", first) + } +} + +func TestMaybeCompactDoesNothingBelowThreshold(t *testing.T) { + h := history.New() + h.Append(llm.Message{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "hi"}}}) + h.Append(llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.Text{Text: "hello"}}}) + before := h.Snapshot() + + provider := &scriptedProvider{responses: []*llm.Response{textResponse("ok")}} + a := New(provider, tools.NewRegistry(), h, "", 1024, Guards{}, 0, AutoApprove) + a.SetCompactionPolicy(CompactionPolicy{ContextWindow: 1_000_000, TriggerFraction: 0.9, KeepTurns: 2}) + + if _, err := a.Run(context.Background(), "another question"); err != nil { + t.Fatalf("Run() error = %v", err) + } + + got := h.Snapshot() + if len(got) != len(before)+2 { + t.Fatalf("history has %d messages, want %d (before + this turn's user/assistant pair, no compaction)", len(got), len(before)+2) + } + if got[0].Role != before[0].Role { + t.Errorf("messages[0].Role = %q, want the original first message untouched (%q) since nothing should have been compacted", got[0].Role, before[0].Role) + } +} + +func TestMaybeCompactSkippedWhenSummarizationFails(t *testing.T) { + h := history.New() + for i := 0; i < 10; i++ { + h.Append(llm.Message{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "a fairly long question to pad out the estimated token count here"}}}) + h.Append(llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.Text{Text: "a fairly long answer to pad out the estimated token count here too"}}}) + } + before := h.Snapshot() + + provider := &alwaysErrorsProvider{err: errors.New("provider unavailable")} + a := New(provider, tools.NewRegistry(), h, "", 1024, Guards{}, 0, AutoApprove) + a.SetCompactionPolicy(CompactionPolicy{ContextWindow: 100, TriggerFraction: 0.5, KeepTurns: 2}) + + _, err := a.Run(context.Background(), "one more question") + if err == nil { + t.Fatal("Run() error = nil, want the turn's own Chat call to fail too, since this fake provider always errors") + } + + got := h.Snapshot() + if len(got) != len(before)+1 { + t.Fatalf("history has %d messages, want %d (before + the new user turn Run appends before calling Chat)", len(got), len(before)+1) + } + if diffMsg := cmpMessage(got[0], before[0]); diffMsg != "" { + t.Errorf("messages[0] changed after a failed summarization call, want it untouched (compaction should be skipped, not partially applied): %s", diffMsg) + } +} + +func cmpMessage(a, b llm.Message) string { + if a.Role != b.Role { + return fmt.Sprintf("Role: got %q, want %q", a.Role, b.Role) + } + at, aok := a.Blocks[0].(llm.Text) + bt, bok := b.Blocks[0].(llm.Text) + if !aok || !bok || at != bt { + return fmt.Sprintf("Blocks[0]: got %+v, want %+v", a.Blocks[0], b.Blocks[0]) + } + return "" +} + +// alwaysErrorsProvider errors on every call, used to make maybeCompact's +// summarization call fail deterministically. +type alwaysErrorsProvider struct { + err error +} + +func (p *alwaysErrorsProvider) Name() string { return "always-errors" } + +func (p *alwaysErrorsProvider) Chat(context.Context, llm.Request) (*llm.Response, error) { + return nil, p.err +} + +func (p *alwaysErrorsProvider) Stream(context.Context, llm.Request) (<-chan llm.Event, error) { + return nil, p.err +} diff --git a/internal/agent/loop.go b/internal/agent/loop.go new file mode 100644 index 0000000..b5ea2c0 --- /dev/null +++ b/internal/agent/loop.go @@ -0,0 +1,334 @@ +// Package agent implements the core loop: call the provider, run any tool +// calls it returns, feed the results back, and repeat until it stops asking. +package agent + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "strings" + "time" + + "golang.org/x/sync/errgroup" + + "github.com/reno/pico-code/internal/history" + "github.com/reno/pico-code/internal/llm" + "github.com/reno/pico-code/internal/tools" + "github.com/reno/pico-code/internal/ui" +) + +// Guards bounds how long a single Run can keep asking the provider for more +// tool calls. A zero field disables that guard. Repetition detection has no +// field: CLAUDE.md fixes it at 3 identical consecutive tool calls. +type Guards struct { + MaxTurns int + TokenBudget int + WallClock time.Duration +} + +const repetitionThreshold = 3 + +// Agent drives one conversation: a provider to talk to, a tool registry to +// satisfy its tool calls, and the history both read from and append to. +type Agent struct { + provider llm.Provider + tools *tools.Registry + history *history.History + system string + maxTokens int + guards Guards + toolTimeout time.Duration + approver Approver + + turnUsages []llm.Usage + compaction CompactionPolicy +} + +// New returns an Agent ready to run turns against provider, using registry +// to execute any tool call it issues, appending to h. toolTimeout bounds a +// single tool call (0 = unlimited). approver is consulted before running any +// tools.ApprovalRequired call; pass AutoApprove for --yes. +func New(provider llm.Provider, registry *tools.Registry, h *history.History, system string, maxTokens int, guards Guards, toolTimeout time.Duration, approver Approver) *Agent { + return &Agent{provider: provider, tools: registry, history: h, system: system, maxTokens: maxTokens, guards: guards, toolTimeout: toolTimeout, approver: approver} +} + +// Run appends userInput to history, then drives the call-provider / +// run-tools cycle until a reply has no ToolUse blocks (its text is +// returned) or a guard trips (a final explanatory Text message is appended +// and returned instead). Every round's tool calls are answered before +// either exit, so history always satisfies CLAUDE.md invariant 3. +func (a *Agent) Run(ctx context.Context, userInput string) (string, error) { + a.maybeCompact(ctx) + a.history.Append(llm.Message{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: userInput}}}) + + rs := &roundState{start: time.Now()} + for { + resp, err := a.provider.Chat(ctx, a.request()) + if err != nil { + return "", fmt.Errorf("agent: chat: %w", err) + } + if text, done := a.runRound(ctx, resp, rs, nil); done { + return text, nil + } + } +} + +// RunStream behaves like Run but drives each round's provider call through +// Stream instead of Chat, handing the raw event channel to renderer so it +// can display text as it arrives. If renderer also implements +// ui.ToolStatusReporter, it additionally learns about each tool call's +// start and outcome — a step Render alone can't see, since running a tool +// happens after a round's events finish, not during them. +func (a *Agent) RunStream(ctx context.Context, userInput string, renderer ui.Renderer) (string, error) { + a.maybeCompact(ctx) + a.history.Append(llm.Message{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: userInput}}}) + + reporter, _ := renderer.(ui.ToolStatusReporter) + rs := &roundState{start: time.Now()} + for { + ch, err := a.provider.Stream(ctx, a.request()) + if err != nil { + return "", fmt.Errorf("agent: stream: %w", err) + } + resp, err := renderer.Render(ctx, ch) + if err != nil { + return "", fmt.Errorf("agent: render: %w", err) + } + if text, done := a.runRound(ctx, resp, rs, reporter); done { + return text, nil + } + } +} + +func (a *Agent) request() llm.Request { + return llm.Request{ + System: a.system, + Messages: a.history.Snapshot(), + Tools: a.tools.Definitions(), + MaxTokens: a.maxTokens, + } +} + +// roundState carries the bookkeeping Run and RunStream both accumulate +// across rounds of the same turn. +type roundState struct { + start time.Time + turns int + usage llm.Usage + repeatStreak int + lastSignature string +} + +// runRound appends resp to history, runs any tool calls it contains, and +// reports whether the turn is over — either because resp had no tool calls +// or because a guard tripped — along with the text to return in that case. +// Either way it's the turn's last round, so this is also where its total +// Usage (summed across every round) gets recorded for TurnUsages/ +// CumulativeUsage (8.1). +func (a *Agent) runRound(ctx context.Context, resp *llm.Response, rs *roundState, reporter ui.ToolStatusReporter) (string, bool) { + rs.turns++ + rs.usage.InputTokens += resp.Usage.InputTokens + rs.usage.OutputTokens += resp.Usage.OutputTokens + a.history.Append(resp.Message) + + calls := toolUseBlocks(resp.Message) + if len(calls) == 0 { + a.turnUsages = append(a.turnUsages, rs.usage) + return textOf(resp.Message), true + } + + if sig := callSignature(calls); sig == rs.lastSignature { + rs.repeatStreak++ + } else { + rs.lastSignature, rs.repeatStreak = sig, 1 + } + + a.history.Append(llm.Message{Role: llm.RoleUser, Blocks: a.runTools(ctx, calls, reporter)}) + + totalTokens := rs.usage.InputTokens + rs.usage.OutputTokens + if reason, tripped := a.guardTripped(rs.turns, totalTokens, time.Since(rs.start), rs.repeatStreak); tripped { + a.turnUsages = append(a.turnUsages, rs.usage) + return a.stopWithExplanation(reason), true + } + return "", false +} + +// TurnUsages returns each completed turn's total Usage (summed across that +// turn's provider rounds), in the order the turns finished. +func (a *Agent) TurnUsages() []llm.Usage { + return append([]llm.Usage(nil), a.turnUsages...) +} + +// CumulativeUsage sums every completed turn's Usage. +func (a *Agent) CumulativeUsage() llm.Usage { + var total llm.Usage + for _, u := range a.turnUsages { + total.InputTokens += u.InputTokens + total.OutputTokens += u.OutputTokens + } + return total +} + +func (a *Agent) guardTripped(turns, totalTokens int, elapsed time.Duration, repeatStreak int) (string, bool) { + switch { + case a.guards.MaxTurns > 0 && turns >= a.guards.MaxTurns: + return fmt.Sprintf("Stopping: reached the maximum of %d turn(s).", a.guards.MaxTurns), true + case a.guards.TokenBudget > 0 && totalTokens >= a.guards.TokenBudget: + return fmt.Sprintf("Stopping: reached the token budget of %d.", a.guards.TokenBudget), true + case a.guards.WallClock > 0 && elapsed >= a.guards.WallClock: + return fmt.Sprintf("Stopping: exceeded the wall-clock timeout of %s.", a.guards.WallClock), true + case repeatStreak >= repetitionThreshold: + return fmt.Sprintf("Stopping: the same tool call repeated %d times in a row.", repeatStreak), true + default: + return "", false + } +} + +func (a *Agent) stopWithExplanation(reason string) string { + a.history.Append(llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.Text{Text: reason}}}) + return reason +} + +// runTools executes every call concurrently and returns their ToolResult +// blocks in calls' original order, independent of completion order. +func (a *Agent) runTools(ctx context.Context, calls []llm.ToolUse, reporter ui.ToolStatusReporter) []llm.Block { + results := make([]llm.Block, len(calls)) + var g errgroup.Group + for i, call := range calls { + g.Go(func() error { + results[i] = a.runTool(ctx, call, reporter) + return nil + }) + } + _ = g.Wait() + return results +} + +type toolOutcome struct { + out string + err error +} + +// runTool always produces a ToolResult, even if the tool panics, times out, +// or ctx is cancelled while it's running. The tool runs in its own +// goroutine so a badly-behaved implementation that ignores ctx can't block +// this call past toolTimeout or cancellation; if the deadline wins the race, +// the goroutine is simply abandoned rather than waited on. +func (a *Agent) runTool(ctx context.Context, call llm.ToolUse, reporter ui.ToolStatusReporter) llm.Block { + if result, denied := a.checkApproval(ctx, call); denied { + return result + } + + if reporter != nil { + reporter.ToolStarted(call.ID, call.Name, call.Input) + } + result := a.doRunTool(ctx, call) + if reporter != nil { + tr := result.(llm.ToolResult) + reporter.ToolFinished(call.ID, call.Name, tr.Content, tr.IsError) + } + return result +} + +func (a *Agent) doRunTool(ctx context.Context, call llm.ToolUse) llm.Block { + toolCtx := ctx + if a.toolTimeout > 0 { + var cancel context.CancelFunc + toolCtx, cancel = context.WithTimeout(ctx, a.toolTimeout) + defer cancel() + } + + done := make(chan toolOutcome, 1) + go func() { + defer func() { + if r := recover(); r != nil { + done <- toolOutcome{err: fmt.Errorf("tool %q panicked: %v", call.Name, r)} + } + }() + out, err := a.tools.Run(toolCtx, call.Name, call.Input) + done <- toolOutcome{out: out, err: err} + }() + + select { + case o := <-done: + if o.err != nil { + return llm.ToolResult{ToolUseID: call.ID, Content: o.err.Error(), IsError: true} + } + return llm.ToolResult{ToolUseID: call.ID, Content: o.out, IsError: false} + case <-toolCtx.Done(): + return llm.ToolResult{ToolUseID: call.ID, Content: fmt.Sprintf("tool %q: %v", call.Name, toolCtx.Err()), IsError: true} + } +} + +// checkApproval reports whether call was denied, in which case the +// returned ToolResult should be used as-is instead of running the tool. +func (a *Agent) checkApproval(ctx context.Context, call llm.ToolUse) (llm.Block, bool) { + t, err := a.tools.Get(call.Name) + if err != nil { + return nil, false + } + ar, ok := t.(tools.ApprovalRequired) + if !ok || !ar.NeedsApproval() { + return nil, false + } + + preview := "" + if p, ok := t.(tools.Previewable); ok { + if text, err := p.Preview(ctx, call.Input); err == nil { + preview = text + } else { + preview = fmt.Sprintf("(preview unavailable: %v)", err) + } + } + + approved, err := a.approver.Approve(ctx, call.Name, call.Input, preview) + if err != nil { + return llm.ToolResult{ToolUseID: call.ID, Content: fmt.Sprintf("approval error: %v", err), IsError: true}, true + } + if !approved { + return llm.ToolResult{ToolUseID: call.ID, Content: fmt.Sprintf("tool %q call denied by user", call.Name), IsError: true}, true + } + return nil, false +} + +func callSignature(calls []llm.ToolUse) string { + parts := make([]string, len(calls)) + for i, c := range calls { + parts[i] = c.Name + ":" + compactJSON(c.Input) + } + return strings.Join(parts, "|") +} + +func compactJSON(raw json.RawMessage) string { + var buf bytes.Buffer + if err := json.Compact(&buf, raw); err != nil { + return string(raw) + } + return buf.String() +} + +func toolUseBlocks(m llm.Message) []llm.ToolUse { + var calls []llm.ToolUse + for _, b := range m.Blocks { + if tu, ok := b.(llm.ToolUse); ok { + calls = append(calls, tu) + } + } + return calls +} + +func textOf(m llm.Message) string { + var b strings.Builder + for _, blk := range m.Blocks { + t, ok := blk.(llm.Text) + if !ok { + continue + } + if b.Len() > 0 { + b.WriteString("\n") + } + b.WriteString(t.Text) + } + return b.String() +} diff --git a/internal/agent/loop_test.go b/internal/agent/loop_test.go new file mode 100644 index 0000000..3b9c54e --- /dev/null +++ b/internal/agent/loop_test.go @@ -0,0 +1,1021 @@ +package agent + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "github.com/google/go-cmp/cmp" + + "github.com/reno/pico-code/internal/history" + "github.com/reno/pico-code/internal/llm" + "github.com/reno/pico-code/internal/tools" + "github.com/reno/pico-code/internal/ui" +) + +// scriptedProvider returns responses in order, one per Chat call — the +// "scripted fake provider" 4.1's AC calls for. +type scriptedProvider struct { + responses []*llm.Response + calls int +} + +func (s *scriptedProvider) Name() string { return "scripted" } + +func (s *scriptedProvider) Chat(ctx context.Context, _ llm.Request) (*llm.Response, error) { + // A real adapter checks ctx before/while talking to the network + // (CLAUDE.md invariant 6); mirroring that here is what lets a + // cancellation test observe Run exiting cleanly via this path rather + // than via scriptedProvider running out of scripted responses. + if err := ctx.Err(); err != nil { + return nil, err + } + if s.calls >= len(s.responses) { + return nil, fmt.Errorf("scripted: no response scripted for call %d", s.calls+1) + } + resp := s.responses[s.calls] + s.calls++ + return resp, nil +} + +// Stream reuses the same scripted responses as Chat, turning each one into +// the event sequence llm.CollectStream would reassemble back into it, so +// RunStream tests can script conversations the same way Run tests do. +func (s *scriptedProvider) Stream(ctx context.Context, _ llm.Request) (<-chan llm.Event, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if s.calls >= len(s.responses) { + return nil, fmt.Errorf("scripted: no response scripted for call %d", s.calls+1) + } + resp := s.responses[s.calls] + s.calls++ + return llm.StreamEvents(ctx, func(send func(llm.Event) bool) { + for _, b := range resp.Message.Blocks { + switch v := b.(type) { + case llm.Text: + if !send(llm.TextDelta(v)) { + return + } + case llm.ToolUse: + if !send(llm.ToolUseStart{ID: v.ID, Name: v.Name}) { + return + } + if !send(llm.ToolUseDone{ID: v.ID, Input: v.Input}) { + return + } + } + } + send(llm.MessageDone{StopReason: resp.StopReason, Usage: resp.Usage}) + }), nil +} + +// delayedTool sleeps for delay before recording its own name to order (under +// mu) and returning result. Delays let the test prove the loop reassembles +// results in call order even when the slower call was issued first. +type delayedTool struct { + name string + delay time.Duration + result string + mu *sync.Mutex + order *[]string +} + +func (d *delayedTool) Name() string { return d.name } +func (d *delayedTool) Description() string { return "a delayed test tool" } +func (d *delayedTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) } + +func (d *delayedTool) Run(_ context.Context, _ json.RawMessage) (string, error) { + time.Sleep(d.delay) + d.mu.Lock() + *d.order = append(*d.order, d.name) + d.mu.Unlock() + return d.result, nil +} + +var _ tools.Tool = (*delayedTool)(nil) + +// TestRunDrivesScriptedConversationWithParallelTools is 4.1's AC: a scripted +// fake provider drives a 3-round conversation with 2 parallel tools; the +// final history passes Validate(); result order is deterministic. +func TestRunDrivesScriptedConversationWithParallelTools(t *testing.T) { + var mu sync.Mutex + var completionOrder []string + + slow := &delayedTool{name: "slow_tool", delay: 30 * time.Millisecond, result: "slow result", mu: &mu, order: &completionOrder} + fast := &delayedTool{name: "fast_tool", delay: 0, result: "fast result", mu: &mu, order: &completionOrder} + + reg := tools.NewRegistry() + if err := reg.Register(slow); err != nil { + t.Fatalf("Register(slow) error = %v", err) + } + if err := reg.Register(fast); err != nil { + t.Fatalf("Register(fast) error = %v", err) + } + + provider := &scriptedProvider{responses: []*llm.Response{ + { + // Round 1: two parallel calls. slow_tool is called first but + // finishes last, so a naive sequential-append implementation + // would still happen to get this right — the completion-order + // assertion below is what actually catches a bug here. + Message: llm.Message{ + Role: llm.RoleAssistant, + Blocks: []llm.Block{ + llm.ToolUse{ID: "call_1", Name: "slow_tool", Input: json.RawMessage(`{}`)}, + llm.ToolUse{ID: "call_2", Name: "fast_tool", Input: json.RawMessage(`{}`)}, + }, + }, + StopReason: "tool_use", + }, + { + // Round 2: one more tool call. + Message: llm.Message{ + Role: llm.RoleAssistant, + Blocks: []llm.Block{llm.ToolUse{ID: "call_3", Name: "fast_tool", Input: json.RawMessage(`{}`)}}, + }, + StopReason: "tool_use", + }, + { + // Round 3: final text, no tool calls. + Message: llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.Text{Text: "done"}}}, + StopReason: "end_turn", + }, + }} + + h := history.New() + a := New(provider, reg, h, "you are a test agent", 1024, Guards{}, 0, AutoApprove) + + got, err := a.Run(context.Background(), "please help") + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if got != "done" { + t.Errorf("Run() = %q, want %q", got, "done") + } + if provider.calls != 3 { + t.Errorf("provider was called %d times, want 3", provider.calls) + } + + if err := h.Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } + + mu.Lock() + order := append([]string(nil), completionOrder...) + mu.Unlock() + want := []string{"fast_tool", "slow_tool", "fast_tool"} + if diff := cmp.Diff(want, order); diff != "" { + t.Errorf("tool completion order (-want +got):\n%s\n(fast_tool completing before the slow_tool it was issued after proves the round-1 pair ran concurrently)", diff) + } + + msgs := h.Snapshot() + if len(msgs) != 6 { + t.Fatalf("history has %d messages, want 6 (user, assistant, user, assistant, user, assistant)", len(msgs)) + } + + round1Results := msgs[2] + if round1Results.Role != llm.RoleUser { + t.Fatalf("messages[2].Role = %q, want %q", round1Results.Role, llm.RoleUser) + } + if len(round1Results.Blocks) != 2 { + t.Fatalf("messages[2] has %d blocks, want 2", len(round1Results.Blocks)) + } + gotIDs := []string{ + round1Results.Blocks[0].(llm.ToolResult).ToolUseID, + round1Results.Blocks[1].(llm.ToolResult).ToolUseID, + } + wantIDs := []string{"call_1", "call_2"} + if diff := cmp.Diff(wantIDs, gotIDs); diff != "" { + t.Errorf("round-1 ToolResult order (-want +got):\n%s\n(must match call order, not completion order)", diff) + } + if got := round1Results.Blocks[0].(llm.ToolResult).Content; got != "slow result" { + t.Errorf("ToolResult for call_1 = %q, want %q", got, "slow result") + } + if got := round1Results.Blocks[1].(llm.ToolResult).Content; got != "fast result" { + t.Errorf("ToolResult for call_2 = %q, want %q", got, "fast result") + } +} + +func TestRunReturnsTextImmediatelyWhenNoToolCall(t *testing.T) { + provider := &scriptedProvider{responses: []*llm.Response{ + { + Message: llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.Text{Text: "hello there"}}}, + StopReason: "end_turn", + }, + }} + + h := history.New() + a := New(provider, tools.NewRegistry(), h, "", 1024, Guards{}, 0, AutoApprove) + + got, err := a.Run(context.Background(), "hi") + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if got != "hello there" { + t.Errorf("Run() = %q, want %q", got, "hello there") + } + if provider.calls != 1 { + t.Errorf("provider was called %d times, want 1", provider.calls) + } + if err := h.Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } +} + +func TestRunTurnsUnknownToolIntoErrorResult(t *testing.T) { + provider := &scriptedProvider{responses: []*llm.Response{ + { + Message: llm.Message{ + Role: llm.RoleAssistant, + Blocks: []llm.Block{llm.ToolUse{ID: "call_1", Name: "does_not_exist", Input: json.RawMessage(`{}`)}}, + }, + StopReason: "tool_use", + }, + { + Message: llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.Text{Text: "gave up"}}}, + StopReason: "end_turn", + }, + }} + + h := history.New() + a := New(provider, tools.NewRegistry(), h, "", 1024, Guards{}, 0, AutoApprove) + + if _, err := a.Run(context.Background(), "hi"); err != nil { + t.Fatalf("Run() error = %v", err) + } + if err := h.Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } + + msgs := h.Snapshot() + result := msgs[2].Blocks[0].(llm.ToolResult) + if !result.IsError { + t.Error("ToolResult.IsError = false, want true for an unknown tool") + } + if result.ToolUseID != "call_1" { + t.Errorf("ToolResult.ToolUseID = %q, want %q", result.ToolUseID, "call_1") + } +} + +func TestRunPropagatesProviderError(t *testing.T) { + provider := &scriptedProvider{responses: nil} + h := history.New() + a := New(provider, tools.NewRegistry(), h, "", 1024, Guards{}, 0, AutoApprove) + + if _, err := a.Run(context.Background(), "hi"); err == nil { + t.Fatal("Run() error = nil, want the provider's error propagated") + } +} + +// echoTool is a minimal registrable tool for guard tests, which only care +// that a ToolUse gets answered, not what the answer is. +type echoTool struct{} + +func (echoTool) Name() string { return "echo_tool" } +func (echoTool) Description() string { return "a no-op test tool" } +func (echoTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) } +func (echoTool) Run(_ context.Context, _ json.RawMessage) (string, error) { return "ok", nil } + +var _ tools.Tool = echoTool{} + +func toolCallResponse(id string) *llm.Response { + return toolCallResponseNamed(id, "echo_tool") +} + +func toolCallResponseNamed(id, name string) *llm.Response { + return &llm.Response{ + Message: llm.Message{ + Role: llm.RoleAssistant, + Blocks: []llm.Block{llm.ToolUse{ID: id, Name: name, Input: json.RawMessage(`{}`)}}, + }, + StopReason: "tool_use", + } +} + +func textResponse(text string) *llm.Response { + return &llm.Response{ + Message: llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.Text{Text: text}}}, + StopReason: "end_turn", + } +} + +func assertFinalMessageIsExplanation(t *testing.T, h *history.History) { + t.Helper() + msgs := h.Snapshot() + last := msgs[len(msgs)-1] + if last.Role != llm.RoleAssistant || len(last.Blocks) != 1 { + t.Fatalf("final message = %+v, want a single-block assistant explanation", last) + } + if _, ok := last.Blocks[0].(llm.Text); !ok { + t.Errorf("final message block = %T, want llm.Text", last.Blocks[0]) + } +} + +// TestRunStopsAtMaxTurnsGuard is 4.2's AC for the max-turns guard: it trips +// exactly once and leaves valid history. +func TestRunStopsAtMaxTurnsGuard(t *testing.T) { + reg := tools.NewRegistry() + if err := reg.Register(echoTool{}); err != nil { + t.Fatalf("Register() error = %v", err) + } + + // A model that would keep asking for tool calls forever; only 2 + // responses are scripted, so if the guard doesn't trip after turn 2, + // the loop's 3rd Chat call fails with scriptedProvider's own error and + // this test's assertions below catch it. + provider := &scriptedProvider{responses: []*llm.Response{toolCallResponse("call_1"), toolCallResponse("call_2")}} + + h := history.New() + a := New(provider, reg, h, "", 1024, Guards{MaxTurns: 2}, 0, AutoApprove) + + got, err := a.Run(context.Background(), "keep going forever") + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if provider.calls != 2 { + t.Errorf("provider was called %d times, want exactly 2", provider.calls) + } + if !strings.Contains(got, "maximum of 2 turn") { + t.Errorf("Run() = %q, want an explanation mentioning the turn limit", got) + } + if err := h.Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } + assertFinalMessageIsExplanation(t, h) +} + +// TestRunStopsAtTokenBudgetGuard is 4.2's AC for the token-budget guard. +func TestRunStopsAtTokenBudgetGuard(t *testing.T) { + reg := tools.NewRegistry() + if err := reg.Register(echoTool{}); err != nil { + t.Fatalf("Register() error = %v", err) + } + + round := func(id string, tokens int) *llm.Response { + resp := toolCallResponse(id) + resp.Usage = llm.Usage{InputTokens: tokens} + return resp + } + // 10 tokens/round: budget 15 is under cumulative 20 (after round 2) but + // over 10 (after round 1), so the guard should trip after exactly 2 + // rounds, not 1. + provider := &scriptedProvider{responses: []*llm.Response{round("call_1", 10), round("call_2", 10)}} + + h := history.New() + a := New(provider, reg, h, "", 1024, Guards{TokenBudget: 15}, 0, AutoApprove) + + got, err := a.Run(context.Background(), "keep going forever") + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if provider.calls != 2 { + t.Errorf("provider was called %d times, want exactly 2", provider.calls) + } + if !strings.Contains(got, "token budget") { + t.Errorf("Run() = %q, want an explanation mentioning the token budget", got) + } + if err := h.Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } + assertFinalMessageIsExplanation(t, h) +} + +// TestRunStopsAtWallClockGuard is 4.2's AC for the wall-clock guard: a tool +// slower than the configured budget trips it after a single round. +func TestRunStopsAtWallClockGuard(t *testing.T) { + reg := tools.NewRegistry() + slow := &delayedTool{name: "slow_tool", delay: 30 * time.Millisecond, result: "ok", mu: &sync.Mutex{}, order: &[]string{}} + if err := reg.Register(slow); err != nil { + t.Fatalf("Register() error = %v", err) + } + + provider := &scriptedProvider{responses: []*llm.Response{ + { + Message: llm.Message{ + Role: llm.RoleAssistant, + Blocks: []llm.Block{llm.ToolUse{ID: "call_1", Name: "slow_tool", Input: json.RawMessage(`{}`)}}, + }, + StopReason: "tool_use", + }, + }} + + h := history.New() + a := New(provider, reg, h, "", 1024, Guards{WallClock: 15 * time.Millisecond}, 0, AutoApprove) + + got, err := a.Run(context.Background(), "keep going forever") + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if provider.calls != 1 { + t.Errorf("provider was called %d times, want exactly 1", provider.calls) + } + if !strings.Contains(got, "wall-clock timeout") { + t.Errorf("Run() = %q, want an explanation mentioning the wall-clock timeout", got) + } + if err := h.Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } + assertFinalMessageIsExplanation(t, h) +} + +// TestRunStopsAtRepetitionGuard is 4.2's AC for repetition detection: +// CLAUDE.md's fixed threshold of 3 identical consecutive tool calls. +func TestRunStopsAtRepetitionGuard(t *testing.T) { + reg := tools.NewRegistry() + if err := reg.Register(echoTool{}); err != nil { + t.Fatalf("Register() error = %v", err) + } + + identical := func(id string) *llm.Response { + return &llm.Response{ + Message: llm.Message{ + Role: llm.RoleAssistant, + Blocks: []llm.Block{llm.ToolUse{ID: id, Name: "echo_tool", Input: json.RawMessage(`{"n":1}`)}}, + }, + StopReason: "tool_use", + } + } + // Different IDs each round (as a real provider would produce) but the + // same name+arguments — that's what "identical" means here, not ID + // equality. + provider := &scriptedProvider{responses: []*llm.Response{identical("call_1"), identical("call_2"), identical("call_3")}} + + h := history.New() + a := New(provider, reg, h, "", 1024, Guards{}, 0, AutoApprove) + + got, err := a.Run(context.Background(), "keep going forever") + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if provider.calls != 3 { + t.Errorf("provider was called %d times, want exactly 3", provider.calls) + } + if !strings.Contains(got, "repeated 3 times") { + t.Errorf("Run() = %q, want an explanation mentioning the repetition", got) + } + if err := h.Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } + assertFinalMessageIsExplanation(t, h) +} + +// TestRunDoesNotTripRepetitionGuardOnAlternatingCalls guards against a +// too-eager repetition counter (e.g. one that doesn't reset on a change). +func TestRunDoesNotTripRepetitionGuardOnAlternatingCalls(t *testing.T) { + reg := tools.NewRegistry() + if err := reg.Register(echoTool{}); err != nil { + t.Fatalf("Register() error = %v", err) + } + + callWith := func(id string, n int) *llm.Response { + return &llm.Response{ + Message: llm.Message{ + Role: llm.RoleAssistant, + Blocks: []llm.Block{llm.ToolUse{ID: id, Name: "echo_tool", Input: json.RawMessage(fmt.Sprintf(`{"n":%d}`, n))}}, + }, + StopReason: "tool_use", + } + } + done := &llm.Response{ + Message: llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.Text{Text: "finished"}}}, + StopReason: "end_turn", + } + provider := &scriptedProvider{responses: []*llm.Response{ + callWith("call_a", 1), callWith("call_b", 2), callWith("call_c", 1), callWith("call_d", 2), done, + }} + + h := history.New() + a := New(provider, reg, h, "", 1024, Guards{}, 0, AutoApprove) + + got, err := a.Run(context.Background(), "alternate") + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if got != "finished" { + t.Errorf("Run() = %q, want %q (repetition guard must not trip on alternating calls)", got, "finished") + } + if provider.calls != 5 { + t.Errorf("provider was called %d times, want 5", provider.calls) + } +} + +// panicTool always panics — it stands in for a broken tool implementation. +type panicTool struct{} + +func (panicTool) Name() string { return "panic_tool" } +func (panicTool) Description() string { return "a tool that always panics" } +func (panicTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) } +func (panicTool) Run(_ context.Context, _ json.RawMessage) (string, error) { + panic("boom") +} + +var _ tools.Tool = panicTool{} + +// TestRunRecoversFromPanickingTool is 4.3's AC for the panicking-tool case: +// a valid paired result, and the loop continues. +func TestRunRecoversFromPanickingTool(t *testing.T) { + reg := tools.NewRegistry() + if err := reg.Register(panicTool{}); err != nil { + t.Fatalf("Register() error = %v", err) + } + + provider := &scriptedProvider{responses: []*llm.Response{ + toolCallResponseNamed("call_1", "panic_tool"), + textResponse("recovered"), + }} + + h := history.New() + a := New(provider, reg, h, "", 1024, Guards{}, 0, AutoApprove) + + got, err := a.Run(context.Background(), "call the panicking tool") + if err != nil { + t.Fatalf("Run() error = %v, want the loop to survive the panic and continue", err) + } + if got != "recovered" { + t.Errorf("Run() = %q, want %q", got, "recovered") + } + if err := h.Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } + + result := h.Snapshot()[2].Blocks[0].(llm.ToolResult) + if !result.IsError { + t.Error("ToolResult.IsError = false, want true for a panicking tool") + } + if !strings.Contains(result.Content, "panicked") { + t.Errorf("ToolResult.Content = %q, want it to mention the panic", result.Content) + } +} + +// hangingTool ignores ctx entirely and sleeps far longer than any test +// timeout — a genuine hang, not just a slow call, so the only thing that +// can end it is runTool giving up and moving on. +type hangingTool struct{} + +func (hangingTool) Name() string { return "hanging_tool" } +func (hangingTool) Description() string { return "a tool that ignores ctx and never returns in time" } +func (hangingTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) } +func (hangingTool) Run(_ context.Context, _ json.RawMessage) (string, error) { + time.Sleep(200 * time.Millisecond) + return "too late", nil +} + +var _ tools.Tool = hangingTool{} + +// TestRunTimesOutHangingTool is 4.3's AC for the hanging-tool case: a valid +// paired result, and the loop continues, without waiting for the tool. +func TestRunTimesOutHangingTool(t *testing.T) { + reg := tools.NewRegistry() + if err := reg.Register(hangingTool{}); err != nil { + t.Fatalf("Register() error = %v", err) + } + + provider := &scriptedProvider{responses: []*llm.Response{ + toolCallResponseNamed("call_1", "hanging_tool"), + textResponse("gave up waiting"), + }} + + h := history.New() + a := New(provider, reg, h, "", 1024, Guards{}, 20*time.Millisecond, AutoApprove) + + start := time.Now() + got, err := a.Run(context.Background(), "call the hanging tool") + elapsed := time.Since(start) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if got != "gave up waiting" { + t.Errorf("Run() = %q, want %q", got, "gave up waiting") + } + if elapsed > 100*time.Millisecond { + t.Errorf("Run() took %s, want well under the tool's 200ms sleep (the per-tool timeout must not block the loop)", elapsed) + } + if err := h.Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } + + result := h.Snapshot()[2].Blocks[0].(llm.ToolResult) + if !result.IsError { + t.Error("ToolResult.IsError = false, want true for a timed-out tool") + } + if !strings.Contains(result.Content, "deadline exceeded") { + t.Errorf("ToolResult.Content = %q, want it to mention the timeout", result.Content) + } +} + +// blockingTool is well-behaved: it blocks on ctx and returns as soon as ctx +// is cancelled, closing started first so the test knows the tool is +// actually in flight before it cancels. +type blockingTool struct { + started chan struct{} +} + +func (t *blockingTool) Name() string { return "blocking_tool" } +func (t *blockingTool) Description() string { return "a tool that blocks until ctx is done" } +func (t *blockingTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) } +func (t *blockingTool) Run(ctx context.Context, _ json.RawMessage) (string, error) { + close(t.started) + <-ctx.Done() + return "", ctx.Err() +} + +var _ tools.Tool = (*blockingTool)(nil) + +// TestRunHandlesCancellationMidTool is 4.3's AC for the cancelled-tool +// case: a valid paired result, and the loop exits cleanly (Run returns the +// cancellation error rather than hanging or corrupting history). +func TestRunHandlesCancellationMidTool(t *testing.T) { + reg := tools.NewRegistry() + bt := &blockingTool{started: make(chan struct{})} + if err := reg.Register(bt); err != nil { + t.Fatalf("Register() error = %v", err) + } + + provider := &scriptedProvider{responses: []*llm.Response{toolCallResponseNamed("call_1", "blocking_tool")}} + + h := history.New() + a := New(provider, reg, h, "", 1024, Guards{}, 0, AutoApprove) + + ctx, cancel := context.WithCancel(context.Background()) + go func() { + <-bt.started + cancel() + }() + + _, err := a.Run(ctx, "call the blocking tool") + if !errors.Is(err, context.Canceled) { + t.Fatalf("Run() error = %v, want it to wrap context.Canceled", err) + } + if err := h.Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } + + result := h.Snapshot()[2].Blocks[0].(llm.ToolResult) + if !result.IsError { + t.Error("ToolResult.IsError = false, want true for a cancelled tool") + } +} + +// approvalTool implements tools.ApprovalRequired; ran records whether Run +// was actually invoked, so tests can prove a denial skips it. +type approvalTool struct { + ran bool +} + +func (t *approvalTool) Name() string { return "sensitive_tool" } +func (t *approvalTool) Description() string { return "a tool that requires approval" } +func (t *approvalTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) } +func (t *approvalTool) NeedsApproval() bool { return true } +func (t *approvalTool) Run(_ context.Context, _ json.RawMessage) (string, error) { + t.ran = true + return "done", nil +} + +var _ tools.ApprovalRequired = (*approvalTool)(nil) + +type fakeApprover struct { + approve bool + err error + asked bool +} + +func (f *fakeApprover) Approve(_ context.Context, _ string, _ json.RawMessage, _ string) (bool, error) { + f.asked = true + return f.approve, f.err +} + +// TestRunDeniedApprovalSkipsToolAndContinues is 5.3's AC: a denied approval +// returns an error result and the loop continues. +func TestRunDeniedApprovalSkipsToolAndContinues(t *testing.T) { + reg := tools.NewRegistry() + tool := &approvalTool{} + if err := reg.Register(tool); err != nil { + t.Fatalf("Register() error = %v", err) + } + + provider := &scriptedProvider{responses: []*llm.Response{ + toolCallResponseNamed("call_1", "sensitive_tool"), + textResponse("moved on"), + }} + + h := history.New() + approver := &fakeApprover{approve: false} + a := New(provider, reg, h, "", 1024, Guards{}, 0, approver) + + got, err := a.Run(context.Background(), "do the sensitive thing") + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if got != "moved on" { + t.Errorf("Run() = %q, want %q", got, "moved on") + } + if !approver.asked { + t.Error("approver was never consulted") + } + if tool.ran { + t.Error("tool.Run was called despite a denied approval") + } + if err := h.Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } + + result := h.Snapshot()[2].Blocks[0].(llm.ToolResult) + if !result.IsError { + t.Error("ToolResult.IsError = false, want true for a denied approval") + } +} + +func TestRunApprovedCallRunsTheTool(t *testing.T) { + reg := tools.NewRegistry() + tool := &approvalTool{} + if err := reg.Register(tool); err != nil { + t.Fatalf("Register() error = %v", err) + } + + provider := &scriptedProvider{responses: []*llm.Response{ + toolCallResponseNamed("call_1", "sensitive_tool"), + textResponse("finished"), + }} + + h := history.New() + approver := &fakeApprover{approve: true} + a := New(provider, reg, h, "", 1024, Guards{}, 0, approver) + + if _, err := a.Run(context.Background(), "do the sensitive thing"); err != nil { + t.Fatalf("Run() error = %v", err) + } + if !approver.asked { + t.Error("approver was never consulted") + } + if !tool.ran { + t.Error("tool.Run was never called despite an approved approval") + } +} + +func TestRunSkipsApprovalForToolsThatDontNeedIt(t *testing.T) { + reg := tools.NewRegistry() + if err := reg.Register(echoTool{}); err != nil { + t.Fatalf("Register() error = %v", err) + } + + provider := &scriptedProvider{responses: []*llm.Response{toolCallResponse("call_1"), textResponse("done")}} + + h := history.New() + approver := &fakeApprover{approve: false} + a := New(provider, reg, h, "", 1024, Guards{}, 0, approver) + + if _, err := a.Run(context.Background(), "go"); err != nil { + t.Fatalf("Run() error = %v", err) + } + if approver.asked { + t.Error("approver was consulted for a tool that doesn't implement ApprovalRequired") + } +} + +// TestRunDeniedWriteFileLeavesFileByteIdentical is 5.4's AC exercised +// end-to-end through the loop's approval path, using the real +// tools.WriteFileTool rather than a fake. +func TestRunDeniedWriteFileLeavesFileByteIdentical(t *testing.T) { + root := t.TempDir() + path := filepath.Join(root, "existing.txt") + if err := os.WriteFile(path, []byte("original content"), 0o600); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + + sandbox, err := tools.NewSandbox(root, nil) + if err != nil { + t.Fatalf("NewSandbox() error = %v", err) + } + writeTool, err := tools.NewWriteFileTool(sandbox) + if err != nil { + t.Fatalf("NewWriteFileTool() error = %v", err) + } + reg := tools.NewRegistry() + if err := reg.Register(writeTool); err != nil { + t.Fatalf("Register() error = %v", err) + } + + writeInput, _ := json.Marshal(tools.WriteFileInput{Path: "existing.txt", Content: "attacker-controlled overwrite"}) + provider := &scriptedProvider{responses: []*llm.Response{ + { + Message: llm.Message{ + Role: llm.RoleAssistant, + Blocks: []llm.Block{llm.ToolUse{ID: "call_1", Name: "write_file", Input: writeInput}}, + }, + StopReason: "tool_use", + }, + textResponse("ok"), + }} + + h := history.New() + a := New(provider, reg, h, "", 1024, Guards{}, 0, &fakeApprover{approve: false}) + + if _, err := a.Run(context.Background(), "overwrite the file"); err != nil { + t.Fatalf("Run() error = %v", err) + } + + got, err := os.ReadFile(path) + if err != nil { + t.Fatalf("ReadFile() error = %v", err) + } + if string(got) != "original content" { + t.Errorf("file contents = %q, want the original untouched after a denied approval", got) + } +} + +// statusRecordingRenderer wraps a ui.Renderer and records every +// ToolStarted/ToolFinished call it receives, so a test can prove +// RunStream actually reports tool status to a Renderer that asks for it. +type statusRecordingRenderer struct { + ui.Renderer + mu sync.Mutex + events []string +} + +func (r *statusRecordingRenderer) ToolStarted(_, name string, _ json.RawMessage) { + r.mu.Lock() + defer r.mu.Unlock() + r.events = append(r.events, "start:"+name) +} + +func (r *statusRecordingRenderer) ToolFinished(_, name, _ string, isError bool) { + r.mu.Lock() + defer r.mu.Unlock() + status := "ok" + if isError { + status = "error" + } + r.events = append(r.events, fmt.Sprintf("finish:%s:%s", name, status)) +} + +var _ ui.ToolStatusReporter = (*statusRecordingRenderer)(nil) + +// TestRunStreamDrivesScriptedConversationWithToolCall mirrors +// TestRunDrivesScriptedConversationWithParallelTools but through RunStream: +// the provider's Stream (not Chat) drives the round, and a ui.PlainRenderer +// does the rendering. The point is proving RunStream reaches the same +// final state as Run — valid history, right answer, right call count. +func TestRunStreamDrivesScriptedConversationWithToolCall(t *testing.T) { + echo := &delayedTool{name: "echo_tool", delay: 0, result: "echoed", mu: &sync.Mutex{}, order: &[]string{}} + reg := tools.NewRegistry() + if err := reg.Register(echo); err != nil { + t.Fatalf("Register() error = %v", err) + } + + provider := &scriptedProvider{responses: []*llm.Response{ + { + Message: llm.Message{ + Role: llm.RoleAssistant, + Blocks: []llm.Block{llm.ToolUse{ID: "call_1", Name: "echo_tool", Input: json.RawMessage(`{}`)}}, + }, + StopReason: "tool_use", + }, + textResponse("done"), + }} + + h := history.New() + a := New(provider, reg, h, "you are a test agent", 1024, Guards{}, 0, AutoApprove) + + var buf bytes.Buffer + renderer := &statusRecordingRenderer{Renderer: ui.PlainRenderer{Out: &buf}} + got, err := a.RunStream(context.Background(), "please help", renderer) + if err != nil { + t.Fatalf("RunStream() error = %v", err) + } + if got != "done" { + t.Errorf("RunStream() = %q, want %q", got, "done") + } + if provider.calls != 2 { + t.Errorf("provider was called %d times, want 2", provider.calls) + } + if err := h.Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } + + renderer.mu.Lock() + events := append([]string(nil), renderer.events...) + renderer.mu.Unlock() + want := []string{"start:echo_tool", "finish:echo_tool:ok"} + if diff := cmp.Diff(want, events); diff != "" { + t.Errorf("tool status events (-want +got):\n%s", diff) + } +} + +// fakeRenderer lets a test script Render's return value directly, to +// exercise RunStream's error path without a real event stream. +type fakeRenderer struct { + resp *llm.Response + err error +} + +func (f *fakeRenderer) Render(context.Context, <-chan llm.Event) (*llm.Response, error) { + return f.resp, f.err +} + +func TestRunStreamPropagatesRenderError(t *testing.T) { + provider := &scriptedProvider{responses: []*llm.Response{textResponse("unused")}} + h := history.New() + a := New(provider, tools.NewRegistry(), h, "", 1024, Guards{}, 0, AutoApprove) + + wantErr := errors.New("render exploded") + _, err := a.RunStream(context.Background(), "hi", &fakeRenderer{err: wantErr}) + if !errors.Is(err, wantErr) { + t.Fatalf("RunStream() error = %v, want it to wrap %v", err, wantErr) + } +} + +func TestRunStreamStopsAtMaxTurnsGuard(t *testing.T) { + provider := &scriptedProvider{responses: []*llm.Response{ + { + Message: llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.ToolUse{ID: "call_1", Name: "missing_tool", Input: json.RawMessage(`{}`)}}}, + StopReason: "tool_use", + }, + { + Message: llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.ToolUse{ID: "call_2", Name: "missing_tool", Input: json.RawMessage(`{}`)}}}, + StopReason: "tool_use", + }, + }} + + h := history.New() + a := New(provider, tools.NewRegistry(), h, "", 1024, Guards{MaxTurns: 2}, 0, AutoApprove) + + var buf bytes.Buffer + got, err := a.RunStream(context.Background(), "hi", ui.PlainRenderer{Out: &buf}) + if err != nil { + t.Fatalf("RunStream() error = %v", err) + } + if !strings.Contains(got, "maximum of 2 turn") { + t.Errorf("RunStream() = %q, want an explanation naming the max-turns guard", got) + } + if err := h.Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } +} + +// TestCumulativeUsageMatchesSumOfPerTurnUsage is 8.1's AC: cumulative +// totals match the sum of per-turn Usage in a scripted run. The scripted +// run includes a two-round turn (a tool call, then a final reply) so a +// per-turn total that only counted the last round's Usage — not summed +// across the turn's rounds — would also be caught here. +func TestCumulativeUsageMatchesSumOfPerTurnUsage(t *testing.T) { + tool := &delayedTool{name: "echo_tool", result: "ok", mu: &sync.Mutex{}, order: &[]string{}} + reg := tools.NewRegistry() + if err := reg.Register(tool); err != nil { + t.Fatalf("Register() error = %v", err) + } + + provider := &scriptedProvider{responses: []*llm.Response{ + { + Message: llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.ToolUse{ID: "call_1", Name: "echo_tool", Input: json.RawMessage(`{}`)}}}, + StopReason: "tool_use", + Usage: llm.Usage{InputTokens: 10, OutputTokens: 5}, + }, + { + Message: llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.Text{Text: "done with turn one"}}}, + StopReason: "end_turn", + Usage: llm.Usage{InputTokens: 20, OutputTokens: 8}, + }, + { + Message: llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.Text{Text: "done with turn two"}}}, + StopReason: "end_turn", + Usage: llm.Usage{InputTokens: 30, OutputTokens: 12}, + }, + }} + + h := history.New() + a := New(provider, reg, h, "", 1024, Guards{}, 0, AutoApprove) + + if _, err := a.Run(context.Background(), "first turn"); err != nil { + t.Fatalf("Run() 1 error = %v", err) + } + if _, err := a.Run(context.Background(), "second turn"); err != nil { + t.Fatalf("Run() 2 error = %v", err) + } + + turns := a.TurnUsages() + wantTurns := []llm.Usage{ + {InputTokens: 30, OutputTokens: 13}, // turn 1: round1 (10,5) + round2 (20,8) + {InputTokens: 30, OutputTokens: 12}, // turn 2: single round + } + if diff := cmp.Diff(wantTurns, turns); diff != "" { + t.Errorf("TurnUsages() mismatch (-want +got):\n%s", diff) + } + + var wantCumulative llm.Usage + for _, u := range turns { + wantCumulative.InputTokens += u.InputTokens + wantCumulative.OutputTokens += u.OutputTokens + } + if got := a.CumulativeUsage(); got != wantCumulative { + t.Errorf("CumulativeUsage() = %+v, want it to equal the sum of TurnUsages(): %+v", got, wantCumulative) + } +} diff --git a/internal/agent/testdata/golden/summarization_prompt.txt b/internal/agent/testdata/golden/summarization_prompt.txt new file mode 100644 index 0000000..e915e66 --- /dev/null +++ b/internal/agent/testdata/golden/summarization_prompt.txt @@ -0,0 +1,7 @@ +Summarize the earlier part of this conversation between you (the assistant) and a user, so the summary can replace the original messages in your context. Preserve any facts, decisions, file paths, commands, or conclusions that later turns might depend on. Do not include pleasantries or filler. Write the summary as plain prose, not a transcript. + +Conversation to summarize: +user: read config.go and tell me the default port +assistant: [called tool read_file with {"path":"config.go"}] +user: [tool result (ok): const defaultPort = 8080] +assistant: The default port is 8080. \ No newline at end of file diff --git a/internal/config/config.go b/internal/config/config.go index 2c299fc..92ffc69 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -25,6 +25,17 @@ const ( ToolsPrompted ToolsMode = "prompted" ) +// LogLevel selects the minimum severity slog emits. +type LogLevel string + +// Supported log levels. +const ( + LogLevelDebug LogLevel = "debug" + LogLevelInfo LogLevel = "info" + LogLevelWarn LogLevel = "warn" + LogLevelError LogLevel = "error" +) + // ErrUnknownProvider is returned when --provider names a backend pico code // does not implement. var ErrUnknownProvider = errors.New("unknown provider") @@ -33,21 +44,35 @@ var ErrUnknownProvider = errors.New("unknown provider") // native or prompted. var ErrUnknownToolsMode = errors.New("unknown tools mode") +// ErrUnknownLogLevel is returned when --log-level names a level other than +// debug, info, warn, or error. +var ErrUnknownLogLevel = errors.New("unknown log level") + +// ErrModelRequired is returned when --model is empty. Nothing downstream +// defaults it, so an empty model would otherwise reach a provider adapter +// and fail as an opaque transport error (a 404 from Anthropic, a connection +// attempt to Ollama) instead of a clear message at startup. +var ErrModelRequired = errors.New("model is required") + // Flags carries the values collected from the chat subcommand's flags, // already resolved against any environment fallback (e.g. PICO_CODE_PROVIDER) // by the caller. Load only validates and fills in credentials. type Flags struct { - Provider string - Model string - MaxTurns int - TokenBudget int - Workspace string - Yes bool - Tools string - Stream bool - TUI bool - LogLevel string - NumCtx int + Provider string + Model string + MaxTurns int + TokenBudget int + Workspace string + Yes bool + Tools string + Stream bool + TUI bool + LogLevel string + NumCtx int + AllowWrite bool + Session string + AllowCommands []string + ContextWindow int } // Config is the fully resolved, validated configuration for a chat session. @@ -61,13 +86,33 @@ type Config struct { Tools ToolsMode Stream bool TUI bool - LogLevel string + LogLevel LogLevel // NumCtx is the Ollama adapter's context window size (num_ctx). Ollama // silently truncates context if this is left unset, so it always has an // explicit value rather than an optional one the adapter might omit. NumCtx int + // AllowWrite gates registering write_file: false by default, since it's + // the only built-in tool that mutates the workspace. + AllowWrite bool + + // Session names a session to resume (if it already exists) or start + // (if it doesn't); empty means no session persistence. + Session string + + // AllowCommands is the binary allowlist run_command is registered + // with; empty means the tool isn't registered at all, since an empty + // allowlist would make it present but unconditionally useless. + AllowCommands []string + + // ContextWindow is the token threshold history compaction (8.2) + // measures estimated usage against for providers other than Ollama, + // which uses NumCtx instead — Ollama's context window is already an + // explicit, required value, so compaction reuses it rather than + // tracking a second number that could drift from it. + ContextWindow int + // AnthropicAPIKey and OllamaHost come from the environment only; there // is no flag for either, since committing a secret to a shell history // is worse than typing it once into the environment. @@ -85,6 +130,10 @@ func Load(f Flags, getenv func(string) string) (*Config, error) { return nil, fmt.Errorf("%w: %q (want %q or %q)", ErrUnknownProvider, f.Provider, ProviderAnthropic, ProviderOllama) } + if f.Model == "" { + return nil, fmt.Errorf("%w: pass --model with an identifier for %q (e.g. a Claude model name for anthropic, or a locally pulled model tag for ollama)", ErrModelRequired, provider) + } + tools := ToolsMode(f.Tools) switch tools { case ToolsNative, ToolsPrompted: @@ -92,6 +141,13 @@ func Load(f Flags, getenv func(string) string) (*Config, error) { return nil, fmt.Errorf("%w: %q (want %q or %q)", ErrUnknownToolsMode, f.Tools, ToolsNative, ToolsPrompted) } + logLevel := LogLevel(f.LogLevel) + switch logLevel { + case LogLevelDebug, LogLevelInfo, LogLevelWarn, LogLevelError: + default: + return nil, fmt.Errorf("%w: %q (want %q, %q, %q, or %q)", ErrUnknownLogLevel, f.LogLevel, LogLevelDebug, LogLevelInfo, LogLevelWarn, LogLevelError) + } + return &Config{ Provider: provider, Model: f.Model, @@ -102,8 +158,12 @@ func Load(f Flags, getenv func(string) string) (*Config, error) { Tools: tools, Stream: f.Stream, TUI: f.TUI, - LogLevel: f.LogLevel, + LogLevel: logLevel, NumCtx: f.NumCtx, + AllowWrite: f.AllowWrite, + Session: f.Session, + AllowCommands: f.AllowCommands, + ContextWindow: f.ContextWindow, AnthropicAPIKey: getenv("ANTHROPIC_API_KEY"), OllamaHost: getenv("OLLAMA_HOST"), }, nil diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 9764cc1..6086a77 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -96,6 +96,16 @@ func TestLoad(t *testing.T) { getenv: noEnv, wantErr: ErrUnknownProvider, }, + { + name: "empty model", + flags: func() Flags { + f := baseFlags() + f.Model = "" + return f + }(), + getenv: noEnv, + wantErr: ErrModelRequired, + }, { name: "unknown tools mode", flags: func() Flags { @@ -106,6 +116,16 @@ func TestLoad(t *testing.T) { getenv: noEnv, wantErr: ErrUnknownToolsMode, }, + { + name: "unknown log level", + flags: func() Flags { + f := baseFlags() + f.LogLevel = "verbose" + return f + }(), + getenv: noEnv, + wantErr: ErrUnknownLogLevel, + }, } for _, tt := range tests { diff --git a/internal/history/compact.go b/internal/history/compact.go new file mode 100644 index 0000000..a955a9e --- /dev/null +++ b/internal/history/compact.go @@ -0,0 +1,100 @@ +package history + +import ( + "fmt" + + "github.com/reno/pico-code/internal/llm" +) + +// CompactionBoundary returns the index of the first message that must +// survive compaction verbatim: the start of the keepTurns most recent user +// turns. A "turn" starts at a user message carrying a Text block — the +// message the agent loop appends for actual user input, as opposed to the +// user messages carrying only ToolResult blocks that continue the same +// turn. Because a genuine turn can only start once every ToolUse earlier in +// the conversation has already been answered (CLAUDE.md invariant 3), any +// boundary this returns is automatically safe to cut at: it can never land +// between a ToolUse and its ToolResult. +// +// It returns 0 (nothing to compact) if there are keepTurns or fewer turns. +func CompactionBoundary(messages []llm.Message, keepTurns int) int { + if keepTurns <= 0 { + return 0 + } + found := 0 + for i := len(messages) - 1; i >= 0; i-- { + if isTurnStart(messages[i]) { + found++ + if found == keepTurns { + return i + } + } + } + return 0 +} + +func isTurnStart(m llm.Message) bool { + if m.Role != llm.RoleUser { + return false + } + for _, b := range m.Blocks { + if _, ok := b.(llm.Text); ok { + return true + } + } + return false +} + +// Compact replaces every message before the boundary that keeps the last +// keepTurns user turns verbatim with a single synthetic assistant message +// containing summary. The synthetic message is Role assistant (not user) +// so it alternates cleanly with the kept messages, which always start at a +// genuine user turn. It refuses to leave history in a state that fails +// Validate() (CLAUDE.md invariant 3) rather than commit one — compaction is +// the one operation most likely to violate that invariant, since it +// discards history wholesale; the boundary construction above already +// makes that unreachable in practice, but this is the backstop. +func (h *History) Compact(keepTurns int, summary string) error { + boundary := CompactionBoundary(h.messages, keepTurns) + if boundary <= 0 { + return nil + } + + compacted := make([]llm.Message, 0, len(h.messages)-boundary+1) + compacted = append(compacted, llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.Text{Text: summary}}}) + compacted = append(compacted, h.messages[boundary:]...) + + trial := &History{messages: compacted} + if err := trial.Validate(); err != nil { + return fmt.Errorf("history: refusing to compact into an invalid state: %w", err) + } + h.messages = compacted + return nil +} + +// EstimateTokens approximates messages' total size in tokens from +// character length (~4 characters per token, the same rough rule of thumb +// used for Ollama's own estimate in internal/llm/ollama) — good enough to +// decide when to compact, not for billing. +func EstimateTokens(messages []llm.Message) int { + chars := 0 + for _, m := range messages { + for _, b := range m.Blocks { + switch v := b.(type) { + case llm.Text: + chars += len(v.Text) + case llm.ToolUse: + chars += len(v.Input) + case llm.ToolResult: + chars += len(v.Content) + } + } + } + if chars == 0 { + return 0 + } + if n := chars / 4; n > 0 { + return n + } + return 1 +} diff --git a/internal/history/compact_test.go b/internal/history/compact_test.go new file mode 100644 index 0000000..f76cd80 --- /dev/null +++ b/internal/history/compact_test.go @@ -0,0 +1,120 @@ +package history_test + +import ( + "encoding/json" + "fmt" + "testing" + + "github.com/google/go-cmp/cmp" + + "github.com/reno/pico-code/internal/history" + "github.com/reno/pico-code/internal/llm" +) + +// buildTurns returns n user/assistant turns; every third turn also +// includes a tool call and its result, so a compaction boundary has +// ToolUse/ToolResult pairs to potentially (and must not) split. +func buildTurns(n int) []llm.Message { + var msgs []llm.Message + for i := 0; i < n; i++ { + msgs = append(msgs, textMsg(llm.RoleUser, fmt.Sprintf("turn %d question", i))) + if i%3 == 0 { + id := fmt.Sprintf("call_%d", i) + msgs = append(msgs, toolUseMsg(llm.ToolUse{ID: id, Name: "read_file", Input: json.RawMessage(`{}`)})) + msgs = append(msgs, toolResultMsg(llm.ToolResult{ToolUseID: id, Content: "result"})) + msgs = append(msgs, textMsg(llm.RoleAssistant, fmt.Sprintf("turn %d answer after tool", i))) + } else { + msgs = append(msgs, textMsg(llm.RoleAssistant, fmt.Sprintf("turn %d answer", i))) + } + } + return msgs +} + +func TestCompactionBoundaryKeepsExactTurnCount(t *testing.T) { + msgs := buildTurns(10) + boundary := history.CompactionBoundary(msgs, 3) + + if !cmp.Equal(msgs[boundary], textMsg(llm.RoleUser, "turn 7 question")) { + t.Fatalf("message at boundary %d = %+v, want the start of turn 7 (the 3 most recent of 10 turns: 7, 8, 9)", boundary, msgs[boundary]) + } +} + +func TestCompactionBoundaryZeroWhenNotEnoughTurns(t *testing.T) { + msgs := buildTurns(3) + if boundary := history.CompactionBoundary(msgs, 5); boundary != 0 { + t.Errorf("CompactionBoundary() = %d, want 0 when there are fewer turns than keepTurns", boundary) + } +} + +// TestCompact40TurnHistoryWithToolCallsStillValidates is 8.2's AC: +// compacting a 40-turn history with tool calls still passes Validate(), +// and the boundary never splits a ToolUse/ToolResult pair. +func TestCompact40TurnHistoryWithToolCallsStillValidates(t *testing.T) { + h := history.New() + for _, m := range buildTurns(40) { + h.Append(m) + } + if err := h.Validate(); err != nil { + t.Fatalf("Validate() on the uncompacted history error = %v", err) + } + + if err := h.Compact(5, "summary of the first 35 turns"); err != nil { + t.Fatalf("Compact() error = %v", err) + } + if err := h.Validate(); err != nil { + t.Fatalf("Validate() after Compact() error = %v", err) + } + + msgs := h.Snapshot() + if len(msgs) == 0 || msgs[0].Role != llm.RoleAssistant { + t.Fatalf("messages[0] = %+v, want the synthetic assistant summary", msgs[0]) + } + text, ok := msgs[0].Blocks[0].(llm.Text) + if !ok || text.Text != "summary of the first 35 turns" { + t.Errorf("messages[0] text = %+v, want the summary", msgs[0].Blocks[0]) + } +} + +func TestCompactPreservesLastNTurnsVerbatim(t *testing.T) { + all := buildTurns(10) + h := history.New() + for _, m := range all { + h.Append(m) + } + + if err := h.Compact(3, "elided"); err != nil { + t.Fatalf("Compact() error = %v", err) + } + + got := h.Snapshot() + boundary := history.CompactionBoundary(all, 3) + wantTail := all[boundary:] + gotTail := got[1:] // [0] is the synthetic summary message + if diff := cmp.Diff(wantTail, gotTail); diff != "" { + t.Errorf("kept turns mismatch (-want +got):\n%s", diff) + } +} + +func TestCompactNoopsWhenNothingToCompact(t *testing.T) { + h := history.New() + for _, m := range buildTurns(2) { + h.Append(m) + } + before := h.Snapshot() + + if err := h.Compact(5, "unused"); err != nil { + t.Fatalf("Compact() error = %v", err) + } + if diff := cmp.Diff(before, h.Snapshot()); diff != "" { + t.Errorf("Compact() changed history when there was nothing to compact (-before +after):\n%s", diff) + } +} + +func TestEstimateTokens(t *testing.T) { + if got := history.EstimateTokens(nil); got != 0 { + t.Errorf("EstimateTokens(nil) = %d, want 0", got) + } + if got := history.EstimateTokens(buildTurns(1)); got <= 0 { + t.Errorf("EstimateTokens(1 turn) = %d, want a positive estimate", got) + } +} diff --git a/internal/history/history.go b/internal/history/history.go index f332752..935d4f4 100644 --- a/internal/history/history.go +++ b/internal/history/history.go @@ -83,6 +83,34 @@ func (h *History) Save(path string) error { // returning, so a corrupted or hand-edited file can never be replayed in a // violating state. func Load(path string) (*History, error) { + msgs, err := loadMessages(path) + if err != nil { + return nil, err + } + return &History{messages: msgs}, nil +} + +// LoadInto replaces h's messages with the ones stored at path, in place — +// unlike Load, which returns a new History — so callers holding a shared +// *History (e.g. an already-constructed Agent) can point it at a different +// session without reconstructing anything downstream. Like Load, it +// validates before committing: on error h is left completely untouched, +// never partially loaded. +func (h *History) LoadInto(path string) error { + msgs, err := loadMessages(path) + if err != nil { + return err + } + h.messages = msgs + return nil +} + +// Reset clears h back to empty, in place. +func (h *History) Reset() { + h.messages = nil +} + +func loadMessages(path string) ([]llm.Message, error) { data, err := os.ReadFile(path) if err != nil { return nil, fmt.Errorf("history: read %s: %w", path, err) @@ -95,7 +123,7 @@ func Load(path string) (*History, error) { if err := h.Validate(); err != nil { return nil, fmt.Errorf("history: loaded invalid history from %s: %w", path, err) } - return h, nil + return msgs, nil } func toolUseIDs(m llm.Message) []string { diff --git a/internal/history/history_test.go b/internal/history/history_test.go index ca8e98d..0b3b1ca 100644 --- a/internal/history/history_test.go +++ b/internal/history/history_test.go @@ -195,3 +195,61 @@ func TestLoadRejectsCorruptFile(t *testing.T) { t.Fatal("Load() = nil, want an error for a corrupt file") } } + +// TestLoadIntoRoundTrip is 8.3's resume AC in miniature: a session saved, +// then loaded into a different History via LoadInto, has the full prior +// context. +func TestLoadIntoRoundTrip(t *testing.T) { + saved := history.New() + saved.Append(textMsg(llm.RoleUser, "hello")) + saved.Append(textMsg(llm.RoleAssistant, "hi there")) + path := filepath.Join(t.TempDir(), "session.json") + if err := saved.Save(path); err != nil { + t.Fatalf("Save() error = %v", err) + } + + h := history.New() + h.Append(textMsg(llm.RoleUser, "this should be replaced")) + if err := h.LoadInto(path); err != nil { + t.Fatalf("LoadInto() error = %v", err) + } + if diff := cmp.Diff(saved.Snapshot(), h.Snapshot()); diff != "" { + t.Errorf("LoadInto() mismatch (-want +got):\n%s", diff) + } +} + +// TestLoadIntoRejectsCorruptFileWithoutMutatingH is 8.3's AC: a corrupt +// session file fails with a readable error rather than a panic — and here, +// specifically, without leaving the caller's History half-replaced. +func TestLoadIntoRejectsCorruptFileWithoutMutatingH(t *testing.T) { + path := filepath.Join(t.TempDir(), "session.json") + if err := os.WriteFile(path, []byte("not json"), 0o600); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + + h := history.New() + h.Append(textMsg(llm.RoleUser, "keep me")) + before := h.Snapshot() + + if err := h.LoadInto(path); err == nil { + t.Fatal("LoadInto() = nil, want an error for a corrupt file") + } + if diff := cmp.Diff(before, h.Snapshot()); diff != "" { + t.Errorf("h was mutated despite LoadInto() failing (-before +after):\n%s", diff) + } +} + +func TestReset(t *testing.T) { + h := history.New() + h.Append(textMsg(llm.RoleUser, "hello")) + h.Append(textMsg(llm.RoleAssistant, "hi")) + + h.Reset() + + if got := h.Snapshot(); len(got) != 0 { + t.Errorf("Snapshot() after Reset() = %v, want empty", got) + } + if err := h.Validate(); err != nil { + t.Errorf("Validate() after Reset() error = %v", err) + } +} diff --git a/internal/llm/anthropic/conformance_test.go b/internal/llm/anthropic/conformance_test.go new file mode 100644 index 0000000..6a5cd72 --- /dev/null +++ b/internal/llm/anthropic/conformance_test.go @@ -0,0 +1,126 @@ +package anthropic + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/reno/pico-code/internal/llm" + "github.com/reno/pico-code/internal/llm/conformance" +) + +// TestConformance is 6.4's AC: the shared conformance suite passes for the +// Anthropic adapter unchanged. +func TestConformance(t *testing.T) { + conformance.Run(t, conformance.Harness{ + Name: "anthropic", + NewProvider: func(baseURL string) llm.Provider { + return testProvider(baseURL) + }, + Server: conformanceServer, + WantText: "hello", + WantToolUse: llm.ToolUse{Name: "get_weather", Input: json.RawMessage(`{"location":"Paris"}`)}, + }) +} + +func conformanceServer(t *testing.T, scenario conformance.Scenario) *httptest.Server { + t.Helper() + switch scenario { + case conformance.ScenarioSimpleText: + return jsonServer(`{ + "id": "msg_ct1", "type": "message", "role": "assistant", "model": "claude-test-model", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 5, "output_tokens": 2} + }`) + case conformance.ScenarioToolCall: + return jsonServer(`{ + "id": "msg_ct2", "type": "message", "role": "assistant", "model": "claude-test-model", + "content": [{"type": "tool_use", "id": "toolu_ct1", "name": "get_weather", "input": {"location":"Paris"}}], + "stop_reason": "tool_use", + "usage": {"input_tokens": 8, "output_tokens": 4} + }`) + case conformance.ScenarioSimpleTextStream: + return sseServer( + sseEvent("message_start", map[string]any{"type": "message_start", "message": map[string]any{ + "id": "msg_ct1", "type": "message", "role": "assistant", "model": "claude-test-model", + "content": []any{}, "usage": map[string]any{"input_tokens": 5, "output_tokens": 1}, + }}), + sseEvent("content_block_start", map[string]any{"type": "content_block_start", "index": 0, + "content_block": map[string]any{"type": "text", "text": ""}}), + sseEvent("content_block_delta", map[string]any{"type": "content_block_delta", "index": 0, + "delta": map[string]any{"type": "text_delta", "text": "hello"}}), + sseEvent("content_block_stop", map[string]any{"type": "content_block_stop", "index": 0}), + sseEvent("message_delta", map[string]any{"type": "message_delta", + "delta": map[string]any{"stop_reason": "end_turn"}, "usage": map[string]any{"output_tokens": 2}}), + sseEvent("message_stop", map[string]any{"type": "message_stop"}), + ) + case conformance.ScenarioToolCallStream: + return sseServer( + sseEvent("message_start", map[string]any{"type": "message_start", "message": map[string]any{ + "id": "msg_ct2", "type": "message", "role": "assistant", "model": "claude-test-model", + "content": []any{}, "usage": map[string]any{"input_tokens": 8, "output_tokens": 1}, + }}), + sseEvent("content_block_start", map[string]any{"type": "content_block_start", "index": 0, + "content_block": map[string]any{"type": "tool_use", "id": "toolu_ct1", "name": "get_weather", "input": map[string]any{}}}), + sseEvent("content_block_delta", map[string]any{"type": "content_block_delta", "index": 0, + "delta": map[string]any{"type": "input_json_delta", "partial_json": `{"location":"Paris"}`}}), + sseEvent("content_block_stop", map[string]any{"type": "content_block_stop", "index": 0}), + sseEvent("message_delta", map[string]any{"type": "message_delta", + "delta": map[string]any{"stop_reason": "tool_use"}, "usage": map[string]any{"output_tokens": 4}}), + sseEvent("message_stop", map[string]any{"type": "message_stop"}), + ) + case conformance.ScenarioServerError: + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"type":"error","error":{"type":"authentication_error","message":"bad key"}}`)) + })) + case conformance.ScenarioHang: + // Bounded, not truly infinite: a cancelled client doesn't always + // close its TCP connection quickly enough for the server to + // observe via r.Context(), which would otherwise hang + // httptest.Server.Close() waiting for this handler to return. What + // this scenario actually tests is the client (our provider) + // returning promptly on its own ctx cancellation, not the + // server's. + return httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { + select { + case <-r.Context().Done(): + case <-time.After(150 * time.Millisecond): + } + })) + default: + t.Fatalf("conformanceServer: unhandled scenario %v", scenario) + return nil + } +} + +func jsonServer(body string) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(body)) + })) +} + +func sseServer(events ...string) *httptest.Server { + body := strings.Join(events, "") + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(body)) + })) +} + +func sseEvent(name string, data any) string { + b, err := json.Marshal(data) + if err != nil { + panic(err) + } + return fmt.Sprintf("event: %s\ndata: %s\n\n", name, string(b)) +} diff --git a/internal/llm/anthropic/logging_test.go b/internal/llm/anthropic/logging_test.go new file mode 100644 index 0000000..ba82dee --- /dev/null +++ b/internal/llm/anthropic/logging_test.go @@ -0,0 +1,75 @@ +package anthropic + +import ( + "bytes" + "context" + "log/slog" + "net/http" + "net/http/httptest" + "strings" + "testing" + + sdk "github.com/anthropics/anthropic-sdk-go" + "github.com/anthropics/anthropic-sdk-go/option" +) + +func withDefaultLogger(t *testing.T, level slog.Level) *bytes.Buffer { + t.Helper() + var buf bytes.Buffer + prev := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: level}))) + t.Cleanup(func() { slog.SetDefault(prev) }) + return &buf +} + +// TestDebugLoggingDumpsExchangeWithoutCredentials is 9.2's AC: debug output +// of a full turn contains no credentials. +func TestDebugLoggingDumpsExchangeWithoutCredentials(t *testing.T) { + buf := withDefaultLogger(t, slog.LevelDebug) + + const apiKey = "sk-ant-super-secret-test-key" + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(okMessageBody)) + })) + defer srv.Close() + + p := &Provider{ + client: sdk.NewClient(option.WithAPIKey(apiKey), option.WithBaseURL(srv.URL)), + model: "claude-test-model", + apiKey: apiKey, + } + if _, err := p.Chat(context.Background(), testChatRequest()); err != nil { + t.Fatalf("Chat() error = %v", err) + } + + out := buf.String() + if !strings.Contains(out, "anthropic: request") || !strings.Contains(out, "anthropic: response") { + t.Fatalf("debug output missing request/response dump; got:\n%s", out) + } + if strings.Contains(out, apiKey) { + t.Errorf("debug output contains the API key:\n%s", out) + } +} + +// TestDefaultLevelDoesNotDumpRequestsOrResponses is 9.2's other AC: default +// level is silent on the happy path — no request/response body ever +// reaches the log outside --log-level=debug. +func TestDefaultLevelDoesNotDumpRequestsOrResponses(t *testing.T) { + buf := withDefaultLogger(t, slog.LevelInfo) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(okMessageBody)) + })) + defer srv.Close() + + p := testProvider(srv.URL) + if _, err := p.Chat(context.Background(), testChatRequest()); err != nil { + t.Fatalf("Chat() error = %v", err) + } + + if out := buf.String(); strings.Contains(out, "anthropic: request") || strings.Contains(out, "anthropic: response") { + t.Errorf("default (info) level logged a request/response dump, want none:\n%s", out) + } +} diff --git a/internal/llm/anthropic/provider.go b/internal/llm/anthropic/provider.go index 3b0e830..38e79e2 100644 --- a/internal/llm/anthropic/provider.go +++ b/internal/llm/anthropic/provider.go @@ -10,6 +10,7 @@ import ( "github.com/reno/pico-code/internal/config" "github.com/reno/pico-code/internal/llm" + "github.com/reno/pico-code/internal/llm/recordutil" ) func init() { @@ -20,6 +21,9 @@ func init() { type Provider struct { client sdk.Client model string + // apiKey is kept only so --log-level=debug can scrub it by exact value + // in addition to recordutil's generic sk-shaped-literal pattern. + apiKey string } // New constructs a Provider from resolved config. It is registered under @@ -32,6 +36,7 @@ func New(cfg *config.Config) (llm.Provider, error) { return &Provider{ client: sdk.NewClient(option.WithAPIKey(cfg.AnthropicAPIKey)), model: cfg.Model, + apiKey: cfg.AnthropicAPIKey, }, nil } @@ -49,14 +54,11 @@ func (p *Provider) Chat(ctx context.Context, req llm.Request) (*llm.Response, er if err != nil { return nil, fmt.Errorf("anthropic: %w", err) } + recordutil.LogJSON(ctx, "anthropic: request", params, p.apiKey) msg, err := p.client.Messages.New(ctx, params) if err != nil { return nil, mapError(err) } + recordutil.LogJSON(ctx, "anthropic: response", msg, p.apiKey) return fromResponse(msg) } - -// Stream is implemented in phase 6.2. -func (p *Provider) Stream(_ context.Context, _ llm.Request) (<-chan llm.Event, error) { - return nil, errors.New("anthropic: streaming not implemented until phase 6.2") -} diff --git a/internal/llm/anthropic/record_test.go b/internal/llm/anthropic/record_test.go new file mode 100644 index 0000000..52856ca --- /dev/null +++ b/internal/llm/anthropic/record_test.go @@ -0,0 +1,97 @@ +package anthropic + +import ( + "context" + "net/http" + "os" + "path/filepath" + "strings" + "testing" + + sdk "github.com/anthropics/anthropic-sdk-go" + "github.com/anthropics/anthropic-sdk-go/option" + + "github.com/reno/pico-code/internal/llm" + "github.com/reno/pico-code/internal/llm/recordutil" +) + +// TestRecordLiveExchange is 9.1's record half: with RECORD=1 and a real +// ANTHROPIC_API_KEY, it drives one real Chat call and writes the scrubbed +// request/response to testdata/golden/live_exchange, the fixture a +// developer re-recording this adapter's behavior would inspect. It never +// runs in `make test` (no RECORD, no key, no network) — only this +// repository's maintainer running it manually can actually record. +func TestRecordLiveExchange(t *testing.T) { + if !recordutil.Enabled() { + t.Skip("set RECORD=1 (with ANTHROPIC_API_KEY) to re-record this fixture against the live API") + } + apiKey := os.Getenv("ANTHROPIC_API_KEY") + if apiKey == "" { + t.Fatal("RECORD=1 requires ANTHROPIC_API_KEY") + } + + var exchange recordutil.Exchange + recordingClient := &http.Client{Transport: &recordutil.Recorder{ + OnExchange: func(e recordutil.Exchange) { exchange = e }, + }} + client := sdk.NewClient( + option.WithAPIKey(apiKey), + option.WithHTTPClient(recordingClient), + ) + p := &Provider{client: client, model: string(sdk.ModelClaudeHaiku4_5)} + + _, err := p.Chat(context.Background(), llm.Request{ + Messages: []llm.Message{{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "Reply with exactly: ok"}}}}, + MaxTokens: 16, + }) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + + dir := "testdata/golden/live_exchange" + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatalf("MkdirAll() error = %v", err) + } + reqOut := recordutil.Scrub(string(exchange.RequestBody), apiKey) + respOut := recordutil.Scrub(string(exchange.ResponseBody), apiKey) + if err := os.WriteFile(filepath.Join(dir, "request.json"), []byte(reqOut), 0o644); err != nil { + t.Fatalf("WriteFile(request) error = %v", err) + } + if err := os.WriteFile(filepath.Join(dir, "response.json"), []byte(respOut), 0o644); err != nil { + t.Fatalf("WriteFile(response) error = %v", err) + } +} + +// TestGoldenFixturesContainNoSecrets is 9.1's AC: no key ever appears in a +// fixture. It walks every checked-in fixture (recorded or hand-written) +// and greps for an sk-shaped literal and the actual ANTHROPIC_API_KEY value +// if one happens to be set in this environment. +func TestGoldenFixturesContainNoSecrets(t *testing.T) { + assertNoSecretsUnder(t, "testdata/golden") +} + +func assertNoSecretsUnder(t *testing.T, root string) { + t.Helper() + liveKey := os.Getenv("ANTHROPIC_API_KEY") + + err := filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { + if err != nil || d.IsDir() { + return err + } + data, rerr := os.ReadFile(path) + if rerr != nil { + return rerr + } + content := string(data) + if strings.Contains(content, "sk-") { + t.Errorf("%s contains an sk-shaped key literal", path) + } + if liveKey != "" && strings.Contains(content, liveKey) { + t.Errorf("%s contains the live ANTHROPIC_API_KEY value", path) + } + return nil + }) + if err != nil { + t.Fatalf("walking %s: %v", root, err) + } +} diff --git a/internal/llm/anthropic/stream.go b/internal/llm/anthropic/stream.go new file mode 100644 index 0000000..e1be646 --- /dev/null +++ b/internal/llm/anthropic/stream.go @@ -0,0 +1,135 @@ +package anthropic + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + + sdk "github.com/anthropics/anthropic-sdk-go" + + "github.com/reno/pico-code/internal/llm" + "github.com/reno/pico-code/internal/llm/recordutil" +) + +// streamAccumulator tracks the per-block state needed to translate +// Anthropic's SSE events into canonical llm.Events: which content-block +// Index is which tool call's ID, its raw (unparsed) input_json_delta +// fragments, and the running usage/stop-reason that only becomes complete +// once message_delta and message_stop arrive. +type streamAccumulator struct { + toolID map[int64]string + toolBuf map[int64]*bytes.Buffer + + usage llm.Usage + stopReason string +} + +func newStreamAccumulator() *streamAccumulator { + return &streamAccumulator{toolID: map[int64]string{}, toolBuf: map[int64]*bytes.Buffer{}} +} + +// handle translates one SSE event into zero or more canonical events. +func (a *streamAccumulator) handle(ctx context.Context, event sdk.MessageStreamEventUnion) ([]llm.Event, error) { + switch event.Type { + case "message_start": + ms := event.AsMessageStart() + a.usage.InputTokens = int(ms.Message.Usage.InputTokens) + a.usage.OutputTokens = int(ms.Message.Usage.OutputTokens) + return nil, nil + + case "content_block_start": + cbs := event.AsContentBlockStart() + if cbs.ContentBlock.Type != "tool_use" { + return nil, nil + } + a.toolID[cbs.Index] = cbs.ContentBlock.ID + a.toolBuf[cbs.Index] = &bytes.Buffer{} + return []llm.Event{llm.ToolUseStart{ID: cbs.ContentBlock.ID, Name: cbs.ContentBlock.Name}}, nil + + case "content_block_delta": + cbd := event.AsContentBlockDelta() + switch cbd.Delta.Type { + case "text_delta": + return []llm.Event{llm.TextDelta{Text: cbd.Delta.Text}}, nil + case "input_json_delta": + id, ok := a.toolID[cbd.Index] + if !ok { + return nil, fmt.Errorf("anthropic: input_json_delta for unknown block index %d", cbd.Index) + } + a.toolBuf[cbd.Index].WriteString(cbd.Delta.PartialJSON) + return []llm.Event{llm.ToolUseArgsDelta{ID: id, Partial: cbd.Delta.PartialJSON}}, nil + default: + return nil, nil + } + + case "content_block_stop": + cbs := event.AsContentBlockStop() + id, ok := a.toolID[cbs.Index] + if !ok { + return nil, nil + } + raw := a.toolBuf[cbs.Index].Bytes() + if len(raw) == 0 { + raw = []byte("{}") + } + if !json.Valid(raw) { + return nil, fmt.Errorf("anthropic: tool_use %s: accumulated input is not valid JSON: %s", id, raw) + } + return []llm.Event{llm.ToolUseDone{ID: id, Input: json.RawMessage(raw)}}, nil + + case "message_delta": + md := event.AsMessageDelta() + a.stopReason = string(md.Delta.StopReason) + if md.Usage.OutputTokens > 0 { + a.usage.OutputTokens = int(md.Usage.OutputTokens) + } + if md.Usage.InputTokens > 0 { + a.usage.InputTokens = int(md.Usage.InputTokens) + } + return nil, nil + + case "message_stop": + recordutil.LogJSON(ctx, "anthropic: stream response", struct { + StopReason string `json:"stop_reason"` + Usage llm.Usage `json:"usage"` + }{a.stopReason, a.usage}) + return []llm.Event{llm.MessageDone{StopReason: a.stopReason, Usage: a.usage}}, nil + + default: + return nil, nil + } +} + +// Stream sends req and translates the SSE response into canonical events +// via llm.StreamEvents, which owns the close-exactly-once/exit-on-cancel +// channel contract. +func (p *Provider) Stream(ctx context.Context, req llm.Request) (<-chan llm.Event, error) { + params, err := toParams(req, p.model) + if err != nil { + return nil, fmt.Errorf("anthropic: %w", err) + } + recordutil.LogJSON(ctx, "anthropic: stream request", params, p.apiKey) + + return llm.StreamEvents(ctx, func(send func(llm.Event) bool) { + stream := p.client.Messages.NewStreaming(ctx, params) + defer func() { _ = stream.Close() }() + + acc := newStreamAccumulator() + for stream.Next() { + events, err := acc.handle(ctx, stream.Current()) + if err != nil { + send(llm.Error{Err: fmt.Errorf("anthropic: %w", err)}) + return + } + for _, e := range events { + if !send(e) { + return + } + } + } + if err := stream.Err(); err != nil { + send(llm.Error{Err: mapError(err)}) + } + }), nil +} diff --git a/internal/llm/anthropic/stream_test.go b/internal/llm/anthropic/stream_test.go new file mode 100644 index 0000000..de58216 --- /dev/null +++ b/internal/llm/anthropic/stream_test.go @@ -0,0 +1,109 @@ +package anthropic + +import ( + "context" + "net/http" + "net/http/httptest" + "os" + "testing" + + "github.com/google/go-cmp/cmp" + + "github.com/reno/pico-code/internal/llm" +) + +const ( + streamGoldenPath = "testdata/golden/stream_response.sse" + streamEquivGoldenPath = "testdata/golden/stream_response_equivalent.json" +) + +// TestStreamMatchesNonStreamingEquivalent is 6.2's AC: replaying a recorded +// SSE fixture that splits a tool argument across 7 chunks — including one +// split between the "\u00" and "e9" halves of a single \uXXXX escape, i.e. +// mid-character — yields the exact same Message as the non-streaming path +// for the same content. +func TestStreamMatchesNonStreamingEquivalent(t *testing.T) { + sse, err := os.ReadFile(streamGoldenPath) + if err != nil { + t.Fatalf("ReadFile(%s) error = %v", streamGoldenPath, err) + } + streamSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(sse) + })) + defer streamSrv.Close() + + nonStreamBody, err := os.ReadFile(streamEquivGoldenPath) + if err != nil { + t.Fatalf("ReadFile(%s) error = %v", streamEquivGoldenPath, err) + } + nonStreamSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(nonStreamBody) + })) + defer nonStreamSrv.Close() + + req := llm.Request{ + Messages: []llm.Message{{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "what's the weather in the café district?"}}}}, + MaxTokens: 64, + } + + streamP := testProvider(streamSrv.URL) + ch, err := streamP.Stream(context.Background(), req) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + streamed, err := llm.CollectStream(context.Background(), ch) + if err != nil { + t.Fatalf("CollectStream() error = %v", err) + } + + nonStreamP := testProvider(nonStreamSrv.URL) + nonStreamed, err := nonStreamP.Chat(context.Background(), req) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + + if diff := cmp.Diff(nonStreamed, streamed, rawJSONSemanticEqual); diff != "" { + t.Errorf("streamed Response differs from non-streaming equivalent (-nonStreaming +streamed):\n%s", diff) + } + + // Pin the specific split-argument reconstruction directly: this is + // what actually exercises the mid-\uXXXX-escape chunk boundary, since + // the cmp.Diff above only proves semantic equality after decoding. + toolUse, ok := streamed.Message.Blocks[1].(llm.ToolUse) + if !ok { + t.Fatalf("Blocks[1] = %T, want llm.ToolUse", streamed.Message.Blocks[1]) + } + wantInput := `{"location":"caf\u00e9"}` + if string(toolUse.Input) != wantInput { + t.Errorf("reassembled tool input = %s, want the exact accumulated bytes %s", toolUse.Input, wantInput) + } +} + +// TestStreamCancelledContextStopsPromptly proves ctx cancellation reaches +// an in-flight stream (CLAUDE.md invariant 6), reusing the same +// slow-429-forever server pattern as the non-streaming cancellation test. +func TestStreamCancelledContextStopsPromptly(t *testing.T) { + block := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { + <-r.Context().Done() + })) + defer srv.Close() + defer close(block) + + p := testProvider(srv.URL) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + ch, err := p.Stream(ctx, testChatRequest()) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + _, err = llm.CollectStream(context.Background(), ch) + if err == nil { + t.Fatal("CollectStream() error = nil, want an error for a cancelled stream") + } +} diff --git a/internal/llm/anthropic/testdata/golden/stream_response.sse b/internal/llm/anthropic/testdata/golden/stream_response.sse new file mode 100644 index 0000000..0d2a1b5 --- /dev/null +++ b/internal/llm/anthropic/testdata/golden/stream_response.sse @@ -0,0 +1,45 @@ +event: message_start +data: {"message":{"content":[],"id":"msg_stream_1","model":"claude-test-model","role":"assistant","stop_reason":null,"stop_sequence":null,"type":"message","usage":{"input_tokens":25,"output_tokens":1}},"type":"message_start"} + +event: content_block_start +data: {"content_block":{"text":"","type":"text"},"index":0,"type":"content_block_start"} + +event: content_block_delta +data: {"delta":{"text":"Let me look that up.","type":"text_delta"},"index":0,"type":"content_block_delta"} + +event: content_block_stop +data: {"index":0,"type":"content_block_stop"} + +event: content_block_start +data: {"content_block":{"id":"toolu_01","input":{},"name":"get_weather","type":"tool_use"},"index":1,"type":"content_block_start"} + +event: content_block_delta +data: {"delta":{"partial_json":"{\"lo","type":"input_json_delta"},"index":1,"type":"content_block_delta"} + +event: content_block_delta +data: {"delta":{"partial_json":"cati","type":"input_json_delta"},"index":1,"type":"content_block_delta"} + +event: content_block_delta +data: {"delta":{"partial_json":"on\":","type":"input_json_delta"},"index":1,"type":"content_block_delta"} + +event: content_block_delta +data: {"delta":{"partial_json":"\"caf","type":"input_json_delta"},"index":1,"type":"content_block_delta"} + +event: content_block_delta +data: {"delta":{"partial_json":"\\u00","type":"input_json_delta"},"index":1,"type":"content_block_delta"} + +event: content_block_delta +data: {"delta":{"partial_json":"e9","type":"input_json_delta"},"index":1,"type":"content_block_delta"} + +event: content_block_delta +data: {"delta":{"partial_json":"\"}","type":"input_json_delta"},"index":1,"type":"content_block_delta"} + +event: content_block_stop +data: {"index":1,"type":"content_block_stop"} + +event: message_delta +data: {"delta":{"stop_reason":"tool_use","stop_sequence":null},"type":"message_delta","usage":{"output_tokens":15}} + +event: message_stop +data: {"type":"message_stop"} + diff --git a/internal/llm/anthropic/testdata/golden/stream_response_equivalent.json b/internal/llm/anthropic/testdata/golden/stream_response_equivalent.json new file mode 100644 index 0000000..28bea8c --- /dev/null +++ b/internal/llm/anthropic/testdata/golden/stream_response_equivalent.json @@ -0,0 +1,9 @@ +{ + "id": "msg_stream_1", "type": "message", "role": "assistant", "model": "claude-test-model", + "content": [ + {"type": "text", "text": "Let me look that up."}, + {"type": "tool_use", "id": "toolu_01", "name": "get_weather", "input": {"location":"café"}} + ], + "stop_reason": "tool_use", + "usage": {"input_tokens": 25, "output_tokens": 15} +} diff --git a/internal/llm/collect.go b/internal/llm/collect.go new file mode 100644 index 0000000..0f0b0ee --- /dev/null +++ b/internal/llm/collect.go @@ -0,0 +1,64 @@ +package llm + +import ( + "context" + "fmt" + "strings" +) + +// CollectStream drains ch into a Response, reconstructing Message.Blocks in +// the same shape a non-streaming call would produce: TextDelta events +// accumulate into the currently open Text block (a new one starts whenever +// a ToolUseStart/ToolUseDone pair intervenes), and each ToolUseDone +// contributes one ToolUse block using the Name its ToolUseStart announced. +// It returns an error if ch yields an Error event, closes before a +// MessageDone arrives, or ctx is cancelled first. +func CollectStream(ctx context.Context, ch <-chan Event) (*Response, error) { + var blocks []Block + var text strings.Builder + textOpen := false + names := map[string]string{} + + flushText := func() { + if textOpen { + blocks = append(blocks, Text{Text: text.String()}) + text.Reset() + textOpen = false + } + } + + for { + select { + case e, ok := <-ch: + if !ok { + return nil, fmt.Errorf("llm: stream closed before a MessageDone event") + } + switch v := e.(type) { + case TextDelta: + textOpen = true + text.WriteString(v.Text) + case ToolUseStart: + flushText() + names[v.ID] = v.Name + case ToolUseArgsDelta: + // Informational only; the adapter accumulates internally + // and hands back the finished value in ToolUseDone. + case ToolUseDone: + blocks = append(blocks, ToolUse{ID: v.ID, Name: names[v.ID], Input: v.Input}) + case MessageDone: + flushText() + return &Response{ + Message: Message{Role: RoleAssistant, Blocks: blocks}, + StopReason: v.StopReason, + Usage: v.Usage, + }, nil + case Error: + return nil, v.Err + default: + return nil, fmt.Errorf("llm: unknown event type %T", e) + } + case <-ctx.Done(): + return nil, ctx.Err() + } + } +} diff --git a/internal/llm/collect_test.go b/internal/llm/collect_test.go new file mode 100644 index 0000000..2f62f28 --- /dev/null +++ b/internal/llm/collect_test.go @@ -0,0 +1,82 @@ +package llm + +import ( + "context" + "encoding/json" + "errors" + "testing" + + "github.com/google/go-cmp/cmp" +) + +func TestCollectStreamReconstructsBlockBoundaries(t *testing.T) { + ch := StreamEvents(context.Background(), func(send func(Event) bool) { + send(TextDelta{Text: "Let me check. "}) + send(ToolUseStart{ID: "call_1", Name: "get_weather"}) + send(ToolUseArgsDelta{ID: "call_1", Partial: `{"loc`}) + send(ToolUseArgsDelta{ID: "call_1", Partial: `ation":"Paris"}`}) + send(ToolUseDone{ID: "call_1", Input: json.RawMessage(`{"location":"Paris"}`)}) + send(TextDelta{Text: "One moment."}) + send(MessageDone{StopReason: "tool_use", Usage: Usage{InputTokens: 10, OutputTokens: 5}}) + }) + + got, err := CollectStream(context.Background(), ch) + if err != nil { + t.Fatalf("CollectStream() error = %v", err) + } + + want := &Response{ + Message: Message{ + Role: RoleAssistant, + Blocks: []Block{ + Text{Text: "Let me check. "}, + ToolUse{ID: "call_1", Name: "get_weather", Input: json.RawMessage(`{"location":"Paris"}`)}, + Text{Text: "One moment."}, + }, + }, + StopReason: "tool_use", + Usage: Usage{InputTokens: 10, OutputTokens: 5}, + } + if diff := cmp.Diff(want, got); diff != "" { + t.Errorf("CollectStream() mismatch (-want +got):\n%s", diff) + } +} + +func TestCollectStreamPropagatesErrorEvent(t *testing.T) { + boom := errors.New("boom") + ch := StreamEvents(context.Background(), func(send func(Event) bool) { + send(TextDelta{Text: "partial"}) + send(Error{Err: boom}) + }) + + _, err := CollectStream(context.Background(), ch) + if !errors.Is(err, boom) { + t.Fatalf("CollectStream() error = %v, want wrapping %v", err, boom) + } +} + +func TestCollectStreamErrorsIfChannelClosesWithoutMessageDone(t *testing.T) { + ch := StreamEvents(context.Background(), func(send func(Event) bool) { + send(TextDelta{Text: "partial"}) + }) + + _, err := CollectStream(context.Background(), ch) + if err == nil { + t.Fatal("CollectStream() error = nil, want an error for a stream that closed early") + } +} + +func TestCollectStreamRespectsCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + block := make(chan struct{}) + ch := StreamEvents(context.Background(), func(_ func(Event) bool) { + <-block + }) + + cancel() + _, err := CollectStream(ctx, ch) + if !errors.Is(err, context.Canceled) { + t.Fatalf("CollectStream() error = %v, want wrapping context.Canceled", err) + } + close(block) +} diff --git a/internal/llm/conformance/conformance.go b/internal/llm/conformance/conformance.go new file mode 100644 index 0000000..9ed07a7 --- /dev/null +++ b/internal/llm/conformance/conformance.go @@ -0,0 +1,190 @@ +// Package conformance is a single test suite run against every llm.Provider +// adapter, so a behavior guaranteed by the Provider interface (a canonical +// round trip, a tool call decodes correctly, an HTTP error surfaces as a Go +// error, ctx cancellation returns promptly, streaming and non-streaming +// agree) is asserted once instead of once per adapter. Adding a third +// provider only requires building a Harness for it — this package never +// imports a provider SDK. +package conformance + +import ( + "context" + "encoding/json" + "errors" + "net/http/httptest" + "reflect" + "testing" + "time" + + "github.com/google/go-cmp/cmp" + + "github.com/reno/pico-code/internal/llm" +) + +// Scenario names one canned server behavior a Harness must be able to +// serve, so the suite can drive each adapter through the same situations. +type Scenario int + +const ( + // ScenarioSimpleText is a plain-text, non-streaming reply. + ScenarioSimpleText Scenario = iota + // ScenarioSimpleTextStream is ScenarioSimpleText's streaming form. + ScenarioSimpleTextStream + // ScenarioToolCall is a non-streaming reply containing one tool call. + ScenarioToolCall + // ScenarioToolCallStream is ScenarioToolCall's streaming form. + ScenarioToolCallStream + // ScenarioServerError is an HTTP error response. + ScenarioServerError + // ScenarioHang never responds until the request's ctx is done. + ScenarioHang +) + +// Harness is what a provider package supplies to run this suite against +// its own adapter. +type Harness struct { + // Name identifies the provider in subtest names. + Name string + // NewProvider returns a Provider whose requests go to baseURL. + NewProvider func(baseURL string) llm.Provider + // Server returns an httptest.Server serving scenario's canned + // response; the suite closes it. + Server func(t *testing.T, scenario Scenario) *httptest.Server + // WantText is the text ScenarioSimpleText/ScenarioSimpleTextStream + // decode to. + WantText string + // WantToolUse is the call ScenarioToolCall/ScenarioToolCallStream + // decode to (ID is not compared). + WantToolUse llm.ToolUse +} + +// Run drives every conformance check against h. +func Run(t *testing.T, h Harness) { + t.Run(h.Name+"_CanonicalRoundTrip", func(t *testing.T) { testCanonicalRoundTrip(t, h) }) + t.Run(h.Name+"_ToolCall", func(t *testing.T) { testToolCall(t, h) }) + t.Run(h.Name+"_ErrorMapping", func(t *testing.T) { testErrorMapping(t, h) }) + t.Run(h.Name+"_Cancellation", func(t *testing.T) { testCancellation(t, h) }) + t.Run(h.Name+"_StreamNonStreamEquivalence", func(t *testing.T) { testStreamNonStreamEquivalence(t, h) }) +} + +func request() llm.Request { + return llm.Request{ + Messages: []llm.Message{{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "hi"}}}}, + MaxTokens: 32, + } +} + +func testCanonicalRoundTrip(t *testing.T, h Harness) { + srv := h.Server(t, ScenarioSimpleText) + defer srv.Close() + + resp, err := h.NewProvider(srv.URL).Chat(context.Background(), request()) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + if len(resp.Message.Blocks) != 1 { + t.Fatalf("Blocks = %v, want exactly one Text block", resp.Message.Blocks) + } + text, ok := resp.Message.Blocks[0].(llm.Text) + if !ok || text.Text != h.WantText { + t.Errorf("Blocks[0] = %#v, want Text{%q}", resp.Message.Blocks[0], h.WantText) + } +} + +func testToolCall(t *testing.T, h Harness) { + srv := h.Server(t, ScenarioToolCall) + defer srv.Close() + + resp, err := h.NewProvider(srv.URL).Chat(context.Background(), request()) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + assertToolUse(t, resp.Message.Blocks, h.WantToolUse) +} + +func testErrorMapping(t *testing.T, h Harness) { + srv := h.Server(t, ScenarioServerError) + defer srv.Close() + + _, err := h.NewProvider(srv.URL).Chat(context.Background(), request()) + if err == nil { + t.Fatal("Chat() error = nil, want a non-nil error for a server error response") + } +} + +func testCancellation(t *testing.T, h Harness) { + srv := h.Server(t, ScenarioHang) + defer srv.Close() + + ctx, cancel := context.WithCancel(context.Background()) + time.AfterFunc(20*time.Millisecond, cancel) + + start := time.Now() + _, err := h.NewProvider(srv.URL).Chat(ctx, request()) + elapsed := time.Since(start) + + if !errors.Is(err, context.Canceled) { + t.Fatalf("Chat() error = %v, want it to wrap context.Canceled", err) + } + if elapsed > time.Second { + t.Fatalf("Chat() took %s after cancellation, want a prompt return", elapsed) + } +} + +func testStreamNonStreamEquivalence(t *testing.T, h Harness) { + streamSrv := h.Server(t, ScenarioToolCallStream) + defer streamSrv.Close() + nonStreamSrv := h.Server(t, ScenarioToolCall) + defer nonStreamSrv.Close() + + ch, err := h.NewProvider(streamSrv.URL).Stream(context.Background(), request()) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + streamed, err := llm.CollectStream(context.Background(), ch) + if err != nil { + t.Fatalf("CollectStream() error = %v", err) + } + + nonStreamed, err := h.NewProvider(nonStreamSrv.URL).Chat(context.Background(), request()) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + + rawJSONSemanticEqual := cmp.Comparer(func(a, b json.RawMessage) bool { + var va, vb any + if json.Unmarshal(a, &va) != nil || json.Unmarshal(b, &vb) != nil { + return false + } + return reflect.DeepEqual(va, vb) + }) + if diff := cmp.Diff(nonStreamed, streamed, rawJSONSemanticEqual); diff != "" { + t.Errorf("streamed Response differs from non-streaming equivalent (-nonStreaming +streamed):\n%s", diff) + } +} + +func assertToolUse(t *testing.T, blocks []llm.Block, want llm.ToolUse) { + t.Helper() + var got *llm.ToolUse + for _, b := range blocks { + if tu, ok := b.(llm.ToolUse); ok { + got = &tu + } + } + if got == nil { + t.Fatalf("Blocks = %v, want a ToolUse block", blocks) + } + if got.Name != want.Name { + t.Errorf("ToolUse.Name = %q, want %q", got.Name, want.Name) + } + var gotArgs, wantArgs any + if err := json.Unmarshal(got.Input, &gotArgs); err != nil { + t.Fatalf("ToolUse.Input is not valid JSON: %v", err) + } + if err := json.Unmarshal(want.Input, &wantArgs); err != nil { + t.Fatalf("want.Input is not valid JSON: %v", err) + } + if diff := cmp.Diff(wantArgs, gotArgs); diff != "" { + t.Errorf("ToolUse.Input mismatch (-want +got):\n%s", diff) + } +} diff --git a/internal/llm/events.go b/internal/llm/events.go index d830da3..b08fd2e 100644 --- a/internal/llm/events.go +++ b/internal/llm/events.go @@ -1,11 +1,86 @@ package llm +import ( + "context" + "encoding/json" +) + // Event is one item in a Provider's Stream. It is sealed to this package // (via the unexported isEvent method), mirroring Block, so a consumer's -// type switch is exhaustive by construction. Concrete variants (TextDelta, -// ToolUseStart, ToolUseArgsDelta, ToolUseDone, MessageDone, Error) are added -// in phase 6.1 alongside the streaming adapters that produce them; Provider -// only needs the type to exist to declare its Stream method now. +// type switch is exhaustive by construction. type Event interface { isEvent() } + +// TextDelta is an incremental chunk of assistant text. +type TextDelta struct { + Text string +} + +func (TextDelta) isEvent() {} + +// ToolUseStart announces a new tool call; its arguments follow as +// ToolUseArgsDelta events and finish with a matching ToolUseDone. +type ToolUseStart struct { + ID string + Name string +} + +func (ToolUseStart) isEvent() {} + +// ToolUseArgsDelta is a raw fragment of a tool call's JSON arguments. +// Fragments must only be concatenated, never parsed individually — CLAUDE.md: +// "Accumulate the raw string and json.Unmarshal exactly once at block stop." +type ToolUseArgsDelta struct { + ID string + Partial string +} + +func (ToolUseArgsDelta) isEvent() {} + +// ToolUseDone finalizes a tool call with its fully assembled input. +type ToolUseDone struct { + ID string + Input json.RawMessage +} + +func (ToolUseDone) isEvent() {} + +// MessageDone marks the end of the assistant's turn. +type MessageDone struct { + StopReason string + Usage Usage +} + +func (MessageDone) isEvent() {} + +// Error carries a terminal stream failure. A Provider sends at most one, +// always as its last event before the channel closes. +type Error struct { + Err error +} + +func (Error) isEvent() {} + +// StreamEvents runs produce in its own goroutine and returns the channel it +// sends to, closing that channel exactly once when produce returns — +// whether that's because it finished normally or because send reported ctx +// was cancelled. Every Provider.Stream implementation should build its +// channel through this rather than managing goroutine/channel lifecycle +// itself, so the "closed exactly once, exits on cancel" contract lives in +// one place. +func StreamEvents(ctx context.Context, produce func(send func(Event) bool)) <-chan Event { + ch := make(chan Event) + go func() { + defer close(ch) + produce(func(e Event) bool { + select { + case ch <- e: + return true + case <-ctx.Done(): + return false + } + }) + }() + return ch +} diff --git a/internal/llm/events_test.go b/internal/llm/events_test.go new file mode 100644 index 0000000..fccf07f --- /dev/null +++ b/internal/llm/events_test.go @@ -0,0 +1,86 @@ +package llm + +import ( + "context" + "testing" + "time" + + "go.uber.org/goleak" +) + +func TestMain(m *testing.M) { + goleak.VerifyTestMain(m) +} + +// TestStreamEventsConsumedToCompletion is half of 6.1's AC: a fake stream +// consumed to completion closes its channel exactly once, with every event +// delivered in order. +func TestStreamEventsConsumedToCompletion(t *testing.T) { + done := make(chan struct{}) + ch := StreamEvents(context.Background(), func(send func(Event) bool) { + defer close(done) + send(TextDelta{Text: "a"}) + send(TextDelta{Text: "b"}) + send(MessageDone{StopReason: "end_turn"}) + }) + + var got []Event + for e := range ch { + got = append(got, e) + } + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("producer goroutine did not finish") + } + + if len(got) != 3 { + t.Fatalf("got %d events, want 3", len(got)) + } + if _, ok := got[2].(MessageDone); !ok { + t.Errorf("last event = %T, want MessageDone", got[2]) + } + + if _, ok := <-ch; ok { + t.Error("channel yielded another value after range completed") + } +} + +// TestStreamEventsCancelledMidFlight is the other half of 6.1's AC: a +// producer that would otherwise send forever stops and the channel closes +// once ctx is cancelled, with no leaked goroutine (verified by TestMain's +// goleak.VerifyTestMain, and directly here via the finished handshake) and +// no send on a closed channel (the race detector would catch that). +func TestStreamEventsCancelledMidFlight(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + started := make(chan struct{}) + finished := make(chan struct{}) + + ch := StreamEvents(ctx, func(send func(Event) bool) { + defer close(finished) + close(started) + for { + if !send(TextDelta{Text: "x"}) { + return + } + } + }) + + <-started + <-ch + <-ch + cancel() + + select { + case <-finished: + case <-time.After(time.Second): + t.Fatal("producer goroutine did not exit after ctx cancellation") + } + + for { + if _, ok := <-ch; !ok { + break + } + } +} diff --git a/internal/llm/ollama/conformance_test.go b/internal/llm/ollama/conformance_test.go new file mode 100644 index 0000000..bb008ca --- /dev/null +++ b/internal/llm/ollama/conformance_test.go @@ -0,0 +1,96 @@ +package ollama + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/reno/pico-code/internal/llm" + "github.com/reno/pico-code/internal/llm/conformance" +) + +// TestConformance is 6.4's AC: the shared conformance suite passes for the +// Ollama adapter unchanged. +func TestConformance(t *testing.T) { + conformance.Run(t, conformance.Harness{ + Name: "ollama", + NewProvider: func(baseURL string) llm.Provider { + return &Provider{httpClient: http.DefaultClient, baseURL: baseURL, model: "qwen3:8b", numCtx: 4096} + }, + Server: conformanceServer, + WantText: "hello", + WantToolUse: llm.ToolUse{Name: "get_weather", Input: json.RawMessage(`{"location":"Paris"}`)}, + }) +} + +func conformanceServer(t *testing.T, scenario conformance.Scenario) *httptest.Server { + t.Helper() + switch scenario { + case conformance.ScenarioSimpleText: + return jsonServer(`{ + "model": "qwen3:8b", + "message": {"role": "assistant", "content": "hello"}, + "done": true, "done_reason": "stop", + "prompt_eval_count": 5, "eval_count": 2 + }`) + case conformance.ScenarioToolCall: + return jsonServer(`{ + "model": "qwen3:8b", + "message": {"role": "assistant", "content": "", "tool_calls": [ + {"id": "call_ct1", "function": {"name": "get_weather", "arguments": {"location":"Paris"}}} + ]}, + "done": true, "done_reason": "stop", + "prompt_eval_count": 8, "eval_count": 4 + }`) + case conformance.ScenarioSimpleTextStream: + return ndjsonServer( + `{"model":"qwen3:8b","message":{"role":"assistant","content":"hello"},"done":false}`, + `{"model":"qwen3:8b","message":{"role":"assistant","content":""},"done":true,"done_reason":"stop","prompt_eval_count":5,"eval_count":2}`, + ) + case conformance.ScenarioToolCallStream: + return ndjsonServer( + `{"model":"qwen3:8b","message":{"role":"assistant","content":"","tool_calls":[{"id":"call_ct1","function":{"name":"get_weather","arguments":{"location":"Paris"}}}]},"done":false}`, + `{"model":"qwen3:8b","message":{"role":"assistant","content":""},"done":true,"done_reason":"stop","prompt_eval_count":8,"eval_count":4}`, + ) + case conformance.ScenarioServerError: + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":"bad key"}`)) + })) + case conformance.ScenarioHang: + // Bounded rather than truly infinite — see the identical comment + // in the Anthropic conformance harness for why. + return httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { + select { + case <-r.Context().Done(): + case <-time.After(150 * time.Millisecond): + } + })) + default: + t.Fatalf("conformanceServer: unhandled scenario %v", scenario) + return nil + } +} + +func jsonServer(body string) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(body)) + })) +} + +func ndjsonServer(lines ...string) *httptest.Server { + body := "" + for _, l := range lines { + body += l + "\n" + } + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/x-ndjson") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(body)) + })) +} diff --git a/internal/llm/ollama/estimate.go b/internal/llm/ollama/estimate.go new file mode 100644 index 0000000..3d4b078 --- /dev/null +++ b/internal/llm/ollama/estimate.go @@ -0,0 +1,56 @@ +package ollama + +import ( + "strings" + + "github.com/reno/pico-code/internal/llm" +) + +// estimateTokens approximates a token count from character length for use +// when Ollama's response omits real counts (prompt_eval_count/eval_count +// come back 0 whenever it skips evaluation, e.g. a fully cached prompt). +// ~4 characters per token is the same rough rule of thumb OpenAI's +// tokenizer docs use for English text — good enough for a fallback +// estimate, not for billing. +func estimateTokens(s string) int { + if s == "" { + return 0 + } + if n := len(s) / 4; n > 0 { + return n + } + return 1 +} + +// estimateUsage fills in any zero field of reported with a character-based +// estimate derived from req's text and the response's own text, so +// cumulative usage tracking (8.1) keeps making progress against an Ollama +// version or model that omits real counts. +func estimateUsage(req llm.Request, responseText string, reported llm.Usage) llm.Usage { + usage := reported + if usage.InputTokens == 0 { + usage.InputTokens = estimateTokens(requestText(req)) + } + if usage.OutputTokens == 0 { + usage.OutputTokens = estimateTokens(responseText) + } + return usage +} + +func requestText(req llm.Request) string { + var b strings.Builder + b.WriteString(req.System) + for _, m := range req.Messages { + for _, blk := range m.Blocks { + switch v := blk.(type) { + case llm.Text: + b.WriteString(v.Text) + case llm.ToolResult: + b.WriteString(v.Content) + case llm.ToolUse: + b.Write(v.Input) + } + } + } + return b.String() +} diff --git a/internal/llm/ollama/estimate_test.go b/internal/llm/ollama/estimate_test.go new file mode 100644 index 0000000..159706f --- /dev/null +++ b/internal/llm/ollama/estimate_test.go @@ -0,0 +1,56 @@ +package ollama + +import ( + "testing" + + "github.com/reno/pico-code/internal/llm" +) + +func TestEstimateUsageLeavesReportedCountsAlone(t *testing.T) { + req := llm.Request{System: "be concise", Messages: []llm.Message{ + {Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "hello there"}}}, + }} + reported := llm.Usage{InputTokens: 512, OutputTokens: 64} + + got := estimateUsage(req, "hi", reported) + if got != reported { + t.Errorf("estimateUsage() = %+v, want reported counts unchanged: %+v", got, reported) + } +} + +// TestEstimateUsageFillsInZeroCounts is 8.1's AC for Ollama: when a +// response omits its real counts, estimation still produces a non-zero +// usage so cumulative tracking keeps making progress. +func TestEstimateUsageFillsInZeroCounts(t *testing.T) { + req := llm.Request{Messages: []llm.Message{ + {Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "this prompt has more than a few characters in it"}}}, + }} + + got := estimateUsage(req, "a fairly short reply", llm.Usage{}) + if got.InputTokens <= 0 { + t.Errorf("InputTokens = %d, want a positive estimate", got.InputTokens) + } + if got.OutputTokens <= 0 { + t.Errorf("OutputTokens = %d, want a positive estimate", got.OutputTokens) + } +} + +func TestEstimateUsagePartialZeroOnlyFillsMissingField(t *testing.T) { + req := llm.Request{Messages: []llm.Message{ + {Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "hello"}}}, + }} + + got := estimateUsage(req, "a reply with several words in it", llm.Usage{InputTokens: 10}) + if got.InputTokens != 10 { + t.Errorf("InputTokens = %d, want the reported 10 preserved", got.InputTokens) + } + if got.OutputTokens <= 0 { + t.Errorf("OutputTokens = %d, want a positive estimate for the omitted field", got.OutputTokens) + } +} + +func TestEstimateTokensEmptyStringIsZero(t *testing.T) { + if got := estimateTokens(""); got != 0 { + t.Errorf("estimateTokens(\"\") = %d, want 0", got) + } +} diff --git a/internal/llm/ollama/logging_test.go b/internal/llm/ollama/logging_test.go new file mode 100644 index 0000000..757ed0f --- /dev/null +++ b/internal/llm/ollama/logging_test.go @@ -0,0 +1,78 @@ +package ollama + +import ( + "bytes" + "context" + "log/slog" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/reno/pico-code/internal/llm" +) + +func withDefaultLogger(t *testing.T, level slog.Level) *bytes.Buffer { + t.Helper() + var buf bytes.Buffer + prev := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: level}))) + t.Cleanup(func() { slog.SetDefault(prev) }) + return &buf +} + +const loggingTestResponseBody = `{ + "model": "qwen3:8b", + "message": {"role": "assistant", "content": "hi there"}, + "done": true, "done_reason": "stop", + "prompt_eval_count": 10, "eval_count": 5 +}` + +// TestDebugLoggingDumpsExchange is 9.2's AC: debug output of a full turn +// contains no credentials. Ollama has none by default (no API key), so +// this mainly proves the dump actually happens and stays clean of the +// configured host string. +func TestDebugLoggingDumpsExchange(t *testing.T) { + buf := withDefaultLogger(t, slog.LevelDebug) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(loggingTestResponseBody)) + })) + defer srv.Close() + + p := &Provider{httpClient: http.DefaultClient, baseURL: srv.URL, model: "qwen3:8b", numCtx: 4096} + if _, err := p.Chat(context.Background(), llm.Request{ + Messages: []llm.Message{{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "hello"}}}}, + }); err != nil { + t.Fatalf("Chat() error = %v", err) + } + + out := buf.String() + if !strings.Contains(out, "ollama: request") || !strings.Contains(out, "ollama: response") { + t.Fatalf("debug output missing request/response dump; got:\n%s", out) + } +} + +// TestDefaultLevelDoesNotDumpRequestsOrResponses is 9.2's other AC: default +// level is silent on the happy path. +func TestDefaultLevelDoesNotDumpRequestsOrResponses(t *testing.T) { + buf := withDefaultLogger(t, slog.LevelInfo) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(loggingTestResponseBody)) + })) + defer srv.Close() + + p := &Provider{httpClient: http.DefaultClient, baseURL: srv.URL, model: "qwen3:8b", numCtx: 4096} + if _, err := p.Chat(context.Background(), llm.Request{ + Messages: []llm.Message{{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "hello"}}}}, + }); err != nil { + t.Fatalf("Chat() error = %v", err) + } + + if out := buf.String(); strings.Contains(out, "ollama: request") || strings.Contains(out, "ollama: response") { + t.Errorf("default (info) level logged a request/response dump, want none:\n%s", out) + } +} diff --git a/internal/llm/ollama/provider.go b/internal/llm/ollama/provider.go index df0a49f..b7224ad 100644 --- a/internal/llm/ollama/provider.go +++ b/internal/llm/ollama/provider.go @@ -4,7 +4,6 @@ import ( "bytes" "context" "encoding/json" - "errors" "fmt" "io" "log/slog" @@ -15,6 +14,7 @@ import ( "github.com/reno/pico-code/internal/config" "github.com/reno/pico-code/internal/llm" + "github.com/reno/pico-code/internal/llm/recordutil" ) func init() { @@ -80,6 +80,8 @@ func (p *Provider) Chat(ctx context.Context, req llm.Request) (*llm.Response, er return nil, fmt.Errorf("ollama: marshal request: %w", err) } + recordutil.LogBytes(ctx, "ollama: request", body) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL+"/api/chat", bytes.NewReader(body)) if err != nil { return nil, fmt.Errorf("ollama: %w", err) @@ -99,6 +101,7 @@ func (p *Provider) Chat(ctx context.Context, req llm.Request) (*llm.Response, er if res.StatusCode >= http.StatusBadRequest { return nil, fmt.Errorf("ollama: request failed with status %d: %s", res.StatusCode, bytes.TrimSpace(respBody)) } + recordutil.LogBytes(ctx, "ollama: response", respBody) normalized, err := normalizeToolCallArguments(respBody) if err != nil { @@ -110,10 +113,5 @@ func (p *Provider) Chat(ctx context.Context, req llm.Request) (*llm.Response, er return nil, fmt.Errorf("ollama: decode response: %w", err) } - return fromResponse(&chatResp) -} - -// Stream is implemented in phase 6.3. -func (p *Provider) Stream(_ context.Context, _ llm.Request) (<-chan llm.Event, error) { - return nil, errors.New("ollama: streaming not implemented until phase 6.3") + return fromResponse(req, &chatResp) } diff --git a/internal/llm/ollama/record_test.go b/internal/llm/ollama/record_test.go new file mode 100644 index 0000000..8aed7e4 --- /dev/null +++ b/internal/llm/ollama/record_test.go @@ -0,0 +1,83 @@ +package ollama + +import ( + "context" + "net/http" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/reno/pico-code/internal/llm" + "github.com/reno/pico-code/internal/llm/recordutil" +) + +// TestRecordLiveExchange is 9.1's record half: with RECORD=1 and a running +// local Ollama server, it drives one real Chat call and writes the +// request/response to testdata/golden/live_exchange. Ollama needs no API +// key, so there is nothing to scrub beyond OLLAMA_HOST if it's non-default +// — but recordutil.Scrub still runs, in case a custom host string embeds +// something sensitive. It never runs in `make test` (no RECORD set) and +// skips itself if Ollama isn't actually reachable. +func TestRecordLiveExchange(t *testing.T) { + if !recordutil.Enabled() { + t.Skip("set RECORD=1 (with Ollama running locally) to re-record this fixture against the live API") + } + host := os.Getenv("OLLAMA_HOST") + if host == "" { + host = defaultHost + } + model := os.Getenv("PICO_RECORD_OLLAMA_MODEL") + if model == "" { + t.Skip("set PICO_RECORD_OLLAMA_MODEL to a locally pulled model to record against") + } + + var exchange recordutil.Exchange + p := &Provider{ + httpClient: &http.Client{Transport: &recordutil.Recorder{OnExchange: func(e recordutil.Exchange) { exchange = e }}}, + baseURL: host, + model: model, + numCtx: 4096, + } + + _, err := p.Chat(context.Background(), llm.Request{ + Messages: []llm.Message{{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "Reply with exactly: ok"}}}}, + }) + if err != nil { + t.Fatalf("Chat() error = %v (is `ollama serve` running with %q pulled?)", err, model) + } + + dir := "testdata/golden/live_exchange" + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatalf("MkdirAll() error = %v", err) + } + reqOut := recordutil.Scrub(string(exchange.RequestBody), host) + respOut := recordutil.Scrub(string(exchange.ResponseBody), host) + if err := os.WriteFile(filepath.Join(dir, "request.json"), []byte(reqOut), 0o644); err != nil { + t.Fatalf("WriteFile(request) error = %v", err) + } + if err := os.WriteFile(filepath.Join(dir, "response.json"), []byte(respOut), 0o644); err != nil { + t.Fatalf("WriteFile(response) error = %v", err) + } +} + +// TestGoldenFixturesContainNoSecrets is 9.1's AC: no key ever appears in a +// fixture. +func TestGoldenFixturesContainNoSecrets(t *testing.T) { + err := filepath.WalkDir("testdata/golden", func(path string, d os.DirEntry, err error) error { + if err != nil || d.IsDir() { + return err + } + data, rerr := os.ReadFile(path) + if rerr != nil { + return rerr + } + if strings.Contains(string(data), "sk-") { + t.Errorf("%s contains an sk-shaped key literal", path) + } + return nil + }) + if err != nil { + t.Fatalf("walking testdata/golden: %v", err) + } +} diff --git a/internal/llm/ollama/response_test.go b/internal/llm/ollama/response_test.go index a95b60c..8dc2ae7 100644 --- a/internal/llm/ollama/response_test.go +++ b/internal/llm/ollama/response_test.go @@ -26,7 +26,7 @@ func decodeFixture(t *testing.T, path string) *llm.Response { if err := json.Unmarshal(normalized, &chatResp); err != nil { t.Fatalf("Unmarshal(%s) error = %v", path, err) } - got, err := fromResponse(&chatResp) + got, err := fromResponse(llm.Request{}, &chatResp) if err != nil { t.Fatalf("fromResponse(%s) error = %v", path, err) } diff --git a/internal/llm/ollama/stream.go b/internal/llm/ollama/stream.go new file mode 100644 index 0000000..ed90e10 --- /dev/null +++ b/internal/llm/ollama/stream.go @@ -0,0 +1,127 @@ +package ollama + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + + "github.com/ollama/ollama/api" + + "github.com/reno/pico-code/internal/llm" + "github.com/reno/pico-code/internal/llm/recordutil" +) + +// Stream sends req to /api/chat with streaming enabled and translates each +// NDJSON line into canonical events. Unlike Anthropic, Ollama doesn't +// fragment a tool call's arguments across lines — each line's tool_calls +// are already complete — so a tool call becomes an immediate +// ToolUseStart/ToolUseDone pair rather than accumulating deltas. +func (p *Provider) Stream(ctx context.Context, req llm.Request) (<-chan llm.Event, error) { + if err := p.checkToolSupport(ctx, req); err != nil { + return nil, err + } + + ar, err := toRequest(req, p.model, p.numCtx) + if err != nil { + return nil, fmt.Errorf("ollama: %w", err) + } + stream := true + ar.Stream = &stream + + body, err := json.Marshal(ar) + if err != nil { + return nil, fmt.Errorf("ollama: marshal request: %w", err) + } + + recordutil.LogBytes(ctx, "ollama: stream request", body) + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL+"/api/chat", bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("ollama: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + res, err := p.httpClient.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("ollama: %w", err) + } + if res.StatusCode >= http.StatusBadRequest { + defer func() { _ = res.Body.Close() }() + respBody, _ := io.ReadAll(res.Body) + return nil, fmt.Errorf("ollama: request failed with status %d: %s", res.StatusCode, bytes.TrimSpace(respBody)) + } + + return llm.StreamEvents(ctx, func(send func(llm.Event) bool) { + defer func() { _ = res.Body.Close() }() + + toolCallSeq := 0 + var respText strings.Builder + scanner := bufio.NewScanner(res.Body) + scanner.Buffer(nil, 1024*1024) + + for scanner.Scan() { + line := bytes.TrimSpace(scanner.Bytes()) + if len(line) == 0 { + continue + } + + normalized, err := normalizeToolCallArguments(line) + if err != nil { + send(llm.Error{Err: fmt.Errorf("ollama: %w", err)}) + return + } + var chunk api.ChatResponse + if err := json.Unmarshal(normalized, &chunk); err != nil { + send(llm.Error{Err: fmt.Errorf("ollama: decode stream chunk: %w", err)}) + return + } + + if chunk.Message.Content != "" { + respText.WriteString(chunk.Message.Content) + if !send(llm.TextDelta{Text: chunk.Message.Content}) { + return + } + } + + for _, tc := range chunk.Message.ToolCalls { + id := tc.ID + if id == "" { + id = fmt.Sprintf("ollama_call_%d", toolCallSeq) + } + toolCallSeq++ + + input, err := json.Marshal(tc.Function.Arguments) + if err != nil { + send(llm.Error{Err: fmt.Errorf("ollama: tool call arguments: %w", err)}) + return + } + if !send(llm.ToolUseStart{ID: id, Name: tc.Function.Name}) { + return + } + if !send(llm.ToolUseDone{ID: id, Input: input}) { + return + } + } + + if chunk.Done { + recordutil.LogBytes(ctx, "ollama: stream response", line) + send(llm.MessageDone{ + StopReason: chunk.DoneReason, + Usage: estimateUsage(req, respText.String(), llm.Usage{ + InputTokens: chunk.PromptEvalCount, + OutputTokens: chunk.EvalCount, + }), + }) + return + } + } + if err := scanner.Err(); err != nil { + send(llm.Error{Err: fmt.Errorf("ollama: read stream: %w", err)}) + } + }), nil +} diff --git a/internal/llm/ollama/stream_test.go b/internal/llm/ollama/stream_test.go new file mode 100644 index 0000000..2ffdfda --- /dev/null +++ b/internal/llm/ollama/stream_test.go @@ -0,0 +1,95 @@ +package ollama + +import ( + "context" + "net/http" + "net/http/httptest" + "os" + "testing" + + "github.com/google/go-cmp/cmp" + + "github.com/reno/pico-code/internal/llm" +) + +const ( + streamGoldenPath = "testdata/golden/stream_response.ndjson" + streamEquivGoldenPath = "testdata/golden/stream_response_equivalent.json" +) + +// TestStreamMatchesNonStreamingEquivalent is 6.3's assembly half: the same +// conversation replayed through NDJSON streaming (including one line whose +// tool call arguments arrive as a JSON-encoded string, the normalization +// quirk 3.1 already handles for the non-streaming path) produces the exact +// same Response as the non-streaming path. +func TestStreamMatchesNonStreamingEquivalent(t *testing.T) { + ndjson, err := os.ReadFile(streamGoldenPath) + if err != nil { + t.Fatalf("ReadFile(%s) error = %v", streamGoldenPath, err) + } + streamSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/x-ndjson") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(ndjson) + })) + defer streamSrv.Close() + + nonStreamBody, err := os.ReadFile(streamEquivGoldenPath) + if err != nil { + t.Fatalf("ReadFile(%s) error = %v", streamEquivGoldenPath, err) + } + nonStreamSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(nonStreamBody) + })) + defer nonStreamSrv.Close() + + req := llm.Request{ + Messages: []llm.Message{{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "what's the weather in Paris?"}}}}, + Tools: []llm.ToolDefinition{ + {Name: "get_weather", Description: "get the weather", InputSchema: []byte(`{"type":"object"}`)}, + }, + } + + streamP := &Provider{httpClient: http.DefaultClient, baseURL: streamSrv.URL, model: "qwen3:8b", numCtx: 4096} + streamP.toolSupport = true + streamP.toolSupportOnce.Do(func() {}) + ch, err := streamP.Stream(context.Background(), req) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + streamed, err := llm.CollectStream(context.Background(), ch) + if err != nil { + t.Fatalf("CollectStream() error = %v", err) + } + + nonStreamP := &Provider{httpClient: http.DefaultClient, baseURL: nonStreamSrv.URL, model: "qwen3:8b", numCtx: 4096} + nonStreamP.toolSupport = true + nonStreamP.toolSupportOnce.Do(func() {}) + nonStreamed, err := nonStreamP.Chat(context.Background(), req) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + + if diff := cmp.Diff(nonStreamed, streamed); diff != "" { + t.Errorf("streamed Response differs from non-streaming equivalent (-nonStreaming +streamed):\n%s", diff) + } +} + +func TestStreamCancelledContextStopsPromptly(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { + <-r.Context().Done() + })) + defer srv.Close() + + p := &Provider{httpClient: http.DefaultClient, baseURL: srv.URL, model: "qwen3:8b", numCtx: 4096} + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := p.Stream(ctx, llm.Request{Messages: []llm.Message{{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "hi"}}}}}) + if err == nil { + t.Fatal("Stream() error = nil, want an error for a cancelled ctx") + } +} diff --git a/internal/llm/ollama/testdata/golden/stream_response.ndjson b/internal/llm/ollama/testdata/golden/stream_response.ndjson new file mode 100644 index 0000000..39dc1d9 --- /dev/null +++ b/internal/llm/ollama/testdata/golden/stream_response.ndjson @@ -0,0 +1,4 @@ +{"model":"qwen3:8b","created_at":"2024-01-01T00:00:00Z","message":{"role":"assistant","content":"Checking the "},"done":false} +{"model":"qwen3:8b","created_at":"2024-01-01T00:00:01Z","message":{"role":"assistant","content":"weather for you."},"done":false} +{"model":"qwen3:8b","created_at":"2024-01-01T00:00:02Z","message":{"role":"assistant","content":"","tool_calls":[{"id":"call_1","function":{"name":"get_weather","arguments":"{\"location\": \"Paris\"}"}}]},"done":false} +{"model":"qwen3:8b","created_at":"2024-01-01T00:00:03Z","message":{"role":"assistant","content":""},"done":true,"done_reason":"stop","prompt_eval_count":512,"eval_count":64} diff --git a/internal/llm/ollama/testdata/golden/stream_response_equivalent.json b/internal/llm/ollama/testdata/golden/stream_response_equivalent.json new file mode 100644 index 0000000..ae57ae6 --- /dev/null +++ b/internal/llm/ollama/testdata/golden/stream_response_equivalent.json @@ -0,0 +1,21 @@ +{ + "model": "qwen3:8b", + "created_at": "2024-01-01T00:00:03Z", + "message": { + "role": "assistant", + "content": "Checking the weather for you.", + "tool_calls": [ + { + "id": "call_1", + "function": { + "name": "get_weather", + "arguments": {"location": "Paris"} + } + } + ] + }, + "done": true, + "done_reason": "stop", + "prompt_eval_count": 512, + "eval_count": 64 +} diff --git a/internal/llm/ollama/translate.go b/internal/llm/ollama/translate.go index db91294..3dbd2ba 100644 --- a/internal/llm/ollama/translate.go +++ b/internal/llm/ollama/translate.go @@ -168,8 +168,9 @@ func toParameters(raw json.RawMessage) (api.ToolFunctionParameters, error) { } // fromResponse translates a non-streaming ChatResponse into the canonical -// Response. -func fromResponse(resp *api.ChatResponse) (*llm.Response, error) { +// Response. req is only used to estimate token counts when resp's own are +// 0 (see estimateUsage) — translation of the message itself doesn't need it. +func fromResponse(req llm.Request, resp *api.ChatResponse) (*llm.Response, error) { blocks, err := fromMessage(resp.Message) if err != nil { return nil, err @@ -177,10 +178,10 @@ func fromResponse(resp *api.ChatResponse) (*llm.Response, error) { return &llm.Response{ Message: llm.Message{Role: llm.RoleAssistant, Blocks: blocks}, StopReason: resp.DoneReason, - Usage: llm.Usage{ + Usage: estimateUsage(req, resp.Message.Content, llm.Usage{ InputTokens: resp.PromptEvalCount, OutputTokens: resp.EvalCount, - }, + }), }, nil } diff --git a/internal/llm/recordutil/recordutil.go b/internal/llm/recordutil/recordutil.go new file mode 100644 index 0000000..59c6155 --- /dev/null +++ b/internal/llm/recordutil/recordutil.go @@ -0,0 +1,135 @@ +// Package recordutil holds the secret-redaction logic shared by two +// things: the RECORD=1 convention CLAUDE.md's testing section describes +// (with RECORD=1 set and real credentials in the environment, an adapter's +// test proxies one live exchange through a Recorder that captures the raw +// request/response and writes them to disk as a scrubbed golden fixture), +// and --log-level=debug's request/response dump (9.2), which needs the +// exact same guarantee — no credential ever reaches a fixture or a log +// line. Without RECORD=1, the Recorder/Exchange half of this package is +// never exercised — `make test` only replays already-recorded fixtures via +// httptest.Server, so the suite stays offline and key-free. +package recordutil + +import ( + "bytes" + "context" + "encoding/json" + "io" + "log/slog" + "net/http" + "os" + "regexp" + "strings" +) + +// Enabled reports whether RECORD=1 is set. +func Enabled() bool { + return os.Getenv("RECORD") == "1" +} + +// apiKeyPattern matches an API-key-shaped literal (Anthropic's sk-ant-..., +// OpenAI-style sk-..., and similar) so Scrub catches a stray key even if a +// caller forgot to list it explicitly among secrets. +var apiKeyPattern = regexp.MustCompile(`sk-[A-Za-z0-9_-]{10,}`) + +// Scrub redacts every occurrence of each secret in s, plus anything +// matching apiKeyPattern, so a recorded fixture can never contain a live +// credential. +func Scrub(s string, secrets ...string) string { + for _, secret := range secrets { + if secret == "" { + continue + } + s = strings.ReplaceAll(s, secret, "REDACTED") + } + return apiKeyPattern.ReplaceAllString(s, "sk-REDACTED") +} + +// Exchange is one recorded request/response pair, already read into memory +// so a Recorder's caller can inspect and scrub it after the real +// RoundTrip completes. +type Exchange struct { + RequestBody []byte + ResponseBody []byte + StatusCode int +} + +// Recorder wraps an http.RoundTripper, forwarding every request unchanged +// but also calling OnExchange with the raw request/response bytes — +// letting a test capture a live exchange without altering how the request +// is actually sent. +type Recorder struct { + Transport http.RoundTripper + OnExchange func(Exchange) +} + +// RoundTrip implements http.RoundTripper. +func (r *Recorder) RoundTrip(req *http.Request) (*http.Response, error) { + var reqBody []byte + if req.Body != nil { + var err error + reqBody, err = io.ReadAll(req.Body) + if err != nil { + return nil, err + } + req.Body = io.NopCloser(bytes.NewReader(reqBody)) + } + + transport := r.Transport + if transport == nil { + transport = http.DefaultTransport + } + resp, err := transport.RoundTrip(req) + if err != nil { + return nil, err + } + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + _ = resp.Body.Close() + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + + if r.OnExchange != nil { + r.OnExchange(Exchange{RequestBody: reqBody, ResponseBody: respBody, StatusCode: resp.StatusCode}) + } + return resp, nil +} + +// maxDebugBodyBytes caps how much of a request/response body --log-level +// debug prints — long enough to be useful, short enough that a large tool +// result or file read doesn't flood the log. +const maxDebugBodyBytes = 4000 + +// Truncate caps s at max bytes, appending a marker if it was cut. +func Truncate(s string, limit int) string { + if len(s) <= limit { + return s + } + return s[:limit] + "...(truncated)" +} + +// LogJSON debug-logs v marshaled to JSON, scrubbed and truncated — the +// request/response dump 9.2's AC calls for. It skips the marshal work +// entirely when debug logging isn't enabled, so it's cheap to call +// unconditionally on every request. +func LogJSON(ctx context.Context, label string, v any, secrets ...string) { + if !slog.Default().Enabled(ctx, slog.LevelDebug) { + return + } + body, err := json.Marshal(v) + if err != nil { + return + } + LogBytes(ctx, label, body, secrets...) +} + +// LogBytes is LogJSON's counterpart for a caller that already has raw wire +// bytes rather than a Go value to marshal. +func LogBytes(ctx context.Context, label string, body []byte, secrets ...string) { + if !slog.Default().Enabled(ctx, slog.LevelDebug) { + return + } + slog.DebugContext(ctx, label, "body", Truncate(Scrub(string(body), secrets...), maxDebugBodyBytes)) +} diff --git a/internal/llm/recordutil/recordutil_test.go b/internal/llm/recordutil/recordutil_test.go new file mode 100644 index 0000000..f08095d --- /dev/null +++ b/internal/llm/recordutil/recordutil_test.go @@ -0,0 +1,143 @@ +package recordutil + +import ( + "bytes" + "context" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestScrubRedactsExactSecretsAndKeyShapedLiterals(t *testing.T) { + in := `{"header":"Bearer top-secret-value","key":"sk-ant-abc123XYZ456","other":"fine"}` + got := Scrub(in, "top-secret-value") + + if strings.Contains(got, "top-secret-value") { + t.Errorf("Scrub() = %q, still contains the exact secret", got) + } + if strings.Contains(got, "sk-ant-abc123XYZ456") { + t.Errorf("Scrub() = %q, still contains an sk-shaped key literal", got) + } + if !strings.Contains(got, "other") || !strings.Contains(got, "fine") { + t.Errorf("Scrub() = %q, want unrelated content left alone", got) + } +} + +func TestScrubIgnoresEmptySecrets(t *testing.T) { + got := Scrub("hello world", "", "world") + if got != "hello REDACTED" { + t.Errorf("Scrub() = %q, want empty secrets skipped and non-empty ones applied", got) + } +} + +func TestRecorderCapturesExchangeAndForwardsResponseUnchanged(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + if string(body) != `{"hello":"world"}` { + t.Errorf("server saw body = %q", body) + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"ok":true}`)) + })) + defer srv.Close() + + var captured Exchange + client := &http.Client{Transport: &Recorder{ + OnExchange: func(e Exchange) { captured = e }, + }} + + resp, err := client.Post(srv.URL, "application/json", bytes.NewReader([]byte(`{"hello":"world"}`))) + if err != nil { + t.Fatalf("Post() error = %v", err) + } + defer func() { _ = resp.Body.Close() }() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("ReadAll() error = %v", err) + } + if string(respBody) != `{"ok":true}` { + t.Errorf("response body (after recording) = %q, want the server's original body untouched", respBody) + } + + if string(captured.RequestBody) != `{"hello":"world"}` { + t.Errorf("captured.RequestBody = %q", captured.RequestBody) + } + if string(captured.ResponseBody) != `{"ok":true}` { + t.Errorf("captured.ResponseBody = %q", captured.ResponseBody) + } + if captured.StatusCode != http.StatusOK { + t.Errorf("captured.StatusCode = %d, want 200", captured.StatusCode) + } +} + +func TestTruncateLeavesShortStringsAlone(t *testing.T) { + if got := Truncate("short", 100); got != "short" { + t.Errorf("Truncate() = %q, want unchanged", got) + } +} + +func TestTruncateCutsLongStrings(t *testing.T) { + got := Truncate(strings.Repeat("x", 100), 10) + if len(got) <= 10 { + t.Fatalf("Truncate() = %q, want it to still show it was cut (len %d)", got, len(got)) + } + if !strings.HasPrefix(got, strings.Repeat("x", 10)) || !strings.Contains(got, "truncated") { + t.Errorf("Truncate() = %q, want the first 10 bytes followed by a truncation marker", got) + } +} + +func TestLogJSONRedactsSecretsAndRespectsLevel(t *testing.T) { + var buf bytes.Buffer + logger := slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelInfo})) + prevDefault := slog.Default() + slog.SetDefault(logger) + defer slog.SetDefault(prevDefault) + + type payload struct { + APIKey string `json:"api_key"` + } + LogJSON(context.Background(), "test request", payload{APIKey: "sk-ant-verysecret123"}) + if buf.Len() != 0 { + t.Errorf("LogJSON() wrote %q at Info level with debug disabled, want nothing", buf.String()) + } + + debugLogger := slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelDebug})) + slog.SetDefault(debugLogger) + LogJSON(context.Background(), "test request", payload{APIKey: "sk-ant-verysecret123"}) + + out := buf.String() + if strings.Contains(out, "verysecret123") { + t.Errorf("debug log = %q, still contains the secret", out) + } + if !strings.Contains(out, "test request") { + t.Errorf("debug log = %q, want the label present", out) + } +} + +func TestLogBytesRedactsExplicitSecret(t *testing.T) { + var buf bytes.Buffer + prevDefault := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelDebug}))) + defer slog.SetDefault(prevDefault) + + LogBytes(context.Background(), "test response", []byte(`{"token":"plain-secret-value"}`), "plain-secret-value") + + if strings.Contains(buf.String(), "plain-secret-value") { + t.Errorf("debug log = %q, still contains the exact secret", buf.String()) + } +} + +func TestEnabledReflectsRecordEnvVar(t *testing.T) { + t.Setenv("RECORD", "") + if Enabled() { + t.Error("Enabled() = true with RECORD unset, want false") + } + t.Setenv("RECORD", "1") + if !Enabled() { + t.Error("Enabled() = false with RECORD=1, want true") + } +} diff --git a/internal/tools/diff.go b/internal/tools/diff.go new file mode 100644 index 0000000..629fdad --- /dev/null +++ b/internal/tools/diff.go @@ -0,0 +1,81 @@ +package tools + +import ( + "fmt" + "strings" +) + +const maxDiffLines = 2000 + +// lineDiff renders a minimal line-based diff of oldText to newText for an +// approval preview. It falls back to a plain before/after dump for inputs +// too large for the O(n*m) LCS below to be worth computing. +func lineDiff(oldText, newText string) string { + if oldText == newText { + return "(no changes)" + } + + oldLines := strings.Split(oldText, "\n") + newLines := strings.Split(newText, "\n") + if len(oldLines) > maxDiffLines || len(newLines) > maxDiffLines { + return truncateBytes(fmt.Sprintf("--- old (%d lines) ---\n%s\n+++ new (%d lines) +++\n%s", + len(oldLines), oldText, len(newLines), newText), defaultByteBudget) + } + + lcs := longestCommonSubsequence(oldLines, newLines) + + var b strings.Builder + oi, ni, li := 0, 0, 0 + for oi < len(oldLines) || ni < len(newLines) { + switch { + case li < len(lcs) && oi < len(oldLines) && oldLines[oi] == lcs[li] && ni < len(newLines) && newLines[ni] == lcs[li]: + fmt.Fprintf(&b, " %s\n", oldLines[oi]) + oi++ + ni++ + li++ + case oi < len(oldLines) && (li >= len(lcs) || oldLines[oi] != lcs[li]): + fmt.Fprintf(&b, "- %s\n", oldLines[oi]) + oi++ + default: + fmt.Fprintf(&b, "+ %s\n", newLines[ni]) + ni++ + } + } + return truncateBytes(b.String(), defaultByteBudget) +} + +func longestCommonSubsequence(a, b []string) []string { + n, m := len(a), len(b) + dp := make([][]int, n+1) + for i := range dp { + dp[i] = make([]int, m+1) + } + for i := n - 1; i >= 0; i-- { + for j := m - 1; j >= 0; j-- { + switch { + case a[i] == b[j]: + dp[i][j] = dp[i+1][j+1] + 1 + case dp[i+1][j] >= dp[i][j+1]: + dp[i][j] = dp[i+1][j] + default: + dp[i][j] = dp[i][j+1] + } + } + } + + var lcs []string + i, j := 0, 0 + for i < n && j < m { + switch { + case a[i] == b[j]: + lcs = append(lcs, a[i]) + i++ + j++ + case dp[i+1][j] >= dp[i][j+1]: + i++ + default: + j++ + } + } + return lcs +} diff --git a/internal/tools/diff_test.go b/internal/tools/diff_test.go new file mode 100644 index 0000000..08a55cc --- /dev/null +++ b/internal/tools/diff_test.go @@ -0,0 +1,29 @@ +package tools + +import ( + "strings" + "testing" +) + +func TestLineDiffNoChanges(t *testing.T) { + if got := lineDiff("same", "same"); got != "(no changes)" { + t.Errorf("lineDiff() = %q, want %q", got, "(no changes)") + } +} + +func TestLineDiffShowsAddedRemovedAndContext(t *testing.T) { + got := lineDiff("a\nb\nc", "a\nX\nc") + for _, want := range []string{" a", "- b", "+ X", " c"} { + if !strings.Contains(got, want) { + t.Errorf("lineDiff() = %q, want it to contain %q", got, want) + } + } +} + +func TestLineDiffFallsBackForLargeInput(t *testing.T) { + big := strings.Repeat("line\n", maxDiffLines+1) + got := lineDiff("", big) + if !strings.Contains(got, "old (1 lines)") { + t.Errorf("lineDiff() = %q, want the fallback header", got[:200]) + } +} diff --git a/internal/tools/list_dir.go b/internal/tools/list_dir.go new file mode 100644 index 0000000..b94dd68 --- /dev/null +++ b/internal/tools/list_dir.go @@ -0,0 +1,105 @@ +package tools + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "sort" + "strings" +) + +const defaultEntryCap = 500 + +// ListDirInput is list_dir's argument shape. +type ListDirInput struct { + Path string `json:"path" jsonschema:"description=Directory path, relative to the workspace root"` + MaxDepth int `json:"max_depth,omitempty" jsonschema:"description=Maximum recursion depth; 0 lists only the given directory"` +} + +// ListDirTool lists directory entries from within a Sandbox. +type ListDirTool struct { + sandbox *Sandbox + schema json.RawMessage + entryCap int +} + +// NewListDirTool returns a ListDirTool confined to sandbox. +func NewListDirTool(sandbox *Sandbox) (*ListDirTool, error) { + schema, err := GenerateSchema(ListDirInput{}) + if err != nil { + return nil, fmt.Errorf("tools: list_dir schema: %w", err) + } + return &ListDirTool{sandbox: sandbox, schema: schema, entryCap: defaultEntryCap}, nil +} + +// Name implements tools.Tool. +func (t *ListDirTool) Name() string { return "list_dir" } + +// Description implements tools.Tool. +func (t *ListDirTool) Description() string { + return "List directory entries, optionally recursively, from within the workspace." +} + +// Schema implements tools.Tool. +func (t *ListDirTool) Schema() json.RawMessage { return t.schema } + +// Run lists the requested directory up to MaxDepth and entryCap entries, +// output truncated to defaultByteBudget. +func (t *ListDirTool) Run(_ context.Context, input json.RawMessage) (string, error) { + var in ListDirInput + if err := json.Unmarshal(input, &in); err != nil { + return "", fmt.Errorf("tools: list_dir: %w", err) + } + + resolved, err := t.sandbox.Resolve(in.Path) + if err != nil { + return "", err + } + + var lines []string + count := 0 + capped := false + + var walk func(dir string, depth int, prefix string) error + walk = func(dir string, depth int, prefix string) error { + entries, err := os.ReadDir(dir) + if err != nil { + return err + } + sort.Slice(entries, func(i, j int) bool { return entries[i].Name() < entries[j].Name() }) + + for _, e := range entries { + if count >= t.entryCap { + capped = true + return nil + } + rel := prefix + e.Name() + if e.IsDir() { + rel += "/" + } + lines = append(lines, rel) + count++ + + if e.IsDir() && depth < in.MaxDepth { + if err := walk(filepath.Join(dir, e.Name()), depth+1, rel); err != nil { + return err + } + } + } + return nil + } + + if err := walk(resolved, 0, ""); err != nil { + return "", fmt.Errorf("tools: list_dir %q: %w", in.Path, err) + } + + out := strings.Join(lines, "\n") + if capped { + out += fmt.Sprintf("\n... [entry cap of %d reached, output truncated] ...", t.entryCap) + } + return truncateBytes(out, defaultByteBudget), nil +} + +var _ Tool = (*ListDirTool)(nil) diff --git a/internal/tools/list_dir_test.go b/internal/tools/list_dir_test.go new file mode 100644 index 0000000..57b8834 --- /dev/null +++ b/internal/tools/list_dir_test.go @@ -0,0 +1,106 @@ +package tools + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "testing" +) + +func newTestListDirTool(t *testing.T, root string) *ListDirTool { + t.Helper() + sandbox, err := NewSandbox(root, nil) + if err != nil { + t.Fatalf("NewSandbox() error = %v", err) + } + tool, err := NewListDirTool(sandbox) + if err != nil { + t.Fatalf("NewListDirTool() error = %v", err) + } + return tool +} + +func TestListDirRespectsMaxDepth(t *testing.T) { + root := t.TempDir() + mustMkdirAll(t, filepath.Join(root, "a", "b")) + mustWriteFile(t, filepath.Join(root, "top.txt"), "x") + mustWriteFile(t, filepath.Join(root, "a", "mid.txt"), "x") + mustWriteFile(t, filepath.Join(root, "a", "b", "deep.txt"), "x") + + tool := newTestListDirTool(t, root) + + input, _ := json.Marshal(ListDirInput{Path: ".", MaxDepth: 0}) + got, err := tool.Run(context.Background(), input) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if strings.Contains(got, "mid.txt") || strings.Contains(got, "deep.txt") { + t.Errorf("Run() with MaxDepth=0 = %q, want no recursion", got) + } + if !strings.Contains(got, "top.txt") || !strings.Contains(got, "a/") { + t.Errorf("Run() with MaxDepth=0 = %q, want top-level entries", got) + } + + input, _ = json.Marshal(ListDirInput{Path: ".", MaxDepth: 2}) + got, err = tool.Run(context.Background(), input) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + for _, want := range []string{"top.txt", "a/", "a/mid.txt", "a/b/", "a/b/deep.txt"} { + if !strings.Contains(got, want) { + t.Errorf("Run() with MaxDepth=2 = %q, want it to contain %q", got, want) + } + } +} + +func TestListDirEntryCap(t *testing.T) { + root := t.TempDir() + for i := 0; i < 10; i++ { + mustWriteFile(t, filepath.Join(root, fmt.Sprintf("file%02d.txt", i)), "x") + } + + tool := newTestListDirTool(t, root) + tool.entryCap = 3 + + input, _ := json.Marshal(ListDirInput{Path: "."}) + got, err := tool.Run(context.Background(), input) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if !strings.Contains(got, "entry cap of 3 reached") { + t.Errorf("Run() = %q, want an entry-cap marker", got) + } + entryLines := strings.Split(strings.TrimSuffix(got, "\n... [entry cap of 3 reached, output truncated] ..."), "\n") + if len(entryLines) != 3 { + t.Errorf("Run() listed %d entries before the marker, want 3", len(entryLines)) + } +} + +func TestListDirRejectsSandboxEscape(t *testing.T) { + root := t.TempDir() + tool := newTestListDirTool(t, root) + + input, _ := json.Marshal(ListDirInput{Path: "../"}) + _, err := tool.Run(context.Background(), input) + if !errors.Is(err, ErrPathEscapesSandbox) { + t.Fatalf("Run() error = %v, want wrapping %v", err, ErrPathEscapesSandbox) + } +} + +func mustMkdirAll(t *testing.T, path string) { + t.Helper() + if err := os.MkdirAll(path, 0o755); err != nil { + t.Fatalf("MkdirAll(%q) error = %v", path, err) + } +} + +func mustWriteFile(t *testing.T, path, content string) { + t.Helper() + if err := os.WriteFile(path, []byte(content), 0o600); err != nil { + t.Fatalf("WriteFile(%q) error = %v", path, err) + } +} diff --git a/internal/tools/read_file.go b/internal/tools/read_file.go new file mode 100644 index 0000000..48d820f --- /dev/null +++ b/internal/tools/read_file.go @@ -0,0 +1,84 @@ +package tools + +import ( + "context" + "encoding/json" + "fmt" + "os" + "strings" +) + +// ReadFileInput is read_file's argument shape. +type ReadFileInput struct { + Path string `json:"path" jsonschema:"description=Path to the file, relative to the workspace root"` + StartLine int `json:"start_line,omitempty" jsonschema:"description=1-based first line to include; 0 means the start of the file"` + EndLine int `json:"end_line,omitempty" jsonschema:"description=1-based last line to include, inclusive; 0 means the end of the file"` +} + +// ReadFileTool reads a file's contents, or a line range of it, from within +// a Sandbox. +type ReadFileTool struct { + sandbox *Sandbox + schema json.RawMessage +} + +// NewReadFileTool returns a ReadFileTool confined to sandbox. +func NewReadFileTool(sandbox *Sandbox) (*ReadFileTool, error) { + schema, err := GenerateSchema(ReadFileInput{}) + if err != nil { + return nil, fmt.Errorf("tools: read_file schema: %w", err) + } + return &ReadFileTool{sandbox: sandbox, schema: schema}, nil +} + +// Name implements tools.Tool. +func (t *ReadFileTool) Name() string { return "read_file" } + +// Description implements tools.Tool. +func (t *ReadFileTool) Description() string { + return "Read a file's contents, optionally a line range, from within the workspace." +} + +// Schema implements tools.Tool. +func (t *ReadFileTool) Schema() json.RawMessage { return t.schema } + +// Run reads the requested file and returns its contents (or line range), +// truncated to defaultByteBudget. +func (t *ReadFileTool) Run(_ context.Context, input json.RawMessage) (string, error) { + var in ReadFileInput + if err := json.Unmarshal(input, &in); err != nil { + return "", fmt.Errorf("tools: read_file: %w", err) + } + + resolved, err := t.sandbox.Resolve(in.Path) + if err != nil { + return "", err + } + + data, err := os.ReadFile(resolved) + if err != nil { + return "", fmt.Errorf("tools: read_file %q: %w", in.Path, err) + } + + text := string(data) + if in.StartLine > 0 || in.EndLine > 0 { + text = sliceLines(text, in.StartLine, in.EndLine) + } + return truncateBytes(text, defaultByteBudget), nil +} + +func sliceLines(text string, start, end int) string { + lines := strings.Split(text, "\n") + if start <= 0 { + start = 1 + } + if end <= 0 || end > len(lines) { + end = len(lines) + } + if start > end || start > len(lines) { + return "" + } + return strings.Join(lines[start-1:end], "\n") +} + +var _ Tool = (*ReadFileTool)(nil) diff --git a/internal/tools/read_file_test.go b/internal/tools/read_file_test.go new file mode 100644 index 0000000..411bd67 --- /dev/null +++ b/internal/tools/read_file_test.go @@ -0,0 +1,90 @@ +package tools + +import ( + "context" + "encoding/json" + "errors" + "os" + "path/filepath" + "strings" + "testing" +) + +func newTestReadFileTool(t *testing.T, root string) *ReadFileTool { + t.Helper() + sandbox, err := NewSandbox(root, nil) + if err != nil { + t.Fatalf("NewSandbox() error = %v", err) + } + tool, err := NewReadFileTool(sandbox) + if err != nil { + t.Fatalf("NewReadFileTool() error = %v", err) + } + return tool +} + +func TestReadFileReturnsFullContents(t *testing.T) { + root := t.TempDir() + if err := os.WriteFile(filepath.Join(root, "a.txt"), []byte("line1\nline2\nline3"), 0o600); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + tool := newTestReadFileTool(t, root) + + input, _ := json.Marshal(ReadFileInput{Path: "a.txt"}) + got, err := tool.Run(context.Background(), input) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if got != "line1\nline2\nline3" { + t.Errorf("Run() = %q, want full file contents", got) + } +} + +func TestReadFileLineRange(t *testing.T) { + root := t.TempDir() + if err := os.WriteFile(filepath.Join(root, "a.txt"), []byte("line1\nline2\nline3\nline4"), 0o600); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + tool := newTestReadFileTool(t, root) + + input, _ := json.Marshal(ReadFileInput{Path: "a.txt", StartLine: 2, EndLine: 3}) + got, err := tool.Run(context.Background(), input) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if got != "line2\nline3" { + t.Errorf("Run() = %q, want %q", got, "line2\nline3") + } +} + +func TestReadFileTruncatesLargeFile(t *testing.T) { + root := t.TempDir() + big := strings.Repeat("X", 10*1024*1024) + if err := os.WriteFile(filepath.Join(root, "big.txt"), []byte(big), 0o600); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + tool := newTestReadFileTool(t, root) + + input, _ := json.Marshal(ReadFileInput{Path: "big.txt"}) + got, err := tool.Run(context.Background(), input) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if len(got) > defaultByteBudget { + t.Errorf("Run() returned %d bytes, want <= %d", len(got), defaultByteBudget) + } + if !strings.Contains(got, "elided") { + t.Error("Run() output missing elision marker") + } +} + +func TestReadFileRejectsSandboxEscape(t *testing.T) { + root := t.TempDir() + tool := newTestReadFileTool(t, root) + + input, _ := json.Marshal(ReadFileInput{Path: "../outside.txt"}) + _, err := tool.Run(context.Background(), input) + if !errors.Is(err, ErrPathEscapesSandbox) { + t.Fatalf("Run() error = %v, want wrapping %v", err, ErrPathEscapesSandbox) + } +} diff --git a/internal/tools/registry.go b/internal/tools/registry.go index 6c14bf3..1c6a604 100644 --- a/internal/tools/registry.go +++ b/internal/tools/registry.go @@ -5,7 +5,10 @@ import ( "encoding/json" "errors" "fmt" + "sort" "sync" + + "github.com/reno/pico-code/internal/llm" ) // ErrToolNotFound is returned when a name has no registered Tool. @@ -54,6 +57,24 @@ func (r *Registry) Get(name string) (Tool, error) { return t, nil } +// Definitions returns the canonical llm.ToolDefinition for every registered +// tool, sorted by name so the request payload a caller builds from it is +// deterministic across runs. +func (r *Registry) Definitions() []llm.ToolDefinition { + r.mu.RLock() + defer r.mu.RUnlock() + defs := make([]llm.ToolDefinition, 0, len(r.tools)) + for _, t := range r.tools { + defs = append(defs, llm.ToolDefinition{ + Name: t.Name(), + Description: t.Description(), + InputSchema: t.Schema(), + }) + } + sort.Slice(defs, func(i, j int) bool { return defs[i].Name < defs[j].Name }) + return defs +} + // Run looks up name, validates input against its Schema(), and only then // calls Run — so a tool implementation never has to defend against a // malformed or incomplete argument object. diff --git a/internal/tools/run_command.go b/internal/tools/run_command.go new file mode 100644 index 0000000..c216a3a --- /dev/null +++ b/internal/tools/run_command.go @@ -0,0 +1,130 @@ +package tools + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "os/exec" + "strings" + "syscall" + "time" +) + +// ErrCommandNotAllowed is returned when the requested binary is not in the +// configured allowlist. +var ErrCommandNotAllowed = errors.New("tools: command not allowed") + +// ErrShellMetacharacter is returned when the command or an argument +// contains a character that would be meaningful to a shell — run_command +// never invokes one (CLAUDE.md: "no shell interpretation"), so this always +// indicates an attempt to smuggle one in. +var ErrShellMetacharacter = errors.New("tools: shell metacharacter not allowed") + +const shellMetaChars = "|&;<>()$`\\\"'*?[]#~!{}\n" + +// RunCommandInput is run_command's argument shape. +type RunCommandInput struct { + Command string `json:"command" jsonschema:"description=The binary to execute; must be in the configured allowlist"` + Args []string `json:"args,omitempty" jsonschema:"description=Arguments passed to the command, unexpanded"` +} + +// RunCommandTool runs an allowlisted binary directly (exec.CommandContext, +// no shell), with a timeout and output truncation. +type RunCommandTool struct { + schema json.RawMessage + allowlist map[string]bool + timeout time.Duration +} + +// NewRunCommandTool returns a RunCommandTool that only executes binaries +// named in allowlist, each bounded by timeout (0 = unlimited). +func NewRunCommandTool(allowlist []string, timeout time.Duration) (*RunCommandTool, error) { + schema, err := GenerateSchema(RunCommandInput{}) + if err != nil { + return nil, fmt.Errorf("tools: run_command schema: %w", err) + } + allowed := make(map[string]bool, len(allowlist)) + for _, b := range allowlist { + allowed[b] = true + } + return &RunCommandTool{schema: schema, allowlist: allowed, timeout: timeout}, nil +} + +// Name implements tools.Tool. +func (t *RunCommandTool) Name() string { return "run_command" } + +// Description implements tools.Tool. +func (t *RunCommandTool) Description() string { + return "Run an allowlisted command with no shell interpretation." +} + +// Schema implements tools.Tool. +func (t *RunCommandTool) Schema() json.RawMessage { return t.schema } + +// NeedsApproval implements tools.ApprovalRequired: every call to +// run_command needs sign-off. +func (t *RunCommandTool) NeedsApproval() bool { return true } + +// Run executes the requested command, refusing anything outside the +// allowlist or containing a shell metacharacter before starting it. The +// result always reports an exit code — a nonzero exit is data the model +// sees, not a tool failure — except when the process is killed by timeout +// or never starts. +func (t *RunCommandTool) Run(ctx context.Context, input json.RawMessage) (string, error) { + var in RunCommandInput + if err := json.Unmarshal(input, &in); err != nil { + return "", fmt.Errorf("tools: run_command: %w", err) + } + + if !t.allowlist[in.Command] { + return "", fmt.Errorf("%w: %q", ErrCommandNotAllowed, in.Command) + } + if strings.ContainsAny(in.Command, shellMetaChars) { + return "", fmt.Errorf("%w: in command %q", ErrShellMetacharacter, in.Command) + } + for _, arg := range in.Args { + if strings.ContainsAny(arg, shellMetaChars) { + return "", fmt.Errorf("%w: in argument %q", ErrShellMetacharacter, arg) + } + } + + runCtx := ctx + if t.timeout > 0 { + var cancel context.CancelFunc + runCtx, cancel = context.WithTimeout(ctx, t.timeout) + defer cancel() + } + + cmd := exec.CommandContext(runCtx, in.Command, in.Args...) + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + cmd.Cancel = func() error { + return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) + } + var out bytes.Buffer + cmd.Stdout = &out + cmd.Stderr = &out + + runErr := cmd.Run() + + exitCode := 0 + var exitErr *exec.ExitError + switch { + case runCtx.Err() != nil: + // A signal-killed process also surfaces as an *exec.ExitError, so + // this must be checked before that case or a timeout gets + // misreported as a normal nonzero exit. + return truncateBytes(out.String(), defaultByteBudget), fmt.Errorf("tools: run_command %q timed out: %w", in.Command, runCtx.Err()) + case runErr == nil: + case errors.As(runErr, &exitErr): + exitCode = exitErr.ExitCode() + default: + return "", fmt.Errorf("tools: run_command %q: %w", in.Command, runErr) + } + + result := fmt.Sprintf("exit code: %d\n%s", exitCode, truncateBytes(out.String(), defaultByteBudget)) + return result, nil +} + +var _ ApprovalRequired = (*RunCommandTool)(nil) diff --git a/internal/tools/run_command_test.go b/internal/tools/run_command_test.go new file mode 100644 index 0000000..a2ad02d --- /dev/null +++ b/internal/tools/run_command_test.go @@ -0,0 +1,97 @@ +package tools + +import ( + "context" + "encoding/json" + "errors" + "strings" + "testing" + "time" +) + +func TestRunCommandRejectsBlockedBinary(t *testing.T) { + tool, err := NewRunCommandTool([]string{"echo"}, time.Second) + if err != nil { + t.Fatalf("NewRunCommandTool() error = %v", err) + } + input, _ := json.Marshal(RunCommandInput{Command: "rm", Args: []string{"-rf", "/"}}) + + _, err = tool.Run(context.Background(), input) + if !errors.Is(err, ErrCommandNotAllowed) { + t.Fatalf("Run() error = %v, want wrapping %v", err, ErrCommandNotAllowed) + } +} + +func TestRunCommandRejectsShellMetacharacters(t *testing.T) { + tool, err := NewRunCommandTool([]string{"echo"}, time.Second) + if err != nil { + t.Fatalf("NewRunCommandTool() error = %v", err) + } + input, _ := json.Marshal(RunCommandInput{Command: "echo", Args: []string{"$(whoami)"}}) + + _, err = tool.Run(context.Background(), input) + if !errors.Is(err, ErrShellMetacharacter) { + t.Fatalf("Run() error = %v, want wrapping %v", err, ErrShellMetacharacter) + } +} + +func TestRunCommandReportsExitCode(t *testing.T) { + tool, err := NewRunCommandTool([]string{"sh"}, time.Second) + if err != nil { + t.Fatalf("NewRunCommandTool() error = %v", err) + } + input, _ := json.Marshal(RunCommandInput{Command: "sh", Args: []string{"-c", "exit 3"}}) + + got, err := tool.Run(context.Background(), input) + if err != nil { + t.Fatalf("Run() error = %v, want a nonzero exit reported as data, not a tool error", err) + } + if !strings.Contains(got, "exit code: 3") { + t.Errorf("Run() = %q, want it to report exit code 3", got) + } +} + +func TestRunCommandTruncatesOutput(t *testing.T) { + // No shell means no pipes either, so this generates 100000 bytes via a + // single allowlisted binary and its flags rather than "yes | head". + tool, err := NewRunCommandTool([]string{"head"}, 5*time.Second) + if err != nil { + t.Fatalf("NewRunCommandTool() error = %v", err) + } + input, _ := json.Marshal(RunCommandInput{Command: "head", Args: []string{"-c", "100000", "/dev/zero"}}) + + got, err := tool.Run(context.Background(), input) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if len(got) > defaultByteBudget+64 { + t.Errorf("Run() output length = %d, want roughly within the %d byte budget", len(got), defaultByteBudget) + } + if !strings.Contains(got, "elided") { + t.Errorf("Run() = %q, want an elision marker", got[:200]) + } +} + +// TestRunCommandTimeoutKillsTheProcess proves the timeout actually +// terminates a hung child rather than Run() blocking until it exits +// naturally. Setpgid+Cancel (see run_command.go) target the negative PID +// specifically so descendants die too, but exercising an actual +// grandchild would need a non-shell helper binary to spawn one — this +// verifies the direct child dies promptly instead. +func TestRunCommandTimeoutKillsTheProcess(t *testing.T) { + tool, err := NewRunCommandTool([]string{"sleep"}, 50*time.Millisecond) + if err != nil { + t.Fatalf("NewRunCommandTool() error = %v", err) + } + input, _ := json.Marshal(RunCommandInput{Command: "sleep", Args: []string{"30"}}) + + start := time.Now() + _, err = tool.Run(context.Background(), input) + elapsed := time.Since(start) + if err == nil { + t.Fatal("Run() error = nil, want a timeout error") + } + if elapsed > 2*time.Second { + t.Errorf("Run() took %s, want it to return promptly once the 50ms timeout fires, not wait out the 30s sleep", elapsed) + } +} diff --git a/internal/tools/sandbox.go b/internal/tools/sandbox.go new file mode 100644 index 0000000..dc4b95e --- /dev/null +++ b/internal/tools/sandbox.go @@ -0,0 +1,117 @@ +package tools + +import ( + "errors" + "fmt" + "os" + "path/filepath" + "strings" +) + +// ErrPathEscapesSandbox is returned by Sandbox.Resolve for a path that +// falls outside the workspace root, directly or via a symlink. +var ErrPathEscapesSandbox = errors.New("tools: path escapes workspace sandbox") + +// ErrPathDenied is returned by Sandbox.Resolve for a path matching a deny +// glob. +var ErrPathDenied = errors.New("tools: path is denied") + +// defaultDenyGlobs are always denied, regardless of caller-supplied globs — +// CLAUDE.md: "Deny reads of .env, .git/config, *.pem, id_*". +var defaultDenyGlobs = []string{".env", ".git/config", "*.pem", "id_*"} + +// Sandbox confines filesystem tools to a workspace root, shared by every +// tool that touches disk. +type Sandbox struct { + root string + denyGlobs []string +} + +// NewSandbox resolves root to an absolute, symlink-free path and returns a +// Sandbox scoped to it. denyGlobs are checked in addition to +// defaultDenyGlobs. +func NewSandbox(root string, denyGlobs []string) (*Sandbox, error) { + abs, err := filepath.Abs(root) + if err != nil { + return nil, fmt.Errorf("tools: sandbox root %q: %w", root, err) + } + resolved, err := filepath.EvalSymlinks(abs) + if err != nil { + return nil, fmt.Errorf("tools: sandbox root %q: %w", root, err) + } + all := make([]string, 0, len(defaultDenyGlobs)+len(denyGlobs)) + all = append(all, defaultDenyGlobs...) + all = append(all, denyGlobs...) + return &Sandbox{root: resolved, denyGlobs: all}, nil +} + +// Resolve validates path (relative to the sandbox root, or absolute) and +// returns its resolved absolute form. It rejects anything that escapes the +// root — including via a symlink whose target lies outside it — or matches +// a deny glob. path need not exist: only its existing ancestors are +// symlink-resolved, so a not-yet-created write target still resolves +// correctly. +func (s *Sandbox) Resolve(path string) (string, error) { + if path == "" { + return "", errors.New("tools: empty path") + } + + joined := path + if !filepath.IsAbs(path) { + joined = filepath.Join(s.root, path) + } + joined = filepath.Clean(joined) + + resolved, err := resolveExistingSymlinks(joined) + if err != nil { + return "", fmt.Errorf("tools: resolve %q: %w", path, err) + } + if !s.contains(resolved) { + return "", fmt.Errorf("%w: %q", ErrPathEscapesSandbox, path) + } + + if s.denied(resolved) { + return "", fmt.Errorf("%w: %q", ErrPathDenied, path) + } + + return resolved, nil +} + +func (s *Sandbox) contains(p string) bool { + return p == s.root || strings.HasPrefix(p, s.root+string(filepath.Separator)) +} + +func (s *Sandbox) denied(resolved string) bool { + rel, err := filepath.Rel(s.root, resolved) + if err != nil { + return false + } + base := filepath.Base(resolved) + for _, pattern := range s.denyGlobs { + if ok, _ := filepath.Match(pattern, base); ok { + return true + } + if ok, _ := filepath.Match(pattern, rel); ok { + return true + } + } + return false +} + +// resolveExistingSymlinks walks up from path to the longest existing +// ancestor, symlink-resolves that ancestor, and rejoins the remaining +// (possibly not-yet-existing) suffix. +func resolveExistingSymlinks(path string) (string, error) { + if _, err := os.Lstat(path); err == nil { + return filepath.EvalSymlinks(path) + } + parent := filepath.Dir(path) + if parent == path { + return path, nil + } + resolvedParent, err := resolveExistingSymlinks(parent) + if err != nil { + return "", err + } + return filepath.Join(resolvedParent, filepath.Base(path)), nil +} diff --git a/internal/tools/sandbox_test.go b/internal/tools/sandbox_test.go new file mode 100644 index 0000000..bed11dc --- /dev/null +++ b/internal/tools/sandbox_test.go @@ -0,0 +1,114 @@ +package tools + +import ( + "errors" + "os" + "path/filepath" + "testing" +) + +func TestSandboxResolveRejections(t *testing.T) { + root := t.TempDir() + if err := os.WriteFile(filepath.Join(root, ".env"), []byte("SECRET=1"), 0o600); err != nil { + t.Fatalf("WriteFile(.env) error = %v", err) + } + if err := os.MkdirAll(filepath.Join(root, "sub"), 0o755); err != nil { + t.Fatalf("MkdirAll(sub) error = %v", err) + } + if err := os.WriteFile(filepath.Join(root, "sub", "ok.txt"), []byte("hi"), 0o600); err != nil { + t.Fatalf("WriteFile(sub/ok.txt) error = %v", err) + } + + outside := t.TempDir() + if err := os.WriteFile(filepath.Join(outside, "secret.txt"), []byte("nope"), 0o600); err != nil { + t.Fatalf("WriteFile(outside) error = %v", err) + } + if err := os.Symlink(filepath.Join(outside, "secret.txt"), filepath.Join(root, "escape_link")); err != nil { + t.Fatalf("Symlink() error = %v", err) + } + + sandbox, err := NewSandbox(root, []string{"custom_deny*"}) + if err != nil { + t.Fatalf("NewSandbox() error = %v", err) + } + if err := os.WriteFile(filepath.Join(root, "custom_deny.txt"), []byte("x"), 0o600); err != nil { + t.Fatalf("WriteFile(custom_deny.txt) error = %v", err) + } + + tests := []struct { + name string + path string + wantErr error + }{ + {"dot-dot escape", "../outside.txt", ErrPathEscapesSandbox}, + {"absolute path outside root", filepath.Join(outside, "secret.txt"), ErrPathEscapesSandbox}, + {"symlink escapes root", "escape_link", ErrPathEscapesSandbox}, + {"default deny glob", ".env", ErrPathDenied}, + {"caller deny glob", "custom_deny.txt", ErrPathDenied}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := sandbox.Resolve(tt.path) + if !errors.Is(err, tt.wantErr) { + t.Fatalf("Resolve(%q) error = %v, want wrapping %v", tt.path, err, tt.wantErr) + } + if err.Error() == "" { + t.Fatal("expected a descriptive error message") + } + }) + } +} + +func TestSandboxResolveAcceptsPathsInsideRoot(t *testing.T) { + root := t.TempDir() + if err := os.MkdirAll(filepath.Join(root, "sub"), 0o755); err != nil { + t.Fatalf("MkdirAll() error = %v", err) + } + if err := os.WriteFile(filepath.Join(root, "sub", "ok.txt"), []byte("hi"), 0o600); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + + sandbox, err := NewSandbox(root, nil) + if err != nil { + t.Fatalf("NewSandbox() error = %v", err) + } + + tests := []string{"sub/ok.txt", filepath.Join(root, "sub", "ok.txt"), "sub/../sub/ok.txt"} + for _, path := range tests { + t.Run(path, func(t *testing.T) { + resolved, err := sandbox.Resolve(path) + if err != nil { + t.Fatalf("Resolve(%q) error = %v", path, err) + } + want, err := filepath.EvalSymlinks(filepath.Join(root, "sub", "ok.txt")) + if err != nil { + t.Fatalf("EvalSymlinks() error = %v", err) + } + if resolved != want { + t.Errorf("Resolve(%q) = %q, want %q", path, resolved, want) + } + }) + } +} + +func TestSandboxResolveAllowsNotYetExistingPathInsideRoot(t *testing.T) { + root := t.TempDir() + sandbox, err := NewSandbox(root, nil) + if err != nil { + t.Fatalf("NewSandbox() error = %v", err) + } + + resolved, err := sandbox.Resolve("new/nested/file.txt") + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + resolvedRoot, err := filepath.EvalSymlinks(root) + if err != nil { + t.Fatalf("EvalSymlinks() error = %v", err) + } + want := filepath.Join(resolvedRoot, "new", "nested", "file.txt") + if resolved != want { + t.Errorf("Resolve() = %q, want %q", resolved, want) + } +} diff --git a/internal/tools/tool.go b/internal/tools/tool.go index f50380d..05e0e55 100644 --- a/internal/tools/tool.go +++ b/internal/tools/tool.go @@ -18,3 +18,20 @@ type Tool interface { Schema() json.RawMessage Run(ctx context.Context, input json.RawMessage) (string, error) } + +// ApprovalRequired is implemented by a Tool whose side effects need +// explicit user sign-off before Run executes. CLAUDE.md invariant 5 makes +// the agent loop the sole owner of approval policy, so a Tool only +// declares whether it needs approval — it never prompts. +type ApprovalRequired interface { + Tool + NeedsApproval() bool +} + +// Previewable is implemented by a Tool that can describe its effect ahead +// of Run — e.g. write_file's diff against the file's current contents — +// so the loop can show it alongside an approval prompt. +type Previewable interface { + Tool + Preview(ctx context.Context, input json.RawMessage) (string, error) +} diff --git a/internal/tools/truncate.go b/internal/tools/truncate.go new file mode 100644 index 0000000..521c40e --- /dev/null +++ b/internal/tools/truncate.go @@ -0,0 +1,22 @@ +package tools + +import "fmt" + +const defaultByteBudget = 32 * 1024 + +// truncateBytes returns s unchanged if it's within budget; otherwise a +// head+tail slice with an elision marker in between, kept to budget so a +// large tool output can never blow the context window. +func truncateBytes(s string, budget int) string { + if len(s) <= budget { + return s + } + marker := fmt.Sprintf("\n\n... [%d bytes elided] ...\n\n", len(s)-budget) + keep := budget - len(marker) + if keep < 0 { + keep = 0 + } + head := keep / 2 + tail := keep - head + return s[:head] + marker + s[len(s)-tail:] +} diff --git a/internal/tools/truncate_test.go b/internal/tools/truncate_test.go new file mode 100644 index 0000000..bd1070c --- /dev/null +++ b/internal/tools/truncate_test.go @@ -0,0 +1,32 @@ +package tools + +import ( + "strings" + "testing" +) + +func TestTruncateBytesUnderBudgetUnchanged(t *testing.T) { + s := "hello world" + if got := truncateBytes(s, 1024); got != s { + t.Errorf("truncateBytes() = %q, want unchanged %q", got, s) + } +} + +func TestTruncateBytesOverBudget(t *testing.T) { + s := strings.Repeat("A", 10*1024*1024) + budget := 32 * 1024 + + got := truncateBytes(s, budget) + if len(got) > budget { + t.Fatalf("truncateBytes() length = %d, want <= %d", len(got), budget) + } + if !strings.Contains(got, "elided") { + t.Errorf("truncateBytes() = %q, want it to contain an elision marker", got[:200]) + } + if !strings.HasPrefix(got, "AAAA") { + t.Error("truncateBytes() should keep a head slice of the original content") + } + if !strings.HasSuffix(got, "AAAA") { + t.Error("truncateBytes() should keep a tail slice of the original content") + } +} diff --git a/internal/tools/write_file.go b/internal/tools/write_file.go new file mode 100644 index 0000000..a1a9a27 --- /dev/null +++ b/internal/tools/write_file.go @@ -0,0 +1,109 @@ +package tools + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" +) + +// WriteFileInput is write_file's argument shape. +type WriteFileInput struct { + Path string `json:"path" jsonschema:"description=Path to the file, relative to the workspace root"` + Content string `json:"content" jsonschema:"description=The file's new full contents"` +} + +// WriteFileTool atomically writes a file's full contents within a Sandbox. +type WriteFileTool struct { + sandbox *Sandbox + schema json.RawMessage +} + +// NewWriteFileTool returns a WriteFileTool confined to sandbox. +func NewWriteFileTool(sandbox *Sandbox) (*WriteFileTool, error) { + schema, err := GenerateSchema(WriteFileInput{}) + if err != nil { + return nil, fmt.Errorf("tools: write_file schema: %w", err) + } + return &WriteFileTool{sandbox: sandbox, schema: schema}, nil +} + +// Name implements tools.Tool. +func (t *WriteFileTool) Name() string { return "write_file" } + +// Description implements tools.Tool. +func (t *WriteFileTool) Description() string { + return "Write a file's full contents atomically, creating or overwriting it, within the workspace." +} + +// Schema implements tools.Tool. +func (t *WriteFileTool) Schema() json.RawMessage { return t.schema } + +// NeedsApproval implements tools.ApprovalRequired: every write needs +// sign-off. +func (t *WriteFileTool) NeedsApproval() bool { return true } + +// Preview implements tools.Previewable: a line diff of the current file +// contents (empty if it doesn't exist yet) against the proposed content. +func (t *WriteFileTool) Preview(_ context.Context, input json.RawMessage) (string, error) { + var in WriteFileInput + if err := json.Unmarshal(input, &in); err != nil { + return "", fmt.Errorf("tools: write_file preview: %w", err) + } + resolved, err := t.sandbox.Resolve(in.Path) + if err != nil { + return "", err + } + old, err := os.ReadFile(resolved) + if err != nil && !os.IsNotExist(err) { + return "", fmt.Errorf("tools: write_file preview %q: %w", in.Path, err) + } + return lineDiff(string(old), in.Content), nil +} + +// Run writes in.Content to in.Path atomically: it writes to a temp file in +// the same directory, then renames it over the target, so a reader never +// observes a partially written file and an interrupted run leaves the +// original untouched. +func (t *WriteFileTool) Run(ctx context.Context, input json.RawMessage) (string, error) { + var in WriteFileInput + if err := json.Unmarshal(input, &in); err != nil { + return "", fmt.Errorf("tools: write_file: %w", err) + } + + resolved, err := t.sandbox.Resolve(in.Path) + if err != nil { + return "", err + } + + dir := filepath.Dir(resolved) + tmp, err := os.CreateTemp(dir, ".pico-write-*") + if err != nil { + return "", fmt.Errorf("tools: write_file %q: create temp: %w", in.Path, err) + } + tmpPath := tmp.Name() + defer os.Remove(tmpPath) //nolint:errcheck // no-op once the rename below succeeds + + if _, err := tmp.WriteString(in.Content); err != nil { + _ = tmp.Close() + return "", fmt.Errorf("tools: write_file %q: %w", in.Path, err) + } + if err := tmp.Close(); err != nil { + return "", fmt.Errorf("tools: write_file %q: %w", in.Path, err) + } + + if err := ctx.Err(); err != nil { + return "", fmt.Errorf("tools: write_file %q: %w", in.Path, err) + } + + if err := os.Rename(tmpPath, resolved); err != nil { + return "", fmt.Errorf("tools: write_file %q: rename: %w", in.Path, err) + } + return fmt.Sprintf("wrote %d bytes to %s", len(in.Content), in.Path), nil +} + +var ( + _ ApprovalRequired = (*WriteFileTool)(nil) + _ Previewable = (*WriteFileTool)(nil) +) diff --git a/internal/tools/write_file_test.go b/internal/tools/write_file_test.go new file mode 100644 index 0000000..2ca14cb --- /dev/null +++ b/internal/tools/write_file_test.go @@ -0,0 +1,161 @@ +package tools + +import ( + "context" + "encoding/json" + "errors" + "os" + "path/filepath" + "strings" + "testing" +) + +func newTestWriteFileTool(t *testing.T, root string) *WriteFileTool { + t.Helper() + sandbox, err := NewSandbox(root, nil) + if err != nil { + t.Fatalf("NewSandbox() error = %v", err) + } + tool, err := NewWriteFileTool(sandbox) + if err != nil { + t.Fatalf("NewWriteFileTool() error = %v", err) + } + return tool +} + +func TestWriteFileCreatesNewFile(t *testing.T) { + root := t.TempDir() + tool := newTestWriteFileTool(t, root) + + input, _ := json.Marshal(WriteFileInput{Path: "new.txt", Content: "hello"}) + if _, err := tool.Run(context.Background(), input); err != nil { + t.Fatalf("Run() error = %v", err) + } + + got, err := os.ReadFile(filepath.Join(root, "new.txt")) + if err != nil { + t.Fatalf("ReadFile() error = %v", err) + } + if string(got) != "hello" { + t.Errorf("file contents = %q, want %q", got, "hello") + } +} + +func TestWriteFileOverwritesExistingFile(t *testing.T) { + root := t.TempDir() + path := filepath.Join(root, "existing.txt") + if err := os.WriteFile(path, []byte("old"), 0o600); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + tool := newTestWriteFileTool(t, root) + + input, _ := json.Marshal(WriteFileInput{Path: "existing.txt", Content: "new"}) + if _, err := tool.Run(context.Background(), input); err != nil { + t.Fatalf("Run() error = %v", err) + } + + got, err := os.ReadFile(path) + if err != nil { + t.Fatalf("ReadFile() error = %v", err) + } + if string(got) != "new" { + t.Errorf("file contents = %q, want %q", got, "new") + } +} + +func TestWriteFileRejectsSandboxEscape(t *testing.T) { + root := t.TempDir() + tool := newTestWriteFileTool(t, root) + + input, _ := json.Marshal(WriteFileInput{Path: "../outside.txt", Content: "x"}) + _, err := tool.Run(context.Background(), input) + if !errors.Is(err, ErrPathEscapesSandbox) { + t.Fatalf("Run() error = %v, want wrapping %v", err, ErrPathEscapesSandbox) + } +} + +// TestWriteFileInterruptedRunLeavesOriginalUntouched is 5.4's atomicity AC: +// if ctx is cancelled after the temp file is written but before the commit +// (rename), the write is aborted and the target file is byte-identical to +// what it was before the call. +func TestWriteFileInterruptedRunLeavesOriginalUntouched(t *testing.T) { + root := t.TempDir() + path := filepath.Join(root, "existing.txt") + if err := os.WriteFile(path, []byte("original"), 0o600); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + tool := newTestWriteFileTool(t, root) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + input, _ := json.Marshal(WriteFileInput{Path: "existing.txt", Content: "interrupted write"}) + _, err := tool.Run(ctx, input) + if err == nil { + t.Fatal("Run() error = nil, want an error for a cancelled ctx") + } + + got, err := os.ReadFile(path) + if err != nil { + t.Fatalf("ReadFile() error = %v", err) + } + if string(got) != "original" { + t.Errorf("file contents = %q, want the original %q untouched", got, "original") + } + + entries, err := os.ReadDir(root) + if err != nil { + t.Fatalf("ReadDir() error = %v", err) + } + for _, e := range entries { + if strings.HasPrefix(e.Name(), ".pico-write-") { + t.Errorf("leftover temp file %q was not cleaned up", e.Name()) + } + } +} + +func TestWriteFilePreviewShowsDiff(t *testing.T) { + root := t.TempDir() + path := filepath.Join(root, "existing.txt") + if err := os.WriteFile(path, []byte("line1\nline2\nline3"), 0o600); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + tool := newTestWriteFileTool(t, root) + + input, _ := json.Marshal(WriteFileInput{Path: "existing.txt", Content: "line1\nCHANGED\nline3"}) + got, err := tool.Preview(context.Background(), input) + if err != nil { + t.Fatalf("Preview() error = %v", err) + } + if !strings.Contains(got, "- line2") { + t.Errorf("Preview() = %q, want it to show the removed line", got) + } + if !strings.Contains(got, "+ CHANGED") { + t.Errorf("Preview() = %q, want it to show the added line", got) + } + if !strings.Contains(got, " line1") || !strings.Contains(got, " line3") { + t.Errorf("Preview() = %q, want unchanged lines shown as context", got) + } +} + +func TestWriteFilePreviewForNewFile(t *testing.T) { + root := t.TempDir() + tool := newTestWriteFileTool(t, root) + + input, _ := json.Marshal(WriteFileInput{Path: "new.txt", Content: "hello"}) + got, err := tool.Preview(context.Background(), input) + if err != nil { + t.Fatalf("Preview() error = %v", err) + } + if !strings.Contains(got, "+ hello") { + t.Errorf("Preview() = %q, want the new content shown as added", got) + } +} + +func TestWriteFileNeedsApproval(t *testing.T) { + root := t.TempDir() + tool := newTestWriteFileTool(t, root) + if !tool.NeedsApproval() { + t.Error("NeedsApproval() = false, want true") + } +} diff --git a/internal/ui/model.go b/internal/ui/model.go new file mode 100644 index 0000000..2488b87 --- /dev/null +++ b/internal/ui/model.go @@ -0,0 +1,341 @@ +package ui + +import ( + "context" + "fmt" + "strings" + + "github.com/charmbracelet/bubbles/spinner" + "github.com/charmbracelet/bubbles/textarea" + "github.com/charmbracelet/bubbles/viewport" + tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/glamour" +) + +// sessionState is the TUI's coarse state machine: idle waits for input, +// streaming/toolRunning reflect where the current turn is, and approval +// suspends everything else until a modal is answered. +type sessionState int + +const ( + stateIdle sessionState = iota + stateStreaming + stateToolRunning + stateApproval +) + +// toolBlock is one tool call's rendered status within the current turn. +type toolBlock struct { + id, name string + status string // "running", "ok", "error" + output string +} + +type approvalRequest struct { + toolName string + preview string + resp chan<- bool +} + +// Model is the bubbletea model driving the interactive TUI. It never talks +// to a Provider or the agent loop directly — TUIRenderer and TUIApprover +// (this package) push it messages over the tea.Program those own, keeping +// UI rendering separate from agent policy per CLAUDE.md invariant 5. +type Model struct { + submit chan<- string + + state sessionState + priorState sessionState + + viewport viewport.Model + textarea textarea.Model + spinner spinner.Model + renderer *glamour.TermRenderer + + completed string // finalized, glamour-rendered transcript + liveText strings.Builder + liveTools []toolBlock + cancel context.CancelFunc + + approvalQueue []approvalRequest + approval *approvalRequest + + err error + ready bool + width int +} + +// NewModel returns a Model that sends submitted user input on submit. +func NewModel(submit chan<- string) Model { + ta := textarea.New() + ta.Placeholder = "Ask pico code…" + ta.Focus() + ta.ShowLineNumbers = false + ta.SetHeight(3) + + sp := spinner.New() + sp.Spinner = spinner.Dot + + return Model{ + submit: submit, + textarea: ta, + spinner: sp, + viewport: viewport.New(80, 20), + } +} + +// Init implements tea.Model. +func (m Model) Init() tea.Cmd { + return textarea.Blink +} + +// Update implements tea.Model. +func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.WindowSizeMsg: + return m.handleResize(msg) + case tea.KeyMsg: + return m.handleKey(msg) + case spinner.TickMsg: + if m.state == stateIdle { + return m, nil + } + var cmd tea.Cmd + m.spinner, cmd = m.spinner.Update(msg) + return m, cmd + case turnStartedMsg: + m.state = stateStreaming + m.cancel = msg.cancel + m.liveText.Reset() + m.liveTools = nil + m.err = nil + m.refreshViewport() + return m, m.spinner.Tick + case textDeltaMsg: + m.liveText.WriteString(msg.text) + m.refreshViewport() + return m, nil + case toolStartedMsg: + m.state = stateToolRunning + m.liveTools = append(m.liveTools, toolBlock{id: msg.id, name: msg.name, status: "running"}) + m.refreshViewport() + return m, nil + case toolFinishedMsg: + for i := range m.liveTools { + if m.liveTools[i].id == msg.id { + m.liveTools[i].status = "ok" + if msg.isError { + m.liveTools[i].status = "error" + } + m.liveTools[i].output = msg.output + } + } + if !m.anyToolRunning() { + m.state = stateStreaming + } + m.refreshViewport() + return m, nil + case turnDoneMsg: + m.finalizeTurn(msg.text, msg.err) + return m, nil + case commandOutputMsg: + m.completed += msg.text + m.refreshViewport() + return m, nil + case approvalRequestMsg: + m.approvalQueue = append(m.approvalQueue, approvalRequest{toolName: msg.toolName, preview: msg.preview, resp: msg.resp}) + if m.approval == nil { + m.activateNextApproval() + } + return m, nil + } + return m, nil +} + +func (m *Model) anyToolRunning() bool { + for _, tb := range m.liveTools { + if tb.status == "running" { + return true + } + } + return false +} + +func (m *Model) activateNextApproval() { + if len(m.approvalQueue) == 0 { + m.approval = nil + m.state = m.priorState + return + } + if m.state != stateApproval { + m.priorState = m.state + } + req := m.approvalQueue[0] + m.approvalQueue = m.approvalQueue[1:] + m.approval = &req + m.state = stateApproval +} + +func (m *Model) finalizeTurn(text string, err error) { + m.state = stateIdle + m.cancel = nil + if err != nil { + m.err = err + } else if text != "" { + rendered := text + if m.renderer != nil { + if out, rerr := m.renderer.Render(text); rerr == nil { + rendered = out + } + } + m.completed += rendered + } + m.liveText.Reset() + m.liveTools = nil + m.refreshViewport() + m.textarea.Focus() +} + +func (m Model) handleResize(msg tea.WindowSizeMsg) (Model, tea.Cmd) { + m.width = msg.Width + m.ready = true + m.textarea.SetWidth(msg.Width) + m.viewport.Width = msg.Width + m.viewport.Height = msg.Height - m.textarea.Height() - 2 + if r, err := glamour.NewTermRenderer(glamour.WithAutoStyle(), glamour.WithWordWrap(msg.Width)); err == nil { + m.renderer = r + } + m.refreshViewport() + return m, nil +} + +func (m Model) handleKey(msg tea.KeyMsg) (Model, tea.Cmd) { + if msg.Type == tea.KeyCtrlD { + if m.cancel != nil { + m.cancel() + } + return m, tea.Quit + } + + if m.state == stateApproval { + return m.handleApprovalKey(msg) + } + + if msg.Type == tea.KeyCtrlC { + if m.cancel != nil { + m.cancel() + } + return m, nil + } + + if m.state == stateIdle && msg.Type == tea.KeyEnter { + input := strings.TrimSpace(m.textarea.Value()) + if input == "" { + return m, nil + } + select { + case m.submit <- input: + m.textarea.Reset() + default: + } + return m, nil + } + + if m.state != stateIdle { + var cmd tea.Cmd + m.viewport, cmd = m.viewport.Update(msg) + return m, cmd + } + + var cmds []tea.Cmd + var cmd tea.Cmd + m.textarea, cmd = m.textarea.Update(msg) + cmds = append(cmds, cmd) + m.viewport, cmd = m.viewport.Update(msg) + cmds = append(cmds, cmd) + return m, tea.Batch(cmds...) +} + +func (m Model) handleApprovalKey(msg tea.KeyMsg) (Model, tea.Cmd) { + if m.approval == nil { + return m, nil + } + switch msg.Type { + case tea.KeyRunes: + switch string(msg.Runes) { + case "y", "Y": + m.approval.resp <- true + m.activateNextApproval() + case "n", "N": + m.approval.resp <- false + m.activateNextApproval() + } + case tea.KeyEsc: + m.approval.resp <- false + m.activateNextApproval() + } + return m, nil +} + +// refreshViewport rebuilds the viewport's content from the finalized +// transcript plus whatever the in-progress turn has produced so far. +func (m *Model) refreshViewport() { + var b strings.Builder + b.WriteString(m.completed) + if m.liveText.Len() > 0 || len(m.liveTools) > 0 { + b.WriteString(m.liveText.String()) + b.WriteString("\n") + } + for _, tb := range m.liveTools { + icon := "⏳" + switch tb.status { + case "ok": + icon = "✓" + case "error": + icon = "✗" + } + fmt.Fprintf(&b, "%s %s\n", icon, tb.name) + } + if m.err != nil { + fmt.Fprintf(&b, "\nerror: %v\n", m.err) + } + m.viewport.SetContent(b.String()) + m.viewport.GotoBottom() +} + +// View implements tea.Model. +func (m Model) View() string { + if !m.ready { + return "initializing…" + } + if m.state == stateApproval && m.approval != nil { + return m.viewport.View() + "\n" + m.approvalModalView() + } + + status := "> " + if m.state != stateIdle { + status = m.spinner.View() + " " + m.statusText() + } + return m.viewport.View() + "\n" + status + "\n" + m.textarea.View() +} + +func (m Model) statusText() string { + switch m.state { + case stateStreaming: + return "thinking…" + case stateToolRunning: + return "running tools…" + default: + return "" + } +} + +func (m Model) approvalModalView() string { + var b strings.Builder + fmt.Fprintf(&b, "Approve call to %q? [y/N]\n", m.approval.toolName) + if m.approval.preview != "" { + b.WriteString(m.approval.preview) + b.WriteString("\n") + } + return b.String() +} diff --git a/internal/ui/model_test.go b/internal/ui/model_test.go new file mode 100644 index 0000000..480e36d --- /dev/null +++ b/internal/ui/model_test.go @@ -0,0 +1,220 @@ +package ui + +import ( + "errors" + "strings" + "testing" + + tea "github.com/charmbracelet/bubbletea" +) + +func newTestModel() Model { + m := NewModel(make(chan string, 1)) + rm, _ := m.handleResize(tea.WindowSizeMsg{Width: 80, Height: 24}) + return rm +} + +func update(m Model, msg tea.Msg) Model { + next, _ := m.Update(msg) + return next.(Model) +} + +// TestModelStateTransitionsIdleStreamingToolIdle is 7.2's AC: the model +// layer has unit tests for the idle -> streaming -> tool -> idle cycle. +func TestModelStateTransitionsIdleStreamingToolIdle(t *testing.T) { + m := newTestModel() + if m.state != stateIdle { + t.Fatalf("initial state = %v, want stateIdle", m.state) + } + + m = update(m, turnStartedMsg{cancel: func() {}}) + if m.state != stateStreaming { + t.Fatalf("after turnStartedMsg state = %v, want stateStreaming", m.state) + } + + m = update(m, textDeltaMsg{text: "hello"}) + if m.state != stateStreaming { + t.Fatalf("after textDeltaMsg state = %v, want stateStreaming", m.state) + } + if got := m.liveText.String(); got != "hello" { + t.Errorf("liveText = %q, want %q", got, "hello") + } + + m = update(m, toolStartedMsg{id: "t1", name: "read_file"}) + if m.state != stateToolRunning { + t.Fatalf("after toolStartedMsg state = %v, want stateToolRunning", m.state) + } + + m = update(m, toolFinishedMsg{id: "t1", name: "read_file", output: "contents", isError: false}) + if m.state != stateStreaming { + t.Fatalf("after the only running tool finishes, state = %v, want stateStreaming", m.state) + } + + m = update(m, turnDoneMsg{text: "done"}) + if m.state != stateIdle { + t.Fatalf("after turnDoneMsg state = %v, want stateIdle", m.state) + } + if m.cancel != nil { + t.Error("cancel should be cleared once the turn is done") + } +} + +func TestModelToolRunningPersistsUntilLastParallelToolFinishes(t *testing.T) { + m := newTestModel() + m = update(m, turnStartedMsg{cancel: func() {}}) + m = update(m, toolStartedMsg{id: "t1", name: "a"}) + m = update(m, toolStartedMsg{id: "t2", name: "b"}) + if m.state != stateToolRunning { + t.Fatalf("state = %v, want stateToolRunning", m.state) + } + + m = update(m, toolFinishedMsg{id: "t1", name: "a", output: "ok"}) + if m.state != stateToolRunning { + t.Fatalf("state = %v, want stateToolRunning while t2 is still running", m.state) + } + + m = update(m, toolFinishedMsg{id: "t2", name: "b", output: "ok"}) + if m.state != stateStreaming { + t.Fatalf("state = %v, want stateStreaming once every tool has finished", m.state) + } +} + +func TestModelApprovalRequestSuspendsAndRestoresPriorState(t *testing.T) { + m := newTestModel() + m = update(m, turnStartedMsg{cancel: func() {}}) + + resp := make(chan bool, 1) + m = update(m, approvalRequestMsg{toolName: "run_command", resp: resp}) + if m.state != stateApproval { + t.Fatalf("state = %v, want stateApproval", m.state) + } + if m.approval == nil || m.approval.toolName != "run_command" { + t.Fatalf("approval = %+v, want a pending request for run_command", m.approval) + } + + m = update(m, tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("y")}) + select { + case ok := <-resp: + if !ok { + t.Error("resp = false, want true for a 'y' answer") + } + default: + t.Fatal("expected an answer to be sent on resp") + } + if m.state != stateStreaming { + t.Fatalf("state after answering = %v, want it restored to stateStreaming", m.state) + } +} + +func TestModelApprovalQueueServesOneAtATime(t *testing.T) { + m := newTestModel() + m = update(m, turnStartedMsg{cancel: func() {}}) + + resp1 := make(chan bool, 1) + resp2 := make(chan bool, 1) + m = update(m, approvalRequestMsg{toolName: "first", resp: resp1}) + m = update(m, approvalRequestMsg{toolName: "second", resp: resp2}) + + if m.approval.toolName != "first" { + t.Fatalf("active approval = %q, want %q", m.approval.toolName, "first") + } + + m = update(m, tea.KeyMsg{Type: tea.KeyEsc}) + if ok := <-resp1; ok { + t.Error("Esc should deny, got true") + } + if m.approval == nil || m.approval.toolName != "second" { + t.Fatalf("active approval after answering first = %+v, want %q next", m.approval, "second") + } + if m.state != stateApproval { + t.Fatalf("state = %v, want stateApproval while a second request is queued", m.state) + } +} + +func TestModelTurnDoneWithErrorReturnsToIdleAndRecordsErr(t *testing.T) { + m := newTestModel() + m = update(m, turnStartedMsg{cancel: func() {}}) + wantErr := errors.New("boom") + m = update(m, turnDoneMsg{err: wantErr}) + + if m.state != stateIdle { + t.Fatalf("state = %v, want stateIdle", m.state) + } + if !errors.Is(m.err, wantErr) { + t.Errorf("err = %v, want %v", m.err, wantErr) + } +} + +func TestModelCtrlDQuitsAndCancelsAnyInFlightTurn(t *testing.T) { + m := newTestModel() + cancelled := false + m = update(m, turnStartedMsg{cancel: func() { cancelled = true }}) + + _, cmd := m.Update(tea.KeyMsg{Type: tea.KeyCtrlD}) + if !cancelled { + t.Error("Ctrl+D during a turn should cancel it") + } + if cmd == nil { + t.Fatal("Ctrl+D should return a quit command") + } +} + +func TestModelCtrlCCancelsTurnWithoutQuitting(t *testing.T) { + m := newTestModel() + cancelled := false + m = update(m, turnStartedMsg{cancel: func() { cancelled = true }}) + + next, cmd := m.Update(tea.KeyMsg{Type: tea.KeyCtrlC}) + m = next.(Model) + if !cancelled { + t.Error("Ctrl+C during a turn should cancel it") + } + if cmd != nil { + t.Error("Ctrl+C should not quit the program") + } + if m.state != stateStreaming { + t.Errorf("state after Ctrl+C = %v, want it to stay stateStreaming (cancellation arrives via turnDoneMsg)", m.state) + } +} + +func TestModelEnterSubmitsInputOnlyWhenIdle(t *testing.T) { + ch := make(chan string, 1) + m := NewModel(ch) + m, _ = m.handleResize(tea.WindowSizeMsg{Width: 80, Height: 24}) + m.textarea.SetValue("hello agent") + + m = update(m, tea.KeyMsg{Type: tea.KeyEnter}) + + select { + case got := <-ch: + if got != "hello agent" { + t.Errorf("submitted = %q, want %q", got, "hello agent") + } + default: + t.Fatal("expected Enter in stateIdle to submit the textarea's value") + } + if m.textarea.Value() != "" { + t.Errorf("textarea value after submit = %q, want it cleared", m.textarea.Value()) + } +} + +// TestModelCommandOutputAppendsToTranscriptWithoutChangingState proves a +// slash command's output (sent via CommandOutput/commandOutputMsg) shows +// up in the transcript without going through the turn state machine — no +// spinner, no state change, since cmd/pico's driver never sends +// turnStartedMsg for a command. +func TestModelCommandOutputAppendsToTranscriptWithoutChangingState(t *testing.T) { + m := newTestModel() + if m.state != stateIdle { + t.Fatalf("initial state = %v, want stateIdle", m.state) + } + + m = update(m, commandOutputMsg{text: "> /usage\ncumulative: 5 input, 2 output token(s)\n"}) + + if m.state != stateIdle { + t.Errorf("state after commandOutputMsg = %v, want it to stay stateIdle", m.state) + } + if !strings.Contains(m.completed, "cumulative: 5 input, 2 output token(s)") { + t.Errorf("completed = %q, want it to contain the command output", m.completed) + } +} diff --git a/internal/ui/plain.go b/internal/ui/plain.go new file mode 100644 index 0000000..36b7959 --- /dev/null +++ b/internal/ui/plain.go @@ -0,0 +1,101 @@ +package ui + +import ( + "context" + "encoding/json" + "fmt" + "io" + "strings" + + "github.com/reno/pico-code/internal/llm" +) + +// PlainRenderer writes assistant text to Out as it arrives and nothing +// else beyond a couple of one-line tool status markers — no ANSI escapes, +// no cursor movement — so it is safe piped, redirected, or read by another +// program. It is the default whenever stdout is not a terminal or --tui is +// off. +type PlainRenderer struct { + Out io.Writer +} + +// Render implements Renderer. It reconstructs Message.Blocks the same way +// llm.CollectStream does, writing each TextDelta's text to Out as it +// arrives rather than only after the fact. +func (p PlainRenderer) Render(ctx context.Context, events <-chan llm.Event) (*llm.Response, error) { + var blocks []llm.Block + var text strings.Builder + textOpen := false + wroteAny := false + names := map[string]string{} + + flushText := func() { + if textOpen { + blocks = append(blocks, llm.Text{Text: text.String()}) + text.Reset() + textOpen = false + } + } + + for { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case e, ok := <-events: + if !ok { + return nil, fmt.Errorf("ui: event stream closed before a MessageDone event") + } + switch v := e.(type) { + case llm.TextDelta: + textOpen = true + text.WriteString(v.Text) + if _, err := fmt.Fprint(p.Out, v.Text); err != nil { + return nil, fmt.Errorf("ui: write output: %w", err) + } + wroteAny = true + case llm.ToolUseStart: + flushText() + names[v.ID] = v.Name + case llm.ToolUseArgsDelta: + // Informational only; ToolUseDone carries the finished value. + case llm.ToolUseDone: + blocks = append(blocks, llm.ToolUse{ID: v.ID, Name: names[v.ID], Input: v.Input}) + case llm.MessageDone: + flushText() + if wroteAny { + if _, err := fmt.Fprintln(p.Out); err != nil { + return nil, fmt.Errorf("ui: write output: %w", err) + } + } + return &llm.Response{ + Message: llm.Message{Role: llm.RoleAssistant, Blocks: blocks}, + StopReason: v.StopReason, + Usage: v.Usage, + }, nil + case llm.Error: + return nil, v.Err + default: + return nil, fmt.Errorf("ui: unknown event type %T", e) + } + } + } +} + +// ToolStarted implements ToolStatusReporter. +func (p PlainRenderer) ToolStarted(_, name string, _ json.RawMessage) { + _, _ = fmt.Fprintf(p.Out, "→ %s\n", name) +} + +// ToolFinished implements ToolStatusReporter. +func (p PlainRenderer) ToolFinished(_, name, _ string, isError bool) { + if isError { + _, _ = fmt.Fprintf(p.Out, "✗ %s\n", name) + return + } + _, _ = fmt.Fprintf(p.Out, "✓ %s\n", name) +} + +var ( + _ Renderer = PlainRenderer{} + _ ToolStatusReporter = PlainRenderer{} +) diff --git a/internal/ui/plain_test.go b/internal/ui/plain_test.go new file mode 100644 index 0000000..2b3f32f --- /dev/null +++ b/internal/ui/plain_test.go @@ -0,0 +1,134 @@ +package ui + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "strings" + "testing" + "time" + + "github.com/google/go-cmp/cmp" + + "github.com/reno/pico-code/internal/llm" +) + +func send(ch chan llm.Event, events ...llm.Event) { + for _, e := range events { + ch <- e + } + close(ch) +} + +// TestPlainRendererWritesTextAsItStreamsWithNoANSI is 7.1's AC: piped +// output is clean and ANSI-free. +func TestPlainRendererWritesTextAsItStreamsWithNoANSI(t *testing.T) { + ch := make(chan llm.Event) + go send(ch, + llm.TextDelta{Text: "hel"}, + llm.TextDelta{Text: "lo"}, + llm.MessageDone{StopReason: "end_turn", Usage: llm.Usage{InputTokens: 3, OutputTokens: 2}}, + ) + + var buf bytes.Buffer + resp, err := (PlainRenderer{Out: &buf}).Render(context.Background(), ch) + if err != nil { + t.Fatalf("Render() error = %v", err) + } + + if got := buf.String(); got != "hello\n" { + t.Errorf("Out = %q, want %q", got, "hello\n") + } + if strings.ContainsRune(buf.String(), '\x1b') { + t.Errorf("Out contains an ANSI escape byte: %q", buf.String()) + } + + want := &llm.Response{ + Message: llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.Text{Text: "hello"}}}, + StopReason: "end_turn", + Usage: llm.Usage{InputTokens: 3, OutputTokens: 2}, + } + if diff := cmp.Diff(want, resp); diff != "" { + t.Errorf("Response mismatch (-want +got):\n%s", diff) + } +} + +func TestPlainRendererReconstructsToolUseBlock(t *testing.T) { + ch := make(chan llm.Event) + go send(ch, + llm.TextDelta{Text: "checking"}, + llm.ToolUseStart{ID: "t1", Name: "get_weather"}, + llm.ToolUseArgsDelta{ID: "t1", Partial: `{"location"`}, + llm.ToolUseArgsDelta{ID: "t1", Partial: `:"Paris"}`}, + llm.ToolUseDone{ID: "t1", Input: json.RawMessage(`{"location":"Paris"}`)}, + llm.MessageDone{StopReason: "tool_use"}, + ) + + var buf bytes.Buffer + resp, err := (PlainRenderer{Out: &buf}).Render(context.Background(), ch) + if err != nil { + t.Fatalf("Render() error = %v", err) + } + + want := []llm.Block{ + llm.Text{Text: "checking"}, + llm.ToolUse{ID: "t1", Name: "get_weather", Input: json.RawMessage(`{"location":"Paris"}`)}, + } + if diff := cmp.Diff(want, resp.Message.Blocks); diff != "" { + t.Errorf("Blocks mismatch (-want +got):\n%s", diff) + } + if got := buf.String(); got != "checking\n" { + t.Errorf("Out = %q, want %q", got, "checking\n") + } +} + +func TestPlainRendererPropagatesErrorEvent(t *testing.T) { + wantErr := errors.New("boom") + ch := make(chan llm.Event) + go send(ch, llm.Error{Err: wantErr}) + + _, err := (PlainRenderer{Out: &bytes.Buffer{}}).Render(context.Background(), ch) + if !errors.Is(err, wantErr) { + t.Fatalf("Render() error = %v, want it to wrap %v", err, wantErr) + } +} + +func TestPlainRendererErrorsIfChannelClosesWithoutMessageDone(t *testing.T) { + ch := make(chan llm.Event) + go send(ch, llm.TextDelta{Text: "hi"}) + + _, err := (PlainRenderer{Out: &bytes.Buffer{}}).Render(context.Background(), ch) + if err == nil { + t.Fatal("Render() error = nil, want an error for a stream closed without MessageDone") + } +} + +func TestPlainRendererRespectsCancellation(t *testing.T) { + ch := make(chan llm.Event) // never sends + ctx, cancel := context.WithCancel(context.Background()) + time.AfterFunc(10*time.Millisecond, cancel) + + start := time.Now() + _, err := (PlainRenderer{Out: &bytes.Buffer{}}).Render(ctx, ch) + elapsed := time.Since(start) + + if !errors.Is(err, context.Canceled) { + t.Fatalf("Render() error = %v, want it to wrap context.Canceled", err) + } + if elapsed > time.Second { + t.Fatalf("Render() took %s after cancellation, want a prompt return", elapsed) + } +} + +func TestPlainRendererToolStatusHasNoANSI(t *testing.T) { + var buf bytes.Buffer + r := PlainRenderer{Out: &buf} + r.ToolStarted("t1", "read_file", json.RawMessage(`{"path":"x"}`)) + r.ToolFinished("t1", "read_file", "contents", false) + r.ToolFinished("t1", "read_file", "boom", true) + + if strings.ContainsRune(buf.String(), '\x1b') { + t.Errorf("tool status output contains an ANSI escape byte: %q", buf.String()) + } +} diff --git a/internal/ui/renderer.go b/internal/ui/renderer.go new file mode 100644 index 0000000..3042ff7 --- /dev/null +++ b/internal/ui/renderer.go @@ -0,0 +1,30 @@ +// Package ui renders an agent turn for a person: it turns a Provider's +// event stream (and the agent loop's tool-execution outcomes, which never +// appear on that stream since running a tool is loop policy, not provider +// transport) into terminal output. Two implementations exist: Plain, the +// default for piped/non-interactive use, and a bubbletea TUI. +package ui + +import ( + "context" + "encoding/json" + + "github.com/reno/pico-code/internal/llm" +) + +// Renderer turns one turn's provider event stream into user-facing output, +// returning the assembled Response for the agent loop to append to +// history — the same contract llm.CollectStream has, plus the rendering +// side effect. +type Renderer interface { + Render(ctx context.Context, events <-chan llm.Event) (*llm.Response, error) +} + +// ToolStatusReporter is an optional Renderer capability for showing a tool +// call's progress. The agent loop calls it, when implemented, around +// actually running a tool — a step that happens after Render returns, so it +// cannot be expressed as an llm.Event. +type ToolStatusReporter interface { + ToolStarted(id, name string, input json.RawMessage) + ToolFinished(id, name, output string, isError bool) +} diff --git a/internal/ui/tui.go b/internal/ui/tui.go new file mode 100644 index 0000000..2006753 --- /dev/null +++ b/internal/ui/tui.go @@ -0,0 +1,156 @@ +package ui + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + tea "github.com/charmbracelet/bubbletea" + + "github.com/reno/pico-code/internal/llm" +) + +// The bubbletea messages driving Model.Update. They are unexported: only +// TUIRenderer and TUIApprover in this package construct them, sending each +// through the *tea.Program the caller owns — cmd/pico never touches these +// directly, only wires Model, TUIRenderer, and TUIApprover together. +type ( + turnStartedMsg struct{ cancel context.CancelFunc } + textDeltaMsg struct{ text string } + toolStartedMsg struct { + id, name string + input json.RawMessage + } + toolFinishedMsg struct { + id, name, output string + isError bool + } + turnDoneMsg struct { + text string + err error + } + commandOutputMsg struct{ text string } + approvalRequestMsg struct { + toolName string + input json.RawMessage + preview string + resp chan<- bool + } +) + +// TUIRenderer implements Renderer and ToolStatusReporter by forwarding +// every event to a *tea.Program as a message, so Model.Update can render +// it — the same block-reconstruction PlainRenderer does, plus the side +// effect of a Send instead of a Fprint. +type TUIRenderer struct { + Program *tea.Program +} + +// Render implements Renderer. +func (r *TUIRenderer) Render(ctx context.Context, events <-chan llm.Event) (*llm.Response, error) { + var blocks []llm.Block + var text strings.Builder + textOpen := false + names := map[string]string{} + + flushText := func() { + if textOpen { + blocks = append(blocks, llm.Text{Text: text.String()}) + text.Reset() + textOpen = false + } + } + + for { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case e, ok := <-events: + if !ok { + return nil, fmt.Errorf("ui: event stream closed before a MessageDone event") + } + switch v := e.(type) { + case llm.TextDelta: + textOpen = true + text.WriteString(v.Text) + r.Program.Send(textDeltaMsg{text: v.Text}) + case llm.ToolUseStart: + flushText() + names[v.ID] = v.Name + case llm.ToolUseArgsDelta: + case llm.ToolUseDone: + blocks = append(blocks, llm.ToolUse{ID: v.ID, Name: names[v.ID], Input: v.Input}) + case llm.MessageDone: + flushText() + return &llm.Response{ + Message: llm.Message{Role: llm.RoleAssistant, Blocks: blocks}, + StopReason: v.StopReason, + Usage: v.Usage, + }, nil + case llm.Error: + return nil, v.Err + default: + return nil, fmt.Errorf("ui: unknown event type %T", e) + } + } + } +} + +// ToolStarted implements ToolStatusReporter. +func (r *TUIRenderer) ToolStarted(id, name string, input json.RawMessage) { + r.Program.Send(toolStartedMsg{id: id, name: name, input: input}) +} + +// ToolFinished implements ToolStatusReporter. +func (r *TUIRenderer) ToolFinished(id, name, output string, isError bool) { + r.Program.Send(toolFinishedMsg{id: id, name: name, output: output, isError: isError}) +} + +var ( + _ Renderer = (*TUIRenderer)(nil) + _ ToolStatusReporter = (*TUIRenderer)(nil) +) + +// TUIApprover implements the agent.Approver contract (matched structurally +// to avoid this package importing internal/agent) by showing a modal in +// the TUI and blocking until it's answered or ctx is cancelled. Approve +// runs on the agent loop's tool goroutine, never on the tea.Program's own +// goroutine, so blocking here does not freeze the UI. +type TUIApprover struct { + Program *tea.Program +} + +// Approve sends an approval request to the running Model and waits for its +// answer. +func (a *TUIApprover) Approve(ctx context.Context, toolName string, input json.RawMessage, preview string) (bool, error) { + resp := make(chan bool, 1) + a.Program.Send(approvalRequestMsg{toolName: toolName, input: input, preview: preview, resp: resp}) + select { + case ok := <-resp: + return ok, nil + case <-ctx.Done(): + return false, ctx.Err() + } +} + +// TurnStarted announces a turn's start (and its cancel func, for Ctrl+C) to +// the running Model, letting cmd/pico's driver loop do so without +// constructing the unexported message type itself. TurnDone is its +// counterpart for a turn's end. +func TurnStarted(program *tea.Program, cancel context.CancelFunc) { + program.Send(turnStartedMsg{cancel: cancel}) +} + +// TurnDone announces a turn's result, mirroring TurnStarted. +func TurnDone(program *tea.Program, text string, err error) { + program.Send(turnDoneMsg{text: text, err: err}) +} + +// CommandOutput appends text (a slash command's result, e.g. /usage or +// /save) to the transcript. Unlike a turn, a command never touches the +// spinner/streaming state — cmd/pico's driver loop handles it inline, +// without ever sending TurnStarted for it. +func CommandOutput(program *tea.Program, text string) { + program.Send(commandOutputMsg{text: text}) +}