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..55741f0d --- /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/compaction_preserve.go b/internal/agent/compaction_preserve.go index b3a0da05..6423a5b2 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,123 @@ 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 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{} + sequence := make([]string, 0) + 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 + } + sequence = append(sequence, path) + } + } + order := lastSeenOrder(sequence) + 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 last-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 +} + +// 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. +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 +431,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 +474,12 @@ 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 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. tools := mergeSkillEntries(preservedToolsToEntries(priorState.Tools), loadedToolSchemas(middle)) @@ -360,19 +493,76 @@ 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 } +// 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 { + 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..c8fe2b73 --- /dev/null +++ b/internal/agent/compaction_recent_edits_test.go @@ -0,0 +1,121 @@ +package agent + +import ( + "fmt" + "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 last-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) 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"}, {Path: "z.go", Note: "z"}}} + older := preservedEditsToEntries(prior.RecentEdits) + newer := []skillEntry{{name: "a.go", body: "new note"}, {name: "b.go", body: "added"}} + + 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 != "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/guardrails.go b/internal/agent/guardrails.go index 089dc818..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 @@ -43,8 +49,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 @@ -354,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. @@ -427,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 @@ -485,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/loop.go b/internal/agent/loop.go index 0874f85a..fc693534 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" @@ -145,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 @@ -153,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 @@ -340,6 +350,16 @@ 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. + // 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 // tool-call turn normalizes to "" and clears any prior reason. @@ -2462,6 +2482,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)) @@ -2505,7 +2536,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 @@ -2537,13 +2568,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 @@ -2577,6 +2608,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..078ebc4d --- /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") + } +} 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) + } +} diff --git a/internal/agent/reconnect.go b/internal/agent/reconnect.go index 641366a9..54b339ec 100644 --- a/internal/agent/reconnect.go +++ b/internal/agent/reconnect.go @@ -3,9 +3,11 @@ package agent import ( "context" "fmt" + "math/rand" "strings" "time" + "github.com/Gitlawb/zero/internal/errhint" "github.com/Gitlawb/zero/internal/zeroruntime" ) @@ -17,10 +19,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 +85,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 +108,19 @@ func shouldReconnect(ctx context.Context, err error) bool { if isContextLimitError(msg) || isImageRejectionError(err) { return false } + // 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", @@ -106,8 +131,6 @@ func shouldReconnect(ctx context.Context, err error) bool { "timed out", "temporarily unavailable", "i/o timeout", - "503", - "502", "server closed", "unexpected end", } { @@ -118,14 +141,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..84bd3a3f 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,13 @@ 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 (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)) { @@ -112,6 +123,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) { diff --git a/internal/cli/exec.go b/internal/cli/exec.go index d996e9a1..2e2f6279 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" @@ -949,6 +950,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()) + } + }) +} diff --git a/internal/errhint/errhint.go b/internal/errhint/errhint.go new file mode 100644 index 00000000..d52f5c66 --- /dev/null +++ b/internal/errhint/errhint.go @@ -0,0 +1,167 @@ +// 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 +) + +// 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. 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") || containsStatusCode(m, "401", "403"): + return Auth + case containsAny(m, "rate limit", "rate_limit", "too many requests", "quota", + "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"): + 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 +} + +// 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 { + 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) + 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 new file mode 100644 index 00000000..bab3c990 --- /dev/null +++ b/internal/errhint/errhint_test.go @@ -0,0 +1,94 @@ +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", "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", "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", "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}, + + // 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}, + {"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) { + 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("provider request error: 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("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) == "" { + 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/providers/factory.go b/internal/providers/factory.go index 1d9195c7..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, @@ -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"}}`) 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) + } +} diff --git a/internal/tools/bash.go b/internal/tools/bash.go index 090594fd..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,26 +164,31 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS Meta: meta, } } - markLikelySandboxDenial(meta, plan, exitCode, stdout.String(), stderrText) + markLikelySandboxDenial(meta, plan, exitCode, stdoutText, stderrText) + outText, errText, truncated := budgetBashCapture(stdoutText, stdout.total, stderrText, stderrTotal, meta) return Result{ - Status: StatusError, - Output: formatBashOutputWithShellHint(commandText, stdout.String(), stderrText, exitCode, meta), - Meta: meta, + Status: StatusError, + Output: formatBashOutputWithShellHint(commandText, outText, errText, exitCode, meta), + Truncated: truncated, + Meta: meta, } } - markLikelySandboxDenial(meta, plan, exitCode, stdout.String(), stderrText) + markLikelySandboxDenial(meta, plan, exitCode, stdoutText, stderrText) + outText, errText, truncated := budgetBashCapture(stdoutText, stdout.total, stderrText, stderrTotal, meta) if meta[SandboxLikelyDeniedMeta] == "true" { return Result{ - Status: StatusError, - Output: formatBashOutputWithShellHint(commandText, stdout.String(), stderrText, exitCode, meta), - Meta: meta, + Status: StatusError, + Output: formatBashOutputWithShellHint(commandText, outText, errText, exitCode, meta), + Truncated: truncated, + Meta: meta, } } return Result{ - Status: StatusOK, - Output: formatBashOutput(stdout.String(), stderrText, exitCode), - Meta: meta, + Status: StatusOK, + Output: formatBashOutput(outText, errText, exitCode), + Truncated: truncated, + Meta: meta, } } @@ -345,6 +358,124 @@ 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, bool) { + 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) + meta["raw_bytes"] = strconv.Itoa(outRaw + errRaw) + meta["emitted_bytes"] = strconv.Itoa(emitted) + meta["estimated_tokens"] = strconv.Itoa(estimatedTokensFromBytes(emitted)) + if truncated { + meta["truncated"] = "true" + } + } + 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) { + // 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", 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), total, 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..adb3ab0d --- /dev/null +++ b/internal/tools/bash_budget_test.go @@ -0,0 +1,127 @@ +package tools + +import ( + "fmt" + "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, 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) + } + 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, _, 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") + } + 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)) + } +} + +// 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") 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) diff --git a/internal/tui/coalesce.go b/internal/tui/coalesce.go new file mode 100644 index 00000000..26b3c56e --- /dev/null +++ b/internal/tui/coalesce.go @@ -0,0 +1,109 @@ +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 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) + // 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 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, + afterFunc: func(fn func()) coalesceTimer { + return time.AfterFunc(streamCoalesceInterval, fn) + }, + } +} + +// 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 — both under the lock so nothing can interleave between them. + c.drainAndForwardLocked() + c.forward(msg) + return + } + + // 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 { + c.drainAndForwardLocked() + } + c.runID = text.runID + c.buf = append(c.buf, text.delta...) + if c.timer == nil { + c.timer = c.afterFunc(c.flush) + } +} + +// 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 { + return + } + if c.timer != nil { + c.timer.Stop() + c.timer = nil + } + msg := agentTextMsg{runID: c.runID, delta: string(c.buf)} + c.buf = c.buf[:0] + c.forward(msg) +} diff --git a/internal/tui/coalesce_test.go b/internal/tui/coalesce_test.go new file mode 100644 index 00000000..9157b3e9 --- /dev/null +++ b/internal/tui/coalesce_test.go @@ -0,0 +1,182 @@ +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 +} + +// 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 (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) + } + + 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]) + } +} + +// 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) { + 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/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/commands.go b/internal/tui/commands.go index d3108b9f..b7edab32 100644 --- a/internal/tui/commands.go +++ b/internal/tui/commands.go @@ -41,6 +41,10 @@ const ( commandAddDir commandSelfCorrect commandTurns + commandRetry + commandEdit + commandCopy + commandExport commandNew commandUnknown ) @@ -209,9 +213,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, }, { @@ -257,6 +261,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/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/flush_test.go b/internal/tui/flush_test.go index 63ed4118..f43982cf 100644 --- a/internal/tui/flush_test.go +++ b/internal/tui/flush_test.go @@ -121,7 +121,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) @@ -133,7 +133,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/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/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/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 acfafc9c..1f43c8db 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" @@ -75,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). @@ -260,13 +266,30 @@ 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 + // 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 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 @@ -767,6 +790,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 } @@ -896,9 +924,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() @@ -1570,7 +1611,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 @@ -1853,20 +1894,23 @@ 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 - // 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, }) } - m.streamingText = "" + m.streamingText = nil m.streamingReasoning = "" m.streamingReasoningExpanded = false // Roll the completed run's wall-time into the session's rolling average so @@ -2041,14 +2085,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() @@ -2779,7 +2823,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) != "" { @@ -2843,7 +2887,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" @@ -3891,6 +3935,62 @@ 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 + } + // 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) == "" { + 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: + 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 @@ -3905,6 +4005,15 @@ 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. 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{ @@ -4092,7 +4201,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."}) @@ -4113,7 +4222,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 @@ -4190,12 +4299,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, @@ -4653,6 +4768,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 85eb242e..6563fb9f 100644 --- a/internal/tui/model_test.go +++ b/internal/tui/model_test.go @@ -2301,7 +2301,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/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/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) + } +} 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 + } + } +} 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/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/resume_scope_test.go b/internal/tui/resume_scope_test.go new file mode 100644 index 00000000..5f3d7847 --- /dev/null +++ b/internal/tui/resume_scope_test.go @@ -0,0 +1,31 @@ +package tui + +import ( + "runtime" + "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}, + // 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) { + 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/retry_edit_test.go b/internal/tui/retry_edit_test.go new file mode 100644 index 00000000..bd14f3c3 --- /dev/null +++ b/internal/tui/retry_edit_test.go @@ -0,0 +1,96 @@ +package tui + +import ( + "context" + "testing" + + tea "charm.land/bubbletea/v2" + + "github.com/Gitlawb/zero/internal/zeroruntime" +) + +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) + } +} + +// /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) { + 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) + } +} + +// /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/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{ diff --git a/internal/tui/session.go b/internal/tui/session.go index 76344876..9e3ed43c 100644 --- a/internal/tui/session.go +++ b/internal/tui/session.go @@ -4,6 +4,8 @@ import ( "encoding/json" "errors" "fmt" + "path/filepath" + "runtime" "strings" "time" @@ -80,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 != "" { @@ -220,9 +228,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 +331,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 +378,52 @@ 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 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) { + continue + } + if metas[i].EventCount == 0 { + continue + } + if !m.sessionHasResumableContent(metas[i].SessionID) { + continue + } + 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. 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 + } + 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 { events, err := m.sessionStore.ReadEvents(sessionID) if err != nil { diff --git a/internal/tui/session_controls.go b/internal/tui/session_controls.go index 5e3610c2..4ce15fb5 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 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.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 diff --git a/internal/tui/transcript_export.go b/internal/tui/transcript_export.go new file mode 100644 index 00000000..5c8e34be --- /dev/null +++ b/internal/tui/transcript_export.go @@ -0,0 +1,83 @@ +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) + } + // 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 new file mode 100644 index 00000000..efa55dce --- /dev/null +++ b/internal/tui/transcript_export_test.go @@ -0,0 +1,103 @@ +package tui + +import ( + "os" + "path/filepath" + "runtime" + "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) + } + 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. 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. + 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) + } +} 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/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) + } +} 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)