Compare commits

...
Author SHA1 Message Date
Kevin Codex 07af8c451d fix: pattern-scrub spill files and gate finalization on pending diagnostics
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.
2026-07-05 23:06:10 +08:00
Kevin Codex fc89f2a36d fix(test): use a real workspace root in async diagnostics tests
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.
2026-07-05 22:47:03 +08:00
Kevin Codex 4e74b65dd9 fix: address review — max-turns diagnostics drain, capture-gap marker, test hardening
- 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.
2026-07-05 22:18:26 +08:00
Kevin Codex 541acf9b43 perf(agent): collect post-edit LSP diagnostics off the tool-call path
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.
2026-07-05 22:14:31 +08:00
Kevin Codex 3fa8b39672 feat(tools): universal token-based output ceiling with spill-to-disk
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.
2026-07-05 22:14:31 +08:00
13 changed files with 988 additions and 36 deletions

No files matched your search

+185
View File
@@ -0,0 +1,185 @@
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
// 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"
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 {
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 ""
}
diagnostics.mu.Lock()
busy := diagnostics.working
diagnostics.mu.Unlock()
if busy != nil {
timer := time.NewTimer(timeout)
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)
}
+338
View File
@@ -0,0 +1,338 @@
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) {
// 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 {
mu.Lock()
checked = append(checked, absPath)
mu.Unlock()
if filepath.Base(absPath) == "clean.go" {
return ""
}
return "ERR " + filepath.Base(absPath)
}
diagnostics := newAsyncDiagnostics(check, root)
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, root+string(filepath.Separator)) {
t.Fatalf("check received %q, want path under workspace root %q", path, root)
}
}
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},
},
},
}
root := t.TempDir()
var checkedPath string
var mu sync.Mutex
result, err := Run(context.Background(), "fix main.go", provider, Options{
Registry: registry,
Cwd: root,
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(root, "main.go"); checkedPath != want {
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: t.TempDir(),
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")
}
}
// 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")
}
}
+54 -6
View File
@@ -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
@@ -474,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
@@ -590,11 +620,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)
}
}
@@ -677,6 +710,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 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,
})
}
if answer, finalMessages, finishReason := finalAnswerAfterMaxTurns(ctx, provider, messages, options); strings.TrimSpace(answer) != "" {
result.FinalAnswer = answer
result.FinishReason = finishReason
@@ -1100,8 +1144,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,
+7 -5
View File
@@ -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,
+4 -3
View File
@@ -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)
+67 -8
View File
@@ -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, 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
} 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,43 @@ 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. 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")
combined.WriteString(stdout)
if !strings.HasSuffix(stdout, "\n") {
combined.WriteString("\n")
}
combined.WriteString("### stderr\n")
combined.WriteString(stderr)
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
+63 -7
View File
@@ -2,11 +2,23 @@ package tools
import (
"fmt"
"os"
"strconv"
"strings"
"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{}
@@ -26,8 +38,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) {
setTestTempDir(t)
head := "FIRST_LINE_MARKER\n"
tail := "\nLAST_LINE_MARKER"
big := head + strings.Repeat("x", bashOutputBudgetBytes) + tail
@@ -47,9 +61,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,10 +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) {
// 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)
@@ -114,7 +144,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) {
@@ -124,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)
}
}
+81
View File
@@ -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
}
+140
View File
@@ -0,0 +1,140 @@
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) {
setTestTempDir(t)
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) {
setTestTempDir(t)
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),
NewReadMinifiedFileTool(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")
}
}
+13 -2
View File
@@ -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
+6 -2
View File
@@ -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 ""
+27
View File
@@ -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")
+3 -3
View File
@@ -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)
}