From 05306b99a34cc657ea2fddc41f8737495949f7b6 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 18:01:27 +0530 Subject: [PATCH 01/32] feat(tui): actionable hints on provider/model errors Classify a failed turn's provider error into auth / rate-limit / connectivity / model-not-found / context-overflow and render a one-line next step below the red error row, instead of dumping an identical raw blob for every failure mode. The classifier lives in a new internal/errhint package so both the TUI error row and the CLI exec path can share it (TUIHint references slash commands, CLIHint references zero subcommands). Classification is a conservative string heuristic keyed off providerio.ClassifiedError's prefixes plus lower-level DNS/TLS/timeout/context-length signatures, since the numeric HTTP status is gone by the time the error reaches a UI surface. --- internal/errhint/errhint.go | 109 +++++++++++++++++++++++++++++++ internal/errhint/errhint_test.go | 80 +++++++++++++++++++++++ internal/tui/error_hint_test.go | 27 ++++++++ internal/tui/model.go | 6 +- internal/tui/rendering.go | 6 ++ internal/tui/transcript.go | 1 + 6 files changed, 228 insertions(+), 1 deletion(-) create mode 100644 internal/errhint/errhint.go create mode 100644 internal/errhint/errhint_test.go create mode 100644 internal/tui/error_hint_test.go diff --git a/internal/errhint/errhint.go b/internal/errhint/errhint.go new file mode 100644 index 00000000..e29eaae3 --- /dev/null +++ b/internal/errhint/errhint.go @@ -0,0 +1,109 @@ +// Package errhint classifies provider/model failures into a few user-actionable +// categories and turns them into a one-line "next step" hint. +// +// Provider errors already arrive with a classified string prefix from +// providerio.ClassifiedError ("auth error:", "rate limit error:", …); lower-level +// failures (DNS, TLS, timeouts, context-length) arrive as raw driver or library +// messages. Classify matches both so the interactive error row (TUI) and the +// `zero exec` provider-error path (CLI) can append one concrete next step instead +// of dumping an identical red blob for every failure mode. +package errhint + +import "strings" + +// Category buckets a provider/model failure into a small set of classes that each +// map to a distinct recovery action. +type Category int + +const ( + // Unknown means the error didn't match any known signature; callers should + // emit no hint rather than guess. + Unknown Category = iota + Auth + RateLimit + Connectivity + ModelNotFound + ContextOverflow +) + +// Classify buckets err by scanning its message for known signatures. It is a +// deliberately conservative string heuristic — the provider layer has already +// discarded the numeric HTTP status by the time the error reaches a UI surface, +// so the message is all we have. Order matters: more specific signatures are +// tested before broader ones (e.g. "context length" as overflow before the +// generic "timeout" as connectivity). +func Classify(err error) Category { + if err == nil { + return Unknown + } + m := strings.ToLower(err.Error()) + switch { + case containsAny(m, "auth error:", "unauthorized", "api key", "api_key", "invalid_api_key", + "authentication", "permission denied", "forbidden", "401", "403"): + return Auth + case containsAny(m, "rate limit", "rate_limit", "too many requests", "quota", + "resource_exhausted", "overloaded", "429", "529"): + return RateLimit + case containsAny(m, "context length", "context window", "maximum context", "context_length_exceeded", + "too many tokens", "prompt is too long", "reduce the length", "maximum context length"): + return ContextOverflow + case containsAny(m, "model not found", "model_not_found", "does not exist", "unknown model", + "no such model", "unsupported model", "invalid model", "model is not"): + return ModelNotFound + case containsAny(m, "dial tcp", "no such host", "connection refused", "network is unreachable", + "i/o timeout", "context deadline exceeded", "tls handshake", "connection reset", + "unexpected eof", "lookup ", "timeout"): + return Connectivity + default: + return Unknown + } +} + +// TUIHint returns a one-line hint referencing interactive slash commands, or "" +// when the category is Unknown. Meant to sit under the raw error in the live +// error row. +func TUIHint(err error) string { + switch Classify(err) { + case Auth: + return "API key rejected — run /provider to re-check your credentials" + case RateLimit: + return "Rate limited — wait a moment, or switch model with /model" + case Connectivity: + return "Can't reach the provider — run /doctor --connectivity" + case ModelNotFound: + return "Model unavailable — pick another with /model" + case ContextOverflow: + return "Context window full — run /compact to free space" + default: + return "" + } +} + +// CLIHint returns a one-line hint referencing `zero …` subcommands, or "" when the +// category is Unknown. Meant for the non-interactive `zero exec` error path, where +// slash commands don't apply. +func CLIHint(err error) string { + switch Classify(err) { + case Auth: + return "API key rejected — run `zero auth` or set the provider's API key, then retry" + case RateLimit: + return "Rate limited — wait a moment, or switch model with --model" + case Connectivity: + return "Can't reach the provider — run `zero doctor`" + case ModelNotFound: + return "Model unavailable — run `zero doctor` or pick another with --model" + case ContextOverflow: + return "Context window full — shorten the prompt or start a fresh session" + default: + return "" + } +} + +func containsAny(haystack string, needles ...string) bool { + for _, n := range needles { + if strings.Contains(haystack, n) { + return true + } + } + return false +} diff --git a/internal/errhint/errhint_test.go b/internal/errhint/errhint_test.go new file mode 100644 index 00000000..2a6c2866 --- /dev/null +++ b/internal/errhint/errhint_test.go @@ -0,0 +1,80 @@ +package errhint + +import ( + "errors" + "strings" + "testing" +) + +func TestClassify(t *testing.T) { + cases := []struct { + name string + msg string + want Category + }{ + {"nil-ish empty", "", Unknown}, + {"providerio auth prefix", "auth error: your API key is missing or invalid — run `zero auth`", Auth}, + {"raw 401", "provider request error: 401 Unauthorized", Auth}, + {"invalid api key", "Error: invalid_api_key: incorrect key provided", Auth}, + {"rate limit prefix", "rate limit error: 429 too many requests", RateLimit}, + {"overloaded", "provider error: model is overloaded, please retry", RateLimit}, + {"resource exhausted gemini", "rpc error: code = ResourceExhausted desc = quota exceeded", RateLimit}, + {"context length openai", "This model's maximum context length is 128000 tokens", ContextOverflow}, + {"context_length_exceeded", "provider request error: context_length_exceeded", ContextOverflow}, + {"model not found", "provider request error: 404 the model `gpt-9` does not exist", ModelNotFound}, + {"unknown model", "unknown model: sonnet-99", ModelNotFound}, + {"dns failure", "Post \"https://api...\": dial tcp: lookup api.foo.com: no such host", Connectivity}, + {"connection refused", "dial tcp 127.0.0.1:443: connection refused", Connectivity}, + {"deadline", "context deadline exceeded (Client.Timeout exceeded)", Connectivity}, + {"gibberish", "provider error: something totally unexpected happened", Unknown}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := Classify(errors.New(tc.msg)); got != tc.want { + t.Fatalf("Classify(%q) = %v, want %v", tc.msg, got, tc.want) + } + }) + } +} + +func TestClassifyNil(t *testing.T) { + if got := Classify(nil); got != Unknown { + t.Fatalf("Classify(nil) = %v, want Unknown", got) + } +} + +// Context-overflow must win over the connectivity "timeout" catch-all: a message +// mentioning both "context length" and a timeout is an overflow, not a network +// problem. +func TestContextOverflowBeatsConnectivity(t *testing.T) { + err := errors.New("maximum context length exceeded; request timeout") + if got := Classify(err); got != ContextOverflow { + t.Fatalf("Classify = %v, want ContextOverflow", got) + } +} + +func TestHintsPresentForKnownCategories(t *testing.T) { + known := []error{ + errors.New("auth error: bad key"), + errors.New("rate limit error: 429"), + errors.New("dial tcp: no such host"), + errors.New("model does not exist"), + errors.New("maximum context length is 128000 tokens"), + } + for _, err := range known { + if h := TUIHint(err); strings.TrimSpace(h) == "" { + t.Fatalf("TUIHint(%q) empty, want a hint", err) + } + if h := CLIHint(err); strings.TrimSpace(h) == "" { + t.Fatalf("CLIHint(%q) empty, want a hint", err) + } + } + // Unknown yields no hint on either surface. + unknown := errors.New("provider error: mystery") + if h := TUIHint(unknown); h != "" { + t.Fatalf("TUIHint(unknown) = %q, want empty", h) + } + if h := CLIHint(unknown); h != "" { + t.Fatalf("CLIHint(unknown) = %q, want empty", h) + } +} diff --git a/internal/tui/error_hint_test.go b/internal/tui/error_hint_test.go new file mode 100644 index 00000000..b58a4e21 --- /dev/null +++ b/internal/tui/error_hint_test.go @@ -0,0 +1,27 @@ +package tui + +import ( + "strings" + "testing" +) + +// A hinted error row renders the raw error and the faint one-line next step +// below it; an unhinted row renders just the error. +func TestRenderErrorRowShowsHint(t *testing.T) { + withHint := renderErrorRow(transcriptRow{ + kind: rowError, + text: "auth error: your API key is missing or invalid", + hint: "API key rejected — run /provider to re-check your credentials", + }, 80) + if !strings.Contains(withHint, "auth error:") { + t.Fatalf("expected raw error text in output, got:\n%s", withHint) + } + if !strings.Contains(withHint, "/provider") { + t.Fatalf("expected hint referencing /provider, got:\n%s", withHint) + } + + noHint := renderErrorRow(transcriptRow{kind: rowError, text: "provider error: mystery"}, 80) + if strings.Contains(noHint, "→") { + t.Fatalf("unhinted error must not render a hint arrow, got:\n%s", noHint) + } +} diff --git a/internal/tui/model.go b/internal/tui/model.go index 766f9ab9..480d601c 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -19,6 +19,7 @@ import ( "github.com/Gitlawb/zero/internal/agent" "github.com/Gitlawb/zero/internal/config" "github.com/Gitlawb/zero/internal/doctor" + "github.com/Gitlawb/zero/internal/errhint" "github.com/Gitlawb/zero/internal/lsp" internalmcp "github.com/Gitlawb/zero/internal/mcp" "github.com/Gitlawb/zero/internal/modelregistry" @@ -1847,10 +1848,13 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) { m.transcript = appendTranscriptRow(m.transcript, transcriptRow{kind: rowAssistant, text: text}) } // The error row terminates the turn, so it carries the done-line - // metadata a final assistant row would have carried. + // metadata a final assistant row would have carried. A recognized + // provider failure (auth/rate-limit/connectivity/…) also carries a + // one-line next step so the user isn't left staring at a raw blob. m.transcript = appendTranscriptRow(m.transcript, transcriptRow{ kind: rowError, text: msg.err.Error(), + hint: errhint.TUIHint(msg.err), final: true, turnTools: msg.turnTools, turnElapsed: msg.turnElapsed, diff --git a/internal/tui/rendering.go b/internal/tui/rendering.go index 48e34da7..4569a255 100644 --- a/internal/tui/rendering.go +++ b/internal/tui/rendering.go @@ -987,6 +987,12 @@ func isDoctorResultHeading(value string) bool { func renderErrorRow(row transcriptRow, width int) string { note := noteBox(row.text, width, zeroTheme.cardErr, zeroTheme.red) + // A recognized failure carries a one-line next step. Render it just below the + // red box in the faint metadata style so it reads as guidance, not more error + // text (and to avoid nesting ANSI styles inside noteBox's per-line red wrap). + if hint := strings.TrimSpace(row.hint); hint != "" { + note += "\n" + fitStyledLine(zeroTheme.faint.Render("→ "+hint), width) + } return note } diff --git a/internal/tui/transcript.go b/internal/tui/transcript.go index 431af311..a31f7b5a 100644 --- a/internal/tui/transcript.go +++ b/internal/tui/transcript.go @@ -35,6 +35,7 @@ type transcriptRow struct { tool string // tool name, for tool call/result rows status tools.Status // result status, for tool result rows detail string // raw multi-line output (e.g. a diff to render as a card) + hint string // one-line actionable hint, rendered faintly below error rows arg string // secondary argument hint (pattern/command), for tool call rows runID int // owning run, for tool call rows (0 = rehydrated/unknown) permission *agent.PermissionEvent From 0a164e08fb0b2e09acbdf057b9ff461d1c6d686c Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 18:07:46 +0530 Subject: [PATCH 02/32] feat(tools): budget bash output with head+tail truncation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit bash was the one tool that returned command output with no byte cap — 'cat large.log' or a verbose test run could dump megabytes into context and force compaction, while every other read/search tool already applies a budget. Cap stdout and stderr at 96KB each, keeping the head and tail of an oversized stream (build/test failures usually land at the tail) and dropping the middle behind a marker that points at redirect-to-file + read_file. Record raw_bytes/emitted_bytes/estimated_tokens/truncated in Meta like the other tools. Sandbox-denial and shell-issue detection still run on the full raw output before budgeting. --- internal/tools/bash.go | 53 ++++++++++++++++++++++++++-- internal/tools/bash_budget_test.go | 55 ++++++++++++++++++++++++++++++ 2 files changed, 105 insertions(+), 3 deletions(-) create mode 100644 internal/tools/bash_budget_test.go diff --git a/internal/tools/bash.go b/internal/tools/bash.go index 090594fd..b0329771 100644 --- a/internal/tools/bash.go +++ b/internal/tools/bash.go @@ -157,24 +157,26 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS } } markLikelySandboxDenial(meta, plan, exitCode, stdout.String(), stderrText) + outText, errText := budgetBashOutput(stdout.String(), stderrText, meta) return Result{ Status: StatusError, - Output: formatBashOutputWithShellHint(commandText, stdout.String(), stderrText, exitCode, meta), + Output: formatBashOutputWithShellHint(commandText, outText, errText, exitCode, meta), Meta: meta, } } markLikelySandboxDenial(meta, plan, exitCode, stdout.String(), stderrText) + outText, errText := budgetBashOutput(stdout.String(), stderrText, meta) if meta[SandboxLikelyDeniedMeta] == "true" { return Result{ Status: StatusError, - Output: formatBashOutputWithShellHint(commandText, stdout.String(), stderrText, exitCode, meta), + Output: formatBashOutputWithShellHint(commandText, outText, errText, exitCode, meta), Meta: meta, } } return Result{ Status: StatusOK, - Output: formatBashOutput(stdout.String(), stderrText, exitCode), + Output: formatBashOutput(outText, errText, exitCode), Meta: meta, } } @@ -345,6 +347,51 @@ 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 + +// budgetBashOutput truncates stdout and stderr to bashOutputBudgetBytes each, +// keeping the head and tail of anything larger, and records raw/emitted byte +// counts plus a truncated flag in meta (mirroring outputBudgetMeta's shape for +// the read/search tools). Detection that needs the full output (sandbox-denial +// scanning) must run on the raw strings before this is applied. +func budgetBashOutput(stdout string, stderr string, meta map[string]string) (string, string) { + outText, outRaw, outTrunc := truncateHeadTail(stdout, bashOutputBudgetBytes) + errText, errRaw, errTrunc := truncateHeadTail(stderr, bashOutputBudgetBytes) + if meta != nil { + emitted := len(outText) + len(errText) + meta["raw_bytes"] = strconv.Itoa(outRaw + errRaw) + meta["emitted_bytes"] = strconv.Itoa(emitted) + meta["estimated_tokens"] = strconv.Itoa(estimatedTokensFromBytes(emitted)) + if outTrunc || errTrunc { + meta["truncated"] = "true" + } + } + return outText, errText +} + +// truncateHeadTail keeps the first and last halves of value when it exceeds +// maxBytes, dropping the middle behind a marker. Returns the possibly-truncated +// text, the raw byte length, and whether truncation happened. +func truncateHeadTail(value string, maxBytes int) (string, int, bool) { + raw := len(value) + if maxBytes <= 0 || raw <= maxBytes { + return value, raw, false + } + marker := fmt.Sprintf("\n[zero] output truncated: %d bytes omitted from the middle — redirect to a file and read_file a range for the full text]\n", raw-maxBytes) + budget := maxBytes - len(marker) + if budget < 0 { + budget = 0 + } + head := budget / 2 + tail := budget - head + return utf8Prefix(value, head) + marker + utf8Suffix(value, tail), raw, true +} + func formatBashOutputWithShellHint(command string, stdout string, stderr string, exitCode int, meta map[string]string) string { output := formatBashOutput(stdout, stderr, exitCode) if issue := detectShellOutputIssue(command, stdout+"\n"+stderr, runtime.GOOS); issue != nil { diff --git a/internal/tools/bash_budget_test.go b/internal/tools/bash_budget_test.go new file mode 100644 index 00000000..1fdfceb7 --- /dev/null +++ b/internal/tools/bash_budget_test.go @@ -0,0 +1,55 @@ +package tools + +import ( + "strconv" + "strings" + "testing" +) + +// Small output passes through untouched and records raw==emitted, no truncated flag. +func TestBudgetBashOutputSmallPassesThrough(t *testing.T) { + meta := map[string]string{} + out, errStr := budgetBashOutput("hello\n", "warn\n", meta) + if out != "hello\n" || errStr != "warn\n" { + t.Fatalf("small output altered: out=%q err=%q", out, errStr) + } + if meta["truncated"] == "true" { + t.Fatalf("small output must not be flagged truncated: %v", meta) + } + if meta["raw_bytes"] != strconv.Itoa(len("hello\n")+len("warn\n")) { + t.Fatalf("raw_bytes wrong: %v", meta) + } +} + +// Oversized stdout is truncated head+tail: both the first and last lines survive, +// the middle is dropped behind a marker, and meta is flagged. +func TestBudgetBashOutputTruncatesHeadAndTail(t *testing.T) { + head := "FIRST_LINE_MARKER\n" + tail := "\nLAST_LINE_MARKER" + big := head + strings.Repeat("x", bashOutputBudgetBytes) + tail + + meta := map[string]string{} + out, _ := budgetBashOutput(big, "", meta) + + if !strings.Contains(out, "FIRST_LINE_MARKER") { + t.Fatalf("head lost after truncation") + } + if !strings.Contains(out, "LAST_LINE_MARKER") { + t.Fatalf("tail lost after truncation (failures live at the tail)") + } + if !strings.Contains(out, "output truncated") { + t.Fatalf("expected a truncation marker, got:\n%s", out[:200]) + } + if len(out) > bashOutputBudgetBytes { + t.Fatalf("emitted %d bytes exceeds budget %d", len(out), bashOutputBudgetBytes) + } + if meta["truncated"] != "true" { + t.Fatalf("expected truncated=true, got %v", meta) + } + if meta["raw_bytes"] != strconv.Itoa(len(big)) { + t.Fatalf("raw_bytes = %s, want %d", meta["raw_bytes"], len(big)) + } + if got, _ := strconv.Atoi(meta["emitted_bytes"]); got != len(out) { + t.Fatalf("emitted_bytes = %s, want %d", meta["emitted_bytes"], len(out)) + } +} From f1692cb412518782ed7f9c6648d89957704788be Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 18:13:22 +0530 Subject: [PATCH 03/32] fix(gemini): retry 401 with refreshed OAuth token The Gemini provider used providerio.SendWithRetry (429/503/529 only), so an expired OAuth token surfaced as a terminal 'API key not valid' auth error while OpenAI and Anthropic silently recovered via SendWithAuthRetry's force-refresh-and-replay on 401. Plumb OAuthResolver through gemini.Options/Provider and the factory (exactly as the other two providers already do) and switch the stream send to SendWithAuthRetry. With a nil resolver (API-key users) behavior is unchanged. --- internal/providers/factory.go | 1 + internal/providers/gemini/provider.go | 24 ++++++---- internal/providers/gemini/provider_test.go | 51 ++++++++++++++++++++++ 3 files changed, 68 insertions(+), 8 deletions(-) diff --git a/internal/providers/factory.go b/internal/providers/factory.go index 1d9195c7..d432ca2f 100644 --- a/internal/providers/factory.go +++ b/internal/providers/factory.go @@ -91,6 +91,7 @@ func New(profile config.ProviderProfile, options Options) (zeroruntime.Provider, AuthScheme: profile.AuthScheme, AuthHeaderValue: profile.AuthHeaderValue, CustomHeaders: profile.CustomHeaders, + OAuthResolver: options.OAuthResolver, MaxTokens: resolved.maxOutputTokens, HTTPClient: options.HTTPClient, UserAgent: options.UserAgent, diff --git a/internal/providers/gemini/provider.go b/internal/providers/gemini/provider.go index cb6bb49e..e002db87 100644 --- a/internal/providers/gemini/provider.go +++ b/internal/providers/gemini/provider.go @@ -49,6 +49,10 @@ type Options struct { CustomHeaders map[string]string HTTPClient *http.Client UserAgent string + // OAuthResolver, when set, supplies an OAuth bearer credential per request and + // is retried once with a forced token refresh after an upstream 401 (matching + // the OpenAI and Anthropic providers). Nil falls back to plain API-key auth. + OAuthResolver providerio.TokenResolver // StreamIdleTimeout aborts the stream if no data arrives for this long. // When unset, Zero uses providerio.ResolveStreamIdleTimeout — the // ZERO_STREAM_IDLE_TIMEOUT override or providerio.DefaultStreamIdleTimeout. @@ -67,6 +71,7 @@ type Provider struct { customHeaders map[string]string httpClient *http.Client userAgent string + oauthResolver providerio.TokenResolver streamIdleTimeout time.Duration } @@ -95,6 +100,7 @@ func New(options Options) (*Provider, error) { customHeaders: providerio.CopyHeaders(options.CustomHeaders), httpClient: providerio.HTTPClient(options.HTTPClient), userAgent: options.UserAgent, + oauthResolver: options.OAuthResolver, streamIdleTimeout: providerio.ResolveStreamIdleTimeout(options.StreamIdleTimeout), }, nil } @@ -127,20 +133,22 @@ func (provider *Provider) stream(ctx context.Context, body []byte, events chan<- streamCtx, cancelStream := context.WithCancel(ctx) defer cancelStream() - response, err := providerio.SendWithRetry(streamCtx, provider.httpClient, http.MethodPost, provider.streamURL(), body, func(request *http.Request) { - request.Header.Set("Content-Type", "application/json") - if provider.userAgent != "" { - request.Header.Set("User-Agent", provider.userAgent) - } - providerio.ApplyAuthHeaders(request, providerio.AuthHeaders{ + response, err := providerio.SendWithAuthRetry(streamCtx, provider.httpClient, http.MethodPost, provider.streamURL(), body, + providerio.AuthHeaders{ APIKey: provider.apiKey, DefaultAuthHeader: "x-goog-api-key", AuthHeader: provider.authHeader, AuthScheme: provider.authScheme, AuthHeaderValue: provider.authHeaderValue, CustomHeaders: provider.customHeaders, - }) - }, 0) + }, + provider.oauthResolver, + func(request *http.Request) { + request.Header.Set("Content-Type", "application/json") + if provider.userAgent != "" { + request.Header.Set("User-Agent", provider.userAgent) + } + }, 0) if err != nil { providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventError, Error: provider.redact("provider stream error: " + err.Error())}) return diff --git a/internal/providers/gemini/provider_test.go b/internal/providers/gemini/provider_test.go index 205df682..7c83c714 100644 --- a/internal/providers/gemini/provider_test.go +++ b/internal/providers/gemini/provider_test.go @@ -370,6 +370,57 @@ func TestStreamCompletionClassifiesHTTPAndPromptBlockErrors(t *testing.T) { } } +// A 401 with an OAuth resolver is retried once with a force-refreshed token; the +// replayed request carries the refreshed bearer and succeeds. +func TestStreamCompletionRetries401WithRefreshedToken(t *testing.T) { + var attempts int + var secondAuth string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attempts++ + if attempts == 1 { + http.Error(w, `{"error":{"message":"token expired"}}`, http.StatusUnauthorized) + return + } + secondAuth = r.Header.Get("Authorization") + writeSSE(w, `{}`) + })) + defer server.Close() + + var forceRefreshOnRetry bool + resolver := func(ctx context.Context, forceRefresh bool) (string, string, bool, error) { + if forceRefresh { + forceRefreshOnRetry = true + return "Authorization", "Bearer refreshed", true, nil + } + return "Authorization", "Bearer stale", true, nil + } + + provider, err := New(Options{ + APIKey: "sk-google", + BaseURL: server.URL, + Model: "gemini-test", + OAuthResolver: resolver, + }) + if err != nil { + t.Fatalf("New returned error: %v", err) + } + stream, err := provider.StreamCompletion(context.Background(), validRequest()) + if err != nil { + t.Fatalf("StreamCompletion returned error: %v", err) + } + drain(stream) + + if attempts != 2 { + t.Fatalf("attempts = %d, want 2 (initial 401 + one refreshed retry)", attempts) + } + if !forceRefreshOnRetry { + t.Fatalf("resolver was not called with forceRefresh on the retry") + } + if secondAuth != "Bearer refreshed" { + t.Fatalf("retry Authorization = %q, want refreshed bearer", secondAuth) + } +} + func TestStreamCompletionEmitsStreamErrorObject(t *testing.T) { provider := newTestProviderWithKey(t, "sk-google", func(w http.ResponseWriter, r *http.Request) { writeSSE(w, `{"error":{"code":429,"message":"stream failed sk-google","status":"RESOURCE_EXHAUSTED"}}`) From f13cbc3453134ceac502900ced4a67077c346435 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 18:17:56 +0530 Subject: [PATCH 04/32] fix(errhint): gate classification on provider-origin marker Address CodeRabbit review on #1: agentResponseMsg.err can also carry local failures (a tool's 'permission denied', a 'file does not exist', a config error), and broad substrings would have attached a bogus /provider or /model hint to those. Classify now returns Unknown unless the message carries a provider marker the provider layer always attaches (auth error: / rate limit error: / provider error: / provider request error: / provider stream error:). Local errors never draw a provider hint; provider errors sub-classify as before. Added local-failure test cases proving they stay Unknown. --- internal/errhint/errhint.go | 27 ++++++++++++++++++++++++--- internal/errhint/errhint_test.go | 31 +++++++++++++++++++------------ 2 files changed, 43 insertions(+), 15 deletions(-) diff --git a/internal/errhint/errhint.go b/internal/errhint/errhint.go index e29eaae3..b8516d6d 100644 --- a/internal/errhint/errhint.go +++ b/internal/errhint/errhint.go @@ -26,17 +26,38 @@ const ( ContextOverflow ) +// providerMarkers are the prefixes the provider layer attaches to every +// provider-originated failure: providerio.ClassifiedError emits "auth error:", +// "rate limit error:", "provider error:", and "provider request error:", and the +// streaming paths wrap transport/decoding failures as "provider stream error:". +// A UI surface's error can also be a *local* failure (a tool's "permission +// denied", a "file does not exist", a config error), so Classify only proceeds +// past this gate for messages that are recognizably from the provider — otherwise +// a broad substring like "does not exist" would attach a bogus /model hint to an +// unrelated local error. +var providerMarkers = []string{ + "auth error:", + "rate limit error:", + "provider error:", + "provider request error:", + "provider stream error:", +} + // Classify buckets err by scanning its message for known signatures. It is a // deliberately conservative string heuristic — the provider layer has already // discarded the numeric HTTP status by the time the error reaches a UI surface, -// so the message is all we have. Order matters: more specific signatures are -// tested before broader ones (e.g. "context length" as overflow before the -// generic "timeout" as connectivity). +// so the message is all we have. It first gates on a provider-origin marker (see +// providerMarkers) so local failures never draw a provider hint, then sub-classifies. +// Order matters: more specific signatures are tested before broader ones (e.g. +// "context length" as overflow before the generic "timeout" as connectivity). func Classify(err error) Category { if err == nil { return Unknown } m := strings.ToLower(err.Error()) + if !containsAny(m, providerMarkers...) { + return Unknown + } switch { case containsAny(m, "auth error:", "unauthorized", "api key", "api_key", "invalid_api_key", "authentication", "permission denied", "forbidden", "401", "403"): diff --git a/internal/errhint/errhint_test.go b/internal/errhint/errhint_test.go index 2a6c2866..8f6fd0fa 100644 --- a/internal/errhint/errhint_test.go +++ b/internal/errhint/errhint_test.go @@ -15,18 +15,25 @@ func TestClassify(t *testing.T) { {"nil-ish empty", "", Unknown}, {"providerio auth prefix", "auth error: your API key is missing or invalid — run `zero auth`", Auth}, {"raw 401", "provider request error: 401 Unauthorized", Auth}, - {"invalid api key", "Error: invalid_api_key: incorrect key provided", Auth}, + {"invalid api key", "provider request error: invalid_api_key: incorrect key provided", Auth}, {"rate limit prefix", "rate limit error: 429 too many requests", RateLimit}, {"overloaded", "provider error: model is overloaded, please retry", RateLimit}, - {"resource exhausted gemini", "rpc error: code = ResourceExhausted desc = quota exceeded", RateLimit}, - {"context length openai", "This model's maximum context length is 128000 tokens", ContextOverflow}, + {"resource exhausted gemini", "provider stream error: rpc error: code = ResourceExhausted desc = quota exceeded", RateLimit}, + {"context length openai", "provider request error: this model's maximum context length is 128000 tokens", ContextOverflow}, {"context_length_exceeded", "provider request error: context_length_exceeded", ContextOverflow}, {"model not found", "provider request error: 404 the model `gpt-9` does not exist", ModelNotFound}, - {"unknown model", "unknown model: sonnet-99", ModelNotFound}, - {"dns failure", "Post \"https://api...\": dial tcp: lookup api.foo.com: no such host", Connectivity}, - {"connection refused", "dial tcp 127.0.0.1:443: connection refused", Connectivity}, - {"deadline", "context deadline exceeded (Client.Timeout exceeded)", Connectivity}, - {"gibberish", "provider error: something totally unexpected happened", Unknown}, + {"unknown model", "provider request error: unknown model: sonnet-99", ModelNotFound}, + {"dns failure", "provider stream error: Post \"https://api...\": dial tcp: lookup api.foo.com: no such host", Connectivity}, + {"connection refused", "provider stream error: dial tcp 127.0.0.1:443: connection refused", Connectivity}, + {"deadline", "provider stream error: context deadline exceeded (Client.Timeout exceeded)", Connectivity}, + {"provider-marked but no sub-signature", "provider error: something totally unexpected happened", Unknown}, + + // Local (non-provider) failures must NOT be classified — no provider marker, + // so no bogus recovery hint attaches to them. + {"local fs permission denied", "open /etc/shadow: permission denied", Unknown}, + {"local file missing", "stat foo.go: no such file or directory", Unknown}, + {"local model config typo", "unknown model: sonnet-99", Unknown}, + {"tool timeout local", "tool bash timed out after 600s", Unknown}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { @@ -47,7 +54,7 @@ func TestClassifyNil(t *testing.T) { // mentioning both "context length" and a timeout is an overflow, not a network // problem. func TestContextOverflowBeatsConnectivity(t *testing.T) { - err := errors.New("maximum context length exceeded; request timeout") + err := errors.New("provider request error: maximum context length exceeded; request timeout") if got := Classify(err); got != ContextOverflow { t.Fatalf("Classify = %v, want ContextOverflow", got) } @@ -57,9 +64,9 @@ func TestHintsPresentForKnownCategories(t *testing.T) { known := []error{ errors.New("auth error: bad key"), errors.New("rate limit error: 429"), - errors.New("dial tcp: no such host"), - errors.New("model does not exist"), - errors.New("maximum context length is 128000 tokens"), + errors.New("provider stream error: dial tcp: no such host"), + errors.New("provider request error: model does not exist"), + errors.New("provider request error: maximum context length is 128000 tokens"), } for _, err := range known { if h := TUIHint(err); strings.TrimSpace(h) == "" { From fb94121b35b60757bfeef306113eef71c791d215 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 18:22:19 +0530 Subject: [PATCH 05/32] feat(cli): actionable hint on text-mode exec provider errors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit zero exec printed provider failures as a bare '[zero] ' with no next step. Append a one-line hint (reusing the shared errhint classifier) for recognized provider failures — 'run `zero auth`', 'run `zero doctor`', etc. The provider-origin gate means non-provider codes (sandbox_error, mcp_error) and JSON/stream-json output are untouched. --- internal/cli/exec.go | 9 +++++ internal/cli/exec_provider_hint_test.go | 50 +++++++++++++++++++++++++ 2 files changed, 59 insertions(+) create mode 100644 internal/cli/exec_provider_hint_test.go diff --git a/internal/cli/exec.go b/internal/cli/exec.go index 069d9e79..7db7c828 100644 --- a/internal/cli/exec.go +++ b/internal/cli/exec.go @@ -12,6 +12,7 @@ import ( "github.com/Gitlawb/zero/internal/agent" "github.com/Gitlawb/zero/internal/config" + "github.com/Gitlawb/zero/internal/errhint" "github.com/Gitlawb/zero/internal/imageinput" "github.com/Gitlawb/zero/internal/lsp" "github.com/Gitlawb/zero/internal/modelregistry" @@ -948,6 +949,14 @@ func writeExecProviderError(stdout io.Writer, stderr io.Writer, format execOutpu if _, err := fmt.Fprintf(stderr, "[zero] %s\n", message); err != nil { return exitCrash } + // Append a one-line next step for recognized provider failures (auth / + // rate-limit / connectivity / …). The classifier gates on a provider-origin + // marker, so non-provider codes (sandbox_error, mcp_error) never draw a hint. + if hint := errhint.CLIHint(errors.New(message)); hint != "" { + if _, err := fmt.Fprintf(stderr, "[zero] %s\n", hint); err != nil { + return exitCrash + } + } return exitProvider } diff --git a/internal/cli/exec_provider_hint_test.go b/internal/cli/exec_provider_hint_test.go new file mode 100644 index 00000000..54e9db63 --- /dev/null +++ b/internal/cli/exec_provider_hint_test.go @@ -0,0 +1,50 @@ +package cli + +import ( + "bytes" + "strings" + "testing" +) + +// A recognized provider error in text mode prints the raw message plus a one-line +// actionable hint; a non-provider error (or JSON mode) prints no hint. +func TestWriteExecProviderErrorAppendsHint(t *testing.T) { + t.Run("provider auth error gets a hint", func(t *testing.T) { + var stdout, stderr bytes.Buffer + code := writeExecProviderError(&stdout, &stderr, execOutputText, "provider_error", + "auth error: your API key is missing or invalid") + if code != exitProvider { + t.Fatalf("exit code = %d, want exitProvider", code) + } + out := stderr.String() + if !strings.Contains(out, "auth error:") { + t.Fatalf("expected raw message, got %q", out) + } + if !strings.Contains(out, "zero auth") { + t.Fatalf("expected an actionable hint referencing `zero auth`, got %q", out) + } + }) + + t.Run("non-provider error gets no hint", func(t *testing.T) { + var stdout, stderr bytes.Buffer + writeExecProviderError(&stdout, &stderr, execOutputText, "sandbox_error", + "sandbox setup failed: permission denied") + out := stderr.String() + // Exactly one "[zero]" line — no spurious hint attached to a local error. + if n := strings.Count(out, "[zero]"); n != 1 { + t.Fatalf("expected exactly one [zero] line for a non-provider error, got %d:\n%s", n, out) + } + }) + + t.Run("json mode never appends a hint line", func(t *testing.T) { + var stdout, stderr bytes.Buffer + writeExecProviderError(&stdout, &stderr, execOutputJSON, "provider_error", + "auth error: bad key") + if stderr.Len() != 0 { + t.Fatalf("json mode must not write to stderr, got %q", stderr.String()) + } + if strings.Contains(stdout.String(), "zero auth") { + t.Fatalf("json mode must not inject a hint into the structured payload, got %q", stdout.String()) + } + }) +} From ab3130360fb56a6abe3e5ab8d699516dccd78751 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 18:26:09 +0530 Subject: [PATCH 06/32] fix(tools): set Result.Truncated for budgeted bash output Address CodeRabbit review on #3: budgetBashOutput recorded truncation in meta["truncated"] but never set the Result.Truncated struct field, so consumers reading the struct (like glob/web_fetch/read_minified do) would miss truncated bash results. Return the bool from budgetBashOutput and set Result.Truncated on every return path. --- internal/tools/bash.go | 32 +++++++++++++++++------------- internal/tools/bash_budget_test.go | 10 ++++++++-- 2 files changed, 26 insertions(+), 16 deletions(-) diff --git a/internal/tools/bash.go b/internal/tools/bash.go index b0329771..f1cde32e 100644 --- a/internal/tools/bash.go +++ b/internal/tools/bash.go @@ -157,27 +157,30 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS } } markLikelySandboxDenial(meta, plan, exitCode, stdout.String(), stderrText) - outText, errText := budgetBashOutput(stdout.String(), stderrText, meta) + outText, errText, truncated := budgetBashOutput(stdout.String(), stderrText, meta) return Result{ - Status: StatusError, - Output: formatBashOutputWithShellHint(commandText, outText, errText, exitCode, meta), - Meta: meta, + Status: StatusError, + Output: formatBashOutputWithShellHint(commandText, outText, errText, exitCode, meta), + Truncated: truncated, + Meta: meta, } } markLikelySandboxDenial(meta, plan, exitCode, stdout.String(), stderrText) - outText, errText := budgetBashOutput(stdout.String(), stderrText, meta) + outText, errText, truncated := budgetBashOutput(stdout.String(), stderrText, meta) if meta[SandboxLikelyDeniedMeta] == "true" { return Result{ - Status: StatusError, - Output: formatBashOutputWithShellHint(commandText, outText, errText, exitCode, meta), - Meta: meta, + Status: StatusError, + Output: formatBashOutputWithShellHint(commandText, outText, errText, exitCode, meta), + Truncated: truncated, + Meta: meta, } } return Result{ - Status: StatusOK, - Output: formatBashOutput(outText, errText, exitCode), - Meta: meta, + Status: StatusOK, + Output: formatBashOutput(outText, errText, exitCode), + Truncated: truncated, + Meta: meta, } } @@ -359,19 +362,20 @@ const bashOutputBudgetBytes = 96 * 1024 // counts plus a truncated flag in meta (mirroring outputBudgetMeta's shape for // the read/search tools). Detection that needs the full output (sandbox-denial // scanning) must run on the raw strings before this is applied. -func budgetBashOutput(stdout string, stderr string, meta map[string]string) (string, string) { +func budgetBashOutput(stdout string, stderr string, meta map[string]string) (string, string, bool) { outText, outRaw, outTrunc := truncateHeadTail(stdout, bashOutputBudgetBytes) errText, errRaw, errTrunc := truncateHeadTail(stderr, bashOutputBudgetBytes) + truncated := outTrunc || errTrunc if meta != nil { emitted := len(outText) + len(errText) meta["raw_bytes"] = strconv.Itoa(outRaw + errRaw) meta["emitted_bytes"] = strconv.Itoa(emitted) meta["estimated_tokens"] = strconv.Itoa(estimatedTokensFromBytes(emitted)) - if outTrunc || errTrunc { + if truncated { meta["truncated"] = "true" } } - return outText, errText + return outText, errText, truncated } // truncateHeadTail keeps the first and last halves of value when it exceeds diff --git a/internal/tools/bash_budget_test.go b/internal/tools/bash_budget_test.go index 1fdfceb7..45047003 100644 --- a/internal/tools/bash_budget_test.go +++ b/internal/tools/bash_budget_test.go @@ -9,10 +9,13 @@ import ( // Small output passes through untouched and records raw==emitted, no truncated flag. func TestBudgetBashOutputSmallPassesThrough(t *testing.T) { meta := map[string]string{} - out, errStr := budgetBashOutput("hello\n", "warn\n", meta) + out, errStr, truncated := budgetBashOutput("hello\n", "warn\n", meta) if out != "hello\n" || errStr != "warn\n" { t.Fatalf("small output altered: out=%q err=%q", out, errStr) } + if truncated { + t.Fatalf("small output must report truncated=false") + } if meta["truncated"] == "true" { t.Fatalf("small output must not be flagged truncated: %v", meta) } @@ -29,8 +32,11 @@ func TestBudgetBashOutputTruncatesHeadAndTail(t *testing.T) { big := head + strings.Repeat("x", bashOutputBudgetBytes) + tail meta := map[string]string{} - out, _ := budgetBashOutput(big, "", meta) + out, _, truncated := budgetBashOutput(big, "", meta) + if !truncated { + t.Fatalf("oversized output must report truncated=true") + } if !strings.Contains(out, "FIRST_LINE_MARKER") { t.Fatalf("head lost after truncation") } From ef956fc82beeac882332dfc8af27bf65beec360e Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 18:36:08 +0530 Subject: [PATCH 07/32] perf(tui): reuse one LSP manager across prompts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit runAgentWithOptions built a fresh lsp.NewManager per run and shut it down when the run returned, so gopls (and every other language server) cold-started on the first edit of every turn — 200ms-2s of latency the manager's own design (long-lived, reused servers) exists to avoid. Build one manager per session in newModel (cheap: servers start lazily on first Check) and reuse it across runs; the SelfCorrector still gets a fresh checker wrapping the shared manager each run. Torn down with a short deadline in quit(). Runs fall back to a per-run manager when the session manager is absent (cwd unknown, or a directly-built test model). --- internal/tui/lsp_manager_test.go | 31 +++++++++++++ internal/tui/model.go | 77 ++++++++++++++++++++++---------- 2 files changed, 84 insertions(+), 24 deletions(-) create mode 100644 internal/tui/lsp_manager_test.go diff --git a/internal/tui/lsp_manager_test.go b/internal/tui/lsp_manager_test.go new file mode 100644 index 00000000..5ce1b1d6 --- /dev/null +++ b/internal/tui/lsp_manager_test.go @@ -0,0 +1,31 @@ +package tui + +import ( + "context" + "testing" +) + +// The LSP manager is built once at construction (when cwd is known) and reused — +// a fresh manager per run would cold-start gopls on the first edit of every turn. +func TestNewModelBuildsSharedLSPManager(t *testing.T) { + m := newModel(context.Background(), Options{Cwd: t.TempDir()}) + if m.lspManager == nil { + t.Fatal("expected a session-long lspManager when cwd is set") + } + // The pointer is stable across model copies (Bubble Tea passes the model by + // value every Update), so every run shares the same warm servers. + copied := m + if copied.lspManager != m.lspManager { + t.Fatal("lspManager pointer must survive a model copy") + } +} + +// shutdownLSPManager must be nil-safe: a model built without a cwd has no manager, +// and quitting it must not panic. +func TestShutdownLSPManagerNilSafe(t *testing.T) { + m := model{} + if m.lspManager != nil { + t.Fatal("expected no manager on a zero-value model") + } + m.shutdownLSPManager() // must not panic +} diff --git a/internal/tui/model.go b/internal/tui/model.go index 480d601c..1143c772 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -76,24 +76,29 @@ type model struct { discoverProviderModels func(context.Context, config.ProviderProfile) ([]providermodeldiscovery.Model, error) discoverOllamaContextWindow func(ctx context.Context, baseURL string, model string) (int, error) registry *tools.Registry - sessionStore *sessions.Store - sandboxStore *sandbox.GrantStore - mcpConfig config.MCPConfig - mcpPermissionStore *internalmcp.PermissionStore - mcpTokenStore *internalmcp.TokenStore - mcpCommand func(context.Context, []string) MCPCommandResult - sandboxSetupCommand func(context.Context) SandboxSetupCommandResult - mcpViewStateCache MCPViewState - mcpViewStateReady bool - mcpCommandSeq int - mcpCommandCancel context.CancelFunc - sandboxSetupSeq int - sandboxSetupInFlight bool - doctorCommandSeq int - doctorInFlight bool - doctorFrame int - activeSession sessions.Metadata - sessionEvents []sessions.Event + // lspManager is created once per session and reused across prompts so gopls (and + // other language servers) stay warm — a fresh manager per run would cold-start + // the server on the first edit of every turn. Nil when cwd is unknown; runs then + // fall back to a per-run manager. Torn down in quit(). + lspManager *lsp.Manager + sessionStore *sessions.Store + sandboxStore *sandbox.GrantStore + mcpConfig config.MCPConfig + mcpPermissionStore *internalmcp.PermissionStore + mcpTokenStore *internalmcp.TokenStore + mcpCommand func(context.Context, []string) MCPCommandResult + sandboxSetupCommand func(context.Context) SandboxSetupCommandResult + mcpViewStateCache MCPViewState + mcpViewStateReady bool + mcpCommandSeq int + mcpCommandCancel context.CancelFunc + sandboxSetupSeq int + sandboxSetupInFlight bool + doctorCommandSeq int + doctorInFlight bool + doctorFrame int + activeSession sessions.Metadata + sessionEvents []sessions.Event // titledSessions records session ids for which a model-generated title has // already been attempted this process, so a finished turn re-fires the title // generator at most once per session (even before its async result lands). @@ -758,6 +763,11 @@ func newModel(ctx context.Context, options Options) model { // Streaming text always renders statically at base ink (the disabled path in // styleStreamingLine), so no accent glow and no per-line fade ticks. m.fadeDisabled = true + // One session-long LSP manager (cheap to build — servers start lazily on the + // first Check), reused across prompts so gopls stays warm between turns. + if cwd != "" { + m.lspManager = lsp.NewManager(cwd) + } m.refreshMCPViewState() return m } @@ -887,9 +897,22 @@ func (m model) noBlockingModal() bool { func (m model) quit() (tea.Model, tea.Cmd) { m.stopPRWatcher() m.stopAllBackgroundTerminalSessions() + m.shutdownLSPManager() return m, tea.Quit } +// shutdownLSPManager gracefully stops the session-long language servers on exit. +// Best-effort with a short deadline so a slow server can't hang the quit; the +// servers are our child processes and would be reaped on exit regardless. +func (m model) shutdownLSPManager() { + if m.lspManager == nil { + return + } + shutdownCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + _ = m.lspManager.Shutdown(shutdownCtx) +} + func (m model) handleCtrlC() (tea.Model, tea.Cmd) { if !m.pending && m.composerValue() != "" && m.noBlockingModal() && !m.transcriptDetailed && !m.subchat.active { m.clearComposer() @@ -4158,12 +4181,18 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str // matching exec; the per-turn lsp.Manager is torn down when this run // returns; auto-fix vs report-only follows the active permission mode. if !runOptions.specDraft && options.Cwd != "" { - lspManager := lsp.NewManager(options.Cwd) - defer func() { - shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - _ = lspManager.Shutdown(shutdownCtx) - }() + // Prefer the session-long manager (kept warm across prompts). Only when it + // is absent — e.g. cwd was unknown at construction, or a test built the + // model directly — fall back to a per-run manager that is shut down here. + lspManager := m.lspManager + if lspManager == nil { + lspManager = lsp.NewManager(options.Cwd) + defer func() { + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = lspManager.Shutdown(shutdownCtx) + }() + } options.SelfCorrect = agent.NewSelfCorrector(options.Cwd, agent.NewLSPDiagnosticsChecker(lspManager), agent.NewProjectVerifier(options.Cwd), agent.SelfCorrectConfig{ Enabled: true, IncludeTests: m.selfCorrectTests, From e6bc357d47aa4b4813b3f2da4ee1f64e81f3ae69 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 18:50:56 +0530 Subject: [PATCH 08/32] perf(tui): accumulate streamingText as []byte (O(1) append) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit m.streamingText += msg.delta was O(len) per delta -> O(n²) across a long generation (a 10k-token code gen allocates hundreds of MB of intermediate strings and the UI gets progressively laggier). Accumulate into a []byte with append (O(1) amortized) and read via streamingTextString(). A []byte rather than strings.Builder because the TUI model is copied by value on every Update, which would trip strings.Builder's copy check; []byte is nil-safe and copies cleanly like the other slice fields. --- internal/tui/flush_test.go | 4 ++-- internal/tui/model.go | 32 ++++++++++++++++++--------- internal/tui/model_test.go | 2 +- internal/tui/rendering_lime_test.go | 14 ++++++------ internal/tui/syntax_highlight_test.go | 6 ++--- internal/tui/transcript_selection.go | 2 +- internal/tui/working_status_test.go | 8 +++---- 7 files changed, 40 insertions(+), 28 deletions(-) diff --git a/internal/tui/flush_test.go b/internal/tui/flush_test.go index 8cb90a8a..75ca6d31 100644 --- a/internal/tui/flush_test.go +++ b/internal/tui/flush_test.go @@ -118,7 +118,7 @@ func TestEscCancellationLeavesVisibleMarker(t *testing.T) { m.pending = true m.activeRunID = 5 m.runCancel = func() {} - m.streamingText = "half an answer" + m.streamingText = []byte("half an answer") updated, _ := m.Update(testKey(tea.KeyEsc)) next := updated.(model) @@ -130,7 +130,7 @@ func TestEscCancellationLeavesVisibleMarker(t *testing.T) { if !transcriptContains(next.transcript, "half an answer") { t.Fatalf("expected the partial streamed answer to be preserved, got %#v", next.transcript) } - if next.streamingText != "" { + if len(next.streamingText) != 0 { t.Fatal("streaming text should be cleared after cancel") } } diff --git a/internal/tui/model.go b/internal/tui/model.go index 1143c772..3634cfc5 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -265,7 +265,12 @@ type model struct { historyIdx int historyDraft string - streamingText string // live assistant text for the current segment + // streamingText is the live assistant text for the current segment, accumulated + // as []byte so each delta is an O(1) amortized append instead of the O(n²) that + // string += delta incurs across a long generation. Read via streamingTextString(). + // A []byte (not strings.Builder) because the model is copied by value on every + // Update, which would trip strings.Builder's copy check. + streamingText []byte streamingReasoning string // live provider reasoning for the current segment streamingReasoningExpanded bool // turnStreamedRunes accumulates every reasoning+answer rune streamed in the @@ -1584,7 +1589,7 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) { // Streaming text means any in-progress tool call has finished — clear the // live "writing" block so it doesn't linger over new prose. m.clearStreamingToolCall() - m.streamingText += msg.delta + m.streamingText = append(m.streamingText, msg.delta...) m.turnStreamedRunes += utf8.RuneCountInString(msg.delta) // recordStreamingDelta appends a time.Time to lineAges for every // newline in the delta and bumps lastStreamActivity. It also @@ -1867,7 +1872,7 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) { if row, ok := reasoningTranscriptRow("", msg.runID, m.streamingReasoning); ok { m.transcript = appendTranscriptRow(m.transcript, row) } - if text := strings.TrimRight(m.streamingText, "\n"); strings.TrimSpace(text) != "" { + if text := strings.TrimRight(m.streamingTextString(), "\n"); strings.TrimSpace(text) != "" { m.transcript = appendTranscriptRow(m.transcript, transcriptRow{kind: rowAssistant, text: text}) } // The error row terminates the turn, so it carries the done-line @@ -1883,7 +1888,7 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) { turnElapsed: msg.turnElapsed, }) } - m.streamingText = "" + m.streamingText = nil m.streamingReasoning = "" m.streamingReasoningExpanded = false if msg.specReview != nil { @@ -2048,14 +2053,14 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) { m.streamingReasoning = "" m.streamingReasoningExpanded = false } - if text := strings.TrimRight(m.streamingText, "\n"); strings.TrimSpace(text) != "" { + if text := strings.TrimRight(m.streamingTextString(), "\n"); strings.TrimSpace(text) != "" { m.transcript = appendTranscriptRow(m.transcript, transcriptRow{kind: rowAssistant, text: text}) // This interim narration is the agent explaining what it's about to // do — attribute it to the active plan step so the step-detail card // can replay the agent's own account of the work. m = m.captureStepNarration(text) } - m.streamingText = "" + m.streamingText = nil // The tool call has finalized into its card — drop the live "writing" // preview so it doesn't linger or duplicate beneath the card. m.clearStreamingToolCall() @@ -2786,7 +2791,7 @@ func (m model) liveReasoningBodyCap() int { } func (m model) interimBlock(width int) string { - text := strings.TrimRight(m.streamingText, "\n") + text := strings.TrimRight(m.streamingTextString(), "\n") reasoning := strings.TrimRight(m.streamingReasoning, "\n") blocks := []string{} if strings.TrimSpace(reasoning) != "" { @@ -2850,7 +2855,7 @@ func (m model) spinnerGlyph() string { // (reasoning, waiting on the model, or a tool in flight). Cheap and robust — no // transcript scan — so it can't misreport on a long, output-less step. func (m model) workingActivity() string { - if strings.TrimSpace(m.streamingText) != "" { + if strings.TrimSpace(m.streamingTextString()) != "" { return "writing" } return "thinking" @@ -4086,7 +4091,7 @@ func (m *model) cancelRun() { if row, ok := reasoningTranscriptRow("", m.activeRunID, m.streamingReasoning); ok { m.transcript = appendTranscriptRow(m.transcript, row) } - if text := strings.TrimRight(m.streamingText, "\n"); strings.TrimSpace(text) != "" { + if text := strings.TrimRight(m.streamingTextString(), "\n"); strings.TrimSpace(text) != "" { m.transcript = appendTranscriptRow(m.transcript, transcriptRow{kind: rowAssistant, text: text}) } m.transcript = appendTranscriptRow(m.transcript, transcriptRow{kind: rowSystem, text: "Run cancelled."}) @@ -4107,7 +4112,7 @@ func (m *model) cancelRun() { m.pendingAskUser = nil // The interim block renders streamingText live; a cancelled run's partial // answer must not leak into (and concatenate with) the next turn's stream. - m.streamingText = "" + m.streamingText = nil m.streamingReasoning = "" m.streamingReasoningExpanded = false // Hard-stop the fade and drop the per-line age map. The next turn's @@ -4640,6 +4645,13 @@ func (m model) sendAgentText(runID int, delta string) { m.runtimeMessageSink(agentTextMsg{runID: runID, delta: delta}) } +// streamingTextString returns the accumulated live assistant text. streamingText +// is stored as []byte for O(1) amortized appends; the conversion here is bounded +// by the segment length, the same cost the renderer already pays. +func (m model) streamingTextString() string { + return string(m.streamingText) +} + func (m model) sendToolCallStreamStart(runID int, id, name string) { if m.runtimeMessageSink == nil { return diff --git a/internal/tui/model_test.go b/internal/tui/model_test.go index 55a4c26a..83e6fb09 100644 --- a/internal/tui/model_test.go +++ b/internal/tui/model_test.go @@ -2298,7 +2298,7 @@ func TestStreamingAssistantAfterToolCardGetsRuleSeparator(t *testing.T) { m.width, m.height = 120, 40 m.altScreen = true m.pending = true - m.streamingText = "Done." + m.streamingText = []byte("Done.") m.transcript = append(m.transcript, transcriptRow{kind: rowUser, text: "run it"}, transcriptRow{kind: rowAssistant, text: "I'll run it first."}, diff --git a/internal/tui/rendering_lime_test.go b/internal/tui/rendering_lime_test.go index 261f1ec4..60df7325 100644 --- a/internal/tui/rendering_lime_test.go +++ b/internal/tui/rendering_lime_test.go @@ -224,14 +224,14 @@ func TestCommandCardRowTrimsIndentedActionsLabel(t *testing.T) { func TestInterimBlockShowsStreamingTextWithCursor(t *testing.T) { m := limeTestModel() m.pending = true - m.streamingText = "I'll add a --version flag" + m.streamingText = []byte("I'll add a --version flag") got := plainRender(t, m.interimBlock(96)) if !strings.Contains(got, "I'll add a --version flag") || !strings.Contains(got, "▌") { t.Fatalf("interim block = %q, want streamed text with trailing cursor", got) } // Before the first delta the block falls back to the liveness spinner. - m.streamingText = "" + m.streamingText = nil if got := plainRender(t, m.interimBlock(96)); !strings.Contains(got, "Working") { t.Fatalf("empty interim block = %q, want the liveness label", got) } @@ -240,13 +240,13 @@ func TestInterimBlockShowsStreamingTextWithCursor(t *testing.T) { func TestInterimBlockRendersStreamingMarkdownTable(t *testing.T) { m := limeTestModel() m.pending = true - m.streamingText = strings.Join([]string{ + m.streamingText = []byte(strings.Join([]string{ "Here's the comparison:", "", "| Category | System A | System B |", "|---|---|---|", "| **Label** | Alpha | Beta |", - }, "\n") + }, "\n")) rendered := m.interimBlock(72) got := plainRender(t, rendered) @@ -1528,10 +1528,10 @@ func TestCancelRunClearsStreamingText(t *testing.T) { m := limeTestModel() m.pending = true m.activeRunID = 3 - m.streamingText = "partial answer from a doomed run" + m.streamingText = []byte("partial answer from a doomed run") m.cancelRun() - if m.streamingText != "" { - t.Fatalf("cancelRun must clear streamingText, got %q", m.streamingText) + if len(m.streamingText) != 0 { + t.Fatalf("cancelRun must clear streamingText, got %q", string(m.streamingText)) } } diff --git a/internal/tui/syntax_highlight_test.go b/internal/tui/syntax_highlight_test.go index b78846d9..2f2b9c0f 100644 --- a/internal/tui/syntax_highlight_test.go +++ b/internal/tui/syntax_highlight_test.go @@ -7,7 +7,7 @@ import ( func TestStreamingCodeRendersHighlighted(t *testing.T) { m := model{ - streamingText: "```go\nfunc main() {}\n```", + streamingText: []byte("```go\nfunc main() {}\n```"), pending: true, } out := m.interimBlock(80) @@ -86,7 +86,7 @@ func TestStreamingMarkdownStablePrefixUsesRenderCache(t *testing.T) { func TestStreamingBuffersOpenFencedCodeBlock(t *testing.T) { open := model{ - streamingText: "Here is the script:\n```python\nfrom datetime import datetime\nprint(datetime.now())", + streamingText: []byte("Here is the script:\n```python\nfrom datetime import datetime\nprint(datetime.now())"), pending: true, } openOut := plainRender(t, open.interimBlock(90)) @@ -98,7 +98,7 @@ func TestStreamingBuffersOpenFencedCodeBlock(t *testing.T) { } closed := model{ - streamingText: open.streamingText + "\n```", + streamingText: []byte(string(open.streamingText) + "\n```"), pending: true, } closedOut := closed.interimBlock(90) diff --git a/internal/tui/transcript_selection.go b/internal/tui/transcript_selection.go index f858da4a..2f5828e5 100644 --- a/internal/tui/transcript_selection.go +++ b/internal/tui/transcript_selection.go @@ -366,7 +366,7 @@ func (m model) transcriptBodyItems(width int, emptyOverlay string) []transcriptB } if m.pending { - pendingShowsAssistantText := m.pendingPermission == nil && m.pendingAskUser == nil && strings.TrimSpace(m.streamingText) != "" + pendingShowsAssistantText := m.pendingPermission == nil && m.pendingAskUser == nil && strings.TrimSpace(m.streamingTextString()) != "" if pendingShowsAssistantText && havePreviousKind && shouldRuleBeforeTurn(previousKind, rowAssistant) { items = append(items, transcriptRuleBodyItem(contentWidth, gutter)) } else { diff --git a/internal/tui/working_status_test.go b/internal/tui/working_status_test.go index 846d1da7..92b32acb 100644 --- a/internal/tui/working_status_test.go +++ b/internal/tui/working_status_test.go @@ -191,7 +191,7 @@ func TestInterimBlockShowsWorkingLineWithStreamedText(t *testing.T) { base := time.Date(2026, 6, 14, 10, 0, 0, 0, time.UTC) m.now = func() time.Time { return base.Add(12 * time.Second) } m.turnStartedAt = base - m.streamingText = "partial answer so far" + m.streamingText = []byte("partial answer so far") got := plainRender(t, m.interimBlock(96)) if !strings.Contains(got, "partial answer so far") { @@ -241,7 +241,7 @@ func TestWorkingTokenIndicatorAccumulatesAcrossSegmentClears(t *testing.T) { // Simulate the segment boundary that clears the live buffers, then stream // answer text in the next segment. m.streamingReasoning = "" - m.streamingText = "" + m.streamingText = nil updated, _ = m.Update(agentTextMsg{runID: rid, delta: strings.Repeat("b", 40)}) m = updated.(model) @@ -288,7 +288,7 @@ func TestInterimBlockShowsReasoningPreviewWhileThinking(t *testing.T) { m.now = func() time.Time { return base.Add(90 * time.Second) } m.turnStartedAt = base m.streamingReasoning = "analyzing the layout\nthe patch was corrupt so re-planning the css edits" - m.streamingText = "" // thinking phase: no answer yet + m.streamingText = nil // thinking phase: no answer yet got := plainRender(t, m.interimBlock(96)) if !strings.Contains(got, "re-planning the css edits") { @@ -307,7 +307,7 @@ func TestInterimBlockNoPreviewWhenReasoningExpanded(t *testing.T) { m.now = func() time.Time { return time.Date(2026, 6, 18, 23, 0, 0, 0, time.UTC) } m.streamingReasoningExpanded = true m.streamingReasoning = "only line of reasoning here" - m.streamingText = "" + m.streamingText = nil got := plainRender(t, m.interimBlock(96)) if strings.Count(got, "only line of reasoning here") != 1 { t.Fatalf("reasoning should appear exactly once when expanded (no preview dup):\n%s", got) From acf333bb66e52d9956a97bcae5ff9dc846eb5c7a Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 18:54:48 +0530 Subject: [PATCH 09/32] perf(tui): coalesce streamed text deltas to one frame MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every OnText delta was its own tea.Msg, so a fast provider (100+ tok/s) drove 100+ full Update->View cycles per second, each re-parsing the growing markdown — visible stutter over SSH and wasted CPU. Batch agentTextMsg deltas at the runtime sink over a ~16ms frame and forward them as a single message, decoupling render rate from token rate. Any non-text message flushes pending text first so ordering with tool-call / reasoning / row messages is preserved; the final agentResponseMsg (a tea.Cmd return, not a sink message) is safe because the model already drops deltas for an inactive runID. --- internal/tui/coalesce.go | 93 +++++++++++++++++++++++++ internal/tui/coalesce_test.go | 124 ++++++++++++++++++++++++++++++++++ internal/tui/run.go | 6 +- 3 files changed, 222 insertions(+), 1 deletion(-) create mode 100644 internal/tui/coalesce.go create mode 100644 internal/tui/coalesce_test.go diff --git a/internal/tui/coalesce.go b/internal/tui/coalesce.go new file mode 100644 index 00000000..3d064dc7 --- /dev/null +++ b/internal/tui/coalesce.go @@ -0,0 +1,93 @@ +package tui + +import ( + "sync" + "time" + + tea "charm.land/bubbletea/v2" +) + +// streamCoalesceInterval is roughly one 60fps frame. Assistant-text deltas that +// arrive within this window are merged into a single agentTextMsg, so the render +// rate decouples from the token rate: a fast local model (100+ tok/s) no longer +// forces 100+ full Update→View cycles (each re-parsing the growing markdown) per +// second. Rendering stays smooth regardless of provider speed. +const streamCoalesceInterval = 16 * time.Millisecond + +// textCoalescer batches agentTextMsg deltas before forwarding them to the Bubble +// Tea program. Any OTHER message flushes the pending text first, so ordering +// between streamed prose and tool-call / reasoning / row / usage messages is +// preserved. The turn's final agentResponseMsg does not pass through here (it is +// a tea.Cmd return, not a sink message), but the model drops deltas whose runID +// is no longer active, so a flush that races just past end-of-turn is harmless. +// +// Sink messages originate from the single agent goroutine and so arrive +// serially; the only concurrent caller is the flush timer. The mutex guards the +// buffer/timer against that one race. +type textCoalescer struct { + forward func(tea.Msg) // downstream sink (external sink + program.Send) + + mu sync.Mutex + buf []byte + runID int + timer *time.Timer +} + +func newTextCoalescer(forward func(tea.Msg)) *textCoalescer { + return &textCoalescer{forward: forward} +} + +// send is the coalescing entry point installed as the RuntimeMessageSink. +func (c *textCoalescer) send(msg tea.Msg) { + text, ok := msg.(agentTextMsg) + if !ok { + // Non-text message: flush buffered text first (preserving order), then + // forward it unchanged. + c.flush() + c.forward(msg) + return + } + + c.mu.Lock() + // A delta for a different run than the one buffered: flush the old run's text + // before buffering the new run's. In practice runs are sequential (the prior + // run's end already flushed via a non-text message), so this is belt-and-braces. + if len(c.buf) > 0 && text.runID != c.runID { + pending := c.drainLocked() + c.mu.Unlock() + c.forward(pending) + c.mu.Lock() + } + c.runID = text.runID + c.buf = append(c.buf, text.delta...) + if c.timer == nil { + c.timer = time.AfterFunc(streamCoalesceInterval, c.flush) + } + c.mu.Unlock() +} + +// flush forwards any buffered text as one agentTextMsg. Safe to call from the +// timer goroutine and inline; a no-op when nothing is buffered. +func (c *textCoalescer) flush() { + c.mu.Lock() + if len(c.buf) == 0 { + c.mu.Unlock() + return + } + msg := c.drainLocked() + c.mu.Unlock() + c.forward(msg) +} + +// drainLocked packages the buffer into an agentTextMsg and stops the timer. The +// caller holds c.mu. string(c.buf) copies, so reusing the backing array via [:0] +// is safe. +func (c *textCoalescer) drainLocked() agentTextMsg { + if c.timer != nil { + c.timer.Stop() + c.timer = nil + } + msg := agentTextMsg{runID: c.runID, delta: string(c.buf)} + c.buf = c.buf[:0] + return msg +} diff --git a/internal/tui/coalesce_test.go b/internal/tui/coalesce_test.go new file mode 100644 index 00000000..ee4ef265 --- /dev/null +++ b/internal/tui/coalesce_test.go @@ -0,0 +1,124 @@ +package tui + +import ( + "sync" + "testing" + "time" + + tea "charm.land/bubbletea/v2" +) + +// recorder is a thread-safe sink that captures forwarded messages in order. +type recorder struct { + mu sync.Mutex + msgs []tea.Msg +} + +func (r *recorder) forward(msg tea.Msg) { + r.mu.Lock() + defer r.mu.Unlock() + r.msgs = append(r.msgs, msg) +} + +func (r *recorder) snapshot() []tea.Msg { + r.mu.Lock() + defer r.mu.Unlock() + out := make([]tea.Msg, len(r.msgs)) + copy(out, r.msgs) + return out +} + +// Rapid deltas within one frame collapse into a single agentTextMsg carrying the +// concatenated text. +func TestCoalescerBatchesDeltas(t *testing.T) { + rec := &recorder{} + c := newTextCoalescer(rec.forward) + + c.send(agentTextMsg{runID: 1, delta: "Hel"}) + c.send(agentTextMsg{runID: 1, delta: "lo, "}) + c.send(agentTextMsg{runID: 1, delta: "world"}) + + // Nothing forwarded yet — still buffered within the frame. + if got := rec.snapshot(); len(got) != 0 { + t.Fatalf("deltas should buffer, forwarded %d early: %#v", len(got), got) + } + + c.flush() + + got := rec.snapshot() + if len(got) != 1 { + t.Fatalf("expected 1 coalesced message, got %d: %#v", len(got), got) + } + text, ok := got[0].(agentTextMsg) + if !ok || text.delta != "Hello, world" || text.runID != 1 { + t.Fatalf("coalesced message = %#v, want agentTextMsg{1, \"Hello, world\"}", got[0]) + } +} + +// A non-text message flushes buffered text first, so ordering (text before the +// tool call it precedes) is preserved. +func TestCoalescerFlushesTextBeforeOtherMessages(t *testing.T) { + rec := &recorder{} + c := newTextCoalescer(rec.forward) + + c.send(agentTextMsg{runID: 1, delta: "about to run"}) + c.send(toolCallStreamStartMsg{runID: 1, id: "t1", name: "bash"}) + + got := rec.snapshot() + if len(got) != 2 { + t.Fatalf("expected text then tool-call, got %d: %#v", len(got), got) + } + if text, ok := got[0].(agentTextMsg); !ok || text.delta != "about to run" { + t.Fatalf("first forwarded message must be the flushed text, got %#v", got[0]) + } + if _, ok := got[1].(toolCallStreamStartMsg); !ok { + t.Fatalf("second forwarded message must be the tool-call start, got %#v", got[1]) + } +} + +// A delta for a new run flushes the previous run's buffered text before buffering +// the new run's, so text is never mis-attributed across runs. +func TestCoalescerFlushesOnRunSwitch(t *testing.T) { + rec := &recorder{} + c := newTextCoalescer(rec.forward) + + c.send(agentTextMsg{runID: 1, delta: "run one text"}) + c.send(agentTextMsg{runID: 2, delta: "run two text"}) + c.flush() + + got := rec.snapshot() + if len(got) != 2 { + t.Fatalf("expected one message per run, got %d: %#v", len(got), got) + } + if first := got[0].(agentTextMsg); first.runID != 1 || first.delta != "run one text" { + t.Fatalf("first message = %#v, want run 1 text flushed on switch", got[0]) + } + if second := got[1].(agentTextMsg); second.runID != 2 || second.delta != "run two text" { + t.Fatalf("second message = %#v, want run 2 text", got[1]) + } +} + +// The frame timer flushes buffered text on its own without an explicit flush or a +// following message. +func TestCoalescerTimerFlushes(t *testing.T) { + rec := &recorder{} + c := newTextCoalescer(rec.forward) + + c.send(agentTextMsg{runID: 1, delta: "timer-driven"}) + + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + if len(rec.snapshot()) > 0 { + break + } + time.Sleep(streamCoalesceInterval) + } + + got := rec.snapshot() + if len(got) != 1 { + t.Fatalf("timer should have flushed exactly one message, got %d: %#v", len(got), got) + } + if text := got[0].(agentTextMsg); text.delta != "timer-driven" { + t.Fatalf("timer-flushed message = %#v", got[0]) + } +} diff --git a/internal/tui/run.go b/internal/tui/run.go index a4d73b1c..f2925dea 100644 --- a/internal/tui/run.go +++ b/internal/tui/run.go @@ -25,7 +25,7 @@ func Run(ctx context.Context, options Options) int { externalSink := options.RuntimeMessageSink var program *tea.Program - options.RuntimeMessageSink = func(msg tea.Msg) { + forward := func(msg tea.Msg) { if externalSink != nil { externalSink(msg) } @@ -33,6 +33,10 @@ func Run(ctx context.Context, options Options) int { program.Send(msg) } } + // Coalesce streamed assistant-text deltas to ~one frame each so a fast provider + // can't drive a full Update→View per token; every other message flushes pending + // text first, keeping order intact. + options.RuntimeMessageSink = newTextCoalescer(forward).send options.AltScreen = useAltScreen(options) programOpts := []tea.ProgramOption{ From 01cdc7576fb6db4833fc6d3b48bdca0e10e02e00 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 19:03:55 +0530 Subject: [PATCH 10/32] fix(agent): raise tool-failure stop threshold 4->6 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit guardrails.go halted a run after 4 consecutive same-error tool failures. A corrective hint fires at 2, and a model iterating on a genuinely tricky edit can legitimately fail a couple more times after the hint while converging — stopping at 4 cut those runs short. Raise to 6; the streak still resets the instant the tool succeeds or hits a different error, so only true same-error loops are affected. --- internal/agent/guardrails.go | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/internal/agent/guardrails.go b/internal/agent/guardrails.go index 089dc818..63238216 100644 --- a/internal/agent/guardrails.go +++ b/internal/agent/guardrails.go @@ -43,8 +43,12 @@ const ( toolFailureHintAt = 2 // toolFailureStopAt halts the run after a tool fails this many times in a row // with the same error, so NO model (weak or strong) burns turns looping on a - // bad call. - toolFailureStopAt = 4 + // bad call. Set to 6 (not 4): a corrective hint fires at toolFailureHintAt (2), + // and a model iterating on a genuinely tricky edit can legitimately fail a few + // times after the hint while converging — stopping at 4 cut those runs short. + // The streak still resets the moment the tool succeeds or hits a different + // error, so this only affects true same-error loops. + toolFailureStopAt = 6 // maxContinueNudges bounds how many times the headless completion gate // (Options.RequireCompletionSignal) re-prompts a model that stopped without a From 50d10ea7bbc624c3622c860f66a7901470f00788 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 19:05:07 +0530 Subject: [PATCH 11/32] fix(agent): more reconnect retries with jitter; narrow to transport errors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #19: bump maxStreamReconnects 2->4 and add up-to-50% jitter on the exponential backoff (capped at 8s), so a multi-second network blip is ridden out instead of killing the run on a 2s hiccup, and concurrent runs (swarms, cron fleets) don't reconnect in lockstep. #34: drop the '502'/'503' substring matches from shouldReconnect. 503 already exhausted providerio.SendWithRetry (retrying here is a redundant double-retry) and 502 is non-idempotent by providerio's rule (the POST may have been processed). Only genuine transport failures — where no response was received — reconnect now. --- internal/agent/reconnect.go | 43 ++++++++++++++++++++++++++++---- internal/agent/reconnect_test.go | 29 ++++++++++++++++++++- 2 files changed, 66 insertions(+), 6 deletions(-) diff --git a/internal/agent/reconnect.go b/internal/agent/reconnect.go index 641366a9..8e309c1b 100644 --- a/internal/agent/reconnect.go +++ b/internal/agent/reconnect.go @@ -3,6 +3,7 @@ package agent import ( "context" "fmt" + "math/rand" "strings" "time" @@ -17,10 +18,20 @@ import ( // times. We retry ONLY the connect (not a partially-consumed stream), so no // already-forwarded OnText is ever duplicated. const ( - maxStreamReconnects = 2 - streamReconnectBase = 500 * time.Millisecond + // maxStreamReconnects is how many times the connect is re-issued after a + // transient disconnect. 4 (not 2): with jittered exponential backoff this rides + // out a multi-second network blip (~0.5s + 1s + 2s + 4s worst case) instead of + // dying on a 2s hiccup and re-burning every token on a restart. + maxStreamReconnects = 4 + // streamReconnectMax caps a single backoff so the tail attempts don't wait + // minutes on a long outage. + streamReconnectMax = 8 * time.Second ) +// streamReconnectBase is the first-retry delay (doubled each subsequent attempt). +// A var, not a const, so tests can shrink it to keep the exhaustion path fast. +var streamReconnectBase = 500 * time.Millisecond + // reconnectNotifier is called before each retry with the 1-based attempt number // and the max, so the caller can surface a "Reconnecting N/max" notice. Nil is // fine. @@ -73,7 +84,7 @@ func streamWithReconnect(ctx context.Context, provider Provider, request zerorun if notify != nil { notify(attempt, maxStreamReconnects) } - if waitErr := sleepWithContext(ctx, backoffFor(attempt)); waitErr != nil { + if waitErr := sleepWithContext(ctx, jitteredBackoff(attempt)); waitErr != nil { return nil, err // ctx cancelled while waiting; surface the original error } stream, err = provider.StreamCompletion(ctx, request) @@ -96,6 +107,13 @@ func shouldReconnect(ctx context.Context, err error) bool { if isContextLimitError(msg) || isImageRejectionError(err) { return false } + // Transport-level disconnects only. HTTP 5xx statuses are deliberately NOT + // matched here: 503 already exhausted providerio.SendWithRetry (retrying it + // again is a redundant double-retry), and 502/500/504 are non-idempotent by + // providerio's rule (the completion POST may have been processed), so replaying + // the connect risks duplicate billable work. A genuine transport failure (EOF, + // reset, refused, timeout) means no response was received, which is safe to + // reconnect. for _, needle := range []string{ "eof", "connection reset", @@ -106,8 +124,6 @@ func shouldReconnect(ctx context.Context, err error) bool { "timed out", "temporarily unavailable", "i/o timeout", - "503", - "502", "server closed", "unexpected end", } { @@ -118,14 +134,31 @@ func shouldReconnect(ctx context.Context, err error) bool { return false } +// backoffFor is the deterministic exponential base delay for a 1-based attempt, +// capped at streamReconnectMax. Jitter is layered on separately (jitteredBackoff). func backoffFor(attempt int) time.Duration { d := streamReconnectBase for i := 1; i < attempt; i++ { + if d >= streamReconnectMax { + return streamReconnectMax + } d *= 2 } + if d > streamReconnectMax { + d = streamReconnectMax + } return d } +// jitteredBackoff adds up to 50% random jitter on top of backoffFor so concurrent +// runs (swarm members, a cron fleet) that all trip on the same outage don't +// reconnect in lockstep and hammer a recovering endpoint. Never shorter than the +// deterministic base, so backoff still grows attempt over attempt. +func jitteredBackoff(attempt int) time.Duration { + base := backoffFor(attempt) + return base + time.Duration(rand.Int63n(int64(base/2)+1)) +} + func sleepWithContext(ctx context.Context, d time.Duration) error { timer := time.NewTimer(d) defer timer.Stop() diff --git a/internal/agent/reconnect_test.go b/internal/agent/reconnect_test.go index 49e1714f..a7c3d38f 100644 --- a/internal/agent/reconnect_test.go +++ b/internal/agent/reconnect_test.go @@ -44,6 +44,10 @@ func TestStreamWithReconnectRecoversFromTransientDisconnect(t *testing.T) { } func TestStreamWithReconnectGivesUpAfterMax(t *testing.T) { + // Shrink the backoff so exhausting all retries stays fast (real base would + // sleep ~7.5s across 4 attempts). + defer func(orig time.Duration) { streamReconnectBase = orig }(streamReconnectBase) + streamReconnectBase = time.Millisecond // Always fails with a disconnect error → exhausts retries and returns it. p := &flakyProvider{failBefore: 99, failErr: errors.New("connection reset by peer")} _, err := streamWithReconnect(context.Background(), p, zeroruntime.CompletionRequest{}, nil) @@ -87,7 +91,7 @@ func TestShouldReconnectClassification(t *testing.T) { ctx := context.Background() disconnects := []string{ "unexpected EOF", "connection reset by peer", "broken pipe", - "i/o timeout", "503 Service Unavailable", "server closed the connection", + "i/o timeout", "server closed the connection", "connection refused", } for _, m := range disconnects { if !shouldReconnect(ctx, errors.New(m)) { @@ -97,6 +101,11 @@ func TestShouldReconnectClassification(t *testing.T) { notDisconnects := []string{ "context length exceeded", "invalid api key", "model not found", "400 bad request: unsupported parameter", + // HTTP 5xx statuses are handled by providerio.SendWithRetry (503) or are + // non-idempotent (502); the reconnect path must not double-retry them. (A + // "504 Gateway Timeout" still reconnects via the generic "timeout" transport + // signal — an accepted overlap, since a gateway timeout saw no response.) + "503 Service Unavailable", "502 Bad Gateway", } for _, m := range notDisconnects { if shouldReconnect(ctx, errors.New(m)) { @@ -112,6 +121,24 @@ func TestBackoffGrows(t *testing.T) { if backoffFor(2) != 2*streamReconnectBase { t.Fatalf("attempt 2 backoff = %v, want %v", backoffFor(2), 2*streamReconnectBase) } + // The exponential base is capped so late attempts don't wait minutes. + if got := backoffFor(20); got != streamReconnectMax { + t.Fatalf("attempt 20 backoff = %v, want cap %v", got, streamReconnectMax) + } +} + +func TestJitteredBackoffStaysInBounds(t *testing.T) { + // Jitter never drops below the deterministic base and never exceeds base*1.5, + // so backoff still grows attempt over attempt while decorrelating retries. + for attempt := 1; attempt <= 5; attempt++ { + base := backoffFor(attempt) + for i := 0; i < 200; i++ { + got := jitteredBackoff(attempt) + if got < base || got > base+base/2 { + t.Fatalf("attempt %d jittered backoff %v out of [%v, %v]", attempt, got, base, base+base/2) + } + } + } } func TestReconnectNoticeRoutesThroughReasoning(t *testing.T) { From 8a316ea26d29355be2e12ca21fb9156aa1af4e9b Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 19:08:34 +0530 Subject: [PATCH 12/32] fix(agent): reduce stall retries 2->1 to bound stuck-session time MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A no-output stream stall is detected only after the full stream idle timeout (~5min) elapses, and each retry can idle again — so 2 retries left an interactive session frozen for ~15min. Drop to 1 retry: keeps the common single-hiccup recovery while bounding the worst case to ~2x the idle timeout. --- internal/agent/loop.go | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/internal/agent/loop.go b/internal/agent/loop.go index badaf3f7..cc6338ea 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -25,7 +25,12 @@ const maxTurnsFinalAnswerPrompt = "You have reached the tool-turn limit. Do not // WITH NO OUTPUT yet is re-issued on a fresh connection before giving up. Only // the no-output case is retried (a partial turn would duplicate), so this is a // safe recovery for a stalled/dead pooled connection. -const maxStreamStallRetries = 2 +// +// Set to 1 (not 2): each attempt can itself idle for the full stream timeout +// (~5min) before the stall is even detected, so 2 retries left an interactive +// session frozen for ~15min. One retry keeps the common single-hiccup recovery +// while bounding the worst case to ~2× the idle timeout. +const maxStreamStallRetries = 1 const ( toolResultMetaControl = "control" From 36e4cc411fbaa8a2266655e81ad968bcbfb90870 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 19:19:12 +0530 Subject: [PATCH 13/32] fix(tui): serialize coalescer forwarding to preserve message order MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address CodeRabbit review on #7: flush() and the run-switch path drained the text buffer under c.mu but called forward() after releasing it, so a timer-fired text flush could race a concurrent non-text send() and land after the tool-call/reasoning message it should precede. Hold c.mu across drain AND forward (drainAndForwardLocked), so whoever holds the lock delivers atomically and the other caller blocks until it is done — text can never overtake a following non-text message. Added a concurrency stress test (timer racing an inline send) that passes under -race. --- internal/tui/coalesce.go | 44 +++++++++++++++++------------------ internal/tui/coalesce_test.go | 43 ++++++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 22 deletions(-) diff --git a/internal/tui/coalesce.go b/internal/tui/coalesce.go index 3d064dc7..628db1be 100644 --- a/internal/tui/coalesce.go +++ b/internal/tui/coalesce.go @@ -23,7 +23,9 @@ const streamCoalesceInterval = 16 * time.Millisecond // // Sink messages originate from the single agent goroutine and so arrive // serially; the only concurrent caller is the flush timer. The mutex guards the -// buffer/timer against that one race. +// buffer/timer AND is held across the downstream forward, so a timer-fired text +// flush can never overtake a concurrent non-text message: whoever holds the lock +// drains and forwards atomically, and the other caller blocks until it is done. type textCoalescer struct { forward func(tea.Msg) // downstream sink (external sink + program.Send) @@ -39,55 +41,53 @@ func newTextCoalescer(forward func(tea.Msg)) *textCoalescer { // send is the coalescing entry point installed as the RuntimeMessageSink. func (c *textCoalescer) send(msg tea.Msg) { + c.mu.Lock() + defer c.mu.Unlock() + text, ok := msg.(agentTextMsg) if !ok { // Non-text message: flush buffered text first (preserving order), then - // forward it unchanged. - c.flush() + // forward it — both under the lock so nothing can interleave between them. + c.drainAndForwardLocked() c.forward(msg) return } - c.mu.Lock() // A delta for a different run than the one buffered: flush the old run's text // before buffering the new run's. In practice runs are sequential (the prior // run's end already flushed via a non-text message), so this is belt-and-braces. if len(c.buf) > 0 && text.runID != c.runID { - pending := c.drainLocked() - c.mu.Unlock() - c.forward(pending) - c.mu.Lock() + c.drainAndForwardLocked() } c.runID = text.runID c.buf = append(c.buf, text.delta...) if c.timer == nil { c.timer = time.AfterFunc(streamCoalesceInterval, c.flush) } - c.mu.Unlock() } -// flush forwards any buffered text as one agentTextMsg. Safe to call from the -// timer goroutine and inline; a no-op when nothing is buffered. +// flush forwards any buffered text as one agentTextMsg. Runs on the timer +// goroutine; the lock it takes serializes it against send so its output can't be +// reordered around a concurrent non-text message. func (c *textCoalescer) flush() { c.mu.Lock() + defer c.mu.Unlock() + c.drainAndForwardLocked() +} + +// drainAndForwardLocked forwards any buffered text as one agentTextMsg and stops +// the timer, all while the caller holds c.mu — so a text flush and any non-text +// forward are strictly ordered and never interleave. A no-op when nothing is +// buffered. string(c.buf) copies, so reusing the backing array via [:0] is safe. +func (c *textCoalescer) drainAndForwardLocked() { if len(c.buf) == 0 { - c.mu.Unlock() return } - msg := c.drainLocked() - c.mu.Unlock() - c.forward(msg) -} - -// drainLocked packages the buffer into an agentTextMsg and stops the timer. The -// caller holds c.mu. string(c.buf) copies, so reusing the backing array via [:0] -// is safe. -func (c *textCoalescer) drainLocked() agentTextMsg { if c.timer != nil { c.timer.Stop() c.timer = nil } msg := agentTextMsg{runID: c.runID, delta: string(c.buf)} c.buf = c.buf[:0] - return msg + c.forward(msg) } diff --git a/internal/tui/coalesce_test.go b/internal/tui/coalesce_test.go index ee4ef265..34988595 100644 --- a/internal/tui/coalesce_test.go +++ b/internal/tui/coalesce_test.go @@ -98,6 +98,49 @@ func TestCoalescerFlushesOnRunSwitch(t *testing.T) { } } +// Under concurrency (the frame timer firing while a non-text message is sent), +// forwarding must stay serialized and ordered: buffered text is always delivered +// before a following non-text message, never after it. Run with -race to catch +// the interleaving CodeRabbit flagged. +func TestCoalescerOrdersTextBeforeNonTextUnderConcurrency(t *testing.T) { + for iter := 0; iter < 200; iter++ { + rec := &recorder{} + c := newTextCoalescer(rec.forward) + + // Buffer text, then let the timer race against an inline non-text send. + c.send(agentTextMsg{runID: 1, delta: "before-tool"}) + done := make(chan struct{}) + go func() { + c.send(toolCallStreamStartMsg{runID: 1, id: "t1", name: "bash"}) + close(done) + }() + <-done + + // Give the timer a chance to also fire, then settle. + time.Sleep(2 * streamCoalesceInterval) + c.flush() + + got := rec.snapshot() + textIdx, toolIdx := -1, -1 + for i, m := range got { + switch m.(type) { + case agentTextMsg: + if textIdx == -1 { + textIdx = i + } + case toolCallStreamStartMsg: + toolIdx = i + } + } + if textIdx == -1 || toolIdx == -1 { + t.Fatalf("iter %d: missing messages: %#v", iter, got) + } + if textIdx > toolIdx { + t.Fatalf("iter %d: text (%d) forwarded after tool-call (%d): %#v", iter, textIdx, toolIdx, got) + } + } +} + // The frame timer flushes buffered text on its own without an explicit flush or a // following message. func TestCoalescerTimerFlushes(t *testing.T) { From 352808ca07d0d2138611e4bcba4cc5291b462afb Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 19:26:13 +0530 Subject: [PATCH 14/32] feat(tui): accept /compact now MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bare /compact already triggers compaction, but /compact now — what users reach for when the context gauge climbs — hit the usage-error branch. Accept the 'now' keyword and advertise it in the command usage. --- internal/tui/commands.go | 36 ++++++++++++++++++++++++++++++-- internal/tui/compact_now_test.go | 27 ++++++++++++++++++++++++ internal/tui/session_controls.go | 6 ++++-- 3 files changed, 65 insertions(+), 4 deletions(-) create mode 100644 internal/tui/compact_now_test.go diff --git a/internal/tui/commands.go b/internal/tui/commands.go index a952784f..b31c7312 100644 --- a/internal/tui/commands.go +++ b/internal/tui/commands.go @@ -41,6 +41,10 @@ const ( commandAddDir commandSelfCorrect commandTurns + commandRetry + commandEdit + commandCopy + commandExport commandUnknown ) @@ -201,9 +205,9 @@ var commandDefinitions = []commandDefinition{ }, { name: "/compact", - usage: "/compact [status]", + usage: "/compact [status|now]", group: commandGroupSession, - description: "Show or request transcript compaction state.", + description: "Compact the transcript now, or show compaction state (/compact status).", kind: commandCompact, }, { @@ -249,6 +253,34 @@ var commandDefinitions = []commandDefinition{ description: "Show or set the per-run tool-turn budget for this session (raise it for long multi-step tasks).", kind: commandTurns, }, + { + name: "/retry", + usage: "/retry", + group: commandGroupSession, + description: "Resend your last prompt.", + kind: commandRetry, + }, + { + name: "/edit", + usage: "/edit", + group: commandGroupSession, + description: "Recall your last prompt into the composer to edit and resend.", + kind: commandEdit, + }, + { + name: "/copy", + usage: "/copy", + group: commandGroupSession, + description: "Copy the last answer to the clipboard.", + kind: commandCopy, + }, + { + name: "/export", + usage: "/export [path]", + group: commandGroupSession, + description: "Write the conversation transcript to a file.", + kind: commandExport, + }, { name: "/help", usage: "/help", diff --git a/internal/tui/compact_now_test.go b/internal/tui/compact_now_test.go new file mode 100644 index 00000000..c4392c4f --- /dev/null +++ b/internal/tui/compact_now_test.go @@ -0,0 +1,27 @@ +package tui + +import ( + "strings" + "testing" +) + +// "/compact now" is accepted as a manual trigger (not a usage error), while a +// bogus argument still reports usage. +func TestCompactNowAccepted(t *testing.T) { + m := model{} + + _, nowText, _ := m.handleCompactCommand("now") + if strings.Contains(nowText, "usage:") { + t.Fatalf("/compact now must be accepted as a trigger, got usage error: %q", nowText) + } + + _, statusText, _ := m.handleCompactCommand("status") + if strings.Contains(statusText, "usage:") { + t.Fatalf("/compact status must not error, got: %q", statusText) + } + + _, bogusText, _ := m.handleCompactCommand("frobnicate") + if !strings.Contains(bogusText, "usage:") { + t.Fatalf("/compact frobnicate must report usage, got: %q", bogusText) + } +} diff --git a/internal/tui/session_controls.go b/internal/tui/session_controls.go index dfc18f86..8bf09833 100644 --- a/internal/tui/session_controls.go +++ b/internal/tui/session_controls.go @@ -326,8 +326,10 @@ func (m model) handleCompactCommand(args string) (model, string, tea.Cmd) { if args == "status" { return m, m.compactText(false), nil } - if args != "" { - return m, "Compact\nusage: /compact [status]", nil + // Bare "/compact" already triggers compaction; accept "now" too since that is + // what users reach for when they see the context gauge climbing. + if args != "" && args != "now" { + return m, "Compact\nusage: /compact [status|now]", nil } if m.compactInFlight { return m, m.compactText(true), nil From 7740c273f6440d3b480c3f8cfa1448af12e530ed Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 19:27:17 +0530 Subject: [PATCH 15/32] feat(tui): add /retry, /edit, /copy, /export commands MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit /retry resends the last prompt; /edit recalls it into the composer to tweak and resend — both read a new lastPrompt field captured verbatim (pre-expansion) in launchPrompt. /copy puts the last answer on the clipboard via the existing copy machinery; /export writes a plain role-prefixed transcript to a file (timestamped default, or a given path). Real gaps for SSH users with no mouse select. --- internal/tui/model.go | 33 ++++++++++ internal/tui/retry_edit_test.go | 46 +++++++++++++ internal/tui/transcript_export.go | 81 +++++++++++++++++++++++ internal/tui/transcript_export_test.go | 91 ++++++++++++++++++++++++++ 4 files changed, 251 insertions(+) create mode 100644 internal/tui/retry_edit_test.go create mode 100644 internal/tui/transcript_export.go create mode 100644 internal/tui/transcript_export_test.go diff --git a/internal/tui/model.go b/internal/tui/model.go index 3634cfc5..7ba5badf 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -259,6 +259,9 @@ type model struct { headerPrinted bool // Composer input history (shell-style ↑/↓ recall of submitted inputs). + // lastPrompt is the verbatim text of the most recent submitted prompt, so + // /retry can resend it and /edit can recall it into the composer. + lastPrompt string // historyIdx == len(inputHistory) means "not navigating"; historyDraft // preserves whatever was typed before recall started. inputHistory []string @@ -3890,6 +3893,33 @@ func (m model) handleSubmit() (tea.Model, tea.Cmd) { } m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: "$ " + cmdText}) return m, runBashEscape(m.cwd, cmdText) + case commandRetry: + if m.pending { + m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: "Retry\ncannot retry while a run is in progress."}) + return m, nil + } + if strings.TrimSpace(m.lastPrompt) == "" { + m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: "Retry\nno previous prompt to resend."}) + return m, nil + } + return m.launchPrompt(m.lastPrompt) + case commandEdit: + if strings.TrimSpace(m.lastPrompt) == "" { + m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: "Edit\nno previous prompt to recall."}) + return m, nil + } + m.input.SetValue(m.lastPrompt) + return m, nil + case commandCopy: + text := m.lastAssistantAnswer() + if strings.TrimSpace(text) == "" { + m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: "Copy\nno answer to copy yet."}) + return m, nil + } + return m, copyTranscriptSelectionCmd(text) + case commandExport: + m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: m.handleExportCommand(command.text)}) + return m, nil case commandPrompt: if intent, ok := detectMCPSetupIntent(command.text); ok { return m.openMCPAddWizardFromIntent(intent), nil @@ -3904,6 +3934,9 @@ func (m model) handleSubmit() (tea.Model, tea.Cmd) { // composer. Queued prompts use this path too, so session and image behavior // stays identical to immediate submissions. func (m model) launchPrompt(prompt string) (model, tea.Cmd) { + // Remember the verbatim prompt (before specialist/document expansion) so /retry + // and /edit can act on exactly what the user submitted. + m.lastPrompt = prompt m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendUser, text: prompt}) if m.provider == nil { m.transcript = reduceTranscript(m.transcript, transcriptAction{ diff --git a/internal/tui/retry_edit_test.go b/internal/tui/retry_edit_test.go new file mode 100644 index 00000000..b4d43a4e --- /dev/null +++ b/internal/tui/retry_edit_test.go @@ -0,0 +1,46 @@ +package tui + +import ( + "context" + "testing" + + tea "charm.land/bubbletea/v2" +) + +func TestCommandsRegistered(t *testing.T) { + for _, name := range []string{"/retry", "/edit", "/copy", "/export"} { + if cmd, ok := resolveCommand(name); !ok { + t.Fatalf("%s should be a registered command", name) + } else if cmd.name != name { + t.Fatalf("resolveCommand(%q) = %q", name, cmd.name) + } + } +} + +// /edit recalls the last prompt into the composer for editing. +func TestEditRecallsLastPrompt(t *testing.T) { + m := newModel(context.Background(), Options{}) + m.lastPrompt = "refactor the parser" + m.input.SetValue("/edit") + + updated, _ := m.Update(testKey(tea.KeyEnter)) + next := updated.(model) + + if got := next.composerValue(); got != "refactor the parser" { + t.Fatalf("/edit should recall last prompt into composer, got %q", got) + } +} + +// /retry with no prior prompt reports that there's nothing to resend rather than +// launching an empty run. +func TestRetryWithoutPriorPromptIsNoOp(t *testing.T) { + m := newModel(context.Background(), Options{}) + m.input.SetValue("/retry") + + updated, _ := m.Update(testKey(tea.KeyEnter)) + next := updated.(model) + + if !transcriptContains(next.transcript, "no previous prompt") { + t.Fatalf("/retry with no history should note there's nothing to resend, got %#v", next.transcript) + } +} diff --git a/internal/tui/transcript_export.go b/internal/tui/transcript_export.go new file mode 100644 index 00000000..4bea03e7 --- /dev/null +++ b/internal/tui/transcript_export.go @@ -0,0 +1,81 @@ +package tui + +import ( + "fmt" + "os" + "path/filepath" + "strings" +) + +// lastAssistantAnswer returns the text of the most recent assistant row — the +// last final answer if one exists, else the last assistant row of any kind. Empty +// when the conversation has no assistant text yet. +func (m model) lastAssistantAnswer() string { + lastAny := "" + for i := len(m.transcript) - 1; i >= 0; i-- { + row := m.transcript[i] + if row.kind != rowAssistant { + continue + } + if row.final { + return row.text + } + if lastAny == "" { + lastAny = row.text + } + } + return lastAny +} + +// plainTranscriptText renders the conversation as a plain, role-prefixed text +// document for /export — the readable content (user prompts, assistant answers, +// system notes, errors), skipping tool-call/permission UI noise. +func (m model) plainTranscriptText() string { + var b strings.Builder + for _, row := range m.transcript { + var prefix string + switch row.kind { + case rowUser: + prefix = "you: " + case rowAssistant: + prefix = "zero: " + case rowSystem: + prefix = "· " + case rowError: + prefix = "error: " + default: + continue + } + text := strings.TrimRight(row.text, "\n") + if strings.TrimSpace(text) == "" { + continue + } + b.WriteString(prefix) + b.WriteString(text) + b.WriteString("\n\n") + } + return b.String() +} + +// handleExportCommand writes the transcript to a file and returns a status +// message. With no argument it derives a timestamped filename in the workspace; +// a relative path is resolved against the workspace root. +func (m model) handleExportCommand(args string) string { + body := m.plainTranscriptText() + if strings.TrimSpace(body) == "" { + return "Export\nnothing to export yet." + } + + path := strings.TrimSpace(args) + if path == "" { + stamp := m.now().Format("20060102-150405") + path = fmt.Sprintf("zero-transcript-%s.txt", stamp) + } + if !filepath.IsAbs(path) && m.cwd != "" { + path = filepath.Join(m.cwd, path) + } + if err := os.WriteFile(path, []byte(body), 0o644); err != nil { + return "Export\nfailed to write " + path + ": " + err.Error() + } + return "Export\nwrote transcript to " + path +} diff --git a/internal/tui/transcript_export_test.go b/internal/tui/transcript_export_test.go new file mode 100644 index 00000000..960990ca --- /dev/null +++ b/internal/tui/transcript_export_test.go @@ -0,0 +1,91 @@ +package tui + +import ( + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +func TestLastAssistantAnswerPrefersFinal(t *testing.T) { + m := model{transcript: []transcriptRow{ + {kind: rowUser, text: "hi"}, + {kind: rowAssistant, text: "interim narration"}, + {kind: rowAssistant, text: "the final answer", final: true}, + {kind: rowSystem, text: "worked for 3s"}, + }} + if got := m.lastAssistantAnswer(); got != "the final answer" { + t.Fatalf("lastAssistantAnswer = %q, want the final answer", got) + } + + // With no final row, falls back to the most recent assistant row. + m2 := model{transcript: []transcriptRow{ + {kind: rowAssistant, text: "first"}, + {kind: rowAssistant, text: "second"}, + }} + if got := m2.lastAssistantAnswer(); got != "second" { + t.Fatalf("fallback lastAssistantAnswer = %q, want second", got) + } + + // Empty when there's no assistant text. + if got := (model{transcript: []transcriptRow{{kind: rowUser, text: "hi"}}}).lastAssistantAnswer(); got != "" { + t.Fatalf("expected empty answer, got %q", got) + } +} + +func TestPlainTranscriptTextSkipsNoise(t *testing.T) { + m := model{transcript: []transcriptRow{ + {kind: rowUser, text: "add a flag"}, + {kind: rowToolCall, tool: "bash", text: "go build"}, + {kind: rowToolResult, tool: "bash", text: "ok"}, + {kind: rowAssistant, text: "Done — added --version.", final: true}, + {kind: rowError, text: "provider error: boom"}, + }} + out := m.plainTranscriptText() + if !strings.Contains(out, "you: add a flag") || !strings.Contains(out, "zero: Done — added --version.") { + t.Fatalf("export missing conversation text:\n%s", out) + } + if !strings.Contains(out, "error: provider error: boom") { + t.Fatalf("export should include error rows:\n%s", out) + } + if strings.Contains(out, "go build") || strings.Contains(out, "bash") { + t.Fatalf("export should skip tool-call noise:\n%s", out) + } +} + +func TestHandleExportCommandWritesFile(t *testing.T) { + dir := t.TempDir() + m := model{ + cwd: dir, + now: func() time.Time { return time.Date(2026, 7, 4, 9, 30, 0, 0, time.UTC) }, + transcript: []transcriptRow{ + {kind: rowUser, text: "hello"}, + {kind: rowAssistant, text: "hi there", final: true}, + }, + } + + // Explicit relative path resolves against cwd. + msg := m.handleExportCommand("out.txt") + if !strings.Contains(msg, "wrote transcript to") { + t.Fatalf("export status = %q", msg) + } + data, err := os.ReadFile(filepath.Join(dir, "out.txt")) + if err != nil { + t.Fatalf("reading exported file: %v", err) + } + if !strings.Contains(string(data), "you: hello") || !strings.Contains(string(data), "zero: hi there") { + t.Fatalf("exported content = %q", string(data)) + } + + // No-arg export derives a timestamped filename in cwd. + if msg := m.handleExportCommand(""); !strings.Contains(msg, "zero-transcript-20260704-093000.txt") { + t.Fatalf("default export status = %q", msg) + } + + // Nothing to export on an empty conversation. + empty := model{cwd: dir, now: m.now} + if msg := empty.handleExportCommand(""); !strings.Contains(msg, "nothing to export") { + t.Fatalf("empty export status = %q", msg) + } +} From 17485a5cab9b735598743e98c440ffd80e30c37e Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 19:31:42 +0530 Subject: [PATCH 16/32] feat(tools): separator-insensitive tool_search matching tool_search ranked deferred tools by exact substring, so a model that typed 'webfetch' (dropping the underscore) matched nothing and looped. Add a separator-squashing fallback ('web_fetch' -> 'webfetch') scored below an exact substring match, so precise queries still rank first. --- internal/tools/tool_search.go | 21 ++++++++++++++++++++- internal/tools/tool_search_test.go | 23 +++++++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/internal/tools/tool_search.go b/internal/tools/tool_search.go index 85134167..773f4442 100644 --- a/internal/tools/tool_search.go +++ b/internal/tools/tool_search.go @@ -261,10 +261,16 @@ func (tool toolSearchTool) rankByKeyword(query string, deferred []Tool) []Tool { for index, candidate := range deferred { name := strings.ToLower(candidate.Name()) desc := strings.ToLower(candidate.Description()) + nameSquashed := squashSeparators(name) score := 0 for _, keyword := range keywords { - if strings.Contains(name, keyword) { + switch { + case strings.Contains(name, keyword): score += 2 + case strings.Contains(nameSquashed, squashSeparators(keyword)): + // Separator-insensitive fallback: "webfetch" matches "web_fetch". + // Ranked below an exact substring so precise queries still win. + score++ } if strings.Contains(desc, keyword) { score++ @@ -290,6 +296,19 @@ func (tool toolSearchTool) rankByKeyword(query string, deferred []Tool) []Tool { return matches } +// squashSeparators drops separators so "web_fetch", "web-fetch", and "web fetch" +// all normalize to "webfetch". It lets a query missing the separators still match +// a tool name, the common way a model mistypes a tool ("webfetch"). +func squashSeparators(s string) string { + return strings.Map(func(r rune) rune { + switch r { + case '_', '-', ' ', '.': + return -1 + } + return r + }, s) +} + // noMatchMessage reports that nothing loaded and names the available deferred // tools so the model can retry with a valid select: query. func (tool toolSearchTool) noMatchMessage(query string, deferred []Tool) string { diff --git a/internal/tools/tool_search_test.go b/internal/tools/tool_search_test.go index a4270a31..3d07481c 100644 --- a/internal/tools/tool_search_test.go +++ b/internal/tools/tool_search_test.go @@ -184,6 +184,29 @@ func TestToolSearchKeywordRanksByNameThenDescription(t *testing.T) { } } +// A query missing the name's separators still matches: "webfetch" -> web_fetch. +func TestToolSearchKeywordMatchesSeparatorInsensitive(t *testing.T) { + reg := NewRegistry() + reg.Register(searchFakeTool{ + name: "web_fetch", + description: "Fetch a URL.", + parameters: Schema{Type: "object", AdditionalProperties: false}, + }) + reg.Register(searchFakeTool{ + name: "stock_quote", + description: "Get a stock price.", + parameters: Schema{Type: "object", AdditionalProperties: false}, + }) + tool := NewToolSearchTool(reg).(optionsAwareTool) + + result := tool.RunWithOptions(context.Background(), + map[string]any{"query": "webfetch"}, RunOptions{}) + + if got := result.Meta["load_tools"]; got != "web_fetch" { + t.Fatalf("load_tools = %q, want web_fetch (separator-insensitive match)", got) + } +} + func TestToolSearchKeywordExcludesNonMatches(t *testing.T) { reg := newDeferredFixtureRegistry() tool := NewToolSearchTool(reg).(optionsAwareTool) From 01e4fc6c5a497703e26042e5585439e31ed3d8d7 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 19:47:10 +0530 Subject: [PATCH 17/32] feat(tools): signal ask_user dismissal distinctly from a blank field MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Empty answers all rendered as '(no answer provided)', so the model couldn't tell a wholesale dismissal (user closed the prompt without answering anything) from one field left blank amid real answers — and might invent a default. FormatAskUserAnswers now flags a full dismissal up front as a skip and marks individual empties '(left blank)' vs '(skipped)'. --- internal/tools/ask_user.go | 25 +++++++++++++++++-- internal/tools/ask_user_format_test.go | 34 ++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) create mode 100644 internal/tools/ask_user_format_test.go diff --git a/internal/tools/ask_user.go b/internal/tools/ask_user.go index b1ba3328..cbd5f227 100644 --- a/internal/tools/ask_user.go +++ b/internal/tools/ask_user.go @@ -272,15 +272,36 @@ func coerceAskUserOptions(optionsValue, descsValue any) (labels []string, descri // FormatAskUserAnswers renders question/answer pairs into a clear, model-readable // block. Missing answers are surfaced explicitly so the model never silently // treats an unanswered question as answered. +// +// It distinguishes two shapes of "empty" the model would otherwise conflate: a +// wholesale dismissal (the user closed the prompt without answering ANYTHING) is +// flagged up front as a skip, so the model doesn't invent a default; a single +// blank field amid other answers is marked "(left blank)". func FormatAskUserAnswers(questions []AskUserQuestion, answers []string) string { - lines := make([]string, 0, len(questions)*3) + anyAnswered := false + for index := range questions { + if index < len(answers) && strings.TrimSpace(answers[index]) != "" { + anyAnswered = true + break + } + } + + lines := make([]string, 0, len(questions)*3+2) + if len(questions) > 0 && !anyAnswered { + lines = append(lines, "[note] The user dismissed this prompt without answering. Treat this as a skip: do not assume a default — ask again more specifically, or proceed only if the intent is already unambiguous.") + lines = append(lines, "") + } for index, question := range questions { answer := "" if index < len(answers) { answer = strings.TrimSpace(answers[index]) } if answer == "" { - answer = "(no answer provided)" + if anyAnswered { + answer = "(left blank)" + } else { + answer = "(skipped)" + } } lines = append(lines, fmt.Sprintf("%d. [question] %s", index+1, question.Question)) lines = append(lines, "[answer] "+answer) diff --git a/internal/tools/ask_user_format_test.go b/internal/tools/ask_user_format_test.go new file mode 100644 index 00000000..731ccde3 --- /dev/null +++ b/internal/tools/ask_user_format_test.go @@ -0,0 +1,34 @@ +package tools + +import ( + "strings" + "testing" +) + +func TestFormatAskUserAnswersDistinguishesSkipFromBlank(t *testing.T) { + questions := []AskUserQuestion{ + {Question: "Which database?"}, + {Question: "Migrate existing data?"}, + } + + // Wholesale dismissal (nothing answered): flagged as a skip up front. + dismissed := FormatAskUserAnswers(questions, []string{"", ""}) + if !strings.Contains(dismissed, "dismissed this prompt") || !strings.Contains(dismissed, "Treat this as a skip") { + t.Fatalf("all-empty answers must be flagged as a skip, got:\n%s", dismissed) + } + if strings.Contains(dismissed, "left blank") { + t.Fatalf("a full dismissal should read as skipped, not left-blank:\n%s", dismissed) + } + if !strings.Contains(dismissed, "(skipped)") { + t.Fatalf("dismissed questions should render (skipped):\n%s", dismissed) + } + + // Partial answer: the blank one is "left blank", no wholesale-skip note. + partial := FormatAskUserAnswers(questions, []string{"postgres", ""}) + if strings.Contains(partial, "dismissed this prompt") { + t.Fatalf("a partial answer is not a dismissal:\n%s", partial) + } + if !strings.Contains(partial, "postgres") || !strings.Contains(partial, "(left blank)") { + t.Fatalf("partial answers should show the answer and mark the blank one left-blank:\n%s", partial) + } +} From 52855eb1a18cdc23d91cb5d27aede6ed1f71d340 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 20:04:33 +0530 Subject: [PATCH 18/32] feat(tui): fuzzy ranking in model/session pickers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pickers filtered by flat strings.Contains with no ranking, so finding 'sonnet 4.5' among 50+ models meant an exact prefix or scrolling. Rank matches (exact < prefix < contains < subsequence, reusing fuzzySubsequenceGap) so the closest match lands on top, and 'snt45' now matches 'Sonnet 4.5'. Groups stay contiguous — ordered by their best match, never split — so the grouped model picker still renders one header per provider. Covers both /model and /resume (shared commandPicker). --- internal/tui/picker.go | 75 +++++++++++++++++++++++++++++-- internal/tui/picker_fuzzy_test.go | 72 +++++++++++++++++++++++++++++ 2 files changed, 143 insertions(+), 4 deletions(-) create mode 100644 internal/tui/picker_fuzzy_test.go diff --git a/internal/tui/picker.go b/internal/tui/picker.go index 29d0deb0..cda22a65 100644 --- a/internal/tui/picker.go +++ b/internal/tui/picker.go @@ -102,16 +102,83 @@ func (p *commandPicker) applyQuery() { p.selected = clampInt(p.selected, 0, maxInt(0, len(p.items)-1)) return } - filtered := make([]pickerItem, 0, len(source)) - for _, item := range source { - if strings.Contains(strings.ToLower(strings.Join([]string{item.Group, item.Label, item.Value, item.Meta}, " ")), query) { - filtered = append(filtered, item) + + // Rank matches (exact < prefix < contains < subsequence) instead of a flat + // substring filter, so the closest match to "sonnet 4.5" lands at the top + // rather than buried behind scrolling. Groups stay contiguous — a group is + // ordered by its best-matching item and never split into two header blocks — + // so the grouped model picker still renders one header per provider. + type entry struct { + item pickerItem + score int + order int + } + groupFirst := map[string]int{} + groupBest := map[string]int{} + entries := make([]entry, 0, len(source)) + for index, item := range source { + score, ok := scorePickerItem(item, query) + if !ok { + continue + } + if _, seen := groupFirst[item.Group]; !seen { + groupFirst[item.Group] = index + } + if best, seen := groupBest[item.Group]; !seen || score < best { + groupBest[item.Group] = score + } + entries = append(entries, entry{item: item, score: score, order: index}) + } + sort.SliceStable(entries, func(a, b int) bool { + ga, gb := entries[a].item.Group, entries[b].item.Group + if ga != gb { + // Most-relevant group first; ties keep original group appearance order so + // a group is never scattered across the list. + if groupBest[ga] != groupBest[gb] { + return groupBest[ga] < groupBest[gb] + } + return groupFirst[ga] < groupFirst[gb] } + // Within a group: best match first, then original order. + if entries[a].score != entries[b].score { + return entries[a].score < entries[b].score + } + return entries[a].order < entries[b].order + }) + filtered := make([]pickerItem, 0, len(entries)) + for _, e := range entries { + filtered = append(filtered, e.item) } p.items = filtered p.selected = 0 } +// scorePickerItem ranks an item against a lowercased query; lower is better, and +// ok is false when it doesn't match at all. Tiers mirror scoreFileSuggestion: +// exact/prefix/contains on the label (what the user reads) beat matches deeper in +// the joined haystack, and a fuzzy subsequence is the last-resort match. +func scorePickerItem(item pickerItem, query string) (int, bool) { + label := strings.ToLower(item.Label) + hay := strings.ToLower(strings.Join([]string{item.Group, item.Label, item.Value, item.Meta}, " ")) + switch { + case label == query: + return 0, true + case strings.HasPrefix(label, query): + return 20, true + case strings.Contains(label, query): + return 40, true + case strings.HasPrefix(hay, query): + return 60, true + case strings.Contains(hay, query): + return 80, true + default: + if gap, ok := fuzzySubsequenceGap(hay, query); ok { + return 120 + gap, true + } + return 0, false + } +} + // newModelPicker lists active (non-deprecated) models, preselecting the active // one. Returns nil when the catalog is unavailable so the caller falls back to // the plain status text. diff --git a/internal/tui/picker_fuzzy_test.go b/internal/tui/picker_fuzzy_test.go new file mode 100644 index 00000000..652b3e69 --- /dev/null +++ b/internal/tui/picker_fuzzy_test.go @@ -0,0 +1,72 @@ +package tui + +import ( + "testing" +) + +func newFuzzyTestPicker(items []pickerItem) *commandPicker { + return &commandPicker{ + kind: pickerModel, + items: append([]pickerItem{}, items...), + allItems: append([]pickerItem{}, items...), + } +} + +// A closer match ranks above a merely-containing one, so the best result lands at +// the top instead of in original order. +func TestPickerRanksBestMatchFirst(t *testing.T) { + p := newFuzzyTestPicker([]pickerItem{ + {Label: "Claude Opus 4.5", Value: "opus-4.5"}, + {Label: "Claude Sonnet 4.5", Value: "sonnet-4.5"}, + {Label: "Sonnet", Value: "sonnet"}, // exact-ish prefix match should win + }) + p.query = "sonnet" + p.applyQuery() + + if len(p.items) != 2 { + t.Fatalf("expected 2 matches for 'sonnet', got %d: %#v", len(p.items), p.items) + } + if p.items[0].Value != "sonnet" { + t.Fatalf("exact/prefix match should rank first, got %q", p.items[0].Value) + } +} + +// A subsequence query (non-contiguous characters) still matches when no substring does. +func TestPickerFuzzySubsequenceMatches(t *testing.T) { + p := newFuzzyTestPicker([]pickerItem{ + {Label: "Claude Sonnet 4.5", Value: "sonnet-4.5"}, + {Label: "GPT-5", Value: "gpt-5"}, + }) + p.query = "snt45" // subsequence of "sonnet 4.5", not a substring + p.applyQuery() + + if len(p.items) != 1 || p.items[0].Value != "sonnet-4.5" { + t.Fatalf("subsequence query should match sonnet-4.5, got %#v", p.items) + } +} + +// Ranking keeps each group contiguous: a filtered result never splits one group +// into two separate blocks, so group headers still render once per group. +func TestPickerRankingKeepsGroupsContiguous(t *testing.T) { + p := newFuzzyTestPicker([]pickerItem{ + {Group: "openai", Label: "gpt-5 mini", Value: "a"}, + {Group: "anthropic", Label: "sonnet 4.5", Value: "b"}, + {Group: "openai", Label: "gpt-5 codex", Value: "c"}, + {Group: "anthropic", Label: "opus 4.5", Value: "d"}, + }) + p.query = "5" // matches both openai models and the "4.5" anthropic ones + p.applyQuery() + + // Every item's group must appear as one contiguous run. + seen := map[string]bool{} + last := "" + for _, item := range p.items { + if item.Group != last { + if seen[item.Group] { + t.Fatalf("group %q split into non-contiguous blocks: %#v", item.Group, p.items) + } + seen[item.Group] = true + last = item.Group + } + } +} From 22bebbaf3b9647f2ac8dd28d0bdc7ec3f8e2c83b Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 20:39:21 +0530 Subject: [PATCH 19/32] perf(agent): cache per-tool schema render across turns MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit partitionTools re-ran the recursive schema->map conversion (schemaToRuntimeMap) for every tool on every turn — redundant, since a tool's advertised name/description/schema is stable for the run. Add a per-run definition cache keyed by tool name (nil-safe; the plain partitionTools entrypoint and tests pass nil for fresh renders). The partitioning itself (visibility, deferral, ordering) still recomputes every turn — it must, because a tool's deferred state can flip mid-run (swarm tools un-defer once a swarm is active). Only the expensive schema render is memoized. tool_search is excluded by its callers (dynamic description) so it never poisons the cache. --- internal/agent/loop.go | 42 +++++++++++++++-- internal/agent/partition_cache_test.go | 63 ++++++++++++++++++++++++++ 2 files changed, 101 insertions(+), 4 deletions(-) create mode 100644 internal/agent/partition_cache_test.go diff --git a/internal/agent/loop.go b/internal/agent/loop.go index 65b7e896..fa6ecb11 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -150,6 +150,11 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) // has already been demanded this run, so it fires at most once. acceptanceRequested := false + // toolDefCache memoizes each tool's rendered JSON-schema definition across + // turns (a tool's advertised schema is stable for the run), so partitionTools + // doesn't re-run the recursive schema→map conversion for every tool every turn. + toolDefCache := map[string]zeroruntime.ToolDefinition{} + result := Result{Messages: copyMessages(messages)} for turn := 0; turn < maxTurns; turn++ { result.Turns = turn + 1 @@ -158,7 +163,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) // the tool-definition tokens (they ride on every request) in its estimate. // partitionTools depends only on registry/permissions/options/loaded, not on // the messages, so computing it before compaction is safe. - exposed, _ := partitionTools(registry, permissionMode, options, loaded) + exposed, _ := partitionToolsCached(registry, permissionMode, options, loaded, toolDefCache) // PROACTIVE compaction: if the history is approaching the model's // context window, summarize the oldest middle before building the @@ -2467,6 +2472,17 @@ func permissionActionFromSandbox(action sandbox.Action) PermissionAction { // exposed. The exposed slice is alpha-sorted by name, matching the legacy order // so the inactive path is stable. func partitionTools(registry *tools.Registry, permissionMode PermissionMode, options Options, loaded map[string]bool) ([]zeroruntime.ToolDefinition, string) { + return partitionToolsCached(registry, permissionMode, options, loaded, nil) +} + +// partitionToolsCached is partitionTools with an optional per-tool definition +// cache. The partitioning itself (visibility, deferral, ordering) is recomputed +// every call — it must be, because a tool's deferred state can flip mid-run (e.g. +// swarm tools un-defer once a swarm is active). Only the expensive part — +// rendering each tool's JSON-schema parameters — is memoized by tool name, since a +// tool's advertised name/description/schema is stable for the run. defCache nil +// disables caching (used by tests and the plain partitionTools entrypoint). +func partitionToolsCached(registry *tools.Registry, permissionMode PermissionMode, options Options, loaded map[string]bool, defCache map[string]zeroruntime.ToolDefinition) ([]zeroruntime.ToolDefinition, string) { registeredTools := registry.All() visible := make([]tools.Tool, 0, len(registeredTools)) @@ -2510,7 +2526,7 @@ func partitionTools(registry *tools.Registry, permissionMode PermissionMode, opt if tool.Name() == tools.ToolSearchToolName { continue } - definitions = append(definitions, runtimeToolDefinition(tool)) + definitions = append(definitions, cachedRuntimeToolDefinition(defCache, tool)) } sort.Slice(definitions, func(left int, right int) bool { return definitions[left].Name < definitions[right].Name @@ -2542,13 +2558,13 @@ func partitionTools(registry *tools.Registry, permissionMode PermissionMode, opt } if tools.IsDeferred(tool) { if loaded[name] { - loadedTail = append(loadedTail, runtimeToolDefinition(tool)) + loadedTail = append(loadedTail, cachedRuntimeToolDefinition(defCache, tool)) } else { hiddenTools = append(hiddenTools, tool) } continue } - eager = append(eager, runtimeToolDefinition(tool)) + eager = append(eager, cachedRuntimeToolDefinition(defCache, tool)) } sort.Slice(eager, func(left int, right int) bool { return eager[left].Name < eager[right].Name @@ -2582,6 +2598,24 @@ func partitionTools(registry *tools.Registry, permissionMode PermissionMode, opt return definitions, discovery } +// cachedRuntimeToolDefinition returns the tool's rendered definition, reusing a +// cached render when defCache holds one for this tool name. A tool's advertised +// definition is stable across a run, so caching skips the recursive schema→map +// conversion (schemaToRuntimeMap) that would otherwise run for every tool on every +// turn. tool_search is excluded by its callers (its description is dynamic), so it +// never poisons the cache. A nil cache computes fresh. +func cachedRuntimeToolDefinition(defCache map[string]zeroruntime.ToolDefinition, tool tools.Tool) zeroruntime.ToolDefinition { + if defCache == nil { + return runtimeToolDefinition(tool) + } + if def, ok := defCache[tool.Name()]; ok { + return def + } + def := runtimeToolDefinition(tool) + defCache[tool.Name()] = def + return def +} + // runtimeToolDefinition renders a tool's advertised definition (name, description, // JSON-schema parameters) as sent to the provider. func runtimeToolDefinition(tool tools.Tool) zeroruntime.ToolDefinition { diff --git a/internal/agent/partition_cache_test.go b/internal/agent/partition_cache_test.go new file mode 100644 index 00000000..9983afb3 --- /dev/null +++ b/internal/agent/partition_cache_test.go @@ -0,0 +1,63 @@ +package agent + +import ( + "context" + "reflect" + "testing" + + "github.com/Gitlawb/zero/internal/tools" + "github.com/Gitlawb/zero/internal/zeroruntime" +) + +// countingSchemaTool records how many times its Parameters() schema is read, so a +// test can prove the definition cache avoids re-rendering across turns. +type countingSchemaTool struct { + name string + calls *int +} + +func (t countingSchemaTool) Name() string { return t.name } +func (t countingSchemaTool) Description() string { return "counts schema reads" } +func (t countingSchemaTool) Parameters() tools.Schema { + *t.calls++ + return tools.Schema{Type: "object", AdditionalProperties: false, Properties: map[string]tools.PropertySchema{ + "x": {Type: "string"}, + }} +} +func (t countingSchemaTool) Safety() tools.Safety { + return tools.Safety{SideEffect: tools.SideEffectRead, Permission: tools.PermissionAllow} +} +func (t countingSchemaTool) Run(_ context.Context, _ map[string]any) tools.Result { + return tools.Result{Status: tools.StatusOK} +} + +// The cache renders a tool's schema once and reuses it across calls, and its +// output is identical to the uncached path. +func TestPartitionToolsCacheRendersOnceAndMatchesUncached(t *testing.T) { + calls := 0 + registry := tools.NewRegistry() + registry.Register(countingSchemaTool{name: "alpha", calls: &calls}) + registry.Register(countingSchemaTool{name: "beta", calls: &calls}) + + uncached, _ := partitionTools(registry, PermissionModeAuto, Options{}, map[string]bool{}) + + base := calls + cache := map[string]zeroruntime.ToolDefinition{} + first, _ := partitionToolsCached(registry, PermissionModeAuto, Options{}, map[string]bool{}, cache) + rendersAfterFirst := calls - base + second, _ := partitionToolsCached(registry, PermissionModeAuto, Options{}, map[string]bool{}, cache) + rendersAfterSecond := calls - base + + if rendersAfterFirst != 2 { + t.Fatalf("first cached call should render 2 tools once each, rendered %d times", rendersAfterFirst) + } + if rendersAfterSecond != 2 { + t.Fatalf("second cached call must reuse the cache (no new schema reads), total renders = %d", rendersAfterSecond) + } + if !reflect.DeepEqual(first, second) { + t.Fatal("cached calls must return identical definitions") + } + if !reflect.DeepEqual(first, uncached) { + t.Fatal("cached output must match the uncached partitionTools output") + } +} From 368af701b1e1780c0c961063c96a21d6fcc5df8f Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 20:45:35 +0530 Subject: [PATCH 20/32] perf(agent): calibrate compaction estimate against real prompt tokens MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The byte/4 heuristic (ApproxTextTokens) over-counts code-heavy content ~15-20%, so with triggerRatio=0.7 compaction fired at ~60% of true capacity — premature summarizer calls that degrade quality. Each turn now folds (rawEstimate, provider InputTokens) into an EMA calibration ratio, clamped per-sample to a sane band; maybeCompact and the reactive path scale their estimates by it. Turn 1 uses the raw estimate (no data yet); later turns compact near real capacity. --- internal/agent/compaction.go | 45 +++++++++++++++++-- internal/agent/compaction_calibrate_test.go | 50 +++++++++++++++++++++ internal/agent/loop.go | 8 ++++ 3 files changed, 99 insertions(+), 4 deletions(-) create mode 100644 internal/agent/compaction_calibrate_test.go diff --git a/internal/agent/compaction.go b/internal/agent/compaction.go index 1ee45d1f..75e353cf 100644 --- a/internal/agent/compaction.go +++ b/internal/agent/compaction.go @@ -325,6 +325,43 @@ type compactionState struct { // OnText is deliberately NOT forwarded (compaction stays invisible to the user), // but its token COST must still be counted so usage reports and budgets include it. onUsage func(Usage) + + // calibrationRatio scales the raw byte/4 token estimate toward the provider's + // real prompt-token count. ApproxTextTokens over-counts code-heavy content by + // ~15-20%, which would trip compaction early (at ~60% of true capacity). It + // starts at 1.0 and converges via an EMA as each turn reports actual usage, so + // later turns compact nearer to real capacity. Zero is treated as 1.0. + calibrationRatio float64 +} + +// calibrate folds one turn's (rawEstimate, actualPromptTokens) sample into the +// running calibration ratio. A single sample is clamped to a sane band so an +// outlier (a huge cache-read turn, a provider-overhead spike) can't skew the +// estimate enough to disable or thrash compaction. +func (state *compactionState) calibrate(rawEstimate int, actualPromptTokens int) { + if !state.enabled || rawEstimate <= 0 || actualPromptTokens <= 0 { + return + } + sample := float64(actualPromptTokens) / float64(rawEstimate) + if sample < 0.5 { + sample = 0.5 + } else if sample > 2.0 { + sample = 2.0 + } + if state.calibrationRatio <= 0 { + state.calibrationRatio = 1.0 + } + const alpha = 0.3 // weight on the newest sample; smooths jitter across turns + state.calibrationRatio = state.calibrationRatio*(1-alpha) + sample*alpha +} + +// calibratedTokens applies the learned ratio to a raw estimate. Before any sample +// arrives (ratio unset) it returns the raw estimate unchanged. +func (state *compactionState) calibratedTokens(raw int) int { + if state.calibrationRatio <= 0 { + return raw + } + return int(float64(raw) * state.calibrationRatio) } func newCompactionState(options Options) *compactionState { @@ -353,7 +390,7 @@ func (state *compactionState) maybeCompact( // the messages; both the threshold check and the shrink check below use the // same term so they stay consistent. toolTokens := estimateToolDefTokens(tools) - size := estimateTokens(messages) + toolTokens + size := state.calibratedTokens(estimateTokens(messages) + toolTokens) if size <= state.threshold { return messages } @@ -370,7 +407,7 @@ func (state *compactionState) maybeCompact( // entirely and preserve recent turns verbatim. if pruned, reclaimed := pruneStaleToolOutput(messages, state.preserveLast); reclaimed > 0 { messages = pruned - size = estimateTokens(messages) + toolTokens + size = state.calibratedTokens(estimateTokens(messages) + toolTokens) if size <= state.threshold { state.lowWaterMark = size return messages @@ -386,7 +423,7 @@ func (state *compactionState) maybeCompact( // later turn) can try again; we never drop messages on failure here. return messages } - newSize := estimateTokens(compacted) + toolTokens + newSize := state.calibratedTokens(estimateTokens(compacted) + toolTokens) if newSize >= size { // Compaction did not actually shrink anything (e.g. nothing to // summarize). Leave the history untouched and don't churn next turn. @@ -444,7 +481,7 @@ func (state *compactionState) recover( // the SAME combined (messages + tool-defs) domain maybeCompact uses, so the // proactive shrink-guard compares like with like. state.reactiveAttempted = true - state.lowWaterMark = estimateTokens(result) + estimateToolDefTokens(tools) + state.lowWaterMark = state.calibratedTokens(estimateTokens(result) + estimateToolDefTokens(tools)) return result, true, nil } diff --git a/internal/agent/compaction_calibrate_test.go b/internal/agent/compaction_calibrate_test.go new file mode 100644 index 00000000..445861d2 --- /dev/null +++ b/internal/agent/compaction_calibrate_test.go @@ -0,0 +1,50 @@ +package agent + +import "testing" + +func TestCalibrationConvergesTowardActual(t *testing.T) { + state := &compactionState{enabled: true} + + // Before any sample, the estimate is unchanged. + if got := state.calibratedTokens(1000); got != 1000 { + t.Fatalf("uncalibrated tokens = %d, want 1000", got) + } + + // The estimator over-counts: our raw estimate is 1000 but the provider reports + // 850 (a 0.85 ratio). Feeding that sample repeatedly should pull the ratio + // toward 0.85, so a 1000-token raw estimate calibrates downward. + for range 20 { + state.calibrate(1000, 850) + } + got := state.calibratedTokens(1000) + if got >= 1000 || got < 820 || got > 880 { + t.Fatalf("calibrated tokens = %d, want ~850 after convergence", got) + } +} + +func TestCalibrateIgnoresDegenerateSamples(t *testing.T) { + state := &compactionState{enabled: true} + state.calibrate(0, 500) // zero estimate + state.calibrate(500, 0) // zero actual + if state.calibrationRatio != 0 { + t.Fatalf("degenerate samples must not move the ratio, got %v", state.calibrationRatio) + } + // Disabled state never calibrates. + disabled := &compactionState{enabled: false} + disabled.calibrate(1000, 850) + if disabled.calibrationRatio != 0 { + t.Fatalf("disabled compaction must not calibrate, got %v", disabled.calibrationRatio) + } +} + +func TestCalibrateClampsOutliers(t *testing.T) { + state := &compactionState{enabled: true} + // A wild outlier (10x) is clamped to 2.0, so one bad sample can't blow up the + // ratio and disable compaction. + for range 50 { + state.calibrate(1000, 10000) + } + if state.calibrationRatio > 2.0 { + t.Fatalf("ratio should be clamped at 2.0, got %v", state.calibrationRatio) + } +} diff --git a/internal/agent/loop.go b/internal/agent/loop.go index fa6ecb11..ad3a6f90 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -174,6 +174,9 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) Tools: exposed, ReasoningEffort: options.ReasoningEffort, } + // Raw (uncalibrated) estimate of this request's input, paired below with the + // provider's reported prompt tokens to calibrate future compaction triggers. + requestTokenEstimate := estimateTokens(messages) + estimateToolDefTokens(exposed) // Report the per-category context budget for this turn so a surface can // show utilization. Opt-in: a no-op when OnContext is unset. @@ -350,6 +353,11 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) } } + // Calibrate the compaction token estimator against the provider's real + // prompt-token count for the request we just sent, so later turns trigger + // compaction near true capacity instead of ~15% early on code-heavy history. + compactor.calibrate(requestTokenEstimate, collected.Usage.InputTokens) + // Carry the turn's terminal stop reason so a final answer cut off at the // output token cap (or by a content filter) is reported as truncated. A // tool-call turn normalizes to "" and clears any prior reason. From aac88855127be7e26f2a21f86b822322dc52d929 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 20:50:34 +0530 Subject: [PATCH 21/32] style: gofmt the new agent test files The CI gofmt check flagged the #22/#25 test files (struct-field and switch alignment). No behavior change. --- internal/agent/compaction_calibrate_test.go | 4 ++-- internal/agent/partition_cache_test.go | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/internal/agent/compaction_calibrate_test.go b/internal/agent/compaction_calibrate_test.go index 445861d2..55741f0d 100644 --- a/internal/agent/compaction_calibrate_test.go +++ b/internal/agent/compaction_calibrate_test.go @@ -24,8 +24,8 @@ func TestCalibrationConvergesTowardActual(t *testing.T) { func TestCalibrateIgnoresDegenerateSamples(t *testing.T) { state := &compactionState{enabled: true} - state.calibrate(0, 500) // zero estimate - state.calibrate(500, 0) // zero actual + state.calibrate(0, 500) // zero estimate + state.calibrate(500, 0) // zero actual if state.calibrationRatio != 0 { t.Fatalf("degenerate samples must not move the ratio, got %v", state.calibrationRatio) } diff --git a/internal/agent/partition_cache_test.go b/internal/agent/partition_cache_test.go index 9983afb3..078ebc4d 100644 --- a/internal/agent/partition_cache_test.go +++ b/internal/agent/partition_cache_test.go @@ -16,7 +16,7 @@ type countingSchemaTool struct { calls *int } -func (t countingSchemaTool) Name() string { return t.name } +func (t countingSchemaTool) Name() string { return t.name } func (t countingSchemaTool) Description() string { return "counts schema reads" } func (t countingSchemaTool) Parameters() tools.Schema { *t.calls++ From de333d687546303ae3809e7c993cfcd1a869c0be Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 21:06:28 +0530 Subject: [PATCH 22/32] feat(agent): turn-based plan staleness reminder MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The stale-plan nudge only fired after 10 tool calls since the last update_plan, so a plan that drifts stale across many low-tool-call turns (one call per turn) took many turns to catch — the 'plan set turn 1 stays pending forever' case. Add a turn-count complement: after stalePlanTurnThreshold (8) turns without an update AND with items still pending, fire the same one-shot reminder. Gated on pending items so a fully-completed plan is never nagged; the existing tool-call trigger is unchanged. (The roadmap's fuzzy 'tool calls unrelated to plan items' relevance match is deliberately omitted — too false-positive-prone.) --- internal/agent/guardrails.go | 30 ++++++++++--- internal/agent/plan_staleness_test.go | 64 +++++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 6 deletions(-) create mode 100644 internal/agent/plan_staleness_test.go diff --git a/internal/agent/guardrails.go b/internal/agent/guardrails.go index 63238216..926ec38a 100644 --- a/internal/agent/guardrails.go +++ b/internal/agent/guardrails.go @@ -22,6 +22,12 @@ const ( // calls have executed since the last update_plan call. staleToolCallThreshold = 10 + // stalePlanTurnThreshold injects the same one-shot reminder once this many + // turns have passed since the last update_plan while plan items are still + // pending — the turn-based complement to staleToolCallThreshold, catching a + // plan that drifts stale across many turns that each make few tool calls. + stalePlanTurnThreshold = 8 + // toolOnlyProgressReminderAt injects a one-shot progress nudge after this // many consecutive turns contain tool calls but no visible assistant text. // It does not stop the run; it tells the model to synthesize what it already @@ -358,8 +364,13 @@ type guardState struct { emptyTurns int totalToolCalls int toolCallsSincePlanUpdate int - planEverCalled bool - notCalledReminderSent bool + // turnsSincePlanUpdate counts turns (not individual tool calls) since the last + // update_plan, so a plan that goes stale across many low-tool-call turns is + // still caught — the tool-call counter alone can take many turns to trip when + // the model makes only one call per turn. + turnsSincePlanUpdate int + planEverCalled bool + notCalledReminderSent bool // staleReminderSent records whether the stale reminder has already fired for // the current stale interval. It is cleared when a plan update opens a new // interval, making the reminder one-shot per interval rather than per turn. @@ -431,11 +442,16 @@ func (state *guardState) observeTurn(collected zeroruntime.CollectedStream) (sto state.toolOnlyReminderSent = false } + // One turn has passed; the plan-update below resets this to 0 when the model + // refreshes the plan this turn. + state.turnsSincePlanUpdate++ + for _, call := range collected.ToolCalls { state.totalToolCalls++ if call.Name == planToolName { state.planEverCalled = true state.toolCallsSincePlanUpdate = 0 + state.turnsSincePlanUpdate = 0 // A fresh plan update opens a new stale interval. state.staleReminderSent = false // Record how many items remain so the completion gate knows whether @@ -489,10 +505,12 @@ func (state *guardState) progressReminder() string { // number of turns completed so far). func (state *guardState) planReminder(turn int) string { // STALE reminder takes priority: a long run without a plan update is the - // stronger signal. One-shot per stale interval. - if state.planEverCalled && - !state.staleReminderSent && - state.toolCallsSincePlanUpdate >= staleToolCallThreshold { + // stronger signal. Fires on either the tool-call streak OR a turn streak with + // pending items (so a plan drifting stale across many low-call turns is caught, + // while a fully-completed plan is left alone). One-shot per stale interval. + if state.planEverCalled && !state.staleReminderSent && + (state.toolCallsSincePlanUpdate >= staleToolCallThreshold || + (state.turnsSincePlanUpdate >= stalePlanTurnThreshold && state.planItemsPending > 0)) { state.staleReminderSent = true return planStaleReminder(state.toolCallsSincePlanUpdate) } diff --git a/internal/agent/plan_staleness_test.go b/internal/agent/plan_staleness_test.go new file mode 100644 index 00000000..4e406643 --- /dev/null +++ b/internal/agent/plan_staleness_test.go @@ -0,0 +1,64 @@ +package agent + +import ( + "strings" + "testing" + + "github.com/Gitlawb/zero/internal/zeroruntime" +) + +func planUpdateTurn(status string) zeroruntime.CollectedStream { + return zeroruntime.CollectedStream{ToolCalls: []zeroruntime.ToolCall{ + {Name: "update_plan", Arguments: `{"plan":[{"content":"step one","status":"` + status + `"}]}`}, + }} +} + +// Turn-based staleness: a plan that goes many turns without an update — while +// items are still pending — draws the stale reminder even though few tool calls +// have run (the tool-call streak alone would not have tripped). +func TestPlanReminderFiresOnTurnStaleness(t *testing.T) { + state := newGuardState() + state.observeTurn(planUpdateTurn("in_progress")) + + // Text-only turns advance the turn counter without adding tool calls. + for range stalePlanTurnThreshold { + state.observeTurn(zeroruntime.CollectedStream{Text: "still working…"}) + } + + if state.toolCallsSincePlanUpdate >= staleToolCallThreshold { + t.Fatalf("precondition: the tool-call trigger must not fire (%d calls)", state.toolCallsSincePlanUpdate) + } + if got := state.planReminder(stalePlanTurnThreshold + 2); !strings.Contains(got, planStaleReminderMarker) { + t.Fatalf("expected a turn-based stale reminder, got %q", got) + } +} + +// A fully-completed plan is NOT stale: no pending items means the turn-based +// trigger stays quiet even after many turns, so a finished plan isn't nagged. +func TestPlanReminderSkipsTurnStalenessWhenComplete(t *testing.T) { + state := newGuardState() + state.observeTurn(planUpdateTurn("completed")) + + for range stalePlanTurnThreshold + 3 { + state.observeTurn(zeroruntime.CollectedStream{Text: "wrapping up…"}) + } + + if got := state.planReminder(stalePlanTurnThreshold + 5); got != "" { + t.Fatalf("a completed plan must not draw a stale reminder, got %q", got) + } +} + +// A plan update resets the turn counter, so the stale reminder does not fire +// right after a fresh update. +func TestPlanUpdateResetsTurnStaleness(t *testing.T) { + state := newGuardState() + state.observeTurn(planUpdateTurn("in_progress")) + for range stalePlanTurnThreshold - 1 { + state.observeTurn(zeroruntime.CollectedStream{Text: "working…"}) + } + // Refresh the plan — resets the turn streak. + state.observeTurn(planUpdateTurn("in_progress")) + if got := state.planReminder(stalePlanTurnThreshold + 1); got != "" { + t.Fatalf("a just-refreshed plan must not be stale, got %q", got) + } +} From 706415fe2cca6976a46f3da04ad357f0a7c81466 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 21:27:01 +0530 Subject: [PATCH 23/32] feat(agent): preserve recent-edits summary across compaction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Compaction preserved the plan, loaded skills, tool schemas, and project instructions, but not file edits — so after the editing turns were elided the model knew it had changed a file only vaguely, and would re-read to rediscover its own footprint. Extract write_file/edit_file targets from the elided middle with a one-line note from each tool result, merge them across repeated compactions (newer note per path wins), and carry them in the preserved-state JSON (capped at 20 paths / 160 bytes per note). --- internal/agent/compaction_preserve.go | 145 +++++++++++++++++- .../agent/compaction_recent_edits_test.go | 75 +++++++++ 2 files changed, 216 insertions(+), 4 deletions(-) create mode 100644 internal/agent/compaction_recent_edits_test.go diff --git a/internal/agent/compaction_preserve.go b/internal/agent/compaction_preserve.go index b3a0da05..66654773 100644 --- a/internal/agent/compaction_preserve.go +++ b/internal/agent/compaction_preserve.go @@ -20,6 +20,16 @@ const ( toolNameUpdatePlan = "update_plan" toolNameToolSearch = "tool_search" toolNameSkill = "skill" + toolNameWriteFile = "write_file" + toolNameEditFile = "edit_file" +) + +// maxRecentEdits caps how many edited files are carried across a compaction, and +// maxEditNoteBytes caps each edit's one-line note, so a long editing run can't +// bloat the preserved-state block. +const ( + maxRecentEdits = 20 + maxEditNoteBytes = 160 ) const ( @@ -100,12 +110,105 @@ func formatPlanArguments(arguments string) string { } // skillEntry is a named preserved body. It began as loaded-skill state and is -// reused for loaded tools and project instruction blocks. +// reused for loaded tools, project instruction blocks, and recent file edits. type skillEntry struct { name string body string } +// recentEdits returns the files mutated by write_file/edit_file calls in messages +// — latest note per path, in first-seen order — as skillEntry{name: path, body: +// note}. After compaction elides the editing turns, this tells the model WHAT it +// changed in each file (from the tool's result) so it need not re-read to +// rediscover its own footprint. Capped at maxRecentEdits paths. +func recentEdits(messages []zeroruntime.Message) []skillEntry { + pathByID := map[string]string{} + order := make([]string, 0) + seen := map[string]bool{} + for _, message := range messages { + for _, call := range message.ToolCalls { + if call.Name != toolNameWriteFile && call.Name != toolNameEditFile { + continue + } + path := editPathFromArguments(call.Arguments) + if path == "" { + continue + } + if call.ID != "" { + pathByID[call.ID] = path + } + if !seen[path] { + seen[path] = true + order = append(order, path) + } + } + } + if len(order) == 0 { + return nil + } + + noteByPath := map[string]string{} + for _, message := range messages { + if message.Role != zeroruntime.MessageRoleTool || message.ToolCallID == "" { + continue + } + if path, ok := pathByID[message.ToolCallID]; ok { + noteByPath[path] = editNote(message.Content) + } + } + + // Keep the most recent maxRecentEdits paths (the tail of first-seen order). + if len(order) > maxRecentEdits { + order = order[len(order)-maxRecentEdits:] + } + entries := make([]skillEntry, 0, len(order)) + for _, path := range order { + entries = append(entries, skillEntry{name: path, body: noteByPath[path]}) + } + return entries +} + +// editPathFromArguments pulls the target file path from a write_file/edit_file +// call's JSON arguments (path/file/file_path/filename aliases). Returns "" on +// malformed arguments or no path. +func editPathFromArguments(arguments string) string { + var parsed struct { + Path string `json:"path"` + File string `json:"file"` + FilePath string `json:"file_path"` + Filename string `json:"filename"` + } + if err := json.Unmarshal([]byte(strings.TrimSpace(arguments)), &parsed); err != nil { + return "" + } + for _, candidate := range []string{parsed.Path, parsed.File, parsed.FilePath, parsed.Filename} { + if s := strings.TrimSpace(candidate); s != "" { + return s + } + } + return "" +} + +// editNote reduces an edit tool's result to a single short line (its first +// non-empty line, byte-capped on a rune boundary) for the preserved summary. +func editNote(content string) string { + for _, line := range strings.Split(content, "\n") { + s := strings.TrimSpace(line) + if s == "" { + continue + } + if len(s) > maxEditNoteBytes { + limit := maxEditNoteBytes + for limit > 0 && !utf8.RuneStart(s[limit]) { + limit-- + } + s = s[:limit] + "…" + } + return s + } + return "" +} + // loadedSkills returns the skills loaded via the skill tool in messages — the // latest body per name, in first-seen order — matching each skill tool call to // its tool result by id. @@ -310,11 +413,17 @@ func capBody(body string) string { // preservedState is the JSON shape of the carried-across-compaction block. type preservedState struct { Plan string `json:"plan,omitempty"` + RecentEdits []preservedEdit `json:"recent_edits,omitempty"` Tools []preservedTool `json:"tools,omitempty"` Skills []preservedSkill `json:"skills,omitempty"` ProjectInstructions []preservedInstruction `json:"project_instructions,omitempty"` } +type preservedEdit struct { + Path string `json:"path"` + Note string `json:"note,omitempty"` +} + type preservedTool struct { Name string `json:"name"` Body string `json:"body"` @@ -347,6 +456,11 @@ func appendPreservedState(summary string, middle []zeroruntime.Message) string { plan = priorState.Plan } + // Recent edits: merge edits preserved earlier with fresh write_file/edit_file + // results in middle (newer note per path wins), so the model keeps an + // authoritative record of what it changed across repeated compactions. + edits := capRecentEdits(mergeSkillEntries(preservedEditsToEntries(priorState.RecentEdits), recentEdits(middle))) + // Tools: preserve deferred tool_search schemas from the transcript. Fresh // loads override older carried copies by name. tools := mergeSkillEntries(preservedToolsToEntries(priorState.Tools), loadedToolSchemas(middle)) @@ -360,19 +474,42 @@ func appendPreservedState(summary string, middle []zeroruntime.Message) string { projectInstructionEntries(middle), ) - if block := formatPreservedState(plan, tools, skills, instructions); block != "" { + if block := formatPreservedState(plan, edits, tools, skills, instructions); block != "" { summary += "\n\n" + block } return summary } +// capRecentEdits bounds the preserved edit list to the most recent maxRecentEdits +// entries after a merge, so repeated compactions can't grow it without limit. +func capRecentEdits(entries []skillEntry) []skillEntry { + if len(entries) > maxRecentEdits { + return entries[len(entries)-maxRecentEdits:] + } + return entries +} + +func preservedEditsToEntries(edits []preservedEdit) []skillEntry { + entries := make([]skillEntry, 0, len(edits)) + for _, e := range edits { + if e.Path == "" { + continue + } + entries = append(entries, skillEntry{name: e.Path, body: e.Note}) + } + return entries +} + // formatPreservedState renders state as the labelled, single-line // JSON block. Returns "" when there is nothing to preserve. -func formatPreservedState(plan string, tools, skills, instructions []skillEntry) string { - if plan == "" && len(tools) == 0 && len(skills) == 0 && len(instructions) == 0 { +func formatPreservedState(plan string, edits, tools, skills, instructions []skillEntry) string { + if plan == "" && len(edits) == 0 && len(tools) == 0 && len(skills) == 0 && len(instructions) == 0 { return "" } state := preservedState{Plan: plan} + for _, e := range edits { + state.RecentEdits = append(state.RecentEdits, preservedEdit{Path: e.name, Note: e.body}) + } for _, t := range tools { state.Tools = append(state.Tools, preservedTool{Name: t.name, Body: t.body}) } diff --git a/internal/agent/compaction_recent_edits_test.go b/internal/agent/compaction_recent_edits_test.go new file mode 100644 index 00000000..42d3011f --- /dev/null +++ b/internal/agent/compaction_recent_edits_test.go @@ -0,0 +1,75 @@ +package agent + +import ( + "strings" + "testing" + + "github.com/Gitlawb/zero/internal/zeroruntime" +) + +// recentEdits extracts each mutated file's path and a one-line note from the +// matching tool result, latest note per path in first-seen order. +func TestRecentEditsExtractsPathsAndNotes(t *testing.T) { + messages := []zeroruntime.Message{ + {Role: zeroruntime.MessageRoleAssistant, ToolCalls: []zeroruntime.ToolCall{ + {ID: "e1", Name: "write_file", Arguments: `{"path":"internal/foo.go","content":"package foo"}`}, + }}, + {Role: zeroruntime.MessageRoleTool, ToolCallID: "e1", Content: "Wrote internal/foo.go (12 lines)"}, + {Role: zeroruntime.MessageRoleAssistant, ToolCalls: []zeroruntime.ToolCall{ + {ID: "e2", Name: "edit_file", Arguments: `{"path":"internal/bar.go","old_string":"a","new_string":"b"}`}, + }}, + {Role: zeroruntime.MessageRoleTool, ToolCallID: "e2", Content: "Applied edit to internal/bar.go"}, + } + + edits := recentEdits(messages) + if len(edits) != 2 { + t.Fatalf("expected 2 edited files, got %d: %#v", len(edits), edits) + } + if edits[0].name != "internal/foo.go" || !strings.Contains(edits[0].body, "12 lines") { + t.Fatalf("first edit = %#v, want foo.go with its note", edits[0]) + } + if edits[1].name != "internal/bar.go" || !strings.Contains(edits[1].body, "Applied edit") { + t.Fatalf("second edit = %#v, want bar.go with its note", edits[1]) + } +} + +// After compaction elides the editing turns, the preserved-state block still +// names the edited files and what changed, so the model needn't re-read them. +func TestCompactionPreservesRecentEdits(t *testing.T) { + messages := []zeroruntime.Message{ + {Role: zeroruntime.MessageRoleSystem, Content: "system"}, + {Role: zeroruntime.MessageRoleUser, Content: "add a flag"}, + {Role: zeroruntime.MessageRoleAssistant, Content: "editing", ToolCalls: []zeroruntime.ToolCall{ + {ID: "e1", Name: "write_file", Arguments: `{"path":"cmd/main.go","content":"..."}`}, + }}, + {Role: zeroruntime.MessageRoleTool, ToolCallID: "e1", Content: "Wrote cmd/main.go (adds --version flag)"}, + {Role: zeroruntime.MessageRoleAssistant, Content: "done"}, + {Role: zeroruntime.MessageRoleUser, Content: "continue"}, + {Role: zeroruntime.MessageRoleAssistant, Content: "continuing"}, + } + summary := compactStateConversation(t, messages) + + state := parsePreservedStateBlock(summary) + if len(state.RecentEdits) != 1 { + t.Fatalf("expected 1 preserved edit, got %#v", state.RecentEdits) + } + if state.RecentEdits[0].Path != "cmd/main.go" || !strings.Contains(state.RecentEdits[0].Note, "--version") { + t.Fatalf("preserved edit = %#v, want cmd/main.go + its note", state.RecentEdits[0]) + } +} + +// A fresh edit note for a path overrides the one carried from an earlier +// compaction (newer wins), rather than duplicating the path. +func TestRecentEditsMergeNewerWins(t *testing.T) { + prior := preservedState{RecentEdits: []preservedEdit{{Path: "a.go", Note: "old note"}}} + older := preservedEditsToEntries(prior.RecentEdits) + newer := []skillEntry{{name: "a.go", body: "new note"}, {name: "b.go", body: "added"}} + + merged := mergeSkillEntries(older, newer) + if len(merged) != 2 { + t.Fatalf("expected a.go merged (not duplicated) + b.go, got %#v", merged) + } + if merged[0].name != "a.go" || merged[0].body != "new note" { + t.Fatalf("a.go should take the newer note, got %#v", merged[0]) + } +} From 097f3433038e4e9b2bcffa04d1708fd24d33cbfd Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 21:39:43 +0530 Subject: [PATCH 24/32] feat(tui): workspace-scoped /resume + fewer event reads /resume listed sessions from every project globally, and read the full event file for each of them to test for resumable content. Filter the picker (and /resume latest) by the current workspace Cwd, checked BEFORE the per-session read so a large global history no longer pays N full file reads to build one workspace's list; also skip zero-event sessions via metadata without any read. Sessions with no recorded Cwd (older runs) and an unknown current workspace stay visible so nothing is hidden that can't be confidently placed. Explicit /resume still resolves any session. --- internal/tui/resume_scope_test.go | 25 ++++++++++++++++ internal/tui/session.go | 50 +++++++++++++++++++++++++++++-- 2 files changed, 72 insertions(+), 3 deletions(-) create mode 100644 internal/tui/resume_scope_test.go diff --git a/internal/tui/resume_scope_test.go b/internal/tui/resume_scope_test.go new file mode 100644 index 00000000..a2519007 --- /dev/null +++ b/internal/tui/resume_scope_test.go @@ -0,0 +1,25 @@ +package tui + +import "testing" + +func TestSessionMatchesWorkspace(t *testing.T) { + cases := []struct { + name string + sessionCwd string + workspace string + wantVisible bool + }{ + {"same workspace", "/home/u/proj", "/home/u/proj", true}, + {"trailing slash normalizes", "/home/u/proj/", "/home/u/proj", true}, + {"different workspace hidden", "/home/u/other", "/home/u/proj", false}, + {"session with no cwd stays visible", "", "/home/u/proj", true}, + {"unknown current workspace keeps all", "/home/u/other", "", true}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := sessionMatchesWorkspace(tc.sessionCwd, tc.workspace); got != tc.wantVisible { + t.Fatalf("sessionMatchesWorkspace(%q, %q) = %v, want %v", tc.sessionCwd, tc.workspace, got, tc.wantVisible) + } + }) + } +} diff --git a/internal/tui/session.go b/internal/tui/session.go index 76344876..b9b18745 100644 --- a/internal/tui/session.go +++ b/internal/tui/session.go @@ -4,6 +4,7 @@ import ( "encoding/json" "errors" "fmt" + "path/filepath" "strings" "time" @@ -220,9 +221,11 @@ func (m model) sessionPrompt(prompt string) string { func (m model) resolveResumeSession(args string) (*sessions.Metadata, error) { if strings.EqualFold(args, "latest") { - // Latest *resumable* conversation, so "latest" never lands on a child or - // spec sub-run. An explicit `/resume ` below still resolves any session. - latest, err := m.sessionStore.LatestResumable() + // Latest *resumable* conversation IN THIS WORKSPACE, so "latest" never lands + // on a child, a spec sub-run, or a session from another project (matching the + // workspace-scoped picker). An explicit `/resume ` below still resolves + // any session regardless of workspace. + latest, err := m.latestResumableInWorkspace() if err != nil { return nil, err } @@ -321,6 +324,18 @@ func (m model) newSessionPicker() *commandPicker { now := m.now() items := make([]pickerItem, 0, len(metas)) for _, meta := range metas { + // Workspace-scoped: hide sessions from other project directories so /resume + // lists this workspace's history, not every project's. Checked BEFORE the + // per-session event read below, so a large global history doesn't pay 50 + // full file reads to build one workspace's list. Sessions with no recorded + // Cwd (older runs) stay visible rather than vanishing. + if !sessionMatchesWorkspace(meta.Cwd, m.cwd) { + continue + } + // A zero-event session has nothing to resume — skip it without a file read. + if meta.EventCount == 0 { + continue + } // Skip empty/failed runs (no assistant output, no tool calls) — e.g. the // same prompt retried while the model wasn't responding. They have nothing // to resume and otherwise flood the list with identical rows. Still on disk. @@ -356,6 +371,35 @@ func (m model) newSessionPicker() *commandPicker { // resuming: a tool call/result, or a non-user message with real content (not the // no-output guardrail stop). Empty/failed runs return false and are hidden from // the picker (they stay on disk). Errors fail open (the session is kept). +// latestResumableInWorkspace returns the most-recently-updated resumable session +// belonging to the current workspace, or nil when none exist. ListResumable is +// ordered latest-first, so the first workspace match is the latest. +func (m model) latestResumableInWorkspace() (*sessions.Metadata, error) { + metas, err := m.sessionStore.ListResumable() + if err != nil { + return nil, err + } + for i := range metas { + if sessionMatchesWorkspace(metas[i].Cwd, m.cwd) { + return &metas[i], nil + } + } + return nil, nil +} + +// sessionMatchesWorkspace reports whether a session recorded in sessionCwd +// belongs to the current workspaceCwd. A session with no recorded Cwd (older +// runs) or an unknown current workspace is kept visible rather than filtered out, +// so the scoping never hides history it can't confidently place elsewhere. +func sessionMatchesWorkspace(sessionCwd, workspaceCwd string) bool { + sessionCwd = strings.TrimSpace(sessionCwd) + workspaceCwd = strings.TrimSpace(workspaceCwd) + if sessionCwd == "" || workspaceCwd == "" { + return true + } + return filepath.Clean(sessionCwd) == filepath.Clean(workspaceCwd) +} + func (m model) sessionHasResumableContent(sessionID string) bool { events, err := m.sessionStore.ReadEvents(sessionID) if err != nil { From 8765aa736e13c9d9bb39010174a4ea92152391ee Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 21:45:11 +0530 Subject: [PATCH 25/32] feat(tui): warn on model switch that drops staged images MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Staged images were silently dropped at submit when the active model had no vision support — the user only learned after sending. Chips already render above the composer (renderAttachmentChips); add an immediate amber warning line to the /model switch status (both same-provider and cross-provider paths) when images are staged and the newly selected model can't accept them, so the drop is surfaced at switch time. --- internal/tui/command_center.go | 9 ++++++++- internal/tui/image_attach.go | 12 ++++++++++++ internal/tui/vision_drop_test.go | 23 +++++++++++++++++++++++ 3 files changed, 43 insertions(+), 1 deletion(-) create mode 100644 internal/tui/vision_drop_test.go diff --git a/internal/tui/command_center.go b/internal/tui/command_center.go index a9f00e84..7931d7ac 100644 --- a/internal/tui/command_center.go +++ b/internal/tui/command_center.go @@ -482,6 +482,9 @@ func (m model) handleModelCommand(args string) (model, string) { lines = append(lines, target.notice) } lines = append(lines, summary) + if warn := m.visionDropWarning(); warn != "" { + lines = append(lines, warn) + } return m, strings.Join(lines, "\n") } @@ -536,7 +539,11 @@ func (m model) switchProviderModel(providerName, modelID string) (model, string, cmds = append(cmds, cmd) } } - return m, fmt.Sprintf("Model\nSwitched to %s · %s", target.Name, target.Model), tea.Batch(cmds...) + status := fmt.Sprintf("Model\nSwitched to %s · %s", target.Name, target.Model) + if warn := m.visionDropWarning(); warn != "" { + status += "\n" + warn + } + return m, status, tea.Batch(cmds...) } // profileWithCredential fills a profile's APIKey for provider construction the same diff --git a/internal/tui/image_attach.go b/internal/tui/image_attach.go index 3b855d22..b7902415 100644 --- a/internal/tui/image_attach.go +++ b/internal/tui/image_attach.go @@ -299,6 +299,18 @@ func renderImageChips(labels []string) string { // and staged documents, e.g. "[Image #1] [Image #2] [Doc #1]". Returns "" when // nothing is staged. Numbered (not named) so a long screenshot path never shows // in the composer. +// visionDropWarning returns a one-line notice when images are staged but the +// (now active) model can't accept them, so switching to a non-vision model warns +// the user immediately at switch time instead of silently dropping the images at +// submit. Empty when there is nothing staged or the model supports vision. +func (m model) visionDropWarning() string { + if len(m.pendingImages) == 0 || m.modelSupportsVisionTUI() { + return "" + } + return fmt.Sprintf("⚠ %d staged image(s) will be dropped — %s has no vision support.", + len(m.pendingImages), displayValue(m.modelName, "the active model")) +} + func renderAttachmentChips(imageLabels []string, docs []pendingDocument) string { chips := make([]string, 0, len(imageLabels)+len(docs)) for i := range imageLabels { diff --git a/internal/tui/vision_drop_test.go b/internal/tui/vision_drop_test.go new file mode 100644 index 00000000..15887d7e --- /dev/null +++ b/internal/tui/vision_drop_test.go @@ -0,0 +1,23 @@ +package tui + +import ( + "strings" + "testing" + + "github.com/Gitlawb/zero/internal/zeroruntime" +) + +func TestVisionDropWarning(t *testing.T) { + // No staged images: never warns, regardless of model. + if got := (model{}).visionDropWarning(); got != "" { + t.Fatalf("no images should give no warning, got %q", got) + } + + // Staged images + a model with no vision support (empty model name qualifies): + // warn immediately, naming the count. + withImg := model{pendingImages: make([]zeroruntime.ImageBlock, 2)} + warn := withImg.visionDropWarning() + if !strings.Contains(warn, "will be dropped") || !strings.Contains(warn, "2 staged") { + t.Fatalf("expected a drop warning naming the count, got %q", warn) + } +} From 76af0aff0d8d39ca411759b967adef3a5e280d61 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 21:52:10 +0530 Subject: [PATCH 26/32] feat(tui): inline recovery affordance on setup errors A first-run setup error rendered as a bare red 'error: ' with no pointer to the fix. Append a faint, stage-tailored recovery line (edit the endpoint/name/key above then Enter, or pick another model) that names only keys the stage's footer confirms, so the guidance is always accurate. Empty for stages whose footer already makes recovery obvious. --- internal/tui/onboarding.go | 26 +++++++++++++++++++ internal/tui/onboarding_affordance_test.go | 29 ++++++++++++++++++++++ 2 files changed, 55 insertions(+) create mode 100644 internal/tui/onboarding_affordance_test.go diff --git a/internal/tui/onboarding.go b/internal/tui/onboarding.go index f0f726cb..c201697d 100644 --- a/internal/tui/onboarding.go +++ b/internal/tui/onboarding.go @@ -1151,6 +1151,29 @@ func (m model) setupProvider() SetupProviderOption { return m.setup.providers[index] } +// setupErrorAffordance returns a faint one-line recovery hint tailored to the +// current setup stage, so a first-run error points the user at the inline fix +// instead of leaving them staring at a red line. Empty for stages whose footer +// already makes recovery obvious. Each hint names only keys the stage's footer +// confirms, so the guidance is always accurate. +func (m model) setupErrorAffordance() string { + switch m.setup.stage { + case setupStageEndpoint: + return "→ edit the endpoint above, then Enter to retry (← to go back)" + case setupStageCredentials: + if m.setupCredentialInputActive() { + return "→ re-enter the key above, then Enter to retry (← to go back)" + } + case setupStageName: + return "→ edit the name above, then Enter to continue (← to go back)" + case setupStageModel: + if !m.setup.modelLoad { + return "→ pick another model with ↑/↓, or type to search" + } + } + return "" +} + func (m model) setupView(width int) string { if width <= 0 { width = defaultStartupWidth @@ -1159,6 +1182,9 @@ func (m model) setupView(width int) string { content := m.setupStageLines(width, height) if m.setup.err != "" { content = append(content, "", zeroTheme.red.Render("error: "+m.setup.err)) + if hint := m.setupErrorAffordance(); hint != "" { + content = append(content, zeroTheme.faint.Render(hint)) + } } progress := m.setupProgressText() footer := m.setupFooter() diff --git a/internal/tui/onboarding_affordance_test.go b/internal/tui/onboarding_affordance_test.go new file mode 100644 index 00000000..90d219a0 --- /dev/null +++ b/internal/tui/onboarding_affordance_test.go @@ -0,0 +1,29 @@ +package tui + +import ( + "strings" + "testing" +) + +// Each input/choice stage's error affordance names an accurate recovery step; +// stages without a tailored hint return empty. +func TestSetupErrorAffordance(t *testing.T) { + endpoint := model{} + endpoint.setup.stage = setupStageEndpoint + if got := endpoint.setupErrorAffordance(); !strings.Contains(got, "endpoint") || !strings.Contains(got, "Enter") { + t.Fatalf("endpoint affordance = %q", got) + } + + name := model{} + name.setup.stage = setupStageName + if got := name.setupErrorAffordance(); !strings.Contains(got, "name") { + t.Fatalf("name affordance = %q", got) + } + + // A stage with no tailored hint (method chooser) returns empty. + method := model{} + method.setup.stage = setupStageMethod + if got := method.setupErrorAffordance(); got != "" { + t.Fatalf("method stage should have no affordance, got %q", got) + } +} From 6ba43ff52af64c95aebdfdaffdf2ecc988a085d4 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 22:16:17 +0530 Subject: [PATCH 27/32] fix: address CodeRabbit review (5 findings) - loop.go: recompute the compaction calibration estimate at calibration time, not request-build time, so a reactive-compaction reissue that shrinks the request feeds the EMA the sample that actually completed. - errhint: match bare HTTP status codes (401/403/429/529) only at digit boundaries, so an incidental number ('4290ms', 'id 14015') isn't mis-bucketed as auth/rate-limit. - bash: drop the stray unmatched ']' in the output-truncation marker. - tui /retry: apply the same exiting/compactInFlight guards commandPrompt has, so a retry can't race compactResultMsg's rewrite of session state. - tui /export: write the transcript 0o600 (it may contain echoed secrets), not world/group-readable 0o644. --- internal/agent/loop.go | 10 ++++---- internal/errhint/errhint.go | 32 ++++++++++++++++++++++++-- internal/errhint/errhint_test.go | 7 ++++++ internal/tools/bash.go | 2 +- internal/tui/model.go | 14 +++++++++++ internal/tui/retry_edit_test.go | 22 ++++++++++++++++++ internal/tui/transcript_export.go | 4 +++- internal/tui/transcript_export_test.go | 10 +++++++- 8 files changed, 92 insertions(+), 9 deletions(-) diff --git a/internal/agent/loop.go b/internal/agent/loop.go index ad3a6f90..fc693534 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -174,9 +174,6 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) Tools: exposed, ReasoningEffort: options.ReasoningEffort, } - // Raw (uncalibrated) estimate of this request's input, paired below with the - // provider's reported prompt tokens to calibrate future compaction triggers. - requestTokenEstimate := estimateTokens(messages) + estimateToolDefTokens(exposed) // Report the per-category context budget for this turn so a surface can // show utilization. Opt-in: a no-op when OnContext is unset. @@ -356,7 +353,12 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) // Calibrate the compaction token estimator against the provider's real // prompt-token count for the request we just sent, so later turns trigger // compaction near true capacity instead of ~15% early on code-heavy history. - compactor.calibrate(requestTokenEstimate, collected.Usage.InputTokens) + // Recompute the estimate HERE (not at request-build time): a reactive + // compaction may have replaced `messages` with a smaller set and re-sent, so + // this reflects the request that actually produced collected.Usage. The + // assistant reply is appended below, after this, so `messages` is still the + // sent request. + compactor.calibrate(estimateTokens(messages)+estimateToolDefTokens(exposed), collected.Usage.InputTokens) // Carry the turn's terminal stop reason so a final answer cut off at the // output token cap (or by a content filter) is reported as truncated. A diff --git a/internal/errhint/errhint.go b/internal/errhint/errhint.go index b8516d6d..a4562032 100644 --- a/internal/errhint/errhint.go +++ b/internal/errhint/errhint.go @@ -60,10 +60,10 @@ func Classify(err error) Category { } switch { case containsAny(m, "auth error:", "unauthorized", "api key", "api_key", "invalid_api_key", - "authentication", "permission denied", "forbidden", "401", "403"): + "authentication", "permission denied", "forbidden") || containsStatusCode(m, "401", "403"): return Auth case containsAny(m, "rate limit", "rate_limit", "too many requests", "quota", - "resource_exhausted", "overloaded", "429", "529"): + "resource_exhausted", "overloaded") || containsStatusCode(m, "429", "529"): return RateLimit case containsAny(m, "context length", "context window", "maximum context", "context_length_exceeded", "too many tokens", "prompt is too long", "reduce the length", "maximum context length"): @@ -128,3 +128,31 @@ func containsAny(haystack string, needles ...string) bool { } return false } + +// containsStatusCode reports whether haystack contains any of the given HTTP +// status codes as a standalone number — not embedded in a longer digit run like +// "completed in 4290ms" or "request id 14015" — so an incidental number can't be +// mis-bucketed as an auth/rate-limit failure. +func containsStatusCode(haystack string, codes ...string) bool { + for _, code := range codes { + for from := 0; ; { + rel := strings.Index(haystack[from:], code) + if rel < 0 { + break + } + pos := from + rel + beforeOK := pos == 0 || !isASCIIDigit(haystack[pos-1]) + end := pos + len(code) + afterOK := end >= len(haystack) || !isASCIIDigit(haystack[end]) + if beforeOK && afterOK { + return true + } + from = pos + 1 + } + } + return false +} + +func isASCIIDigit(b byte) bool { + return b >= '0' && b <= '9' +} diff --git a/internal/errhint/errhint_test.go b/internal/errhint/errhint_test.go index 8f6fd0fa..bab3c990 100644 --- a/internal/errhint/errhint_test.go +++ b/internal/errhint/errhint_test.go @@ -28,6 +28,13 @@ func TestClassify(t *testing.T) { {"deadline", "provider stream error: context deadline exceeded (Client.Timeout exceeded)", Connectivity}, {"provider-marked but no sub-signature", "provider error: something totally unexpected happened", Unknown}, + // An incidental number that merely contains a status-code digit run must NOT + // be mis-bucketed (digit-boundary check), even behind a provider marker. + {"incidental latency number", "provider error: request completed in 4290ms then failed", Unknown}, + {"incidental request id", "provider error: request id 14015 failed unexpectedly", Unknown}, + {"standalone 401 still auth", "provider request error: got 401 from upstream", Auth}, + {"standalone 429 still rate limit", "provider error: 429 returned", RateLimit}, + // Local (non-provider) failures must NOT be classified — no provider marker, // so no bogus recovery hint attaches to them. {"local fs permission denied", "open /etc/shadow: permission denied", Unknown}, diff --git a/internal/tools/bash.go b/internal/tools/bash.go index f1cde32e..50af2125 100644 --- a/internal/tools/bash.go +++ b/internal/tools/bash.go @@ -386,7 +386,7 @@ func truncateHeadTail(value string, maxBytes int) (string, int, bool) { if maxBytes <= 0 || raw <= maxBytes { return value, raw, false } - marker := fmt.Sprintf("\n[zero] output truncated: %d bytes omitted from the middle — redirect to a file and read_file a range for the full text]\n", raw-maxBytes) + marker := fmt.Sprintf("\n[zero] output truncated: %d bytes omitted from the middle — redirect to a file and read_file a range for the full text\n", raw-maxBytes) budget := maxBytes - len(marker) if budget < 0 { budget = 0 diff --git a/internal/tui/model.go b/internal/tui/model.go index 4dce86f1..52580271 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -3927,10 +3927,24 @@ func (m model) handleSubmit() (tea.Model, tea.Cmd) { m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: "$ " + cmdText}) return m, runBashEscape(m.cwd, cmdText) case commandRetry: + // /retry launches a run, so it needs the same guards a normal prompt gets: + // never start one while exiting (would strand the shutdown flush) or during + // compaction (would race compactResultMsg's wholesale rewrite of + // transcript/sessionEvents and silently drop events). + if m.exiting { + return m, nil + } if m.pending { m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: "Retry\ncannot retry while a run is in progress."}) return m, nil } + if m.compactInFlight { + m.transcript = reduceTranscript(m.transcript, transcriptAction{ + kind: actionAppendSystem, + text: "Retry\nstatus: warning\nCompaction is running. Retry once it finishes.", + }) + return m, nil + } if strings.TrimSpace(m.lastPrompt) == "" { m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: "Retry\nno previous prompt to resend."}) return m, nil diff --git a/internal/tui/retry_edit_test.go b/internal/tui/retry_edit_test.go index b4d43a4e..3ba4a8db 100644 --- a/internal/tui/retry_edit_test.go +++ b/internal/tui/retry_edit_test.go @@ -44,3 +44,25 @@ func TestRetryWithoutPriorPromptIsNoOp(t *testing.T) { t.Fatalf("/retry with no history should note there's nothing to resend, got %#v", next.transcript) } } + +// /retry must not launch a run while compaction is rewriting session state — the +// same guard a normal prompt has. +func TestRetryBlockedDuringCompaction(t *testing.T) { + m := newModel(context.Background(), Options{}) + m.lastPrompt = "do the thing" + m.compactInFlight = true + m.input.SetValue("/retry") + + updated, cmd := m.Update(testKey(tea.KeyEnter)) + next := updated.(model) + + if next.pending { + t.Fatal("/retry must not start a run during compaction") + } + if cmd != nil { + t.Fatal("/retry during compaction must not return a run command") + } + if !transcriptContains(next.transcript, "Compaction is running") { + t.Fatalf("/retry during compaction should warn, got %#v", next.transcript) + } +} diff --git a/internal/tui/transcript_export.go b/internal/tui/transcript_export.go index 4bea03e7..5c8e34be 100644 --- a/internal/tui/transcript_export.go +++ b/internal/tui/transcript_export.go @@ -74,7 +74,9 @@ func (m model) handleExportCommand(args string) string { if !filepath.IsAbs(path) && m.cwd != "" { path = filepath.Join(m.cwd, path) } - if err := os.WriteFile(path, []byte(body), 0o644); err != nil { + // 0o600: a transcript can include tool/bash output that echoed secrets, so it + // must not be world/group-readable on a shared machine. + if err := os.WriteFile(path, []byte(body), 0o600); err != nil { return "Export\nfailed to write " + path + ": " + err.Error() } return "Export\nwrote transcript to " + path diff --git a/internal/tui/transcript_export_test.go b/internal/tui/transcript_export_test.go index 960990ca..1fa419bb 100644 --- a/internal/tui/transcript_export_test.go +++ b/internal/tui/transcript_export_test.go @@ -70,13 +70,21 @@ func TestHandleExportCommandWritesFile(t *testing.T) { if !strings.Contains(msg, "wrote transcript to") { t.Fatalf("export status = %q", msg) } - data, err := os.ReadFile(filepath.Join(dir, "out.txt")) + outPath := filepath.Join(dir, "out.txt") + data, err := os.ReadFile(outPath) if err != nil { t.Fatalf("reading exported file: %v", err) } if !strings.Contains(string(data), "you: hello") || !strings.Contains(string(data), "zero: hi there") { t.Fatalf("exported content = %q", string(data)) } + // The transcript may contain secrets echoed in tool output; it must not be + // world/group-readable. + if info, err := os.Stat(outPath); err != nil { + t.Fatalf("stat exported file: %v", err) + } else if perm := info.Mode().Perm(); perm != 0o600 { + t.Fatalf("exported file mode = %o, want 0600", perm) + } // No-arg export derives a timestamped filename in cwd. if msg := m.handleExportCommand(""); !strings.Contains(msg, "zero-transcript-20260704-093000.txt") { From 94af345ad86e260be8fce56ed53a0dd712e81d03 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 22:22:18 +0530 Subject: [PATCH 28/32] test(tui): skip export file-mode assertion on Windows os.WriteFile ignores Unix permission bits on Windows and Stat reports 0666, so the 0o600 assertion failed there. The production 0o600 is honored on Linux/macOS (where it matters); guard the check with runtime.GOOS so Windows smoke passes. --- internal/tui/transcript_export_test.go | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/internal/tui/transcript_export_test.go b/internal/tui/transcript_export_test.go index 1fa419bb..efa55dce 100644 --- a/internal/tui/transcript_export_test.go +++ b/internal/tui/transcript_export_test.go @@ -3,6 +3,7 @@ package tui import ( "os" "path/filepath" + "runtime" "strings" "testing" "time" @@ -79,11 +80,14 @@ func TestHandleExportCommandWritesFile(t *testing.T) { t.Fatalf("exported content = %q", string(data)) } // The transcript may contain secrets echoed in tool output; it must not be - // world/group-readable. - if info, err := os.Stat(outPath); err != nil { - t.Fatalf("stat exported file: %v", err) - } else if perm := info.Mode().Perm(); perm != 0o600 { - t.Fatalf("exported file mode = %o, want 0600", perm) + // world/group-readable. Unix perm bits are meaningless on Windows (os.WriteFile + // ignores them and Stat reports 0666), so only assert the mode where it applies. + if runtime.GOOS != "windows" { + if info, err := os.Stat(outPath); err != nil { + t.Fatalf("stat exported file: %v", err) + } else if perm := info.Mode().Perm(); perm != 0o600 { + t.Fatalf("exported file mode = %o, want 0600", perm) + } } // No-arg export derives a timestamped filename in cwd. From f7503d5e83cf887b6e8b6eb1ac65616118199ec0 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 4 Jul 2026 23:18:55 +0530 Subject: [PATCH 29/32] fix(tui): align /resume latest filters + case-insensitive workspace match on Windows Address CodeRabbit review on #21: - latestResumableInWorkspace now applies the SAME filters as the picker (skip zero-event metadata and empty/failed runs), so /resume latest can't land on a session the picker intentionally hides. - sessionMatchesWorkspace compares paths case-insensitively on Windows (its filesystem is), so the same workspace spelled with different casing no longer hides resumable sessions; other platforms stay case-sensitive. --- internal/tui/resume_scope_test.go | 8 +++++++- internal/tui/session.go | 28 +++++++++++++++++++++++----- 2 files changed, 30 insertions(+), 6 deletions(-) diff --git a/internal/tui/resume_scope_test.go b/internal/tui/resume_scope_test.go index a2519007..5f3d7847 100644 --- a/internal/tui/resume_scope_test.go +++ b/internal/tui/resume_scope_test.go @@ -1,6 +1,9 @@ package tui -import "testing" +import ( + "runtime" + "testing" +) func TestSessionMatchesWorkspace(t *testing.T) { cases := []struct { @@ -14,6 +17,9 @@ func TestSessionMatchesWorkspace(t *testing.T) { {"different workspace hidden", "/home/u/other", "/home/u/proj", false}, {"session with no cwd stays visible", "", "/home/u/proj", true}, {"unknown current workspace keeps all", "/home/u/other", "", true}, + // Casing: matched case-insensitively on Windows (its FS is), case-sensitively + // elsewhere. This exercises the platform-specific branch on each OS's CI. + {"case differs only", "/home/U/Proj", "/home/u/proj", runtime.GOOS == "windows"}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { diff --git a/internal/tui/session.go b/internal/tui/session.go index b9b18745..7fa1f541 100644 --- a/internal/tui/session.go +++ b/internal/tui/session.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "path/filepath" + "runtime" "strings" "time" @@ -373,16 +374,26 @@ func (m model) newSessionPicker() *commandPicker { // the picker (they stay on disk). Errors fail open (the session is kept). // latestResumableInWorkspace returns the most-recently-updated resumable session // belonging to the current workspace, or nil when none exist. ListResumable is -// ordered latest-first, so the first workspace match is the latest. +// ordered latest-first, so the first qualifying match is the latest. It applies +// the SAME filters as newSessionPicker — workspace membership, non-empty +// metadata, and real resumable content — so `/resume latest` never lands on a +// zero-event or empty/failed run the picker would have hidden. func (m model) latestResumableInWorkspace() (*sessions.Metadata, error) { metas, err := m.sessionStore.ListResumable() if err != nil { return nil, err } for i := range metas { - if sessionMatchesWorkspace(metas[i].Cwd, m.cwd) { - return &metas[i], nil + if !sessionMatchesWorkspace(metas[i].Cwd, m.cwd) { + continue + } + if metas[i].EventCount == 0 { + continue } + if !m.sessionHasResumableContent(metas[i].SessionID) { + continue + } + return &metas[i], nil } return nil, nil } @@ -390,14 +401,21 @@ func (m model) latestResumableInWorkspace() (*sessions.Metadata, error) { // sessionMatchesWorkspace reports whether a session recorded in sessionCwd // belongs to the current workspaceCwd. A session with no recorded Cwd (older // runs) or an unknown current workspace is kept visible rather than filtered out, -// so the scoping never hides history it can't confidently place elsewhere. +// so the scoping never hides history it can't confidently place elsewhere. On +// Windows the comparison is case-insensitive, since the filesystem is and the +// same workspace can be spelled with different casing (C:\Proj vs c:\proj). func sessionMatchesWorkspace(sessionCwd, workspaceCwd string) bool { sessionCwd = strings.TrimSpace(sessionCwd) workspaceCwd = strings.TrimSpace(workspaceCwd) if sessionCwd == "" || workspaceCwd == "" { return true } - return filepath.Clean(sessionCwd) == filepath.Clean(workspaceCwd) + a := filepath.Clean(sessionCwd) + b := filepath.Clean(workspaceCwd) + if runtime.GOOS == "windows" { + return strings.EqualFold(a, b) + } + return a == b } func (m model) sessionHasResumableContent(sessionID string) bool { From a78414f84e8b89a6842f42c53f8f6b4d7d9d4ee3 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sun, 5 Jul 2026 10:26:13 +0530 Subject: [PATCH 30/32] fix(tui): /retry resends the last prompt's image/PDF attachments /retry relaunched only lastPrompt's text, but launchPrompt clears the pending image/document queues once a turn is sent. A vision- or PDF-backed prompt that failed would silently retry as text-only and answer a different task. Snapshot the consumed attachments at launch and re-stage them on /retry so launchPrompt rebuilds an identical request (document preamble, images, and the submit-time vision re-check). startNewSession drops the snapshot so a post-/new /retry cannot leak the previous session's attachments. Also correct the OAuthResolver factory comment, which omitted the Gemini provider it is now passed to. --- internal/providers/factory.go | 2 +- internal/tui/image_attach_test.go | 73 +++++++++++++++++++++++++++++++ internal/tui/model.go | 24 +++++++++- internal/tui/new_session_test.go | 9 ++++ internal/tui/session.go | 6 +++ 5 files changed, 112 insertions(+), 2 deletions(-) diff --git a/internal/providers/factory.go b/internal/providers/factory.go index d432ca2f..57f0997c 100644 --- a/internal/providers/factory.go +++ b/internal/providers/factory.go @@ -24,7 +24,7 @@ type Options struct { ModelRegistry *modelregistry.Registry // OAuthResolver, when set, lets the provider authenticate model calls with an // OAuth bearer token (preferred over the API key). nil => API-key auth only. - // Applied to the OpenAI and Anthropic providers. + // Applied to the OpenAI, Anthropic, and Google (Gemini) providers. OAuthResolver providerio.TokenResolver // OAuthLoginKey is the credential-store key OAuthResolver bound to (empty when // there is no OAuth login). It is the SAME selection the bearer resolver made, diff --git a/internal/tui/image_attach_test.go b/internal/tui/image_attach_test.go index b25f3f40..d4754028 100644 --- a/internal/tui/image_attach_test.go +++ b/internal/tui/image_attach_test.go @@ -483,3 +483,76 @@ func TestRenderAttachmentChips(t *testing.T) { t.Fatalf("chip row %q should not show file names", got) } } + +// TestRetryResendsAttachments guards the /retry contract: it must resend the exact +// request the last prompt carried — including its staged image and PDF document — +// not a degraded text-only prompt. launchPrompt clears the pending queues once a +// turn is sent, so without re-staging the remembered snapshot a vision/document- +// backed turn that failed would silently retry as a different task. +func TestRetryResendsAttachments(t *testing.T) { + root := t.TempDir() + + captured := make(chan []zeroruntime.ImageBlock, 1) + provider := &fakeProvider{events: []zeroruntime.StreamEvent{ + {Type: zeroruntime.StreamEventText, Content: "ok"}, + {Type: zeroruntime.StreamEventDone}, + }} + m := newModel(context.Background(), Options{ + Cwd: root, + ProviderName: "openai", + ModelName: "gpt-4.1", + Provider: provider, + Registry: tools.NewRegistry(), + SessionStore: testSessionStore(t), + }) + m.agentOptions.OnText = func(string) {} + m.captureRunImages = func(imgs []zeroruntime.ImageBlock) { captured <- imgs } + + // State left after a prior vision+PDF prompt was submitted: the pending queues + // are cleared, but the remembered snapshot survives so /retry can reproduce it. + m.lastPrompt = "describe both" + m.lastImages = []zeroruntime.ImageBlock{{MediaType: "image/png", Data: []byte{0x89, 'P', 'N', 'G'}}} + m.lastImageLabels = []string{"photo.png"} + m.lastDocuments = []pendingDocument{{label: "spec.pdf", text: "Top secret design notes"}} + + m.input.SetValue("/retry") + updated, cmd := m.handleSubmit() + next := updated.(model) + if cmd == nil { + t.Fatal("/retry with a remembered prompt should launch a run") + } + // The retried turn re-consumes and then clears the queues, exactly like a fresh + // submit; the snapshot must survive so a second /retry stays reproducible. + if len(next.pendingImages) != 0 || len(next.pendingImageLabels) != 0 || len(next.pendingDocuments) != 0 { + t.Fatalf("retry must clear pending queues after resend, got imgs=%d labels=%d docs=%d", + len(next.pendingImages), len(next.pendingImageLabels), len(next.pendingDocuments)) + } + if len(next.lastImages) != 1 || len(next.lastDocuments) != 1 { + t.Fatalf("retry must keep the snapshot for a subsequent retry, got imgs=%d docs=%d", + len(next.lastImages), len(next.lastDocuments)) + } + + updated, _ = next.Update(execCmd(cmd)) + _ = updated.(model) + + select { + case imgs := <-captured: + if len(imgs) != 1 || imgs[0].MediaType != "image/png" { + t.Fatalf("retried run must resend the remembered image, got %#v", imgs) + } + default: + t.Fatal("expected the retried run to receive the remembered image") + } + + if len(provider.requests) != 1 { + t.Fatalf("expected one provider request, got %d", len(provider.requests)) + } + msgs := provider.requests[0].Messages + last := msgs[len(msgs)-1] + if !strings.Contains(last.Content, "Top secret design notes") { + t.Fatalf("retried prompt should re-prepend the document text, got:\n%s", last.Content) + } + if !strings.Contains(last.Content, "describe both") { + t.Fatalf("retried prompt should include the remembered user text, got:\n%s", last.Content) + } +} diff --git a/internal/tui/model.go b/internal/tui/model.go index 52580271..d5f4f6d3 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -269,6 +269,15 @@ type model struct { // lastPrompt is the verbatim text of the most recent submitted prompt, so // /retry can resend it and /edit can recall it into the composer. lastPrompt string + // lastImages/lastImageLabels/lastDocuments remember the attachments consumed + // by the most recent submitted prompt. launchPrompt clears the pending queues + // once a turn is sent, so /retry re-stages these to reproduce the exact same + // request — otherwise a vision/PDF-backed prompt would silently retry as + // text-only and answer a different task. They share the underlying image bytes + // with the sent turn (never mutated in place), so no deep copy is needed. + lastImages []zeroruntime.ImageBlock + lastImageLabels []string + lastDocuments []pendingDocument // historyIdx == len(inputHistory) means "not navigating"; historyDraft // preserves whatever was typed before recall started. inputHistory []string @@ -3949,6 +3958,13 @@ func (m model) handleSubmit() (tea.Model, tea.Cmd) { m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: "Retry\nno previous prompt to resend."}) return m, nil } + // Re-stage the attachments the last prompt carried so launchPrompt rebuilds + // an identical request (document preamble + images + vision re-check). Without + // this the queues are empty and /retry would resend a text-only prompt, + // silently dropping the image/PDF context and answering a different task. + m.pendingImages = m.lastImages + m.pendingImageLabels = m.lastImageLabels + m.pendingDocuments = m.lastDocuments return m.launchPrompt(m.lastPrompt) case commandEdit: if strings.TrimSpace(m.lastPrompt) == "" { @@ -3982,8 +3998,14 @@ func (m model) handleSubmit() (tea.Model, tea.Cmd) { // stays identical to immediate submissions. func (m model) launchPrompt(prompt string) (model, tea.Cmd) { // Remember the verbatim prompt (before specialist/document expansion) so /retry - // and /edit can act on exactly what the user submitted. + // and /edit can act on exactly what the user submitted. Snapshot the staged + // attachments too: launchPrompt clears the pending queues below, so /retry + // re-stages these to resend an identical vision/PDF-backed request rather than + // a degraded text-only one. m.lastPrompt = prompt + m.lastImages = m.pendingImages + m.lastImageLabels = m.pendingImageLabels + m.lastDocuments = m.pendingDocuments m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendUser, text: prompt}) if m.provider == nil { m.transcript = reduceTranscript(m.transcript, transcriptAction{ diff --git a/internal/tui/new_session_test.go b/internal/tui/new_session_test.go index 468be855..f225951f 100644 --- a/internal/tui/new_session_test.go +++ b/internal/tui/new_session_test.go @@ -20,6 +20,10 @@ func TestStartNewSessionResetsState(t *testing.T) { m.pendingImageLabels = []string{"pic.png"} m.pendingDocuments = []pendingDocument{{label: "doc.pdf"}} m.queuedMessage = "queued" + // The /retry attachment snapshot is prior-session state too and must not survive. + m.lastImages = make([]zeroruntime.ImageBlock, 1) + m.lastImageLabels = []string{"pic.png"} + m.lastDocuments = []pendingDocument{{label: "doc.pdf"}} next := m.startNewSession() @@ -41,6 +45,11 @@ func TestStartNewSessionResetsState(t *testing.T) { t.Fatalf("startNewSession must clear staged input, got images=%d labels=%d docs=%d queued=%q", len(next.pendingImages), len(next.pendingImageLabels), len(next.pendingDocuments), next.queuedMessage) } + // The /retry snapshot must not leak the previous session's attachments. + if len(next.lastImages) != 0 || len(next.lastImageLabels) != 0 || len(next.lastDocuments) != 0 { + t.Fatalf("startNewSession must clear the retry snapshot, got images=%d labels=%d docs=%d", + len(next.lastImages), len(next.lastImageLabels), len(next.lastDocuments)) + } } func TestNewCommandStartsFreshSession(t *testing.T) { diff --git a/internal/tui/session.go b/internal/tui/session.go index 7fa1f541..9e3ed43c 100644 --- a/internal/tui/session.go +++ b/internal/tui/session.go @@ -82,6 +82,12 @@ func (m model) startNewSession() model { m.pendingImageLabels = nil m.pendingDocuments = nil m.queuedMessage = "" + // The remembered /retry attachment snapshot belongs to the previous session + // too — dropping it keeps a post-/new /retry from re-staging old images or + // documents. lastPrompt is composer history (like inputHistory) and stays. + m.lastImages = nil + m.lastImageLabels = nil + m.lastDocuments = nil note := "Started a new session." if previousID != "" { From ab82c7b72f5315148d30c9275a0d817af9fd0298 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sun, 5 Jul 2026 11:43:14 +0530 Subject: [PATCH 31/32] fix: address jatmn review (reconnect 5xx, /edit attachments, edit-cap, coalescer test) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - reconnect: exclude HTTP 5xx (500/502/503/504) before the transport substring match so "504 Gateway Timeout" no longer slips through on the generic "timeout" needle. A gateway timeout is non-idempotent — the completion POST may already have reached the model — so replaying the connect risked duplicate billable work. Reuses a digit-boundary status matcher exported from errhint (HasStatusCode). Test asserts 504/500 stay non-reconnectable. - tui /edit: re-stage the remembered image/PDF snapshot alongside the recalled prompt text, mirroring the /retry fix — editing a vision- or document-backed prompt no longer silently resends a text-only version. - compaction recent-edits: order edits by LAST edit, not first, and move a re-touched path to the newest position on merge, so the maxRecentEdits tail cap keeps the file the model most recently edited instead of dropping it. Adds mergeRecentEdits (edit-specific) and a capped-list regression. - coalescer test: inject the frame timer via afterFunc so TestCoalescerBatchesDeltas proves buffering deterministically instead of racing the real 16ms wall-clock timer. --- internal/agent/compaction_preserve.go | 75 ++++++++++++++++--- .../agent/compaction_recent_edits_test.go | 62 +++++++++++++-- internal/agent/reconnect.go | 21 ++++-- internal/agent/reconnect_test.go | 8 +- internal/errhint/errhint.go | 9 +++ internal/tui/coalesce.go | 22 +++++- internal/tui/coalesce_test.go | 17 ++++- internal/tui/model.go | 8 ++ internal/tui/retry_edit_test.go | 28 +++++++ 9 files changed, 217 insertions(+), 33 deletions(-) diff --git a/internal/agent/compaction_preserve.go b/internal/agent/compaction_preserve.go index 66654773..6423a5b2 100644 --- a/internal/agent/compaction_preserve.go +++ b/internal/agent/compaction_preserve.go @@ -117,14 +117,17 @@ type skillEntry struct { } // recentEdits returns the files mutated by write_file/edit_file calls in messages -// — latest note per path, in first-seen order — as skillEntry{name: path, body: +// — latest note per path, in last-seen order — as skillEntry{name: path, body: // note}. After compaction elides the editing turns, this tells the model WHAT it // changed in each file (from the tool's result) so it need not re-read to // rediscover its own footprint. Capped at maxRecentEdits paths. +// +// Ordering is by LAST edit, not first: re-editing a file moves it to the newest +// position so the tail cap below keeps the file the model most recently touched +// rather than an earlier, now-stale entry. func recentEdits(messages []zeroruntime.Message) []skillEntry { pathByID := map[string]string{} - order := make([]string, 0) - seen := map[string]bool{} + sequence := make([]string, 0) for _, message := range messages { for _, call := range message.ToolCalls { if call.Name != toolNameWriteFile && call.Name != toolNameEditFile { @@ -137,12 +140,10 @@ func recentEdits(messages []zeroruntime.Message) []skillEntry { if call.ID != "" { pathByID[call.ID] = path } - if !seen[path] { - seen[path] = true - order = append(order, path) - } + sequence = append(sequence, path) } } + order := lastSeenOrder(sequence) if len(order) == 0 { return nil } @@ -157,7 +158,7 @@ func recentEdits(messages []zeroruntime.Message) []skillEntry { } } - // Keep the most recent maxRecentEdits paths (the tail of first-seen order). + // Keep the most recent maxRecentEdits paths (the tail of last-seen order). if len(order) > maxRecentEdits { order = order[len(order)-maxRecentEdits:] } @@ -168,6 +169,23 @@ func recentEdits(messages []zeroruntime.Message) []skillEntry { return entries } +// lastSeenOrder dedupes paths keeping each at its LAST occurrence, preserving +// chronological order otherwise. A re-edited path therefore lands at the newest +// position rather than staying pinned to where it first appeared. +func lastSeenOrder(paths []string) []string { + lastIdx := make(map[string]int, len(paths)) + for i, p := range paths { + lastIdx[p] = i + } + order := make([]string, 0, len(lastIdx)) + for i, p := range paths { + if lastIdx[p] == i { + order = append(order, p) + } + } + return order +} + // editPathFromArguments pulls the target file path from a write_file/edit_file // call's JSON arguments (path/file/file_path/filename aliases). Returns "" on // malformed arguments or no path. @@ -457,9 +475,10 @@ func appendPreservedState(summary string, middle []zeroruntime.Message) string { } // Recent edits: merge edits preserved earlier with fresh write_file/edit_file - // results in middle (newer note per path wins), so the model keeps an - // authoritative record of what it changed across repeated compactions. - edits := capRecentEdits(mergeSkillEntries(preservedEditsToEntries(priorState.RecentEdits), recentEdits(middle))) + // results in middle (newer note per path wins AND moves to the newest + // position), so the tail cap keeps the files the model most recently touched + // across repeated compactions rather than dropping a just-re-edited file. + edits := capRecentEdits(mergeRecentEdits(preservedEditsToEntries(priorState.RecentEdits), recentEdits(middle))) // Tools: preserve deferred tool_search schemas from the transcript. Fresh // loads override older carried copies by name. @@ -480,6 +499,40 @@ func appendPreservedState(summary string, middle []zeroruntime.Message) string { return summary } +// mergeRecentEdits overlays fresh edits onto edits preserved by an earlier +// compaction. Unlike mergeSkillEntries (which keeps refreshed entries in their +// original slot), a path touched again by a fresh edit MOVES to the newest +// position and takes the fresh note, so capRecentEdits — which keeps only the +// most-recent tail — never drops a file the model just re-edited. Paths present +// only in the older set keep their relative order, ahead of the fresh ones. +func mergeRecentEdits(older, newer []skillEntry) []skillEntry { + if len(newer) == 0 { + return older + } + freshBody := make(map[string]string, len(newer)) + freshOrder := make([]string, 0, len(newer)) + for _, e := range newer { + if _, ok := freshBody[e.name]; !ok { + freshOrder = append(freshOrder, e.name) + } + freshBody[e.name] = e.body + } + merged := make([]skillEntry, 0, len(older)+len(freshOrder)) + // Older entries that were NOT re-edited keep their position, oldest first. + for _, e := range older { + if _, refreshed := freshBody[e.name]; refreshed { + continue // re-appended at its newest position below + } + merged = append(merged, e) + } + // Fresh edits follow, in last-seen order, each carrying the fresh note — so + // the most recently touched files sit at the tail the cap preserves. + for _, name := range freshOrder { + merged = append(merged, skillEntry{name: name, body: freshBody[name]}) + } + return merged +} + // capRecentEdits bounds the preserved edit list to the most recent maxRecentEdits // entries after a merge, so repeated compactions can't grow it without limit. func capRecentEdits(entries []skillEntry) []skillEntry { diff --git a/internal/agent/compaction_recent_edits_test.go b/internal/agent/compaction_recent_edits_test.go index 42d3011f..c8fe2b73 100644 --- a/internal/agent/compaction_recent_edits_test.go +++ b/internal/agent/compaction_recent_edits_test.go @@ -1,6 +1,7 @@ package agent import ( + "fmt" "strings" "testing" @@ -8,7 +9,7 @@ import ( ) // recentEdits extracts each mutated file's path and a one-line note from the -// matching tool result, latest note per path in first-seen order. +// matching tool result, latest note per path in last-seen order. func TestRecentEditsExtractsPathsAndNotes(t *testing.T) { messages := []zeroruntime.Message{ {Role: zeroruntime.MessageRoleAssistant, ToolCalls: []zeroruntime.ToolCall{ @@ -59,17 +60,62 @@ func TestCompactionPreservesRecentEdits(t *testing.T) { } // A fresh edit note for a path overrides the one carried from an earlier -// compaction (newer wins), rather than duplicating the path. +// compaction (newer wins) and moves the path to the newest position rather than +// duplicating it or leaving it pinned to its old slot. func TestRecentEditsMergeNewerWins(t *testing.T) { - prior := preservedState{RecentEdits: []preservedEdit{{Path: "a.go", Note: "old note"}}} + prior := preservedState{RecentEdits: []preservedEdit{{Path: "a.go", Note: "old note"}, {Path: "z.go", Note: "z"}}} older := preservedEditsToEntries(prior.RecentEdits) newer := []skillEntry{{name: "a.go", body: "new note"}, {name: "b.go", body: "added"}} - merged := mergeSkillEntries(older, newer) - if len(merged) != 2 { - t.Fatalf("expected a.go merged (not duplicated) + b.go, got %#v", merged) + merged := mergeRecentEdits(older, newer) + // z.go (untouched) keeps its lead; a.go is re-edited so it takes the new note + // and moves behind z.go, ahead of the brand-new b.go: [z.go, a.go, b.go]. + if len(merged) != 3 { + t.Fatalf("expected z.go + refreshed a.go + b.go, got %#v", merged) } - if merged[0].name != "a.go" || merged[0].body != "new note" { - t.Fatalf("a.go should take the newer note, got %#v", merged[0]) + if merged[0].name != "z.go" { + t.Fatalf("untouched z.go should stay first, got %#v", merged) + } + if merged[1].name != "a.go" || merged[1].body != "new note" { + t.Fatalf("re-edited a.go should move to the newest edits with the new note, got %#v", merged[1]) + } + if merged[2].name != "b.go" { + t.Fatalf("brand-new b.go should be last, got %#v", merged[2]) + } +} + +// Regression: once more than maxRecentEdits distinct files are tracked, re-editing +// an early file must move it into the capped tail rather than leaving it to be +// dropped — otherwise the next compaction omits exactly the file the model most +// recently touched. +func TestRecentEditsCapKeepsReeditedFile(t *testing.T) { + // Earlier compaction preserved exactly maxRecentEdits files, f0 the oldest. + older := make([]skillEntry, 0, maxRecentEdits) + for i := 0; i < maxRecentEdits; i++ { + older = append(older, skillEntry{name: fmt.Sprintf("f%d.go", i), body: "old"}) + } + // Fresh window re-edits the OLDEST file (f0) and adds a brand-new file, so the + // merged list exceeds the cap and something must be dropped. + newer := []skillEntry{{name: "f0.go", body: "re-edited"}, {name: "fresh.go", body: "new"}} + + merged := capRecentEdits(mergeRecentEdits(older, newer)) + if len(merged) != maxRecentEdits { + t.Fatalf("cap should keep %d entries, got %d", maxRecentEdits, len(merged)) + } + + var f0 *skillEntry + for i := range merged { + if merged[i].name == "f0.go" { + f0 = &merged[i] + } + if merged[i].name == "f1.go" { + t.Fatalf("f1.go (now the least-recently touched) should have been dropped, got %#v", merged) + } + } + if f0 == nil { + t.Fatalf("re-edited f0.go must survive the cap, got %#v", merged) + } + if f0.body != "re-edited" { + t.Fatalf("surviving f0.go should carry the fresh note, got %q", f0.body) } } diff --git a/internal/agent/reconnect.go b/internal/agent/reconnect.go index 8e309c1b..54b339ec 100644 --- a/internal/agent/reconnect.go +++ b/internal/agent/reconnect.go @@ -7,6 +7,7 @@ import ( "strings" "time" + "github.com/Gitlawb/zero/internal/errhint" "github.com/Gitlawb/zero/internal/zeroruntime" ) @@ -107,13 +108,19 @@ func shouldReconnect(ctx context.Context, err error) bool { if isContextLimitError(msg) || isImageRejectionError(err) { return false } - // Transport-level disconnects only. HTTP 5xx statuses are deliberately NOT - // matched here: 503 already exhausted providerio.SendWithRetry (retrying it - // again is a redundant double-retry), and 502/500/504 are non-idempotent by - // providerio's rule (the completion POST may have been processed), so replaying - // the connect risks duplicate billable work. A genuine transport failure (EOF, - // reset, refused, timeout) means no response was received, which is safe to - // reconnect. + // HTTP 5xx statuses are non-reconnectable and must be excluded BEFORE the + // transport substring match below — otherwise "504 Gateway Timeout" would slip + // through on the generic "timeout" needle. 503 already exhausted + // providerio.SendWithRetry (retrying is a redundant double-retry); 500/502/504 + // are non-idempotent by providerio's rule — the completion POST may already + // have reached the model before the upstream/gateway gave up, so replaying the + // connect risks duplicate billable work. Digit-boundary matched so an + // incidental "504" inside a latency/id number is not mistaken for a status. + if errhint.HasStatusCode(msg, "500", "502", "503", "504") { + return false + } + // Transport-level disconnects only. A genuine transport failure (EOF, reset, + // refused, timeout) means no response was received, which is safe to reconnect. for _, needle := range []string{ "eof", "connection reset", diff --git a/internal/agent/reconnect_test.go b/internal/agent/reconnect_test.go index a7c3d38f..84bd3a3f 100644 --- a/internal/agent/reconnect_test.go +++ b/internal/agent/reconnect_test.go @@ -102,10 +102,12 @@ func TestShouldReconnectClassification(t *testing.T) { "context length exceeded", "invalid api key", "model not found", "400 bad request: unsupported parameter", // HTTP 5xx statuses are handled by providerio.SendWithRetry (503) or are - // non-idempotent (502); the reconnect path must not double-retry them. (A - // "504 Gateway Timeout" still reconnects via the generic "timeout" transport - // signal — an accepted overlap, since a gateway timeout saw no response.) + // non-idempotent (500/502/504); the reconnect path must not double-retry + // them. "504 Gateway Timeout" must NOT slip through on the generic "timeout" + // transport signal — a gateway timeout is non-idempotent, since the upstream + // may have already processed the completion POST before giving up. "503 Service Unavailable", "502 Bad Gateway", + "504 Gateway Timeout", "500 Internal Server Error", } for _, m := range notDisconnects { if shouldReconnect(ctx, errors.New(m)) { diff --git a/internal/errhint/errhint.go b/internal/errhint/errhint.go index a4562032..d52f5c66 100644 --- a/internal/errhint/errhint.go +++ b/internal/errhint/errhint.go @@ -134,6 +134,15 @@ func containsAny(haystack string, needles ...string) bool { // "completed in 4290ms" or "request id 14015" — so an incidental number can't be // mis-bucketed as an auth/rate-limit failure. func containsStatusCode(haystack string, codes ...string) bool { + return HasStatusCode(haystack, codes...) +} + +// HasStatusCode reports whether haystack contains any of the given HTTP status +// codes as a standalone number — not embedded in a longer digit run like +// "completed in 4290ms" or "request id 14015". Exported so other packages (e.g. +// the reconnect classifier) can gate on a status code without re-implementing +// the digit-boundary check. +func HasStatusCode(haystack string, codes ...string) bool { for _, code := range codes { for from := 0; ; { rel := strings.Index(haystack[from:], code) diff --git a/internal/tui/coalesce.go b/internal/tui/coalesce.go index 628db1be..26b3c56e 100644 --- a/internal/tui/coalesce.go +++ b/internal/tui/coalesce.go @@ -28,15 +28,31 @@ const streamCoalesceInterval = 16 * time.Millisecond // drains and forwards atomically, and the other caller blocks until it is done. type textCoalescer struct { forward func(tea.Msg) // downstream sink (external sink + program.Send) + // afterFunc schedules fn to run after one frame interval and returns a + // stoppable timer. Defaults to a real time.AfterFunc(streamCoalesceInterval, …); + // tests swap in a controllable timer so flush timing is deterministic instead of + // racing the 16ms wall clock. + afterFunc func(fn func()) coalesceTimer mu sync.Mutex buf []byte runID int - timer *time.Timer + timer coalesceTimer +} + +// coalesceTimer is the subset of *time.Timer the coalescer needs. Abstracted +// behind afterFunc so a test can substitute a timer it controls. +type coalesceTimer interface { + Stop() bool } func newTextCoalescer(forward func(tea.Msg)) *textCoalescer { - return &textCoalescer{forward: forward} + return &textCoalescer{ + forward: forward, + afterFunc: func(fn func()) coalesceTimer { + return time.AfterFunc(streamCoalesceInterval, fn) + }, + } } // send is the coalescing entry point installed as the RuntimeMessageSink. @@ -62,7 +78,7 @@ func (c *textCoalescer) send(msg tea.Msg) { c.runID = text.runID c.buf = append(c.buf, text.delta...) if c.timer == nil { - c.timer = time.AfterFunc(streamCoalesceInterval, c.flush) + c.timer = c.afterFunc(c.flush) } } diff --git a/internal/tui/coalesce_test.go b/internal/tui/coalesce_test.go index 34988595..9157b3e9 100644 --- a/internal/tui/coalesce_test.go +++ b/internal/tui/coalesce_test.go @@ -28,17 +28,32 @@ func (r *recorder) snapshot() []tea.Msg { return out } +// manualTimer is a coalesceTimer that never fires on its own, so a test can prove +// buffering deterministically without racing the real 16ms frame timer: the flush +// happens only when the test calls c.flush() explicitly. +type manualTimer struct{ stopped bool } + +func (m *manualTimer) Stop() bool { + prev := m.stopped + m.stopped = true + return !prev +} + // Rapid deltas within one frame collapse into a single agentTextMsg carrying the // concatenated text. func TestCoalescerBatchesDeltas(t *testing.T) { rec := &recorder{} c := newTextCoalescer(rec.forward) + // Swap in a timer that never auto-fires so the "still buffered" assertion below + // cannot flake if the scheduler pauses this goroutine past the frame interval. + c.afterFunc = func(func()) coalesceTimer { return &manualTimer{} } c.send(agentTextMsg{runID: 1, delta: "Hel"}) c.send(agentTextMsg{runID: 1, delta: "lo, "}) c.send(agentTextMsg{runID: 1, delta: "world"}) - // Nothing forwarded yet — still buffered within the frame. + // Nothing forwarded yet — still buffered within the frame (the timer we injected + // never fires; only the explicit flush below delivers the text). if got := rec.snapshot(); len(got) != 0 { t.Fatalf("deltas should buffer, forwarded %d early: %#v", len(got), got) } diff --git a/internal/tui/model.go b/internal/tui/model.go index d5f4f6d3..1f43c8db 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -3971,6 +3971,14 @@ func (m model) handleSubmit() (tea.Model, tea.Cmd) { m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: "Edit\nno previous prompt to recall."}) return m, nil } + // Re-stage the remembered attachments alongside the recalled text so an + // edited resend carries the same image/PDF context — the reappearing chip + // row is the visible confirmation. Without this, editing a vision- or + // document-backed prompt would silently submit a text-only version and + // answer a different task (the same gap /retry guards against). + m.pendingImages = m.lastImages + m.pendingImageLabels = m.lastImageLabels + m.pendingDocuments = m.lastDocuments m.input.SetValue(m.lastPrompt) return m, nil case commandCopy: diff --git a/internal/tui/retry_edit_test.go b/internal/tui/retry_edit_test.go index 3ba4a8db..bd14f3c3 100644 --- a/internal/tui/retry_edit_test.go +++ b/internal/tui/retry_edit_test.go @@ -5,6 +5,8 @@ import ( "testing" tea "charm.land/bubbletea/v2" + + "github.com/Gitlawb/zero/internal/zeroruntime" ) func TestCommandsRegistered(t *testing.T) { @@ -31,6 +33,32 @@ func TestEditRecallsLastPrompt(t *testing.T) { } } +// /edit must re-stage the remembered attachments alongside the recalled text, so an +// edited resend of a vision/PDF-backed prompt carries the same image/document +// context instead of silently sending a text-only version. +func TestEditRestagesAttachments(t *testing.T) { + m := newModel(context.Background(), Options{ModelName: "gpt-4.1"}) + m.lastPrompt = "describe the diagram" + m.lastImages = []zeroruntime.ImageBlock{{MediaType: "image/png"}} + m.lastImageLabels = []string{"diagram.png"} + m.lastDocuments = []pendingDocument{{label: "spec.pdf", text: "notes"}} + m.input.SetValue("/edit") + + updated, _ := m.Update(testKey(tea.KeyEnter)) + next := updated.(model) + + if got := next.composerValue(); got != "describe the diagram" { + t.Fatalf("/edit should recall the prompt text, got %q", got) + } + if len(next.pendingImages) != 1 || len(next.pendingImageLabels) != 1 || next.pendingImageLabels[0] != "diagram.png" { + t.Fatalf("/edit should re-stage the remembered image, got imgs=%d labels=%v", + len(next.pendingImages), next.pendingImageLabels) + } + if len(next.pendingDocuments) != 1 || next.pendingDocuments[0].label != "spec.pdf" { + t.Fatalf("/edit should re-stage the remembered document, got %#v", next.pendingDocuments) + } +} + // /retry with no prior prompt reports that there's nothing to resend rather than // launching an empty run. func TestRetryWithoutPriorPromptIsNoOp(t *testing.T) { From 192b8e36bd40420331c3724cb6d9cfef7fe88dd6 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sun, 5 Jul 2026 12:15:26 +0530 Subject: [PATCH 32/32] fix(tools): bound bash output capture in memory, not just model-visible text MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously stdout/stderr streamed into unbounded bytes.Buffers and were only truncated to the 96 KiB budget AFTER the command finished, so a runaway command (cat huge.log, yes) could grow Zero's memory until it stalled or OOMed before budgetBashOutput ever ran. Replace the capture with boundedBuffer, an io.Writer that keeps only the head+tail each stream will surface and discards the middle as it arrives, while counting the true total for the truncation marker. Memory is now bounded to ~head+2×tail per stream regardless of output size. truncateHeadTail is refactored onto a total-aware core so the string path (and its tests) are unchanged, and budgetBashCapture reports the true raw_bytes even though only a bounded slice was held. Adds boundedBuffer/capture unit tests and an end-to-end test that streams ~5× the budget and asserts bounded emission with an accurate raw_bytes. --- internal/tools/bash.go | 114 ++++++++++++++++++++++++----- internal/tools/bash_budget_test.go | 66 +++++++++++++++++ internal/tools/bash_tool_test.go | 38 ++++++++++ 3 files changed, 201 insertions(+), 17 deletions(-) diff --git a/internal/tools/bash.go b/internal/tools/bash.go index 50af2125..138561fe 100644 --- a/internal/tools/bash.go +++ b/internal/tools/bash.go @@ -1,7 +1,6 @@ package tools import ( - "bytes" "context" "errors" "fmt" @@ -124,10 +123,14 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS defer plan.Cleanup() addSandboxMeta(meta, plan) - var stdout bytes.Buffer - var stderr bytes.Buffer - command.Stdout = &stdout - command.Stderr = &stderr + // Bound the capture so a command with runaway output (`cat huge.log`, `yes`) + // 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) + command.Stdout = stdout + command.Stderr = stderr // Kill the shell as a process group on timeout and bound the post-kill I/O // wait, so a backgrounded child cannot outlive the command or hang Run(). @@ -139,7 +142,12 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS err = command.Run() exitCode := commandExitCode(err) meta["exit_code"] = strconv.Itoa(exitCode) - stderrText := appendSandboxBlocks(stderr.String(), monitor.Stop()) + stdoutText := stdout.retained() + stderrRetained := stderr.retained() + stderrText := appendSandboxBlocks(stderrRetained, monitor.Stop()) + // Sandbox blocks are extra model-visible stderr bytes appended after capture; + // count them toward the true total so the budget/marker stay accurate. + stderrTotal := stderr.total + (len(stderrText) - len(stderrRetained)) if errors.Is(commandCtx.Err(), context.DeadlineExceeded) { return Result{ @@ -156,8 +164,8 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS Meta: meta, } } - markLikelySandboxDenial(meta, plan, exitCode, stdout.String(), stderrText) - outText, errText, truncated := budgetBashOutput(stdout.String(), stderrText, meta) + markLikelySandboxDenial(meta, plan, exitCode, stdoutText, stderrText) + outText, errText, truncated := budgetBashCapture(stdoutText, stdout.total, stderrText, stderrTotal, meta) return Result{ Status: StatusError, Output: formatBashOutputWithShellHint(commandText, outText, errText, exitCode, meta), @@ -166,8 +174,8 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS } } - markLikelySandboxDenial(meta, plan, exitCode, stdout.String(), stderrText) - outText, errText, truncated := budgetBashOutput(stdout.String(), stderrText, meta) + markLikelySandboxDenial(meta, plan, exitCode, stdoutText, stderrText) + outText, errText, truncated := budgetBashCapture(stdoutText, stdout.total, stderrText, stderrTotal, meta) if meta[SandboxLikelyDeniedMeta] == "true" { return Result{ Status: StatusError, @@ -363,8 +371,17 @@ const bashOutputBudgetBytes = 96 * 1024 // the read/search tools). Detection that needs the full output (sandbox-denial // scanning) must run on the raw strings before this is applied. func budgetBashOutput(stdout string, stderr string, meta map[string]string) (string, string, bool) { - outText, outRaw, outTrunc := truncateHeadTail(stdout, bashOutputBudgetBytes) - errText, errRaw, errTrunc := truncateHeadTail(stderr, bashOutputBudgetBytes) + return budgetBashCapture(stdout, len(stdout), stderr, len(stderr), meta) +} + +// budgetBashCapture is budgetBashOutput for the streaming-capture path: outTotal +// and errTotal are the true byte counts (from boundedBuffer.total), which may +// exceed the retained strings when the middle was dropped during capture. Meta's +// raw_bytes therefore reflects everything the command produced, not just what was +// kept in memory. +func budgetBashCapture(out string, outTotal int, errStr string, errTotal int, meta map[string]string) (string, string, bool) { + outText, outRaw, outTrunc := truncateHeadTailWithTotal(out, outTotal, bashOutputBudgetBytes) + errText, errRaw, errTrunc := truncateHeadTailWithTotal(errStr, errTotal, bashOutputBudgetBytes) truncated := outTrunc || errTrunc if meta != nil { emitted := len(outText) + len(errText) @@ -378,22 +395,85 @@ func budgetBashOutput(stdout string, stderr string, meta map[string]string) (str return outText, errText, truncated } +// 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 +// memory: the middle is discarded as it arrives instead of buffered whole and then +// truncated. total records the full size for the truncation marker even though the +// middle is never held. Not safe for concurrent writes; exec drives Stdout and +// Stderr from separate goroutines, so each stream gets its own buffer. +type boundedBuffer struct { + head []byte + headCap int + tail []byte + tailCap int + total int +} + +func newBoundedBuffer(headCap, tailCap int) *boundedBuffer { + return &boundedBuffer{headCap: headCap, tailCap: tailCap} +} + +func (b *boundedBuffer) Write(p []byte) (int, error) { + n := len(p) + b.total += n + // Fill the head until it reaches headCap; the head is written once and frozen. + if len(b.head) < b.headCap { + take := b.headCap - len(b.head) + if take > len(p) { + take = len(p) + } + b.head = append(b.head, p[:take]...) + p = p[take:] + } + // Anything past the head feeds a rolling tail that keeps only the last tailCap + // bytes. Compact once it grows past 2×tailCap so the append stays amortized O(1) + // while memory never exceeds ~2×tailCap. + if len(p) > 0 && b.tailCap > 0 { + b.tail = append(b.tail, p...) + if len(b.tail) > 2*b.tailCap { + b.tail = append(b.tail[:0], b.tail[len(b.tail)-b.tailCap:]...) + } + } + return n, nil +} + +// retained returns the kept head+tail bytes (marker-less) as a string. The junction +// between head and tail lands in the middle, which the display budget trims away; +// callers that need the true size read total separately. +func (b *boundedBuffer) retained() string { + if len(b.tail) > b.tailCap { + // Not yet compacted since the last overflow; expose only the last tailCap. + return string(b.head) + string(b.tail[len(b.tail)-b.tailCap:]) + } + return string(b.head) + string(b.tail) +} + // truncateHeadTail keeps the first and last halves of value when it exceeds // maxBytes, dropping the middle behind a marker. Returns the possibly-truncated // text, the raw byte length, and whether truncation happened. func truncateHeadTail(value string, maxBytes int) (string, int, bool) { - raw := len(value) - if maxBytes <= 0 || raw <= maxBytes { - return value, raw, false + // value holds the whole stream, so its length is the true raw size. + return truncateHeadTailWithTotal(value, len(value), maxBytes) +} + +// truncateHeadTailWithTotal head+tail-truncates value to maxBytes, using total — +// the full original byte count — for the "N bytes omitted" marker and the raw +// count. total may exceed len(value) when the middle was already discarded during +// bounded capture (boundedBuffer): value then holds only the retained head+tail, +// and this trims it to the display budget while still reporting the true total. +func truncateHeadTailWithTotal(value string, total, maxBytes int) (string, int, bool) { + if maxBytes <= 0 || total <= maxBytes { + return value, total, false } - marker := fmt.Sprintf("\n[zero] output truncated: %d bytes omitted from the middle — redirect to a file and read_file a range for the full text\n", raw-maxBytes) + marker := fmt.Sprintf("\n[zero] output truncated: %d bytes omitted from the middle — redirect to a file and read_file a range for the full text\n", total-maxBytes) budget := maxBytes - len(marker) if budget < 0 { budget = 0 } head := budget / 2 tail := budget - head - return utf8Prefix(value, head) + marker + utf8Suffix(value, tail), raw, true + return utf8Prefix(value, head) + marker + utf8Suffix(value, tail), total, true } func formatBashOutputWithShellHint(command string, stdout string, stderr string, exitCode int, meta map[string]string) string { diff --git a/internal/tools/bash_budget_test.go b/internal/tools/bash_budget_test.go index 45047003..adb3ab0d 100644 --- a/internal/tools/bash_budget_test.go +++ b/internal/tools/bash_budget_test.go @@ -1,6 +1,7 @@ package tools import ( + "fmt" "strconv" "strings" "testing" @@ -59,3 +60,68 @@ func TestBudgetBashOutputTruncatesHeadAndTail(t *testing.T) { t.Fatalf("emitted_bytes = %s, want %d", meta["emitted_bytes"], len(out)) } } + +// boundedBuffer must keep memory bounded to head+tail while counting every byte, +// so a runaway command can't OOM Zero before truncation runs. It keeps the very +// first bytes (head) and the very last bytes (tail), discarding the middle. +func TestBoundedBufferKeepsHeadAndTailBounded(t *testing.T) { + b := newBoundedBuffer(8, 8) + total := 0 + for i := 0; i < 1000; i++ { + chunk := []byte(fmt.Sprintf("%05d.", i)) // 6 bytes each + n, err := b.Write(chunk) + if err != nil || n != len(chunk) { + t.Fatalf("Write = (%d, %v), want (%d, nil)", n, err, len(chunk)) + } + total += len(chunk) + } + if b.total != total { + t.Fatalf("total = %d, want %d (every byte must be counted)", b.total, total) + } + // Retained memory is bounded regardless of how much streamed through. + if len(b.head) > 8 { + t.Fatalf("head grew past its cap: %d", len(b.head)) + } + if len(b.tail) > 2*8 { + t.Fatalf("tail grew past 2×cap (compaction failed): %d", len(b.tail)) + } + r := b.retained() + if len(r) > 8+8 { + t.Fatalf("retained %d bytes, want <= headCap+tailCap", len(r)) + } + if !strings.HasPrefix(r, "00000.") { + t.Fatalf("retained head lost, got prefix %q", r[:min(6, len(r))]) + } + if !strings.HasSuffix(r, "00999.") { + t.Fatalf("retained tail lost, got %q", r) + } +} + +// 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 + + meta := map[string]string{} + out, _, truncated := budgetBashCapture(retained, total, "", 0, meta) + + if !truncated { + t.Fatal("oversized capture must report truncated=true") + } + 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 { + t.Fatalf("emitted %d bytes exceeds budget %d", len(out), bashOutputBudgetBytes) + } + if meta["raw_bytes"] != strconv.Itoa(total) { + t.Fatalf("raw_bytes = %s, want the true total %d", meta["raw_bytes"], total) + } + // The marker must cite the true omitted count (total-budget), not retained-budget. + 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))]) + } +} diff --git a/internal/tools/bash_tool_test.go b/internal/tools/bash_tool_test.go index 11073d5d..69d15e60 100644 --- a/internal/tools/bash_tool_test.go +++ b/internal/tools/bash_tool_test.go @@ -8,6 +8,7 @@ import ( "os" "path/filepath" "runtime" + "strconv" "strings" "testing" "time" @@ -240,6 +241,43 @@ func TestBashToolRunsCommandInWorkspace(t *testing.T) { } } +// A command with runaway output must not be buffered whole in memory: the capture +// is bounded to head+tail, yet raw_bytes still reports the true (much larger) size. +// If capture were unbounded this would balloon Zero's memory before truncation. +func TestBashToolBoundsRunawayOutputCapture(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("uses a POSIX shell pipeline (yes | head)") + } + const produced = 500000 // ~5× the 96 KiB budget + + result := NewBashTool(t.TempDir()).Run(context.Background(), map[string]any{ + "command": fmt.Sprintf("yes ABCDEFGH | head -c %d", produced), + }) + + if result.Status != StatusOK { + t.Fatalf("expected ok status, got %s: %s", result.Status, result.Output) + } + if !result.Truncated || result.Meta["truncated"] != "true" { + t.Fatalf("runaway output should be flagged truncated, meta=%v", result.Meta) + } + if !strings.Contains(result.Output, "output truncated") { + t.Fatalf("expected a truncation marker in output") + } + // The emitted (model-visible) text is capped near the budget... + if len(result.Output) > bashOutputBudgetBytes+1024 { + t.Fatalf("emitted %d bytes far exceeds budget %d", len(result.Output), bashOutputBudgetBytes) + } + // ...but raw_bytes reflects the full stream, proving the total was counted while + // only a bounded slice was ever held (raw >> the ~2×budget retention cap). + raw, err := strconv.Atoi(result.Meta["raw_bytes"]) + if err != nil { + t.Fatalf("raw_bytes not an int: %q", result.Meta["raw_bytes"]) + } + if raw < produced { + t.Fatalf("raw_bytes = %d, want >= %d (full stream counted)", raw, produced) + } +} + func TestBashToolUsesRequestedCwd(t *testing.T) { root := t.TempDir() nested := filepath.Join(root, "nested")