mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-08 01:09:46 -10:00
Compare commits
5
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
07af8c451d | ||
|
|
fc89f2a36d | ||
|
|
4e74b65dd9 | ||
|
|
541acf9b43 | ||
|
|
3fa8b39672 |
No files matched your search
@@ -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)
|
||||
}
|
||||
@@ -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
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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 ""
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user