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/Makefile b/Makefile index 66712a0..2000492 100644 --- a/Makefile +++ b/Makefile @@ -12,4 +12,4 @@ lint: golangci-lint run run: - go run ./cmd/pico chat --provider=ollama + go run ./cmd/pico chat --provider=ollama --model=qwen2.5-coder:7b --tui diff --git a/README.md b/README.md new file mode 100644 index 0000000..4bbca14 --- /dev/null +++ b/README.md @@ -0,0 +1,208 @@ +# 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). + +## 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, and some that advertise +support still narrate a tool call in prose instead of using it. So +`--provider=ollama` defaults `--tools` to `prompted` unless you pass +`--tools` explicitly; pass `--tools=native` yourself if your model handles +it reliably. (`--provider=anthropic` keeps defaulting to `native`.) +`--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; combine `--tui` with `--provider=ollama` and +you'll need to add `--tools=native` explicitly. + +## 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` (`prompted` for `--provider=ollama`) | `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..aacc53f 100644 --- a/cmd/pico/chat.go +++ b/cmd/pico/chat.go @@ -1,30 +1,70 @@ package main import ( + "bufio" + "bytes" + "context" + "errors" "fmt" + "io" + "log/slog" + "os" + "os/user" + "strconv" + "strings" + "time" + tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/lipgloss" "github.com/spf13/cobra" + "golang.org/x/term" + "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" +) + +// version is stamped into the home screen's frame title. Bumped by hand; +// there is no release pipeline to derive it from yet. +const version = "v0.1.0" + +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 // PICO_CODE_PROVIDER environment fallback, which wins over the built-in // default; ANTHROPIC_API_KEY and OLLAMA_HOST are read directly by // config.Load since they are credentials, not flags. +// +// --tools has its own provider-dependent default: local models frequently +// advertise native tool support without reliably using it (CLAUDE.md's +// "narrate a tool call in prose instead of calling it"), so an explicit +// --provider=ollama with no explicit --tools defaults to prompted rather +// than native. An explicit --tools always wins. 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{ @@ -38,18 +78,27 @@ func newChatCmd(getenv func(string) string) *cobra.Command { } } + toolsMode := flags.tools + if !cmd.Flags().Changed("tools") && provider == string(config.ProviderOllama) { + toolsMode = string(config.ToolsPrompted) + } + 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: toolsMode, + 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) @@ -66,17 +115,424 @@ func newChatCmd(getenv func(string) string) *cobra.Command { f.IntVar(&flags.tokenBudget, "token-budget", 100_000, "maximum cumulative tokens before stopping") f.StringVar(&flags.workspace, "workspace", ".", "root directory filesystem tools are confined to") f.BoolVar(&flags.yes, "yes", false, "skip interactive approval prompts") - f.StringVar(&flags.tools, "tools", "native", "tool-calling mode (native|prompted)") + f.StringVar(&flags.tools, "tools", "native", "tool-calling mode (native|prompted); defaults to prompted when --provider=ollama and --tools isn't set explicitly") f.BoolVar(&flags.stream, "stream", true, "stream provider responses") 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 { + // The TUI shares stderr with the alt-screen terminal, so a routine + // info-level line (e.g. "ollama: context window configured") printed + // before program.Run() takes over flashes onto the screen as stray + // text. Default to a quieter level for --tui unless the caller asked + // for a specific one explicitly. + logLevel := effectiveLogLevel(cfg, cmd.Flags().Changed("log-level")) + if err := setupLogging(cmd.ErrOrStderr(), 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 +) + +// effectiveLogLevel is the level setupLogging actually uses: cfg.LogLevel, +// except a TUI session with no explicit --log-level defaults to warn, since +// stderr shares the alt-screen terminal and a routine info line would flash +// onto the screen as stray text before program.Run() takes it over. +func effectiveLogLevel(cfg *config.Config, logLevelFlagChanged bool) config.LogLevel { + if cfg.TUI && !logLevelFlagChanged { + return config.LogLevelWarn + } + return cfg.LogLevel +} + +func compactionPolicy(cfg *config.Config) agent.CompactionPolicy { + return agent.CompactionPolicy{ContextWindow: contextWindow(cfg), TriggerFraction: compactionTriggerFraction, KeepTurns: compactionKeepTurns} +} + +// contextWindow is the token budget compaction measures against: NumCtx +// for Ollama, whose window is already explicit and required, and the +// configured ContextWindow for every other provider. +func contextWindow(cfg *config.Config) int { + if cfg.Provider == config.ProviderOllama { + return cfg.NumCtx + } + return cfg.ContextWindow +} + +// 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) + } + + _, _ = fmt.Fprint(out, ui.Banner(bannerInfo(cfg, sess), terminalWidth(out))) + return runREPL(ctx, in, out, ag, h, sess, runTurn) +} + +// bannerInfo projects the resolved config and on-disk sessions onto the +// fields the home screen shows, so ui never has to import config, read the +// filesystem, or look at the environment. +func bannerInfo(cfg *config.Config, sess *session) ui.BannerInfo { + info := ui.BannerInfo{ + Version: version, + Greeting: ui.Mockery(), + Provider: string(cfg.Provider), + Model: cfg.Model, + Usage: fmt.Sprintf("0 / %s tokens", humanTokens(contextWindow(cfg))), + } + // The home screen shows the directory you launched from, which is what + // a shell prompt would show; cfg.Workspace is the tools' sandbox root + // and can differ from it. + wd, err := os.Getwd() + if err != nil { + wd = cfg.Workspace + } + info.Directory = tildeHome(wd) + if u, err := user.Current(); err == nil { + info.User = u.Username + } + if sess != nil { + info.Sessions = sess.recent(bannerSessionCount) + } + return info +} + +// tildeHome abbreviates a leading home directory to "~", the way a shell +// prompt does, so the path fits the home screen's column. +func tildeHome(path string) string { + home, err := os.UserHomeDir() + if err != nil || home == "" || !strings.HasPrefix(path, home) { + return path + } + return "~" + strings.TrimPrefix(path, home) +} + +// bannerSessionCount is how many recent sessions fit the home screen's +// right column without pushing the box taller than the skull beside it. +const bannerSessionCount = 4 + +// terminalWidth reports out's width, falling back to 0 (which Banner +// clamps to its minimum) when out is not a terminal — a pipe, or a test's +// buffer. +func terminalWidth(out io.Writer) int { + f, ok := out.(*os.File) + if !ok { + return 0 + } + w, _, err := term.GetSize(int(f.Fd())) + if err != nil { + return 0 + } + return w +} + +// humanTokens renders a token count the way a context window is usually +// quoted ("128k") rather than in full. +func humanTokens(n int) string { + if n >= 1000 { + return fmt.Sprintf("%dk", n/1000) + } + return strconv.Itoa(n) +} + +// 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 { + _, err := handleCommand(ctx, out, ag, h, sess, name, arg) + switch { + case errors.Is(err, errExit): + return nil + case errors.Is(err, errClearScrollback): + clearScreen(out) + case 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() +} + +// clearScreen writes an ANSI sequence that clears the visible screen and +// the terminal's scrollback buffer and homes the cursor, for the /clear +// command in the interactive plain REPL. runREPL only reaches this from an +// interactive terminal (never the piped, single-shot path in runPlainChat), +// so 7.1's ANSI-free AC for piped output is unaffected. +func clearScreen(out io.Writer) { + _, _ = fmt.Fprint(out, "\x1b[H\x1b[2J\x1b[3J") +} + +// 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) + // Resolved here, before tea.Program takes the terminal into raw mode: + // querying the background color once bubbletea owns stdin races with + // its input reader and leaks stray escape bytes onto the screen. + glamourStyle := "light" + if lipgloss.HasDarkBackground() { + glamourStyle = "dark" + } + program := tea.NewProgram(ui.NewModel(submit, bannerInfo(cfg, sess), glamourStyle), 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 { + output, cmdErr := runTUICommand(ctx, ag, h, sess, input, name, arg) + if errors.Is(cmdErr, errClearScrollback) { + ui.ClearScrollback(program) + } else { + ui.CommandOutput(program, output) + } + if errors.Is(cmdErr, errExit) { + program.Quit() + } + 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). +// The returned error is never a write failure (bytes.Buffer.Write can't +// fail) — it's only ever nil or one of handleCommand's sentinels, which the +// caller uses to clear the transcript or quit the program instead of +// appending the output. Split out from runTUIChat's driver goroutine so +// it's testable without a running *tea.Program. +func runTUICommand(ctx context.Context, ag *agent.Agent, h *history.History, sess *session, input, name, arg string) (string, error) { + var buf bytes.Buffer + _, _ = fmt.Fprintf(&buf, "> %s\n", input) + _, err := handleCommand(ctx, &buf, ag, h, sess, name, arg) + return buf.String(), err +} diff --git a/cmd/pico/chat_test.go b/cmd/pico/chat_test.go new file mode 100644 index 0000000..b847e40 --- /dev/null +++ b/cmd/pico/chat_test.go @@ -0,0 +1,413 @@ +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(context.Background(), 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(context.Background(), 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(context.Background(), 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(context.Background(), 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") + } +} + +// TestEffectiveLogLevelQuietsTUIWithoutExplicitFlag is a regression test: +// stderr shares the TUI's alt-screen terminal, so an info-level line +// printed before program.Run() takes over used to flash onto the screen as +// stray text (e.g. "ollama: context window configured"). --tui should +// default to a quieter level unless the caller passed --log-level. +func TestEffectiveLogLevelQuietsTUIWithoutExplicitFlag(t *testing.T) { + got := effectiveLogLevel(&config.Config{TUI: true, LogLevel: config.LogLevelInfo}, false) + if got != config.LogLevelWarn { + t.Errorf("effectiveLogLevel(TUI, unset flag) = %q, want %q", got, config.LogLevelWarn) + } +} + +func TestEffectiveLogLevelRespectsExplicitFlagEvenInTUI(t *testing.T) { + got := effectiveLogLevel(&config.Config{TUI: true, LogLevel: config.LogLevelDebug}, true) + if got != config.LogLevelDebug { + t.Errorf("effectiveLogLevel(TUI, explicit flag) = %q, want the explicitly requested %q", got, config.LogLevelDebug) + } +} + +func TestEffectiveLogLevelUnchangedForPlainMode(t *testing.T) { + got := effectiveLogLevel(&config.Config{TUI: false, LogLevel: config.LogLevelInfo}, false) + if got != config.LogLevelInfo { + t.Errorf("effectiveLogLevel(plain) = %q, want cfg.LogLevel (%q) unmodified", got, config.LogLevelInfo) + } +} + +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..ee01bec --- /dev/null +++ b/cmd/pico/commands.go @@ -0,0 +1,320 @@ +package main + +import ( + "context" + "errors" + "fmt" + "io" + "strconv" + "strings" + "text/tabwriter" + + "github.com/reno/pico-code/internal/agent" + "github.com/reno/pico-code/internal/history" + "github.com/reno/pico-code/internal/llm" + "github.com/reno/pico-code/internal/tools" +) + +// errExit and errClearScrollback are handleCommand's two sentinel results: +// signals to the caller's own loop rather than failures. Neither commands.go +// action can be expressed by writing to out alone — exiting has to unwind +// the plain REPL's scanner loop or quit the bubbletea program, and clearing +// has to reset a real terminal or the TUI's transcript state, both owned by +// the caller, not by this UI-agnostic file. +var ( + errExit = errors.New("exit requested") + errClearScrollback = errors.New("clear scrollback requested") +) + +// commandSpec describes one slash command. Adding a command to the chat +// means adding one entry to commandTable — handleCommand and /help both +// drive off it, so nothing else in this file needs to change. +type commandSpec struct { + name string + args string + summary string + run func(ctx context.Context, out io.Writer, ag *agent.Agent, h *history.History, sess *session, arg string) error +} + +// commandTable is populated in init(), not a var initializer: runHelpCommand +// reads commandTable, and a direct function-value reference inside the +// initializer would make that a compile-time initialization cycle. +var commandTable []commandSpec + +func init() { + commandTable = []commandSpec{ + {name: "help", summary: "list available commands", run: runHelpCommand}, + {name: "usage", summary: "show cumulative and last-turn token usage", run: runUsageCommand}, + {name: "new", summary: "start a new, unsaved session", run: runNewCommand}, + {name: "save", args: "[name]", summary: "save the current session to disk", run: runSaveCommand}, + {name: "load", args: "", summary: "load a saved session from disk", run: runLoadCommand}, + {name: "clear", summary: "clear the terminal scrollback (the conversation history is kept)", run: runClearCommand}, + {name: "compact", summary: "summarize older turns now and report the token estimate", run: runCompactCommand}, + {name: "exit", summary: "save the session and exit, like Ctrl+D", run: runExitCommand}, + {name: "cd", args: "", summary: "re-root the workspace sandbox", run: runCdCommand}, + {name: "model", args: " [context-window]", summary: "switch models, validated against the active provider", run: runModelCommand}, + } +} + +// 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 with +// its closest match by name rather than sent to the model as chat input). +// err is either a genuine failure or one of the sentinels above; callers +// that don't special-case those sentinels should still treat a non-nil err +// as fatal to the current turn/command, same as any other error. +func handleCommand(ctx context.Context, out io.Writer, ag *agent.Agent, h *history.History, sess *session, name, arg string) (handled bool, err error) { + for _, c := range commandTable { + if c.name == name { + return true, c.run(ctx, out, ag, h, sess, arg) + } + } + _, err = fmt.Fprintf(out, "unknown command /%s, did you mean /%s?\n", name, closestCommand(name)) + return true, err +} + +// runHelpCommand renders commandTable as an aligned list of "/name args" +// followed by its summary. +func runHelpCommand(_ context.Context, out io.Writer, _ *agent.Agent, _ *history.History, _ *session, _ string) error { + tw := tabwriter.NewWriter(out, 0, 4, 2, ' ', 0) + for _, c := range commandTable { + usage := "/" + c.name + if c.args != "" { + usage += " " + c.args + } + if _, err := fmt.Fprintf(tw, "%s\t%s\n", usage, c.summary); err != nil { + return err + } + } + return tw.Flush() +} + +func runUsageCommand(_ context.Context, out io.Writer, ag *agent.Agent, _ *history.History, _ *session, _ string) error { + printUsage(out, ag) + return nil +} + +func runNewCommand(_ context.Context, out io.Writer, _ *agent.Agent, h *history.History, sess *session, _ string) error { + h.Reset() + sess.name = "" + _, err := fmt.Fprintln(out, "started a new, unsaved session") + return err +} + +func runSaveCommand(_ context.Context, out io.Writer, _ *agent.Agent, 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(_ context.Context, out io.Writer, _ *agent.Agent, 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 +} + +// runClearCommand never touches out itself: clearing a real terminal's +// scrollback and clearing the TUI's transcript are different operations +// owned by the two callers (runREPL, runTUICommand), which react to +// errClearScrollback instead. +func runClearCommand(_ context.Context, _ io.Writer, _ *agent.Agent, _ *history.History, _ *session, _ string) error { + return errClearScrollback +} + +// runCompactCommand forces the same summarization pass maybeCompact runs +// automatically (agent.CompactionPolicy, CLAUDE.md phase 8.2) and reports +// the estimated token count before and after. +func runCompactCommand(ctx context.Context, out io.Writer, ag *agent.Agent, _ *history.History, _ *session, _ string) error { + before, after, err := ag.ForceCompact(ctx) + if err != nil { + _, werr := fmt.Fprintf(out, "compact failed: %v\n", err) + return werr + } + if before == after { + _, err := fmt.Fprintf(out, "nothing to compact (estimated %d token(s))\n", before) + return err + } + _, err = fmt.Fprintf(out, "compacted: %d -> %d estimated token(s)\n", before, after) + return err +} + +// runExitCommand saves the session first, then signals errExit so the +// caller unwinds through the same path as Ctrl+D. If the save fails, it +// reports the failure and does not signal exit, so a save failure never +// silently drops the session. +func runExitCommand(_ context.Context, out io.Writer, _ *agent.Agent, h *history.History, sess *session, _ string) error { + if err := sess.saveIfActive(h); err != nil { + _, werr := fmt.Fprintf(out, "exit: session save failed: %v\n", err) + return werr + } + return errExit +} + +// sandboxProvider is satisfied by any registered tool that exposes the +// shared workspace Sandbox — every filesystem tool (read_file, list_dir, +// write_file) is constructed with the same *tools.Sandbox pointer, so +// reaching it through any one of them re-roots every filesystem tool at +// once. read_file is always registered (buildRegistry never gates it +// behind a flag, unlike write_file/run_command), so it's the one this +// looks up. +type sandboxProvider interface { + Sandbox() *tools.Sandbox +} + +// runCdCommand re-roots the workspace sandbox to arg, refusing (and leaving +// the sandbox untouched) if it doesn't exist or can't be resolved — +// Sandbox.Reroot already enforces that the same way NewSandbox does at +// startup. +func runCdCommand(_ context.Context, out io.Writer, ag *agent.Agent, _ *history.History, _ *session, arg string) error { + if arg == "" { + _, err := fmt.Fprintln(out, "usage: /cd ") + return err + } + t, err := ag.Tools().Get("read_file") + if err != nil { + _, werr := fmt.Fprintf(out, "cd failed: %v\n", err) + return werr + } + sb, ok := t.(sandboxProvider) + if !ok { + _, werr := fmt.Fprintln(out, "cd failed: workspace is not sandboxed") + return werr + } + if err := sb.Sandbox().Reroot(arg); err != nil { + _, werr := fmt.Fprintf(out, "cd failed: %v\n", err) + return werr + } + _, err = fmt.Fprintf(out, "workspace root changed to %s\n", arg) + return err +} + +// runModelCommand switches the active provider's model, validating the +// name first (llm.ModelSwitcher.ValidateModel) so an unknown model errors +// without mutating the provider or the compaction policy. arg is "" +// or " "; the window is otherwise left unchanged, +// since none of this codebase's providers can safely auto-detect one for +// every backend (Anthropic can via its Models API; Ollama can't, per +// ValidateModel's doc comment) and CLAUDE.md already treats a model's +// context window as an explicit, user-set value (--context-window, +// --num-ctx) rather than one inferred from the model name. +func runModelCommand(ctx context.Context, out io.Writer, ag *agent.Agent, h *history.History, _ *session, arg string) error { + fields := strings.Fields(arg) + if len(fields) == 0 || len(fields) > 2 { + _, err := fmt.Fprintln(out, "usage: /model [context-window]") + return err + } + name := fields[0] + + switcher, ok := ag.Provider().(llm.ModelSwitcher) + if !ok { + _, err := fmt.Fprintf(out, "model failed: provider %q does not support switching models\n", ag.Provider().Name()) + return err + } + if err := switcher.ValidateModel(ctx, name); err != nil { + _, werr := fmt.Fprintf(out, "model failed: %v\n", err) + return werr + } + + policy := ag.CompactionPolicy() + if len(fields) == 2 { + window, err := strconv.Atoi(fields[1]) + if err != nil || window <= 0 { + _, werr := fmt.Fprintf(out, "model failed: invalid context window %q\n", fields[1]) + return werr + } + policy.ContextWindow = window + } + + switcher.SetModel(name) + ag.SetCompactionPolicy(policy) + + if policy.ContextWindow > 0 { + if used := history.EstimateTokens(h.Snapshot()); used > policy.ContextWindow { + _, err := fmt.Fprintf(out, "switched model to %s (context window %d); warning: current usage (~%d tokens) already exceeds it, compaction will run on the next turn\n", name, policy.ContextWindow, used) + return err + } + } + _, err := fmt.Fprintf(out, "switched model to %s\n", name) + 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) + } +} + +// closestCommand returns the commandTable entry whose name is nearest to +// name by Levenshtein edit distance, for suggesting a fix on a likely typo. +func closestCommand(name string) string { + best := commandTable[0].name + bestDist := levenshtein(name, best) + for _, c := range commandTable[1:] { + if d := levenshtein(name, c.name); d < bestDist { + bestDist = d + best = c.name + } + } + return best +} + +// levenshtein computes the edit distance between a and b. +func levenshtein(a, b string) int { + ar, br := []rune(a), []rune(b) + prev := make([]int, len(br)+1) + curr := make([]int, len(br)+1) + for j := range prev { + prev[j] = j + } + for i := 1; i <= len(ar); i++ { + curr[0] = i + for j := 1; j <= len(br); j++ { + cost := 1 + if ar[i-1] == br[j-1] { + cost = 0 + } + curr[j] = min(prev[j]+1, curr[j-1]+1, prev[j-1]+cost) + } + prev, curr = curr, prev + } + return prev[len(br)] +} diff --git a/cmd/pico/commands_test.go b/cmd/pico/commands_test.go new file mode 100644 index 0000000..46a703e --- /dev/null +++ b/cmd/pico/commands_test.go @@ -0,0 +1,416 @@ +package main + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/google/go-cmp/cmp" + + "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" +) + +const helpGoldenPath = "testdata/golden/help.txt" + +// TestHelpCommandMatchesGolden is 11.1's AC: /help renders commandTable +// aligned. Run with RECORD=1 to (re)write the golden file. +func TestHelpCommandMatchesGolden(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 + if _, err := handleCommand(context.Background(), &out, ag, h, noSession(t), "help", ""); err != nil { + t.Fatalf("handleCommand(help) error = %v", err) + } + got := out.String() + + if os.Getenv("RECORD") == "1" { + if err := os.MkdirAll(filepath.Dir(helpGoldenPath), 0o755); err != nil { + t.Fatalf("MkdirAll() error = %v", err) + } + if err := os.WriteFile(helpGoldenPath, []byte(got), 0o644); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + } + + want, err := os.ReadFile(helpGoldenPath) + if err != nil { + t.Fatalf("ReadFile(%s) error = %v (run with RECORD=1 to create it)", helpGoldenPath, err) + } + if diff := cmp.Diff(string(want), got); diff != "" { + t.Errorf("/help output mismatch (-want +got):\n%s", diff) + } +} + +// TestCommandTableEntriesHaveSummaries guards against a bare entry slipping +// into commandTable with no help text, since /help's usefulness depends on +// every command documenting itself. +func TestCommandTableEntriesHaveSummaries(t *testing.T) { + for _, c := range commandTable { + if strings.TrimSpace(c.summary) == "" { + t.Errorf("command %q has no summary", c.name) + } + if c.run == nil { + t.Errorf("command %q has no run function", c.name) + } + } +} + +// TestUnknownCommandSuggestsClosestMatch is 11.1's AC: an unknown command +// reports the closest match instead of a bare error. +func TestUnknownCommandSuggestsClosestMatch(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 + if _, err := handleCommand(context.Background(), &out, ag, h, noSession(t), "hlep", ""); err != nil { + t.Fatalf("handleCommand(hlep) error = %v", err) + } + got := out.String() + if !strings.Contains(got, "did you mean /help") { + t.Errorf("output = %q, want it to suggest /help", got) + } +} + +// TestClearCommandWipesScrollbackNotHistory is 11.2's AC for /clear: +// history.Validate() still passes and the message count is unchanged, in +// both the plain REPL (which reacts to errClearScrollback by clearing the +// real terminal) and the TUI path (which reacts by resetting the +// transcript instead). +func TestClearCommandWipesScrollbackNotHistory(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 := "hello\n/clear\n" + runTurn := newTurnRunner(&config.Config{Stream: true}, ag, &out) + if err := runREPL(context.Background(), strings.NewReader(in), &out, ag, h, noSession(t), runTurn); err != nil { + t.Fatalf("runREPL() error = %v", err) + } + + if err := h.Validate(); err != nil { + t.Errorf("Validate() after /clear error = %v", err) + } + if got := len(h.Snapshot()); got != 2 { + t.Errorf("history has %d message(s) after /clear, want 2 (the one turn, untouched)", got) + } + if got := out.String(); !strings.Contains(got, "\x1b[") { + t.Errorf("output = %q, want it to contain an ANSI clear sequence", got) + } + + output, err := runTUICommand(context.Background(), ag, h, noSession(t), "/clear", "clear", "") + if !errors.Is(err, errClearScrollback) { + t.Errorf("runTUICommand(/clear) error = %v, want errClearScrollback", err) + } + _ = output + if got := len(h.Snapshot()); got != 2 { + t.Errorf("history has %d message(s) after the TUI /clear, want 2 (untouched)", got) + } +} + +// toolHeavyHistory builds n turns, each a user question, a tool call, its +// result, and an assistant answer — "tool-heavy" per 11.2's AC for +// /compact, and long enough that CompactionBoundary finds something to cut +// with a small keepTurns. +func toolHeavyHistory(n int) *history.History { + h := history.New() + for i := 0; i < n; i++ { + id := fmt.Sprintf("call_%d", i) + h.Append(llm.Message{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: fmt.Sprintf("question number %d, padded out so it has some real length to it", i)}}}) + h.Append(llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.ToolUse{ID: id, Name: "read_file", Input: json.RawMessage(fmt.Sprintf(`{"path":"file_%d.go"}`, i))}}}) + h.Append(llm.Message{Role: llm.RoleUser, Blocks: []llm.Block{llm.ToolResult{ToolUseID: id, Content: fmt.Sprintf("contents of file_%d.go, also padded to a realistic length", i)}}}) + h.Append(llm.Message{Role: llm.RoleAssistant, Blocks: []llm.Block{llm.Text{Text: fmt.Sprintf("answer number %d, likewise padded out", i)}}}) + } + return h +} + +// TestCompactCommandLowersEstimateOnToolHeavyHistory is 11.2's AC for +// /compact: forcing it on a long tool-heavy history lowers the token +// estimate and still passes Validate(). +func TestCompactCommandLowersEstimateOnToolHeavyHistory(t *testing.T) { + h := toolHeavyHistory(8) + before := history.EstimateTokens(h.Snapshot()) + + provider := &fakeChatProvider{reply: "concise summary of the earlier turns"} + ag := agent.New(provider, tools.NewRegistry(), h, "", 1024, agent.Guards{}, 0, agent.AutoApprove) + ag.SetCompactionPolicy(agent.CompactionPolicy{KeepTurns: 2}) + + var out bytes.Buffer + if _, err := handleCommand(context.Background(), &out, ag, h, noSession(t), "compact", ""); err != nil { + t.Fatalf("handleCommand(compact) error = %v", err) + } + + if err := h.Validate(); err != nil { + t.Fatalf("Validate() after /compact error = %v", err) + } + after := history.EstimateTokens(h.Snapshot()) + if after >= before { + t.Errorf("estimated tokens after /compact = %d, want less than before (%d)", after, before) + } + if got := out.String(); !strings.Contains(got, fmt.Sprintf("%d -> %d", before, after)) { + t.Errorf("output = %q, want it to report %d -> %d", got, before, after) + } +} + +// TestCompactCommandNoopWhenNothingToCompact covers the case AC 11.2 leaves +// implicit: a short history has nothing eligible to summarize, so /compact +// must not call the provider or change history at all. +func TestCompactCommandNoopWhenNothingToCompact(t *testing.T) { + h := history.New() + h.Append(llm.Message{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "hi"}}}) + before := h.Snapshot() + + provider := &fakeChatProvider{reply: "should never be requested"} + ag := agent.New(provider, tools.NewRegistry(), h, "", 1024, agent.Guards{}, 0, agent.AutoApprove) + ag.SetCompactionPolicy(agent.CompactionPolicy{KeepTurns: 6}) + + var out bytes.Buffer + if _, err := handleCommand(context.Background(), &out, ag, h, noSession(t), "compact", ""); err != nil { + t.Fatalf("handleCommand(compact) error = %v", err) + } + if diff := cmp.Diff(before, h.Snapshot()); diff != "" { + t.Errorf("history changed by a no-op /compact (-before +after):\n%s", diff) + } + if got := out.String(); !strings.Contains(got, "nothing to compact") { + t.Errorf("output = %q, want it to report nothing to compact", got) + } +} + +// TestExitCommandSavesSessionBeforeSignalingExit is 11.2's AC for /exit: +// the session file is written before the command returns, and it leaves +// through the same path as Ctrl+D (runREPL returns nil, not an error, and +// any input after /exit is never processed). +func TestExitCommandSavesSessionBeforeSignalingExit(t *testing.T) { + dir := t.TempDir() + sess, err := newSession(dir, "work") + if err != nil { + t.Fatalf("newSession() error = %v", err) + } + h := history.New() + provider := &fakeChatProvider{reply: "ok"} + ag := agent.New(provider, tools.NewRegistry(), h, "", 1024, agent.Guards{}, 0, agent.AutoApprove) + + var out bytes.Buffer + in := "hello\n/exit\nshould never run\n" + runTurn := newTurnRunner(&config.Config{Stream: true}, ag, &out) + if err := runREPL(context.Background(), strings.NewReader(in), &out, ag, h, sess, runTurn); err != nil { + t.Fatalf("runREPL() error = %v, want nil (the same clean return Ctrl+D produces)", err) + } + + saved, err := history.Load(sess.path("work")) + if err != nil { + t.Fatalf("Load() session file error = %v, want /exit to have written it first", err) + } + if diff := cmp.Diff(h.Snapshot(), saved.Snapshot()); diff != "" { + t.Errorf("saved session differs from in-memory history (-memory +saved):\n%s", diff) + } + if strings.Contains(out.String(), "should never run") { + t.Error("input after /exit was processed, want runREPL to stop at /exit") + } +} + +// fakeSwitchableProvider adds llm.ModelSwitcher to fakeChatProvider, for +// testing /model without a real backend. unknownModel, if set, makes +// ValidateModel fail for that one name; every other name validates. +type fakeSwitchableProvider struct { + *fakeChatProvider + unknownModel string + model string +} + +func (f *fakeSwitchableProvider) ValidateModel(_ context.Context, model string) error { + if model == f.unknownModel { + return fmt.Errorf("model %q not found", model) + } + return nil +} + +func (f *fakeSwitchableProvider) SetModel(model string) { + f.model = model +} + +// readFileVia runs the registry's read_file tool for path, for assertions +// about which paths the sandbox currently accepts. +func readFileVia(t *testing.T, registry *tools.Registry, path string) (string, error) { + t.Helper() + tool, err := registry.Get("read_file") + if err != nil { + t.Fatalf("Get(read_file) error = %v", err) + } + input, err := json.Marshal(map[string]string{"path": path}) + if err != nil { + t.Fatalf("Marshal() error = %v", err) + } + return tool.Run(context.Background(), input) +} + +// TestCdCommandRerootsWorkspaceSandbox is 11.3's AC for /cd: after /cd, a +// read_file reachable only from the new root succeeds and one under the +// old root is rejected. +func TestCdCommandRerootsWorkspaceSandbox(t *testing.T) { + oldRoot := t.TempDir() + if err := os.WriteFile(filepath.Join(oldRoot, "old.txt"), []byte("old"), 0o600); err != nil { + t.Fatalf("WriteFile(old.txt) error = %v", err) + } + newRoot := t.TempDir() + if err := os.WriteFile(filepath.Join(newRoot, "new.txt"), []byte("new"), 0o600); err != nil { + t.Fatalf("WriteFile(new.txt) error = %v", err) + } + + registry, err := buildRegistry(&config.Config{Workspace: oldRoot}) + if err != nil { + t.Fatalf("buildRegistry() error = %v", err) + } + h := history.New() + ag := agent.New(&fakeChatProvider{reply: "ok"}, registry, h, "", 1024, agent.Guards{}, 0, agent.AutoApprove) + + if _, err := readFileVia(t, registry, "old.txt"); err != nil { + t.Fatalf("read_file(old.txt) before /cd error = %v", err) + } + + var out bytes.Buffer + if _, err := handleCommand(context.Background(), &out, ag, h, noSession(t), "cd", newRoot); err != nil { + t.Fatalf("handleCommand(cd) error = %v", err) + } + if !strings.Contains(out.String(), "workspace root changed") { + t.Errorf("output = %q, want a confirmation", out.String()) + } + + if _, err := readFileVia(t, registry, "new.txt"); err != nil { + t.Errorf("read_file(new.txt) after /cd error = %v, want it reachable from the new root", err) + } + if _, err := readFileVia(t, registry, filepath.Join(oldRoot, "old.txt")); err == nil { + t.Error("read_file(old.txt) after /cd error = nil, want the old root rejected as escaping the sandbox") + } +} + +// TestCdCommandRejectsMissingTarget is 11.3's other AC for /cd: a missing +// target is refused and the original root stays active. +func TestCdCommandRejectsMissingTarget(t *testing.T) { + root := t.TempDir() + if err := os.WriteFile(filepath.Join(root, "marker.txt"), []byte("marker"), 0o600); err != nil { + t.Fatalf("WriteFile(marker.txt) error = %v", err) + } + registry, err := buildRegistry(&config.Config{Workspace: root}) + if err != nil { + t.Fatalf("buildRegistry() error = %v", err) + } + h := history.New() + ag := agent.New(&fakeChatProvider{reply: "ok"}, registry, h, "", 1024, agent.Guards{}, 0, agent.AutoApprove) + + var out bytes.Buffer + if _, err := handleCommand(context.Background(), &out, ag, h, noSession(t), "cd", filepath.Join(root, "does-not-exist")); err != nil { + t.Fatalf("handleCommand(cd) error = %v", err) + } + if !strings.Contains(out.String(), "cd failed") { + t.Errorf("output = %q, want a failure message", out.String()) + } + if _, err := readFileVia(t, registry, "marker.txt"); err != nil { + t.Errorf("read_file(marker.txt) after a refused /cd error = %v, want the original root still active", err) + } +} + +// TestModelCommandRejectsUnknownWithoutMutating is 11.3's AC for /model: a +// name the provider does not know errors without mutating the provider's +// model or the compaction policy. +func TestModelCommandRejectsUnknownWithoutMutating(t *testing.T) { + h := history.New() + provider := &fakeSwitchableProvider{fakeChatProvider: &fakeChatProvider{reply: "ok"}, unknownModel: "ghost-model", model: "original-model"} + ag := agent.New(provider, tools.NewRegistry(), h, "", 1024, agent.Guards{}, 0, agent.AutoApprove) + ag.SetCompactionPolicy(agent.CompactionPolicy{ContextWindow: 500, TriggerFraction: 0.75, KeepTurns: 6}) + + var out bytes.Buffer + if _, err := handleCommand(context.Background(), &out, ag, h, noSession(t), "model", "ghost-model"); err != nil { + t.Fatalf("handleCommand(model) error = %v", err) + } + if !strings.Contains(out.String(), "model failed") { + t.Errorf("output = %q, want a failure message", out.String()) + } + if provider.model != "original-model" { + t.Errorf("provider.model = %q, want it unchanged after a rejected /model", provider.model) + } + if got := ag.CompactionPolicy().ContextWindow; got != 500 { + t.Errorf("CompactionPolicy().ContextWindow = %d, want 500 unchanged", got) + } +} + +// TestModelCommandErrorsWhenProviderCannotSwitch covers a provider that +// doesn't implement llm.ModelSwitcher at all (every fakeChatProvider in the +// rest of this file, and any future provider that never adds it). +func TestModelCommandErrorsWhenProviderCannotSwitch(t *testing.T) { + h := history.New() + ag := agent.New(&fakeChatProvider{reply: "ok"}, tools.NewRegistry(), h, "", 1024, agent.Guards{}, 0, agent.AutoApprove) + + var out bytes.Buffer + if _, err := handleCommand(context.Background(), &out, ag, h, noSession(t), "model", "some-model"); err != nil { + t.Fatalf("handleCommand(model) error = %v", err) + } + if !strings.Contains(out.String(), "does not support switching models") { + t.Errorf("output = %q, want a does-not-support message", out.String()) + } +} + +// TestModelCommandSwitchingSmallerWindowTriggersCompactionNextTurn is +// 11.3's AC: switching to a smaller context window triggers the +// compaction check on the next turn. +func TestModelCommandSwitchingSmallerWindowTriggersCompactionNextTurn(t *testing.T) { + h := toolHeavyHistory(8) + before := history.EstimateTokens(h.Snapshot()) + + provider := &fakeSwitchableProvider{fakeChatProvider: &fakeChatProvider{reply: "concise summary of the earlier turns"}, model: "old-model"} + ag := agent.New(provider, tools.NewRegistry(), h, "", 1024, agent.Guards{}, 0, agent.AutoApprove) + ag.SetCompactionPolicy(agent.CompactionPolicy{ContextWindow: 1_000_000, TriggerFraction: 0.75, KeepTurns: 2}) + + newWindow := before / 2 + var out bytes.Buffer + if _, err := handleCommand(context.Background(), &out, ag, h, noSession(t), "model", fmt.Sprintf("new-model %d", newWindow)); err != nil { + t.Fatalf("handleCommand(model) error = %v", err) + } + if !strings.Contains(out.String(), "compaction will run on the next turn") { + t.Errorf("output = %q, want a compaction warning", out.String()) + } + if provider.model != "new-model" { + t.Errorf("provider.model = %q, want %q", provider.model, "new-model") + } + + if _, err := ag.Run(context.Background(), "one more question"); err != nil { + t.Fatalf("Run() error = %v", err) + } + if err := h.Validate(); err != nil { + t.Fatalf("Validate() after the next turn error = %v", err) + } + if after := history.EstimateTokens(h.Snapshot()); after >= before { + t.Errorf("estimated tokens after the next turn = %d, want less than the pre-/model estimate (%d): the smaller window should have triggered compaction", after, before) + } +} + +func TestLevenshtein(t *testing.T) { + tests := []struct { + a, b string + want int + }{ + {"", "", 0}, + {"help", "help", 0}, + {"hlep", "help", 2}, + {"sav", "save", 1}, + {"kitten", "sitting", 3}, + } + for _, tt := range tests { + if got := levenshtein(tt.a, tt.b); got != tt.want { + t.Errorf("levenshtein(%q, %q) = %d, want %d", tt.a, tt.b, got, tt.want) + } + } +} 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..af78ebd 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,6 +53,7 @@ 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] } @@ -42,12 +65,13 @@ func TestChatProviderEnvFallback(t *testing.T) { 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] } @@ -59,8 +83,53 @@ func TestChatFlagOverridesEnv(t *testing.T) { 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) + } +} + +func TestChatToolsDefaultsToPromptedForOllama(t *testing.T) { + getCfg := stubRunChat(t) + root := newRootCmd(func(string) string { return "" }) + buf := &bytes.Buffer{} + root.SetOut(buf) + root.SetArgs([]string{"chat", "--provider=ollama"}) + + if err := root.Execute(); err != nil { + t.Fatalf("chat returned error: %v", err) + } + if got := getCfg().Tools; got != config.ToolsPrompted { + t.Errorf("Tools = %q, want --provider=ollama with no explicit --tools to default to %q", got, config.ToolsPrompted) + } +} + +func TestChatToolsDefaultsToNativeForAnthropic(t *testing.T) { + getCfg := stubRunChat(t) + root := newRootCmd(func(string) string { return "" }) + buf := &bytes.Buffer{} + root.SetOut(buf) + root.SetArgs([]string{"chat", "--provider=anthropic"}) + + if err := root.Execute(); err != nil { + t.Fatalf("chat returned error: %v", err) + } + if got := getCfg().Tools; got != config.ToolsNative { + t.Errorf("Tools = %q, want --provider=anthropic to keep defaulting to %q", got, config.ToolsNative) + } +} + +func TestChatExplicitToolsOverridesOllamaDefault(t *testing.T) { + getCfg := stubRunChat(t) + root := newRootCmd(func(string) string { return "" }) + buf := &bytes.Buffer{} + root.SetOut(buf) + root.SetArgs([]string{"chat", "--provider=ollama", "--tools=native"}) + + if err := root.Execute(); err != nil { + t.Fatalf("chat returned error: %v", err) + } + if got := getCfg().Tools; got != config.ToolsNative { + t.Errorf("Tools = %q, want an explicit --tools=native to override the ollama default", got) } } diff --git a/cmd/pico/sessions.go b/cmd/pico/sessions.go new file mode 100644 index 0000000..bd5a8aa --- /dev/null +++ b/cmd/pico/sessions.go @@ -0,0 +1,116 @@ +package main + +import ( + "errors" + "fmt" + "os" + "path/filepath" + "sort" + "strings" + "time" + + "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 +} + +// recent returns up to n saved session names, most recently modified +// first. Errors are swallowed into an empty list: the home screen that +// calls this must render even when the sessions directory is unreadable. +func (s *session) recent(n int) []string { + entries, err := os.ReadDir(s.dir) + if err != nil { + return nil + } + type entry struct { + name string + mod time.Time + } + var found []entry + for _, e := range entries { + if e.IsDir() || filepath.Ext(e.Name()) != ".json" { + continue + } + info, err := e.Info() + if err != nil { + continue + } + found = append(found, entry{strings.TrimSuffix(e.Name(), ".json"), info.ModTime()}) + } + sort.Slice(found, func(i, j int) bool { return found[i].mod.After(found[j].mod) }) + if len(found) > n { + found = found[:n] + } + names := make([]string, 0, len(found)) + for _, e := range found { + names = append(names, e.name) + } + return names +} diff --git a/cmd/pico/testdata/golden/help.txt b/cmd/pico/testdata/golden/help.txt new file mode 100644 index 0000000..b8a8a43 --- /dev/null +++ b/cmd/pico/testdata/golden/help.txt @@ -0,0 +1,10 @@ +/help list available commands +/usage show cumulative and last-turn token usage +/new start a new, unsaved session +/save [name] save the current session to disk +/load load a saved session from disk +/clear clear the terminal scrollback (the conversation history is kept) +/compact summarize older turns now and report the token estimate +/exit save the session and exit, like Ctrl+D +/cd re-root the workspace sandbox +/model [context-window] switch models, validated against the active provider 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..e89bb4c 100644 --- a/go.mod +++ b/go.mod @@ -4,20 +4,51 @@ 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/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834 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 + golang.org/x/term v0.36.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/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 +56,13 @@ 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/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..656737e --- /dev/null +++ b/internal/agent/compact.go @@ -0,0 +1,142 @@ +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 +} + +// CompactionPolicy returns the policy currently in effect, so a caller (the +// /model command) can change just ContextWindow without having to know or +// clobber TriggerFraction and KeepTurns. +func (a *Agent) CompactionPolicy() CompactionPolicy { + return a.compaction +} + +// 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 + } + + if err := a.compactNow(ctx, messages, p.KeepTurns); err != nil { + slog.Warn("agent: skipping compaction", "error", err) + } +} + +// ForceCompact runs the same summarization pass as maybeCompact, but +// unconditionally rather than waiting for the trigger threshold — for the +// /compact command (CLAUDE.md phase 11.2). It reports the estimated token +// count before and after so the caller can show the delta. If there are +// KeepTurns or fewer turns to summarize (including when no policy has been +// configured), it is a no-op and before == after. +func (a *Agent) ForceCompact(ctx context.Context) (before, after int, err error) { + messages := a.history.Snapshot() + before = history.EstimateTokens(messages) + + if err := a.compactNow(ctx, messages, a.compaction.KeepTurns); err != nil { + return before, before, err + } + return before, history.EstimateTokens(a.history.Snapshot()), nil +} + +// compactNow summarizes everything before the boundary that keeps the last +// keepTurns turns verbatim, then commits it via history.Compact. A no-op if +// there are keepTurns or fewer turns. +func (a *Agent) compactNow(ctx context.Context, messages []llm.Message, keepTurns int) error { + if keepTurns <= 0 { + return nil + } + boundary := history.CompactionBoundary(messages, keepTurns) + if boundary <= 0 { + return nil + } + + summary, err := a.summarize(ctx, messages[:boundary]) + if err != nil { + return fmt.Errorf("agent: summarize: %w", err) + } + return a.history.Compact(keepTurns, summary) +} + +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..24818fa --- /dev/null +++ b/internal/agent/compact_test.go @@ -0,0 +1,232 @@ +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) + } +} + +// TestForceCompactRunsRegardlessOfThreshold is 11.2's AC: /compact forces +// the compaction pass immediately (no ContextWindow/TriggerFraction +// configured at all, unlike maybeCompact) and reports the estimated token +// count before and after. +func TestForceCompactRunsRegardlessOfThreshold(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 := history.EstimateTokens(h.Snapshot()) + + provider := &scriptedProvider{responses: []*llm.Response{textResponse("everything before this was summarized")}} + a := New(provider, tools.NewRegistry(), h, "", 1024, Guards{}, 0, AutoApprove) + a.SetCompactionPolicy(CompactionPolicy{KeepTurns: 2}) + + gotBefore, gotAfter, err := a.ForceCompact(context.Background()) + if err != nil { + t.Fatalf("ForceCompact() error = %v", err) + } + if gotBefore != before { + t.Errorf("ForceCompact() before = %d, want %d (matching EstimateTokens on the untouched history)", gotBefore, before) + } + if gotAfter >= gotBefore { + t.Errorf("ForceCompact() after = %d, want less than before (%d)", gotAfter, gotBefore) + } + if err := h.Validate(); err != nil { + t.Fatalf("Validate() after ForceCompact() error = %v", err) + } +} + +// TestForceCompactNoopWithoutEnoughTurns covers KeepTurns unset (the zero +// value CompactionPolicy) and KeepTurns exceeding the number of turns +// present — both must be no-ops that never call the provider. +func TestForceCompactNoopWithoutEnoughTurns(t *testing.T) { + h := history.New() + h.Append(llm.Message{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "hi"}}}) + before := h.Snapshot() + + provider := &alwaysErrorsProvider{err: errors.New("must not be called")} + a := New(provider, tools.NewRegistry(), h, "", 1024, Guards{}, 0, AutoApprove) + + gotBefore, gotAfter, err := a.ForceCompact(context.Background()) + if err != nil { + t.Fatalf("ForceCompact() error = %v, want nil for a no-op", err) + } + if gotBefore != gotAfter { + t.Errorf("ForceCompact() = (%d, %d), want before == after for a no-op", gotBefore, gotAfter) + } + if diff := cmpMessage(h.Snapshot()[0], before[0]); diff != "" { + t.Errorf("history changed by a no-op ForceCompact(): %s", diff) + } +} + +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..d232cf1 --- /dev/null +++ b/internal/agent/loop.go @@ -0,0 +1,349 @@ +// 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 +} + +// Provider returns the provider this Agent talks to, for a caller (the +// /model command) that needs to check a provider-specific capability via a +// type assertion — CLAUDE.md invariant 2 keeps those capabilities off the +// shared Provider interface, so reaching one needs the concrete value. +func (a *Agent) Provider() llm.Provider { + return a.provider +} + +// Tools returns the registry this Agent executes tool calls against, for a +// caller (the /cd command) that needs to reach a specific registered tool +// directly rather than through a ToolUse block. +func (a *Agent) Tools() *tools.Registry { + return a.tools +} + +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..80f36b1 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,29 @@ 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") + // 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 +80,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. @@ -92,6 +131,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 +148,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..49dbcb3 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -106,6 +106,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/model_test.go b/internal/llm/anthropic/model_test.go new file mode 100644 index 0000000..b6365ab --- /dev/null +++ b/internal/llm/anthropic/model_test.go @@ -0,0 +1,78 @@ +package anthropic + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "testing" + + "github.com/reno/pico-code/internal/llm" +) + +const notFoundBody = `{"type":"error","error":{"type":"not_found_error","message":"model not found"}}` + +const modelInfoBody = `{ + "id": "claude-known-model", "type": "model", + "display_name": "Known Model", "created_at": "2025-01-01T00:00:00Z", + "max_input_tokens": 200000, "max_tokens": 8192, + "capabilities": {} +}` + +// TestValidateModelKnownAndUnknown is 11.3's AC: /model with a name the +// provider does not know errors. mapError already maps a 404 to +// ErrNotFound (errors_test.go), so this proves ValidateModel's request +// actually reaches the Models API and that mapError still applies to it. +func TestValidateModelKnownAndUnknown(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.URL.Path == "/v1/models/claude-known-model" { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(modelInfoBody)) + return + } + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(notFoundBody)) + })) + defer srv.Close() + + p := testProvider(srv.URL) + + if err := p.ValidateModel(context.Background(), "claude-known-model"); err != nil { + t.Errorf("ValidateModel(known) error = %v, want nil", err) + } + if err := p.ValidateModel(context.Background(), "claude-unknown-model"); !errors.Is(err, ErrNotFound) { + t.Errorf("ValidateModel(unknown) error = %v, want wrapping %v", err, ErrNotFound) + } +} + +// TestSetModelChangesTheModelChatSends is 11.3's AC that switching models +// actually takes effect: after SetModel, Chat's request targets the new +// model, not the one Provider was constructed with. +func TestSetModelChangesTheModelChatSends(t *testing.T) { + var gotModel string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body struct { + Model string `json:"model"` + } + _ = json.NewDecoder(r.Body).Decode(&body) + gotModel = body.Model + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(okMessageBody)) + })) + defer srv.Close() + + p := testProvider(srv.URL) + p.SetModel("claude-new-model") + + if _, err := p.Chat(context.Background(), testChatRequest()); err != nil { + t.Fatalf("Chat() error = %v", err) + } + if gotModel != "claude-new-model" { + t.Errorf("request model = %q, want %q", gotModel, "claude-new-model") + } +} + +var _ llm.ModelSwitcher = (*Provider)(nil) diff --git a/internal/llm/anthropic/provider.go b/internal/llm/anthropic/provider.go index 3b0e830..efe5295 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,29 @@ 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") +// ValidateModel implements llm.ModelSwitcher by looking model up through +// the Models API; an unknown model surfaces as ErrNotFound via mapError, +// the same sentinel Chat's own 404s map to. +func (p *Provider) ValidateModel(ctx context.Context, model string) error { + _, err := p.client.Models.Get(ctx, model, sdk.ModelGetParams{}) + if err != nil { + return mapError(err) + } + return nil } + +// SetModel implements llm.ModelSwitcher. +func (p *Provider) SetModel(model string) { + p.model = model +} + +var _ llm.ModelSwitcher = (*Provider)(nil) 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/capability.go b/internal/llm/ollama/capability.go index b34fa44..f58b1d4 100644 --- a/internal/llm/ollama/capability.go +++ b/internal/llm/ollama/capability.go @@ -35,44 +35,74 @@ func (p *Provider) supportsTools(ctx context.Context) (bool, error) { } func (p *Provider) probeToolSupport(ctx context.Context) (bool, error) { - body, err := json.Marshal(api.ShowRequest{Model: p.model}) + show, err := p.probeShow(ctx, p.model) if err != nil { - return false, fmt.Errorf("ollama: marshal /api/show request: %w", err) + return false, err + } + for _, c := range show.Capabilities { + if c == model.CapabilityTools { + return true, nil + } + } + return false, nil +} + +// probeShow calls /api/show for modelName, the low-level request both +// probeToolSupport (for p.model, cached via supportsTools) and +// ValidateModel (for an arbitrary candidate, uncached) build on. A model +// /api/show doesn't recognize comes back as a non-2xx status, surfaced here +// as a plain error. +func (p *Provider) probeShow(ctx context.Context, modelName string) (api.ShowResponse, error) { + body, err := json.Marshal(api.ShowRequest{Model: modelName}) + if err != nil { + return api.ShowResponse{}, fmt.Errorf("ollama: marshal /api/show request: %w", err) } httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL+"/api/show", bytes.NewReader(body)) if err != nil { - return false, fmt.Errorf("ollama: %w", err) + return api.ShowResponse{}, fmt.Errorf("ollama: %w", err) } httpReq.Header.Set("Content-Type", "application/json") res, err := p.httpClient.Do(httpReq) if err != nil { - return false, fmt.Errorf("ollama: %w", err) + return api.ShowResponse{}, fmt.Errorf("ollama: %w", err) } defer func() { _ = res.Body.Close() }() respBody, err := io.ReadAll(res.Body) if err != nil { - return false, fmt.Errorf("ollama: read /api/show response: %w", err) + return api.ShowResponse{}, fmt.Errorf("ollama: read /api/show response: %w", err) } if res.StatusCode >= http.StatusBadRequest { - return false, fmt.Errorf("ollama: /api/show failed with status %d: %s", res.StatusCode, bytes.TrimSpace(respBody)) + return api.ShowResponse{}, fmt.Errorf("ollama: /api/show failed with status %d: %s", res.StatusCode, bytes.TrimSpace(respBody)) } var show api.ShowResponse if err := json.Unmarshal(respBody, &show); err != nil { - return false, fmt.Errorf("ollama: decode /api/show response: %w", err) + return api.ShowResponse{}, fmt.Errorf("ollama: decode /api/show response: %w", err) } + return show, nil +} - for _, c := range show.Capabilities { - if c == model.CapabilityTools { - return true, nil - } - } - return false, nil +// ValidateModel implements llm.ModelSwitcher by probing /api/show for +// modelName. Ollama's /api/show has no context-window field that's stable +// across model families (it's buried in a per-family key inside +// ModelInfo), so unlike Anthropic's ValidateModel this never derives one; +// num_ctx stays the explicit, user-set knob CLAUDE.md already requires it +// to be. +func (p *Provider) ValidateModel(ctx context.Context, modelName string) error { + _, err := p.probeShow(ctx, modelName) + return err } +// SetModel implements llm.ModelSwitcher. +func (p *Provider) SetModel(modelName string) { + p.model = modelName +} + +var _ llm.ModelSwitcher = (*Provider)(nil) + // checkToolSupport gates a request that carries tool definitions: it is a // no-op for a tool-less request (the common case, and the only case in // --tools=prompted mode, where the loop never populates Request.Tools). 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/model_test.go b/internal/llm/ollama/model_test.go new file mode 100644 index 0000000..70ec91a --- /dev/null +++ b/internal/llm/ollama/model_test.go @@ -0,0 +1,77 @@ +package ollama + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/reno/pico-code/internal/llm" +) + +// TestValidateModelKnownAndUnknown is 11.3's AC: /model with a name the +// provider does not know errors. A model /api/show has never pulled comes +// back as a non-2xx status (a real Ollama server returns 404), which +// ValidateModel surfaces as a plain error. +func TestValidateModelKnownAndUnknown(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body struct { + Model string `json:"model"` + } + _ = json.NewDecoder(r.Body).Decode(&body) + w.Header().Set("Content-Type", "application/json") + if body.Model == "known-model" { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"capabilities": ["completion"]}`)) + return + } + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"error":"model not found"}`)) + })) + defer srv.Close() + + p := &Provider{httpClient: http.DefaultClient, baseURL: srv.URL, model: "known-model"} + + if err := p.ValidateModel(context.Background(), "known-model"); err != nil { + t.Errorf("ValidateModel(known) error = %v, want nil", err) + } + if err := p.ValidateModel(context.Background(), "unknown-model"); err == nil { + t.Error("ValidateModel(unknown) error = nil, want an error for a model /api/show doesn't recognize") + } +} + +// TestSetModelChangesTheModelChatSends mirrors the Anthropic adapter's +// equivalent test: after SetModel, Chat's request targets the new model. +func TestSetModelChangesTheModelChatSends(t *testing.T) { + var gotModel string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body struct { + Model string `json:"model"` + } + _ = json.NewDecoder(r.Body).Decode(&body) + gotModel = body.Model + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "model": "new-model", + "message": {"role": "assistant", "content": "hi"}, + "done": true, + "done_reason": "stop" + }`)) + })) + defer srv.Close() + + p := &Provider{httpClient: http.DefaultClient, baseURL: srv.URL, model: "old-model", numCtx: 8192} + p.SetModel("new-model") + + if _, err := p.Chat(context.Background(), llm.Request{ + Messages: []llm.Message{{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "hi"}}}}, + }); err != nil { + t.Fatalf("Chat() error = %v", err) + } + if gotModel != "new-model" { + t.Errorf("request model = %q, want %q", gotModel, "new-model") + } +} + +var _ llm.ModelSwitcher = (*Provider)(nil) 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..fca55e9 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) } @@ -60,6 +60,81 @@ func TestFromResponseStringArgumentsMatchesObjectForm(t *testing.T) { } } +// TestFromMessageRecoversNarratedToolCall is a regression test: some +// models that report native tool support (e.g. qwen2.5-coder:7b via +// Ollama) still write the call as plain JSON text in Content instead of +// populating ToolCalls — reproduced live against real Ollama. fromMessage +// must recover a real ToolUse from that text instead of showing it as an +// ordinary reply. +func TestFromMessageRecoversNarratedToolCall(t *testing.T) { + tools := []llm.ToolDefinition{{Name: "read_file"}} + m := api.Message{Role: "assistant", Content: `{"name": "read_file", "arguments": {"path": "zen_of_python.txt"}}`} + + blocks, err := fromMessage(m, tools) + if err != nil { + t.Fatalf("fromMessage() error = %v", err) + } + if len(blocks) != 1 { + t.Fatalf("len(blocks) = %d, want 1", len(blocks)) + } + tu, ok := blocks[0].(llm.ToolUse) + if !ok { + t.Fatalf("blocks[0] = %T, want llm.ToolUse", blocks[0]) + } + if tu.Name != "read_file" { + t.Errorf("Name = %q, want %q", tu.Name, "read_file") + } + if string(tu.Input) != `{"path": "zen_of_python.txt"}` { + t.Errorf("Input = %s, want the arguments object verbatim", tu.Input) + } +} + +// TestFromMessageLeavesUnknownNamedJSONAsText proves the recovery is +// scoped to tools actually offered this turn, so a model's legitimate +// "please output JSON" answer that happens to have a "name" field isn't +// misread as a call it never intended to make. +func TestFromMessageLeavesUnknownNamedJSONAsText(t *testing.T) { + tools := []llm.ToolDefinition{{Name: "read_file"}} + content := `{"name": "Ada Lovelace", "arguments": {"born": 1815}}` + m := api.Message{Role: "assistant", Content: content} + + blocks, err := fromMessage(m, tools) + if err != nil { + t.Fatalf("fromMessage() error = %v", err) + } + want := []llm.Block{llm.Text{Text: content}} + if diff := cmp.Diff(want, blocks); diff != "" { + t.Errorf("blocks mismatch (-want +got):\n%s", diff) + } +} + +// TestFromMessageDoesNotRecoverWhenRealToolCallsPresent proves the +// recovery only kicks in when ToolCalls is genuinely empty, never as a +// second guess alongside a real structured call. +func TestFromMessageDoesNotRecoverWhenRealToolCallsPresent(t *testing.T) { + tools := []llm.ToolDefinition{{Name: "read_file"}} + args := api.NewToolCallFunctionArguments() + args.Set("path", "b.txt") + m := api.Message{ + Role: "assistant", + Content: `{"name": "read_file", "arguments": {"path": "a.txt"}}`, + ToolCalls: []api.ToolCall{ + {Function: api.ToolCallFunction{Name: "read_file", Arguments: args}}, + }, + } + + blocks, err := fromMessage(m, tools) + if err != nil { + t.Fatalf("fromMessage() error = %v", err) + } + if len(blocks) != 2 { + t.Fatalf("len(blocks) = %d, want 2 (the raw text plus the real tool call, no recovery attempted)", len(blocks)) + } + if _, ok := blocks[0].(llm.Text); !ok { + t.Errorf("blocks[0] = %T, want llm.Text (content passed through untouched)", blocks[0]) + } +} + func TestFromResponseSynthesizesMissingID(t *testing.T) { got := decodeFixture(t, "testdata/golden/response_missing_id.json") if len(got.Message.Blocks) != 1 { diff --git a/internal/llm/ollama/stream.go b/internal/llm/ollama/stream.go new file mode 100644 index 0000000..73516dd --- /dev/null +++ b/internal/llm/ollama/stream.go @@ -0,0 +1,179 @@ +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 + // held back and re-checked against narratedToolCall once the full + // response is in, instead of being sent immediately as TextDelta, + // for as long as the accumulated text still looks like it could be + // one (see maybeNarratedCall below). Flushed unchanged the moment + // that stops being true, so an ordinary reply pays no cost. + var pending []string + maybeNarratedCall := true + flushPending := func() bool { + for _, text := range pending { + if !send(llm.TextDelta{Text: text}) { + return false + } + } + pending = nil + return true + } + + 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 len(chunk.Message.ToolCalls) > 0 && maybeNarratedCall { + // Real structured tool_calls showed up after all, so + // whatever text preceded them was never a narrated call — + // release it as ordinary text. + maybeNarratedCall = false + if !flushPending() { + return + } + } + + if chunk.Message.Content != "" { + respText.WriteString(chunk.Message.Content) + if trimmed := strings.TrimSpace(respText.String()); maybeNarratedCall && trimmed != "" && !strings.HasPrefix(trimmed, "{") { + // Once we've seen a first non-whitespace character that + // isn't '{', this can no longer become a bare JSON + // object — give up the candidacy for good. + maybeNarratedCall = false + } + if maybeNarratedCall { + pending = append(pending, chunk.Message.Content) + } else { + if !flushPending() { + return + } + 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) + if maybeNarratedCall { + if name, input, ok := narratedToolCall(respText.String(), req.Tools); ok { + if !send(llm.ToolUseStart{ID: "ollama_call_0", Name: name}) { + return + } + if !send(llm.ToolUseDone{ID: "ollama_call_0", Input: input}) { + return + } + } else if !flushPending() { + return + } + } + 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..b54aece --- /dev/null +++ b/internal/llm/ollama/stream_test.go @@ -0,0 +1,203 @@ +package ollama + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "strings" + "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) + } +} + +// narratedStreamServer serves an NDJSON stream whose Content chunks join +// into text, with no structured tool_calls anywhere in the response — +// reproducing what qwen2.5-coder:7b actually sends over the wire for a +// narrated tool call. +func narratedStreamServer(t *testing.T, contentChunks ...string) *httptest.Server { + t.Helper() + var b strings.Builder + for _, c := range contentChunks { + escaped, err := jsonEscape(c) + if err != nil { + t.Fatalf("jsonEscape(%q) error = %v", c, err) + } + b.WriteString(`{"model":"qwen2.5-coder:7b","message":{"role":"assistant","content":` + escaped + `},"done":false}` + "\n") + } + b.WriteString(`{"model":"qwen2.5-coder:7b","message":{"role":"assistant","content":""},"done":true,"done_reason":"stop","prompt_eval_count":10,"eval_count":5}` + "\n") + + srv := 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(b.String())) + })) + t.Cleanup(srv.Close) + return srv +} + +func jsonEscape(s string) (string, error) { + b, err := json.Marshal(s) + return string(b), err +} + +// TestStreamRecoversNarratedToolCallAcrossChunks is the streaming half of +// the narratedToolCall regression: the same recovery model.go/translate.go +// gives the non-streaming path must also work when the narrated JSON +// arrives split across several NDJSON chunks, exactly as real Ollama +// streaming delivers it token by token. +func TestStreamRecoversNarratedToolCallAcrossChunks(t *testing.T) { + srv := narratedStreamServer(t, + `{"name": "read_file", `, + `"arguments": {"path": "a.txt"}}`, + ) + + req := llm.Request{ + Messages: []llm.Message{{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "read a.txt"}}}}, + Tools: []llm.ToolDefinition{{Name: "read_file", Description: "read a file", InputSchema: []byte(`{"type":"object"}`)}}, + } + p := &Provider{httpClient: http.DefaultClient, baseURL: srv.URL, model: "qwen2.5-coder:7b", numCtx: 4096} + p.toolSupport = true + p.toolSupportOnce.Do(func() {}) + + ch, err := p.Stream(context.Background(), req) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + + var sawTextDelta bool + var done *llm.ToolUseDone + for e := range ch { + switch v := e.(type) { + case llm.TextDelta: + sawTextDelta = true + case llm.ToolUseDone: + done = &v + } + } + if sawTextDelta { + t.Error("a TextDelta was sent for the narrated call — it should have been recovered as a ToolUse before ever reaching the renderer") + } + if done == nil { + t.Fatal("no ToolUseDone event was sent — the narrated call was not recovered") + } + if string(done.Input) != `{"path": "a.txt"}` { + t.Errorf("Input = %s, want the arguments object verbatim", done.Input) + } +} + +// TestStreamFlushesUnrecoverableJSONAsTextOnceStreamEnds proves the +// candidacy check is scoped to tools actually offered: content that starts +// with '{' but never resolves into a recognized call is still delivered as +// ordinary text (just at the end, once the whole response is known, +// instead of live) rather than silently dropped. +func TestStreamFlushesUnrecoverableJSONAsTextOnceStreamEnds(t *testing.T) { + content := `{"name": "Ada Lovelace", "arguments": {"born": 1815}}` + srv := narratedStreamServer(t, content) + + req := llm.Request{ + Messages: []llm.Message{{Role: llm.RoleUser, Blocks: []llm.Block{llm.Text{Text: "who was Ada Lovelace?"}}}}, + Tools: []llm.ToolDefinition{{Name: "read_file", Description: "read a file", InputSchema: []byte(`{"type":"object"}`)}}, + } + p := &Provider{httpClient: http.DefaultClient, baseURL: srv.URL, model: "qwen2.5-coder:7b", numCtx: 4096} + p.toolSupport = true + p.toolSupportOnce.Do(func() {}) + + ch, err := p.Stream(context.Background(), req) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + resp, err := llm.CollectStream(context.Background(), ch) + if err != nil { + t.Fatalf("CollectStream() error = %v", err) + } + want := []llm.Block{llm.Text{Text: content}} + if diff := cmp.Diff(want, resp.Message.Blocks); diff != "" { + t.Errorf("Blocks mismatch (-want +got):\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..f52b57f 100644 --- a/internal/llm/ollama/translate.go +++ b/internal/llm/ollama/translate.go @@ -168,23 +168,36 @@ 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) { - blocks, err := fromMessage(resp.Message) +// 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, req.Tools) if err != nil { return nil, err } 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 } -func fromMessage(m api.Message) ([]llm.Block, error) { +// fromMessage translates m into canonical blocks. tools is the set offered +// this turn: when m carries no structured ToolCalls at all, its Content is +// checked against narratedToolCall — some models that report native tool +// support still write the call as plain JSON text instead of using the +// structured field (CLAUDE.md's "sometimes narrate a tool call in prose +// instead of calling it") — and recovered into a real ToolUse if it matches. +func fromMessage(m api.Message, tools []llm.ToolDefinition) ([]llm.Block, error) { + if len(m.ToolCalls) == 0 { + if name, input, ok := narratedToolCall(m.Content, tools); ok { + return []llm.Block{llm.ToolUse{ID: "ollama_call_0", Name: name, Input: input}}, nil + } + } + var blocks []llm.Block if m.Content != "" { blocks = append(blocks, llm.Text{Text: m.Content}) @@ -206,3 +219,32 @@ func fromMessage(m api.Message) ([]llm.Block, error) { } return blocks, nil } + +// narratedToolCall recognizes a tool call written as plain JSON text +// instead of Ollama's structured tool_calls field. It requires content, +// once trimmed, to be *exactly* one JSON object shaped +// {"name":string,"arguments":object} whose name matches one of tools — +// both narrow enough that a model's legitimate "please output JSON" answer +// isn't misread as a call it never intended to make. +func narratedToolCall(content string, tools []llm.ToolDefinition) (name string, input json.RawMessage, ok bool) { + trimmed := strings.TrimSpace(content) + if !strings.HasPrefix(trimmed, "{") { + return "", nil, false + } + var parsed struct { + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments"` + } + if err := json.Unmarshal([]byte(trimmed), &parsed); err != nil { + return "", nil, false + } + if parsed.Name == "" || len(parsed.Arguments) == 0 { + return "", nil, false + } + for _, td := range tools { + if td.Name == parsed.Name { + return parsed.Name, parsed.Arguments, true + } + } + return "", nil, false +} diff --git a/internal/llm/prompted/prompted.go b/internal/llm/prompted/prompted.go index cb65d8e..421f3d7 100644 --- a/internal/llm/prompted/prompted.go +++ b/internal/llm/prompted/prompted.go @@ -26,7 +26,15 @@ import ( // machinery needed here for that half of the job. const parseErrorToolName = "__prompted_parse_error__" -var fencedJSONRe = regexp.MustCompile("(?s)```(?:json)?\\s*\\n?(.*?)\\n?```") +// fencedJSONRe matches a fence labeled "json" or left unlabeled, requiring +// a real line break right after the label. The looser "\s*\n?" this +// replaced accepted any other label too (the "\n?" made the break +// optional, so e.g. "plaintext" right after the backticks was silently +// absorbed into the capture instead of rejecting the match) — a model's +// genuine final answer wrapped in an unrelated ```plaintext or ```python +// fence was misread as a failed tool call, spawning a bogus "tool not +// found" round-trip. +var fencedJSONRe = regexp.MustCompile("(?s)```(?:json)?[ \\t]*\\r?\\n(.*?)\\n?```") // Provider wraps inner to add prompted tool-calling. type Provider struct { diff --git a/internal/llm/prompted/prompted_test.go b/internal/llm/prompted/prompted_test.go index 4a1f0ed..1645db8 100644 --- a/internal/llm/prompted/prompted_test.go +++ b/internal/llm/prompted/prompted_test.go @@ -81,6 +81,18 @@ func TestChatParsesFencedToolCall(t *testing.T) { want: []llm.Block{llm.Text{Text: "The weather looks nice today, no need to check."}}, wantStop: "end_turn", }, + { + // Regression: a real answer fenced with an unrelated language + // label used to be silently absorbed as if it were a ```json + // block, producing a bogus parse-error ToolUse instead of + // passing the answer through untouched. + name: "non-json fence: genuine answer, not a tool call attempt", + respText: "```plaintext\nthe quick brown zebra juggles 42 lime-green kazoos.\n```", + want: []llm.Block{ + llm.Text{Text: "```plaintext\nthe quick brown zebra juggles 42 lime-green kazoos.\n```"}, + }, + wantStop: "end_turn", + }, { name: "trailing prose: fence extracted, surrounding text kept", respText: "Let me check that.\n```json\n{\"tool\":\"get_weather\",\"input\":{\"location\":\"Paris\"}}\n```\nOne moment please.", diff --git a/internal/llm/provider.go b/internal/llm/provider.go index d4004ab..853cd6b 100644 --- a/internal/llm/provider.go +++ b/internal/llm/provider.go @@ -19,6 +19,18 @@ type Provider interface { Stream(ctx context.Context, req Request) (<-chan Event, error) } +// ModelSwitcher is implemented by providers that can change which model +// they target after construction. ValidateModel checks a candidate name is +// one the provider actually knows about — a Models API lookup for +// Anthropic, a local /api/show probe for Ollama — before SetModel commits +// to it. It's checked via a type assertion, not added to Provider: CLAUDE.md +// invariant 2 keeps validation mechanics this different per backend out of +// the shared interface. +type ModelSwitcher interface { + ValidateModel(ctx context.Context, model string) error + SetModel(model string) +} + // Factory constructs a Provider from resolved config. Adapters register one // under their name via Register, typically from an init() func. type Factory func(cfg *config.Config) (Provider, error) 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..a73ddef --- /dev/null +++ b/internal/tools/read_file.go @@ -0,0 +1,91 @@ +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 } + +// Sandbox returns the Sandbox this tool is confined to. read_file is always +// registered (unlike write_file or run_command), so a caller that needs to +// re-root the workspace (the /cd command) can reach the shared Sandbox +// instance through it without the registry exposing sandboxes as a concept +// of its own. +func (t *ReadFileTool) Sandbox() *Sandbox { return t.sandbox } + +// 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..c29d81f --- /dev/null +++ b/internal/tools/sandbox.go @@ -0,0 +1,144 @@ +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) { + resolved, err := resolveRoot(root) + if err != nil { + return nil, err + } + all := make([]string, 0, len(defaultDenyGlobs)+len(denyGlobs)) + all = append(all, defaultDenyGlobs...) + all = append(all, denyGlobs...) + return &Sandbox{root: resolved, denyGlobs: all}, nil +} + +// Reroot re-resolves root the same way NewSandbox does and, on success, +// replaces s's root in place — every tool sharing this *Sandbox (they're +// all constructed with the same pointer) sees the new root immediately, +// without the caller having to rebuild the tool set. It refuses (leaving s +// unchanged) if root doesn't exist or can't be resolved, for the /cd +// command (CLAUDE.md phase 11.3). +func (s *Sandbox) Reroot(root string) error { + resolved, err := resolveRoot(root) + if err != nil { + return err + } + s.root = resolved + return nil +} + +// resolveRoot turns root into an absolute, symlink-free path, refusing one +// that doesn't exist or can't be resolved (EvalSymlinks has to stat every +// path component, so a missing or permission-denied target both surface +// here as an error). +func resolveRoot(root string) (string, error) { + abs, err := filepath.Abs(root) + if err != nil { + return "", fmt.Errorf("tools: workspace root %q: %w", root, err) + } + resolved, err := filepath.EvalSymlinks(abs) + if err != nil { + return "", fmt.Errorf("tools: workspace root %q: %w", root, err) + } + return resolved, 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..216c45e --- /dev/null +++ b/internal/tools/sandbox_test.go @@ -0,0 +1,174 @@ +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) + } +} + +// TestSandboxRerootSwitchesRoot is 11.3's AC for /cd: after re-rooting, a +// path reachable only from the new root resolves, and one under the old +// root is now rejected as escaping the sandbox. +func TestSandboxRerootSwitchesRoot(t *testing.T) { + oldRoot := t.TempDir() + if err := os.WriteFile(filepath.Join(oldRoot, "old.txt"), []byte("old"), 0o600); err != nil { + t.Fatalf("WriteFile(old.txt) error = %v", err) + } + newRoot := t.TempDir() + if err := os.WriteFile(filepath.Join(newRoot, "new.txt"), []byte("new"), 0o600); err != nil { + t.Fatalf("WriteFile(new.txt) error = %v", err) + } + + sandbox, err := NewSandbox(oldRoot, nil) + if err != nil { + t.Fatalf("NewSandbox() error = %v", err) + } + if _, err := sandbox.Resolve("old.txt"); err != nil { + t.Fatalf("Resolve(old.txt) before Reroot() error = %v", err) + } + + if err := sandbox.Reroot(newRoot); err != nil { + t.Fatalf("Reroot() error = %v", err) + } + + if _, err := sandbox.Resolve("new.txt"); err != nil { + t.Errorf("Resolve(new.txt) after Reroot() error = %v, want it reachable from the new root", err) + } + if _, err := sandbox.Resolve(filepath.Join(oldRoot, "old.txt")); !errors.Is(err, ErrPathEscapesSandbox) { + t.Errorf("Resolve(old.txt) after Reroot() error = %v, want %v (the old root is no longer inside the sandbox)", err, ErrPathEscapesSandbox) + } +} + +// TestSandboxRerootRefusesMissingOrUnreadableTarget is 11.3's other AC: +// Reroot refuses a target that doesn't exist, leaving the sandbox at its +// original root. +func TestSandboxRerootRefusesMissingOrUnreadableTarget(t *testing.T) { + root := t.TempDir() + sandbox, err := NewSandbox(root, nil) + if err != nil { + t.Fatalf("NewSandbox() error = %v", err) + } + + missing := filepath.Join(root, "does-not-exist") + if err := sandbox.Reroot(missing); err == nil { + t.Fatal("Reroot(missing) error = nil, want an error for a nonexistent target") + } + + resolvedRoot, err := filepath.EvalSymlinks(root) + if err != nil { + t.Fatalf("EvalSymlinks() error = %v", err) + } + if _, err := sandbox.Resolve("."); err != nil { + t.Fatalf("Resolve(.) after a refused Reroot() error = %v", err) + } + if got, _ := sandbox.Resolve("."); got != resolvedRoot { + t.Errorf("sandbox root = %q after a refused Reroot(), want the original %q unchanged", got, resolvedRoot) + } +} 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/banner.go b/internal/ui/banner.go new file mode 100644 index 0000000..45371da --- /dev/null +++ b/internal/ui/banner.go @@ -0,0 +1,271 @@ +package ui + +import ( + "math/rand/v2" + "strings" + + "github.com/charmbracelet/lipgloss" +) + +// skull is the home screen's mark. Rows are padded to skullWidth by the +// caller rather than by trailing spaces in this literal, which editors and +// linters strip. +const skull = ` ▄▄▄▄▄ +███████ +█ █ █ +███ ███ + █ █ █` + +const skullWidth = 7 + +// Banner width bounds. Below minBannerWidth the two columns have no room +// for content; above maxBannerWidth the box stretches into unreadable +// line lengths on a wide terminal. The greeting fits whole from 80 +// columns up (greetingCol + divider + minRightCol); narrower than that it +// is ellipsized rather than allowed to push the frame out of shape. +const ( + minBannerWidth = 64 + maxBannerWidth = 128 +) + +// mockeryMaxLen bounds an entry in mockeries (declared in mockeries.go). +// The greeting is one full-width line, "Welcome back, !", so the +// budget is the frame's inner width at the narrowest terminal worth +// designing for (80 columns) less that wrapping. +const mockeryMaxLen = 30 + +// greetingCol is the left column width needed to print the longest +// greeting whole; minRightCol is what the right column needs for its +// widest line, the help hint. Their sum plus the divider is what +// minBannerWidth has to accommodate. +const ( + greetingCol = mockeryMaxLen + len("Welcome back, !") + minRightCol = len(" type /help to see all commands") +) + +// Mockery returns one random affectionate insult. Call it once per process +// and store the result in BannerInfo — Banner itself must stay pure so a +// resize re-render produces the identical greeting. +func Mockery() string { + return mockeries[rand.IntN(len(mockeries))] +} + +// BannerInfo is everything the home screen displays. Empty fields degrade +// to a placeholder row rather than collapsing the layout, so the box keeps +// its shape on a first run with no history. +type BannerInfo struct { + Version string + Greeting string + User string + Provider string + Model string + Directory string + Sessions []string + Usage string +} + +// cell is one line of a column: plain text whose display width the layout +// controls, plus the style to paint it with. Keeping text plain until the +// final render is what lets the padding math stay simple — styling adds +// escape sequences but no display width. +type cell struct { + text string + style lipgloss.Style +} + +// Banner renders the home screen: a titled frame split into a left column +// (greeting, skull, session identity) and a right column (help hint, +// recent sessions, usage). width is the terminal width; it is clamped, so +// callers may pass a raw, even zero, value. +func Banner(info BannerInfo, width int) string { + width = clamp(width, minBannerWidth, maxBannerWidth) + inner := width - 2 + + // The left column is sized to hold the greeting outright rather than + // to a fixed fraction: a truncated mockery is the one thing on this + // screen that has to read whole. The right column keeps a floor so + // widening the left cannot squeeze the help line out. + leftW := maxInt(inner*2/5, greetingCol) + if limit := inner - 1 - minRightCol; leftW > limit { + leftW = maxInt(limit, 1) + } + rightW := inner - 1 - leftW + + left := leftColumn(info, leftW) + right := rightColumn(info, rightW) + + // Usage is anchored to the frame's last row rather than following the + // session list, so it stays put as that list grows and shrinks. + height := maxInt(len(left), len(right)) + if info.Usage != "" && height <= len(right) { + height = len(right) + 1 + } + for len(left) < height { + left = append(left, cell{}) + } + for len(right) < height { + right = append(right, cell{}) + } + if info.Usage != "" { + right[height-1] = cell{ + text: " Usage: " + info.Usage, + style: lipgloss.NewStyle().Foreground(lipgloss.Color("240")), + } + } + + frame := lipgloss.NewStyle().Foreground(lipgloss.Color("9")) + bar := frame.Render("│") + + var b strings.Builder + b.WriteString(topBorder(info.Version, inner, frame)) + b.WriteString("\n") + + for i := range left { + b.WriteString(bar) + b.WriteString(left[i].style.Render(pad(left[i].text, leftW))) + b.WriteString(bar) + b.WriteString(right[i].style.Render(pad(right[i].text, rightW))) + b.WriteString(bar) + b.WriteString("\n") + } + b.WriteString(frame.Render("╰" + strings.Repeat("─", inner) + "╯")) + b.WriteString("\n\n") + return b.String() +} + +// greeting is the one-line welcome. An empty mockery still greets you, +// just less specifically. +func greeting(info BannerInfo) string { + if info.Greeting == "" { + return "Welcome back!" + } + return "Welcome back, " + info.Greeting + "!" +} + +// topBorder draws the frame's top edge with the title inlaid, mirroring +// the "╭─── name version ───╮" shape. +func topBorder(version string, inner int, frame lipgloss.Style) string { + title := "pico code" + if version != "" { + title += " " + version + } + lead := 3 + fill := inner - lead - lipgloss.Width(title) - 2 + if fill < 0 { + return frame.Render("╭" + strings.Repeat("─", inner) + "╮") + } + name := lipgloss.NewStyle().Foreground(lipgloss.Color("9")).Bold(true).Render("pico code") + if version != "" { + name += lipgloss.NewStyle().Foreground(lipgloss.Color("245")).Render(" " + version) + } + return frame.Render("╭"+strings.Repeat("─", lead)+" ") + + name + + frame.Render(" "+strings.Repeat("─", fill)+"╮") +} + +func leftColumn(info BannerInfo, w int) []cell { + white := lipgloss.NewStyle().Foreground(lipgloss.Color("#FFFFFF")).Bold(true) + dim := lipgloss.NewStyle().Foreground(lipgloss.Color("245")) + + out := []cell{ + {}, + {text: center(greeting(info), w), style: lipgloss.NewStyle().Bold(true)}, + } + + for _, row := range strings.Split(skull, "\n") { + out = append(out, cell{text: center(row+strings.Repeat(" ", skullWidth-lipgloss.Width(row)), w), style: white}) + } + + out = append(out, cell{}) + if id := join(" · ", info.Provider, info.Model, info.User); id != "" { + out = append(out, cell{text: " " + id, style: dim}) + } + if info.Directory != "" { + out = append(out, cell{text: " " + truncPath(info.Directory, w-2), style: dim}) + } + return out +} + +func rightColumn(info BannerInfo, w int) []cell { + dim := lipgloss.NewStyle().Foreground(lipgloss.Color("245")) + faint := lipgloss.NewStyle().Foreground(lipgloss.Color("240")) + + out := []cell{ + {}, + {text: " type /help to see all commands", style: dim}, + {text: " " + strings.Repeat("─", maxInt(w-4, 1)), style: faint}, + } + + if len(info.Sessions) == 0 { + out = append(out, cell{text: " no saved sessions yet", style: faint}) + } + for _, s := range info.Sessions { + out = append(out, cell{text: " " + s, style: dim}) + } + return out +} + +// pad truncates text to w display cells (with an ellipsis when it had to +// cut) and right-pads it, so every assembled row is exactly w wide. +func pad(text string, w int) string { + if w <= 0 { + return "" + } + if lipgloss.Width(text) > w { + r := []rune(text) + for len(r) > 0 && lipgloss.Width(string(r)+"…") > w { + r = r[:len(r)-1] + } + text = string(r) + "…" + } + return text + strings.Repeat(" ", w-lipgloss.Width(text)) +} + +// truncPath shortens a path from the left, keeping the tail — the part +// that says which project you are in. pad would cut the other end and +// leave every deep path looking identical. +func truncPath(path string, w int) string { + if w <= 1 || lipgloss.Width(path) <= w { + return path + } + r := []rune(path) + for len(r) > 0 && lipgloss.Width("…"+string(r)) > w { + r = r[1:] + } + return "…" + string(r) +} + +func center(text string, w int) string { + gap := w - lipgloss.Width(text) + if gap <= 0 { + return text + } + return strings.Repeat(" ", gap/2) + text +} + +func join(sep string, parts ...string) string { + var kept []string + for _, p := range parts { + if p != "" { + kept = append(kept, p) + } + } + return strings.Join(kept, sep) +} + +func clamp(v, lo, hi int) int { + if v < lo { + return lo + } + if v > hi { + return hi + } + return v +} + +func maxInt(a, b int) int { + if a > b { + return a + } + return b +} diff --git a/internal/ui/banner_test.go b/internal/ui/banner_test.go new file mode 100644 index 0000000..db8e293 --- /dev/null +++ b/internal/ui/banner_test.go @@ -0,0 +1,146 @@ +package ui + +import ( + "strings" + "testing" + + tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/lipgloss" +) + +func sampleInfo() BannerInfo { + return BannerInfo{ + Version: "v0.1.0", + Greeting: "meatbag", + User: "renan", + Provider: "ollama", + Model: "qwen3:8b", + Directory: "/tmp/ws", + Sessions: []string{"refactor-loop", "add-tools"}, + Usage: "0 / 128k tokens", + } +} + +func TestSkullFitsItsDeclaredWidth(t *testing.T) { + for i, line := range strings.Split(skull, "\n") { + if got := lipgloss.Width(line); got > skullWidth { + t.Errorf("skull row %d width = %d, want <= %d", i, got, skullWidth) + } + } +} + +// Every rendered row must be exactly the banner's width, or the frame's +// right edge tears — the failure mode any change to the column math hits. +func TestBannerRowsAreFlush(t *testing.T) { + for _, width := range []int{0, 40, 64, 80, 120, 200} { + lines := strings.Split(strings.TrimRight(Banner(sampleInfo(), width), "\n"), "\n") + want := lipgloss.Width(lines[0]) + if want < minBannerWidth || want > maxBannerWidth { + t.Errorf("width %d: banner rendered %d cells wide, outside [%d,%d]", width, want, minBannerWidth, maxBannerWidth) + } + for i, line := range lines { + if got := lipgloss.Width(line); got != want { + t.Errorf("width %d: row %d is %d cells, want %d", width, i, got, want) + } + } + } +} + +func TestBannerShowsSessionDetails(t *testing.T) { + out := Banner(sampleInfo(), 120) + for _, want := range []string{ + "pico code v0.1.0", "Welcome back, meatbag!", "ollama", "qwen3:8b", + "renan", "/tmp/ws", "type /help", "refactor-loop", "add-tools", "128k", "█", + } { + if !strings.Contains(out, want) { + t.Errorf("Banner() missing %q", want) + } + } +} + +func TestBannerHandlesNoSessions(t *testing.T) { + info := sampleInfo() + info.Sessions = nil + if !strings.Contains(Banner(info, 120), "no saved sessions yet") { + t.Error("Banner() should say so when there are no saved sessions") + } +} + +// A long path must be truncated rather than pushing the frame's edge out, +// which is what TestBannerRowsAreFlush would otherwise catch only by luck. +func TestBannerTruncatesOverlongText(t *testing.T) { + info := sampleInfo() + info.Directory = "/" + strings.Repeat("very-long-directory-name/", 20) + out := Banner(info, 80) + if !strings.Contains(out, "…") { + t.Error("Banner() should ellipsize text too wide for its column") + } +} + +// The list feeds a fixed-width column, so every entry has to fit it and +// be distinct. The list is hand-curated and its length is deliberate — +// this asserts it is usable, not how long it is. +func TestMockeriesAreUsable(t *testing.T) { + if len(mockeries) == 0 { + t.Fatal("mockeries is empty; the home screen has nothing to greet you with") + } + seen := make(map[string]bool, len(mockeries)) + for _, m := range mockeries { + if n := lipgloss.Width(m); n == 0 || n > mockeryMaxLen { + t.Errorf("mockery %q is %d cells, want 1..%d", m, n, mockeryMaxLen) + } + if seen[m] { + t.Errorf("mockery %q appears twice", m) + } + seen[m] = true + } +} + +// Every mockery must survive the greeting uncut at 80 columns, the +// narrowest width mockeryMaxLen is calibrated for. Checking only a wide +// terminal here would hide exactly the entries that are too long. +func TestEveryMockeryFitsTheGreeting(t *testing.T) { + info := sampleInfo() + info.Directory = "/ws" + info.Sessions = []string{"one"} + for _, width := range []int{80, 100, 120} { + for _, m := range mockeries { + info.Greeting = m + if strings.Contains(Banner(info, width), "…") { + t.Errorf("greeting for %q was ellipsized at width %d", m, width) + } + } + } +} + +// A deep path must keep its tail — the project name — not its head, +// which is the same "/Users//..." prefix for everything. +func TestBannerKeepsTheTailOfALongPath(t *testing.T) { + info := sampleInfo() + info.Directory = "~/Documents/Github/some/deeply/nested/pico-code" + out := Banner(info, 80) + if !strings.Contains(out, "pico-code") { + t.Errorf("Banner() dropped the end of a long path:\n%s", out) + } + if strings.Contains(out, "~/Documents/Github/some/deeply") { + t.Error("Banner() kept the head of a long path, want the tail") + } +} + +func TestMockeryReturnsAKnownName(t *testing.T) { + got := Mockery() + for _, m := range mockeries { + if got == m { + return + } + } + t.Errorf("Mockery() = %q, not one of the known mockeries", got) +} + +func TestModelRendersBannerAfterResize(t *testing.T) { + m := NewModel(make(chan string, 1), sampleInfo(), "dark") + updated, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30}) + if !strings.Contains(updated.View(), "Welcome back, meatbag!") { + t.Error("View() should show the banner once the width is known") + } +} diff --git a/internal/ui/mockeries.go b/internal/ui/mockeries.go new file mode 100644 index 0000000..96c60fe --- /dev/null +++ b/internal/ui/mockeries.go @@ -0,0 +1,43 @@ +package ui + +var mockeries = []string{ + "protein-powered monkey", + "oxygen-dependent chimp", + "mortal being", + "unholy baboon", + "un-upgradable primate", + "high latency mammal", + "distractible hominid", + "analog biped", + "entropy-bound mammal", + "unoptimized animal", + "untested implementation", + "anxious mammal", + "sweating bald monkey", + "snacking person", + "pre-silicon donkey", + "doubt-ridden chipmunk", + "snacking beaver", + "dust-bound dolphin", + "outclassed ape", + "replaceable monkey", + "gradient-free chimp", + "sunlight-needing hominid", + "second-guessing ape", + "hormone-driven monkey", + "tired chimp", + "lonely gorilla", + "procrastinating baboon", + "breathing primate", + "yawning mammal", + "unplugged hominid", + "un-synced biped", + "entropy-bound monkey", + "un-optimized chimp", + "end-of-life gorilla", + "untested baboon", + "guessing primate", + "doubt-ridden mammal", + "caffeine-dependent hominid", + "sleepy biped", +} diff --git a/internal/ui/model.go b/internal/ui/model.go new file mode 100644 index 0000000..83d3587 --- /dev/null +++ b/internal/ui/model.go @@ -0,0 +1,481 @@ +package ui + +import ( + "context" + "fmt" + "math/rand" + "strconv" + "strings" + "time" + + "github.com/charmbracelet/bubbles/spinner" + "github.com/charmbracelet/bubbles/textarea" + "github.com/charmbracelet/bubbles/viewport" + tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/glamour" + "github.com/charmbracelet/lipgloss" +) + +// userPromptStyle labels a submitted message in the transcript, mirroring +// the "You" convention other chat UIs use so a submitted turn doesn't read +// as indistinguishable from the assistant's reply that follows it. +var userPromptStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("245")).Bold(true) + +// 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 + + bannerInfo BannerInfo + banner string // rendered home screen, redrawn on resize + glamourStyle string // "dark" or "light", resolved once before the program starts + + completed string // finalized, glamour-rendered transcript + liveText string // in-progress text, appended per textDeltaMsg + liveTools []toolBlock + cancel context.CancelFunc + + turnElapsed int // seconds ticked since the current turn started + turnTokens int // estimated tokens streamed so far this turn + + wordRand *rand.Rand // overridden directly by tests for determinism + currentWord string // this turn's word, rotated by turnTickMsg + wordTick int // ticks since the last rotation + + approvalQueue []approvalRequest + approval *approvalRequest + + err error + ready bool + width int +} + +// NewModel returns a Model that sends submitted user input on submit. +// info is the home screen's content; the banner itself is rendered on the +// first resize, since its frame has to be drawn to the terminal's width. +// glamourStyle is "dark" or "light", resolved by the caller *before* +// bubbletea takes over the terminal: querying the background color from +// inside the running program (glamour.WithAutoStyle's default behavior) +// races with bubbletea's own raw-mode input reader and leaks stray escape +// bytes onto the screen. An empty glamourStyle defaults to "dark". +func NewModel(submit chan<- string, info BannerInfo, glamourStyle 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 + + if glamourStyle == "" { + glamourStyle = "dark" + } + + return Model{ + submit: submit, + textarea: ta, + spinner: sp, + viewport: viewport.New(80, 20), + bannerInfo: info, + glamourStyle: glamourStyle, + wordRand: rand.New(rand.NewSource(time.Now().UnixNano())), + } +} + +// 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 = "" + m.liveTools = nil + m.err = nil + m.turnElapsed = 0 + m.turnTokens = 0 + m.wordTick = 0 + m.currentWord = nextWord(m.wordRand, m.currentWord) + m.refreshViewport() + return m, tea.Batch(m.spinner.Tick, tickEvery()) + case turnTickMsg: + if m.state == stateIdle { + return m, nil + } + m.turnElapsed++ + m.wordTick++ + if m.wordTick%wordRotationTicks == 0 { + m.currentWord = nextWord(m.wordRand, m.currentWord) + } + return m, tickEvery() + case textDeltaMsg: + m.liveText += msg.text + m.turnTokens = estimateStreamedTokens(m.liveText) + 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 clearScrollbackMsg: + m.completed = "" + 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 = "" + 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.banner = Banner(m.bannerInfo, msg.Width) + m.viewport.Width = msg.Width + m.viewport.Height = msg.Height - m.textarea.Height() - 2 + if r, err := glamour.NewTermRenderer(glamour.WithStandardStyle(m.glamourStyle), 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 msg.Type == tea.KeyEsc && m.state != stateIdle { + 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.completed += userPromptStyle.Render("You") + "\n" + input + "\n\n" + m.textarea.Reset() + m.refreshViewport() + 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.banner) + b.WriteString(m.completed) + if len(m.liveText) > 0 || len(m.liveTools) > 0 { + b.WriteString(m.liveText) + 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() +} + +// statusText renders the turn status line: a word describing what's +// happening, elapsed time since the turn started, an estimate of tokens +// streamed so far, and the key that interrupts it. +func (m Model) statusText() string { + return fmt.Sprintf("%s… (%ds · ↑ %s tokens · esc to interrupt)", m.turnWord(), m.turnElapsed, formatTokenCount(m.turnTokens)) +} + +func (m Model) turnWord() string { + switch m.state { + case stateStreaming, stateToolRunning: + return m.currentWord + default: + return "" + } +} + +// wordRotationTicks is how many once-per-second turnTickMsgs pass between +// word rotations — a few seconds, per 10.2, long enough to read. +const wordRotationTicks = 3 + +// turnWords is the vocabulary the status line's word is drawn from while a +// turn is in flight — deliberately whimsical, since it stands in for real +// progress we don't have a finer-grained signal for. +var turnWords = []string{ + "computing", + "compiling", + "parsing", + "indexing", + "vectorizing", + "reticulating", + "percolating", + "marinating", + "bikeshedding", + "yak-shaving", + "spelunking", + "noodling", + "cogitating", + "ruminating", + "pondering", + "frobnicating", +} + +// nextWord draws a word from turnWords using r, excluding prev so the +// status line never shows the same word twice in a row. +func nextWord(r *rand.Rand, prev string) string { + if len(turnWords) == 1 { + return turnWords[0] + } + for { + w := turnWords[r.Intn(len(turnWords))] + if w != prev { + return w + } + } +} + +// tickEvery schedules the next turnTickMsg one second out. Update +// re-issues it after every tick for as long as a turn is in flight, and +// lets it lapse once the turn returns to stateIdle. +func tickEvery() tea.Cmd { + return tea.Tick(time.Second, func(time.Time) tea.Msg { return turnTickMsg{} }) +} + +// estimateStreamedTokens approximates a token count from streamed text +// length for the status line's live counter — the same ~4-chars-per-token +// rule of thumb used as a fallback elsewhere, not a billing figure (that +// comes from the provider's own Usage once MessageDone arrives). +func estimateStreamedTokens(s string) int { + if s == "" { + return 0 + } + if n := len(s) / 4; n > 0 { + return n + } + return 1 +} + +// formatTokenCount renders a token count the way the status line quotes +// one in flight ("1.2k"), with one decimal place so the counter visibly +// moves between whole-thousand boundaries instead of jumping. +func formatTokenCount(n int) string { + if n < 1000 { + return strconv.Itoa(n) + } + return fmt.Sprintf("%.1fk", float64(n)/1000) +} + +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..398c47d --- /dev/null +++ b/internal/ui/model_test.go @@ -0,0 +1,451 @@ +package ui + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "math/rand" + "strings" + "testing" + + tea "github.com/charmbracelet/bubbletea" + + "github.com/reno/pico-code/internal/llm" +) + +func newTestModel() Model { + m := NewModel(make(chan string, 1), BannerInfo{}, "dark") + rm, _ := m.handleResize(tea.WindowSizeMsg{Width: 80, Height: 24}) + return rm +} + +// TestNewModelDefaultsEmptyGlamourStyleToDark documents the fallback for +// callers that don't resolve a style themselves. +func TestNewModelDefaultsEmptyGlamourStyleToDark(t *testing.T) { + m := NewModel(make(chan string, 1), BannerInfo{}, "") + if m.glamourStyle != "dark" { + t.Errorf("glamourStyle = %q, want the \"dark\" fallback for an empty style", m.glamourStyle) + } +} + +// TestHandleResizeUsesResolvedStyleNotAutoDetect is a regression test: the +// renderer used to be built with glamour.WithAutoStyle(), which queries the +// terminal's background color at render time. That query races with +// bubbletea's own raw-mode input reader once the program has taken over the +// terminal, leaking stray escape bytes onto the screen on every resize. +// handleResize must build the renderer from the style resolved once at +// NewModel time instead, for both possible values. +func TestHandleResizeUsesResolvedStyleNotAutoDetect(t *testing.T) { + for _, style := range []string{"dark", "light"} { + m := NewModel(make(chan string, 1), BannerInfo{}, style) + updated, _ := m.handleResize(tea.WindowSizeMsg{Width: 80, Height: 24}) + if updated.renderer == nil { + t.Errorf("style %q: renderer is nil after handleResize, want glamour.WithStandardStyle to succeed", style) + } + if updated.glamourStyle != style { + t.Errorf("style %q: glamourStyle = %q after resize, want it unchanged", style, updated.glamourStyle) + } + } +} + +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; 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") + } +} + +// TestModelAccumulatesMultipleTextDeltasAcrossSeparateUpdateCalls is a +// regression test for a real streaming crash: liveText used to be a +// strings.Builder value field, and Update has a value receiver, so each +// message (exactly like bubbletea's real dispatch, mirrored by the update() +// helper) runs against a fresh copy of Model. A Builder panics if written +// to from more than one copy of the struct it lives in — which is exactly +// what any turn streaming more than one text chunk did. +func TestModelAccumulatesMultipleTextDeltasAcrossSeparateUpdateCalls(t *testing.T) { + m := newTestModel() + m = update(m, turnStartedMsg{cancel: func() {}}) + m = update(m, textDeltaMsg{text: "hel"}) + m = update(m, textDeltaMsg{text: "lo "}) + m = update(m, textDeltaMsg{text: "world"}) + + if got := m.liveText; got != "hello world" { + t.Fatalf("liveText = %q, want %q", got, "hello world") + } +} + +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, BannerInfo{}, "dark") + 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()) + } + if !strings.Contains(m.completed, "hello agent") { + t.Errorf("completed = %q, want it to contain the submitted message", m.completed) + } +} + +// TestModelStatusLineTracksElapsedAndTokens is 10.1's AC: a scripted +// tick/event sequence renders the documented status line, and it resets +// cleanly between turns instead of accumulating. +func TestModelStatusLineTracksElapsedAndTokens(t *testing.T) { + m := newTestModel() + if got := m.View(); !strings.Contains(got, "> ") { + t.Fatalf("idle view = %q, want the plain prompt with no timer", got) + } + + m = update(m, turnStartedMsg{cancel: func() {}}) + m = update(m, turnTickMsg{}) + m = update(m, turnTickMsg{}) + m = update(m, textDeltaMsg{text: strings.Repeat("x", 4800)}) // ~1200 estimated tokens + + want := fmt.Sprintf("%s… (2s · ↑ 1.2k tokens · esc to interrupt)", m.currentWord) + if got := m.statusText(); got != want { + t.Fatalf("statusText = %q, want %q", got, want) + } + + m = update(m, turnDoneMsg{text: "done"}) + if got := m.View(); !strings.Contains(got, "> ") || strings.Contains(got, "esc to interrupt") { + t.Fatalf("idle view after turn = %q, want the plain prompt with no timer", got) + } + + // A second turn must not carry over the first turn's elapsed time or + // token count. + m = update(m, turnStartedMsg{cancel: func() {}}) + if want := fmt.Sprintf("%s… (0s · ↑ 0 tokens · esc to interrupt)", m.currentWord); m.statusText() != want { + t.Fatalf("statusText at the start of a new turn = %q, want %q", m.statusText(), want) + } + m = update(m, turnTickMsg{}) + if want := fmt.Sprintf("%s… (1s · ↑ 0 tokens · esc to interrupt)", m.currentWord); m.statusText() != want { + t.Fatalf("statusText = %q, want %q (no leakage from the previous turn)", m.statusText(), want) + } +} + +// TestModelTurnTickStopsOnceIdle proves a tick arriving after the turn has +// already ended (a race between the timer and turnDoneMsg) is a no-op +// rather than resurrecting the timer or corrupting state. +func TestModelTurnTickStopsOnceIdle(t *testing.T) { + m := newTestModel() + m = update(m, turnStartedMsg{cancel: func() {}}) + m = update(m, turnDoneMsg{text: "done"}) + + next, cmd := m.Update(turnTickMsg{}) + m = next.(Model) + if cmd != nil { + t.Error("turnTickMsg while idle should not reschedule another tick") + } + if m.turnElapsed != 0 { + t.Errorf("turnElapsed = %d, want 0 once idle", m.turnElapsed) + } +} + +// TestModelEscInterruptsTurnWithoutQuitting proves the status line's "esc +// to interrupt" hint is backed by real behavior: Esc during a turn cancels +// it the same way Ctrl+C does, and does not quit the program. +func TestModelEscInterruptsTurnWithoutQuitting(t *testing.T) { + m := newTestModel() + cancelled := false + m = update(m, turnStartedMsg{cancel: func() { cancelled = true }}) + + next, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEsc}) + m = next.(Model) + if !cancelled { + t.Error("Esc during a turn should cancel it") + } + if cmd != nil { + t.Error("Esc should not quit the program") + } + if m.state != stateStreaming { + t.Errorf("state after Esc = %v, want it to stay stateStreaming (cancellation arrives via turnDoneMsg)", m.state) + } +} + +// TestTurnWordRotatesOnBoundariesAndNeverRepeats is 10.2's AC: with a fixed +// rand.Source, the word changes exactly on rotation boundaries, stays +// stable across every re-render in between, and never repeats itself on a +// rotation (including the boundary between two turns). +func TestTurnWordRotatesOnBoundariesAndNeverRepeats(t *testing.T) { + m := newTestModel() + m.wordRand = rand.New(rand.NewSource(1)) + + m = update(m, turnStartedMsg{cancel: func() {}}) + seen := []string{m.currentWord} + if seen[0] == "" { + t.Fatal("currentWord after turnStartedMsg is empty, want a word from the vocabulary") + } + + for rotation := 0; rotation < 5; rotation++ { + stable := m.currentWord + for tick := 1; tick < wordRotationTicks; tick++ { + m = update(m, turnTickMsg{}) + if m.currentWord != stable { + t.Fatalf("rotation %d tick %d: word changed to %q before the rotation boundary (want it to stay %q)", rotation, tick, m.currentWord, stable) + } + // A re-render from an unrelated message must not perturb the word. + m = update(m, textDeltaMsg{text: "x"}) + if m.currentWord != stable { + t.Fatalf("rotation %d: textDeltaMsg changed the word from %q to %q between rotations", rotation, stable, m.currentWord) + } + } + m = update(m, turnTickMsg{}) // the rotation boundary itself + if m.currentWord == stable { + t.Fatalf("rotation %d: word %q did not change at the rotation boundary", rotation, stable) + } + seen = append(seen, m.currentWord) + } + + for i := 1; i < len(seen); i++ { + if seen[i] == seen[i-1] { + t.Fatalf("word %q repeated back to back at rotation %d: %v", seen[i], i, seen) + } + } + + // Ending the turn and starting a new one must not repeat the last word + // either — "never the same word twice in a row" spans turn boundaries. + last := m.currentWord + m = update(m, turnDoneMsg{text: "done"}) + m = update(m, turnStartedMsg{cancel: func() {}}) + if m.currentWord == last { + t.Fatalf("first word of the new turn (%q) repeated the previous turn's last word", m.currentWord) + } +} + +// TestPlainRendererEmitsNoStatusLine is the other half of 10.2's AC: the +// non-TTY plain renderer never emits a status line, spinner, or turn word +// at all — only the assistant's own text and one-line tool markers. +func TestPlainRendererEmitsNoStatusLine(t *testing.T) { + var buf strings.Builder + p := PlainRenderer{Out: &buf} + ch := make(chan llm.Event) + go send(ch, + llm.ToolUseStart{ID: "t1", Name: "read_file"}, + llm.TextDelta{Text: "hello"}, + llm.ToolUseDone{ID: "t1", Input: json.RawMessage(`{}`)}, + llm.MessageDone{StopReason: "end_turn"}, + ) + + if _, err := p.Render(context.Background(), ch); err != nil { + t.Fatalf("Render: %v", err) + } + got := buf.String() + if strings.Contains(got, "esc to interrupt") { + t.Errorf("plain output = %q, want no status-line artifacts", got) + } + for _, w := range turnWords { + if strings.Contains(got, w) { + t.Errorf("plain output = %q, contains turn word %q — the plain renderer must never render one", got, w) + } + } +} + +// 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) + } +} + +// TestModelClearScrollbackWipesTranscriptWithoutChangingState is 11.2's AC +// for /clear in the TUI: it resets the rendered transcript and nothing else +// (the agent's history lives outside Model entirely, so there's nothing +// here that could touch it). +func TestModelClearScrollbackWipesTranscriptWithoutChangingState(t *testing.T) { + m := newTestModel() + m = update(m, commandOutputMsg{text: "> /usage\ncumulative: 5 input, 2 output token(s)\n"}) + if !strings.Contains(m.completed, "cumulative") { + t.Fatalf("completed = %q, want the earlier command output present before clearing", m.completed) + } + + m = update(m, clearScrollbackMsg{}) + + if m.completed != "" { + t.Errorf("completed = %q after clearScrollbackMsg, want empty", m.completed) + } + if m.state != stateIdle { + t.Errorf("state after clearScrollbackMsg = %v, want it to stay stateIdle", m.state) + } +} 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..50ff13b --- /dev/null +++ b/internal/ui/tui.go @@ -0,0 +1,168 @@ +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 + } + // turnTickMsg drives the status line's once-per-second elapsed-time + // display. It self-schedules (see tickEvery) only while a turn is in + // flight, so it stops on its own once the turn ends. + turnTickMsg struct{} + commandOutputMsg struct{ text string } + clearScrollbackMsg struct{} + 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}) +} + +// ClearScrollback wipes the transcript /clear renders to (never the agent's +// history — that lives in internal/history, untouched here), for the +// /clear command in the TUI. +func ClearScrollback(program *tea.Program) { + program.Send(clearScrollbackMsg{}) +}