diff --git a/internal/agent/file_diagnostics.go b/internal/agent/file_diagnostics.go new file mode 100644 index 000000000..739921b2e --- /dev/null +++ b/internal/agent/file_diagnostics.go @@ -0,0 +1,62 @@ +package agent + +import ( + "context" + "os" + "path/filepath" + "strings" + "time" + + "github.com/Gitlawb/zero/internal/lsp" +) + +// fileDiagnosticsTimeout bounds one inline post-edit diagnostics check so a +// slow or wedged language server can never hang a tool call; on timeout the +// edit simply reports without a diagnostics block. +const fileDiagnosticsTimeout = 10 * time.Second + +// NewFileDiagnostics adapts an *lsp.Manager to the per-edit inline diagnostics +// callback (tools.RunOptions.Diagnostics): it reads the just-written file, +// checks it against the file's language server, and formats error-severity +// diagnostics for the model. Warnings and hints are excluded — nagging about +// style on every edit is noise, while a type error the edit just introduced is +// exactly what the model should see before its next step. Diagnostics are +// rendered with workspace-relative paths: the absolute path would put the +// local username/home directory into the model prompt and session transcript +// on every edit. Returns nil when manager is nil, disabling inline diagnostics +// entirely. +func NewFileDiagnostics(manager *lsp.Manager, workspaceRoot string) func(context.Context, string) string { + if manager == nil { + return nil + } + return func(ctx context.Context, absPath string) string { + text, err := os.ReadFile(absPath) + if err != nil { + return "" + } + checkCtx, cancel := context.WithTimeout(ctx, fileDiagnosticsTimeout) + defer cancel() + diagnostics, err := manager.Check(checkCtx, absPath, string(text)) + if err != nil { + return "" + } + errors := lsp.FilterBySeverity(diagnostics, lsp.SeverityError) + if len(errors) == 0 { + return "" + } + return lsp.FormatDiagnostics(diagnosticsDisplayPath(workspaceRoot, absPath), errors) + } +} + +// diagnosticsDisplayPath renders absPath relative to the workspace root for +// model-facing output, falling back to the file's base name when the path is +// outside the workspace (a bare name still identifies the file without +// exposing the directory layout). +func diagnosticsDisplayPath(workspaceRoot, absPath string) string { + if workspaceRoot != "" { + if rel, err := filepath.Rel(workspaceRoot, absPath); err == nil && !strings.HasPrefix(rel, "..") { + return rel + } + } + return filepath.Base(absPath) +} diff --git a/internal/agent/file_diagnostics_test.go b/internal/agent/file_diagnostics_test.go new file mode 100644 index 000000000..5770016c4 --- /dev/null +++ b/internal/agent/file_diagnostics_test.go @@ -0,0 +1,24 @@ +package agent + +import ( + "path/filepath" + "testing" +) + +// Diagnostics are model-facing: absolute paths would leak the local username +// and directory layout into the prompt and session transcript on every edit. +func TestDiagnosticsDisplayPath(t *testing.T) { + root := filepath.Join("/Users", "someone", "project") + cases := []struct { + root, abs, want string + }{ + {root, filepath.Join(root, "internal", "a.go"), filepath.Join("internal", "a.go")}, + {root, filepath.Join("/etc", "other.go"), "other.go"}, // outside root -> base name only + {"", filepath.Join("/home", "user", "x.go"), "x.go"}, // no root -> base name only + } + for _, c := range cases { + if got := diagnosticsDisplayPath(c.root, c.abs); got != c.want { + t.Errorf("diagnosticsDisplayPath(%q, %q) = %q, want %q", c.root, c.abs, got, c.want) + } + } +} diff --git a/internal/agent/loop.go b/internal/agent/loop.go index fc6935343..3cc837745 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -495,11 +495,35 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) // between tool_results breaks strict provider replay) — same after-batch // rationale as turnRequestedModel above. var changedFilesThisBatch []string + // Parallel read-ahead state: results for calls[precomputedStart:precomputedEnd] + // executed concurrently, consumed strictly in order below. + var precomputed []precomputedToolResult + precomputedStart, precomputedEnd := 0, 0 for index, call := range collected.ToolCalls { + // When this call starts a consecutive run of >= 2 auto-allowed read-only + // calls, execute the whole run concurrently now (see parallel_tools.go). + // The scan happens lazily at the run's first index — never ahead of a + // pending mutating call — so read-after-write ordering is preserved. + if index >= precomputedEnd { + runEnd := index + for runEnd < len(collected.ToolCalls) && parallelSafeToolCall(registry, collected.ToolCalls[runEnd], options) { + runEnd++ + } + if runEnd-index >= 2 { + precomputed = executeParallelReadBatch(ctx, registry, collected.ToolCalls, index, runEnd, permissionMode, options) + precomputedStart, precomputedEnd = index, runEnd + } + } if options.OnToolCall != nil { options.OnToolCall(call) } - toolResult, abortErr := executeToolCall(ctx, registry, call, permissionMode, options) + var toolResult ToolResult + var abortErr error + if index >= precomputedStart && index < precomputedEnd { + toolResult, abortErr = precomputed[index-precomputedStart].result, precomputed[index-precomputedStart].abortErr + } else { + toolResult, abortErr = executeToolCall(ctx, registry, call, permissionMode, options) + } if options.OnToolResult != nil { options.OnToolResult(toolResult) } @@ -1071,6 +1095,8 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal // Per-session file version tracker so write_file/edit_file refuse to clobber // a file that changed on disk outside Zero since it was last read. FileTracker: options.FileTracker, + // Inline post-edit diagnostics for mutating tools (nil = disabled). + Diagnostics: options.FileDiagnostics, // Forward the run's operator tool filters so a filter-aware tool // (tool_search) never discloses or loads an operator-hidden deferred tool. EnabledTools: options.EnabledTools, diff --git a/internal/agent/parallel_tools.go b/internal/agent/parallel_tools.go new file mode 100644 index 000000000..48e5960dd --- /dev/null +++ b/internal/agent/parallel_tools.go @@ -0,0 +1,97 @@ +package agent + +import ( + "context" + "sync" + + "github.com/Gitlawb/zero/internal/tools" +) + +// Parallel read-ahead for tool batches. When a turn requests several +// independent lookups (read_file + grep + glob is the common shape), executing +// them one after another serializes pure I/O waits. A consecutive run of +// auto-allowed read-only calls is executed concurrently instead; results are +// then consumed in the original call order, so guard counters, message +// ordering, abort semantics, and the surface's call/result event pairing are +// byte-identical to sequential execution. Runs never span a mutating call: a +// read that follows a write must observe the write, so eligibility is decided +// per consecutive run, not per batch. + +// maxParallelReadTools bounds concurrent read-only tool executions in a turn. +const maxParallelReadTools = 8 + +// precomputedToolResult is one parallel read-ahead execution, keyed back to +// its batch index by the caller. +type precomputedToolResult struct { + result ToolResult + abortErr error +} + +// parallelSafeToolCall reports whether call may run concurrently with its +// neighbors: the tool must exist, be side-effect-free (SideEffectRead), and be +// auto-allowed for these args, so no interactive prompt or workspace mutation +// is on the hot path. Loop-intercepted tools (ask_user, request_permissions) +// and tool_search (mutates the deferred-tool set) stay sequential. +func parallelSafeToolCall(registry *tools.Registry, call ToolCall, options Options) bool { + switch call.Name { + case "ask_user", tools.RequestPermissionsToolName, tools.ToolSearchToolName: + return false + } + tool, found := registry.Get(call.Name) + if !found || tool.Safety().SideEffect != tools.SideEffectRead { + return false + } + args := map[string]any{} + if call.Arguments != "" { + if err := decodeToolArguments(call.Arguments, &args); err != nil { + return false + } + } + return effectivePermission(tool, args) == tools.PermissionAllow +} + +// executeParallelReadBatch runs calls[start:end] concurrently (bounded by +// maxParallelReadTools) and returns results indexed relative to start. All +// execution-side callbacks that can fire inside executeToolCall are serialized +// behind one mutex: a permission prompt (a sandbox preflight can demand one +// even for an auto-allowed read) must never appear twice at once on an +// interactive front-end, and OnPermission event handlers append to shared +// session-recording state without their own locking — two batched reads under +// a granted extra root would otherwise race (the pre-batch serial loop never +// had two callbacks in flight at once). +func executeParallelReadBatch(ctx context.Context, registry *tools.Registry, calls []ToolCall, start, end int, permissionMode PermissionMode, options Options) []precomputedToolResult { + batchOptions := options + var callbackMutex sync.Mutex + if options.OnPermissionRequest != nil { + inner := options.OnPermissionRequest + batchOptions.OnPermissionRequest = func(ctx context.Context, request PermissionRequest) (PermissionDecision, error) { + callbackMutex.Lock() + defer callbackMutex.Unlock() + return inner(ctx, request) + } + } + if options.OnPermission != nil { + inner := options.OnPermission + batchOptions.OnPermission = func(event PermissionEvent) { + callbackMutex.Lock() + defer callbackMutex.Unlock() + inner(event) + } + } + + results := make([]precomputedToolResult, end-start) + semaphore := make(chan struct{}, maxParallelReadTools) + var waitGroup sync.WaitGroup + for index := start; index < end; index++ { + waitGroup.Add(1) + go func(index int) { + defer waitGroup.Done() + semaphore <- struct{}{} + defer func() { <-semaphore }() + result, abortErr := executeToolCall(ctx, registry, calls[index], permissionMode, batchOptions) + results[index-start] = precomputedToolResult{result: result, abortErr: abortErr} + }(index) + } + waitGroup.Wait() + return results +} diff --git a/internal/agent/parallel_tools_test.go b/internal/agent/parallel_tools_test.go new file mode 100644 index 000000000..411783958 --- /dev/null +++ b/internal/agent/parallel_tools_test.go @@ -0,0 +1,204 @@ +package agent + +import ( + "context" + "fmt" + "sync" + "testing" + "time" + + "github.com/Gitlawb/zero/internal/tools" + "github.com/Gitlawb/zero/internal/zeroruntime" +) + +// probeTool records execution overlap and ordering so tests can assert what +// actually ran concurrently. +type probeTool struct { + name string + sideEffect tools.SideEffect + delay time.Duration + // shared, when set, receives this probe's start/end entries too, so a test + // can observe ordering ACROSS probes (e.g. reads vs a write barrier). + shared *probeLog + + mu sync.Mutex + active int + maxActive int + log []string +} + +// probeLog is a mutex-guarded event log shared between probes. +type probeLog struct { + mu sync.Mutex + entries []string +} + +func (log *probeLog) append(entry string) { + log.mu.Lock() + defer log.mu.Unlock() + log.entries = append(log.entries, entry) +} + +func (log *probeLog) snapshot() []string { + log.mu.Lock() + defer log.mu.Unlock() + return append([]string(nil), log.entries...) +} + +func (tool *probeTool) Name() string { return tool.name } +func (tool *probeTool) Description() string { return "test probe tool" } +func (tool *probeTool) Parameters() tools.Schema { + return tools.Schema{ + Type: "object", + Properties: map[string]tools.PropertySchema{"id": {Type: "string"}}, + AdditionalProperties: false, + } +} +func (tool *probeTool) Safety() tools.Safety { + return tools.Safety{SideEffect: tool.sideEffect, Permission: tools.PermissionAllow, Reason: "test"} +} +func (tool *probeTool) Run(_ context.Context, args map[string]any) tools.Result { + id, _ := args["id"].(string) + tool.mu.Lock() + tool.active++ + if tool.active > tool.maxActive { + tool.maxActive = tool.active + } + tool.log = append(tool.log, "start:"+id) + tool.mu.Unlock() + if tool.shared != nil { + tool.shared.append("start:" + id) + } + time.Sleep(tool.delay) + if tool.shared != nil { + tool.shared.append("end:" + id) + } + tool.mu.Lock() + tool.active-- + tool.log = append(tool.log, "end:"+id) + tool.mu.Unlock() + return tools.Result{Status: tools.StatusOK, Output: "probe " + id} +} + +func probeCallEvents(callID, toolName, id string) []zeroruntime.StreamEvent { + return []zeroruntime.StreamEvent{ + {Type: zeroruntime.StreamEventToolCallStart, ToolCallID: callID, ToolName: toolName}, + {Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: callID, ArgumentsFragment: fmt.Sprintf(`{"id":%q}`, id)}, + {Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: callID}, + } +} + +func TestParallelSafeToolCall(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(&probeTool{name: "probe_read", sideEffect: tools.SideEffectRead}) + registry.Register(&probeTool{name: "probe_write", sideEffect: tools.SideEffectWrite}) + + call := func(name, args string) ToolCall { return ToolCall{ID: "c", Name: name, Arguments: args} } + if !parallelSafeToolCall(registry, call("probe_read", `{"id":"a"}`), Options{}) { + t.Fatal("auto-allowed read tool must be parallel-safe") + } + if parallelSafeToolCall(registry, call("probe_write", `{"id":"a"}`), Options{}) { + t.Fatal("mutating tool must not be parallel-safe") + } + if parallelSafeToolCall(registry, call("unknown_tool", `{}`), Options{}) { + t.Fatal("unknown tool must not be parallel-safe") + } + if parallelSafeToolCall(registry, call("probe_read", `{"id":`), Options{}) { + t.Fatal("undecodable arguments must not be parallel-safe") + } + if parallelSafeToolCall(registry, call("ask_user", `{}`), Options{}) { + t.Fatal("loop-intercepted tools must stay sequential") + } +} + +func TestRunExecutesConsecutiveReadsConcurrently(t *testing.T) { + probe := &probeTool{name: "probe_read", sideEffect: tools.SideEffectRead, delay: 60 * time.Millisecond} + registry := tools.NewRegistry() + registry.Register(probe) + + turnOne := append(probeCallEvents("call-1", "probe_read", "a"), probeCallEvents("call-2", "probe_read", "b")...) + turnOne = append(turnOne, probeCallEvents("call-3", "probe_read", "c")...) + turnOne = append(turnOne, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}) + provider := &mockProvider{ + turns: [][]zeroruntime.StreamEvent{ + turnOne, + {{Type: zeroruntime.StreamEventText, Content: "done"}, {Type: zeroruntime.StreamEventDone}}, + }, + } + + var results []ToolResult + _, err := Run(context.Background(), "probe", provider, Options{ + Registry: registry, + OnToolResult: func(result ToolResult) { results = append(results, result) }, + }) + if err != nil { + t.Fatal(err) + } + if probe.maxActive < 2 { + t.Fatalf("consecutive read-only calls must overlap, max concurrency was %d", probe.maxActive) + } + // Results must still be recorded in original call order. + if len(results) != 3 || results[0].ToolCallID != "call-1" || results[1].ToolCallID != "call-2" || results[2].ToolCallID != "call-3" { + t.Fatalf("tool results out of order: %#v", results) + } + messages := provider.requests[1].Messages + var toolOrder []string + for _, message := range messages { + if message.Role == zeroruntime.MessageRoleTool { + toolOrder = append(toolOrder, message.ToolCallID) + } + } + if len(toolOrder) != 3 || toolOrder[0] != "call-1" || toolOrder[1] != "call-2" || toolOrder[2] != "call-3" { + t.Fatalf("recorded tool messages out of order: %v", toolOrder) + } +} + +func TestRunParallelReadsNeverSpanMutatingCall(t *testing.T) { + shared := &probeLog{} + read := &probeTool{name: "probe_read", sideEffect: tools.SideEffectRead, delay: 30 * time.Millisecond, shared: shared} + write := &probeTool{name: "probe_write", sideEffect: tools.SideEffectWrite, shared: shared} + registry := tools.NewRegistry() + registry.Register(read) + registry.Register(write) + + turnOne := append(probeCallEvents("call-1", "probe_read", "r1"), probeCallEvents("call-2", "probe_read", "r2")...) + turnOne = append(turnOne, probeCallEvents("call-3", "probe_write", "w")...) + turnOne = append(turnOne, probeCallEvents("call-4", "probe_read", "r3")...) + turnOne = append(turnOne, probeCallEvents("call-5", "probe_read", "r4")...) + turnOne = append(turnOne, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}) + provider := &mockProvider{ + turns: [][]zeroruntime.StreamEvent{ + turnOne, + {{Type: zeroruntime.StreamEventText, Content: "done"}, {Type: zeroruntime.StreamEventDone}}, + }, + } + + _, err := Run(context.Background(), "probe", provider, Options{Registry: registry}) + if err != nil { + t.Fatal(err) + } + + // Cross-probe ordering on the SHARED log: the write must start only after + // both first-batch reads finished, and both second-batch reads must start + // only after the write finished — batches never cross a mutating call. + log := shared.snapshot() + index := func(entry string) int { + for i, e := range log { + if e == entry { + return i + } + } + t.Fatalf("entry %q missing from shared log: %v", entry, log) + return -1 + } + writeStart, writeEnd := index("start:w"), index("end:w") + if firstBatchMaxEnd := max(index("end:r1"), index("end:r2")); writeStart < firstBatchMaxEnd { + t.Fatalf("write started before the first read batch finished: %v", log) + } + if secondBatchMinStart := min(index("start:r3"), index("start:r4")); secondBatchMinStart < writeEnd { + t.Fatalf("second read batch started before the write finished: %v", log) + } + if read.maxActive < 2 { + t.Fatalf("reads within a batch must overlap, max concurrency was %d", read.maxActive) + } +} diff --git a/internal/agent/types.go b/internal/agent/types.go index 2deb33a5c..dd3968e38 100644 --- a/internal/agent/types.go +++ b/internal/agent/types.go @@ -285,6 +285,12 @@ type Options struct { // ceiling and the autonomy gate. nil disables it entirely (the loop is // byte-identical to before). One instance per run — it holds attempt state. SelfCorrect *SelfCorrector + // FileDiagnostics, when set, is passed to mutating tools so edit_file / + // write_file append error-severity language diagnostics for the file they + // just wrote directly to their tool output — the model sees an error it + // introduced in the same turn instead of a later verification pass. Build + // one with NewFileDiagnostics. nil disables inline diagnostics. + FileDiagnostics func(ctx context.Context, absPath string) string // RequireCompletionSignal gates run completion for HEADLESS exec. Without it, // any assistant turn that produces text but no tool call is accepted as the diff --git a/internal/cli/app.go b/internal/cli/app.go index 7ab5aadf6..ccb58891d 100644 --- a/internal/cli/app.go +++ b/internal/cli/app.go @@ -19,6 +19,7 @@ import ( "github.com/Gitlawb/zero/internal/hooks" "github.com/Gitlawb/zero/internal/localcontrol" "github.com/Gitlawb/zero/internal/mcp" + "github.com/Gitlawb/zero/internal/modelregistry" "github.com/Gitlawb/zero/internal/observability" "github.com/Gitlawb/zero/internal/plugins" "github.com/Gitlawb/zero/internal/providerhealth" @@ -229,6 +230,12 @@ func runWithDeps(args []string, stdout io.Writer, stderr io.Writer, deps appDeps defer observability.Recover(observability.DefaultCrashDir(), "cli", stderr, &exitCode) deps = fillAppDeps(deps) + // CLI runs opt into the models.dev overlay (cached live context limits and + // pricing on top of the curated catalog). Explicitly enabled here — and only + // here — so library consumers and hermetic tests are never perturbed by a + // cache file on the machine. The refresh itself is fired in exec/TUI startup. + modelregistry.EnableModelsDevOverlay() + addDirs, args, err := splitLeadingAddDirFlags(args) if err != nil { return writeAppError(stderr, err.Error(), 1) @@ -544,6 +551,11 @@ func runInteractiveTUI(stderr io.Writer, deps appDeps, permissionMode agent.Perm } func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode agent.PermissionMode, addDirs []string, theme string, forceSetup bool) int { + // Refresh the models.dev pricing/limits cache in the background when stale; + // the overlay is read at registry construction from the cache file, so this + // benefits the next run and never blocks or fails this one. + go func() { _ = modelregistry.RefreshModelsDevCache(context.Background()) }() + workspaceRoot, err := deps.getwd() if err != nil { return writeAppError(stderr, "failed to resolve workspace: "+err.Error(), 1) diff --git a/internal/cli/exec.go b/internal/cli/exec.go index 2e2f62796..562ab9311 100644 --- a/internal/cli/exec.go +++ b/internal/cli/exec.go @@ -136,6 +136,11 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in return exitSuccess } + // Refresh the models.dev pricing/limits cache in the background when stale; + // the overlay is read at registry construction from the cache file, so this + // benefits the next run and never blocks or fails this one. + go func() { _ = modelregistry.RefreshModelsDevCache(context.Background()) }() + // A mode seeds model/effort/max-turns/tool filters as a preset. Expand it up // front — before tool-filter validation and the --list-tools branch — so a // mode-injected tool filter is validated and reflected in --list-tools, and a @@ -500,9 +505,10 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in // by default the corrector is nil, leaving the agent loop byte-identical. When // on we verify with both the workspace test plan and LSP diagnostics over the // changed files; the autonomy gate inside the corrector still decides whether - // failures auto-fix or just report. lspShutdown tears down any language-server - // sessions the LSP half spawned (no-op when self-correct is off). - selfCorrector, lspShutdown := newExecSelfCorrector(options.selfCorrect, workspaceRoot, options.autonomy) + // failures auto-fix or just report. fileDiagnostics (always on, lazy) gives + // edit_file/write_file inline error diagnostics for the file they just wrote. + // lspShutdown tears down any language-server sessions either half spawned. + selfCorrector, fileDiagnostics, lspShutdown := newExecSelfCorrector(options.selfCorrect, workspaceRoot, options.autonomy) defer lspShutdown() result, err := agent.Run(runCtx, agentPrompt, provider, agent.Options{ MaxTurns: resolved.MaxTurns, @@ -526,6 +532,7 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in PermissionMode: permissionMode, Autonomy: options.autonomy, SelfCorrect: selfCorrector, + FileDiagnostics: fileDiagnostics, // Headless exec: don't accept a no-tool-call turn as "done" while work // clearly remains (pending plan items / a mid-step continuation cue) — // nudge to continue, and finalize as INCOMPLETE rather than false success @@ -701,23 +708,28 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in // unless a changed file's language actually has one installed on PATH. The // returned cleanup shuts that manager down (terminating any spawned server // sessions); it is a no-op when self-correct is off. -func newExecSelfCorrector(enabled bool, workspaceRoot string, autonomy string) (*agent.SelfCorrector, func()) { +func newExecSelfCorrector(enabled bool, workspaceRoot string, autonomy string) (*agent.SelfCorrector, func(context.Context, string) string, func()) { + // The manager is created regardless of --self-correct: it also backs the + // always-on inline post-edit diagnostics (agent.NewFileDiagnostics), and it + // stays lazy — no language server is spawned unless an edited file's + // language actually has one installed on PATH. + manager := lsp.NewManager(workspaceRoot) + cleanup := func() { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = manager.Shutdown(ctx) + } + fileDiagnostics := agent.NewFileDiagnostics(manager, workspaceRoot) if !enabled { - return nil, func() {} + return nil, fileDiagnostics, cleanup } - manager := lsp.NewManager(workspaceRoot) corrector := agent.NewSelfCorrector(workspaceRoot, agent.NewLSPDiagnosticsChecker(manager), agent.NewProjectVerifier(workspaceRoot), agent.SelfCorrectConfig{ Enabled: true, IncludeTests: true, IncludeLSP: true, Autonomy: autonomy, }) - cleanup := func() { - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - _ = manager.Shutdown(ctx) - } - return corrector, cleanup + return corrector, fileDiagnostics, cleanup } func deferredEligibleCount(registry *tools.Registry, permissionMode agent.PermissionMode, enabledTools []string, disabledTools []string) int { diff --git a/internal/doctor/hardening.go b/internal/doctor/hardening.go index ea3a8f280..baf21e04c 100644 --- a/internal/doctor/hardening.go +++ b/internal/doctor/hardening.go @@ -131,29 +131,45 @@ func doctorSandboxPolicy(cfg config.SandboxConfig) sandbox.Policy { // lspServersCheck reports which language servers ZERO would use are present on // PATH. Missing servers are not a failure — ZERO degrades to text-only edits for // those languages — so the worst status is WARN, and each missing server gets an -// actionable install command keyed by its binary name. +// actionable install command keyed by its binary name. Only the tier-1 servers +// (lsp.CoreServerBinaries) drive the status: the long-tail servers configured +// for breadth are listed informationally under missingOptional, because warning +// about a missing zls on a machine with no Zig code would be permanent noise. func lspServersCheck(lookup func(string) (string, error)) Check { if lookup == nil { lookup = exec.LookPath } + core := map[string]bool{} + for _, binary := range lsp.CoreServerBinaries() { + core[binary] = true + } present := map[string]any{} missing := map[string]any{} + missingOptional := map[string]any{} for _, binary := range lsp.ServerBinaries() { if _, err := lookup(binary); err == nil { present[binary] = "on PATH" continue } - missing[binary] = lspRemedy(binary) + if core[binary] { + missing[binary] = lspRemedy(binary) + } else { + missingOptional[binary] = lspRemedy(binary) + } + } + details := map[string]any{"present": present} + if len(missingOptional) > 0 { + details["missingOptional"] = missingOptional } if len(missing) == 0 { - return check("lsp.servers", "LSP servers", StatusPass, "All known language servers are available on PATH.", map[string]any{ - "present": present, - }) + message := "All core language servers are available on PATH." + if len(missingOptional) > 0 { + message = fmt.Sprintf("All core language servers are available on PATH (%d optional server(s) not installed).", len(missingOptional)) + } + return check("lsp.servers", "LSP servers", StatusPass, message, details) } - return check("lsp.servers", "LSP servers", StatusWarn, fmt.Sprintf("%d language server(s) missing from PATH; affected files degrade to text-only edits.", len(missing)), map[string]any{ - "present": present, - "missing": missing, - }) + details["missing"] = missing + return check("lsp.servers", "LSP servers", StatusWarn, fmt.Sprintf("%d language server(s) missing from PATH; affected files degrade to text-only edits.", len(missing)), details) } // lspRemedy returns an actionable install command for a missing language-server diff --git a/internal/lsp/registry.go b/internal/lsp/registry.go index 23b40be6a..0331db6fe 100644 --- a/internal/lsp/registry.go +++ b/internal/lsp/registry.go @@ -9,26 +9,134 @@ import ( // serverCommands maps a file extension to the language-server command (argv) ZERO // will spawn. The first element is the binary looked up on PATH; missing binaries -// are not an error, the agent just degrades to text-only for that file. +// are not an error, the agent just degrades to text-only for that file. Every +// command is the language's community-standard server invoked in stdio mode, so +// a wrong guess can never spawn something surprising — only fail the PATH check. var serverCommands = map[string][]string{ - ".go": {"gopls", "serve"}, - ".ts": {"typescript-language-server", "--stdio"}, - ".tsx": {"typescript-language-server", "--stdio"}, - ".js": {"typescript-language-server", "--stdio"}, - ".jsx": {"typescript-language-server", "--stdio"}, - ".py": {"pyright-langserver", "--stdio"}, - ".rs": {"rust-analyzer"}, + ".go": {"gopls", "serve"}, + ".ts": {"typescript-language-server", "--stdio"}, + ".tsx": {"typescript-language-server", "--stdio"}, + ".mts": {"typescript-language-server", "--stdio"}, + ".cts": {"typescript-language-server", "--stdio"}, + ".js": {"typescript-language-server", "--stdio"}, + ".jsx": {"typescript-language-server", "--stdio"}, + ".mjs": {"typescript-language-server", "--stdio"}, + ".cjs": {"typescript-language-server", "--stdio"}, + ".py": {"pyright-langserver", "--stdio"}, + ".pyi": {"pyright-langserver", "--stdio"}, + ".rs": {"rust-analyzer"}, + ".c": {"clangd"}, + ".h": {"clangd"}, + ".cpp": {"clangd"}, + ".cc": {"clangd"}, + ".cxx": {"clangd"}, + ".hpp": {"clangd"}, + ".java": {"jdtls"}, + ".kt": {"kotlin-language-server"}, + ".kts": {"kotlin-language-server"}, + ".rb": {"ruby-lsp"}, + ".php": {"intelephense", "--stdio"}, + ".zig": {"zls"}, + ".lua": {"lua-language-server"}, + ".ex": {"elixir-ls"}, + ".exs": {"elixir-ls"}, + ".hs": {"haskell-language-server-wrapper", "--lsp"}, + ".swift": {"sourcekit-lsp"}, + ".ml": {"ocamllsp"}, + ".mli": {"ocamllsp"}, + ".scala": {"metals"}, + ".clj": {"clojure-lsp"}, + ".cljs": {"clojure-lsp"}, + ".cljc": {"clojure-lsp"}, + ".dart": {"dart", "language-server"}, + ".gleam": {"gleam", "lsp"}, + ".nix": {"nixd"}, + ".tf": {"terraform-ls", "serve"}, + ".sh": {"bash-language-server", "start"}, + ".bash": {"bash-language-server", "start"}, + ".yaml": {"yaml-language-server", "--stdio"}, + ".yml": {"yaml-language-server", "--stdio"}, + ".json": {"vscode-json-language-server", "--stdio"}, + ".css": {"vscode-css-language-server", "--stdio"}, + ".scss": {"vscode-css-language-server", "--stdio"}, + ".less": {"vscode-css-language-server", "--stdio"}, + ".html": {"vscode-html-language-server", "--stdio"}, + ".svelte": {"svelteserver", "--stdio"}, + ".vue": {"vue-language-server", "--stdio"}, + ".astro": {"astro-ls", "--stdio"}, } // languageIDs maps a file extension to the LSP languageId used in didOpen. var languageIDs = map[string]string{ - ".go": "go", - ".ts": "typescript", - ".tsx": "typescriptreact", - ".js": "javascript", - ".jsx": "javascriptreact", - ".py": "python", - ".rs": "rust", + ".go": "go", + ".ts": "typescript", + ".tsx": "typescriptreact", + ".mts": "typescript", + ".cts": "typescript", + ".js": "javascript", + ".jsx": "javascriptreact", + ".mjs": "javascript", + ".cjs": "javascript", + ".py": "python", + ".pyi": "python", + ".rs": "rust", + ".c": "c", + ".h": "c", + ".cpp": "cpp", + ".cc": "cpp", + ".cxx": "cpp", + ".hpp": "cpp", + ".java": "java", + ".kt": "kotlin", + ".kts": "kotlin", + ".rb": "ruby", + ".php": "php", + ".zig": "zig", + ".lua": "lua", + ".ex": "elixir", + ".exs": "elixir", + ".hs": "haskell", + ".swift": "swift", + ".ml": "ocaml", + ".mli": "ocaml", + ".scala": "scala", + ".clj": "clojure", + ".cljs": "clojure", + ".cljc": "clojure", + ".dart": "dart", + ".gleam": "gleam", + ".nix": "nix", + ".tf": "terraform", + ".sh": "shellscript", + ".bash": "shellscript", + ".yaml": "yaml", + ".yml": "yaml", + ".json": "json", + ".css": "css", + ".scss": "scss", + ".less": "less", + ".html": "html", + ".svelte": "svelte", + ".vue": "vue", + ".astro": "astro", +} + +// coreServerBinaries are the tier-1 servers for the languages agents hit most; +// `zero doctor` treats their absence as warn-worthy. The long tail of servers +// configured above for breadth is reported informationally only — warning about +// a missing zls on a machine with no Zig code would be permanent noise. +var coreServerBinaries = []string{ + "gopls", + "typescript-language-server", + "pyright-langserver", + "rust-analyzer", +} + +// CoreServerBinaries returns the tier-1 server binaries (sorted copy). +func CoreServerBinaries() []string { + binaries := append([]string(nil), coreServerBinaries...) + sort.Strings(binaries) + return binaries } // ServerBinaries returns the unique set of language-server binaries ZERO may diff --git a/internal/modelregistry/catalog.go b/internal/modelregistry/catalog.go index ed4d3a68b..378c88924 100644 --- a/internal/modelregistry/catalog.go +++ b/internal/modelregistry/catalog.go @@ -63,6 +63,10 @@ func DefaultModelEntries() []ModelEntry { googleModel("gemini-2.5-flash-lite", "Gemini 2.5 Flash-Lite", "gemini-2.5-flash-lite", ModelStatusActive, []string{"google:gemini-2.5-flash-lite", "gemini-flash-lite"}, ContextLimits{ContextWindow: 1_048_576, MaxOutputTokens: 65_536}, ModelCost{InputPerMillion: 0.1, CachedInputPerMillion: 0.01, OutputPerMillion: 0.4}, []ModelCapability{ModelCapabilityVision, ModelCapabilityJSONMode, ModelCapabilityReasoning, ModelCapabilityLongContext}, standardReasoningEfforts(), "Google low-cost Flash model for background routing and summaries."), } decorateModelDepth(entries) + // Overlay volatile facts (context limits, base pricing) from a cached + // models.dev snapshot when one is present and fresh — see modelsdev.go. + // Identity fields (ids, aliases, patterns, deprecations) stay curated. + entries = applyModelsDevOverrides(entries, cachedModelsDevProviders()) return cloneModelEntries(entries) } diff --git a/internal/modelregistry/modelsdev.go b/internal/modelregistry/modelsdev.go new file mode 100644 index 000000000..0f49cd236 --- /dev/null +++ b/internal/modelregistry/modelsdev.go @@ -0,0 +1,269 @@ +package modelregistry + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "sync" + "sync/atomic" + "time" +) + +// Live models.dev overlay for the curated catalog. The hand-maintained +// DefaultModelEntries list is the source of truth for identity — ids, aliases, +// match patterns, deprecations, escalation targets — but its VOLATILE facts +// (context window, max output tokens, per-million pricing) go stale between +// releases. When a cached snapshot of https://models.dev/api.json is present, +// those fields are refreshed from it at registry construction; everything else +// stays curated. The overlay never adds models (an auto-added entry would lack +// aliases, match patterns, and provider wiring) and never touches the network +// on the registry hot path — fetching happens only in the explicit background +// refresh, cached to disk with a TTL. + +const ( + modelsDevDefaultURL = "https://models.dev/api.json" + // modelsDevRefreshAfter is how old the cache may get before a background + // refresh re-fetches it. + modelsDevRefreshAfter = 24 * time.Hour + // modelsDevMaxAge is the oldest cache still applied as an overlay. Beyond + // this the curated catalog (updated with the binary) is likely fresher than + // the snapshot, so a stale file is ignored rather than trusted. + modelsDevMaxAge = 7 * 24 * time.Hour + modelsDevFetchLimit = 32 << 20 // 32MiB guard on the response body + modelsDevFetchWindow = 15 * time.Second +) + +// modelsDevModel is the subset of a models.dev model record the overlay uses. +type modelsDevModel struct { + Limit struct { + Context int `json:"context"` + Output int `json:"output"` + } `json:"limit"` + Cost struct { + Input float64 `json:"input"` + Output float64 `json:"output"` + CacheRead float64 `json:"cache_read"` + CacheWrite float64 `json:"cache_write"` + } `json:"cost"` +} + +// modelsDevProvider matches one provider object in api.json. +type modelsDevProvider struct { + Models map[string]modelsDevModel `json:"models"` +} + +// parseModelsDev decodes an api.json document into provider-slug -> api-model +// -> record. api.json is a top-level object keyed by provider id. +func parseModelsDev(data []byte) (map[string]map[string]modelsDevModel, error) { + var doc map[string]modelsDevProvider + if err := json.Unmarshal(data, &doc); err != nil { + return nil, fmt.Errorf("modelregistry: parse models.dev: %w", err) + } + providers := make(map[string]map[string]modelsDevModel, len(doc)) + for slug, provider := range doc { + if len(provider.Models) == 0 { + continue + } + providers[slug] = provider.Models + } + if len(providers) == 0 { + return nil, fmt.Errorf("modelregistry: models.dev document has no providers") + } + return providers, nil +} + +// modelsDevSlugs maps an entry's provider kind to the models.dev provider ids +// worth checking. Only first-party slugs are used: router entries on models.dev +// carry generic or stale numbers. +func modelsDevSlugs(kind ProviderKind) []string { + switch kind { + case ProviderAnthropic: + return []string{"anthropic"} + case ProviderOpenAI: + return []string{"openai"} + case ProviderGoogle: + return []string{"google", "google-vertex"} + default: + return nil + } +} + +// applyModelsDevOverrides refreshes each curated entry's context limits and +// base pricing from the snapshot, when the snapshot knows the model. Tiered +// pricing is never touched: models.dev has no tier data, and mixing a live +// base rate with curated tiers would misprice the tier boundaries. +func applyModelsDevOverrides(entries []ModelEntry, providers map[string]map[string]modelsDevModel) []ModelEntry { + if len(providers) == 0 { + return entries + } + for i := range entries { + entry := &entries[i] + var record modelsDevModel + found := false + for _, slug := range modelsDevSlugs(entry.Provider) { + if models, ok := providers[slug]; ok { + if candidate, ok := models[strings.TrimSpace(entry.APIModel)]; ok { + record = candidate + found = true + break + } + } + } + if !found { + continue + } + if record.Limit.Context > 0 { + entry.ContextLimits.ContextWindow = record.Limit.Context + } + if record.Limit.Output > 0 { + entry.ContextLimits.MaxOutputTokens = record.Limit.Output + } + if len(entry.Cost.Tiers) == 0 && record.Cost.Input > 0 && record.Cost.Output > 0 { + entry.Cost.InputPerMillion = record.Cost.Input + entry.Cost.OutputPerMillion = record.Cost.Output + if record.Cost.CacheRead > 0 { + entry.Cost.CachedInputPerMillion = record.Cost.CacheRead + } + if record.Cost.CacheWrite > 0 { + entry.Cost.CacheWritePerMillion = record.Cost.CacheWrite + } + entry.Cost.Source = "models.dev/api.json (cached)" + } + } + return entries +} + +// modelsDevCachePath returns the on-disk cache location. ZERO_MODELS_CACHE_PATH +// overrides it (used by tests and unusual setups). +func modelsDevCachePath() (string, error) { + if override := strings.TrimSpace(os.Getenv("ZERO_MODELS_CACHE_PATH")); override != "" { + return override, nil + } + base, err := os.UserCacheDir() + if err != nil { + return "", err + } + return filepath.Join(base, "zero", "modelsdev.json"), nil +} + +var ( + modelsDevOnce sync.Once + modelsDevCached map[string]map[string]modelsDevModel + modelsDevEnabled atomic.Bool +) + +// EnableModelsDevOverlay opts the process into applying the cached models.dev +// snapshot on top of the curated catalog. The CLI entrypoint calls it; library +// consumers and tests that never do get the curated catalog byte-identical to +// before, so hermetic tests can't be perturbed by a cache file on the machine. +// ZERO_DISABLE_MODELS_FETCH disables both the overlay and the fetch. +func EnableModelsDevOverlay() { + modelsDevEnabled.Store(true) +} + +// cachedModelsDevProviders loads the cached snapshot once per process. Not +// enabled, missing, stale (> modelsDevMaxAge), or malformed all yield nil and +// the curated catalog is used untouched. Read once deliberately: +// DefaultRegistry is called on hot paths (pickers, cost views) and must not +// re-stat the file every time; a background refresh benefits the NEXT process. +func cachedModelsDevProviders() map[string]map[string]modelsDevModel { + if !modelsDevEnabled.Load() || strings.TrimSpace(os.Getenv("ZERO_DISABLE_MODELS_FETCH")) != "" { + return nil + } + modelsDevOnce.Do(func() { + path, err := modelsDevCachePath() + if err != nil { + return + } + info, err := os.Stat(path) + if err != nil || time.Since(info.ModTime()) > modelsDevMaxAge { + return + } + data, err := os.ReadFile(path) + if err != nil { + return + } + providers, err := parseModelsDev(data) + if err != nil { + return + } + modelsDevCached = providers + }) + return modelsDevCached +} + +// resetModelsDevCacheForTest clears the process-level cache memoization and +// disables the overlay. +func resetModelsDevCacheForTest() { + modelsDevOnce = sync.Once{} + modelsDevCached = nil + modelsDevEnabled.Store(false) +} + +// RefreshModelsDevCache fetches models.dev/api.json into the on-disk cache +// when the cache is missing or older than modelsDevRefreshAfter. It is safe to +// call fire-and-forget from startup (use a goroutine); it never affects the +// current process's registry (see cachedModelsDevProviders). Disabled entirely +// by ZERO_DISABLE_MODELS_FETCH. The URL can be overridden with ZERO_MODELS_URL. +func RefreshModelsDevCache(ctx context.Context) error { + if strings.TrimSpace(os.Getenv("ZERO_DISABLE_MODELS_FETCH")) != "" { + return nil + } + path, err := modelsDevCachePath() + if err != nil { + return err + } + if info, err := os.Stat(path); err == nil && time.Since(info.ModTime()) < modelsDevRefreshAfter { + return nil + } + + url := strings.TrimSpace(os.Getenv("ZERO_MODELS_URL")) + if url == "" { + url = modelsDevDefaultURL + } + fetchCtx, cancel := context.WithTimeout(ctx, modelsDevFetchWindow) + defer cancel() + request, err := http.NewRequestWithContext(fetchCtx, http.MethodGet, url, nil) + if err != nil { + return err + } + request.Header.Set("User-Agent", "zero-models-refresh/0.1") + response, err := http.DefaultClient.Do(request) + if err != nil { + return err + } + defer response.Body.Close() + if response.StatusCode != http.StatusOK { + return fmt.Errorf("modelregistry: models.dev fetch: HTTP %d", response.StatusCode) + } + data, err := io.ReadAll(io.LimitReader(response.Body, modelsDevFetchLimit)) + if err != nil { + return err + } + // Validate before persisting: a bad body must never clobber a good cache. + if _, err := parseModelsDev(data); err != nil { + return err + } + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + return err + } + temp, err := os.CreateTemp(filepath.Dir(path), "modelsdev-*.json") + if err != nil { + return err + } + if _, err := temp.Write(data); err != nil { + _ = temp.Close() + _ = os.Remove(temp.Name()) + return err + } + if err := temp.Close(); err != nil { + _ = os.Remove(temp.Name()) + return err + } + return os.Rename(temp.Name(), path) +} diff --git a/internal/modelregistry/modelsdev_test.go b/internal/modelregistry/modelsdev_test.go new file mode 100644 index 000000000..78ade7d80 --- /dev/null +++ b/internal/modelregistry/modelsdev_test.go @@ -0,0 +1,228 @@ +package modelregistry + +import ( + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "testing" + "time" +) + +const sampleModelsDev = `{ + "anthropic": { + "id": "anthropic", + "models": { + "claude-sonnet-4-5-20250929": { + "limit": {"context": 1000000, "output": 64000}, + "cost": {"input": 3.5, "output": 17.5, "cache_read": 0.35, "cache_write": 4.4} + } + } + }, + "google": { + "id": "google", + "models": { + "gemini-2.5-pro": { + "limit": {"context": 2097152, "output": 65536}, + "cost": {"input": 9.99, "output": 9.99} + } + } + } +}` + +func TestParseModelsDev(t *testing.T) { + providers, err := parseModelsDev([]byte(sampleModelsDev)) + if err != nil { + t.Fatal(err) + } + record, ok := providers["anthropic"]["claude-sonnet-4-5-20250929"] + if !ok { + t.Fatal("expected anthropic sonnet record") + } + if record.Limit.Context != 1_000_000 || record.Cost.Input != 3.5 { + t.Fatalf("unexpected record: %+v", record) + } + if _, err := parseModelsDev([]byte(`{}`)); err == nil { + t.Fatal("empty document must be rejected") + } + if _, err := parseModelsDev([]byte(`not json`)); err == nil { + t.Fatal("malformed document must be rejected") + } +} + +func TestApplyModelsDevOverrides(t *testing.T) { + // Point the cache at a non-existent file so DefaultModelEntries returns the + // pure curated catalog, then apply the sample snapshot explicitly. + t.Setenv("ZERO_MODELS_CACHE_PATH", filepath.Join(t.TempDir(), "absent.json")) + resetModelsDevCacheForTest() + t.Cleanup(resetModelsDevCacheForTest) + + providers, err := parseModelsDev([]byte(sampleModelsDev)) + if err != nil { + t.Fatal(err) + } + entries := applyModelsDevOverrides(DefaultModelEntries(), providers) + + var sonnet, geminiPro, opus ModelEntry + for _, entry := range entries { + switch entry.ID { + case "claude-sonnet-4.5": + sonnet = entry + case "gemini-2.5-pro": + geminiPro = entry + case "claude-opus-4.1": + opus = entry + } + } + + // Known model: limits and base pricing refreshed from the snapshot. + if sonnet.ContextLimits.ContextWindow != 1_000_000 || sonnet.ContextLimits.MaxOutputTokens != 64_000 { + t.Fatalf("sonnet limits not overridden: %+v", sonnet.ContextLimits) + } + if sonnet.Cost.InputPerMillion != 3.5 || sonnet.Cost.OutputPerMillion != 17.5 || sonnet.Cost.CacheWritePerMillion != 4.4 { + t.Fatalf("sonnet cost not overridden: %+v", sonnet.Cost) + } + if sonnet.Cost.Source != "models.dev/api.json (cached)" { + t.Fatalf("sonnet cost source not marked: %q", sonnet.Cost.Source) + } + + // Tiered pricing is curated: limits refresh, cost must NOT (gemini-2.5-pro + // has curated tiers and the snapshot's flat 9.99 would misprice them). + if geminiPro.ContextLimits.ContextWindow != 2_097_152 { + t.Fatalf("gemini limits not overridden: %+v", geminiPro.ContextLimits) + } + if geminiPro.Cost.InputPerMillion == 9.99 || len(geminiPro.Cost.Tiers) == 0 { + t.Fatalf("tiered cost must stay curated: %+v", geminiPro.Cost) + } + + // Model absent from the snapshot: untouched. + if opus.ContextLimits.ContextWindow != 200_000 || opus.Cost.InputPerMillion != 15 { + t.Fatalf("opus must be untouched: %+v %+v", opus.ContextLimits, opus.Cost) + } +} + +func TestRefreshModelsDevCacheFetchesAndCaches(t *testing.T) { + hits := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits++ + _, _ = w.Write([]byte(sampleModelsDev)) + })) + defer server.Close() + + cachePath := filepath.Join(t.TempDir(), "modelsdev.json") + t.Setenv("ZERO_MODELS_CACHE_PATH", cachePath) + t.Setenv("ZERO_MODELS_URL", server.URL) + t.Setenv("ZERO_DISABLE_MODELS_FETCH", "") + + if err := RefreshModelsDevCache(t.Context()); err != nil { + t.Fatal(err) + } + if hits != 1 { + t.Fatalf("expected 1 fetch, got %d", hits) + } + if _, err := os.Stat(cachePath); err != nil { + t.Fatalf("cache file missing: %v", err) + } + // Fresh cache: second call must not re-fetch. + if err := RefreshModelsDevCache(t.Context()); err != nil { + t.Fatal(err) + } + if hits != 1 { + t.Fatalf("fresh cache must skip the fetch, got %d hits", hits) + } +} + +func TestRefreshModelsDevCacheRejectsBadBodyWithoutClobbering(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte("not json")) + })) + defer server.Close() + + cachePath := filepath.Join(t.TempDir(), "modelsdev.json") + if err := os.WriteFile(cachePath, []byte(sampleModelsDev), 0o644); err != nil { + t.Fatal(err) + } + stale := time.Now().Add(-48 * time.Hour) + if err := os.Chtimes(cachePath, stale, stale); err != nil { + t.Fatal(err) + } + t.Setenv("ZERO_MODELS_CACHE_PATH", cachePath) + t.Setenv("ZERO_MODELS_URL", server.URL) + + if err := RefreshModelsDevCache(t.Context()); err == nil { + t.Fatal("bad body must return an error") + } + content, err := os.ReadFile(cachePath) + if err != nil || string(content) != sampleModelsDev { + t.Fatal("bad fetch must not clobber the existing cache") + } +} + +func TestRefreshModelsDevCacheDisabledByEnv(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + t.Fatal("fetch must not happen when disabled") + })) + defer server.Close() + t.Setenv("ZERO_MODELS_CACHE_PATH", filepath.Join(t.TempDir(), "modelsdev.json")) + t.Setenv("ZERO_MODELS_URL", server.URL) + t.Setenv("ZERO_DISABLE_MODELS_FETCH", "1") + if err := RefreshModelsDevCache(t.Context()); err != nil { + t.Fatal(err) + } +} + +func TestCachedModelsDevProvidersIgnoresStaleCache(t *testing.T) { + cachePath := filepath.Join(t.TempDir(), "modelsdev.json") + if err := os.WriteFile(cachePath, []byte(sampleModelsDev), 0o644); err != nil { + t.Fatal(err) + } + stale := time.Now().Add(-modelsDevMaxAge - time.Hour) + if err := os.Chtimes(cachePath, stale, stale); err != nil { + t.Fatal(err) + } + t.Setenv("ZERO_MODELS_CACHE_PATH", cachePath) + resetModelsDevCacheForTest() + t.Cleanup(resetModelsDevCacheForTest) + EnableModelsDevOverlay() + + if providers := cachedModelsDevProviders(); providers != nil { + t.Fatal("stale cache must be ignored") + } +} + +func TestCachedModelsDevProvidersRequiresOptIn(t *testing.T) { + cachePath := filepath.Join(t.TempDir(), "modelsdev.json") + if err := os.WriteFile(cachePath, []byte(sampleModelsDev), 0o644); err != nil { + t.Fatal(err) + } + t.Setenv("ZERO_MODELS_CACHE_PATH", cachePath) + resetModelsDevCacheForTest() + t.Cleanup(resetModelsDevCacheForTest) + + // Without EnableModelsDevOverlay a fresh, valid cache must still be ignored: + // library consumers and hermetic tests get the pure curated catalog. + if providers := cachedModelsDevProviders(); providers != nil { + t.Fatal("overlay must be opt-in") + } +} + +func TestDefaultModelEntriesAppliesFreshCache(t *testing.T) { + cachePath := filepath.Join(t.TempDir(), "modelsdev.json") + if err := os.WriteFile(cachePath, []byte(sampleModelsDev), 0o644); err != nil { + t.Fatal(err) + } + t.Setenv("ZERO_MODELS_CACHE_PATH", cachePath) + resetModelsDevCacheForTest() + t.Cleanup(resetModelsDevCacheForTest) + EnableModelsDevOverlay() + + for _, entry := range DefaultModelEntries() { + if entry.ID == "claude-sonnet-4.5" { + if entry.ContextLimits.ContextWindow != 1_000_000 { + t.Fatalf("fresh cache must overlay the registry: %+v", entry.ContextLimits) + } + return + } + } + t.Fatal("claude-sonnet-4.5 not found") +} diff --git a/internal/providers/anthropic/cache_breakpoints_test.go b/internal/providers/anthropic/cache_breakpoints_test.go new file mode 100644 index 000000000..40f984233 --- /dev/null +++ b/internal/providers/anthropic/cache_breakpoints_test.go @@ -0,0 +1,96 @@ +package anthropic + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/Gitlawb/zero/internal/zeroruntime" +) + +// The final two messages must carry a cache_control breakpoint on their last +// cacheable block so the conversation transcript — not just system prompt and +// tools — is a prompt-cache hit on the next turn. Earlier messages must stay +// unmarked (Anthropic caps breakpoints at 4 per request: system, tools, and +// these two). +func TestAnthropicRequestMarksLastTwoMessagesForCaching(t *testing.T) { + var gotBody map[string]any + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := json.NewDecoder(r.Body).Decode(&gotBody); err != nil { + t.Fatalf("decode request body: %v", err) + } + writeSSEEvent(w, "message_stop", `{"type":"message_stop"}`) + })) + defer server.Close() + + provider, err := New(Options{BaseURL: server.URL + "/", Model: "claude-test"}) + if err != nil { + t.Fatalf("New returned error: %v", err) + } + stream, err := provider.StreamCompletion(context.Background(), zeroruntime.CompletionRequest{ + Messages: []zeroruntime.Message{ + {Role: zeroruntime.MessageRoleUser, Content: "first"}, + { + Role: zeroruntime.MessageRoleAssistant, + Content: "calling a tool", + ToolCalls: []zeroruntime.ToolCall{{ID: "toolu_1", Name: "grep", Arguments: `{}`}}, + }, + {Role: zeroruntime.MessageRoleTool, Content: "result", ToolCallID: "toolu_1"}, + }, + }) + if err != nil { + t.Fatalf("StreamCompletion returned error: %v", err) + } + drain(stream) + + messages := gotBody["messages"].([]any) + if len(messages) != 3 { + t.Fatalf("expected 3 wire messages, got %d: %#v", len(messages), messages) + } + lastBlockHasCache := func(message any) bool { + blocks, ok := message.(map[string]any)["content"].([]any) + if !ok || len(blocks) == 0 { + return false + } + last := blocks[len(blocks)-1].(map[string]any) + control, ok := last["cache_control"].(map[string]any) + return ok && control["type"] == "ephemeral" + } + if lastBlockHasCache(messages[0]) { + t.Fatalf("first message must not carry a breakpoint: %#v", messages[0]) + } + if !lastBlockHasCache(messages[1]) { + t.Fatalf("second-to-last message must carry a breakpoint: %#v", messages[1]) + } + if !lastBlockHasCache(messages[2]) { + t.Fatalf("last message must carry a breakpoint: %#v", messages[2]) + } +} + +// A string-content message must be converted to a block array so it can carry +// the breakpoint, and thinking blocks must never carry cache_control (the API +// rejects them) — the marker goes on the last cacheable block instead. +func TestApplyMessageCacheBreakpointsSkipsThinkingBlocks(t *testing.T) { + messages := []anthropicMessage{ + {Role: "user", Content: "plain string"}, + {Role: "assistant", Content: []map[string]any{ + {"type": "text", "text": "answer"}, + {"type": "thinking", "thinking": "hmm", "signature": "sig"}, + }}, + } + applyMessageCacheBreakpoints(messages) + + userBlocks := contentBlocks(messages[0].Content) + if len(userBlocks) != 1 || userBlocks[0]["cache_control"] == nil { + t.Fatalf("string content must become a marked block: %#v", messages[0].Content) + } + assistantBlocks := contentBlocks(messages[1].Content) + if assistantBlocks[1]["cache_control"] != nil { + t.Fatalf("thinking block must not be marked: %#v", assistantBlocks[1]) + } + if assistantBlocks[0]["cache_control"] == nil { + t.Fatalf("text block must carry the breakpoint instead: %#v", assistantBlocks[0]) + } +} diff --git a/internal/providers/anthropic/provider.go b/internal/providers/anthropic/provider.go index c072567be..5c0774373 100644 --- a/internal/providers/anthropic/provider.go +++ b/internal/providers/anthropic/provider.go @@ -405,9 +405,35 @@ func (provider *Provider) anthropicRequest(request zeroruntime.CompletionRequest } mapped.Tools[len(mapped.Tools)-1].CacheControl = &cacheControl{Type: cacheEphemeral} } + applyMessageCacheBreakpoints(mapped.Messages) return mapped, nil } +// applyMessageCacheBreakpoints marks the last content block of the final two +// messages with cache_control so the conversation transcript is cached +// turn-over-turn, not just the system prompt and tool definitions. Two message +// breakpoints (not one) keep the previous turn's prefix a cache hit while the +// newest suffix is being written. Anthropic allows at most 4 breakpoints per +// request: system + tools use two above, these use the remaining two. Thinking +// blocks cannot carry cache_control, so the marker goes on the last block that +// can. +func applyMessageCacheBreakpoints(messages []anthropicMessage) { + marked := 0 + for i := len(messages) - 1; i >= 0 && marked < 2; i-- { + blocks := contentBlocks(messages[i].Content) + for j := len(blocks) - 1; j >= 0; j-- { + blockType, _ := blocks[j]["type"].(string) + if blockType == "thinking" || blockType == "redacted_thinking" { + continue + } + blocks[j]["cache_control"] = map[string]any{"type": cacheEphemeral} + messages[i].Content = blocks + marked++ + break + } + } +} + func mapMessages(messages []zeroruntime.Message) (string, []anthropicMessage, error) { systemParts := []string{} mapped := []anthropicMessage{} diff --git a/internal/providers/providerio/retry.go b/internal/providers/providerio/retry.go index 10edc9a0d..256ac2a17 100644 --- a/internal/providers/providerio/retry.go +++ b/internal/providers/providerio/retry.go @@ -23,12 +23,18 @@ import ( // (billable) work. Only the INITIAL request is ever in scope; once the response // body starts streaming it is never re-issued. -const defaultMaxRetryAttempts = 3 +const defaultMaxRetryAttempts = 6 // maxBackoff caps a single backoff wait so a hostile or buggy Retry-After can't // stall the agent for minutes. const maxBackoff = 30 * time.Second +// retryBackoffBase is the first wait when the server supplied no Retry-After. +// Rate-limit windows are measured in seconds, not milliseconds: retrying a 429 +// after 400ms almost always burns the attempt while still limited, so the +// schedule is 2s, 4s, 8s, 16s, then maxBackoff. A var so tests can shrink it. +var retryBackoffBase = 2 * time.Second + // SendWithRetry issues an HTTP request, retrying ONLY the safe-to-replay server // responses (429 and 503, see ShouldRetryStatus) up to maxAttempts — backing off // between tries and honoring a server Retry-After header and context @@ -110,17 +116,11 @@ func ShouldRetryStatus(code int) bool { } // Backoff waits before retry attempt N (1-based), returning false if the context -// is cancelled during the wait. The wait is attempt*400ms unless the server -// supplied a (positive) Retry-After, and is capped at maxBackoff. +// is cancelled during the wait. A server-supplied (positive) Retry-After wins; +// otherwise the wait doubles from retryBackoffBase per attempt. Either way the +// wait is capped at maxBackoff. func Backoff(ctx context.Context, attempt int, retryAfter time.Duration) bool { - wait := time.Duration(attempt) * 400 * time.Millisecond - if retryAfter > 0 { - wait = retryAfter - } - if wait > maxBackoff { - wait = maxBackoff - } - timer := time.NewTimer(wait) + timer := time.NewTimer(backoffWait(attempt, retryAfter)) defer timer.Stop() select { case <-ctx.Done(): @@ -130,6 +130,27 @@ func Backoff(ctx context.Context, attempt int, retryAfter time.Duration) bool { } } +// backoffWait computes the wait before retry attempt N (1-based): Retry-After +// when supplied, else exponential from retryBackoffBase, both capped at +// maxBackoff. The exponent is clamped so a large attempt count cannot overflow. +func backoffWait(attempt int, retryAfter time.Duration) time.Duration { + wait := retryAfter + if wait <= 0 { + exponent := attempt - 1 + if exponent > 5 { + exponent = 5 + } + if exponent < 0 { + exponent = 0 + } + wait = retryBackoffBase * time.Duration(1< maxBackoff { + wait = maxBackoff + } + return wait +} + // RetryAfter parses a response's Retry-After header (delay-seconds or an HTTP // date) into a positive duration, or 0 when absent/unparseable. The result is // capped at maxBackoff by Backoff. diff --git a/internal/providers/providerio/retry_test.go b/internal/providers/providerio/retry_test.go index e83774e90..5d20dc82a 100644 --- a/internal/providers/providerio/retry_test.go +++ b/internal/providers/providerio/retry_test.go @@ -94,7 +94,40 @@ func TestBackoffWaitsThenReturnsTrue(t *testing.T) { } } +// shrinkBackoff makes retry waits negligible for the duration of a test. +func shrinkBackoff(t *testing.T) { + t.Helper() + saved := retryBackoffBase + retryBackoffBase = time.Millisecond + t.Cleanup(func() { retryBackoffBase = saved }) +} + +func TestBackoffWaitSchedule(t *testing.T) { + // Without Retry-After the wait doubles per attempt from 2s and caps at 30s; + // a supplied Retry-After wins but is capped too. + cases := []struct { + attempt int + retryAfter time.Duration + want time.Duration + }{ + {1, 0, 2 * time.Second}, + {2, 0, 4 * time.Second}, + {3, 0, 8 * time.Second}, + {4, 0, 16 * time.Second}, + {5, 0, 30 * time.Second}, // 32s capped + {50, 0, 30 * time.Second}, // clamped exponent, no overflow + {1, 7 * time.Second, 7 * time.Second}, + {1, 5 * time.Minute, 30 * time.Second}, // hostile Retry-After capped + } + for _, c := range cases { + if got := backoffWait(c.attempt, c.retryAfter); got != c.want { + t.Errorf("backoffWait(%d, %v) = %v, want %v", c.attempt, c.retryAfter, got, c.want) + } + } +} + func TestSendWithRetryRetriesThenSucceeds(t *testing.T) { + shrinkBackoff(t) var hits int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if atomic.AddInt32(&hits, 1) == 1 { @@ -140,6 +173,7 @@ func TestSendWithRetryReturnsNonRetryableImmediately(t *testing.T) { } func TestSendWithRetryReturnsLastResponseAfterMaxAttempts(t *testing.T) { + shrinkBackoff(t) var hits int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { atomic.AddInt32(&hits, 1) diff --git a/internal/tools/edit_file.go b/internal/tools/edit_file.go index b93a88d4c..55c427eea 100644 --- a/internal/tools/edit_file.go +++ b/internal/tools/edit_file.go @@ -2,6 +2,7 @@ package tools import ( "context" + "errors" "fmt" "os" "strings" @@ -44,7 +45,7 @@ func (tool editFileTool) Run(ctx context.Context, args map[string]any) Result { return tool.RunWithOptions(ctx, args, RunOptions{}) } -func (tool editFileTool) RunWithOptions(_ context.Context, args map[string]any, options RunOptions) Result { +func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any, options RunOptions) Result { requestedPath, err := aliasedStringArg(args, []string{"path", "file", "file_path", "filename"}, "", true, false) if err != nil { return errorResult("Error: Invalid arguments for edit_file: " + err.Error()) @@ -98,6 +99,35 @@ func (tool editFileTool) RunWithOptions(_ context.Context, args map[string]any, } } + // Fuzzy fallback: when the exact string (and its CRLF translation) is absent, + // run a cascade of tolerant matchers (trimmed lines, block anchors, collapsed + // whitespace, indentation drift, escape normalization) to locate the span the + // model intended. Only a span that occurs literally in the file is accepted, + // so the replacement applied below is still exact. + if occurrences == 0 { + findOld, findNew := oldString, newString + if strings.Contains(content, "\r\n") && !strings.Contains(findOld, "\r\n") { + findOld = strings.ReplaceAll(findOld, "\n", "\r\n") + findNew = strings.ReplaceAll(findNew, "\n", "\r\n") + } + search, ferr := fuzzyEditMatch(content, findOld, replaceAll) + switch { + case ferr == nil: + oldString = search + // The model's new_string was written at old_string's (mismatched) + // indentation; re-shape it to the span actually being replaced so a + // tolerant match never strips indentation or a trailing CR. + newString = adaptReplacementToSpan(search, findOld, findNew) + occurrences = strings.Count(content, search) + case errors.Is(ferr, errEditFuzzyAmbiguous): + return errorResult("Error: old_string matches multiple locations in " + relativePath + " even after fuzzy matching. Provide more surrounding context to make the match unique, or pass replace_all: true.") + case errors.Is(ferr, errEditFuzzyNotFound): + // Fall through to the exact-match error below. + default: + return errorResult("Error editing " + relativePath + ": " + ferr.Error()) + } + } + if occurrences == 0 { return errorResult("Error: Could not find the exact string to replace in " + relativePath + ". The old_string must match the file byte-for-byte.") } @@ -120,6 +150,10 @@ func (tool editFileTool) RunWithOptions(_ context.Context, args map[string]any, if err := os.WriteFile(absolutePath, []byte(updated), 0o644); err != nil { return errorResult("Error writing " + relativePath + ": " + err.Error()) } + // Optional format-on-write (ZERO_FORMAT_ON_WRITE). Must run BEFORE the + // FileTracker re-baseline: recording pre-format content would make the very + // next edit look like an external modification and trip the conflict guard. + updated = maybeFormatWrittenFile(ctx, absolutePath, updated) // Re-baseline to the content we just wrote so subsequent edits in this session // compare against the current on-disk state, not the pre-edit version. newInfo, _ := os.Stat(absolutePath) @@ -130,6 +164,7 @@ func (tool editFileTool) RunWithOptions(_ context.Context, args map[string]any, suffix = "s" } summary := fmt.Sprintf("Successfully edited %s (replaced %d occurrence%s).", relativePath, replacedCount, suffix) + summary += inlineDiagnostics(ctx, options, absolutePath, relativePath) result := okResult(summary) result.ChangedFiles = []string{relativePath} // Card-only preview (Display.Preview): the model's Output stays the one-line diff --git a/internal/tools/edit_replacers.go b/internal/tools/edit_replacers.go new file mode 100644 index 000000000..ff168e78b --- /dev/null +++ b/internal/tools/edit_replacers.go @@ -0,0 +1,554 @@ +package tools + +// Fuzzy fallback matching for edit_file. When the model's old_string fails to +// match byte-for-byte (drifted indentation, collapsed whitespace, escaped +// characters, a slightly-misremembered middle line), these replacers propose +// candidate spans that plausibly correspond to what the model intended. Only a +// candidate that occurs literally in the file is accepted, so the replacement +// itself is always exact even when the match was tolerant. +// +// Strategy cascade ported from opencode's edit tool, whose replacers were in +// turn distilled from Cline's diff-apply evals and gemini-cli's edit corrector. + +import ( + "errors" + "regexp" + "sort" + "strings" +) + +// editReplacer proposes candidate spans for find inside content. Candidates +// are re-validated by fuzzyEditMatch before use. +type editReplacer func(content, find string) []string + +var ( + errEditFuzzyNotFound = errors.New("no fuzzy match for old_string") + errEditFuzzyAmbiguous = errors.New("fuzzy match for old_string is ambiguous") +) + +// Minimum average middle-line similarity for block-anchor matches. +const editAnchorSimilarityThreshold = 0.65 + +// fuzzyEditMatch runs the replacer cascade and returns the exact span of +// content to replace. When replaceAll is false the span must be unique in +// content; an ambiguous candidate is skipped in favor of later candidates and +// only reported if nothing unique is found. A span wildly larger than +// old_string is refused outright rather than risking a destructive edit. +func fuzzyEditMatch(content, find string, replaceAll bool) (string, error) { + replacers := []editReplacer{ + lineTrimmedReplacer, + blockAnchorReplacer, + whitespaceNormalizedReplacer, + indentationFlexibleReplacer, + escapeNormalizedReplacer, + trimmedBoundaryReplacer, + contextAwareReplacer, + } + found := false + for _, replacer := range replacers { + // Collect the replacer's DISTINCT candidate spans that literally occur in + // content. Two or more distinct spans from one strategy (e.g. duplicate + // blocks at different indentation, each occurring once) mean the model's + // intent is genuinely ambiguous — silently editing the first would be a + // wrong-span write, so it is rejected instead. + var candidates []string + seen := map[string]bool{} + for _, search := range replacer(content, find) { + if search == "" || seen[search] { + continue + } + seen[search] = true + if strings.Index(content, search) < 0 { + continue + } + candidates = append(candidates, search) + } + if len(candidates) == 0 { + continue + } + found = true + if !replaceAll && len(candidates) > 1 { + return "", errEditFuzzyAmbiguous + } + search := candidates[0] + if isDisproportionateEditMatch(search, find) { + return "", errors.New("refusing replacement because the matched span is much larger than old_string; re-read the file and provide the full exact old_string for the intended replacement") + } + if replaceAll { + return search, nil + } + if strings.Index(content, search) == strings.LastIndex(content, search) { + return search, nil + } + // The single candidate occurs at multiple positions; a later, stricter + // strategy may still resolve a unique span, so keep cascading. + } + if !found { + return "", errEditFuzzyNotFound + } + return "", errEditFuzzyAmbiguous +} + +// isDisproportionateEditMatch guards against anchor-style replacers matching a +// span far larger than the text the model asked to replace (e.g. first/last +// line anchors bridging hundreds of unrelated lines). +func isDisproportionateEditMatch(search, find string) bool { + findLines := strings.Count(find, "\n") + 1 + searchLines := strings.Count(search, "\n") + 1 + limit := findLines + 3 + if findLines*2 > limit { + limit = findLines * 2 + } + if searchLines >= limit { + return true + } + if findLines == 1 { + return false + } + searchTrimmed := len(strings.TrimSpace(search)) + findTrimmed := len(strings.TrimSpace(find)) + byteLimit := findTrimmed + 500 + if findTrimmed*4 > byteLimit { + byteLimit = findTrimmed * 4 + } + return searchTrimmed > byteLimit +} + +// splitFindLines splits find into lines, dropping the trailing empty line a +// trailing newline produces so windows align with real content lines. +func splitFindLines(find string) []string { + lines := strings.Split(find, "\n") + if len(lines) > 0 && lines[len(lines)-1] == "" { + lines = lines[:len(lines)-1] + } + return lines +} + +// lineTrimmedReplacer matches when every line equals the corresponding +// old_string line after trimming surrounding whitespace. +func lineTrimmedReplacer(content, find string) []string { + contentLines := strings.Split(content, "\n") + findLines := splitFindLines(find) + if len(findLines) == 0 { + return nil + } + var candidates []string + for i := 0; i+len(findLines) <= len(contentLines); i++ { + matches := true + for j := range findLines { + if strings.TrimSpace(contentLines[i+j]) != strings.TrimSpace(findLines[j]) { + matches = false + break + } + } + if matches { + candidates = append(candidates, strings.Join(contentLines[i:i+len(findLines)], "\n")) + } + } + return candidates +} + +// blockAnchorReplacer anchors on the first and last lines (trimmed) and +// accepts the block when the middle lines average >= the similarity threshold +// (Levenshtein-based), tolerating a slightly misremembered interior. +func blockAnchorReplacer(content, find string) []string { + findLines := splitFindLines(find) + if len(findLines) < 3 { + return nil + } + contentLines := strings.Split(content, "\n") + firstAnchor := strings.TrimSpace(findLines[0]) + lastAnchor := strings.TrimSpace(findLines[len(findLines)-1]) + searchBlockSize := len(findLines) + maxLineDelta := searchBlockSize / 4 + if maxLineDelta < 1 { + maxLineDelta = 1 + } + + type span struct{ start, end int } + var candidates []span + for i := 0; i < len(contentLines); i++ { + if strings.TrimSpace(contentLines[i]) != firstAnchor { + continue + } + // Only the first occurrence of the last anchor after this start counts, + // mirroring the reference implementation: a farther-away closing line is + // assumed to close a different block. + for j := i + 2; j < len(contentLines); j++ { + if strings.TrimSpace(contentLines[j]) != lastAnchor { + continue + } + actualBlockSize := j - i + 1 + delta := actualBlockSize - searchBlockSize + if delta < 0 { + delta = -delta + } + if delta <= maxLineDelta { + candidates = append(candidates, span{start: i, end: j}) + } + break + } + } + if len(candidates) == 0 { + return nil + } + + middleSimilarity := func(candidate span) float64 { + actualBlockSize := candidate.end - candidate.start + 1 + linesToCheck := searchBlockSize - 2 + if actualBlockSize-2 < linesToCheck { + linesToCheck = actualBlockSize - 2 + } + if linesToCheck <= 0 { + return 1.0 + } + similarity := 0.0 + for j := 1; j < searchBlockSize-1 && j < actualBlockSize-1; j++ { + contentLine := strings.TrimSpace(contentLines[candidate.start+j]) + findLine := strings.TrimSpace(findLines[j]) + maxLen := len(contentLine) + if len(findLine) > maxLen { + maxLen = len(findLine) + } + if maxLen == 0 { + continue + } + distance := levenshtein(contentLine, findLine) + similarity += 1 - float64(distance)/float64(maxLen) + } + return similarity / float64(linesToCheck) + } + + // Return EVERY candidate that clears the similarity threshold, best first. + // Picking only the best would hide competing blocks from fuzzyEditMatch's + // distinct-candidate ambiguity check and could silently edit the wrong + // block; with all qualifying spans surfaced, one clear winner still + // resolves (single candidate) while two plausible blocks are rejected as + // ambiguous. replaceAll consumers take the first (most similar) span. + type scoredSpan struct { + span span + similarity float64 + } + var qualifying []scoredSpan + for _, candidate := range candidates { + if s := middleSimilarity(candidate); s >= editAnchorSimilarityThreshold { + qualifying = append(qualifying, scoredSpan{span: candidate, similarity: s}) + } + } + if len(qualifying) == 0 { + return nil + } + sort.SliceStable(qualifying, func(i, j int) bool { + return qualifying[i].similarity > qualifying[j].similarity + }) + spans := make([]string, 0, len(qualifying)) + for _, scored := range qualifying { + spans = append(spans, strings.Join(contentLines[scored.span.start:scored.span.end+1], "\n")) + } + return spans +} + +var editWhitespaceRun = regexp.MustCompile(`\s+`) + +func normalizeEditWhitespace(text string) string { + return strings.TrimSpace(editWhitespaceRun.ReplaceAllString(text, " ")) +} + +// whitespaceNormalizedReplacer matches after collapsing all whitespace runs to +// a single space: full lines, sub-line spans (via a word-boundary regex), and +// multi-line windows. +func whitespaceNormalizedReplacer(content, find string) []string { + normalizedFind := normalizeEditWhitespace(find) + if normalizedFind == "" { + return nil + } + contentLines := strings.Split(content, "\n") + var candidates []string + var subLinePattern *regexp.Regexp + for _, line := range contentLines { + normalizedLine := normalizeEditWhitespace(line) + if normalizedLine == normalizedFind { + candidates = append(candidates, line) + continue + } + if !strings.Contains(normalizedLine, normalizedFind) { + continue + } + if subLinePattern == nil { + words := strings.Fields(find) + quoted := make([]string, len(words)) + for i, word := range words { + quoted[i] = regexp.QuoteMeta(word) + } + pattern, err := regexp.Compile(strings.Join(quoted, `\s+`)) + if err != nil { + continue + } + subLinePattern = pattern + } + if match := subLinePattern.FindString(line); match != "" { + candidates = append(candidates, match) + } + } + + findLines := strings.Split(find, "\n") + if len(findLines) > 1 { + for i := 0; i+len(findLines) <= len(contentLines); i++ { + block := strings.Join(contentLines[i:i+len(findLines)], "\n") + if normalizeEditWhitespace(block) == normalizedFind { + candidates = append(candidates, block) + } + } + } + return candidates +} + +// stripCommonIndentation removes the minimum leading-whitespace width shared +// by all non-empty lines, so blocks match regardless of their nesting depth. +func stripCommonIndentation(text string) string { + lines := strings.Split(text, "\n") + minIndent := -1 + for _, line := range lines { + if strings.TrimSpace(line) == "" { + continue + } + indent := len(line) - len(strings.TrimLeft(line, " \t")) + if minIndent < 0 || indent < minIndent { + minIndent = indent + } + } + if minIndent <= 0 { + return text + } + stripped := make([]string, len(lines)) + for i, line := range lines { + if strings.TrimSpace(line) == "" { + stripped[i] = line + continue + } + stripped[i] = line[minIndent:] + } + return strings.Join(stripped, "\n") +} + +func indentationFlexibleReplacer(content, find string) []string { + normalizedFind := stripCommonIndentation(find) + contentLines := strings.Split(content, "\n") + findLines := strings.Split(find, "\n") + var candidates []string + for i := 0; i+len(findLines) <= len(contentLines); i++ { + block := strings.Join(contentLines[i:i+len(findLines)], "\n") + if stripCommonIndentation(block) == normalizedFind { + candidates = append(candidates, block) + } + } + return candidates +} + +var editEscapeSequence = regexp.MustCompile("\\\\(n|t|r|'|\"|`|\\\\|\\n|\\$)") + +// unescapeEditString undoes one level of string escaping (\n, \t, \", \\, a +// backslash-newline continuation, \$) — the model sometimes reproduces file +// content as it appeared inside a quoted string literal. +func unescapeEditString(text string) string { + return editEscapeSequence.ReplaceAllStringFunc(text, func(match string) string { + switch match[1:] { + case "n": + return "\n" + case "t": + return "\t" + case "r": + return "\r" + case "\n": + return "\n" + default: + // ', ", `, \, $ all unescape to themselves. + return match[1:] + } + }) +} + +func escapeNormalizedReplacer(content, find string) []string { + unescapedFind := unescapeEditString(find) + var candidates []string + if strings.Contains(content, unescapedFind) { + candidates = append(candidates, unescapedFind) + } + contentLines := strings.Split(content, "\n") + findLines := strings.Split(unescapedFind, "\n") + for i := 0; i+len(findLines) <= len(contentLines); i++ { + block := strings.Join(contentLines[i:i+len(findLines)], "\n") + if unescapeEditString(block) == unescapedFind { + candidates = append(candidates, block) + } + } + return candidates +} + +// trimmedBoundaryReplacer tolerates stray leading/trailing whitespace (often +// blank lines) around an otherwise-exact old_string. +func trimmedBoundaryReplacer(content, find string) []string { + trimmedFind := strings.TrimSpace(find) + if trimmedFind == find || trimmedFind == "" { + return nil + } + var candidates []string + if strings.Contains(content, trimmedFind) { + candidates = append(candidates, trimmedFind) + } + contentLines := strings.Split(content, "\n") + findLines := strings.Split(find, "\n") + for i := 0; i+len(findLines) <= len(contentLines); i++ { + block := strings.Join(contentLines[i:i+len(findLines)], "\n") + if strings.TrimSpace(block) == trimmedFind { + candidates = append(candidates, block) + } + } + return candidates +} + +// contextAwareReplacer anchors on the first and last lines and accepts an +// equal-length block when at least half of its non-empty middle lines match +// after trimming — a cheaper, stricter cousin of blockAnchorReplacer. +func contextAwareReplacer(content, find string) []string { + findLines := splitFindLines(find) + if len(findLines) < 3 { + return nil + } + contentLines := strings.Split(content, "\n") + firstAnchor := strings.TrimSpace(findLines[0]) + lastAnchor := strings.TrimSpace(findLines[len(findLines)-1]) + var candidates []string + for i := 0; i < len(contentLines); i++ { + if strings.TrimSpace(contentLines[i]) != firstAnchor { + continue + } + for j := i + 2; j < len(contentLines); j++ { + if strings.TrimSpace(contentLines[j]) != lastAnchor { + continue + } + if j-i+1 == len(findLines) { + matching, nonEmpty := 0, 0 + for k := 1; k < len(findLines)-1; k++ { + blockLine := strings.TrimSpace(contentLines[i+k]) + findLine := strings.TrimSpace(findLines[k]) + if blockLine == "" && findLine == "" { + continue + } + nonEmpty++ + if blockLine == findLine { + matching++ + } + } + if nonEmpty == 0 || float64(matching)/float64(nonEmpty) >= 0.5 { + candidates = append(candidates, strings.Join(contentLines[i:j+1], "\n")) + } + } + break + } + } + return candidates +} + +// adaptReplacementToSpan re-shapes the model's replacement to the span a +// tolerant matcher resolved. When old_string only matched after normalization, +// new_string was written at old_string's (wrong) shape, so applying it raw +// would strip the file's indentation or drop a trailing CR: +// +// 1. Uniform re-indent: when every span line equals delta + the corresponding +// find line's indentation (the line-trimmed / indentation-flexible shapes), +// the same delta is prepended to every non-blank replacement line. Any +// line that breaks the uniform-delta relationship disables the shift — +// block-anchor matches with a drifted interior are left untouched. +// 2. Trailing CR: a span from a CRLF file ends mid-line at "\r" (candidates +// are built by joining lines split on "\n"); the replacement gets the same +// trailing "\r" so the file's CRLF pairs stay intact. +func adaptReplacementToSpan(span, find, replacement string) string { + if delta, ok := uniformIndentDelta(span, find); ok && delta != "" { + lines := strings.Split(replacement, "\n") + for i, line := range lines { + if strings.TrimSpace(line) == "" { + continue + } + lines[i] = delta + line + } + replacement = strings.Join(lines, "\n") + } + if strings.HasSuffix(span, "\r") && !strings.HasSuffix(replacement, "\r") { + replacement += "\r" + } + return replacement +} + +// uniformIndentDelta returns the indentation prefix that, prepended to every +// non-blank find line, yields the corresponding span line's indentation. ok is +// false when line counts differ, any line pair disagrees on the delta, or the +// span is not simply a uniformly deeper-indented copy of find. +func uniformIndentDelta(span, find string) (string, bool) { + spanLines := strings.Split(span, "\n") + findLines := splitFindLines(find) + if len(spanLines) != len(findLines) { + return "", false + } + leadingWhitespace := func(line string) string { + return line[:len(line)-len(strings.TrimLeft(line, " \t"))] + } + delta := "" + haveDelta := false + for i := range findLines { + spanLine := strings.TrimSuffix(spanLines[i], "\r") + findLine := strings.TrimSuffix(findLines[i], "\r") + if strings.TrimSpace(spanLine) == "" && strings.TrimSpace(findLine) == "" { + continue + } + spanIndent := leadingWhitespace(spanLine) + findIndent := leadingWhitespace(findLine) + if !strings.HasSuffix(spanIndent, findIndent) { + return "", false + } + lineDelta := spanIndent[:len(spanIndent)-len(findIndent)] + if !haveDelta { + delta = lineDelta + haveDelta = true + continue + } + if lineDelta != delta { + return "", false + } + } + return delta, haveDelta +} + +// levenshtein computes edit distance with a two-row rolling matrix. +func levenshtein(a, b string) int { + if a == "" { + return len(b) + } + if b == "" { + return len(a) + } + previous := make([]int, len(b)+1) + current := make([]int, len(b)+1) + for j := range previous { + previous[j] = j + } + for i := 1; i <= len(a); i++ { + current[0] = i + for j := 1; j <= len(b); j++ { + cost := 1 + if a[i-1] == b[j-1] { + cost = 0 + } + minimum := previous[j] + 1 + if current[j-1]+1 < minimum { + minimum = current[j-1] + 1 + } + if previous[j-1]+cost < minimum { + minimum = previous[j-1] + cost + } + current[j] = minimum + } + previous, current = current, previous + } + return previous[len(b)] +} diff --git a/internal/tools/edit_replacers_test.go b/internal/tools/edit_replacers_test.go new file mode 100644 index 000000000..b82a34cf0 --- /dev/null +++ b/internal/tools/edit_replacers_test.go @@ -0,0 +1,369 @@ +package tools + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" +) + +func runEdit(t *testing.T, dir, initial string, args map[string]any) (Result, string) { + t.Helper() + path := filepath.Join(dir, "target.txt") + if err := os.WriteFile(path, []byte(initial), 0o644); err != nil { + t.Fatal(err) + } + if args["path"] == nil { + args["path"] = "target.txt" + } + result := NewEditFileTool(dir).Run(context.Background(), args) + after, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + return result, string(after) +} + +func TestEditExactMatchStillPreferred(t *testing.T) { + // The fast path must keep byte-exact semantics: an exact match wins even + // when fuzzy strategies would also find candidates elsewhere. + result, after := runEdit(t, t.TempDir(), "alpha\nbeta\ngamma\n", map[string]any{ + "old_string": "beta", + "new_string": "delta", + }) + if result.Status != StatusOK { + t.Fatalf("expected ok, got %q", result.Output) + } + if after != "alpha\ndelta\ngamma\n" { + t.Fatalf("unexpected content: %q", after) + } +} + +func TestEditFuzzyLineTrimmed(t *testing.T) { + // Model reproduced the lines without the file's leading indentation; the + // line-trimmed strategy must find the real indented span and preserve the + // file's own indentation on the replaced region boundary. + initial := "func a() {\n\tx := 1\n\treturn x\n}\n" + result, after := runEdit(t, t.TempDir(), initial, map[string]any{ + "old_string": "x := 1\nreturn x", + "new_string": "y := 2\nreturn y", + }) + if result.Status != StatusOK { + t.Fatalf("expected ok, got %q", result.Output) + } + // new_string was written at old_string's outdented shape; the span's + // indentation must be re-applied so the file keeps its tabs. + if !strings.Contains(after, "\ty := 2\n\treturn y\n}") { + t.Fatalf("unexpected content: %q", after) + } + if strings.Contains(after, "x := 1") { + t.Fatalf("old span survived: %q", after) + } +} + +func TestEditFuzzyWhitespaceNormalizedSingleLine(t *testing.T) { + // Whitespace runs collapsed: "x := 1" in the file, single spaces from + // the model. + initial := "start\nx := 1\nend\n" + result, after := runEdit(t, t.TempDir(), initial, map[string]any{ + "old_string": "x := 1", + "new_string": "x := 2", + }) + if result.Status != StatusOK { + t.Fatalf("expected ok, got %q", result.Output) + } + if !strings.Contains(after, "x := 2") || strings.Contains(after, ":= 1") { + t.Fatalf("unexpected content: %q", after) + } +} + +func TestEditFuzzyBlockAnchorToleratesMiddleDrift(t *testing.T) { + // First and last lines anchor the block; one interior line differs slightly + // (comment text drifted). Levenshtein similarity keeps it above 0.65. + initial := strings.Join([]string{ + "func handler(w http.ResponseWriter, r *http.Request) {", + "\t// write the response body to the client", + "\tw.WriteHeader(http.StatusOK)", + "\tfmt.Fprint(w, \"done\")", + "}", + "", + }, "\n") + find := strings.Join([]string{ + "func handler(w http.ResponseWriter, r *http.Request) {", + "\t// write the response body to client", + "\tw.WriteHeader(http.StatusOK)", + "\tfmt.Fprint(w, \"done\")", + "}", + }, "\n") + result, after := runEdit(t, t.TempDir(), initial, map[string]any{ + "old_string": find, + "new_string": "func handler(w http.ResponseWriter, r *http.Request) {\n\tw.WriteHeader(http.StatusNoContent)\n}", + }) + if result.Status != StatusOK { + t.Fatalf("expected ok, got %q", result.Output) + } + if !strings.Contains(after, "StatusNoContent") || strings.Contains(after, "StatusOK") { + t.Fatalf("unexpected content: %q", after) + } +} + +func TestEditFuzzyIndentationFlexible(t *testing.T) { + // The whole block sits one nesting level deeper in the file than in the + // model's old_string, with interior relative indentation preserved. + initial := "if ok {\n\t\tfor i := range xs {\n\t\t\tsum += xs[i]\n\t\t}\n}\n" + result, after := runEdit(t, t.TempDir(), initial, map[string]any{ + "old_string": "for i := range xs {\n\tsum += xs[i]\n}", + "new_string": "sum += total(xs)", + }) + if result.Status != StatusOK { + t.Fatalf("expected ok, got %q", result.Output) + } + // The replacement lands at the span's real depth (two tabs), not at + // old_string's outdented depth. + if !strings.Contains(after, "\t\tsum += total(xs)") || strings.Contains(after, "range xs") { + t.Fatalf("unexpected content: %q", after) + } +} + +func TestEditFuzzyCRLFLineTrimmedPreservesCarriageReturns(t *testing.T) { + // CRLF file + a multi-line outdented old_string: the resolved span carries + // tabs and trailing CRs; the replacement must inherit both so the file's + // CRLF pairs and indentation survive the edit. + initial := "func a() {\r\n\tx := 1\r\n\treturn x\r\n}\r\n" + result, after := runEdit(t, t.TempDir(), initial, map[string]any{ + "old_string": "x := 1\nreturn x", + "new_string": "y := 2\nreturn y", + }) + if result.Status != StatusOK { + t.Fatalf("expected ok, got %q", result.Output) + } + if !strings.Contains(after, "\ty := 2\r\n\treturn y\r\n}") { + t.Fatalf("unexpected content: %q", after) + } + if strings.Contains(after, "\n\treturn y\n") { + t.Fatalf("bare LF introduced into CRLF file: %q", after) + } +} + +func TestUniformIndentDelta(t *testing.T) { + cases := []struct { + span, find string + want string + ok bool + }{ + {"\tx := 1\n\treturn x", "x := 1\nreturn x", "\t", true}, + {"\t\tfor {\n\t\t\tgo()\n\t\t}", "for {\n\tgo()\n}", "\t\t", true}, + {"x := 1", "x := 1", "", true}, // no shift needed + {"\tx := 1\n return x", "x := 1\nreturn x", "", false}, // mixed delta + {"\ta\n\tb\n\tc", "a\nb", "", false}, // line-count mismatch + {" a", "\ta", "", false}, // find indent not a suffix + } + for _, c := range cases { + got, ok := uniformIndentDelta(c.span, c.find) + if ok != c.ok || (ok && got != c.want) { + t.Errorf("uniformIndentDelta(%q, %q) = %q,%v want %q,%v", c.span, c.find, got, ok, c.want, c.ok) + } + } +} + +func TestAdaptReplacementLeavesNonUniformSpansAlone(t *testing.T) { + // Block-anchor style match with a drifted interior: no uniform delta, so + // the replacement must pass through untouched. + span := "\tfunc h() {\n\t\t// drifted comment\n\t}" + find := "func h() {\n// different comment\n}" + if got := adaptReplacementToSpan(span, find, "replacement()"); got != "replacement()" { + t.Fatalf("non-uniform span must not re-indent: %q", got) + } +} + +func TestEditFuzzyEscapeNormalized(t *testing.T) { + // Model reproduced the line as it would appear inside a quoted string + // literal, with escaped quotes. + initial := "console.log(\"hello world\")\n" + result, after := runEdit(t, t.TempDir(), initial, map[string]any{ + "old_string": `console.log(\"hello world\")`, + "new_string": `console.log("goodbye")`, + }) + if result.Status != StatusOK { + t.Fatalf("expected ok, got %q", result.Output) + } + if !strings.Contains(after, `console.log("goodbye")`) { + t.Fatalf("unexpected content: %q", after) + } +} + +func TestEditFuzzyTrimmedBoundary(t *testing.T) { + // Stray whitespace around an otherwise-exact old_string; " two " is not a + // substring of the file, so only the fuzzy cascade can resolve it. + initial := "one\ntwo\nthree\n" + result, after := runEdit(t, t.TempDir(), initial, map[string]any{ + "old_string": " two ", + "new_string": "2", + }) + if result.Status != StatusOK { + t.Fatalf("expected ok, got %q", result.Output) + } + if after != "one\n2\nthree\n" { + t.Fatalf("unexpected content: %q", after) + } +} + +func TestTrimmedBoundaryReplacerYieldsTrimmedSpan(t *testing.T) { + candidates := trimmedBoundaryReplacer("one\ntwo\nthree\n", "\ntwo\n") + if len(candidates) == 0 || candidates[0] != "two" { + t.Fatalf("expected trimmed candidate \"two\", got %v", candidates) + } + if trimmedBoundaryReplacer("anything", "already-trimmed") != nil { + t.Fatal("already-trimmed find must yield no candidates") + } +} + +// Two identical blocks that only match old_string after indentation-tolerant +// matching: exact search finds nothing, fuzzy finds both. +const ambiguousBlocks = "\tif x {\n\t\tgo()\n\t}\nmid\n\tif x {\n\t\tgo()\n\t}\n" + +func TestEditFuzzyAmbiguousReportsError(t *testing.T) { + // No strategy can disambiguate identical blocks, so the tool must refuse + // rather than guess, and the file must be untouched. + result, after := runEdit(t, t.TempDir(), ambiguousBlocks, map[string]any{ + "old_string": "if x {\n\tgo()\n}", + "new_string": "stop()", + }) + if result.Status != StatusError || !strings.Contains(result.Output, "multiple locations") { + t.Fatalf("expected ambiguity error, got %q", result.Output) + } + if after != ambiguousBlocks { + t.Fatalf("file must be unchanged, got %q", after) + } +} + +func TestEditFuzzyReplaceAll(t *testing.T) { + // replace_all applies the fuzzy-resolved span at every occurrence. + result, after := runEdit(t, t.TempDir(), ambiguousBlocks, map[string]any{ + "old_string": "if x {\n\tgo()\n}", + "new_string": "stop()", + "replace_all": true, + }) + if result.Status != StatusOK { + t.Fatalf("expected ok, got %q", result.Output) + } + if strings.Count(after, "stop()") != 2 || strings.Contains(after, "go()") { + t.Fatalf("unexpected content: %q", after) + } +} + +func TestEditFuzzyNotFoundKeepsExactError(t *testing.T) { + result, after := runEdit(t, t.TempDir(), "alpha\n", map[string]any{ + "old_string": "omega", + "new_string": "x", + }) + if result.Status != StatusError || !strings.Contains(result.Output, "Could not find the exact string") { + t.Fatalf("expected not-found error, got %q", result.Output) + } + if after != "alpha\n" { + t.Fatalf("file must be unchanged, got %q", after) + } +} + +func TestEditFuzzyCRLFFile(t *testing.T) { + // CRLF file + LF old_string with indentation drift: the CRLF translation + // feeds the cascade and the replacement preserves CRLF endings. + initial := "func a() {\r\n\tx := 1\r\n}\r\n" + result, after := runEdit(t, t.TempDir(), initial, map[string]any{ + "old_string": "x := 1", + "new_string": "x := 2", + }) + if result.Status != StatusOK { + t.Fatalf("expected ok, got %q", result.Output) + } + if !strings.Contains(after, "\tx := 2\r\n") { + t.Fatalf("unexpected content: %q", after) + } +} + +func TestIsDisproportionateEditMatch(t *testing.T) { + // A candidate span that dwarfs old_string must be refused: anchors bridging + // unrelated code would otherwise delete it all. + big := strings.Repeat("line\n", 10) + if !isDisproportionateEditMatch(big, "a\nb\nc") { + t.Fatal("10-line span for 3-line find must be disproportionate") + } + if isDisproportionateEditMatch("a\nb\nc\nd", "a\nb\nc") { + t.Fatal("one extra line within tolerance must be allowed") + } + if isDisproportionateEditMatch("single line", "single") { + t.Fatal("single-line finds are exempt from the byte-length guard") + } + if !isDisproportionateEditMatch("xx\n"+strings.Repeat("y", 900)+"\nzz", "xx\nab\nzz") { + t.Fatal("byte-length blowup must be disproportionate") + } +} + +func TestLevenshtein(t *testing.T) { + cases := []struct { + a, b string + want int + }{ + {"", "", 0}, + {"abc", "", 3}, + {"", "abc", 3}, + {"kitten", "sitting", 3}, + {"same", "same", 0}, + } + for _, c := range cases { + if got := levenshtein(c.a, c.b); got != c.want { + t.Fatalf("levenshtein(%q,%q) = %d, want %d", c.a, c.b, got, c.want) + } + } +} + +func TestEditFuzzyDistinctCandidatesAreAmbiguous(t *testing.T) { + // Two same-content blocks at DIFFERENT indentation: each resolved span is + // distinct and occurs exactly once, so the old literal-uniqueness check + // passed and silently edited the first block. Distinct-shaped candidates + // from one strategy must instead report ambiguity and leave the file alone. + initial := "\tif x {\n\t\tgo()\n\t}\nmid\n\t\tif x {\n\t\t\tgo()\n\t\t}\n" + result, after := runEdit(t, t.TempDir(), initial, map[string]any{ + "old_string": "if x {\n\tgo()\n}", + "new_string": "stop()", + }) + if result.Status != StatusError || !strings.Contains(result.Output, "multiple locations") { + t.Fatalf("distinct fuzzy candidates must be ambiguous, got %q", result.Output) + } + if after != initial { + t.Fatalf("file must be unchanged, got %q", after) + } +} + +func TestEditFuzzyBlockAnchorTwoPlausibleBlocksAmbiguous(t *testing.T) { + // Two blocks share the same first/last anchor lines and BOTH interiors sit + // above the similarity threshold. Picking the "best" one would silently + // edit a block the model may not have meant; both must surface so the + // cascade reports ambiguity and the file stays untouched. + initial := strings.Join([]string{ + "func setup(cfg Config) {", + "\tvalue := compute(alpha, beta, gamma)", + "}", + "", + "func setup(cfg Config) {", + "\tvalue := compute(alpha, delta, gamma)", + "}", + "", + }, "\n") + find := strings.Join([]string{ + "func setup(cfg Config) {", + "\tvalue := compute(alpha, omega, gamma)", + "}", + }, "\n") + result, after := runEdit(t, t.TempDir(), initial, map[string]any{ + "old_string": find, + "new_string": "func setup(cfg Config) {}", + }) + if result.Status != StatusError || !strings.Contains(result.Output, "multiple locations") { + t.Fatalf("two plausible anchor blocks must be ambiguous, got %q", result.Output) + } + if after != initial { + t.Fatalf("file must be unchanged, got %q", after) + } +} diff --git a/internal/tools/exec_command.go b/internal/tools/exec_command.go index cfc4082fb..0feed7e1a 100644 --- a/internal/tools/exec_command.go +++ b/internal/tools/exec_command.go @@ -908,6 +908,15 @@ func formatExecCommandOutput(output string, sessionID int, exited bool, exitCode } func truncateExecOutput(output string, maxOutputTokens int) (string, bool) { + return truncateExecOutputSpill(output, maxOutputTokens, "exec_command") +} + +// truncateExecOutputSpill keeps a head/tail window of the output within the +// token budget and, on truncation, spills the full output to disk so the model +// can grep/read the elided middle instead of re-running the command with a +// bigger budget. The spill is best-effort: when it fails the notice simply +// omits the file hint. +func truncateExecOutputSpill(output string, maxOutputTokens int, toolName string) (string, bool) { if maxOutputTokens <= 0 { maxOutputTokens = defaultMaxOutputTokens } @@ -915,9 +924,13 @@ func truncateExecOutput(output string, maxOutputTokens int) (string, bool) { if len(output) <= maxBytes { return output, false } + notice := "\n[zero] output truncated\n" + if spillPath := spillTruncatedOutput(toolName, output); spillPath != "" { + notice = "\n[zero] output truncated — full output saved to " + spillPath + " (grep or read_file it instead of re-running)\n" + } head := maxBytes / 2 tail := maxBytes - head - return utf8Prefix(output, head) + "\n[zero] output truncated\n" + utf8Suffix(output, tail), true + return utf8Prefix(output, head) + notice + utf8Suffix(output, tail), true } func utf8Prefix(value string, maxBytes int) string { diff --git a/internal/tools/format_on_write.go b/internal/tools/format_on_write.go new file mode 100644 index 000000000..12e842906 --- /dev/null +++ b/internal/tools/format_on_write.go @@ -0,0 +1,97 @@ +package tools + +import ( + "context" + "os" + "os/exec" + "path/filepath" + "strings" + "time" +) + +// Format-on-write for the mutating file tools. When enabled, a successful +// edit_file/write_file runs the language's standard formatter on the file it +// just wrote, so the model's output always lands in project-canonical style +// and never fails a CI format check it cannot see. Off by default (set +// ZERO_FORMAT_ON_WRITE=1): auto-reformatting changes bytes the model did not +// write, which strict workflows may not want. +// +// Ordering matters: formatting runs BEFORE the FileTracker re-baseline, and +// the caller records the POST-format content. Formatting after the baseline +// would make the very next edit look like an external modification and trip +// the conflict guard. + +// formatOnWriteTimeout bounds one formatter run; a wedged formatter must never +// hang a tool call. On timeout the unformatted write stands. +const formatOnWriteTimeout = 10 * time.Second + +// formatterCommands maps a file extension to the formatter argv; the file path +// is appended as the final argument. Only in-place, config-respecting, +// community-standard formatters — a missing binary silently skips formatting. +var formatterCommands = map[string][]string{ + ".go": {"gofmt", "-w"}, + ".rs": {"rustfmt"}, + ".py": {"ruff", "format", "--quiet"}, + ".ts": {"prettier", "--log-level", "silent", "--write"}, + ".tsx": {"prettier", "--log-level", "silent", "--write"}, + ".js": {"prettier", "--log-level", "silent", "--write"}, + ".jsx": {"prettier", "--log-level", "silent", "--write"}, + ".json": {"prettier", "--log-level", "silent", "--write"}, + ".css": {"prettier", "--log-level", "silent", "--write"}, + ".scss": {"prettier", "--log-level", "silent", "--write"}, + ".html": {"prettier", "--log-level", "silent", "--write"}, + ".md": {"prettier", "--log-level", "silent", "--write"}, + ".yaml": {"prettier", "--log-level", "silent", "--write"}, + ".yml": {"prettier", "--log-level", "silent", "--write"}, + ".zig": {"zig", "fmt"}, + ".dart": {"dart", "format"}, + ".tf": {"terraform", "fmt"}, + ".gleam": {"gleam", "format"}, + ".sh": {"shfmt", "-w"}, + ".bash": {"shfmt", "-w"}, + ".c": {"clang-format", "-i"}, + ".h": {"clang-format", "-i"}, + ".cpp": {"clang-format", "-i"}, + ".hpp": {"clang-format", "-i"}, + ".cc": {"clang-format", "-i"}, + ".kt": {"ktlint", "-F"}, + ".swift": {"swiftformat"}, + ".lua": {"stylua"}, +} + +// formatOnWriteEnabled reports whether the opt-in env toggle is set. +func formatOnWriteEnabled() bool { + value := strings.TrimSpace(os.Getenv("ZERO_FORMAT_ON_WRITE")) + return value != "" && value != "0" && !strings.EqualFold(value, "false") +} + +// maybeFormatWrittenFile runs the configured formatter for absolutePath (when +// enabled and on PATH) and returns the file's content afterwards. Best-effort +// throughout: any failure — no formatter, formatter error, timeout, unreadable +// result — returns writtenContent so the caller's state matches the last write +// it performed itself. +func maybeFormatWrittenFile(ctx context.Context, absolutePath string, writtenContent string) string { + if !formatOnWriteEnabled() { + return writtenContent + } + command, ok := formatterCommands[strings.ToLower(filepath.Ext(absolutePath))] + if !ok { + return writtenContent + } + if _, err := exec.LookPath(command[0]); err != nil { + return writtenContent + } + formatCtx, cancel := context.WithTimeout(ctx, formatOnWriteTimeout) + defer cancel() + arguments := append(append([]string(nil), command[1:]...), absolutePath) + formatter := exec.CommandContext(formatCtx, command[0], arguments...) + formatter.Dir = filepath.Dir(absolutePath) + if err := formatter.Run(); err != nil { + return writtenContent + } + formatted, err := os.ReadFile(absolutePath) + if err != nil { + return writtenContent + } + return string(formatted) +} diff --git a/internal/tools/format_on_write_test.go b/internal/tools/format_on_write_test.go new file mode 100644 index 000000000..76a0960e2 --- /dev/null +++ b/internal/tools/format_on_write_test.go @@ -0,0 +1,82 @@ +package tools + +import ( + "context" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" +) + +// gofmt ships with the Go toolchain, so it is the one formatter guaranteed to +// exist wherever these tests run. +func requireGofmt(t *testing.T) { + t.Helper() + if _, err := exec.LookPath("gofmt"); err != nil { + t.Skip("gofmt not on PATH") + } +} + +func TestFormatOnWriteDisabledByDefault(t *testing.T) { + requireGofmt(t) + t.Setenv("ZERO_FORMAT_ON_WRITE", "") + dir := t.TempDir() + ugly := "package a\n\nfunc A( ) { }\n" + + result := NewWriteFileTool(dir).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{ + "path": "a.go", + "content": ugly, + }, RunOptions{}) + if result.Status != StatusOK { + t.Fatalf("write failed: %q", result.Output) + } + onDisk, err := os.ReadFile(filepath.Join(dir, "a.go")) + if err != nil { + t.Fatal(err) + } + if string(onDisk) != ugly { + t.Fatalf("formatting must be off by default, got %q", onDisk) + } +} + +func TestFormatOnWriteFormatsAndKeepsTrackerConsistent(t *testing.T) { + requireGofmt(t) + t.Setenv("ZERO_FORMAT_ON_WRITE", "1") + dir := t.TempDir() + tracker := NewFileTracker() + + write := NewWriteFileTool(dir).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{ + "path": "a.go", + "content": "package a\n\nfunc A( ) { }\n", + }, RunOptions{FileTracker: tracker}) + if write.Status != StatusOK { + t.Fatalf("write failed: %q", write.Output) + } + onDisk, err := os.ReadFile(filepath.Join(dir, "a.go")) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(onDisk), "func A() {") { + t.Fatalf("expected gofmt-formatted content, got %q", onDisk) + } + + // The tracker must have been re-baselined to the POST-format content: a + // follow-up edit must not trip the external-modification conflict guard. + edit := NewEditFileTool(dir).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{ + "path": "a.go", + "old_string": "func A() {", + "new_string": "func B() {", + }, RunOptions{FileTracker: tracker}) + if edit.Status != StatusOK { + t.Fatalf("follow-up edit must not conflict after formatting: %q", edit.Output) + } +} + +func TestFormatOnWriteSkipsUnknownExtensions(t *testing.T) { + t.Setenv("ZERO_FORMAT_ON_WRITE", "1") + content := maybeFormatWrittenFile(context.Background(), filepath.Join(t.TempDir(), "notes.xyz"), "raw text") + if content != "raw text" { + t.Fatalf("unknown extension must pass through: %q", content) + } +} diff --git a/internal/tools/inline_diagnostics.go b/internal/tools/inline_diagnostics.go new file mode 100644 index 000000000..83495eda3 --- /dev/null +++ b/internal/tools/inline_diagnostics.go @@ -0,0 +1,18 @@ +package tools + +import "context" + +// inlineDiagnostics renders the post-write diagnostics block a mutating tool +// appends to its output. Empty when no Diagnostics callback is wired, the file +// is clean, or no language server is available — the tool output is then +// byte-identical to the pre-diagnostics behavior. +func inlineDiagnostics(ctx context.Context, options RunOptions, absolutePath, relativePath string) string { + if options.Diagnostics == nil { + return "" + } + block := options.Diagnostics(ctx, absolutePath) + if block == "" { + return "" + } + return "\n\nDiagnostics in " + relativePath + " after this change (fix any errors you introduced):\n" + block +} diff --git a/internal/tools/inline_diagnostics_test.go b/internal/tools/inline_diagnostics_test.go new file mode 100644 index 000000000..caa73f8c3 --- /dev/null +++ b/internal/tools/inline_diagnostics_test.go @@ -0,0 +1,53 @@ +package tools + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" +) + +// A wired Diagnostics callback must surface its block in the tool output for +// both mutating tools, and a clean file (empty block) must leave the output +// byte-identical to the pre-diagnostics behavior. +func TestMutatingToolsAppendInlineDiagnostics(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "a.go"), []byte("package a\n"), 0o644); err != nil { + t.Fatal(err) + } + diagnostics := func(_ context.Context, absPath string) string { + return absPath + ":1:1: error: undefined: x" + } + + editResult := NewEditFileTool(dir).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{ + "path": "a.go", + "old_string": "package a", + "new_string": "package b", + }, RunOptions{Diagnostics: diagnostics}) + if editResult.Status != StatusOK { + t.Fatalf("edit failed: %q", editResult.Output) + } + if !strings.Contains(editResult.Output, "Diagnostics in a.go after this change") || !strings.Contains(editResult.Output, "undefined: x") { + t.Fatalf("edit output missing diagnostics block: %q", editResult.Output) + } + + writeResult := NewWriteFileTool(dir).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{ + "path": "b.go", + "content": "package b\n", + }, RunOptions{Diagnostics: diagnostics}) + if writeResult.Status != StatusOK { + t.Fatalf("write failed: %q", writeResult.Output) + } + if !strings.Contains(writeResult.Output, "Diagnostics in b.go after this change") { + t.Fatalf("write output missing diagnostics block: %q", writeResult.Output) + } + + clean := NewWriteFileTool(dir).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{ + "path": "c.go", + "content": "package c\n", + }, RunOptions{Diagnostics: func(context.Context, string) string { return "" }}) + if strings.Contains(clean.Output, "Diagnostics") { + t.Fatalf("clean file must not gain a diagnostics block: %q", clean.Output) + } +} diff --git a/internal/tools/registry.go b/internal/tools/registry.go index d1b9de856..0dcd8a35b 100644 --- a/internal/tools/registry.go +++ b/internal/tools/registry.go @@ -40,6 +40,12 @@ type RunOptions struct { // tool-call progress in the specialist card. nil is a no-op (the default // for every non-Task tool). Progress func(streamjson.Event) + // Diagnostics, when set, returns a formatted language-diagnostics block for + // a file a mutating tool just wrote ("" when clean or no server available). + // edit_file/write_file append it to their output so the model sees an error + // it introduced in the same turn instead of waiting for a later verification + // pass. nil disables inline diagnostics. + Diagnostics func(ctx context.Context, absPath string) string } type sandboxAwareTool interface { diff --git a/internal/tools/spill.go b/internal/tools/spill.go new file mode 100644 index 000000000..f7f49092d --- /dev/null +++ b/internal/tools/spill.go @@ -0,0 +1,141 @@ +package tools + +import ( + "fmt" + "os" + "path/filepath" + "strings" + "time" + + "github.com/Gitlawb/zero/internal/redaction" +) + +// Spill-to-disk for truncated tool output. When a command produces more than +// its context budget, the tail was previously simply gone — the model's only +// recourse was re-running the (possibly expensive or non-idempotent) command +// with a bigger budget. Instead, the full output is written to a cache file +// whose path is included in the truncation notice, so the model can grep or +// read_file the remainder. Spilling is best-effort: on any error truncation +// behaves exactly as before, just without the file hint. + +// spillRetention is how long spilled outputs are kept; files older than this +// are opportunistically removed whenever a new spill happens. +const spillRetention = 7 * 24 * time.Hour + +// spillRootPath returns the spill directory path without creating or checking +// it. Used by the read-path resolver to recognize spill files as readable. +func spillRootPath() string { + name := "zero-tool-output" + if uid := os.Getuid(); uid >= 0 { + name = fmt.Sprintf("zero-tool-output-%d", uid) + } + return filepath.Join(os.TempDir(), name) +} + +// resolveSpillReadPath reports whether requestedPath is a file inside the +// spill directory and, if so, returns its verified absolute path. Spill files +// are zero-created, per-uid owned, and already secret-redacted, so letting the +// scoped read tools open them is what makes the truncation notice's +// "read_file it" recovery actually work. Symlinks are resolved and the result +// must still be inside the spill dir, so a planted link cannot smuggle an +// out-of-scope file through this gate. +func resolveSpillReadPath(requestedPath string) (string, bool) { + if requestedPath == "" || !filepath.IsAbs(requestedPath) { + return "", false + } + root := spillRootPath() + cleaned := filepath.Clean(requestedPath) + if cleaned == root || !strings.HasPrefix(cleaned, root+string(filepath.Separator)) { + return "", false + } + resolved, err := filepath.EvalSymlinks(cleaned) + if err != nil { + return "", false + } + resolvedRoot, err := filepath.EvalSymlinks(root) + if err != nil { + return "", false + } + if !strings.HasPrefix(resolved, resolvedRoot+string(filepath.Separator)) { + return "", false + } + return resolved, true +} + +// spillDir returns the per-user spill directory, creating it on first use. +// Hardened for shared temp dirs (Linux /tmp): the name carries the uid so +// users cannot collide, and a pre-existing path is only accepted when it is a +// real directory (not a symlink that would redirect spills) owned by the +// current user. Any doubt fails the spill — it is best-effort anyway. +func spillDir() (string, error) { + dir := spillRootPath() + if err := os.MkdirAll(dir, 0o700); err != nil { + return "", err + } + // MkdirAll follows symlinks and leaves an existing directory untouched, so + // verify what is actually at the path. + info, err := os.Lstat(dir) + if err != nil { + return "", err + } + if !info.IsDir() { + return "", fmt.Errorf("spill path %s is not a directory", dir) + } + if err := checkSpillDirOwner(info); err != nil { + return "", err + } + return dir, nil +} + +// spillTruncatedOutput writes the full pre-truncation output to the spill +// directory and returns the file path, or "" when spilling fails. Output is +// scrubbed with the same configured-key redaction the registry applies at the +// tool boundary, so a spilled file never holds a secret the transcript would +// have hidden. +func spillTruncatedOutput(toolName, output string) string { + dir, err := spillDir() + if err != nil { + return "" + } + sweepSpillDir(dir) + prefix := strings.Map(func(r rune) rune { + if r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '_' || r == '-' { + return r + } + return '-' + }, toolName) + file, err := os.CreateTemp(dir, prefix+"-*.txt") + if err != nil { + return "" + } + defer file.Close() + scrubbed := redaction.RedactString(output, redaction.Options{}) + if _, err := file.WriteString(scrubbed); err != nil { + _ = os.Remove(file.Name()) + return "" + } + return file.Name() +} + +// sweepSpillDir removes spill files older than spillRetention. Best-effort: +// errors are ignored, the directory is small, and a sweep runs only when a new +// spill is about to happen anyway. +func sweepSpillDir(dir string) { + entries, err := os.ReadDir(dir) + if err != nil { + return + } + cutoff := time.Now().Add(-spillRetention) + for _, entry := range entries { + if entry.IsDir() { + continue + } + info, err := entry.Info() + if err != nil { + continue + } + if info.ModTime().Before(cutoff) { + _ = os.Remove(filepath.Join(dir, entry.Name())) + } + } +} diff --git a/internal/tools/spill_hardening_test.go b/internal/tools/spill_hardening_test.go new file mode 100644 index 000000000..77042288d --- /dev/null +++ b/internal/tools/spill_hardening_test.go @@ -0,0 +1,91 @@ +package tools + +import ( + "context" + "fmt" + "os" + "path/filepath" + "runtime" + "strings" + "testing" +) + +func spillDirName() string { + if uid := os.Getuid(); uid >= 0 { + return fmt.Sprintf("zero-tool-output-%d", uid) + } + return "zero-tool-output" +} + +func TestSpillDirRejectsSymlink(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("symlink creation needs privileges on Windows") + } + tmp := t.TempDir() + t.Setenv("TMPDIR", tmp) + elsewhere := filepath.Join(tmp, "elsewhere") + if err := os.Mkdir(elsewhere, 0o700); err != nil { + t.Fatal(err) + } + // An attacker pre-creates the spill path as a symlink to redirect spills. + if err := os.Symlink(elsewhere, filepath.Join(tmp, spillDirName())); err != nil { + t.Fatal(err) + } + if path := spillTruncatedOutput("bash", "sensitive output"); path != "" { + t.Fatalf("spill must refuse a symlinked directory, wrote %s", path) + } + entries, _ := os.ReadDir(elsewhere) + if len(entries) != 0 { + t.Fatalf("nothing may land behind the symlink: %v", entries) + } +} + +func TestSpillDirAcceptsOwnedDirectory(t *testing.T) { + t.Setenv("TMPDIR", t.TempDir()) + if path := spillTruncatedOutput("bash", "ok"); path == "" { + t.Fatal("spill must work in a clean per-user temp dir") + } +} + +func TestResolveSpillReadPathContainment(t *testing.T) { + t.Setenv("TMPDIR", t.TempDir()) + spillPath := spillTruncatedOutput("exec_command", "spilled body") + if spillPath == "" { + t.Fatal("spill must succeed in a clean temp dir") + } + resolved, ok := resolveSpillReadPath(spillPath) + if !ok || resolved == "" { + t.Fatalf("a real spill file must resolve, got ok=%v", ok) + } + if _, ok := resolveSpillReadPath(filepath.Join(spillRootPath(), "..", "escape.txt")); ok { + t.Fatal("path traversal out of the spill dir must be rejected") + } + if _, ok := resolveSpillReadPath("/etc/hosts"); ok { + t.Fatal("paths outside the spill dir must be rejected") + } + if _, ok := resolveSpillReadPath(spillRootPath()); ok { + t.Fatal("the spill dir itself is not a readable file target") + } + if runtime.GOOS != "windows" { + // A symlink planted inside the spill dir pointing outside must not leak. + link := filepath.Join(spillRootPath(), "sneaky-link") + if err := os.Symlink("/etc/hosts", link); err == nil { + if _, ok := resolveSpillReadPath(link); ok { + t.Fatal("symlink escaping the spill dir must be rejected") + } + } + } +} + +func TestReadFileCanReadSpillFile(t *testing.T) { + t.Setenv("TMPDIR", t.TempDir()) + spillPath := spillTruncatedOutput("exec_command", "line one\nline two\n") + if spillPath == "" { + t.Fatal("spill must succeed") + } + workspace := t.TempDir() + result := NewReadFileTool(workspace).Run(context.Background(), map[string]any{"path": spillPath}) + if result.Status != StatusOK || !strings.Contains(result.Output, "line two") { + t.Fatalf("read_file must be able to follow the truncation notice: %s %q", result.Status, result.Output) + } +} diff --git a/internal/tools/spill_owner_unix.go b/internal/tools/spill_owner_unix.go new file mode 100644 index 000000000..88fc12eeb --- /dev/null +++ b/internal/tools/spill_owner_unix.go @@ -0,0 +1,24 @@ +//go:build !windows + +package tools + +import ( + "fmt" + "os" + "syscall" +) + +// checkSpillDirOwner rejects a spill directory not owned by the current user: +// on a shared /tmp another user could have pre-created the path and would then +// control its lifetime (deletion/renaming) even though the 0600 spill files +// keep their contents private. +func checkSpillDirOwner(info os.FileInfo) error { + stat, ok := info.Sys().(*syscall.Stat_t) + if !ok { + return nil + } + if int(stat.Uid) != os.Geteuid() { + return fmt.Errorf("spill directory is owned by uid %d, not the current user", stat.Uid) + } + return nil +} diff --git a/internal/tools/spill_owner_windows.go b/internal/tools/spill_owner_windows.go new file mode 100644 index 000000000..da6885289 --- /dev/null +++ b/internal/tools/spill_owner_windows.go @@ -0,0 +1,12 @@ +//go:build windows + +package tools + +import "os" + +// checkSpillDirOwner is a no-op on Windows: %TEMP% is per-user by default and +// there is no portable uid to compare. The Lstat symlink/dir check in spillDir +// still applies. +func checkSpillDirOwner(os.FileInfo) error { + return nil +} diff --git a/internal/tools/spill_test.go b/internal/tools/spill_test.go new file mode 100644 index 000000000..a04e498a0 --- /dev/null +++ b/internal/tools/spill_test.go @@ -0,0 +1,82 @@ +package tools + +import ( + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +func TestTruncateExecOutputSpillWritesFullOutput(t *testing.T) { + t.Setenv("TMPDIR", t.TempDir()) + long := strings.Repeat("line of build output\n", 5000) // ~100KB > 40KB budget + + truncated, wasTruncated := truncateExecOutputSpill(long, defaultMaxOutputTokens, "bash") + if !wasTruncated { + t.Fatal("output over budget must truncate") + } + if !strings.Contains(truncated, "full output saved to ") { + t.Fatalf("truncation notice missing spill path: %q", truncated[:200]) + } + // Extract the path and verify the file holds the complete output. + start := strings.Index(truncated, "full output saved to ") + len("full output saved to ") + end := strings.Index(truncated[start:], " (grep") + spillPath := truncated[start : start+end] + content, err := os.ReadFile(spillPath) + if err != nil { + t.Fatalf("spill file unreadable: %v", err) + } + if string(content) != long { + t.Fatalf("spill file must hold the full output: got %d bytes, want %d", len(content), len(long)) + } +} + +func TestTruncateExecOutputUnderBudgetUnchanged(t *testing.T) { + output, wasTruncated := truncateExecOutputSpill("short output", defaultMaxOutputTokens, "bash") + if wasTruncated || output != "short output" { + t.Fatalf("under-budget output must pass through untouched: %q", output) + } +} + +func TestSweepSpillDirRemovesOnlyOldFiles(t *testing.T) { + dir := t.TempDir() + oldFile := filepath.Join(dir, "bash-old.txt") + newFile := filepath.Join(dir, "bash-new.txt") + for _, path := range []string{oldFile, newFile} { + if err := os.WriteFile(path, []byte("x"), 0o600); err != nil { + t.Fatal(err) + } + } + stale := time.Now().Add(-spillRetention - time.Hour) + if err := os.Chtimes(oldFile, stale, stale); err != nil { + t.Fatal(err) + } + + sweepSpillDir(dir) + + if _, err := os.Stat(oldFile); !os.IsNotExist(err) { + t.Fatal("stale spill file must be removed") + } + if _, err := os.Stat(newFile); err != nil { + t.Fatal("fresh spill file must survive the sweep") + } +} + +func TestSpillTruncatedOutputWritesFile(t *testing.T) { + t.Setenv("TMPDIR", t.TempDir()) + path := spillTruncatedOutput("exec_command", "some output body") + if path == "" { + t.Fatal("spill must return a file path") + } + if base := filepath.Base(path); !strings.HasPrefix(base, "exec_command-") { + t.Fatalf("spill file name must carry the tool prefix: %s", base) + } + content, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + if string(content) != "some output body" { + t.Fatalf("unexpected spill content: %q", content) + } +} diff --git a/internal/tools/web_fetch.go b/internal/tools/web_fetch.go index 252c69550..db207eb19 100644 --- a/internal/tools/web_fetch.go +++ b/internal/tools/web_fetch.go @@ -18,9 +18,14 @@ import ( ) const ( - defaultWebFetchMaxBytes = 64 * 1024 - maxWebFetchMaxBytes = 512 * 1024 - webFetchTimeout = 10 * time.Second + // Raw-body budgets. HTML responses are converted to markdown before they + // reach the model (see web_fetch_markdown.go), so a generous raw budget + // does not translate into a generous context cost — conversion typically + // shrinks a page by an order of magnitude. 64KiB of raw HTML often held + // nothing but a page's , which starved research tasks. + defaultWebFetchMaxBytes = 256 * 1024 + maxWebFetchMaxBytes = 2 * 1024 * 1024 + webFetchTimeout = 30 * time.Second webFetchRedirectLimit = 5 webFetchErrorBodyLimit = 4 * 1024 webFetchPublicOnlyHint = "web_fetch only supports public remote HTTP/HTTPS URLs. For localhost or private network URLs, use bash with curl so sandbox network permission can apply." @@ -115,11 +120,17 @@ func newWebFetchToolWithClientAndResolver(client *http.Client, resolver webFetch }, "max_bytes": { Type: "integer", - Description: "Maximum response body bytes to return.", + Description: "Maximum raw response body bytes to download before conversion.", Default: defaultWebFetchMaxBytes, Minimum: intPtr(1), Maximum: intPtr(maxWebFetchMaxBytes), }, + "format": { + Type: "string", + Description: "auto (default): HTML responses are converted to compact markdown, everything else is returned as-is. raw: never convert. markdown: force conversion.", + Enum: []string{"auto", "raw", "markdown"}, + Default: "auto", + }, }, Required: []string{"url"}, AdditionalProperties: false, @@ -170,6 +181,18 @@ func (tool webFetchTool) run(ctx context.Context, args map[string]any) Result { if err != nil { return errorResult("Error: Invalid arguments for web_fetch: " + err.Error()) } + format, err := stringArg(args, "format", "auto", false) + if err != nil { + return errorResult("Error: Invalid arguments for web_fetch: " + err.Error()) + } + format = strings.ToLower(strings.TrimSpace(format)) + switch format { + case "", "auto": + format = "auto" + case "raw", "markdown": + default: + return errorResult(`Error: Invalid arguments for web_fetch: format must be "auto", "raw", or "markdown".`) + } if err := validateWebFetchURLBeforePermission(rawURL); err != nil { return errorResult("Error: Unsafe URL for web_fetch: " + err.Error()) } @@ -224,14 +247,28 @@ func (tool webFetchTool) run(ctx context.Context, args map[string]any) Result { } body = redactWebFetchText(body) - output := strings.Join([]string{ + // HTML responses are converted to compact markdown by default: raw HTML is + // mostly markup, so conversion typically shrinks the page by an order of + // magnitude and the model reads content instead of boilerplate. format=raw + // is the escape hatch for pages the converter mangles. + converted := false + if format == "markdown" || (format == "auto" && looksLikeHTML(contentType, body)) { + if markdown := htmlToMarkdown(body); markdown != "" { + body = markdown + converted = true + } + } + + headers := []string{ "URL: " + finalURL, "Status: " + webFetchStatusLine(response), "Content-Type: " + firstNonEmptyString(contentType, "unknown"), "Bytes: " + strconv.Itoa(len(body)), - "", - body, - }, "\n") + } + if converted { + headers = append(headers, "Converted: html -> markdown (pass format: \"raw\" for the original HTML)") + } + output := strings.Join(append(headers, "", body), "\n") return Result{ Status: StatusOK, @@ -243,6 +280,7 @@ func (tool webFetchTool) run(ctx context.Context, args map[string]any) Result { "content_type": contentType, "bytes": strconv.Itoa(len(body)), "truncated": strconv.FormatBool(truncated), + "converted": strconv.FormatBool(converted), }, } } diff --git a/internal/tools/web_fetch_markdown.go b/internal/tools/web_fetch_markdown.go new file mode 100644 index 000000000..9db993e99 --- /dev/null +++ b/internal/tools/web_fetch_markdown.go @@ -0,0 +1,92 @@ +package tools + +import ( + "fmt" + "html" + "regexp" + "strings" +) + +// HTML-to-markdown conversion for web_fetch. Raw HTML is mostly boilerplate — +// on a typical page well under a tenth of the bytes are readable content — so +// returning it raw wastes nearly the whole response budget on markup. The +// converter is deliberately dependency-free (regex passes + stdlib entity +// unescaping, no HTML parser): the output feeds a language model, not a +// browser, so best-effort structure (headings, links, list markers) is enough, +// and the format=raw escape hatch covers pages it mangles. + +var ( + webFetchDropBlockRe = regexp.MustCompile(`(?is)<(script|style|head|noscript|svg|template|iframe)\b[^>]*>.*?`) + webFetchCommentRe = regexp.MustCompile(`(?s)`) + webFetchAnchorRe = regexp.MustCompile(`(?is)]*?href\s*=\s*["']?([^"'\s>]+)["']?[^>]*>(.*?)`) + webFetchHeadingRe = regexp.MustCompile(`(?is)]*>(.*?)`) + webFetchListItemRe = regexp.MustCompile(`(?i)]*>`) + webFetchBreakRe = regexp.MustCompile(`(?i)<(?:br|hr)\s*/?>`) + webFetchBlockTagRe = regexp.MustCompile(`(?i)]*>`) + webFetchAnyTagRe = regexp.MustCompile(`(?s)<[^>]*>`) + webFetchSpaceRunRe = regexp.MustCompile(`[ \t]+`) + webFetchBlankRunRe = regexp.MustCompile(`\n{3,}`) +) + +// htmlToMarkdown converts an HTML document to compact markdown-flavored text: +// headings become #-prefixed lines, anchors become [text](href), list items +// become "- " bullets, block boundaries become blank lines, all other markup +// is stripped, and entities are unescaped. +func htmlToMarkdown(body string) string { + text := webFetchCommentRe.ReplaceAllString(body, " ") + text = webFetchDropBlockRe.ReplaceAllString(text, " ") + text = webFetchHeadingRe.ReplaceAllStringFunc(text, func(match string) string { + groups := webFetchHeadingRe.FindStringSubmatch(match) + level := int(groups[1][0] - '0') + if level < 1 || level > 6 { + level = 1 + } + title := strings.TrimSpace(webFetchAnyTagRe.ReplaceAllString(groups[2], " ")) + return "\n\n" + strings.Repeat("#", level) + " " + title + "\n\n" + }) + text = webFetchAnchorRe.ReplaceAllStringFunc(text, func(match string) string { + groups := webFetchAnchorRe.FindStringSubmatch(match) + href := groups[1] + label := strings.TrimSpace(webFetchAnyTagRe.ReplaceAllString(groups[2], " ")) + if label == "" { + return " " + } + if strings.HasPrefix(href, "#") || strings.HasPrefix(strings.ToLower(href), "javascript:") { + return label + } + return fmt.Sprintf("[%s](%s)", label, href) + }) + text = webFetchListItemRe.ReplaceAllString(text, "\n- ") + text = webFetchBreakRe.ReplaceAllString(text, "\n") + text = webFetchBlockTagRe.ReplaceAllString(text, "\n\n") + text = webFetchAnyTagRe.ReplaceAllString(text, " ") + text = html.UnescapeString(text) + + // Whitespace normalization: collapse space runs, trim line edges, cap blank + // runs at one empty line. + lines := strings.Split(text, "\n") + for i, line := range lines { + lines[i] = strings.TrimSpace(webFetchSpaceRunRe.ReplaceAllString(line, " ")) + } + text = strings.Join(lines, "\n") + text = webFetchBlankRunRe.ReplaceAllString(text, "\n\n") + return strings.TrimSpace(text) +} + +// looksLikeHTML reports whether a response should be treated as an HTML +// document, by Content-Type first and a body sniff as fallback (servers +// mislabel HTML as text/plain often enough to matter). +func looksLikeHTML(contentType, body string) bool { + lowered := strings.ToLower(contentType) + if strings.Contains(lowered, "text/html") || strings.Contains(lowered, "application/xhtml") { + return true + } + if lowered != "" && !strings.Contains(lowered, "text/plain") { + return false + } + head := strings.ToLower(strings.TrimSpace(body)) + if len(head) > 512 { + head = head[:512] + } + return strings.Contains(head, " +ignored + + +

Main Title

+

First paragraph with a docs link and & entity.

+
  • alpha
  • beta
+ +

Section

+

Body with extra spaces.

+` + +func TestHTMLToMarkdown(t *testing.T) { + markdown := htmlToMarkdown(sampleHTML) + + for _, want := range []string{ + "# Main Title", + "## Section", + "[docs link](https://example.com/docs)", + "- alpha", + "- beta", + "with a", // paragraph text survives + "& entity", + } { + if !strings.Contains(markdown, want) { + t.Fatalf("markdown missing %q:\n%s", want, markdown) + } + } + for _, banned := range []string{"tracking", "color:red", "hidden comment", "

", "", "ignored"} { + if strings.Contains(markdown, banned) { + t.Fatalf("markdown must not contain %q:\n%s", banned, markdown) + } + } + if strings.Contains(markdown, "extra spaces") { + t.Fatalf("space runs must collapse:\n%s", markdown) + } + if strings.Contains(markdown, "\n\n\n") { + t.Fatalf("blank-line runs must be capped:\n%s", markdown) + } +} + +func TestLooksLikeHTML(t *testing.T) { + cases := []struct { + contentType string + body string + want bool + }{ + {"text/html; charset=utf-8", "anything", true}, + {"application/xhtml+xml", "anything", true}, + {"application/json", `{"html":""}`, false}, + {"text/plain", "", true}, // mislabeled HTML sniffed + {"text/plain", "just text", false}, + {"", "x", true}, + } + for _, c := range cases { + if got := looksLikeHTML(c.contentType, c.body); got != c.want { + t.Errorf("looksLikeHTML(%q, %.20q) = %v, want %v", c.contentType, c.body, got, c.want) + } + } +} + +func TestWebFetchConvertsHTMLByDefault(t *testing.T) { + tool := newWebFetchToolWithClient(webFetchTestClient(func(request *http.Request) (*http.Response, error) { + return webFetchTestResponse(request, http.StatusOK, "text/html; charset=utf-8", sampleHTML), nil + })) + + result := tool.Run(context.Background(), map[string]any{"url": "https://example.com/page"}) + if result.Status != StatusOK { + t.Fatalf("expected ok, got %s: %s", result.Status, result.Output) + } + if !strings.Contains(result.Output, "# Main Title") || strings.Contains(result.Output, "

") { + t.Fatalf("HTML must convert to markdown by default:\n%s", result.Output) + } + if result.Meta["converted"] != "true" { + t.Fatalf("converted meta must be true: %#v", result.Meta) + } +} + +func TestWebFetchRawFormatSkipsConversion(t *testing.T) { + tool := newWebFetchToolWithClient(webFetchTestClient(func(request *http.Request) (*http.Response, error) { + return webFetchTestResponse(request, http.StatusOK, "text/html", sampleHTML), nil + })) + + result := tool.Run(context.Background(), map[string]any{"url": "https://example.com/page", "format": "raw"}) + if result.Status != StatusOK { + t.Fatalf("expected ok, got %s: %s", result.Status, result.Output) + } + if !strings.Contains(result.Output, "

Main Title

") { + t.Fatalf("format=raw must keep original HTML:\n%s", result.Output) + } + if result.Meta["converted"] != "false" { + t.Fatalf("converted meta must be false: %#v", result.Meta) + } +} + +func TestWebFetchLeavesNonHTMLUntouched(t *testing.T) { + tool := newWebFetchToolWithClient(webFetchTestClient(func(request *http.Request) (*http.Response, error) { + return webFetchTestResponse(request, http.StatusOK, "application/json", `{"key":"value"}`), nil + })) + + result := tool.Run(context.Background(), map[string]any{"url": "https://example.com/api"}) + if result.Status != StatusOK { + t.Fatalf("expected ok, got %s: %s", result.Status, result.Output) + } + if !strings.Contains(result.Output, `{"key":"value"}`) { + t.Fatalf("non-HTML must pass through untouched:\n%s", result.Output) + } +} + +func TestWebFetchRejectsInvalidFormat(t *testing.T) { + tool := newWebFetchToolWithClient(webFetchTestClient(func(request *http.Request) (*http.Response, error) { + return webFetchTestResponse(request, http.StatusOK, "text/plain", "x"), nil + })) + result := tool.Run(context.Background(), map[string]any{"url": "https://example.com", "format": "xml"}) + if result.Status != StatusError || !strings.Contains(result.Output, "format must be") { + t.Fatalf("invalid format must be rejected: %s %s", result.Status, result.Output) + } +} diff --git a/internal/tools/workspace.go b/internal/tools/workspace.go index 1a51e08c3..8b8322b32 100644 --- a/internal/tools/workspace.go +++ b/internal/tools/workspace.go @@ -254,6 +254,14 @@ func scopedReadRoots(workspaceRoot string, scope PathScope) ([]string, error) { } func resolveScopedReadPath(workspaceRoot string, scope PathScope, requestedPath string) (string, string, error) { + // Spill files (truncated tool output saved under the per-uid temp dir) are + // readable regardless of scope: the truncation notice tells the model to + // read_file/grep them, which must actually work. resolveSpillReadPath + // verifies containment after symlink resolution, so this cannot be used to + // reach anything outside the spill dir. + if spillPath, ok := resolveSpillReadPath(requestedPath); ok { + return spillPath, spillPath, nil + } if requestedPath == "" || !filepath.IsAbs(requestedPath) || scope == nil { return resolveWorkspacePath(workspaceRoot, requestedPath) } diff --git a/internal/tools/write_file.go b/internal/tools/write_file.go index 875d35c46..e60802e7d 100644 --- a/internal/tools/write_file.go +++ b/internal/tools/write_file.go @@ -44,7 +44,7 @@ func (tool writeFileTool) Run(ctx context.Context, args map[string]any) Result { return tool.RunWithOptions(ctx, args, RunOptions{}) } -func (tool writeFileTool) RunWithOptions(_ context.Context, args map[string]any, options RunOptions) Result { +func (tool writeFileTool) RunWithOptions(ctx context.Context, args map[string]any, options RunOptions) Result { requestedPath, err := aliasedStringArg(args, []string{"path", "file", "file_path", "filename"}, "", true, false) if err != nil { return errorResult("Error: Invalid arguments for write_file: " + err.Error()) @@ -110,6 +110,10 @@ func (tool writeFileTool) RunWithOptions(_ context.Context, args map[string]any, if err := os.WriteFile(absolutePath, []byte(content), 0o644); err != nil { return errorResult("Error writing file " + relativePath + ": " + err.Error()) } + // Optional format-on-write (ZERO_FORMAT_ON_WRITE). Must run BEFORE the + // FileTracker baseline: recording pre-format content would make the very + // next edit look like an external modification and trip the conflict guard. + content = maybeFormatWrittenFile(ctx, absolutePath, content) // Baseline the freshly written content so a later edit/overwrite in this // session compares against what is now on disk. newInfo, _ := os.Stat(absolutePath) @@ -126,6 +130,7 @@ func (tool writeFileTool) RunWithOptions(_ context.Context, args map[string]any, lines++ } summary := fmt.Sprintf("%s %s (%d lines).", verb, relativePath, lines) + summary += inlineDiagnostics(ctx, options, absolutePath, relativePath) result := okResult(summary) result.ChangedFiles = []string{relativePath} // Card-only preview: a real unified diff (all-green for a create, red/green for diff --git a/internal/tui/model.go b/internal/tui/model.go index 81d4a5eed..aa1d61a91 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -4323,6 +4323,10 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str IncludeLSP: true, Autonomy: selfCorrectAutonomyForMode(options.PermissionMode), }) + // Inline post-edit diagnostics: edit_file/write_file append error + // diagnostics for the file they just wrote to their own output, so the + // model sees a break in the same turn. Shares the run's lazy manager. + options.FileDiagnostics = agent.NewFileDiagnostics(lspManager, options.Cwd) } // Some providers synthesize tool-call ids that repeat within a run (e.g.