From 3fa8b396724d7bb79e8efbf118a48a237338f34c Mon Sep 17 00:00:00 2001 From: Kevin Codex Date: Sun, 5 Jul 2026 19:09:49 +0800 Subject: [PATCH 1/5] feat(tools): universal token-based output ceiling with spill-to-disk MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cap every tool result at the registry boundary (16k estimated tokens, ZERO_TOOL_OUTPUT_CEILING_TOKENS-tunable, 0 disables) unless the tool manages its own deliberate budget. Head+tail truncation, full output spilled to the existing redacted spill dir and re-readable via grep/read_file, so nothing is lost — only deferred. This closes the three unbounded paths a single call could flood the context window through (and then re-bill every following turn): - web_fetch: up to 256KiB default / 2MiB max of non-HTML body - skill: full SKILL.md returned uncapped - MCP-served tools: no budget at all Self-budgeting tools opt out via an unexported marker interface (an MCP server can never exempt itself): bash, exec_command (model-raisable), read_file, read_minified_file, grep, glob, list_directory. Browser snapshots deliberately stay under the net. bash itself is tightened to match the unified token model: emit budget 96KiB -> 32KiB per stream (~24k -> ~8k tokens worst case per stream), while capture retention stays at 96KiB head+tail so the new spill file covers 6x more of the output than the transcript shows. --- internal/tools/bash.go | 58 +++++++++-- internal/tools/bash_budget_test.go | 25 ++++- internal/tools/output_ceiling.go | 81 +++++++++++++++ internal/tools/output_ceiling_test.go | 139 ++++++++++++++++++++++++++ internal/tools/registry.go | 15 ++- 5 files changed, 305 insertions(+), 13 deletions(-) create mode 100644 internal/tools/output_ceiling.go create mode 100644 internal/tools/output_ceiling_test.go diff --git a/internal/tools/bash.go b/internal/tools/bash.go index c93dcbd12..ba9af9ffd 100644 --- a/internal/tools/bash.go +++ b/internal/tools/bash.go @@ -126,8 +126,8 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS // can't grow Zero's memory before truncation: only the head+tail each stream // will ever surface to the model are retained, the middle is discarded as it // streams, and the true size is counted for the truncation marker. - stdout := newBoundedBuffer(bashOutputBudgetBytes, bashOutputBudgetBytes) - stderr := newBoundedBuffer(bashOutputBudgetBytes, bashOutputBudgetBytes) + stdout := newBoundedBuffer(bashCaptureBudgetBytes, bashCaptureBudgetBytes) + stderr := newBoundedBuffer(bashCaptureBudgetBytes, bashCaptureBudgetBytes) command.Stdout = stdout command.Stderr = stderr @@ -370,12 +370,21 @@ func formatBashOutput(stdout string, stderr string, exitCode int) string { return strings.Join(parts, "\n") } -// bashOutputBudgetBytes caps each of stdout/stderr shown to the model. bash is the -// one tool that can emit unbounded output (`cat large.log`, `find /`, verbose test -// runs); every other read/search tool already budgets its output. Head+tail -// truncation keeps both the start and the end of an oversized stream, since -// build/test failures usually surface at the tail. -const bashOutputBudgetBytes = 96 * 1024 +// bashOutputBudgetBytes caps each of stdout/stderr shown to the model (32 KiB +// ≈ 8k tokens per stream). bash is the one tool that can emit unbounded output +// (`cat large.log`, `find /`, verbose test runs); a chatty build log at the old +// 96 KiB budget injected ~24k tokens per call and was then re-billed on every +// following turn. Head+tail truncation keeps both the start and the end of an +// oversized stream, since build/test failures usually surface at the tail, and +// the full captured text is spilled to disk so nothing is lost — the model +// greps/reads the spill instead of re-running the command. +const bashOutputBudgetBytes = 32 * 1024 + +// bashCaptureBudgetBytes bounds the head and tail retained IN MEMORY per +// stream during capture. Deliberately larger than the emit budget: the extra +// retained bytes never reach the transcript, but they make the spill file — +// the model's recovery path — cover 6× more of the output. +const bashCaptureBudgetBytes = 96 * 1024 // budgetBashOutput truncates stdout and stderr to bashOutputBudgetBytes each, // keeping the head and tail of anything larger, and records raw/emitted byte @@ -395,6 +404,19 @@ func budgetBashCapture(out string, outTotal int, errStr string, errTotal int, me outText, outRaw, outTrunc := truncateHeadTailWithTotal(out, outTotal, bashOutputBudgetBytes) errText, errRaw, errTrunc := truncateHeadTailWithTotal(errStr, errTotal, bashOutputBudgetBytes) truncated := outTrunc || errTrunc + if truncated { + if spillPath := spillBashStreams(out, errStr); spillPath != "" { + hint := "\n[zero] captured output saved to " + spillPath + " (grep or read_file it instead of re-running)" + if errTrunc { + errText += hint + } else { + outText += hint + } + if meta != nil { + meta["spill_path"] = spillPath + } + } + } if meta != nil { emitted := len(outText) + len(errText) meta["raw_bytes"] = strconv.Itoa(outRaw + errRaw) @@ -407,6 +429,26 @@ func budgetBashCapture(out string, outTotal int, errStr string, errTotal int, me return outText, errText, truncated } +// spillBashStreams writes the retained (capture-bounded) stdout and stderr to +// the spill directory as one sectioned file and returns its path, or "" when +// spilling fails. The spill holds up to bashCaptureBudgetBytes of head and +// tail per stream — everything zero kept in memory, which is more than the +// emit budget shows the model but not necessarily the whole output (the +// capture middle is discarded as it streams; the in-text truncation marker +// reports the true omitted size). +func spillBashStreams(stdout, stderr string) string { + var combined strings.Builder + combined.Grow(len(stdout) + len(stderr) + 64) + combined.WriteString("### stdout\n") + combined.WriteString(stdout) + if !strings.HasSuffix(stdout, "\n") { + combined.WriteString("\n") + } + combined.WriteString("### stderr\n") + combined.WriteString(stderr) + return spillTruncatedOutput("bash", combined.String()) +} + // boundedBuffer is an io.Writer that retains at most headCap bytes from the start // and tailCap bytes from the end of a stream while counting the total written, so // a command emitting unbounded output (`cat huge.log`, `yes`) cannot grow Zero's diff --git a/internal/tools/bash_budget_test.go b/internal/tools/bash_budget_test.go index adb3ab0db..9d40173ae 100644 --- a/internal/tools/bash_budget_test.go +++ b/internal/tools/bash_budget_test.go @@ -2,6 +2,7 @@ package tools import ( "fmt" + "os" "strconv" "strings" "testing" @@ -26,8 +27,10 @@ func TestBudgetBashOutputSmallPassesThrough(t *testing.T) { } // Oversized stdout is truncated head+tail: both the first and last lines survive, -// the middle is dropped behind a marker, and meta is flagged. +// the middle is dropped behind a marker, meta is flagged, and the captured text +// is spilled to a re-readable file. func TestBudgetBashOutputTruncatesHeadAndTail(t *testing.T) { + t.Setenv("TMPDIR", t.TempDir()) head := "FIRST_LINE_MARKER\n" tail := "\nLAST_LINE_MARKER" big := head + strings.Repeat("x", bashOutputBudgetBytes) + tail @@ -47,9 +50,23 @@ func TestBudgetBashOutputTruncatesHeadAndTail(t *testing.T) { if !strings.Contains(out, "output truncated") { t.Fatalf("expected a truncation marker, got:\n%s", out[:200]) } - if len(out) > bashOutputBudgetBytes { + // The spill hint (path length varies) is the only allowed overage past the budget. + if len(out) > bashOutputBudgetBytes+512 { t.Fatalf("emitted %d bytes exceeds budget %d", len(out), bashOutputBudgetBytes) } + if !strings.Contains(out, "captured output saved to ") { + t.Fatalf("truncated bash output must carry a spill hint:\n%s", out[len(out)-300:]) + } + if meta["spill_path"] == "" { + t.Fatal("spill path missing from meta") + } + content, err := os.ReadFile(meta["spill_path"]) + if err != nil { + t.Fatalf("spill file unreadable: %v", err) + } + if !strings.Contains(string(content), "### stdout") || !strings.Contains(string(content), strings.Repeat("x", 1024)) { + t.Fatalf("spill must hold the sectioned captured streams (got %d bytes)", len(content)) + } if meta["truncated"] != "true" { t.Fatalf("expected truncated=true, got %v", meta) } @@ -100,6 +117,7 @@ func TestBoundedBufferKeepsHeadAndTailBounded(t *testing.T) { // budgetBashCapture reports the TRUE total (not the retained size) in the marker // and raw_bytes, even though only a bounded head+tail was ever held in memory. func TestBudgetBashCaptureReportsTrueTotal(t *testing.T) { + t.Setenv("TMPDIR", t.TempDir()) // Retained head+tail as boundedBuffer would hand over; the real command produced // far more than was kept. retained := "HEAD_START" + strings.Repeat("y", bashOutputBudgetBytes) + "TAIL_END" @@ -114,7 +132,8 @@ func TestBudgetBashCaptureReportsTrueTotal(t *testing.T) { if !strings.Contains(out, "HEAD_START") || !strings.Contains(out, "TAIL_END") { t.Fatalf("head/tail lost after budgeting:\n%s", out[:min(120, len(out))]) } - if len(out) > bashOutputBudgetBytes { + // Spill hint is the only allowed overage past the budget. + if len(out) > bashOutputBudgetBytes+512 { t.Fatalf("emitted %d bytes exceeds budget %d", len(out), bashOutputBudgetBytes) } if meta["raw_bytes"] != strconv.Itoa(total) { diff --git a/internal/tools/output_ceiling.go b/internal/tools/output_ceiling.go new file mode 100644 index 000000000..b911b6950 --- /dev/null +++ b/internal/tools/output_ceiling.go @@ -0,0 +1,81 @@ +package tools + +import ( + "os" + "strconv" + "strings" +) + +// Universal output ceiling. Every tool result leaving the registry boundary is +// capped at a single token-denominated budget unless the tool manages its own +// deliberate budget (selfBudgeting below). This is the safety net for tools +// with no cap of their own — web_fetch bodies, skill files, MCP server tools, +// browser snapshots, anything added later — so no single call can flood the +// context window with output that is then re-billed on every following turn. +// Truncation keeps head+tail and spills the full text to disk (re-readable via +// grep/read_file), so nothing is lost, only deferred. + +// defaultOutputCeilingTokens is the per-call ceiling in estimated tokens +// (bytes/4). 16k tokens = 64 KiB — deliberately equal to the search-tool +// budget: generous enough for a large fetched document, small enough that one +// call cannot eat a third of a small context window. +const defaultOutputCeilingTokens = 16_000 + +// outputCeilingEnv overrides the ceiling (in tokens). Zero or negative +// disables the ceiling entirely; unset or unparsable keeps the default. +const outputCeilingEnv = "ZERO_TOOL_OUTPUT_CEILING_TOKENS" + +// selfBudgeting marks a tool that enforces its own deliberate output budget — +// possibly model-raisable (exec_command) — which the registry ceiling must not +// second-guess. The method is unexported on purpose: only tools in this +// package can opt out, so an MCP-served tool can never exempt itself. +type selfBudgeting interface{ managesOutputBudget() } + +// The exemption list, kept in one place. Each of these applies its own budget +// before returning: bash (bashOutputBudgetBytes per stream + spill), +// exec_command (model-raisable token budget + spill), read tools (128 KiB), +// search tools (64 KiB). +func (bashTool) managesOutputBudget() {} +func (execCommandTool) managesOutputBudget() {} +func (readFileTool) managesOutputBudget() {} +func (readMinifiedFileTool) managesOutputBudget() {} +func (grepTool) managesOutputBudget() {} +func (globTool) managesOutputBudget() {} +func (listDirectoryTool) managesOutputBudget() {} + +func resolveOutputCeilingTokens() int { + raw := strings.TrimSpace(os.Getenv(outputCeilingEnv)) + if raw == "" { + return defaultOutputCeilingTokens + } + parsed, err := strconv.Atoi(raw) + if err != nil { + return defaultOutputCeilingTokens + } + return parsed +} + +// enforceOutputCeiling truncates an over-ceiling result to head+tail within +// the token budget, spilling the full (already redaction-scrubbed) output to +// disk with a recovery hint. Runs after scrubResultSecrets so the transcript +// and the spill file agree on what was hidden. +func enforceOutputCeiling(toolName string, result Result) Result { + ceiling := resolveOutputCeilingTokens() + if ceiling <= 0 { + return result + } + if len(result.Output) <= ceiling*4 { + return result + } + rawBytes := len(result.Output) + truncated, _ := truncateExecOutputSpill(result.Output, ceiling, toolName) + result.Output = truncated + result.Truncated = true + if result.Meta == nil { + result.Meta = map[string]string{} + } + result.Meta["raw_bytes"] = strconv.Itoa(rawBytes) + result.Meta["emitted_bytes"] = strconv.Itoa(len(result.Output)) + result.Meta["estimated_tokens"] = strconv.Itoa(estimatedTokensFromBytes(len(result.Output))) + return result +} diff --git a/internal/tools/output_ceiling_test.go b/internal/tools/output_ceiling_test.go new file mode 100644 index 000000000..40995a7fb --- /dev/null +++ b/internal/tools/output_ceiling_test.go @@ -0,0 +1,139 @@ +package tools + +import ( + "context" + "os" + "strconv" + "strings" + "testing" +) + +// ceilingFakeTool is a minimal tool with no output budget of its own — the +// stand-in for an MCP-served tool or any future builtin that forgets to cap +// its output. +type ceilingFakeTool struct { + baseTool + output string +} + +func (tool ceilingFakeTool) Run(context.Context, map[string]any) Result { + return okResult(tool.output) +} + +// ceilingExemptTool is the same fake but marked self-budgeting. +type ceilingExemptTool struct{ ceilingFakeTool } + +func (ceilingExemptTool) managesOutputBudget() {} + +func newCeilingFakeTool(name, output string) ceilingFakeTool { + return ceilingFakeTool{ + baseTool: baseTool{name: name, safety: readOnlySafety("test tool")}, + output: output, + } +} + +// An unbudgeted tool's oversized output is capped at the universal ceiling, +// head and tail survive, and the full output is spilled to a re-readable file. +func TestOutputCeilingCapsUnbudgetedTool(t *testing.T) { + t.Setenv("TMPDIR", t.TempDir()) + big := "HEAD_MARK\n" + strings.Repeat("z", defaultOutputCeilingTokens*4*3) + "\nTAIL_MARK" + + registry := NewRegistry() + registry.Register(newCeilingFakeTool("fake_big", big)) + result := registry.Run(context.Background(), "fake_big", map[string]any{}) + + if !result.Truncated { + t.Fatal("over-ceiling output must report truncated=true") + } + maxBytes := defaultOutputCeilingTokens * 4 + if len(result.Output) > maxBytes+512 { + t.Fatalf("emitted %d bytes, ceiling is %d", len(result.Output), maxBytes) + } + if !strings.Contains(result.Output, "HEAD_MARK") || !strings.Contains(result.Output, "TAIL_MARK") { + t.Fatal("head/tail lost at the ceiling") + } + if result.Meta["raw_bytes"] != strconv.Itoa(len(big)) { + t.Fatalf("raw_bytes = %s, want %d", result.Meta["raw_bytes"], len(big)) + } + if !strings.Contains(result.Output, "full output saved to ") { + t.Fatalf("ceiling truncation must include a spill hint: %q", result.Output[:200]) + } + start := strings.Index(result.Output, "full output saved to ") + len("full output saved to ") + end := strings.Index(result.Output[start:], " (grep") + content, err := os.ReadFile(result.Output[start : start+end]) + if err != nil { + t.Fatalf("spill file unreadable: %v", err) + } + if string(content) != big { + t.Fatalf("spill must hold the full output: got %d bytes, want %d", len(content), len(big)) + } +} + +// A self-budgeting tool's output passes the boundary untouched even when it +// exceeds the universal ceiling — its own (possibly model-raised) budget wins. +func TestOutputCeilingSkipsSelfBudgetingTool(t *testing.T) { + big := strings.Repeat("z", defaultOutputCeilingTokens*4*2) + registry := NewRegistry() + registry.Register(ceilingExemptTool{newCeilingFakeTool("fake_exempt", big)}) + + result := registry.Run(context.Background(), "fake_exempt", map[string]any{}) + if result.Truncated || result.Output != big { + t.Fatalf("self-budgeting tool must bypass the ceiling (truncated=%v, %d bytes)", result.Truncated, len(result.Output)) + } +} + +// Under-ceiling output is untouched and gains no budget metadata. +func TestOutputCeilingSmallOutputUntouched(t *testing.T) { + registry := NewRegistry() + registry.Register(newCeilingFakeTool("fake_small", "hello")) + + result := registry.Run(context.Background(), "fake_small", map[string]any{}) + if result.Truncated || result.Output != "hello" { + t.Fatalf("small output altered: %+v", result) + } + if _, ok := result.Meta["raw_bytes"]; ok { + t.Fatal("small output must not gain budget metadata") + } +} + +// ZERO_TOOL_OUTPUT_CEILING_TOKENS tightens, loosens, or disables the ceiling. +func TestOutputCeilingEnvOverride(t *testing.T) { + t.Setenv("TMPDIR", t.TempDir()) + big := strings.Repeat("z", 100*1024) + + t.Setenv(outputCeilingEnv, "1000") + registry := NewRegistry() + registry.Register(newCeilingFakeTool("fake_env", big)) + result := registry.Run(context.Background(), "fake_env", map[string]any{}) + if !result.Truncated || len(result.Output) > 1000*4+512 { + t.Fatalf("ceiling override ignored: truncated=%v, %d bytes", result.Truncated, len(result.Output)) + } + + t.Setenv(outputCeilingEnv, "0") + result = registry.Run(context.Background(), "fake_env", map[string]any{}) + if result.Truncated || len(result.Output) != len(big) { + t.Fatalf("ceiling=0 must disable the net: truncated=%v, %d bytes", result.Truncated, len(result.Output)) + } +} + +// The exemption list matches the tools that really do budget themselves, and +// the unbudgeted ones (web_fetch, skill, browser) stay under the net. +func TestSelfBudgetingExemptionList(t *testing.T) { + dir := t.TempDir() + exempt := []Tool{ + NewBashTool(dir), + NewExecCommandTool(dir, newExecSessionManager()), + NewReadFileTool(dir), + NewGrepTool(dir), + NewGlobTool(dir), + NewListDirectoryTool(dir), + } + for _, tool := range exempt { + if _, ok := tool.(selfBudgeting); !ok { + t.Errorf("%s must be exempt from the output ceiling", tool.Name()) + } + } + if _, ok := NewWebFetchTool().(selfBudgeting); ok { + t.Error("web_fetch must NOT be exempt — the ceiling is its only budget") + } +} diff --git a/internal/tools/registry.go b/internal/tools/registry.go index 0dcd8a35b..c41c40d04 100644 --- a/internal/tools/registry.go +++ b/internal/tools/registry.go @@ -123,13 +123,24 @@ func (registry *Registry) Run(ctx context.Context, name string, args map[string] func (registry *Registry) RunWithOptions(ctx context.Context, name string, args map[string]any, options RunOptions) (result Result) { // Every return path passes through scrubResultSecrets exactly once, so denial, // permission, and unknown-tool error messages (which can echo secret-bearing - // args/paths) are redacted at the boundary just like tool output. - defer func() { result = scrubResultSecrets(result) }() + // args/paths) are redacted at the boundary just like tool output. The output + // ceiling runs after the scrub so the transcript and the spill file agree on + // what was hidden. + ceilingExempt := false + defer func() { + result = scrubResultSecrets(result) + if !ceilingExempt { + result = enforceOutputCeiling(name, result) + } + }() tool, ok := registry.Get(name) if !ok { return errorResult(`Error: Unknown tool "` + name + `".`) } + if _, ok := tool.(selfBudgeting); ok { + ceilingExempt = true + } if rejecter, ok := tool.(PrePermissionRejecter); ok { if res, rejected := rejecter.RejectBeforePermission(args); rejected { return res From 541acf9b4359c88a431ecf5d0a88bd18065291ae Mon Sep 17 00:00:00 2001 From: Kevin Codex Date: Sun, 5 Jul 2026 19:18:13 +0800 Subject: [PATCH 2/5] perf(agent): collect post-edit LSP diagnostics off the tool-call path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Inline post-edit diagnostics blocked every edit_file/write_file result on the language server: a >=300ms quiet-period debounce that reset on each publish, seconds on the first edit per language, capped at 10s — the largest recurring latency tax in an edit-heavy session. Diagnostics are now collected in the background: a file is enqueued the moment a mutating tool finishes, a single worker checks it while the rest of the batch executes, and the loop drains completed results just before building the NEXT provider request, appending errors as a user nudge. The model always issues another request after an edit turn (to read its tool results), so it still sees introduced errors at the same decision point — no tool call ever stalls on the language server. A file re-edited before its check runs is checked once in its final state (the checker re-reads the file); re-edited after, it is re-queued and the newer result wins. A drain that finds the worker still busy gives up after 3s and delivers on a later turn — results stay accurate because any further edit re-enqueues the file. The tools-side inline mechanism (tools.RunOptions.Diagnostics) is kept for direct API callers; the loop just no longer wires it. --- internal/agent/async_diagnostics.go | 168 +++++++++++++++++ internal/agent/async_diagnostics_test.go | 223 +++++++++++++++++++++++ internal/agent/loop.go | 36 +++- internal/agent/types.go | 12 +- internal/cli/exec.go | 7 +- internal/tui/model.go | 6 +- 6 files changed, 435 insertions(+), 17 deletions(-) create mode 100644 internal/agent/async_diagnostics.go create mode 100644 internal/agent/async_diagnostics_test.go diff --git a/internal/agent/async_diagnostics.go b/internal/agent/async_diagnostics.go new file mode 100644 index 000000000..a1a0dd9b9 --- /dev/null +++ b/internal/agent/async_diagnostics.go @@ -0,0 +1,168 @@ +package agent + +import ( + "context" + "path/filepath" + "sort" + "strings" + "sync" + "time" +) + +// Post-edit LSP diagnostics used to run synchronously inside edit_file / +// write_file: every edit blocked its tool result on a fresh-publish wait with +// a quiet-period debounce (≥300ms floor that reset on each publish, 10s cap) — +// the largest recurring latency tax in an edit-heavy session. Collection is +// now asynchronous: a file is enqueued the moment a mutating tool finishes, a +// single worker checks it while the rest of the tool batch (and any +// self-correct verification) runs, and the loop drains completed results just +// before building the NEXT provider request, appending errors as a user-role +// nudge. The model always issues another request after an edit turn — it has +// to read its tool results — so it still sees introduced errors at the same +// decision point as before, without any tool call stalling on the language +// server. +// +// A file re-edited before its check runs is checked once in its final state +// (the checker re-reads the file); a file re-edited after its check completed +// is re-enqueued and the newer result wins. When the worker is still busy at +// drain time, the drain gives up after a short wait and the results are +// delivered on the following turn instead — they stay accurate because any +// further edit re-enqueues the file. + +// asyncDiagnosticsDrainTimeout bounds how long a turn waits on the in-flight +// check before deferring delivery to the next turn. Most checks finish well +// under it (debounce floor + analysis); a cold language server on a large +// repo does not get to stall the loop the way the old inline 10s cap could. +// A var so tests can shorten the wait. +var asyncDiagnosticsDrainTimeout = 3 * time.Second + +// asyncDiagnosticsNudge prefixes the drained diagnostics blocks; phrased like +// the old inline block so the model reacts the same way. +const asyncDiagnosticsNudge = "Diagnostics after your recent edits (fix any errors you introduced):\n" + +type asyncDiagnostics struct { + check func(context.Context, string) string + workspaceRoot string + + mu sync.Mutex + queue []string + queued map[string]bool + results map[string]string + // working is non-nil while the worker goroutine runs and is closed when it + // exits, so drain can wait for quiescence without polling. + working chan struct{} +} + +// newAsyncDiagnostics returns nil when check is nil, and every method on a nil +// receiver is a no-op — callers wire it unconditionally. +func newAsyncDiagnostics(check func(context.Context, string) string, workspaceRoot string) *asyncDiagnostics { + if check == nil { + return nil + } + return &asyncDiagnostics{ + check: check, + workspaceRoot: workspaceRoot, + queued: map[string]bool{}, + results: map[string]string{}, + } +} + +// enqueue schedules changed files for a background check. Paths are +// workspace-relative from Result.ChangedFiles (absolute when under a granted +// extra write root), mirroring SelfCorrector's resolution. +func (diagnostics *asyncDiagnostics) enqueue(ctx context.Context, files []string) { + if diagnostics == nil || len(files) == 0 { + return + } + diagnostics.mu.Lock() + defer diagnostics.mu.Unlock() + for _, file := range files { + path := diagnostics.absPath(file) + if diagnostics.queued[path] { + continue + } + diagnostics.queued[path] = true + // A completed result for this file describes a pre-edit state now. + delete(diagnostics.results, path) + diagnostics.queue = append(diagnostics.queue, path) + } + if diagnostics.working == nil && len(diagnostics.queue) > 0 { + done := make(chan struct{}) + diagnostics.working = done + go diagnostics.work(ctx, done) + } +} + +func (diagnostics *asyncDiagnostics) work(ctx context.Context, done chan struct{}) { + defer close(done) + for { + diagnostics.mu.Lock() + if len(diagnostics.queue) == 0 || ctx.Err() != nil { + diagnostics.working = nil + diagnostics.mu.Unlock() + return + } + path := diagnostics.queue[0] + diagnostics.queue = diagnostics.queue[1:] + // Un-mark before checking so an edit that lands DURING the check + // re-enqueues the file and it is checked again in its final state. + delete(diagnostics.queued, path) + diagnostics.mu.Unlock() + + block := diagnostics.check(ctx, path) + + diagnostics.mu.Lock() + if block != "" { + diagnostics.results[path] = block + } else { + delete(diagnostics.results, path) + } + diagnostics.mu.Unlock() + } +} + +// drain waits briefly for the worker to go quiet, then formats and consumes +// all completed error blocks as one nudge. Returns "" when there is nothing +// to report or the worker is still busy (results then surface next turn). +func (diagnostics *asyncDiagnostics) drain(ctx context.Context) string { + if diagnostics == nil { + return "" + } + diagnostics.mu.Lock() + busy := diagnostics.working + diagnostics.mu.Unlock() + if busy != nil { + timer := time.NewTimer(asyncDiagnosticsDrainTimeout) + defer timer.Stop() + select { + case <-busy: + case <-timer.C: + return "" + case <-ctx.Done(): + return "" + } + } + diagnostics.mu.Lock() + defer diagnostics.mu.Unlock() + if len(diagnostics.results) == 0 { + return "" + } + paths := make([]string, 0, len(diagnostics.results)) + for path := range diagnostics.results { + paths = append(paths, path) + } + sort.Strings(paths) + blocks := make([]string, 0, len(paths)) + for _, path := range paths { + blocks = append(blocks, diagnostics.results[path]) + } + diagnostics.results = map[string]string{} + return asyncDiagnosticsNudge + strings.Join(blocks, "\n") +} + +func (diagnostics *asyncDiagnostics) absPath(path string) string { + if filepath.IsAbs(path) { + return path + } + return filepath.Join(diagnostics.workspaceRoot, path) +} diff --git a/internal/agent/async_diagnostics_test.go b/internal/agent/async_diagnostics_test.go new file mode 100644 index 000000000..ca9d62e18 --- /dev/null +++ b/internal/agent/async_diagnostics_test.go @@ -0,0 +1,223 @@ +package agent + +import ( + "context" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "github.com/Gitlawb/zero/internal/tools" + "github.com/Gitlawb/zero/internal/zeroruntime" +) + +func TestAsyncDiagnosticsNilCollectorNoOps(t *testing.T) { + var diagnostics *asyncDiagnostics + if diagnostics != newAsyncDiagnostics(nil, "/ws") { + t.Fatal("nil check must yield a nil collector") + } + diagnostics.enqueue(context.Background(), []string{"a.go"}) + if nudge := diagnostics.drain(context.Background()); nudge != "" { + t.Fatalf("nil collector drain = %q, want empty", nudge) + } +} + +func TestAsyncDiagnosticsCollectsAndDrainsOnce(t *testing.T) { + var mu sync.Mutex + var checked []string + check := func(_ context.Context, absPath string) string { + mu.Lock() + checked = append(checked, absPath) + mu.Unlock() + if filepath.Base(absPath) == "clean.go" { + return "" + } + return "ERR " + filepath.Base(absPath) + } + diagnostics := newAsyncDiagnostics(check, "/ws") + diagnostics.enqueue(context.Background(), []string{"broken.go", "clean.go"}) + + nudge := diagnostics.drain(context.Background()) + if !strings.HasPrefix(nudge, asyncDiagnosticsNudge) { + t.Fatalf("nudge missing prefix: %q", nudge) + } + if !strings.Contains(nudge, "ERR broken.go") || strings.Contains(nudge, "clean.go") { + t.Fatalf("nudge = %q, want broken.go errors only", nudge) + } + mu.Lock() + for _, path := range checked { + if !filepath.IsAbs(path) || !strings.HasPrefix(path, filepath.FromSlash("/ws")) { + t.Fatalf("check received %q, want workspace-absolute path", path) + } + } + mu.Unlock() + if second := diagnostics.drain(context.Background()); second != "" { + t.Fatalf("second drain = %q, want empty (results are consumed)", second) + } +} + +// A check still running at drain time defers delivery to the next drain +// instead of blocking the turn. +func TestAsyncDiagnosticsSlowCheckDefersToNextDrain(t *testing.T) { + previous := asyncDiagnosticsDrainTimeout + asyncDiagnosticsDrainTimeout = 20 * time.Millisecond + defer func() { asyncDiagnosticsDrainTimeout = previous }() + + release := make(chan struct{}) + check := func(_ context.Context, absPath string) string { + <-release + return "ERR " + filepath.Base(absPath) + } + diagnostics := newAsyncDiagnostics(check, "/ws") + diagnostics.enqueue(context.Background(), []string{"slow.go"}) + + start := time.Now() + if nudge := diagnostics.drain(context.Background()); nudge != "" { + t.Fatalf("drain during in-flight check = %q, want empty", nudge) + } + if elapsed := time.Since(start); elapsed > time.Second { + t.Fatalf("drain blocked %v, want bounded by the (shortened) timeout", elapsed) + } + + close(release) + deadline := time.Now().Add(2 * time.Second) + for { + if nudge := diagnostics.drain(context.Background()); nudge != "" { + if !strings.Contains(nudge, "ERR slow.go") { + t.Fatalf("deferred nudge = %q", nudge) + } + return + } + if time.Now().After(deadline) { + t.Fatal("deferred diagnostics never delivered") + } + time.Sleep(5 * time.Millisecond) + } +} + +// Re-enqueueing a file after its check completed replaces the stale result. +func TestAsyncDiagnosticsReEditReplacesResult(t *testing.T) { + var mu sync.Mutex + response := "ERR old" + check := func(context.Context, string) string { + mu.Lock() + defer mu.Unlock() + return response + } + diagnostics := newAsyncDiagnostics(check, "/ws") + + diagnostics.enqueue(context.Background(), []string{"a.go"}) + waitForIdle(t, diagnostics) + mu.Lock() + response = "ERR new" + mu.Unlock() + diagnostics.enqueue(context.Background(), []string{"a.go"}) + + nudge := diagnostics.drain(context.Background()) + if !strings.Contains(nudge, "ERR new") || strings.Contains(nudge, "ERR old") { + t.Fatalf("nudge = %q, want only the re-check result", nudge) + } +} + +func waitForIdle(t *testing.T, diagnostics *asyncDiagnostics) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for { + diagnostics.mu.Lock() + busy := diagnostics.working + diagnostics.mu.Unlock() + if busy == nil { + return + } + if time.Now().After(deadline) { + t.Fatal("worker never went idle") + } + time.Sleep(2 * time.Millisecond) + } +} + +// changedFilesTool is a fake mutating tool reporting a changed file without +// touching the filesystem. +type changedFilesTool struct{} + +func (changedFilesTool) Name() string { return "fake_edit" } +func (changedFilesTool) Description() string { return "test mutating tool" } +func (changedFilesTool) Parameters() tools.Schema { + return tools.Schema{Type: "object", Properties: map[string]tools.PropertySchema{}} +} +func (changedFilesTool) Safety() tools.Safety { + return tools.Safety{SideEffect: tools.SideEffectWrite, Permission: tools.PermissionAllow, Reason: "test"} +} +func (changedFilesTool) Run(context.Context, map[string]any) tools.Result { + return tools.Result{Status: tools.StatusOK, Output: "edited", ChangedFiles: []string{"main.go"}} +} + +// End-to-end through Run: an edit turn does NOT carry diagnostics in the tool +// result (the old inline path), but the next request contains the background +// diagnostics nudge, before the model finalizes. +func TestRunDeliversAsyncDiagnosticsNudgeNextTurn(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(changedFilesTool{}) + provider := &mockProvider{ + turns: [][]zeroruntime.StreamEvent{ + { + {Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call-1", ToolName: "fake_edit"}, + {Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call-1", ArgumentsFragment: `{}`}, + {Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call-1"}, + {Type: zeroruntime.StreamEventDone}, + }, + { + {Type: zeroruntime.StreamEventText, Content: "done"}, + {Type: zeroruntime.StreamEventDone}, + }, + }, + } + var checkedPath string + var mu sync.Mutex + + result, err := Run(context.Background(), "fix main.go", provider, Options{ + Registry: registry, + Cwd: filepath.FromSlash("/ws"), + FileDiagnostics: func(_ context.Context, absPath string) string { + mu.Lock() + checkedPath = absPath + mu.Unlock() + return "main.go:1:1 error: boom" + }, + }) + if err != nil { + t.Fatal(err) + } + if result.FinalAnswer != "done" { + t.Fatalf("final answer = %q", result.FinalAnswer) + } + if len(provider.requests) != 2 { + t.Fatalf("expected 2 provider turns, got %d", len(provider.requests)) + } + + // The tool result itself must NOT block on / embed diagnostics. + for _, message := range provider.requests[1].Messages { + if message.Role == zeroruntime.MessageRoleTool && strings.Contains(message.Content, "boom") { + t.Fatalf("diagnostics leaked into the tool result: %q", message.Content) + } + } + // The nudge must be present in the second request as a user message. + found := false + for _, message := range provider.requests[1].Messages { + if message.Role == zeroruntime.MessageRoleUser && strings.HasPrefix(message.Content, asyncDiagnosticsNudge) { + if !strings.Contains(message.Content, "boom") { + t.Fatalf("nudge missing diagnostics: %q", message.Content) + } + found = true + } + } + if !found { + t.Fatal("no async diagnostics nudge in the follow-up request") + } + mu.Lock() + defer mu.Unlock() + if want := filepath.Join(filepath.FromSlash("/ws"), "main.go"); checkedPath != want { + t.Fatalf("checked path = %q, want %q", checkedPath, want) + } +} diff --git a/internal/agent/loop.go b/internal/agent/loop.go index 6583fb0d7..f162bb56e 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -137,6 +137,12 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) guards := newGuardState() compactor := newCompactionState(options) + // Background post-edit diagnostics: files changed by mutating tools are + // checked off the tool-call critical path and any errors are appended as a + // nudge before the NEXT request (see async_diagnostics.go). nil when no + // FileDiagnostics callback is wired; every method no-ops on nil. + postEditDiagnostics := newAsyncDiagnostics(options.FileDiagnostics, options.Cwd) + // loaded tracks deferred-eligible tools the model has pulled via tool_search // during THIS run. It is consulted by partitionTools each turn to expose a // loaded tool's full schema; it lives only for the run (v1 within-run scope). @@ -159,6 +165,17 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) for turn := 0; turn < maxTurns; turn++ { result.Turns = turn + 1 + // Deliver background post-edit diagnostics from the previous turn's edits + // BEFORE compaction so the nudge is part of the request being budgeted. + // A brief wait at most (asyncDiagnosticsDrainTimeout); an unfinished check + // simply delivers on a later turn. + if nudge := postEditDiagnostics.drain(ctx); nudge != "" { + messages = append(messages, zeroruntime.Message{ + Role: zeroruntime.MessageRoleUser, + Content: nudge, + }) + } + // Build the per-turn tool list first so proactive compaction can include // the tool-definition tokens (they ride on every request) in its estimate. // partitionTools depends only on registry/permissions/options/loaded, not on @@ -590,11 +607,14 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) failureHint = toolFailureHint(call.Name, toolSchemaJSON(registry, call.Name), toolResult.Output) } - // Post-edit self-correction: collect the files this successful mutating - // tool changed; verification runs once over the union after the batch. - // A read-only tool (no ChangedFiles) never contributes. - if options.SelfCorrect != nil && toolResult.Status == tools.StatusOK && len(toolResult.ChangedFiles) > 0 { + // Collect the files this successful mutating tool changed. Self-correct + // verification runs once over the union after the batch; background + // diagnostics start NOW so the language server analyzes while the rest + // of the batch executes. A read-only tool (no ChangedFiles) never + // contributes. + if toolResult.Status == tools.StatusOK && len(toolResult.ChangedFiles) > 0 { changedFilesThisBatch = append(changedFilesThisBatch, toolResult.ChangedFiles...) + postEditDiagnostics.enqueue(ctx, toolResult.ChangedFiles) } } @@ -1100,8 +1120,12 @@ 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, + // Post-edit diagnostics are NOT forwarded to the tools: they used to run + // synchronously inside edit_file/write_file and block every mutating tool + // call on the language server (≥300ms debounce floor, 10s cap). The loop + // now checks changed files in the background and nudges the model before + // its next request instead (see async_diagnostics.go). The tools' inline + // mechanism (tools.RunOptions.Diagnostics) stays for direct API callers. // 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/types.go b/internal/agent/types.go index dd3968e38..1b12e0c78 100644 --- a/internal/agent/types.go +++ b/internal/agent/types.go @@ -285,11 +285,13 @@ 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, when set, checks files changed by mutating tools for + // error-severity language diagnostics IN THE BACKGROUND and appends any + // errors as a nudge before the model's next request — the model still sees + // an error it introduced at its next decision point, but no tool call ever + // blocks on the language server (the old inline path stalled every edit on + // a ≥300ms debounce, 10s cap). Build one with NewFileDiagnostics. nil + // disables post-edit diagnostics. FileDiagnostics func(ctx context.Context, absPath string) string // RequireCompletionSignal gates run completion for HEADLESS exec. Without it, diff --git a/internal/cli/exec.go b/internal/cli/exec.go index b75f729b7..8e3bee6a7 100644 --- a/internal/cli/exec.go +++ b/internal/cli/exec.go @@ -711,9 +711,10 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in // sessions); it is a no-op when self-correct is off. 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. + // always-on background post-edit diagnostics (agent.NewFileDiagnostics, + // collected off the tool-call path by the loop), 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) diff --git a/internal/tui/model.go b/internal/tui/model.go index c27406e11..d8b82f0c8 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -4446,9 +4446,9 @@ 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. + // Background post-edit diagnostics: the loop checks files changed by + // edit_file/write_file off the tool-call path and nudges the model with + // any errors before its next request. Shares the run's lazy manager. options.FileDiagnostics = agent.NewFileDiagnostics(lspManager, options.Cwd) } From 4e74b65dd9731415a25d4c10600fa0bc6837ff11 Mon Sep 17 00:00:00 2001 From: Kevin Codex Date: Sun, 5 Jul 2026 22:18:26 +0800 Subject: [PATCH 3/5] =?UTF-8?q?fix:=20address=20review=20=E2=80=94=20max-t?= =?UTF-8?q?urns=20diagnostics=20drain,=20capture-gap=20marker,=20test=20ha?= =?UTF-8?q?rdening?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Drain background diagnostics before the max-turns final-answer request so an error introduced by the LAST turn's edit is reported in the summary instead of silently dropped (+ regression test). - spillBashStreams marks the head/tail junction with a capture-gap marker when boundedBuffer dropped the middle of a stream, so a spilled log never reads as contiguous when it is not; junction position is asserted in the test. - Tests override TMP/TEMP alongside TMPDIR (os.TempDir reads TMP/TEMP on Windows) via a shared setTestTempDir helper. - TestBudgetBashCaptureReportsTrueTotal now asserts spill_path meta and spill file content; exemption-list test gains read_minified_file. --- internal/agent/async_diagnostics_test.go | 49 ++++++++++++++++++++++++ internal/agent/loop.go | 11 ++++++ internal/tools/bash.go | 27 ++++++++++--- internal/tools/bash_budget_test.go | 49 +++++++++++++++++++++--- internal/tools/output_ceiling_test.go | 5 ++- 5 files changed, 128 insertions(+), 13 deletions(-) diff --git a/internal/agent/async_diagnostics_test.go b/internal/agent/async_diagnostics_test.go index ca9d62e18..3ac09c9e8 100644 --- a/internal/agent/async_diagnostics_test.go +++ b/internal/agent/async_diagnostics_test.go @@ -221,3 +221,52 @@ func TestRunDeliversAsyncDiagnosticsNudgeNextTurn(t *testing.T) { t.Fatalf("checked path = %q, want %q", checkedPath, want) } } + +// When the run hits the maxTurns ceiling right after an edit turn, the +// final-answer request must still carry the diagnostics nudge — otherwise an +// error introduced by the last edit would go unreported in the summary. +func TestRunDrainsAsyncDiagnosticsBeforeMaxTurnsFinalAnswer(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(changedFilesTool{}) + provider := &mockProvider{ + turns: [][]zeroruntime.StreamEvent{ + { + {Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call-1", ToolName: "fake_edit"}, + {Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call-1", ArgumentsFragment: `{}`}, + {Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call-1"}, + {Type: zeroruntime.StreamEventDone}, + }, + { + {Type: zeroruntime.StreamEventText, Content: "summary"}, + {Type: zeroruntime.StreamEventDone}, + }, + }, + } + + result, err := Run(context.Background(), "fix main.go", provider, Options{ + Registry: registry, + Cwd: filepath.FromSlash("/ws"), + MaxTurns: 1, + FileDiagnostics: func(context.Context, string) string { + return "main.go:1:1 error: boom" + }, + }) + if err != nil { + t.Fatal(err) + } + if result.FinalAnswer != "summary" { + t.Fatalf("final answer = %q", result.FinalAnswer) + } + if len(provider.requests) != 2 { + t.Fatalf("expected the max-turns final-answer request, got %d requests", len(provider.requests)) + } + found := false + for _, message := range provider.requests[1].Messages { + if message.Role == zeroruntime.MessageRoleUser && strings.HasPrefix(message.Content, asyncDiagnosticsNudge) { + found = true + } + } + if !found { + t.Fatal("max-turns final-answer request missing the diagnostics nudge") + } +} diff --git a/internal/agent/loop.go b/internal/agent/loop.go index f162bb56e..5fa6f4fb2 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -697,6 +697,17 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) // mid-run, not a natural completion (a finished run returns via the no-tool-call // success path before this). Under the headless completion gate that is INCOMPLETE, // not success, so a run that loops to the turn limit isn't reported as done. + // + // The final-answer call is one more model request, so diagnostics from the + // LAST turn's edits get the same pre-request drain the loop gives every + // other turn — otherwise an error introduced by the final edit would go + // unreported in the summary. + if nudge := postEditDiagnostics.drain(ctx); nudge != "" { + messages = append(messages, zeroruntime.Message{ + Role: zeroruntime.MessageRoleUser, + Content: nudge, + }) + } if answer, finalMessages, finishReason := finalAnswerAfterMaxTurns(ctx, provider, messages, options); strings.TrimSpace(answer) != "" { result.FinalAnswer = answer result.FinishReason = finishReason diff --git a/internal/tools/bash.go b/internal/tools/bash.go index ba9af9ffd..7361c346f 100644 --- a/internal/tools/bash.go +++ b/internal/tools/bash.go @@ -405,7 +405,7 @@ func budgetBashCapture(out string, outTotal int, errStr string, errTotal int, me errText, errRaw, errTrunc := truncateHeadTailWithTotal(errStr, errTotal, bashOutputBudgetBytes) truncated := outTrunc || errTrunc if truncated { - if spillPath := spillBashStreams(out, errStr); spillPath != "" { + if spillPath := spillBashStreams(out, outTotal, errStr, errTotal); spillPath != "" { hint := "\n[zero] captured output saved to " + spillPath + " (grep or read_file it instead of re-running)" if errTrunc { errText += hint @@ -433,10 +433,13 @@ func budgetBashCapture(out string, outTotal int, errStr string, errTotal int, me // the spill directory as one sectioned file and returns its path, or "" when // spilling fails. The spill holds up to bashCaptureBudgetBytes of head and // tail per stream — everything zero kept in memory, which is more than the -// emit budget shows the model but not necessarily the whole output (the -// capture middle is discarded as it streams; the in-text truncation marker -// reports the true omitted size). -func spillBashStreams(stdout, stderr string) string { +// emit budget shows the model but not necessarily the whole output. When the +// capture itself dropped the middle of a stream (total exceeds the retained +// bytes), a gap marker is inserted at the head/tail junction so the spilled +// log never reads as contiguous when it is not. +func spillBashStreams(stdout string, stdoutTotal int, stderr string, stderrTotal int) string { + stdout = sectionWithCaptureGap(stdout, stdoutTotal) + stderr = sectionWithCaptureGap(stderr, stderrTotal) var combined strings.Builder combined.Grow(len(stdout) + len(stderr) + 64) combined.WriteString("### stdout\n") @@ -449,6 +452,20 @@ func spillBashStreams(stdout, stderr string) string { return spillTruncatedOutput("bash", combined.String()) } +// sectionWithCaptureGap marks the point where boundedBuffer dropped the middle +// of a stream. When total exceeds the retained bytes, the retained text is the +// frozen head (bashCaptureBudgetBytes, always full once overflow happened) +// followed immediately by the rolling tail — the junction sits at the head cap, +// snapped back to a rune boundary. +func sectionWithCaptureGap(text string, total int) string { + if total <= len(text) || len(text) <= bashCaptureBudgetBytes { + return text + } + head := utf8Prefix(text, bashCaptureBudgetBytes) + marker := fmt.Sprintf("\n[zero] capture gap: %d bytes omitted from the middle of this stream\n", total-len(text)) + return head + marker + text[len(head):] +} + // boundedBuffer is an io.Writer that retains at most headCap bytes from the start // and tailCap bytes from the end of a stream while counting the total written, so // a command emitting unbounded output (`cat huge.log`, `yes`) cannot grow Zero's diff --git a/internal/tools/bash_budget_test.go b/internal/tools/bash_budget_test.go index 9d40173ae..5c652212d 100644 --- a/internal/tools/bash_budget_test.go +++ b/internal/tools/bash_budget_test.go @@ -8,6 +8,17 @@ import ( "testing" ) +// setTestTempDir points every platform's temp-dir lookup at a per-test +// directory so spill files never land in the real temp dir: os.TempDir reads +// TMPDIR on Unix but TMP/TEMP on Windows. +func setTestTempDir(t *testing.T) { + t.Helper() + dir := t.TempDir() + t.Setenv("TMPDIR", dir) + t.Setenv("TMP", dir) + t.Setenv("TEMP", dir) +} + // Small output passes through untouched and records raw==emitted, no truncated flag. func TestBudgetBashOutputSmallPassesThrough(t *testing.T) { meta := map[string]string{} @@ -30,7 +41,7 @@ func TestBudgetBashOutputSmallPassesThrough(t *testing.T) { // the middle is dropped behind a marker, meta is flagged, and the captured text // is spilled to a re-readable file. func TestBudgetBashOutputTruncatesHeadAndTail(t *testing.T) { - t.Setenv("TMPDIR", t.TempDir()) + setTestTempDir(t) head := "FIRST_LINE_MARKER\n" tail := "\nLAST_LINE_MARKER" big := head + strings.Repeat("x", bashOutputBudgetBytes) + tail @@ -117,11 +128,12 @@ func TestBoundedBufferKeepsHeadAndTailBounded(t *testing.T) { // budgetBashCapture reports the TRUE total (not the retained size) in the marker // and raw_bytes, even though only a bounded head+tail was ever held in memory. func TestBudgetBashCaptureReportsTrueTotal(t *testing.T) { - t.Setenv("TMPDIR", t.TempDir()) - // Retained head+tail as boundedBuffer would hand over; the real command produced - // far more than was kept. - retained := "HEAD_START" + strings.Repeat("y", bashOutputBudgetBytes) + "TAIL_END" - total := 10 * bashOutputBudgetBytes + setTestTempDir(t) + // Retained head+tail as boundedBuffer would hand over after overflow: the + // frozen head is full at bashCaptureBudgetBytes, then the rolling tail; the + // real command produced far more than was kept. + retained := "HEAD_START" + strings.Repeat("y", bashCaptureBudgetBytes) + "TAIL_END" + total := 10 * bashCaptureBudgetBytes meta := map[string]string{} out, _, truncated := budgetBashCapture(retained, total, "", 0, meta) @@ -143,4 +155,29 @@ func TestBudgetBashCaptureReportsTrueTotal(t *testing.T) { if !strings.Contains(out, strconv.Itoa(total-bashOutputBudgetBytes)) { t.Fatalf("marker should cite the true omitted byte count:\n%s", out[:min(200, len(out))]) } + + // The spill must be recorded, sectioned, and carry a capture-gap marker at + // the head/tail junction — the capture dropped the middle, so the spilled + // log must never read as contiguous. + if meta["spill_path"] == "" { + t.Fatal("spill path missing from meta") + } + content, err := os.ReadFile(meta["spill_path"]) + if err != nil { + t.Fatalf("spill file unreadable: %v", err) + } + if !strings.Contains(string(content), "### stdout") { + t.Fatal("spill must be sectioned by stream") + } + gapMarker := fmt.Sprintf("capture gap: %d bytes omitted", total-len(retained)) + gapIndex := strings.Index(string(content), gapMarker) + if gapIndex < 0 { + t.Fatalf("spill missing capture-gap marker %q", gapMarker) + } + // The junction sits right after the frozen head (plus the section header), + // not appended at the end of the stream. + headerOffset := len("### stdout\n") + if gapIndex < headerOffset+bashCaptureBudgetBytes-64 || gapIndex > headerOffset+bashCaptureBudgetBytes+64 { + t.Fatalf("capture-gap marker at offset %d, want ~%d (the head cap)", gapIndex, headerOffset+bashCaptureBudgetBytes) + } } diff --git a/internal/tools/output_ceiling_test.go b/internal/tools/output_ceiling_test.go index 40995a7fb..bf65d5712 100644 --- a/internal/tools/output_ceiling_test.go +++ b/internal/tools/output_ceiling_test.go @@ -35,7 +35,7 @@ func newCeilingFakeTool(name, output string) ceilingFakeTool { // An unbudgeted tool's oversized output is capped at the universal ceiling, // head and tail survive, and the full output is spilled to a re-readable file. func TestOutputCeilingCapsUnbudgetedTool(t *testing.T) { - t.Setenv("TMPDIR", t.TempDir()) + setTestTempDir(t) big := "HEAD_MARK\n" + strings.Repeat("z", defaultOutputCeilingTokens*4*3) + "\nTAIL_MARK" registry := NewRegistry() @@ -98,7 +98,7 @@ func TestOutputCeilingSmallOutputUntouched(t *testing.T) { // ZERO_TOOL_OUTPUT_CEILING_TOKENS tightens, loosens, or disables the ceiling. func TestOutputCeilingEnvOverride(t *testing.T) { - t.Setenv("TMPDIR", t.TempDir()) + setTestTempDir(t) big := strings.Repeat("z", 100*1024) t.Setenv(outputCeilingEnv, "1000") @@ -124,6 +124,7 @@ func TestSelfBudgetingExemptionList(t *testing.T) { NewBashTool(dir), NewExecCommandTool(dir, newExecSessionManager()), NewReadFileTool(dir), + NewReadMinifiedFileTool(dir), NewGrepTool(dir), NewGlobTool(dir), NewListDirectoryTool(dir), From fc89f2a36d617ce47dff5796da131603a6970610 Mon Sep 17 00:00:00 2001 From: Kevin Codex Date: Sun, 5 Jul 2026 22:47:03 +0800 Subject: [PATCH 4/5] fix(test): use a real workspace root in async diagnostics tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit filepath.IsAbs("/ws") is false on Windows (no drive letter), so the collector test's workspace-absolute assertion failed on the windows smoke job. Root the fake workspace at t.TempDir() — absolute on every platform — and assert against it. --- internal/agent/async_diagnostics_test.go | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/internal/agent/async_diagnostics_test.go b/internal/agent/async_diagnostics_test.go index 3ac09c9e8..a6cd83cf5 100644 --- a/internal/agent/async_diagnostics_test.go +++ b/internal/agent/async_diagnostics_test.go @@ -24,6 +24,9 @@ func TestAsyncDiagnosticsNilCollectorNoOps(t *testing.T) { } func TestAsyncDiagnosticsCollectsAndDrainsOnce(t *testing.T) { + // A real absolute root: a literal "/ws" is not filepath.IsAbs on Windows + // (no drive letter), which is exactly what the assertions below check. + root := t.TempDir() var mu sync.Mutex var checked []string check := func(_ context.Context, absPath string) string { @@ -35,7 +38,7 @@ func TestAsyncDiagnosticsCollectsAndDrainsOnce(t *testing.T) { } return "ERR " + filepath.Base(absPath) } - diagnostics := newAsyncDiagnostics(check, "/ws") + diagnostics := newAsyncDiagnostics(check, root) diagnostics.enqueue(context.Background(), []string{"broken.go", "clean.go"}) nudge := diagnostics.drain(context.Background()) @@ -47,8 +50,8 @@ func TestAsyncDiagnosticsCollectsAndDrainsOnce(t *testing.T) { } mu.Lock() for _, path := range checked { - if !filepath.IsAbs(path) || !strings.HasPrefix(path, filepath.FromSlash("/ws")) { - t.Fatalf("check received %q, want workspace-absolute path", path) + if !filepath.IsAbs(path) || !strings.HasPrefix(path, root+string(filepath.Separator)) { + t.Fatalf("check received %q, want path under workspace root %q", path, root) } } mu.Unlock() @@ -173,12 +176,13 @@ func TestRunDeliversAsyncDiagnosticsNudgeNextTurn(t *testing.T) { }, }, } + root := t.TempDir() var checkedPath string var mu sync.Mutex result, err := Run(context.Background(), "fix main.go", provider, Options{ Registry: registry, - Cwd: filepath.FromSlash("/ws"), + Cwd: root, FileDiagnostics: func(_ context.Context, absPath string) string { mu.Lock() checkedPath = absPath @@ -217,7 +221,7 @@ func TestRunDeliversAsyncDiagnosticsNudgeNextTurn(t *testing.T) { } mu.Lock() defer mu.Unlock() - if want := filepath.Join(filepath.FromSlash("/ws"), "main.go"); checkedPath != want { + if want := filepath.Join(root, "main.go"); checkedPath != want { t.Fatalf("checked path = %q, want %q", checkedPath, want) } } @@ -245,7 +249,7 @@ func TestRunDrainsAsyncDiagnosticsBeforeMaxTurnsFinalAnswer(t *testing.T) { result, err := Run(context.Background(), "fix main.go", provider, Options{ Registry: registry, - Cwd: filepath.FromSlash("/ws"), + Cwd: t.TempDir(), MaxTurns: 1, FileDiagnostics: func(context.Context, string) string { return "main.go:1:1 error: boom" From 07af8c451d9c2f054dcc75e3f5853e9ff7c3bcfd Mon Sep 17 00:00:00 2001 From: Kevin Codex Date: Sun, 5 Jul 2026 23:06:10 +0800 Subject: [PATCH 5/5] fix: pattern-scrub spill files and gate finalization on pending diagnostics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups (Vasanthdev2004): - Spill files now pass the pattern-based secret scanner (AWS keys, GitHub/Slack/OpenAI tokens, PEM blocks, JWTs) in addition to the configured-key scrub. A bash spill runs before formatBashOutput's scan, so without this a spilled file held in cleartext exactly the credentials the transcript redacts. Applied centrally in spillTruncatedOutput, covering bash, exec_command, and ceiling spills. - Finalization diagnostics gate: a no-tool-call final answer now drains pending post-edit diagnostics with the full inline-era budget (10s, vs the 3s per-turn wait) before the run is accepted; introduced errors append the nudge and give the model one more turn instead of being silently dropped when no later turn exists. The max-turns summary path uses the same finalization budget. Free for runs that never edited — an idle collector returns immediately. --- internal/agent/async_diagnostics.go | 19 +++++++- internal/agent/async_diagnostics_test.go | 62 ++++++++++++++++++++++++ internal/agent/loop.go | 21 ++++++-- internal/tools/spill.go | 8 ++- internal/tools/spill_test.go | 27 +++++++++++ 5 files changed, 130 insertions(+), 7 deletions(-) diff --git a/internal/agent/async_diagnostics.go b/internal/agent/async_diagnostics.go index a1a0dd9b9..544d07df1 100644 --- a/internal/agent/async_diagnostics.go +++ b/internal/agent/async_diagnostics.go @@ -36,6 +36,13 @@ import ( // A var so tests can shorten the wait. var asyncDiagnosticsDrainTimeout = 3 * time.Second +// asyncDiagnosticsFinalDrainTimeout is the wait used when the run is about to +// FINALIZE (natural completion or the max-turns summary): there is no later +// turn to defer to, so the gate waits out the old inline-era budget rather +// than silently dropping an error the last edit introduced. A var so tests +// can shorten the wait. +var asyncDiagnosticsFinalDrainTimeout = 10 * time.Second + // asyncDiagnosticsNudge prefixes the drained diagnostics blocks; phrased like // the old inline block so the model reacts the same way. const asyncDiagnosticsNudge = "Diagnostics after your recent edits (fix any errors you introduced):\n" @@ -125,6 +132,16 @@ func (diagnostics *asyncDiagnostics) work(ctx context.Context, done chan struct{ // all completed error blocks as one nudge. Returns "" when there is nothing // to report or the worker is still busy (results then surface next turn). func (diagnostics *asyncDiagnostics) drain(ctx context.Context) string { + return diagnostics.drainWithin(ctx, asyncDiagnosticsDrainTimeout) +} + +// drainFinal is drain with the finalization budget: called before the run's +// last model request, where an undelivered result would otherwise be lost. +func (diagnostics *asyncDiagnostics) drainFinal(ctx context.Context) string { + return diagnostics.drainWithin(ctx, asyncDiagnosticsFinalDrainTimeout) +} + +func (diagnostics *asyncDiagnostics) drainWithin(ctx context.Context, timeout time.Duration) string { if diagnostics == nil { return "" } @@ -132,7 +149,7 @@ func (diagnostics *asyncDiagnostics) drain(ctx context.Context) string { busy := diagnostics.working diagnostics.mu.Unlock() if busy != nil { - timer := time.NewTimer(asyncDiagnosticsDrainTimeout) + timer := time.NewTimer(timeout) defer timer.Stop() select { case <-busy: diff --git a/internal/agent/async_diagnostics_test.go b/internal/agent/async_diagnostics_test.go index a6cd83cf5..6dba03419 100644 --- a/internal/agent/async_diagnostics_test.go +++ b/internal/agent/async_diagnostics_test.go @@ -274,3 +274,65 @@ func TestRunDrainsAsyncDiagnosticsBeforeMaxTurnsFinalAnswer(t *testing.T) { t.Fatal("max-turns final-answer request missing the diagnostics nudge") } } + +// When the model finalizes while a slow check is still in flight (the +// per-turn drain missed it), the finalization gate must wait it out and give +// the model one more turn with the nudge instead of dropping the errors. +func TestRunFinalizationGateDeliversLateDiagnostics(t *testing.T) { + previous := asyncDiagnosticsDrainTimeout + asyncDiagnosticsDrainTimeout = time.Millisecond + defer func() { asyncDiagnosticsDrainTimeout = previous }() + + registry := tools.NewRegistry() + registry.Register(changedFilesTool{}) + provider := &mockProvider{ + turns: [][]zeroruntime.StreamEvent{ + { + {Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call-1", ToolName: "fake_edit"}, + {Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call-1", ArgumentsFragment: `{}`}, + {Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call-1"}, + {Type: zeroruntime.StreamEventDone}, + }, + { + {Type: zeroruntime.StreamEventText, Content: "done"}, + {Type: zeroruntime.StreamEventDone}, + }, + { + {Type: zeroruntime.StreamEventText, Content: "done after fix"}, + {Type: zeroruntime.StreamEventDone}, + }, + }, + } + + result, err := Run(context.Background(), "fix main.go", provider, Options{ + Registry: registry, + Cwd: t.TempDir(), + FileDiagnostics: func(context.Context, string) string { + // Slower than the (shortened) per-turn drain, well under the + // finalization budget. + time.Sleep(100 * time.Millisecond) + return "main.go:1:1 error: boom" + }, + }) + if err != nil { + t.Fatal(err) + } + if result.FinalAnswer != "done after fix" { + t.Fatalf("final answer = %q, want the post-nudge answer", result.FinalAnswer) + } + if len(provider.requests) != 3 { + t.Fatalf("expected a third request carrying the late nudge, got %d requests", len(provider.requests)) + } + found := false + for _, message := range provider.requests[2].Messages { + if message.Role == zeroruntime.MessageRoleUser && strings.HasPrefix(message.Content, asyncDiagnosticsNudge) { + if !strings.Contains(message.Content, "boom") { + t.Fatalf("nudge missing diagnostics: %q", message.Content) + } + found = true + } + } + if !found { + t.Fatal("finalization gate did not deliver the late diagnostics nudge") + } +} diff --git a/internal/agent/loop.go b/internal/agent/loop.go index 5fa6f4fb2..dfa0c2519 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -491,6 +491,19 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) continue } } + // Finalization diagnostics gate: edits from this run may still have + // checks in flight — the per-turn drain waits only briefly and defers + // to "a later turn", but a final answer means there is no later turn. + // Wait out the full inline-era budget once; an introduced error gives + // the model one more turn to see (and fix) it instead of being lost. + // Free for runs that never edited: an idle collector returns "". + if nudge := postEditDiagnostics.drainFinal(ctx); nudge != "" { + messages = append(messages, zeroruntime.Message{ + Role: zeroruntime.MessageRoleUser, + Content: nudge, + }) + continue + } result.FinalAnswer = collected.Text result.Messages = copyMessages(messages) return result, nil @@ -699,10 +712,10 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) // not success, so a run that loops to the turn limit isn't reported as done. // // The final-answer call is one more model request, so diagnostics from the - // LAST turn's edits get the same pre-request drain the loop gives every - // other turn — otherwise an error introduced by the final edit would go - // unreported in the summary. - if nudge := postEditDiagnostics.drain(ctx); nudge != "" { + // LAST turn's edits get a drain too — with the finalization budget, since + // there is no later turn to defer to — otherwise an error introduced by + // the final edit would go unreported in the summary. + if nudge := postEditDiagnostics.drainFinal(ctx); nudge != "" { messages = append(messages, zeroruntime.Message{ Role: zeroruntime.MessageRoleUser, Content: nudge, diff --git a/internal/tools/spill.go b/internal/tools/spill.go index f7f49092d..0282c9e49 100644 --- a/internal/tools/spill.go +++ b/internal/tools/spill.go @@ -8,6 +8,7 @@ import ( "time" "github.com/Gitlawb/zero/internal/redaction" + "github.com/Gitlawb/zero/internal/secrets" ) // Spill-to-disk for truncated tool output. When a command produces more than @@ -90,8 +91,10 @@ func spillDir() (string, error) { // 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. +// tool boundary PLUS the pattern-based secret scanner (AWS keys, tokens, PEM +// blocks, JWTs) that bash applies to its model-visible output — a spill runs +// before that formatter, so without the scan here a spilled file would hold +// pattern-matched credentials in cleartext that the transcript hides. func spillTruncatedOutput(toolName, output string) string { dir, err := spillDir() if err != nil { @@ -110,6 +113,7 @@ func spillTruncatedOutput(toolName, output string) string { } defer file.Close() scrubbed := redaction.RedactString(output, redaction.Options{}) + scrubbed, _ = secrets.Redact(scrubbed) if _, err := file.WriteString(scrubbed); err != nil { _ = os.Remove(file.Name()) return "" diff --git a/internal/tools/spill_test.go b/internal/tools/spill_test.go index a04e498a0..bd29d3550 100644 --- a/internal/tools/spill_test.go +++ b/internal/tools/spill_test.go @@ -63,6 +63,33 @@ func TestSweepSpillDirRemovesOnlyOldFiles(t *testing.T) { } } +// A spill happens BEFORE bash's model-facing formatter runs its pattern-based +// secret scan, so the spill itself must scrub pattern-matched credentials — +// otherwise the file would hold in cleartext exactly what the transcript hides. +func TestSpillTruncatedOutputScrubsPatternSecrets(t *testing.T) { + setTestTempDir(t) + githubToken := "ghp_" + strings.Repeat("a", 36) + body := "before\nAKIAIOSFODNN7EXAMPLE\n" + githubToken + "\nafter" + + path := spillTruncatedOutput("bash", body) + if path == "" { + t.Fatal("spill must return a file path") + } + content, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + if strings.Contains(string(content), "AKIAIOSFODNN7EXAMPLE") { + t.Fatal("AWS access key reached the spill file in cleartext") + } + if strings.Contains(string(content), githubToken) { + t.Fatal("GitHub token reached the spill file in cleartext") + } + if !strings.Contains(string(content), "before") || !strings.Contains(string(content), "after") { + t.Fatalf("non-secret content must survive the scrub: %q", content) + } +} + func TestSpillTruncatedOutputWritesFile(t *testing.T) { t.Setenv("TMPDIR", t.TempDir()) path := spillTruncatedOutput("exec_command", "some output body")