mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-08 09:49:34 -10:00
Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
533b41495d | ||
|
|
bb66a44ab3 | ||
|
|
6f97c90316 |
No files matched your search
@@ -325,43 +325,6 @@ 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 {
|
||||
@@ -390,7 +353,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 := state.calibratedTokens(estimateTokens(messages) + toolTokens)
|
||||
size := estimateTokens(messages) + toolTokens
|
||||
if size <= state.threshold {
|
||||
return messages
|
||||
}
|
||||
@@ -407,7 +370,7 @@ func (state *compactionState) maybeCompact(
|
||||
// entirely and preserve recent turns verbatim.
|
||||
if pruned, reclaimed := pruneStaleToolOutput(messages, state.preserveLast); reclaimed > 0 {
|
||||
messages = pruned
|
||||
size = state.calibratedTokens(estimateTokens(messages) + toolTokens)
|
||||
size = estimateTokens(messages) + toolTokens
|
||||
if size <= state.threshold {
|
||||
state.lowWaterMark = size
|
||||
return messages
|
||||
@@ -423,7 +386,7 @@ func (state *compactionState) maybeCompact(
|
||||
// later turn) can try again; we never drop messages on failure here.
|
||||
return messages
|
||||
}
|
||||
newSize := state.calibratedTokens(estimateTokens(compacted) + toolTokens)
|
||||
newSize := 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.
|
||||
@@ -481,7 +444,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 = state.calibratedTokens(estimateTokens(result) + estimateToolDefTokens(tools))
|
||||
state.lowWaterMark = estimateTokens(result) + estimateToolDefTokens(tools)
|
||||
return result, true, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -1,50 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -20,16 +20,6 @@ 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 (
|
||||
@@ -110,123 +100,12 @@ func formatPlanArguments(arguments string) string {
|
||||
}
|
||||
|
||||
// skillEntry is a named preserved body. It began as loaded-skill state and is
|
||||
// reused for loaded tools, project instruction blocks, and recent file edits.
|
||||
// reused for loaded tools and project instruction blocks.
|
||||
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.
|
||||
@@ -431,17 +310,11 @@ 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"`
|
||||
@@ -474,12 +347,6 @@ 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))
|
||||
@@ -493,76 +360,19 @@ func appendPreservedState(summary string, middle []zeroruntime.Message) string {
|
||||
projectInstructionEntries(middle),
|
||||
)
|
||||
|
||||
if block := formatPreservedState(plan, edits, tools, skills, instructions); block != "" {
|
||||
if block := formatPreservedState(plan, 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, edits, tools, skills, instructions []skillEntry) string {
|
||||
if plan == "" && len(edits) == 0 && len(tools) == 0 && len(skills) == 0 && len(instructions) == 0 {
|
||||
func formatPreservedState(plan string, tools, skills, instructions []skillEntry) string {
|
||||
if plan == "" && 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})
|
||||
}
|
||||
|
||||
@@ -1,121 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -1,62 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/lsp"
|
||||
)
|
||||
|
||||
// fileDiagnosticsTimeout bounds one inline post-edit diagnostics check so a
|
||||
// slow or wedged language server can never hang a tool call; on timeout the
|
||||
// edit simply reports without a diagnostics block.
|
||||
const fileDiagnosticsTimeout = 10 * time.Second
|
||||
|
||||
// NewFileDiagnostics adapts an *lsp.Manager to the per-edit inline diagnostics
|
||||
// callback (tools.RunOptions.Diagnostics): it reads the just-written file,
|
||||
// checks it against the file's language server, and formats error-severity
|
||||
// diagnostics for the model. Warnings and hints are excluded — nagging about
|
||||
// style on every edit is noise, while a type error the edit just introduced is
|
||||
// exactly what the model should see before its next step. Diagnostics are
|
||||
// rendered with workspace-relative paths: the absolute path would put the
|
||||
// local username/home directory into the model prompt and session transcript
|
||||
// on every edit. Returns nil when manager is nil, disabling inline diagnostics
|
||||
// entirely.
|
||||
func NewFileDiagnostics(manager *lsp.Manager, workspaceRoot string) func(context.Context, string) string {
|
||||
if manager == nil {
|
||||
return nil
|
||||
}
|
||||
return func(ctx context.Context, absPath string) string {
|
||||
text, err := os.ReadFile(absPath)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
checkCtx, cancel := context.WithTimeout(ctx, fileDiagnosticsTimeout)
|
||||
defer cancel()
|
||||
diagnostics, err := manager.Check(checkCtx, absPath, string(text))
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
errors := lsp.FilterBySeverity(diagnostics, lsp.SeverityError)
|
||||
if len(errors) == 0 {
|
||||
return ""
|
||||
}
|
||||
return lsp.FormatDiagnostics(diagnosticsDisplayPath(workspaceRoot, absPath), errors)
|
||||
}
|
||||
}
|
||||
|
||||
// diagnosticsDisplayPath renders absPath relative to the workspace root for
|
||||
// model-facing output, falling back to the file's base name when the path is
|
||||
// outside the workspace (a bare name still identifies the file without
|
||||
// exposing the directory layout).
|
||||
func diagnosticsDisplayPath(workspaceRoot, absPath string) string {
|
||||
if workspaceRoot != "" {
|
||||
if rel, err := filepath.Rel(workspaceRoot, absPath); err == nil && !strings.HasPrefix(rel, "..") {
|
||||
return rel
|
||||
}
|
||||
}
|
||||
return filepath.Base(absPath)
|
||||
}
|
||||
@@ -1,24 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// Diagnostics are model-facing: absolute paths would leak the local username
|
||||
// and directory layout into the prompt and session transcript on every edit.
|
||||
func TestDiagnosticsDisplayPath(t *testing.T) {
|
||||
root := filepath.Join("/Users", "someone", "project")
|
||||
cases := []struct {
|
||||
root, abs, want string
|
||||
}{
|
||||
{root, filepath.Join(root, "internal", "a.go"), filepath.Join("internal", "a.go")},
|
||||
{root, filepath.Join("/etc", "other.go"), "other.go"}, // outside root -> base name only
|
||||
{"", filepath.Join("/home", "user", "x.go"), "x.go"}, // no root -> base name only
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := diagnosticsDisplayPath(c.root, c.abs); got != c.want {
|
||||
t.Errorf("diagnosticsDisplayPath(%q, %q) = %q, want %q", c.root, c.abs, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -22,12 +22,6 @@ 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
|
||||
@@ -49,12 +43,8 @@ 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. 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
|
||||
// bad call.
|
||||
toolFailureStopAt = 4
|
||||
|
||||
// maxContinueNudges bounds how many times the headless completion gate
|
||||
// (Options.RequireCompletionSignal) re-prompts a model that stopped without a
|
||||
@@ -364,13 +354,8 @@ type guardState struct {
|
||||
emptyTurns int
|
||||
totalToolCalls int
|
||||
toolCallsSincePlanUpdate int
|
||||
// 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
|
||||
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.
|
||||
@@ -442,16 +427,11 @@ 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
|
||||
@@ -505,12 +485,10 @@ 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. 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)) {
|
||||
// stronger signal. One-shot per stale interval.
|
||||
if state.planEverCalled &&
|
||||
!state.staleReminderSent &&
|
||||
state.toolCallsSincePlanUpdate >= staleToolCallThreshold {
|
||||
state.staleReminderSent = true
|
||||
return planStaleReminder(state.toolCallsSincePlanUpdate)
|
||||
}
|
||||
|
||||
+6
-81
@@ -25,12 +25,7 @@ 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.
|
||||
//
|
||||
// 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 maxStreamStallRetries = 2
|
||||
|
||||
const (
|
||||
toolResultMetaControl = "control"
|
||||
@@ -150,11 +145,6 @@ 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
|
||||
@@ -163,7 +153,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, _ := partitionToolsCached(registry, permissionMode, options, loaded, toolDefCache)
|
||||
exposed, _ := partitionTools(registry, permissionMode, options, loaded)
|
||||
|
||||
// PROACTIVE compaction: if the history is approaching the model's
|
||||
// context window, summarize the oldest middle before building the
|
||||
@@ -350,16 +340,6 @@ 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.
|
||||
@@ -495,35 +475,11 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
|
||||
// between tool_results breaks strict provider replay) — same after-batch
|
||||
// rationale as turnRequestedModel above.
|
||||
var changedFilesThisBatch []string
|
||||
// Parallel read-ahead state: results for calls[precomputedStart:precomputedEnd]
|
||||
// executed concurrently, consumed strictly in order below.
|
||||
var precomputed []precomputedToolResult
|
||||
precomputedStart, precomputedEnd := 0, 0
|
||||
for index, call := range collected.ToolCalls {
|
||||
// When this call starts a consecutive run of >= 2 auto-allowed read-only
|
||||
// calls, execute the whole run concurrently now (see parallel_tools.go).
|
||||
// The scan happens lazily at the run's first index — never ahead of a
|
||||
// pending mutating call — so read-after-write ordering is preserved.
|
||||
if index >= precomputedEnd {
|
||||
runEnd := index
|
||||
for runEnd < len(collected.ToolCalls) && parallelSafeToolCall(registry, collected.ToolCalls[runEnd], options) {
|
||||
runEnd++
|
||||
}
|
||||
if runEnd-index >= 2 {
|
||||
precomputed = executeParallelReadBatch(ctx, registry, collected.ToolCalls, index, runEnd, permissionMode, options)
|
||||
precomputedStart, precomputedEnd = index, runEnd
|
||||
}
|
||||
}
|
||||
if options.OnToolCall != nil {
|
||||
options.OnToolCall(call)
|
||||
}
|
||||
var toolResult ToolResult
|
||||
var abortErr error
|
||||
if index >= precomputedStart && index < precomputedEnd {
|
||||
toolResult, abortErr = precomputed[index-precomputedStart].result, precomputed[index-precomputedStart].abortErr
|
||||
} else {
|
||||
toolResult, abortErr = executeToolCall(ctx, registry, call, permissionMode, options)
|
||||
}
|
||||
toolResult, abortErr := executeToolCall(ctx, registry, call, permissionMode, options)
|
||||
if options.OnToolResult != nil {
|
||||
options.OnToolResult(toolResult)
|
||||
}
|
||||
@@ -1095,8 +1051,6 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
|
||||
// Per-session file version tracker so write_file/edit_file refuse to clobber
|
||||
// a file that changed on disk outside Zero since it was last read.
|
||||
FileTracker: options.FileTracker,
|
||||
// Inline post-edit diagnostics for mutating tools (nil = disabled).
|
||||
Diagnostics: options.FileDiagnostics,
|
||||
// Forward the run's operator tool filters so a filter-aware tool
|
||||
// (tool_search) never discloses or loads an operator-hidden deferred tool.
|
||||
EnabledTools: options.EnabledTools,
|
||||
@@ -2508,17 +2462,6 @@ 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))
|
||||
@@ -2562,7 +2505,7 @@ func partitionToolsCached(registry *tools.Registry, permissionMode PermissionMod
|
||||
if tool.Name() == tools.ToolSearchToolName {
|
||||
continue
|
||||
}
|
||||
definitions = append(definitions, cachedRuntimeToolDefinition(defCache, tool))
|
||||
definitions = append(definitions, runtimeToolDefinition(tool))
|
||||
}
|
||||
sort.Slice(definitions, func(left int, right int) bool {
|
||||
return definitions[left].Name < definitions[right].Name
|
||||
@@ -2594,13 +2537,13 @@ func partitionToolsCached(registry *tools.Registry, permissionMode PermissionMod
|
||||
}
|
||||
if tools.IsDeferred(tool) {
|
||||
if loaded[name] {
|
||||
loadedTail = append(loadedTail, cachedRuntimeToolDefinition(defCache, tool))
|
||||
loadedTail = append(loadedTail, runtimeToolDefinition(tool))
|
||||
} else {
|
||||
hiddenTools = append(hiddenTools, tool)
|
||||
}
|
||||
continue
|
||||
}
|
||||
eager = append(eager, cachedRuntimeToolDefinition(defCache, tool))
|
||||
eager = append(eager, runtimeToolDefinition(tool))
|
||||
}
|
||||
sort.Slice(eager, func(left int, right int) bool {
|
||||
return eager[left].Name < eager[right].Name
|
||||
@@ -2634,24 +2577,6 @@ func partitionToolsCached(registry *tools.Registry, permissionMode PermissionMod
|
||||
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 {
|
||||
|
||||
@@ -1,97 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
)
|
||||
|
||||
// Parallel read-ahead for tool batches. When a turn requests several
|
||||
// independent lookups (read_file + grep + glob is the common shape), executing
|
||||
// them one after another serializes pure I/O waits. A consecutive run of
|
||||
// auto-allowed read-only calls is executed concurrently instead; results are
|
||||
// then consumed in the original call order, so guard counters, message
|
||||
// ordering, abort semantics, and the surface's call/result event pairing are
|
||||
// byte-identical to sequential execution. Runs never span a mutating call: a
|
||||
// read that follows a write must observe the write, so eligibility is decided
|
||||
// per consecutive run, not per batch.
|
||||
|
||||
// maxParallelReadTools bounds concurrent read-only tool executions in a turn.
|
||||
const maxParallelReadTools = 8
|
||||
|
||||
// precomputedToolResult is one parallel read-ahead execution, keyed back to
|
||||
// its batch index by the caller.
|
||||
type precomputedToolResult struct {
|
||||
result ToolResult
|
||||
abortErr error
|
||||
}
|
||||
|
||||
// parallelSafeToolCall reports whether call may run concurrently with its
|
||||
// neighbors: the tool must exist, be side-effect-free (SideEffectRead), and be
|
||||
// auto-allowed for these args, so no interactive prompt or workspace mutation
|
||||
// is on the hot path. Loop-intercepted tools (ask_user, request_permissions)
|
||||
// and tool_search (mutates the deferred-tool set) stay sequential.
|
||||
func parallelSafeToolCall(registry *tools.Registry, call ToolCall, options Options) bool {
|
||||
switch call.Name {
|
||||
case "ask_user", tools.RequestPermissionsToolName, tools.ToolSearchToolName:
|
||||
return false
|
||||
}
|
||||
tool, found := registry.Get(call.Name)
|
||||
if !found || tool.Safety().SideEffect != tools.SideEffectRead {
|
||||
return false
|
||||
}
|
||||
args := map[string]any{}
|
||||
if call.Arguments != "" {
|
||||
if err := decodeToolArguments(call.Arguments, &args); err != nil {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return effectivePermission(tool, args) == tools.PermissionAllow
|
||||
}
|
||||
|
||||
// executeParallelReadBatch runs calls[start:end] concurrently (bounded by
|
||||
// maxParallelReadTools) and returns results indexed relative to start. All
|
||||
// execution-side callbacks that can fire inside executeToolCall are serialized
|
||||
// behind one mutex: a permission prompt (a sandbox preflight can demand one
|
||||
// even for an auto-allowed read) must never appear twice at once on an
|
||||
// interactive front-end, and OnPermission event handlers append to shared
|
||||
// session-recording state without their own locking — two batched reads under
|
||||
// a granted extra root would otherwise race (the pre-batch serial loop never
|
||||
// had two callbacks in flight at once).
|
||||
func executeParallelReadBatch(ctx context.Context, registry *tools.Registry, calls []ToolCall, start, end int, permissionMode PermissionMode, options Options) []precomputedToolResult {
|
||||
batchOptions := options
|
||||
var callbackMutex sync.Mutex
|
||||
if options.OnPermissionRequest != nil {
|
||||
inner := options.OnPermissionRequest
|
||||
batchOptions.OnPermissionRequest = func(ctx context.Context, request PermissionRequest) (PermissionDecision, error) {
|
||||
callbackMutex.Lock()
|
||||
defer callbackMutex.Unlock()
|
||||
return inner(ctx, request)
|
||||
}
|
||||
}
|
||||
if options.OnPermission != nil {
|
||||
inner := options.OnPermission
|
||||
batchOptions.OnPermission = func(event PermissionEvent) {
|
||||
callbackMutex.Lock()
|
||||
defer callbackMutex.Unlock()
|
||||
inner(event)
|
||||
}
|
||||
}
|
||||
|
||||
results := make([]precomputedToolResult, end-start)
|
||||
semaphore := make(chan struct{}, maxParallelReadTools)
|
||||
var waitGroup sync.WaitGroup
|
||||
for index := start; index < end; index++ {
|
||||
waitGroup.Add(1)
|
||||
go func(index int) {
|
||||
defer waitGroup.Done()
|
||||
semaphore <- struct{}{}
|
||||
defer func() { <-semaphore }()
|
||||
result, abortErr := executeToolCall(ctx, registry, calls[index], permissionMode, batchOptions)
|
||||
results[index-start] = precomputedToolResult{result: result, abortErr: abortErr}
|
||||
}(index)
|
||||
}
|
||||
waitGroup.Wait()
|
||||
return results
|
||||
}
|
||||
@@ -1,204 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// probeTool records execution overlap and ordering so tests can assert what
|
||||
// actually ran concurrently.
|
||||
type probeTool struct {
|
||||
name string
|
||||
sideEffect tools.SideEffect
|
||||
delay time.Duration
|
||||
// shared, when set, receives this probe's start/end entries too, so a test
|
||||
// can observe ordering ACROSS probes (e.g. reads vs a write barrier).
|
||||
shared *probeLog
|
||||
|
||||
mu sync.Mutex
|
||||
active int
|
||||
maxActive int
|
||||
log []string
|
||||
}
|
||||
|
||||
// probeLog is a mutex-guarded event log shared between probes.
|
||||
type probeLog struct {
|
||||
mu sync.Mutex
|
||||
entries []string
|
||||
}
|
||||
|
||||
func (log *probeLog) append(entry string) {
|
||||
log.mu.Lock()
|
||||
defer log.mu.Unlock()
|
||||
log.entries = append(log.entries, entry)
|
||||
}
|
||||
|
||||
func (log *probeLog) snapshot() []string {
|
||||
log.mu.Lock()
|
||||
defer log.mu.Unlock()
|
||||
return append([]string(nil), log.entries...)
|
||||
}
|
||||
|
||||
func (tool *probeTool) Name() string { return tool.name }
|
||||
func (tool *probeTool) Description() string { return "test probe tool" }
|
||||
func (tool *probeTool) Parameters() tools.Schema {
|
||||
return tools.Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]tools.PropertySchema{"id": {Type: "string"}},
|
||||
AdditionalProperties: false,
|
||||
}
|
||||
}
|
||||
func (tool *probeTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tool.sideEffect, Permission: tools.PermissionAllow, Reason: "test"}
|
||||
}
|
||||
func (tool *probeTool) Run(_ context.Context, args map[string]any) tools.Result {
|
||||
id, _ := args["id"].(string)
|
||||
tool.mu.Lock()
|
||||
tool.active++
|
||||
if tool.active > tool.maxActive {
|
||||
tool.maxActive = tool.active
|
||||
}
|
||||
tool.log = append(tool.log, "start:"+id)
|
||||
tool.mu.Unlock()
|
||||
if tool.shared != nil {
|
||||
tool.shared.append("start:" + id)
|
||||
}
|
||||
time.Sleep(tool.delay)
|
||||
if tool.shared != nil {
|
||||
tool.shared.append("end:" + id)
|
||||
}
|
||||
tool.mu.Lock()
|
||||
tool.active--
|
||||
tool.log = append(tool.log, "end:"+id)
|
||||
tool.mu.Unlock()
|
||||
return tools.Result{Status: tools.StatusOK, Output: "probe " + id}
|
||||
}
|
||||
|
||||
func probeCallEvents(callID, toolName, id string) []zeroruntime.StreamEvent {
|
||||
return []zeroruntime.StreamEvent{
|
||||
{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: callID, ToolName: toolName},
|
||||
{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: callID, ArgumentsFragment: fmt.Sprintf(`{"id":%q}`, id)},
|
||||
{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: callID},
|
||||
}
|
||||
}
|
||||
|
||||
func TestParallelSafeToolCall(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(&probeTool{name: "probe_read", sideEffect: tools.SideEffectRead})
|
||||
registry.Register(&probeTool{name: "probe_write", sideEffect: tools.SideEffectWrite})
|
||||
|
||||
call := func(name, args string) ToolCall { return ToolCall{ID: "c", Name: name, Arguments: args} }
|
||||
if !parallelSafeToolCall(registry, call("probe_read", `{"id":"a"}`), Options{}) {
|
||||
t.Fatal("auto-allowed read tool must be parallel-safe")
|
||||
}
|
||||
if parallelSafeToolCall(registry, call("probe_write", `{"id":"a"}`), Options{}) {
|
||||
t.Fatal("mutating tool must not be parallel-safe")
|
||||
}
|
||||
if parallelSafeToolCall(registry, call("unknown_tool", `{}`), Options{}) {
|
||||
t.Fatal("unknown tool must not be parallel-safe")
|
||||
}
|
||||
if parallelSafeToolCall(registry, call("probe_read", `{"id":`), Options{}) {
|
||||
t.Fatal("undecodable arguments must not be parallel-safe")
|
||||
}
|
||||
if parallelSafeToolCall(registry, call("ask_user", `{}`), Options{}) {
|
||||
t.Fatal("loop-intercepted tools must stay sequential")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunExecutesConsecutiveReadsConcurrently(t *testing.T) {
|
||||
probe := &probeTool{name: "probe_read", sideEffect: tools.SideEffectRead, delay: 60 * time.Millisecond}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(probe)
|
||||
|
||||
turnOne := append(probeCallEvents("call-1", "probe_read", "a"), probeCallEvents("call-2", "probe_read", "b")...)
|
||||
turnOne = append(turnOne, probeCallEvents("call-3", "probe_read", "c")...)
|
||||
turnOne = append(turnOne, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone})
|
||||
provider := &mockProvider{
|
||||
turns: [][]zeroruntime.StreamEvent{
|
||||
turnOne,
|
||||
{{Type: zeroruntime.StreamEventText, Content: "done"}, {Type: zeroruntime.StreamEventDone}},
|
||||
},
|
||||
}
|
||||
|
||||
var results []ToolResult
|
||||
_, err := Run(context.Background(), "probe", provider, Options{
|
||||
Registry: registry,
|
||||
OnToolResult: func(result ToolResult) { results = append(results, result) },
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if probe.maxActive < 2 {
|
||||
t.Fatalf("consecutive read-only calls must overlap, max concurrency was %d", probe.maxActive)
|
||||
}
|
||||
// Results must still be recorded in original call order.
|
||||
if len(results) != 3 || results[0].ToolCallID != "call-1" || results[1].ToolCallID != "call-2" || results[2].ToolCallID != "call-3" {
|
||||
t.Fatalf("tool results out of order: %#v", results)
|
||||
}
|
||||
messages := provider.requests[1].Messages
|
||||
var toolOrder []string
|
||||
for _, message := range messages {
|
||||
if message.Role == zeroruntime.MessageRoleTool {
|
||||
toolOrder = append(toolOrder, message.ToolCallID)
|
||||
}
|
||||
}
|
||||
if len(toolOrder) != 3 || toolOrder[0] != "call-1" || toolOrder[1] != "call-2" || toolOrder[2] != "call-3" {
|
||||
t.Fatalf("recorded tool messages out of order: %v", toolOrder)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunParallelReadsNeverSpanMutatingCall(t *testing.T) {
|
||||
shared := &probeLog{}
|
||||
read := &probeTool{name: "probe_read", sideEffect: tools.SideEffectRead, delay: 30 * time.Millisecond, shared: shared}
|
||||
write := &probeTool{name: "probe_write", sideEffect: tools.SideEffectWrite, shared: shared}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(read)
|
||||
registry.Register(write)
|
||||
|
||||
turnOne := append(probeCallEvents("call-1", "probe_read", "r1"), probeCallEvents("call-2", "probe_read", "r2")...)
|
||||
turnOne = append(turnOne, probeCallEvents("call-3", "probe_write", "w")...)
|
||||
turnOne = append(turnOne, probeCallEvents("call-4", "probe_read", "r3")...)
|
||||
turnOne = append(turnOne, probeCallEvents("call-5", "probe_read", "r4")...)
|
||||
turnOne = append(turnOne, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone})
|
||||
provider := &mockProvider{
|
||||
turns: [][]zeroruntime.StreamEvent{
|
||||
turnOne,
|
||||
{{Type: zeroruntime.StreamEventText, Content: "done"}, {Type: zeroruntime.StreamEventDone}},
|
||||
},
|
||||
}
|
||||
|
||||
_, err := Run(context.Background(), "probe", provider, Options{Registry: registry})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Cross-probe ordering on the SHARED log: the write must start only after
|
||||
// both first-batch reads finished, and both second-batch reads must start
|
||||
// only after the write finished — batches never cross a mutating call.
|
||||
log := shared.snapshot()
|
||||
index := func(entry string) int {
|
||||
for i, e := range log {
|
||||
if e == entry {
|
||||
return i
|
||||
}
|
||||
}
|
||||
t.Fatalf("entry %q missing from shared log: %v", entry, log)
|
||||
return -1
|
||||
}
|
||||
writeStart, writeEnd := index("start:w"), index("end:w")
|
||||
if firstBatchMaxEnd := max(index("end:r1"), index("end:r2")); writeStart < firstBatchMaxEnd {
|
||||
t.Fatalf("write started before the first read batch finished: %v", log)
|
||||
}
|
||||
if secondBatchMinStart := min(index("start:r3"), index("start:r4")); secondBatchMinStart < writeEnd {
|
||||
t.Fatalf("second read batch started before the write finished: %v", log)
|
||||
}
|
||||
if read.maxActive < 2 {
|
||||
t.Fatalf("reads within a batch must overlap, max concurrency was %d", read.maxActive)
|
||||
}
|
||||
}
|
||||
@@ -1,63 +0,0 @@
|
||||
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")
|
||||
}
|
||||
}
|
||||
@@ -1,64 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -3,11 +3,9 @@ package agent
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math/rand"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/errhint"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
@@ -19,20 +17,10 @@ import (
|
||||
// times. We retry ONLY the connect (not a partially-consumed stream), so no
|
||||
// already-forwarded OnText is ever duplicated.
|
||||
const (
|
||||
// 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
|
||||
maxStreamReconnects = 2
|
||||
streamReconnectBase = 500 * time.Millisecond
|
||||
)
|
||||
|
||||
// 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.
|
||||
@@ -85,7 +73,7 @@ func streamWithReconnect(ctx context.Context, provider Provider, request zerorun
|
||||
if notify != nil {
|
||||
notify(attempt, maxStreamReconnects)
|
||||
}
|
||||
if waitErr := sleepWithContext(ctx, jitteredBackoff(attempt)); waitErr != nil {
|
||||
if waitErr := sleepWithContext(ctx, backoffFor(attempt)); waitErr != nil {
|
||||
return nil, err // ctx cancelled while waiting; surface the original error
|
||||
}
|
||||
stream, err = provider.StreamCompletion(ctx, request)
|
||||
@@ -108,19 +96,6 @@ 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",
|
||||
@@ -131,6 +106,8 @@ func shouldReconnect(ctx context.Context, err error) bool {
|
||||
"timed out",
|
||||
"temporarily unavailable",
|
||||
"i/o timeout",
|
||||
"503",
|
||||
"502",
|
||||
"server closed",
|
||||
"unexpected end",
|
||||
} {
|
||||
@@ -141,31 +118,14 @@ 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()
|
||||
|
||||
@@ -44,10 +44,6 @@ 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)
|
||||
@@ -91,7 +87,7 @@ func TestShouldReconnectClassification(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
disconnects := []string{
|
||||
"unexpected EOF", "connection reset by peer", "broken pipe",
|
||||
"i/o timeout", "server closed the connection", "connection refused",
|
||||
"i/o timeout", "503 Service Unavailable", "server closed the connection",
|
||||
}
|
||||
for _, m := range disconnects {
|
||||
if !shouldReconnect(ctx, errors.New(m)) {
|
||||
@@ -101,13 +97,6 @@ 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)) {
|
||||
@@ -123,24 +112,6 @@ 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) {
|
||||
|
||||
@@ -285,12 +285,6 @@ type Options struct {
|
||||
// ceiling and the autonomy gate. nil disables it entirely (the loop is
|
||||
// byte-identical to before). One instance per run — it holds attempt state.
|
||||
SelfCorrect *SelfCorrector
|
||||
// FileDiagnostics, when set, is passed to mutating tools so edit_file /
|
||||
// write_file append error-severity language diagnostics for the file they
|
||||
// just wrote directly to their tool output — the model sees an error it
|
||||
// introduced in the same turn instead of a later verification pass. Build
|
||||
// one with NewFileDiagnostics. nil disables inline diagnostics.
|
||||
FileDiagnostics func(ctx context.Context, absPath string) string
|
||||
|
||||
// RequireCompletionSignal gates run completion for HEADLESS exec. Without it,
|
||||
// any assistant turn that produces text but no tool call is accepted as the
|
||||
|
||||
@@ -19,7 +19,6 @@ import (
|
||||
"github.com/Gitlawb/zero/internal/hooks"
|
||||
"github.com/Gitlawb/zero/internal/localcontrol"
|
||||
"github.com/Gitlawb/zero/internal/mcp"
|
||||
"github.com/Gitlawb/zero/internal/modelregistry"
|
||||
"github.com/Gitlawb/zero/internal/observability"
|
||||
"github.com/Gitlawb/zero/internal/plugins"
|
||||
"github.com/Gitlawb/zero/internal/providerhealth"
|
||||
@@ -230,12 +229,6 @@ func runWithDeps(args []string, stdout io.Writer, stderr io.Writer, deps appDeps
|
||||
defer observability.Recover(observability.DefaultCrashDir(), "cli", stderr, &exitCode)
|
||||
deps = fillAppDeps(deps)
|
||||
|
||||
// CLI runs opt into the models.dev overlay (cached live context limits and
|
||||
// pricing on top of the curated catalog). Explicitly enabled here — and only
|
||||
// here — so library consumers and hermetic tests are never perturbed by a
|
||||
// cache file on the machine. The refresh itself is fired in exec/TUI startup.
|
||||
modelregistry.EnableModelsDevOverlay()
|
||||
|
||||
addDirs, args, err := splitLeadingAddDirFlags(args)
|
||||
if err != nil {
|
||||
return writeAppError(stderr, err.Error(), 1)
|
||||
@@ -551,11 +544,6 @@ func runInteractiveTUI(stderr io.Writer, deps appDeps, permissionMode agent.Perm
|
||||
}
|
||||
|
||||
func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode agent.PermissionMode, addDirs []string, theme string, forceSetup bool) int {
|
||||
// Refresh the models.dev pricing/limits cache in the background when stale;
|
||||
// the overlay is read at registry construction from the cache file, so this
|
||||
// benefits the next run and never blocks or fails this one.
|
||||
go func() { _ = modelregistry.RefreshModelsDevCache(context.Background()) }()
|
||||
|
||||
workspaceRoot, err := deps.getwd()
|
||||
if err != nil {
|
||||
return writeAppError(stderr, "failed to resolve workspace: "+err.Error(), 1)
|
||||
|
||||
+12
-33
@@ -12,7 +12,6 @@ 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"
|
||||
@@ -136,11 +135,6 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
|
||||
return exitSuccess
|
||||
}
|
||||
|
||||
// Refresh the models.dev pricing/limits cache in the background when stale;
|
||||
// the overlay is read at registry construction from the cache file, so this
|
||||
// benefits the next run and never blocks or fails this one.
|
||||
go func() { _ = modelregistry.RefreshModelsDevCache(context.Background()) }()
|
||||
|
||||
// A mode seeds model/effort/max-turns/tool filters as a preset. Expand it up
|
||||
// front — before tool-filter validation and the --list-tools branch — so a
|
||||
// mode-injected tool filter is validated and reflected in --list-tools, and a
|
||||
@@ -505,10 +499,9 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
|
||||
// by default the corrector is nil, leaving the agent loop byte-identical. When
|
||||
// on we verify with both the workspace test plan and LSP diagnostics over the
|
||||
// changed files; the autonomy gate inside the corrector still decides whether
|
||||
// failures auto-fix or just report. fileDiagnostics (always on, lazy) gives
|
||||
// edit_file/write_file inline error diagnostics for the file they just wrote.
|
||||
// lspShutdown tears down any language-server sessions either half spawned.
|
||||
selfCorrector, fileDiagnostics, lspShutdown := newExecSelfCorrector(options.selfCorrect, workspaceRoot, options.autonomy)
|
||||
// failures auto-fix or just report. lspShutdown tears down any language-server
|
||||
// sessions the LSP half spawned (no-op when self-correct is off).
|
||||
selfCorrector, lspShutdown := newExecSelfCorrector(options.selfCorrect, workspaceRoot, options.autonomy)
|
||||
defer lspShutdown()
|
||||
result, err := agent.Run(runCtx, agentPrompt, provider, agent.Options{
|
||||
MaxTurns: resolved.MaxTurns,
|
||||
@@ -532,7 +525,6 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
|
||||
PermissionMode: permissionMode,
|
||||
Autonomy: options.autonomy,
|
||||
SelfCorrect: selfCorrector,
|
||||
FileDiagnostics: fileDiagnostics,
|
||||
// Headless exec: don't accept a no-tool-call turn as "done" while work
|
||||
// clearly remains (pending plan items / a mid-step continuation cue) —
|
||||
// nudge to continue, and finalize as INCOMPLETE rather than false success
|
||||
@@ -708,28 +700,23 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
|
||||
// unless a changed file's language actually has one installed on PATH. The
|
||||
// returned cleanup shuts that manager down (terminating any spawned server
|
||||
// sessions); it is a no-op when self-correct is off.
|
||||
func newExecSelfCorrector(enabled bool, workspaceRoot string, autonomy string) (*agent.SelfCorrector, func(context.Context, string) string, func()) {
|
||||
// The manager is created regardless of --self-correct: it also backs the
|
||||
// always-on inline post-edit diagnostics (agent.NewFileDiagnostics), and it
|
||||
// stays lazy — no language server is spawned unless an edited file's
|
||||
// language actually has one installed on PATH.
|
||||
manager := lsp.NewManager(workspaceRoot)
|
||||
cleanup := func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
_ = manager.Shutdown(ctx)
|
||||
}
|
||||
fileDiagnostics := agent.NewFileDiagnostics(manager, workspaceRoot)
|
||||
func newExecSelfCorrector(enabled bool, workspaceRoot string, autonomy string) (*agent.SelfCorrector, func()) {
|
||||
if !enabled {
|
||||
return nil, fileDiagnostics, cleanup
|
||||
return nil, func() {}
|
||||
}
|
||||
manager := lsp.NewManager(workspaceRoot)
|
||||
corrector := agent.NewSelfCorrector(workspaceRoot, agent.NewLSPDiagnosticsChecker(manager), agent.NewProjectVerifier(workspaceRoot), agent.SelfCorrectConfig{
|
||||
Enabled: true,
|
||||
IncludeTests: true,
|
||||
IncludeLSP: true,
|
||||
Autonomy: autonomy,
|
||||
})
|
||||
return corrector, fileDiagnostics, cleanup
|
||||
cleanup := func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
_ = manager.Shutdown(ctx)
|
||||
}
|
||||
return corrector, cleanup
|
||||
}
|
||||
|
||||
func deferredEligibleCount(registry *tools.Registry, permissionMode agent.PermissionMode, enabledTools []string, disabledTools []string) int {
|
||||
@@ -962,14 +949,6 @@ 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
|
||||
}
|
||||
|
||||
|
||||
@@ -1,50 +0,0 @@
|
||||
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())
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -590,17 +590,9 @@ func applyProviderEnv(cfg *FileConfig, providerKind ProviderKind, env envProfile
|
||||
if profile.Name == "" {
|
||||
profile.Name = string(providerKind)
|
||||
}
|
||||
// A non-official baseURL from the environment signals a proxy/gateway, so promote
|
||||
// the first-party kind to its -compatible transport (which accepts a custom URL).
|
||||
// Env-only by design: a custom baseURL in a project config.json is left as-is and
|
||||
// rejected by Resolve, so an untrusted project can't redirect a first-party key to
|
||||
// another host.
|
||||
if providerKind == ProviderKindOpenAI && baseURL != "" && !isOfficialOpenAIBaseURL(baseURL) {
|
||||
profile.ProviderKind = ProviderKindOpenAICompatible
|
||||
}
|
||||
if providerKind == ProviderKindAnthropic && baseURL != "" && !isOfficialAnthropicBaseURL(baseURL) {
|
||||
profile.ProviderKind = ProviderKindAnthropicCompat
|
||||
}
|
||||
// When the env supplies only credentials (no baseURL) for a provider name that
|
||||
// already exists, don't force the standard transport kind — a same-named
|
||||
// openai-compatible/anthropic-compatible proxy or gateway must keep its kind.
|
||||
|
||||
@@ -617,35 +617,6 @@ func TestResolveUsesOpenAIEnvFallback(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveUsesAnthropicEnvBaseURLAsCompatible(t *testing.T) {
|
||||
// ANTHROPIC_BASE_URL pointing at a proxy/gateway must resolve to an
|
||||
// anthropic-compatible provider (mirroring OPENAI_BASE_URL), not the "requires
|
||||
// official baseURL" rejection — issue #479. The env is user-controlled, so a
|
||||
// custom URL there is a deliberate gateway choice, unlike a project config.json.
|
||||
resolved, err := Resolve(ResolveOptions{
|
||||
Env: map[string]string{
|
||||
"ANTHROPIC_API_KEY": "sk-ant-env",
|
||||
"ANTHROPIC_BASE_URL": "https://gateway.example/anthropic",
|
||||
"ANTHROPIC_MODEL": "claude-custom",
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Resolve() error = %v", err)
|
||||
}
|
||||
if resolved.ActiveProvider != "anthropic" {
|
||||
t.Fatalf("ActiveProvider = %q, want anthropic", resolved.ActiveProvider)
|
||||
}
|
||||
if resolved.Provider.ProviderKind != ProviderKindAnthropicCompat {
|
||||
t.Fatalf("ProviderKind = %q, want anthropic-compatible", resolved.Provider.ProviderKind)
|
||||
}
|
||||
if resolved.Provider.BaseURL != "https://gateway.example/anthropic" {
|
||||
t.Fatalf("BaseURL = %q, want gateway URL", resolved.Provider.BaseURL)
|
||||
}
|
||||
if resolved.Provider.APIKey != "sk-ant-env" || resolved.Provider.Model != "claude-custom" {
|
||||
t.Fatalf("Provider = %#v, want env credentials/model", resolved.Provider)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveUsesOpenAIAPIKeyOnlyWithDefaultModel(t *testing.T) {
|
||||
resolved, err := Resolve(ResolveOptions{
|
||||
Env: map[string]string{
|
||||
|
||||
@@ -131,45 +131,29 @@ func doctorSandboxPolicy(cfg config.SandboxConfig) sandbox.Policy {
|
||||
// lspServersCheck reports which language servers ZERO would use are present on
|
||||
// PATH. Missing servers are not a failure — ZERO degrades to text-only edits for
|
||||
// those languages — so the worst status is WARN, and each missing server gets an
|
||||
// actionable install command keyed by its binary name. Only the tier-1 servers
|
||||
// (lsp.CoreServerBinaries) drive the status: the long-tail servers configured
|
||||
// for breadth are listed informationally under missingOptional, because warning
|
||||
// about a missing zls on a machine with no Zig code would be permanent noise.
|
||||
// actionable install command keyed by its binary name.
|
||||
func lspServersCheck(lookup func(string) (string, error)) Check {
|
||||
if lookup == nil {
|
||||
lookup = exec.LookPath
|
||||
}
|
||||
core := map[string]bool{}
|
||||
for _, binary := range lsp.CoreServerBinaries() {
|
||||
core[binary] = true
|
||||
}
|
||||
present := map[string]any{}
|
||||
missing := map[string]any{}
|
||||
missingOptional := map[string]any{}
|
||||
for _, binary := range lsp.ServerBinaries() {
|
||||
if _, err := lookup(binary); err == nil {
|
||||
present[binary] = "on PATH"
|
||||
continue
|
||||
}
|
||||
if core[binary] {
|
||||
missing[binary] = lspRemedy(binary)
|
||||
} else {
|
||||
missingOptional[binary] = lspRemedy(binary)
|
||||
}
|
||||
}
|
||||
details := map[string]any{"present": present}
|
||||
if len(missingOptional) > 0 {
|
||||
details["missingOptional"] = missingOptional
|
||||
missing[binary] = lspRemedy(binary)
|
||||
}
|
||||
if len(missing) == 0 {
|
||||
message := "All core language servers are available on PATH."
|
||||
if len(missingOptional) > 0 {
|
||||
message = fmt.Sprintf("All core language servers are available on PATH (%d optional server(s) not installed).", len(missingOptional))
|
||||
}
|
||||
return check("lsp.servers", "LSP servers", StatusPass, message, details)
|
||||
return check("lsp.servers", "LSP servers", StatusPass, "All known language servers are available on PATH.", map[string]any{
|
||||
"present": present,
|
||||
})
|
||||
}
|
||||
details["missing"] = missing
|
||||
return check("lsp.servers", "LSP servers", StatusWarn, fmt.Sprintf("%d language server(s) missing from PATH; affected files degrade to text-only edits.", len(missing)), details)
|
||||
return check("lsp.servers", "LSP servers", StatusWarn, fmt.Sprintf("%d language server(s) missing from PATH; affected files degrade to text-only edits.", len(missing)), map[string]any{
|
||||
"present": present,
|
||||
"missing": missing,
|
||||
})
|
||||
}
|
||||
|
||||
// lspRemedy returns an actionable install command for a missing language-server
|
||||
|
||||
@@ -1,167 +0,0 @@
|
||||
// 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'
|
||||
}
|
||||
@@ -1,94 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
+15
-123
@@ -9,134 +9,26 @@ import (
|
||||
|
||||
// serverCommands maps a file extension to the language-server command (argv) ZERO
|
||||
// will spawn. The first element is the binary looked up on PATH; missing binaries
|
||||
// are not an error, the agent just degrades to text-only for that file. Every
|
||||
// command is the language's community-standard server invoked in stdio mode, so
|
||||
// a wrong guess can never spawn something surprising — only fail the PATH check.
|
||||
// are not an error, the agent just degrades to text-only for that file.
|
||||
var serverCommands = map[string][]string{
|
||||
".go": {"gopls", "serve"},
|
||||
".ts": {"typescript-language-server", "--stdio"},
|
||||
".tsx": {"typescript-language-server", "--stdio"},
|
||||
".mts": {"typescript-language-server", "--stdio"},
|
||||
".cts": {"typescript-language-server", "--stdio"},
|
||||
".js": {"typescript-language-server", "--stdio"},
|
||||
".jsx": {"typescript-language-server", "--stdio"},
|
||||
".mjs": {"typescript-language-server", "--stdio"},
|
||||
".cjs": {"typescript-language-server", "--stdio"},
|
||||
".py": {"pyright-langserver", "--stdio"},
|
||||
".pyi": {"pyright-langserver", "--stdio"},
|
||||
".rs": {"rust-analyzer"},
|
||||
".c": {"clangd"},
|
||||
".h": {"clangd"},
|
||||
".cpp": {"clangd"},
|
||||
".cc": {"clangd"},
|
||||
".cxx": {"clangd"},
|
||||
".hpp": {"clangd"},
|
||||
".java": {"jdtls"},
|
||||
".kt": {"kotlin-language-server"},
|
||||
".kts": {"kotlin-language-server"},
|
||||
".rb": {"ruby-lsp"},
|
||||
".php": {"intelephense", "--stdio"},
|
||||
".zig": {"zls"},
|
||||
".lua": {"lua-language-server"},
|
||||
".ex": {"elixir-ls"},
|
||||
".exs": {"elixir-ls"},
|
||||
".hs": {"haskell-language-server-wrapper", "--lsp"},
|
||||
".swift": {"sourcekit-lsp"},
|
||||
".ml": {"ocamllsp"},
|
||||
".mli": {"ocamllsp"},
|
||||
".scala": {"metals"},
|
||||
".clj": {"clojure-lsp"},
|
||||
".cljs": {"clojure-lsp"},
|
||||
".cljc": {"clojure-lsp"},
|
||||
".dart": {"dart", "language-server"},
|
||||
".gleam": {"gleam", "lsp"},
|
||||
".nix": {"nixd"},
|
||||
".tf": {"terraform-ls", "serve"},
|
||||
".sh": {"bash-language-server", "start"},
|
||||
".bash": {"bash-language-server", "start"},
|
||||
".yaml": {"yaml-language-server", "--stdio"},
|
||||
".yml": {"yaml-language-server", "--stdio"},
|
||||
".json": {"vscode-json-language-server", "--stdio"},
|
||||
".css": {"vscode-css-language-server", "--stdio"},
|
||||
".scss": {"vscode-css-language-server", "--stdio"},
|
||||
".less": {"vscode-css-language-server", "--stdio"},
|
||||
".html": {"vscode-html-language-server", "--stdio"},
|
||||
".svelte": {"svelteserver", "--stdio"},
|
||||
".vue": {"vue-language-server", "--stdio"},
|
||||
".astro": {"astro-ls", "--stdio"},
|
||||
".go": {"gopls", "serve"},
|
||||
".ts": {"typescript-language-server", "--stdio"},
|
||||
".tsx": {"typescript-language-server", "--stdio"},
|
||||
".js": {"typescript-language-server", "--stdio"},
|
||||
".jsx": {"typescript-language-server", "--stdio"},
|
||||
".py": {"pyright-langserver", "--stdio"},
|
||||
".rs": {"rust-analyzer"},
|
||||
}
|
||||
|
||||
// languageIDs maps a file extension to the LSP languageId used in didOpen.
|
||||
var languageIDs = map[string]string{
|
||||
".go": "go",
|
||||
".ts": "typescript",
|
||||
".tsx": "typescriptreact",
|
||||
".mts": "typescript",
|
||||
".cts": "typescript",
|
||||
".js": "javascript",
|
||||
".jsx": "javascriptreact",
|
||||
".mjs": "javascript",
|
||||
".cjs": "javascript",
|
||||
".py": "python",
|
||||
".pyi": "python",
|
||||
".rs": "rust",
|
||||
".c": "c",
|
||||
".h": "c",
|
||||
".cpp": "cpp",
|
||||
".cc": "cpp",
|
||||
".cxx": "cpp",
|
||||
".hpp": "cpp",
|
||||
".java": "java",
|
||||
".kt": "kotlin",
|
||||
".kts": "kotlin",
|
||||
".rb": "ruby",
|
||||
".php": "php",
|
||||
".zig": "zig",
|
||||
".lua": "lua",
|
||||
".ex": "elixir",
|
||||
".exs": "elixir",
|
||||
".hs": "haskell",
|
||||
".swift": "swift",
|
||||
".ml": "ocaml",
|
||||
".mli": "ocaml",
|
||||
".scala": "scala",
|
||||
".clj": "clojure",
|
||||
".cljs": "clojure",
|
||||
".cljc": "clojure",
|
||||
".dart": "dart",
|
||||
".gleam": "gleam",
|
||||
".nix": "nix",
|
||||
".tf": "terraform",
|
||||
".sh": "shellscript",
|
||||
".bash": "shellscript",
|
||||
".yaml": "yaml",
|
||||
".yml": "yaml",
|
||||
".json": "json",
|
||||
".css": "css",
|
||||
".scss": "scss",
|
||||
".less": "less",
|
||||
".html": "html",
|
||||
".svelte": "svelte",
|
||||
".vue": "vue",
|
||||
".astro": "astro",
|
||||
}
|
||||
|
||||
// coreServerBinaries are the tier-1 servers for the languages agents hit most;
|
||||
// `zero doctor` treats their absence as warn-worthy. The long tail of servers
|
||||
// configured above for breadth is reported informationally only — warning about
|
||||
// a missing zls on a machine with no Zig code would be permanent noise.
|
||||
var coreServerBinaries = []string{
|
||||
"gopls",
|
||||
"typescript-language-server",
|
||||
"pyright-langserver",
|
||||
"rust-analyzer",
|
||||
}
|
||||
|
||||
// CoreServerBinaries returns the tier-1 server binaries (sorted copy).
|
||||
func CoreServerBinaries() []string {
|
||||
binaries := append([]string(nil), coreServerBinaries...)
|
||||
sort.Strings(binaries)
|
||||
return binaries
|
||||
".go": "go",
|
||||
".ts": "typescript",
|
||||
".tsx": "typescriptreact",
|
||||
".js": "javascript",
|
||||
".jsx": "javascriptreact",
|
||||
".py": "python",
|
||||
".rs": "rust",
|
||||
}
|
||||
|
||||
// ServerBinaries returns the unique set of language-server binaries ZERO may
|
||||
|
||||
@@ -3,6 +3,7 @@ package modelregistry
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
const DefaultModelID = "gpt-4.1"
|
||||
@@ -42,8 +43,20 @@ type ListOptions struct {
|
||||
Capability ModelCapability
|
||||
}
|
||||
|
||||
func DefaultRegistry() (Registry, error) {
|
||||
// defaultRegistry builds the default registry exactly once and caches the result.
|
||||
// The default catalog is static, every accessor returns clones, and a built
|
||||
// Registry is never mutated (concurrent reads of its maps are safe), so a single
|
||||
// shared instance is correct — and it avoids re-cloning all entries and
|
||||
// recompiling every fuzzy-match regex on the per-frame TUI paths (title chip,
|
||||
// context gauge, pickers) that call DefaultRegistry each render.
|
||||
var defaultRegistry = sync.OnceValues(func() (Registry, error) {
|
||||
return NewRegistry(DefaultModelEntries())
|
||||
})
|
||||
|
||||
// DefaultRegistry returns the shared, immutable default model registry, building
|
||||
// it on first use.
|
||||
func DefaultRegistry() (Registry, error) {
|
||||
return defaultRegistry()
|
||||
}
|
||||
|
||||
func DefaultModelEntries() []ModelEntry {
|
||||
@@ -63,10 +76,6 @@ func DefaultModelEntries() []ModelEntry {
|
||||
googleModel("gemini-2.5-flash-lite", "Gemini 2.5 Flash-Lite", "gemini-2.5-flash-lite", ModelStatusActive, []string{"google:gemini-2.5-flash-lite", "gemini-flash-lite"}, ContextLimits{ContextWindow: 1_048_576, MaxOutputTokens: 65_536}, ModelCost{InputPerMillion: 0.1, CachedInputPerMillion: 0.01, OutputPerMillion: 0.4}, []ModelCapability{ModelCapabilityVision, ModelCapabilityJSONMode, ModelCapabilityReasoning, ModelCapabilityLongContext}, standardReasoningEfforts(), "Google low-cost Flash model for background routing and summaries."),
|
||||
}
|
||||
decorateModelDepth(entries)
|
||||
// Overlay volatile facts (context limits, base pricing) from a cached
|
||||
// models.dev snapshot when one is present and fresh — see modelsdev.go.
|
||||
// Identity fields (ids, aliases, patterns, deprecations) stay curated.
|
||||
entries = applyModelsDevOverrides(entries, cachedModelsDevProviders())
|
||||
return cloneModelEntries(entries)
|
||||
}
|
||||
|
||||
@@ -80,11 +89,13 @@ func decorateModelDepth(entries []ModelEntry) {
|
||||
patterns []string
|
||||
deprecation *DeprecationRule
|
||||
upgradeTarget string
|
||||
fastVariant string
|
||||
}{
|
||||
"claude-sonnet-4.5": {
|
||||
defaultEffort: ReasoningEffortMedium,
|
||||
patterns: []string{`(?i)^sonnet[^a-z0-9]*4[.\s]?5$`},
|
||||
upgradeTarget: "claude-opus-4.1",
|
||||
fastVariant: "claude-haiku-4.5",
|
||||
},
|
||||
"claude-opus-4.1": {
|
||||
defaultEffort: ReasoningEffortHigh,
|
||||
@@ -98,6 +109,7 @@ func decorateModelDepth(entries []ModelEntry) {
|
||||
"gemini-2.5-pro": {
|
||||
defaultEffort: ReasoningEffortMedium,
|
||||
patterns: []string{`(?i)^gemini[^a-z0-9]*pro$`},
|
||||
fastVariant: "gemini-2.5-flash",
|
||||
},
|
||||
"gemini-2.5-flash": {
|
||||
defaultEffort: ReasoningEffortLow,
|
||||
@@ -106,6 +118,12 @@ func decorateModelDepth(entries []ModelEntry) {
|
||||
"gemini-2.5-flash-lite": {
|
||||
upgradeTarget: "gemini-2.5-flash",
|
||||
},
|
||||
"gpt-4.1": {
|
||||
fastVariant: "gpt-4.1-mini",
|
||||
},
|
||||
"gpt-4o": {
|
||||
fastVariant: "gpt-4o-mini",
|
||||
},
|
||||
"gpt-4.1-mini": {
|
||||
upgradeTarget: "gpt-4.1",
|
||||
},
|
||||
@@ -147,6 +165,9 @@ func decorateModelDepth(entries []ModelEntry) {
|
||||
if extra.upgradeTarget != "" {
|
||||
entries[index].UpgradeTargetID = extra.upgradeTarget
|
||||
}
|
||||
if extra.fastVariant != "" {
|
||||
entries[index].FastVariantID = extra.fastVariant
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -110,7 +110,14 @@ type ModelEntry struct {
|
||||
// during mid-run escalation. Empty means no escalation target (e.g. top-tier
|
||||
// models). Resolved and availability-checked by Registry.UpgradeTarget.
|
||||
UpgradeTargetID string
|
||||
Description string
|
||||
// FastVariantID, when set, names a genuinely faster same-family model this one
|
||||
// can drop down to (e.g. claude-sonnet-4.5 -> claude-haiku-4.5). Empty means no
|
||||
// fast variant. It is the single source of truth for the /fast command and the
|
||||
// "f" fast-mode marker: Registry.FastVariant walks it forward and BaseVariant
|
||||
// walks it in reverse, so a model is "in fast mode" iff another model names it
|
||||
// here. Resolved and availability-checked like UpgradeTargetID.
|
||||
FastVariantID string
|
||||
Description string
|
||||
}
|
||||
|
||||
// DeprecationRule describes how a deprecated model is phased out and what to use
|
||||
@@ -368,6 +375,18 @@ func NewRegistry(entries []ModelEntry) (Registry, error) {
|
||||
return Registry{}, fmt.Errorf("model %q upgrade target %q does not resolve to a known model", entry.ID, targetID)
|
||||
}
|
||||
}
|
||||
// A non-empty FastVariantID must resolve to a known model, mirroring the
|
||||
// UpgradeTargetID check — otherwise /fast and the "f" fast-mode marker would
|
||||
// silently vanish at runtime on a catalog typo instead of failing loudly here.
|
||||
for _, entry := range registry.models {
|
||||
fastID := strings.TrimSpace(entry.FastVariantID)
|
||||
if fastID == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := registry.Get(fastID); !ok {
|
||||
return Registry{}, fmt.Errorf("model %q fast variant %q does not resolve to a known model", entry.ID, fastID)
|
||||
}
|
||||
}
|
||||
return registry, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -1,269 +0,0 @@
|
||||
package modelregistry
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Live models.dev overlay for the curated catalog. The hand-maintained
|
||||
// DefaultModelEntries list is the source of truth for identity — ids, aliases,
|
||||
// match patterns, deprecations, escalation targets — but its VOLATILE facts
|
||||
// (context window, max output tokens, per-million pricing) go stale between
|
||||
// releases. When a cached snapshot of https://models.dev/api.json is present,
|
||||
// those fields are refreshed from it at registry construction; everything else
|
||||
// stays curated. The overlay never adds models (an auto-added entry would lack
|
||||
// aliases, match patterns, and provider wiring) and never touches the network
|
||||
// on the registry hot path — fetching happens only in the explicit background
|
||||
// refresh, cached to disk with a TTL.
|
||||
|
||||
const (
|
||||
modelsDevDefaultURL = "https://models.dev/api.json"
|
||||
// modelsDevRefreshAfter is how old the cache may get before a background
|
||||
// refresh re-fetches it.
|
||||
modelsDevRefreshAfter = 24 * time.Hour
|
||||
// modelsDevMaxAge is the oldest cache still applied as an overlay. Beyond
|
||||
// this the curated catalog (updated with the binary) is likely fresher than
|
||||
// the snapshot, so a stale file is ignored rather than trusted.
|
||||
modelsDevMaxAge = 7 * 24 * time.Hour
|
||||
modelsDevFetchLimit = 32 << 20 // 32MiB guard on the response body
|
||||
modelsDevFetchWindow = 15 * time.Second
|
||||
)
|
||||
|
||||
// modelsDevModel is the subset of a models.dev model record the overlay uses.
|
||||
type modelsDevModel struct {
|
||||
Limit struct {
|
||||
Context int `json:"context"`
|
||||
Output int `json:"output"`
|
||||
} `json:"limit"`
|
||||
Cost struct {
|
||||
Input float64 `json:"input"`
|
||||
Output float64 `json:"output"`
|
||||
CacheRead float64 `json:"cache_read"`
|
||||
CacheWrite float64 `json:"cache_write"`
|
||||
} `json:"cost"`
|
||||
}
|
||||
|
||||
// modelsDevProvider matches one provider object in api.json.
|
||||
type modelsDevProvider struct {
|
||||
Models map[string]modelsDevModel `json:"models"`
|
||||
}
|
||||
|
||||
// parseModelsDev decodes an api.json document into provider-slug -> api-model
|
||||
// -> record. api.json is a top-level object keyed by provider id.
|
||||
func parseModelsDev(data []byte) (map[string]map[string]modelsDevModel, error) {
|
||||
var doc map[string]modelsDevProvider
|
||||
if err := json.Unmarshal(data, &doc); err != nil {
|
||||
return nil, fmt.Errorf("modelregistry: parse models.dev: %w", err)
|
||||
}
|
||||
providers := make(map[string]map[string]modelsDevModel, len(doc))
|
||||
for slug, provider := range doc {
|
||||
if len(provider.Models) == 0 {
|
||||
continue
|
||||
}
|
||||
providers[slug] = provider.Models
|
||||
}
|
||||
if len(providers) == 0 {
|
||||
return nil, fmt.Errorf("modelregistry: models.dev document has no providers")
|
||||
}
|
||||
return providers, nil
|
||||
}
|
||||
|
||||
// modelsDevSlugs maps an entry's provider kind to the models.dev provider ids
|
||||
// worth checking. Only first-party slugs are used: router entries on models.dev
|
||||
// carry generic or stale numbers.
|
||||
func modelsDevSlugs(kind ProviderKind) []string {
|
||||
switch kind {
|
||||
case ProviderAnthropic:
|
||||
return []string{"anthropic"}
|
||||
case ProviderOpenAI:
|
||||
return []string{"openai"}
|
||||
case ProviderGoogle:
|
||||
return []string{"google", "google-vertex"}
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// applyModelsDevOverrides refreshes each curated entry's context limits and
|
||||
// base pricing from the snapshot, when the snapshot knows the model. Tiered
|
||||
// pricing is never touched: models.dev has no tier data, and mixing a live
|
||||
// base rate with curated tiers would misprice the tier boundaries.
|
||||
func applyModelsDevOverrides(entries []ModelEntry, providers map[string]map[string]modelsDevModel) []ModelEntry {
|
||||
if len(providers) == 0 {
|
||||
return entries
|
||||
}
|
||||
for i := range entries {
|
||||
entry := &entries[i]
|
||||
var record modelsDevModel
|
||||
found := false
|
||||
for _, slug := range modelsDevSlugs(entry.Provider) {
|
||||
if models, ok := providers[slug]; ok {
|
||||
if candidate, ok := models[strings.TrimSpace(entry.APIModel)]; ok {
|
||||
record = candidate
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
continue
|
||||
}
|
||||
if record.Limit.Context > 0 {
|
||||
entry.ContextLimits.ContextWindow = record.Limit.Context
|
||||
}
|
||||
if record.Limit.Output > 0 {
|
||||
entry.ContextLimits.MaxOutputTokens = record.Limit.Output
|
||||
}
|
||||
if len(entry.Cost.Tiers) == 0 && record.Cost.Input > 0 && record.Cost.Output > 0 {
|
||||
entry.Cost.InputPerMillion = record.Cost.Input
|
||||
entry.Cost.OutputPerMillion = record.Cost.Output
|
||||
if record.Cost.CacheRead > 0 {
|
||||
entry.Cost.CachedInputPerMillion = record.Cost.CacheRead
|
||||
}
|
||||
if record.Cost.CacheWrite > 0 {
|
||||
entry.Cost.CacheWritePerMillion = record.Cost.CacheWrite
|
||||
}
|
||||
entry.Cost.Source = "models.dev/api.json (cached)"
|
||||
}
|
||||
}
|
||||
return entries
|
||||
}
|
||||
|
||||
// modelsDevCachePath returns the on-disk cache location. ZERO_MODELS_CACHE_PATH
|
||||
// overrides it (used by tests and unusual setups).
|
||||
func modelsDevCachePath() (string, error) {
|
||||
if override := strings.TrimSpace(os.Getenv("ZERO_MODELS_CACHE_PATH")); override != "" {
|
||||
return override, nil
|
||||
}
|
||||
base, err := os.UserCacheDir()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return filepath.Join(base, "zero", "modelsdev.json"), nil
|
||||
}
|
||||
|
||||
var (
|
||||
modelsDevOnce sync.Once
|
||||
modelsDevCached map[string]map[string]modelsDevModel
|
||||
modelsDevEnabled atomic.Bool
|
||||
)
|
||||
|
||||
// EnableModelsDevOverlay opts the process into applying the cached models.dev
|
||||
// snapshot on top of the curated catalog. The CLI entrypoint calls it; library
|
||||
// consumers and tests that never do get the curated catalog byte-identical to
|
||||
// before, so hermetic tests can't be perturbed by a cache file on the machine.
|
||||
// ZERO_DISABLE_MODELS_FETCH disables both the overlay and the fetch.
|
||||
func EnableModelsDevOverlay() {
|
||||
modelsDevEnabled.Store(true)
|
||||
}
|
||||
|
||||
// cachedModelsDevProviders loads the cached snapshot once per process. Not
|
||||
// enabled, missing, stale (> modelsDevMaxAge), or malformed all yield nil and
|
||||
// the curated catalog is used untouched. Read once deliberately:
|
||||
// DefaultRegistry is called on hot paths (pickers, cost views) and must not
|
||||
// re-stat the file every time; a background refresh benefits the NEXT process.
|
||||
func cachedModelsDevProviders() map[string]map[string]modelsDevModel {
|
||||
if !modelsDevEnabled.Load() || strings.TrimSpace(os.Getenv("ZERO_DISABLE_MODELS_FETCH")) != "" {
|
||||
return nil
|
||||
}
|
||||
modelsDevOnce.Do(func() {
|
||||
path, err := modelsDevCachePath()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
info, err := os.Stat(path)
|
||||
if err != nil || time.Since(info.ModTime()) > modelsDevMaxAge {
|
||||
return
|
||||
}
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
providers, err := parseModelsDev(data)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
modelsDevCached = providers
|
||||
})
|
||||
return modelsDevCached
|
||||
}
|
||||
|
||||
// resetModelsDevCacheForTest clears the process-level cache memoization and
|
||||
// disables the overlay.
|
||||
func resetModelsDevCacheForTest() {
|
||||
modelsDevOnce = sync.Once{}
|
||||
modelsDevCached = nil
|
||||
modelsDevEnabled.Store(false)
|
||||
}
|
||||
|
||||
// RefreshModelsDevCache fetches models.dev/api.json into the on-disk cache
|
||||
// when the cache is missing or older than modelsDevRefreshAfter. It is safe to
|
||||
// call fire-and-forget from startup (use a goroutine); it never affects the
|
||||
// current process's registry (see cachedModelsDevProviders). Disabled entirely
|
||||
// by ZERO_DISABLE_MODELS_FETCH. The URL can be overridden with ZERO_MODELS_URL.
|
||||
func RefreshModelsDevCache(ctx context.Context) error {
|
||||
if strings.TrimSpace(os.Getenv("ZERO_DISABLE_MODELS_FETCH")) != "" {
|
||||
return nil
|
||||
}
|
||||
path, err := modelsDevCachePath()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if info, err := os.Stat(path); err == nil && time.Since(info.ModTime()) < modelsDevRefreshAfter {
|
||||
return nil
|
||||
}
|
||||
|
||||
url := strings.TrimSpace(os.Getenv("ZERO_MODELS_URL"))
|
||||
if url == "" {
|
||||
url = modelsDevDefaultURL
|
||||
}
|
||||
fetchCtx, cancel := context.WithTimeout(ctx, modelsDevFetchWindow)
|
||||
defer cancel()
|
||||
request, err := http.NewRequestWithContext(fetchCtx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
request.Header.Set("User-Agent", "zero-models-refresh/0.1")
|
||||
response, err := http.DefaultClient.Do(request)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode != http.StatusOK {
|
||||
return fmt.Errorf("modelregistry: models.dev fetch: HTTP %d", response.StatusCode)
|
||||
}
|
||||
data, err := io.ReadAll(io.LimitReader(response.Body, modelsDevFetchLimit))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Validate before persisting: a bad body must never clobber a good cache.
|
||||
if _, err := parseModelsDev(data); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
|
||||
return err
|
||||
}
|
||||
temp, err := os.CreateTemp(filepath.Dir(path), "modelsdev-*.json")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := temp.Write(data); err != nil {
|
||||
_ = temp.Close()
|
||||
_ = os.Remove(temp.Name())
|
||||
return err
|
||||
}
|
||||
if err := temp.Close(); err != nil {
|
||||
_ = os.Remove(temp.Name())
|
||||
return err
|
||||
}
|
||||
return os.Rename(temp.Name(), path)
|
||||
}
|
||||
@@ -1,228 +0,0 @@
|
||||
package modelregistry
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
const sampleModelsDev = `{
|
||||
"anthropic": {
|
||||
"id": "anthropic",
|
||||
"models": {
|
||||
"claude-sonnet-4-5-20250929": {
|
||||
"limit": {"context": 1000000, "output": 64000},
|
||||
"cost": {"input": 3.5, "output": 17.5, "cache_read": 0.35, "cache_write": 4.4}
|
||||
}
|
||||
}
|
||||
},
|
||||
"google": {
|
||||
"id": "google",
|
||||
"models": {
|
||||
"gemini-2.5-pro": {
|
||||
"limit": {"context": 2097152, "output": 65536},
|
||||
"cost": {"input": 9.99, "output": 9.99}
|
||||
}
|
||||
}
|
||||
}
|
||||
}`
|
||||
|
||||
func TestParseModelsDev(t *testing.T) {
|
||||
providers, err := parseModelsDev([]byte(sampleModelsDev))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
record, ok := providers["anthropic"]["claude-sonnet-4-5-20250929"]
|
||||
if !ok {
|
||||
t.Fatal("expected anthropic sonnet record")
|
||||
}
|
||||
if record.Limit.Context != 1_000_000 || record.Cost.Input != 3.5 {
|
||||
t.Fatalf("unexpected record: %+v", record)
|
||||
}
|
||||
if _, err := parseModelsDev([]byte(`{}`)); err == nil {
|
||||
t.Fatal("empty document must be rejected")
|
||||
}
|
||||
if _, err := parseModelsDev([]byte(`not json`)); err == nil {
|
||||
t.Fatal("malformed document must be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyModelsDevOverrides(t *testing.T) {
|
||||
// Point the cache at a non-existent file so DefaultModelEntries returns the
|
||||
// pure curated catalog, then apply the sample snapshot explicitly.
|
||||
t.Setenv("ZERO_MODELS_CACHE_PATH", filepath.Join(t.TempDir(), "absent.json"))
|
||||
resetModelsDevCacheForTest()
|
||||
t.Cleanup(resetModelsDevCacheForTest)
|
||||
|
||||
providers, err := parseModelsDev([]byte(sampleModelsDev))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
entries := applyModelsDevOverrides(DefaultModelEntries(), providers)
|
||||
|
||||
var sonnet, geminiPro, opus ModelEntry
|
||||
for _, entry := range entries {
|
||||
switch entry.ID {
|
||||
case "claude-sonnet-4.5":
|
||||
sonnet = entry
|
||||
case "gemini-2.5-pro":
|
||||
geminiPro = entry
|
||||
case "claude-opus-4.1":
|
||||
opus = entry
|
||||
}
|
||||
}
|
||||
|
||||
// Known model: limits and base pricing refreshed from the snapshot.
|
||||
if sonnet.ContextLimits.ContextWindow != 1_000_000 || sonnet.ContextLimits.MaxOutputTokens != 64_000 {
|
||||
t.Fatalf("sonnet limits not overridden: %+v", sonnet.ContextLimits)
|
||||
}
|
||||
if sonnet.Cost.InputPerMillion != 3.5 || sonnet.Cost.OutputPerMillion != 17.5 || sonnet.Cost.CacheWritePerMillion != 4.4 {
|
||||
t.Fatalf("sonnet cost not overridden: %+v", sonnet.Cost)
|
||||
}
|
||||
if sonnet.Cost.Source != "models.dev/api.json (cached)" {
|
||||
t.Fatalf("sonnet cost source not marked: %q", sonnet.Cost.Source)
|
||||
}
|
||||
|
||||
// Tiered pricing is curated: limits refresh, cost must NOT (gemini-2.5-pro
|
||||
// has curated tiers and the snapshot's flat 9.99 would misprice them).
|
||||
if geminiPro.ContextLimits.ContextWindow != 2_097_152 {
|
||||
t.Fatalf("gemini limits not overridden: %+v", geminiPro.ContextLimits)
|
||||
}
|
||||
if geminiPro.Cost.InputPerMillion == 9.99 || len(geminiPro.Cost.Tiers) == 0 {
|
||||
t.Fatalf("tiered cost must stay curated: %+v", geminiPro.Cost)
|
||||
}
|
||||
|
||||
// Model absent from the snapshot: untouched.
|
||||
if opus.ContextLimits.ContextWindow != 200_000 || opus.Cost.InputPerMillion != 15 {
|
||||
t.Fatalf("opus must be untouched: %+v %+v", opus.ContextLimits, opus.Cost)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefreshModelsDevCacheFetchesAndCaches(t *testing.T) {
|
||||
hits := 0
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
hits++
|
||||
_, _ = w.Write([]byte(sampleModelsDev))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
cachePath := filepath.Join(t.TempDir(), "modelsdev.json")
|
||||
t.Setenv("ZERO_MODELS_CACHE_PATH", cachePath)
|
||||
t.Setenv("ZERO_MODELS_URL", server.URL)
|
||||
t.Setenv("ZERO_DISABLE_MODELS_FETCH", "")
|
||||
|
||||
if err := RefreshModelsDevCache(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if hits != 1 {
|
||||
t.Fatalf("expected 1 fetch, got %d", hits)
|
||||
}
|
||||
if _, err := os.Stat(cachePath); err != nil {
|
||||
t.Fatalf("cache file missing: %v", err)
|
||||
}
|
||||
// Fresh cache: second call must not re-fetch.
|
||||
if err := RefreshModelsDevCache(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if hits != 1 {
|
||||
t.Fatalf("fresh cache must skip the fetch, got %d hits", hits)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefreshModelsDevCacheRejectsBadBodyWithoutClobbering(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = w.Write([]byte("not json"))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
cachePath := filepath.Join(t.TempDir(), "modelsdev.json")
|
||||
if err := os.WriteFile(cachePath, []byte(sampleModelsDev), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
stale := time.Now().Add(-48 * time.Hour)
|
||||
if err := os.Chtimes(cachePath, stale, stale); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Setenv("ZERO_MODELS_CACHE_PATH", cachePath)
|
||||
t.Setenv("ZERO_MODELS_URL", server.URL)
|
||||
|
||||
if err := RefreshModelsDevCache(t.Context()); err == nil {
|
||||
t.Fatal("bad body must return an error")
|
||||
}
|
||||
content, err := os.ReadFile(cachePath)
|
||||
if err != nil || string(content) != sampleModelsDev {
|
||||
t.Fatal("bad fetch must not clobber the existing cache")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefreshModelsDevCacheDisabledByEnv(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
t.Fatal("fetch must not happen when disabled")
|
||||
}))
|
||||
defer server.Close()
|
||||
t.Setenv("ZERO_MODELS_CACHE_PATH", filepath.Join(t.TempDir(), "modelsdev.json"))
|
||||
t.Setenv("ZERO_MODELS_URL", server.URL)
|
||||
t.Setenv("ZERO_DISABLE_MODELS_FETCH", "1")
|
||||
if err := RefreshModelsDevCache(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCachedModelsDevProvidersIgnoresStaleCache(t *testing.T) {
|
||||
cachePath := filepath.Join(t.TempDir(), "modelsdev.json")
|
||||
if err := os.WriteFile(cachePath, []byte(sampleModelsDev), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
stale := time.Now().Add(-modelsDevMaxAge - time.Hour)
|
||||
if err := os.Chtimes(cachePath, stale, stale); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Setenv("ZERO_MODELS_CACHE_PATH", cachePath)
|
||||
resetModelsDevCacheForTest()
|
||||
t.Cleanup(resetModelsDevCacheForTest)
|
||||
EnableModelsDevOverlay()
|
||||
|
||||
if providers := cachedModelsDevProviders(); providers != nil {
|
||||
t.Fatal("stale cache must be ignored")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCachedModelsDevProvidersRequiresOptIn(t *testing.T) {
|
||||
cachePath := filepath.Join(t.TempDir(), "modelsdev.json")
|
||||
if err := os.WriteFile(cachePath, []byte(sampleModelsDev), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Setenv("ZERO_MODELS_CACHE_PATH", cachePath)
|
||||
resetModelsDevCacheForTest()
|
||||
t.Cleanup(resetModelsDevCacheForTest)
|
||||
|
||||
// Without EnableModelsDevOverlay a fresh, valid cache must still be ignored:
|
||||
// library consumers and hermetic tests get the pure curated catalog.
|
||||
if providers := cachedModelsDevProviders(); providers != nil {
|
||||
t.Fatal("overlay must be opt-in")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultModelEntriesAppliesFreshCache(t *testing.T) {
|
||||
cachePath := filepath.Join(t.TempDir(), "modelsdev.json")
|
||||
if err := os.WriteFile(cachePath, []byte(sampleModelsDev), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Setenv("ZERO_MODELS_CACHE_PATH", cachePath)
|
||||
resetModelsDevCacheForTest()
|
||||
t.Cleanup(resetModelsDevCacheForTest)
|
||||
EnableModelsDevOverlay()
|
||||
|
||||
for _, entry := range DefaultModelEntries() {
|
||||
if entry.ID == "claude-sonnet-4.5" {
|
||||
if entry.ContextLimits.ContextWindow != 1_000_000 {
|
||||
t.Fatalf("fresh cache must overlay the registry: %+v", entry.ContextLimits)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Fatal("claude-sonnet-4.5 not found")
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
package modelregistry
|
||||
|
||||
import "strings"
|
||||
|
||||
// The base<->fast model mapping is data, not logic: it lives entirely in each
|
||||
// ModelEntry's FastVariantID field (populated in the catalog, validated in
|
||||
// NewRegistry). Adding a pair is a one-line catalog change — never a code change
|
||||
// here. Both directions derive from that one field: FastVariant walks a model to
|
||||
// the fast id it names, BaseVariant scans for the model that names id as its fast
|
||||
// variant. This is why "fast mode" is encoded by model identity and needs no flag.
|
||||
//
|
||||
// Only models with an unambiguous 1:1 fast counterpart carry a FastVariantID. A
|
||||
// model with no clean counterpart (e.g. a top-tier Opus, where both Sonnet and
|
||||
// Haiku would be candidates) is intentionally left unset, so /fast reports "no
|
||||
// fast mode available" for it rather than guessing.
|
||||
|
||||
// FastVariant resolves id to its configured faster same-family counterpart. It
|
||||
// resolves id to a concrete entry, reads that entry's FastVariantID, and returns
|
||||
// the target only when it resolves and is not deprecated. Returns (_, false) when
|
||||
// id is unknown, has no fast variant, or the variant is unavailable/deprecated.
|
||||
// Mirrors Registry.UpgradeTarget.
|
||||
func (registry Registry) FastVariant(id string) (ModelEntry, bool) {
|
||||
source, ok := registry.Resolve(id)
|
||||
if !ok {
|
||||
return ModelEntry{}, false
|
||||
}
|
||||
fastID := strings.TrimSpace(source.FastVariantID)
|
||||
if fastID == "" {
|
||||
return ModelEntry{}, false
|
||||
}
|
||||
target, ok := registry.Resolve(fastID)
|
||||
if !ok || target.Status == ModelStatusDeprecated {
|
||||
return ModelEntry{}, false
|
||||
}
|
||||
return target, true
|
||||
}
|
||||
|
||||
// BaseVariant resolves id to the base model that declares id as its fast variant
|
||||
// — the reverse of FastVariant. A non-empty result therefore means id IS a fast
|
||||
// model: this is exactly how the /fast command derives "currently in fast mode"
|
||||
// without storing any flag. Same resolution and availability rules as FastVariant.
|
||||
func (registry Registry) BaseVariant(id string) (ModelEntry, bool) {
|
||||
source, ok := registry.Resolve(id)
|
||||
if !ok {
|
||||
return ModelEntry{}, false
|
||||
}
|
||||
for _, entry := range registry.models {
|
||||
fastID := strings.TrimSpace(entry.FastVariantID)
|
||||
if fastID == "" {
|
||||
continue
|
||||
}
|
||||
target, ok := registry.Resolve(fastID)
|
||||
if !ok || target.ID != source.ID {
|
||||
continue
|
||||
}
|
||||
if entry.Status == ModelStatusDeprecated {
|
||||
return ModelEntry{}, false
|
||||
}
|
||||
base, ok := registry.Resolve(entry.ID)
|
||||
if !ok {
|
||||
return ModelEntry{}, false
|
||||
}
|
||||
return base, true
|
||||
}
|
||||
return ModelEntry{}, false
|
||||
}
|
||||
|
||||
// HasFastVariant reports whether id has an available fast variant — i.e. whether
|
||||
// `/fast on` would switch to a faster model from here. It is the predicate the
|
||||
// TUI uses to show the "f" fast-mode marker on a model (in the model/provider
|
||||
// pickers and the active-model indicator).
|
||||
func (registry Registry) HasFastVariant(id string) bool {
|
||||
_, ok := registry.FastVariant(id)
|
||||
return ok
|
||||
}
|
||||
@@ -0,0 +1,145 @@
|
||||
package modelregistry
|
||||
|
||||
import "testing"
|
||||
|
||||
func testRegistry(t *testing.T) Registry {
|
||||
t.Helper()
|
||||
registry, err := DefaultRegistry()
|
||||
if err != nil {
|
||||
t.Fatalf("DefaultRegistry() error = %v", err)
|
||||
}
|
||||
return registry
|
||||
}
|
||||
|
||||
// expectedFastPairs is the known-good base->fast mapping. It is a literal
|
||||
// regression guard: if the catalog silently drops or mis-wires a FastVariantID,
|
||||
// the forward/reverse assertions below fail loudly rather than making /fast a
|
||||
// quiet no-op. Keep in sync with the FastVariantID fields in catalog.go.
|
||||
var expectedFastPairs = map[string]string{
|
||||
"claude-sonnet-4.5": "claude-haiku-4.5",
|
||||
"gpt-4.1": "gpt-4.1-mini",
|
||||
"gpt-4o": "gpt-4o-mini",
|
||||
"gemini-2.5-pro": "gemini-2.5-flash",
|
||||
}
|
||||
|
||||
// Every FastVariantID declared in the catalog must resolve to a real, non-self,
|
||||
// non-deprecated model. NewRegistry already rejects an unresolvable id at build
|
||||
// time; this additionally guards the runtime availability contract that /fast and
|
||||
// the "f" marker depend on.
|
||||
func TestCatalogFastVariantsResolve(t *testing.T) {
|
||||
registry := testRegistry(t)
|
||||
for _, entry := range registry.List(ListOptions{IncludeDeprecated: true}) {
|
||||
if entry.FastVariantID == "" {
|
||||
continue
|
||||
}
|
||||
target, ok := registry.Resolve(entry.FastVariantID)
|
||||
if !ok {
|
||||
t.Errorf("model %q fast variant %q does not resolve", entry.ID, entry.FastVariantID)
|
||||
continue
|
||||
}
|
||||
if target.ID == entry.ID {
|
||||
t.Errorf("model %q names itself as its fast variant", entry.ID)
|
||||
}
|
||||
if target.Status == ModelStatusDeprecated {
|
||||
t.Errorf("model %q fast variant %q is deprecated", entry.ID, entry.FastVariantID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// No model may be both a base and a fast, or the two directions become ambiguous
|
||||
// (a model would be "currently fast" AND advertise a fast variant). Derived from
|
||||
// the registry so a bad catalog edit — e.g. chaining mini->nano — is caught.
|
||||
func TestFastVariantPairsUnambiguous(t *testing.T) {
|
||||
registry := testRegistry(t)
|
||||
role := map[string]string{} // canonical id -> "base" | "fast"
|
||||
assign := func(id, r string) {
|
||||
entry, ok := registry.Resolve(id)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if prev, seen := role[entry.ID]; seen && prev != r {
|
||||
t.Fatalf("model %q is declared as both a base and a fast variant", entry.ID)
|
||||
}
|
||||
role[entry.ID] = r
|
||||
}
|
||||
for _, entry := range registry.List(ListOptions{IncludeDeprecated: true}) {
|
||||
if entry.FastVariantID == "" {
|
||||
continue
|
||||
}
|
||||
assign(entry.ID, "base")
|
||||
assign(entry.FastVariantID, "fast")
|
||||
}
|
||||
}
|
||||
|
||||
// The known pairs must be present and wired in both directions, and neither end
|
||||
// may leak into the other role (a base is not itself fast; a fast has no fast).
|
||||
func TestFastVariantAndBaseVariant(t *testing.T) {
|
||||
registry := testRegistry(t)
|
||||
|
||||
for base, fast := range expectedFastPairs {
|
||||
// base -> fast (forward), and the fast id has NO fast variant of its own.
|
||||
if got, ok := registry.FastVariant(base); !ok {
|
||||
t.Errorf("FastVariant(%q) = not found, want the fast variant", base)
|
||||
} else if canon, _ := registry.Resolve(fast); got.ID != canon.ID {
|
||||
t.Errorf("FastVariant(%q).ID = %q, want %q", base, got.ID, canon.ID)
|
||||
}
|
||||
if _, ok := registry.FastVariant(fast); ok {
|
||||
t.Errorf("FastVariant(%q) resolved, but a fast model should have no fast variant", fast)
|
||||
}
|
||||
|
||||
// fast -> base (reverse), and the base id is NOT itself a fast variant.
|
||||
if got, ok := registry.BaseVariant(fast); !ok {
|
||||
t.Errorf("BaseVariant(%q) = not found, want the base variant", fast)
|
||||
} else if canon, _ := registry.Resolve(base); got.ID != canon.ID {
|
||||
t.Errorf("BaseVariant(%q).ID = %q, want %q", fast, got.ID, canon.ID)
|
||||
}
|
||||
if _, ok := registry.BaseVariant(base); ok {
|
||||
t.Errorf("BaseVariant(%q) resolved, but a base model is not a fast variant", base)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// HasFastVariant is the TUI marker predicate: true only for base models with an
|
||||
// available fast variant, false for the fast models themselves, unpaired models,
|
||||
// and unknown ids.
|
||||
func TestHasFastVariant(t *testing.T) {
|
||||
registry := testRegistry(t)
|
||||
for base, fast := range expectedFastPairs {
|
||||
if !registry.HasFastVariant(base) {
|
||||
t.Errorf("HasFastVariant(%q) = false, want true", base)
|
||||
}
|
||||
if registry.HasFastVariant(fast) {
|
||||
t.Errorf("HasFastVariant(%q) = true, want false (a fast model has no fast variant)", fast)
|
||||
}
|
||||
}
|
||||
for _, id := range []string{"claude-opus-4.1", "definitely-not-a-real-model", ""} {
|
||||
if registry.HasFastVariant(id) {
|
||||
t.Errorf("HasFastVariant(%q) = true, want false", id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFastVariantAbsentAndUnknown(t *testing.T) {
|
||||
registry := testRegistry(t)
|
||||
|
||||
// A real model with no configured fast/base variant reports none in both
|
||||
// directions (a valid, handled case — /fast will say "no fast mode available").
|
||||
const noVariant = "claude-opus-4.1"
|
||||
if _, ok := registry.Resolve(noVariant); !ok {
|
||||
t.Skipf("%q not in registry; skipping no-variant case", noVariant)
|
||||
}
|
||||
if _, ok := registry.FastVariant(noVariant); ok {
|
||||
t.Errorf("FastVariant(%q) resolved, want none", noVariant)
|
||||
}
|
||||
if _, ok := registry.BaseVariant(noVariant); ok {
|
||||
t.Errorf("BaseVariant(%q) resolved, want none", noVariant)
|
||||
}
|
||||
|
||||
// An unknown id resolves to nothing in either direction.
|
||||
if _, ok := registry.FastVariant("definitely-not-a-real-model"); ok {
|
||||
t.Error("FastVariant(unknown) resolved, want none")
|
||||
}
|
||||
if _, ok := registry.BaseVariant("definitely-not-a-real-model"); ok {
|
||||
t.Error("BaseVariant(unknown) resolved, want none")
|
||||
}
|
||||
}
|
||||
@@ -1,96 +0,0 @@
|
||||
package anthropic
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// The final two messages must carry a cache_control breakpoint on their last
|
||||
// cacheable block so the conversation transcript — not just system prompt and
|
||||
// tools — is a prompt-cache hit on the next turn. Earlier messages must stay
|
||||
// unmarked (Anthropic caps breakpoints at 4 per request: system, tools, and
|
||||
// these two).
|
||||
func TestAnthropicRequestMarksLastTwoMessagesForCaching(t *testing.T) {
|
||||
var gotBody map[string]any
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if err := json.NewDecoder(r.Body).Decode(&gotBody); err != nil {
|
||||
t.Fatalf("decode request body: %v", err)
|
||||
}
|
||||
writeSSEEvent(w, "message_stop", `{"type":"message_stop"}`)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
provider, err := New(Options{BaseURL: server.URL + "/", Model: "claude-test"})
|
||||
if err != nil {
|
||||
t.Fatalf("New returned error: %v", err)
|
||||
}
|
||||
stream, err := provider.StreamCompletion(context.Background(), zeroruntime.CompletionRequest{
|
||||
Messages: []zeroruntime.Message{
|
||||
{Role: zeroruntime.MessageRoleUser, Content: "first"},
|
||||
{
|
||||
Role: zeroruntime.MessageRoleAssistant,
|
||||
Content: "calling a tool",
|
||||
ToolCalls: []zeroruntime.ToolCall{{ID: "toolu_1", Name: "grep", Arguments: `{}`}},
|
||||
},
|
||||
{Role: zeroruntime.MessageRoleTool, Content: "result", ToolCallID: "toolu_1"},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion returned error: %v", err)
|
||||
}
|
||||
drain(stream)
|
||||
|
||||
messages := gotBody["messages"].([]any)
|
||||
if len(messages) != 3 {
|
||||
t.Fatalf("expected 3 wire messages, got %d: %#v", len(messages), messages)
|
||||
}
|
||||
lastBlockHasCache := func(message any) bool {
|
||||
blocks, ok := message.(map[string]any)["content"].([]any)
|
||||
if !ok || len(blocks) == 0 {
|
||||
return false
|
||||
}
|
||||
last := blocks[len(blocks)-1].(map[string]any)
|
||||
control, ok := last["cache_control"].(map[string]any)
|
||||
return ok && control["type"] == "ephemeral"
|
||||
}
|
||||
if lastBlockHasCache(messages[0]) {
|
||||
t.Fatalf("first message must not carry a breakpoint: %#v", messages[0])
|
||||
}
|
||||
if !lastBlockHasCache(messages[1]) {
|
||||
t.Fatalf("second-to-last message must carry a breakpoint: %#v", messages[1])
|
||||
}
|
||||
if !lastBlockHasCache(messages[2]) {
|
||||
t.Fatalf("last message must carry a breakpoint: %#v", messages[2])
|
||||
}
|
||||
}
|
||||
|
||||
// A string-content message must be converted to a block array so it can carry
|
||||
// the breakpoint, and thinking blocks must never carry cache_control (the API
|
||||
// rejects them) — the marker goes on the last cacheable block instead.
|
||||
func TestApplyMessageCacheBreakpointsSkipsThinkingBlocks(t *testing.T) {
|
||||
messages := []anthropicMessage{
|
||||
{Role: "user", Content: "plain string"},
|
||||
{Role: "assistant", Content: []map[string]any{
|
||||
{"type": "text", "text": "answer"},
|
||||
{"type": "thinking", "thinking": "hmm", "signature": "sig"},
|
||||
}},
|
||||
}
|
||||
applyMessageCacheBreakpoints(messages)
|
||||
|
||||
userBlocks := contentBlocks(messages[0].Content)
|
||||
if len(userBlocks) != 1 || userBlocks[0]["cache_control"] == nil {
|
||||
t.Fatalf("string content must become a marked block: %#v", messages[0].Content)
|
||||
}
|
||||
assistantBlocks := contentBlocks(messages[1].Content)
|
||||
if assistantBlocks[1]["cache_control"] != nil {
|
||||
t.Fatalf("thinking block must not be marked: %#v", assistantBlocks[1])
|
||||
}
|
||||
if assistantBlocks[0]["cache_control"] == nil {
|
||||
t.Fatalf("text block must carry the breakpoint instead: %#v", assistantBlocks[0])
|
||||
}
|
||||
}
|
||||
@@ -405,35 +405,9 @@ func (provider *Provider) anthropicRequest(request zeroruntime.CompletionRequest
|
||||
}
|
||||
mapped.Tools[len(mapped.Tools)-1].CacheControl = &cacheControl{Type: cacheEphemeral}
|
||||
}
|
||||
applyMessageCacheBreakpoints(mapped.Messages)
|
||||
return mapped, nil
|
||||
}
|
||||
|
||||
// applyMessageCacheBreakpoints marks the last content block of the final two
|
||||
// messages with cache_control so the conversation transcript is cached
|
||||
// turn-over-turn, not just the system prompt and tool definitions. Two message
|
||||
// breakpoints (not one) keep the previous turn's prefix a cache hit while the
|
||||
// newest suffix is being written. Anthropic allows at most 4 breakpoints per
|
||||
// request: system + tools use two above, these use the remaining two. Thinking
|
||||
// blocks cannot carry cache_control, so the marker goes on the last block that
|
||||
// can.
|
||||
func applyMessageCacheBreakpoints(messages []anthropicMessage) {
|
||||
marked := 0
|
||||
for i := len(messages) - 1; i >= 0 && marked < 2; i-- {
|
||||
blocks := contentBlocks(messages[i].Content)
|
||||
for j := len(blocks) - 1; j >= 0; j-- {
|
||||
blockType, _ := blocks[j]["type"].(string)
|
||||
if blockType == "thinking" || blockType == "redacted_thinking" {
|
||||
continue
|
||||
}
|
||||
blocks[j]["cache_control"] = map[string]any{"type": cacheEphemeral}
|
||||
messages[i].Content = blocks
|
||||
marked++
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func mapMessages(messages []zeroruntime.Message) (string, []anthropicMessage, error) {
|
||||
systemParts := []string{}
|
||||
mapped := []anthropicMessage{}
|
||||
|
||||
@@ -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, Anthropic, and Google (Gemini) providers.
|
||||
// Applied to the OpenAI and Anthropic 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,7 +91,6 @@ 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,
|
||||
|
||||
@@ -49,10 +49,6 @@ 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.
|
||||
@@ -71,7 +67,6 @@ type Provider struct {
|
||||
customHeaders map[string]string
|
||||
httpClient *http.Client
|
||||
userAgent string
|
||||
oauthResolver providerio.TokenResolver
|
||||
streamIdleTimeout time.Duration
|
||||
}
|
||||
|
||||
@@ -100,7 +95,6 @@ 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
|
||||
}
|
||||
@@ -133,22 +127,20 @@ func (provider *Provider) stream(ctx context.Context, body []byte, events chan<-
|
||||
streamCtx, cancelStream := context.WithCancel(ctx)
|
||||
defer cancelStream()
|
||||
|
||||
response, err := providerio.SendWithAuthRetry(streamCtx, provider.httpClient, http.MethodPost, provider.streamURL(), body,
|
||||
providerio.AuthHeaders{
|
||||
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{
|
||||
APIKey: provider.apiKey,
|
||||
DefaultAuthHeader: "x-goog-api-key",
|
||||
AuthHeader: provider.authHeader,
|
||||
AuthScheme: provider.authScheme,
|
||||
AuthHeaderValue: provider.authHeaderValue,
|
||||
CustomHeaders: provider.customHeaders,
|
||||
},
|
||||
provider.oauthResolver,
|
||||
func(request *http.Request) {
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
if provider.userAgent != "" {
|
||||
request.Header.Set("User-Agent", provider.userAgent)
|
||||
}
|
||||
}, 0)
|
||||
})
|
||||
}, 0)
|
||||
if err != nil {
|
||||
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventError, Error: provider.redact("provider stream error: " + err.Error())})
|
||||
return
|
||||
|
||||
@@ -370,57 +370,6 @@ 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"}}`)
|
||||
|
||||
@@ -23,18 +23,12 @@ import (
|
||||
// (billable) work. Only the INITIAL request is ever in scope; once the response
|
||||
// body starts streaming it is never re-issued.
|
||||
|
||||
const defaultMaxRetryAttempts = 6
|
||||
const defaultMaxRetryAttempts = 3
|
||||
|
||||
// maxBackoff caps a single backoff wait so a hostile or buggy Retry-After can't
|
||||
// stall the agent for minutes.
|
||||
const maxBackoff = 30 * time.Second
|
||||
|
||||
// retryBackoffBase is the first wait when the server supplied no Retry-After.
|
||||
// Rate-limit windows are measured in seconds, not milliseconds: retrying a 429
|
||||
// after 400ms almost always burns the attempt while still limited, so the
|
||||
// schedule is 2s, 4s, 8s, 16s, then maxBackoff. A var so tests can shrink it.
|
||||
var retryBackoffBase = 2 * time.Second
|
||||
|
||||
// SendWithRetry issues an HTTP request, retrying ONLY the safe-to-replay server
|
||||
// responses (429 and 503, see ShouldRetryStatus) up to maxAttempts — backing off
|
||||
// between tries and honoring a server Retry-After header and context
|
||||
@@ -116,11 +110,17 @@ func ShouldRetryStatus(code int) bool {
|
||||
}
|
||||
|
||||
// Backoff waits before retry attempt N (1-based), returning false if the context
|
||||
// is cancelled during the wait. A server-supplied (positive) Retry-After wins;
|
||||
// otherwise the wait doubles from retryBackoffBase per attempt. Either way the
|
||||
// wait is capped at maxBackoff.
|
||||
// is cancelled during the wait. The wait is attempt*400ms unless the server
|
||||
// supplied a (positive) Retry-After, and is capped at maxBackoff.
|
||||
func Backoff(ctx context.Context, attempt int, retryAfter time.Duration) bool {
|
||||
timer := time.NewTimer(backoffWait(attempt, retryAfter))
|
||||
wait := time.Duration(attempt) * 400 * time.Millisecond
|
||||
if retryAfter > 0 {
|
||||
wait = retryAfter
|
||||
}
|
||||
if wait > maxBackoff {
|
||||
wait = maxBackoff
|
||||
}
|
||||
timer := time.NewTimer(wait)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
@@ -130,27 +130,6 @@ func Backoff(ctx context.Context, attempt int, retryAfter time.Duration) bool {
|
||||
}
|
||||
}
|
||||
|
||||
// backoffWait computes the wait before retry attempt N (1-based): Retry-After
|
||||
// when supplied, else exponential from retryBackoffBase, both capped at
|
||||
// maxBackoff. The exponent is clamped so a large attempt count cannot overflow.
|
||||
func backoffWait(attempt int, retryAfter time.Duration) time.Duration {
|
||||
wait := retryAfter
|
||||
if wait <= 0 {
|
||||
exponent := attempt - 1
|
||||
if exponent > 5 {
|
||||
exponent = 5
|
||||
}
|
||||
if exponent < 0 {
|
||||
exponent = 0
|
||||
}
|
||||
wait = retryBackoffBase * time.Duration(1<<exponent)
|
||||
}
|
||||
if wait > maxBackoff {
|
||||
wait = maxBackoff
|
||||
}
|
||||
return wait
|
||||
}
|
||||
|
||||
// RetryAfter parses a response's Retry-After header (delay-seconds or an HTTP
|
||||
// date) into a positive duration, or 0 when absent/unparseable. The result is
|
||||
// capped at maxBackoff by Backoff.
|
||||
|
||||
@@ -94,40 +94,7 @@ func TestBackoffWaitsThenReturnsTrue(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// shrinkBackoff makes retry waits negligible for the duration of a test.
|
||||
func shrinkBackoff(t *testing.T) {
|
||||
t.Helper()
|
||||
saved := retryBackoffBase
|
||||
retryBackoffBase = time.Millisecond
|
||||
t.Cleanup(func() { retryBackoffBase = saved })
|
||||
}
|
||||
|
||||
func TestBackoffWaitSchedule(t *testing.T) {
|
||||
// Without Retry-After the wait doubles per attempt from 2s and caps at 30s;
|
||||
// a supplied Retry-After wins but is capped too.
|
||||
cases := []struct {
|
||||
attempt int
|
||||
retryAfter time.Duration
|
||||
want time.Duration
|
||||
}{
|
||||
{1, 0, 2 * time.Second},
|
||||
{2, 0, 4 * time.Second},
|
||||
{3, 0, 8 * time.Second},
|
||||
{4, 0, 16 * time.Second},
|
||||
{5, 0, 30 * time.Second}, // 32s capped
|
||||
{50, 0, 30 * time.Second}, // clamped exponent, no overflow
|
||||
{1, 7 * time.Second, 7 * time.Second},
|
||||
{1, 5 * time.Minute, 30 * time.Second}, // hostile Retry-After capped
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := backoffWait(c.attempt, c.retryAfter); got != c.want {
|
||||
t.Errorf("backoffWait(%d, %v) = %v, want %v", c.attempt, c.retryAfter, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSendWithRetryRetriesThenSucceeds(t *testing.T) {
|
||||
shrinkBackoff(t)
|
||||
var hits int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if atomic.AddInt32(&hits, 1) == 1 {
|
||||
@@ -173,7 +140,6 @@ func TestSendWithRetryReturnsNonRetryableImmediately(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestSendWithRetryReturnsLastResponseAfterMaxAttempts(t *testing.T) {
|
||||
shrinkBackoff(t)
|
||||
var hits int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
atomic.AddInt32(&hits, 1)
|
||||
|
||||
@@ -87,20 +87,11 @@ func windowsCommandLine(args []string) (string, error) {
|
||||
if len(args) == 0 {
|
||||
return "", errorsNewWindowsProcess("missing command")
|
||||
}
|
||||
out := make([]string, 0, len(args))
|
||||
for _, arg := range args {
|
||||
if strings.ContainsRune(arg, 0) {
|
||||
return "", errorsNewWindowsProcess("command argument contains NUL")
|
||||
}
|
||||
}
|
||||
// Zero's cmd.exe shell invocation needs its final element (the shell
|
||||
// command text) reaching cmd.exe raw, not escaped like a normal argv
|
||||
// element: see WindowsShellCommandLine for why. Every other command this
|
||||
// runner might launch is escaped normally below.
|
||||
if commandLine, ok := windowsShellCommandLineFromArgs(args); ok {
|
||||
return commandLine, nil
|
||||
}
|
||||
out := make([]string, 0, len(args))
|
||||
for _, arg := range args {
|
||||
out = append(out, syscall.EscapeArg(arg))
|
||||
}
|
||||
return strings.Join(out, " "), nil
|
||||
|
||||
@@ -105,56 +105,6 @@ const (
|
||||
WindowsSandboxLevelDisabled WindowsSandboxLevel = "disabled"
|
||||
)
|
||||
|
||||
// WindowsShellArgs returns cmd.exe's argv (after the executable name) for
|
||||
// running commandText as a shell command on Windows: no wrapping quotes and
|
||||
// no /S, so commandText reaches cmd.exe's own /C remainder parsing exactly as
|
||||
// written. See WindowsShellCommandLine for why that matters. Used both as
|
||||
// CommandSpec.Args for the plain (unwrapped) path and, via
|
||||
// windowsShellCommandLineFromArgs, to recognize this same shape once it has
|
||||
// round-tripped through the sandboxed runner's CLI argv.
|
||||
func WindowsShellArgs(commandText string) []string {
|
||||
return []string{"/d", "/c", commandText}
|
||||
}
|
||||
|
||||
// WindowsShellCommandLine builds the raw command line a Windows CreateProcess
|
||||
// call should receive to run commandText as a shell command: "cmd.exe /d /c "
|
||||
// followed by commandText completely unescaped.
|
||||
//
|
||||
// cmd.exe's own /C remainder parsing is not CommandLineToArgvW-compatible: it
|
||||
// applies its own quote-stripping heuristic (documented in `cmd /?`) whenever
|
||||
// the remainder starts with a literal quote, and that heuristic does not
|
||||
// understand backslash-escaped inner quotes the way a normal argv consumer
|
||||
// would. A command like `python -c "print(15 / 3)"` passed through the
|
||||
// standard exec.Cmd Args mechanism gets wrapped in an outer pair of quotes
|
||||
// (because it contains spaces) with its own quotes escaped as \" — exactly
|
||||
// the shape that heuristic corrupts, stripping the outer quotes but leaving
|
||||
// the literal backslashes behind. Passing commandText through raw, exactly as
|
||||
// the model wrote it, means cmd.exe parses it the same way it would if typed
|
||||
// directly at an interactive prompt: a leading quoted executable path like
|
||||
// `"C:\Program Files\foo.exe" arg` is preserved correctly by cmd.exe's own
|
||||
// quote-preserving special case, and a command with no leading quote at all,
|
||||
// like the python example above, is never touched by the stripping heuristic
|
||||
// in the first place.
|
||||
func WindowsShellCommandLine(commandText string) string {
|
||||
return "cmd.exe /d /c " + commandText
|
||||
}
|
||||
|
||||
// windowsShellCommandLineFromArgs reports whether args is exactly
|
||||
// ["cmd.exe"]+WindowsShellArgs(text) for some text and, if so, returns
|
||||
// WindowsShellCommandLine's result for that text. The sandboxed runner path
|
||||
// only has the deserialized argv (WindowsSandboxCommandConfig.Command) to
|
||||
// work from, not the original commandText, so it recovers the raw command
|
||||
// line by recognizing the shape instead.
|
||||
func windowsShellCommandLineFromArgs(args []string) (string, bool) {
|
||||
if len(args) != 4 {
|
||||
return "", false
|
||||
}
|
||||
if !strings.EqualFold(args[0], "cmd.exe") || args[1] != "/d" || args[2] != "/c" {
|
||||
return "", false
|
||||
}
|
||||
return WindowsShellCommandLine(args[3]), true
|
||||
}
|
||||
|
||||
type WindowsSandboxCommandArgsOptions struct {
|
||||
SandboxHome string
|
||||
CommandCWD string
|
||||
|
||||
@@ -147,56 +147,3 @@ func TestRunWindowsSandboxCommandRunnerRejectsInvalidArgs(t *testing.T) {
|
||||
t.Fatalf("stderr = %q, want runner name", stderr.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestWindowsShellArgsHasNoSAndKeepsCommandTextWhole(t *testing.T) {
|
||||
args := WindowsShellArgs(`python -c "print(15 / 3)"`)
|
||||
want := []string{"/d", "/c", `python -c "print(15 / 3)"`}
|
||||
if len(args) != len(want) {
|
||||
t.Fatalf("args = %#v, want %#v", args, want)
|
||||
}
|
||||
for index := range want {
|
||||
if args[index] != want[index] {
|
||||
t.Fatalf("args = %#v, want %#v", args, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWindowsShellCommandLineLeavesCommandTextUnescaped(t *testing.T) {
|
||||
// The whole point of this function: commandText's own quotes must reach
|
||||
// cmd.exe exactly as written, not re-quoted or backslash-escaped the way
|
||||
// a normal argv-escaping function (like syscall.EscapeArg) would treat an
|
||||
// argument containing spaces and embedded quotes.
|
||||
commandText := `python -c "print(15 / 3)"`
|
||||
got := WindowsShellCommandLine(commandText)
|
||||
want := `cmd.exe /d /c python -c "print(15 / 3)"`
|
||||
if got != want {
|
||||
t.Fatalf("WindowsShellCommandLine = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWindowsShellCommandLineFromArgsRecognizesTheShellShape(t *testing.T) {
|
||||
commandLine, ok := windowsShellCommandLineFromArgs([]string{"cmd.exe", "/d", "/c", `echo "hi"`})
|
||||
if !ok {
|
||||
t.Fatal("windowsShellCommandLineFromArgs did not recognize the shell shape")
|
||||
}
|
||||
if want := `cmd.exe /d /c echo "hi"`; commandLine != want {
|
||||
t.Fatalf("commandLine = %q, want %q", commandLine, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWindowsShellCommandLineFromArgsRejectsOtherShapes(t *testing.T) {
|
||||
cases := [][]string{
|
||||
nil,
|
||||
{"cmd.exe"},
|
||||
{"cmd.exe", "/d", "/c"},
|
||||
{"cmd.exe", "/d", "/c", "echo hi", "extra"},
|
||||
{"git.exe", "/d", "/c", "echo hi"},
|
||||
{"cmd.exe", "/s", "/c", "echo hi"},
|
||||
{"cmd.exe", "/d", "/k", "echo hi"},
|
||||
}
|
||||
for _, args := range cases {
|
||||
if _, ok := windowsShellCommandLineFromArgs(args); ok {
|
||||
t.Fatalf("windowsShellCommandLineFromArgs(%#v) matched, want no match", args)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -25,7 +25,6 @@ import (
|
||||
const (
|
||||
sessionTagSpecialist = "specialist"
|
||||
promptFileThresholdBytes = 4 * 1024
|
||||
maxSpecialistDepth = 8 // hard cap to prevent infinite recursion/resource exhaustion via Task
|
||||
)
|
||||
|
||||
const SessionTagSpecialist = sessionTagSpecialist
|
||||
@@ -219,13 +218,6 @@ func (executor Executor) Run(ctx context.Context, params TaskParameters, options
|
||||
if options.CurrentDepth < 0 {
|
||||
return ExecResult{}, fmt.Errorf("current depth cannot be negative")
|
||||
}
|
||||
// >= (not >): this Run call spawns a CHILD at options.CurrentDepth+1, so a
|
||||
// parent already AT the cap must still be rejected here rather than being
|
||||
// allowed to launch one more level before the child's own next call trips
|
||||
// the guard.
|
||||
if options.CurrentDepth >= maxSpecialistDepth {
|
||||
return ExecResult{}, fmt.Errorf("spawning a specialist at depth %d would exceed maximum nesting depth %d", options.CurrentDepth+1, maxSpecialistDepth)
|
||||
}
|
||||
if strings.TrimSpace(params.Prompt) == "" {
|
||||
return ExecResult{}, fmt.Errorf("specialist prompt is required")
|
||||
}
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
package specialist
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
@@ -337,39 +336,6 @@ func TestBuildArgsRejectsInvalidInputs(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunRejectsDepthExceedingMax(t *testing.T) {
|
||||
_, err := (Executor{}).Run(context.Background(), TaskParameters{
|
||||
Prompt: "hi",
|
||||
}, TaskRunOptions{CurrentDepth: maxSpecialistDepth + 1})
|
||||
if err == nil || !strings.Contains(err.Error(), "depth") {
|
||||
t.Fatalf("Run error = %v, want depth error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRunRejectsDepthAtMax covers the boundary: a parent already AT the cap
|
||||
// must be rejected too, since this Run call would launch a child one level
|
||||
// past it (--depth CurrentDepth+1). Only checking ">" here would let that
|
||||
// child start before the guard ever caught it. The guard sits before the
|
||||
// fresh/resume branch in Run, so both call shapes must be proven to reject
|
||||
// here rather than falling through to runFresh/runResume.
|
||||
func TestRunRejectsDepthAtMax(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
params TaskParameters
|
||||
}{
|
||||
{name: "fresh", params: TaskParameters{Prompt: "hi"}},
|
||||
{name: "resume", params: TaskParameters{Prompt: "hi", Resume: "child_session"}},
|
||||
}
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
_, err := (Executor{}).Run(context.Background(), tc.params, TaskRunOptions{CurrentDepth: maxSpecialistDepth})
|
||||
if err == nil || !strings.Contains(err.Error(), "depth") {
|
||||
t.Fatalf("Run error = %v, want depth error", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildArgsRejectsInvalidSessionIDs(t *testing.T) {
|
||||
_, err := (Executor{NewSessionID: func() (string, error) { return "../escape", nil }}).BuildArgs(BuildArgsInput{
|
||||
Prompt: "hi",
|
||||
|
||||
@@ -272,36 +272,15 @@ 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 {
|
||||
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, "")
|
||||
}
|
||||
lines := make([]string, 0, len(questions)*3)
|
||||
for index, question := range questions {
|
||||
answer := ""
|
||||
if index < len(answers) {
|
||||
answer = strings.TrimSpace(answers[index])
|
||||
}
|
||||
if answer == "" {
|
||||
if anyAnswered {
|
||||
answer = "(left blank)"
|
||||
} else {
|
||||
answer = "(skipped)"
|
||||
}
|
||||
answer = "(no answer provided)"
|
||||
}
|
||||
lines = append(lines, fmt.Sprintf("%d. [question] %s", index+1, question.Question))
|
||||
lines = append(lines, "[answer] "+answer)
|
||||
|
||||
@@ -1,34 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
+20
-163
@@ -1,6 +1,7 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
@@ -119,17 +120,14 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS
|
||||
Meta: meta,
|
||||
}
|
||||
}
|
||||
// Release any plan-scoped resources once the command has finished running.
|
||||
defer plan.Cleanup()
|
||||
addSandboxMeta(meta, plan)
|
||||
|
||||
// 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
|
||||
var stdout bytes.Buffer
|
||||
var stderr bytes.Buffer
|
||||
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().
|
||||
@@ -141,12 +139,7 @@ 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)
|
||||
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))
|
||||
stderrText := appendSandboxBlocks(stderr.String(), monitor.Stop())
|
||||
|
||||
if errors.Is(commandCtx.Err(), context.DeadlineExceeded) {
|
||||
return Result{
|
||||
@@ -163,31 +156,26 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS
|
||||
Meta: meta,
|
||||
}
|
||||
}
|
||||
markLikelySandboxDenial(meta, plan, exitCode, stdoutText, stderrText)
|
||||
outText, errText, truncated := budgetBashCapture(stdoutText, stdout.total, stderrText, stderrTotal, meta)
|
||||
markLikelySandboxDenial(meta, plan, exitCode, stdout.String(), stderrText)
|
||||
return Result{
|
||||
Status: StatusError,
|
||||
Output: formatBashOutputWithShellHint(commandText, outText, errText, exitCode, meta),
|
||||
Truncated: truncated,
|
||||
Meta: meta,
|
||||
Status: StatusError,
|
||||
Output: formatBashOutputWithShellHint(commandText, stdout.String(), stderrText, exitCode, meta),
|
||||
Meta: meta,
|
||||
}
|
||||
}
|
||||
|
||||
markLikelySandboxDenial(meta, plan, exitCode, stdoutText, stderrText)
|
||||
outText, errText, truncated := budgetBashCapture(stdoutText, stdout.total, stderrText, stderrTotal, meta)
|
||||
markLikelySandboxDenial(meta, plan, exitCode, stdout.String(), stderrText)
|
||||
if meta[SandboxLikelyDeniedMeta] == "true" {
|
||||
return Result{
|
||||
Status: StatusError,
|
||||
Output: formatBashOutputWithShellHint(commandText, outText, errText, exitCode, meta),
|
||||
Truncated: truncated,
|
||||
Meta: meta,
|
||||
Status: StatusError,
|
||||
Output: formatBashOutputWithShellHint(commandText, stdout.String(), stderrText, exitCode, meta),
|
||||
Meta: meta,
|
||||
}
|
||||
}
|
||||
return Result{
|
||||
Status: StatusOK,
|
||||
Output: formatBashOutput(outText, errText, exitCode),
|
||||
Truncated: truncated,
|
||||
Meta: meta,
|
||||
Status: StatusOK,
|
||||
Output: formatBashOutput(stdout.String(), stderrText, exitCode),
|
||||
Meta: meta,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -233,14 +221,6 @@ func shellIssueBlockResult(issue shellIssue) Result {
|
||||
}
|
||||
}
|
||||
|
||||
// buildBashCommand returns the exec.Cmd and the sandbox plan for running
|
||||
// commandText. On Windows, when the command is not wrapped by the sandbox
|
||||
// engine (plan.Wrapped == false), it also overrides the child's raw command
|
||||
// line so commandText reaches cmd.exe unescaped; see
|
||||
// zeroSandbox.WindowsShellCommandLine for why that matters. The wrapped case
|
||||
// gets the same treatment inside the sandboxed runner process itself
|
||||
// (internal/sandbox/windows_process_windows.go), since that command line is
|
||||
// built there, not here.
|
||||
func buildBashCommand(ctx context.Context, commandText string, absoluteCwd string, engine *zeroSandbox.Engine) (*exec.Cmd, zeroSandbox.CommandPlan, error) {
|
||||
spec := zeroSandbox.CommandSpec{
|
||||
Name: shellExecutable(),
|
||||
@@ -248,11 +228,7 @@ func buildBashCommand(ctx context.Context, commandText string, absoluteCwd strin
|
||||
Dir: absoluteCwd,
|
||||
}
|
||||
if engine != nil {
|
||||
command, plan, err := engine.CommandContext(ctx, spec)
|
||||
if err == nil {
|
||||
applyWindowsShellCommandLine(command, commandText, plan.Wrapped)
|
||||
}
|
||||
return command, plan, err
|
||||
return engine.CommandContext(ctx, spec)
|
||||
}
|
||||
plan := zeroSandbox.CommandPlan{
|
||||
Backend: zeroSandbox.Backend{
|
||||
@@ -266,7 +242,6 @@ func buildBashCommand(ctx context.Context, commandText string, absoluteCwd strin
|
||||
}
|
||||
command := exec.CommandContext(ctx, spec.Name, spec.Args...)
|
||||
command.Dir = spec.Dir
|
||||
applyWindowsShellCommandLine(command, commandText, plan.Wrapped)
|
||||
return command, plan, nil
|
||||
}
|
||||
|
||||
@@ -327,7 +302,7 @@ func shellExecutable() string {
|
||||
|
||||
func shellArguments(command string) []string {
|
||||
if runtime.GOOS == "windows" {
|
||||
return zeroSandbox.WindowsShellArgs(command)
|
||||
return []string{"/d", "/s", "/c", command}
|
||||
}
|
||||
return []string{"-c", command}
|
||||
}
|
||||
@@ -370,124 +345,6 @@ 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 {
|
||||
|
||||
@@ -1,127 +0,0 @@
|
||||
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))])
|
||||
}
|
||||
}
|
||||
@@ -42,7 +42,3 @@ func hardenProcessLifetime(command *exec.Cmd) {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// applyWindowsShellCommandLine is a no-op outside Windows; there is no
|
||||
// cmd.exe command-line-parsing quirk to work around here.
|
||||
func applyWindowsShellCommandLine(command *exec.Cmd, commandText string, wrapped bool) {}
|
||||
@@ -5,10 +5,7 @@ package tools
|
||||
import (
|
||||
"os/exec"
|
||||
"strconv"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
zeroSandbox "github.com/Gitlawb/zero/internal/sandbox"
|
||||
)
|
||||
|
||||
// bashWaitDelay bounds how long Wait blocks for the I/O pipes to drain after the
|
||||
@@ -31,17 +28,3 @@ func hardenProcessLifetime(command *exec.Cmd) {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// applyWindowsShellCommandLine overrides command's raw child command line so
|
||||
// commandText reaches cmd.exe unescaped instead of auto-quoted the way
|
||||
// exec.Cmd would normally encode a single Args element. Skipped when wrapped
|
||||
// is true: the sandbox engine then routes execution through a separate
|
||||
// zero-windows-command-runner process, which builds its own child command
|
||||
// line from scratch (internal/sandbox/windows_process_windows.go) rather than
|
||||
// inheriting whatever this outer exec.Cmd is configured with.
|
||||
func applyWindowsShellCommandLine(command *exec.Cmd, commandText string, wrapped bool) {
|
||||
if wrapped {
|
||||
return
|
||||
}
|
||||
command.SysProcAttr = &syscall.SysProcAttr{CmdLine: zeroSandbox.WindowsShellCommandLine(commandText)}
|
||||
}
|
||||
@@ -8,7 +8,6 @@ import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -18,18 +17,6 @@ import (
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
if len(os.Args) >= 3 && os.Args[1] == "--zero-bash-helper" {
|
||||
if os.Args[2] == "echo-arg" {
|
||||
// Echoes back exactly the one argument it received, to verify a
|
||||
// quoted, space-and-slash-containing argument (the same shape as
|
||||
// `python -c "print(15 / 3)"`) survives the shell invocation
|
||||
// intact instead of being truncated/corrupted.
|
||||
if len(os.Args) < 4 {
|
||||
fmt.Fprintln(os.Stderr, "echo-arg requires exactly one argument")
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Println(os.Args[3])
|
||||
return
|
||||
}
|
||||
runBashToolHelper(os.Args[2])
|
||||
return
|
||||
}
|
||||
@@ -253,43 +240,6 @@ 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")
|
||||
@@ -585,68 +535,6 @@ func TestBashToolAllowsNonInteractiveCommand(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestBashToolPreservesEmbeddedQuotesOnWindows pins a real, previously-broken
|
||||
// case: passing commandText as a normal exec.Cmd argument makes Go wrap it in
|
||||
// an outer pair of quotes (it contains spaces) with its own quotes escaped as
|
||||
// \", and cmd.exe's /C remainder parsing strips the first and last literal
|
||||
// quote character in that remainder without undoing the backslash-escaping -
|
||||
// corrupting a command whose own text contains embedded double quotes,
|
||||
// exactly the shape of `python -c "print(15 / 3)"`, `git commit -m
|
||||
// "message"`, `node -e "..."`. Before the fix, the helper below received a
|
||||
// truncated, mis-quoted argument (starting with a stray literal `"`) instead
|
||||
// of the text between the quotes.
|
||||
func TestBashToolPreservesEmbeddedQuotesOnWindows(t *testing.T) {
|
||||
if runtime.GOOS != "windows" {
|
||||
t.Skip("cmd.exe /S/C quote-remainder parsing is Windows-specific")
|
||||
}
|
||||
root := t.TempDir()
|
||||
executable := shellQuote(os.Args[0])
|
||||
// Space and a slash: the two characters whose position in the original
|
||||
// bug report's SyntaxError ("print(15, truncated right before the /")
|
||||
// showed exactly where the corruption happened.
|
||||
commandText := executable + ` --zero-bash-helper echo-arg "hello / world"`
|
||||
|
||||
result := NewBashTool(root).Run(context.Background(), map[string]any{
|
||||
"command": commandText,
|
||||
})
|
||||
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("expected ok status, got %s: %s", result.Status, result.Output)
|
||||
}
|
||||
if !strings.Contains(result.Output, "hello / world") {
|
||||
t.Fatalf("expected the embedded-quote argument to survive intact, got %q", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBashToolRunsCommandLineForLoopSyntax pins cmd.exe command-line syntax,
|
||||
// not batch-file syntax: a `for %i in (...) do ...` loop with a single
|
||||
// percent sign is valid typed directly at a cmd.exe prompt (or via
|
||||
// `cmd /C "..."`), but requires %%i inside an actual .bat/.cmd FILE, because
|
||||
// batch files perform an extra pass of percent-substitution before FOR ever
|
||||
// runs, consuming the single percent so the loop variable never binds. An
|
||||
// earlier version of this fix ran commandText from a temporary .cmd file
|
||||
// instead of passing it straight through to cmd.exe's /C remainder, which
|
||||
// silently broke single-percent syntax like this.
|
||||
func TestBashToolRunsCommandLineForLoopSyntax(t *testing.T) {
|
||||
if runtime.GOOS != "windows" {
|
||||
t.Skip("cmd.exe command-line vs batch-file parsing is Windows-specific")
|
||||
}
|
||||
root := t.TempDir()
|
||||
|
||||
result := NewBashTool(root).Run(context.Background(), map[string]any{
|
||||
"command": "for %i in (1 2 3) do echo %i",
|
||||
})
|
||||
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("expected ok status, got %s: %s", result.Status, result.Output)
|
||||
}
|
||||
for _, want := range []string{"1", "2", "3"} {
|
||||
if !strings.Contains(result.Output, want) {
|
||||
t.Fatalf("expected the for-loop to expand %%i to %q, got %q", want, result.Output)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func helperCommand(name string) string {
|
||||
executable := shellQuote(os.Args[0])
|
||||
return executable + " --zero-bash-helper " + name
|
||||
|
||||
@@ -2,7 +2,6 @@ package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
@@ -45,7 +44,7 @@ func (tool editFileTool) Run(ctx context.Context, args map[string]any) Result {
|
||||
return tool.RunWithOptions(ctx, args, RunOptions{})
|
||||
}
|
||||
|
||||
func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any, options RunOptions) Result {
|
||||
func (tool editFileTool) RunWithOptions(_ context.Context, args map[string]any, options RunOptions) Result {
|
||||
requestedPath, err := aliasedStringArg(args, []string{"path", "file", "file_path", "filename"}, "", true, false)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for edit_file: " + err.Error())
|
||||
@@ -99,35 +98,6 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
|
||||
}
|
||||
}
|
||||
|
||||
// Fuzzy fallback: when the exact string (and its CRLF translation) is absent,
|
||||
// run a cascade of tolerant matchers (trimmed lines, block anchors, collapsed
|
||||
// whitespace, indentation drift, escape normalization) to locate the span the
|
||||
// model intended. Only a span that occurs literally in the file is accepted,
|
||||
// so the replacement applied below is still exact.
|
||||
if occurrences == 0 {
|
||||
findOld, findNew := oldString, newString
|
||||
if strings.Contains(content, "\r\n") && !strings.Contains(findOld, "\r\n") {
|
||||
findOld = strings.ReplaceAll(findOld, "\n", "\r\n")
|
||||
findNew = strings.ReplaceAll(findNew, "\n", "\r\n")
|
||||
}
|
||||
search, ferr := fuzzyEditMatch(content, findOld, replaceAll)
|
||||
switch {
|
||||
case ferr == nil:
|
||||
oldString = search
|
||||
// The model's new_string was written at old_string's (mismatched)
|
||||
// indentation; re-shape it to the span actually being replaced so a
|
||||
// tolerant match never strips indentation or a trailing CR.
|
||||
newString = adaptReplacementToSpan(search, findOld, findNew)
|
||||
occurrences = strings.Count(content, search)
|
||||
case errors.Is(ferr, errEditFuzzyAmbiguous):
|
||||
return errorResult("Error: old_string matches multiple locations in " + relativePath + " even after fuzzy matching. Provide more surrounding context to make the match unique, or pass replace_all: true.")
|
||||
case errors.Is(ferr, errEditFuzzyNotFound):
|
||||
// Fall through to the exact-match error below.
|
||||
default:
|
||||
return errorResult("Error editing " + relativePath + ": " + ferr.Error())
|
||||
}
|
||||
}
|
||||
|
||||
if occurrences == 0 {
|
||||
return errorResult("Error: Could not find the exact string to replace in " + relativePath + ". The old_string must match the file byte-for-byte.")
|
||||
}
|
||||
@@ -150,10 +120,6 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
|
||||
if err := os.WriteFile(absolutePath, []byte(updated), 0o644); err != nil {
|
||||
return errorResult("Error writing " + relativePath + ": " + err.Error())
|
||||
}
|
||||
// Optional format-on-write (ZERO_FORMAT_ON_WRITE). Must run BEFORE the
|
||||
// FileTracker re-baseline: recording pre-format content would make the very
|
||||
// next edit look like an external modification and trip the conflict guard.
|
||||
updated = maybeFormatWrittenFile(ctx, absolutePath, updated)
|
||||
// Re-baseline to the content we just wrote so subsequent edits in this session
|
||||
// compare against the current on-disk state, not the pre-edit version.
|
||||
newInfo, _ := os.Stat(absolutePath)
|
||||
@@ -164,7 +130,6 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
|
||||
suffix = "s"
|
||||
}
|
||||
summary := fmt.Sprintf("Successfully edited %s (replaced %d occurrence%s).", relativePath, replacedCount, suffix)
|
||||
summary += inlineDiagnostics(ctx, options, absolutePath, relativePath)
|
||||
result := okResult(summary)
|
||||
result.ChangedFiles = []string{relativePath}
|
||||
// Card-only preview (Display.Preview): the model's Output stays the one-line
|
||||
|
||||
@@ -1,554 +0,0 @@
|
||||
package tools
|
||||
|
||||
// Fuzzy fallback matching for edit_file. When the model's old_string fails to
|
||||
// match byte-for-byte (drifted indentation, collapsed whitespace, escaped
|
||||
// characters, a slightly-misremembered middle line), these replacers propose
|
||||
// candidate spans that plausibly correspond to what the model intended. Only a
|
||||
// candidate that occurs literally in the file is accepted, so the replacement
|
||||
// itself is always exact even when the match was tolerant.
|
||||
//
|
||||
// Strategy cascade ported from opencode's edit tool, whose replacers were in
|
||||
// turn distilled from Cline's diff-apply evals and gemini-cli's edit corrector.
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// editReplacer proposes candidate spans for find inside content. Candidates
|
||||
// are re-validated by fuzzyEditMatch before use.
|
||||
type editReplacer func(content, find string) []string
|
||||
|
||||
var (
|
||||
errEditFuzzyNotFound = errors.New("no fuzzy match for old_string")
|
||||
errEditFuzzyAmbiguous = errors.New("fuzzy match for old_string is ambiguous")
|
||||
)
|
||||
|
||||
// Minimum average middle-line similarity for block-anchor matches.
|
||||
const editAnchorSimilarityThreshold = 0.65
|
||||
|
||||
// fuzzyEditMatch runs the replacer cascade and returns the exact span of
|
||||
// content to replace. When replaceAll is false the span must be unique in
|
||||
// content; an ambiguous candidate is skipped in favor of later candidates and
|
||||
// only reported if nothing unique is found. A span wildly larger than
|
||||
// old_string is refused outright rather than risking a destructive edit.
|
||||
func fuzzyEditMatch(content, find string, replaceAll bool) (string, error) {
|
||||
replacers := []editReplacer{
|
||||
lineTrimmedReplacer,
|
||||
blockAnchorReplacer,
|
||||
whitespaceNormalizedReplacer,
|
||||
indentationFlexibleReplacer,
|
||||
escapeNormalizedReplacer,
|
||||
trimmedBoundaryReplacer,
|
||||
contextAwareReplacer,
|
||||
}
|
||||
found := false
|
||||
for _, replacer := range replacers {
|
||||
// Collect the replacer's DISTINCT candidate spans that literally occur in
|
||||
// content. Two or more distinct spans from one strategy (e.g. duplicate
|
||||
// blocks at different indentation, each occurring once) mean the model's
|
||||
// intent is genuinely ambiguous — silently editing the first would be a
|
||||
// wrong-span write, so it is rejected instead.
|
||||
var candidates []string
|
||||
seen := map[string]bool{}
|
||||
for _, search := range replacer(content, find) {
|
||||
if search == "" || seen[search] {
|
||||
continue
|
||||
}
|
||||
seen[search] = true
|
||||
if strings.Index(content, search) < 0 {
|
||||
continue
|
||||
}
|
||||
candidates = append(candidates, search)
|
||||
}
|
||||
if len(candidates) == 0 {
|
||||
continue
|
||||
}
|
||||
found = true
|
||||
if !replaceAll && len(candidates) > 1 {
|
||||
return "", errEditFuzzyAmbiguous
|
||||
}
|
||||
search := candidates[0]
|
||||
if isDisproportionateEditMatch(search, find) {
|
||||
return "", errors.New("refusing replacement because the matched span is much larger than old_string; re-read the file and provide the full exact old_string for the intended replacement")
|
||||
}
|
||||
if replaceAll {
|
||||
return search, nil
|
||||
}
|
||||
if strings.Index(content, search) == strings.LastIndex(content, search) {
|
||||
return search, nil
|
||||
}
|
||||
// The single candidate occurs at multiple positions; a later, stricter
|
||||
// strategy may still resolve a unique span, so keep cascading.
|
||||
}
|
||||
if !found {
|
||||
return "", errEditFuzzyNotFound
|
||||
}
|
||||
return "", errEditFuzzyAmbiguous
|
||||
}
|
||||
|
||||
// isDisproportionateEditMatch guards against anchor-style replacers matching a
|
||||
// span far larger than the text the model asked to replace (e.g. first/last
|
||||
// line anchors bridging hundreds of unrelated lines).
|
||||
func isDisproportionateEditMatch(search, find string) bool {
|
||||
findLines := strings.Count(find, "\n") + 1
|
||||
searchLines := strings.Count(search, "\n") + 1
|
||||
limit := findLines + 3
|
||||
if findLines*2 > limit {
|
||||
limit = findLines * 2
|
||||
}
|
||||
if searchLines >= limit {
|
||||
return true
|
||||
}
|
||||
if findLines == 1 {
|
||||
return false
|
||||
}
|
||||
searchTrimmed := len(strings.TrimSpace(search))
|
||||
findTrimmed := len(strings.TrimSpace(find))
|
||||
byteLimit := findTrimmed + 500
|
||||
if findTrimmed*4 > byteLimit {
|
||||
byteLimit = findTrimmed * 4
|
||||
}
|
||||
return searchTrimmed > byteLimit
|
||||
}
|
||||
|
||||
// splitFindLines splits find into lines, dropping the trailing empty line a
|
||||
// trailing newline produces so windows align with real content lines.
|
||||
func splitFindLines(find string) []string {
|
||||
lines := strings.Split(find, "\n")
|
||||
if len(lines) > 0 && lines[len(lines)-1] == "" {
|
||||
lines = lines[:len(lines)-1]
|
||||
}
|
||||
return lines
|
||||
}
|
||||
|
||||
// lineTrimmedReplacer matches when every line equals the corresponding
|
||||
// old_string line after trimming surrounding whitespace.
|
||||
func lineTrimmedReplacer(content, find string) []string {
|
||||
contentLines := strings.Split(content, "\n")
|
||||
findLines := splitFindLines(find)
|
||||
if len(findLines) == 0 {
|
||||
return nil
|
||||
}
|
||||
var candidates []string
|
||||
for i := 0; i+len(findLines) <= len(contentLines); i++ {
|
||||
matches := true
|
||||
for j := range findLines {
|
||||
if strings.TrimSpace(contentLines[i+j]) != strings.TrimSpace(findLines[j]) {
|
||||
matches = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if matches {
|
||||
candidates = append(candidates, strings.Join(contentLines[i:i+len(findLines)], "\n"))
|
||||
}
|
||||
}
|
||||
return candidates
|
||||
}
|
||||
|
||||
// blockAnchorReplacer anchors on the first and last lines (trimmed) and
|
||||
// accepts the block when the middle lines average >= the similarity threshold
|
||||
// (Levenshtein-based), tolerating a slightly misremembered interior.
|
||||
func blockAnchorReplacer(content, find string) []string {
|
||||
findLines := splitFindLines(find)
|
||||
if len(findLines) < 3 {
|
||||
return nil
|
||||
}
|
||||
contentLines := strings.Split(content, "\n")
|
||||
firstAnchor := strings.TrimSpace(findLines[0])
|
||||
lastAnchor := strings.TrimSpace(findLines[len(findLines)-1])
|
||||
searchBlockSize := len(findLines)
|
||||
maxLineDelta := searchBlockSize / 4
|
||||
if maxLineDelta < 1 {
|
||||
maxLineDelta = 1
|
||||
}
|
||||
|
||||
type span struct{ start, end int }
|
||||
var candidates []span
|
||||
for i := 0; i < len(contentLines); i++ {
|
||||
if strings.TrimSpace(contentLines[i]) != firstAnchor {
|
||||
continue
|
||||
}
|
||||
// Only the first occurrence of the last anchor after this start counts,
|
||||
// mirroring the reference implementation: a farther-away closing line is
|
||||
// assumed to close a different block.
|
||||
for j := i + 2; j < len(contentLines); j++ {
|
||||
if strings.TrimSpace(contentLines[j]) != lastAnchor {
|
||||
continue
|
||||
}
|
||||
actualBlockSize := j - i + 1
|
||||
delta := actualBlockSize - searchBlockSize
|
||||
if delta < 0 {
|
||||
delta = -delta
|
||||
}
|
||||
if delta <= maxLineDelta {
|
||||
candidates = append(candidates, span{start: i, end: j})
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
if len(candidates) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
middleSimilarity := func(candidate span) float64 {
|
||||
actualBlockSize := candidate.end - candidate.start + 1
|
||||
linesToCheck := searchBlockSize - 2
|
||||
if actualBlockSize-2 < linesToCheck {
|
||||
linesToCheck = actualBlockSize - 2
|
||||
}
|
||||
if linesToCheck <= 0 {
|
||||
return 1.0
|
||||
}
|
||||
similarity := 0.0
|
||||
for j := 1; j < searchBlockSize-1 && j < actualBlockSize-1; j++ {
|
||||
contentLine := strings.TrimSpace(contentLines[candidate.start+j])
|
||||
findLine := strings.TrimSpace(findLines[j])
|
||||
maxLen := len(contentLine)
|
||||
if len(findLine) > maxLen {
|
||||
maxLen = len(findLine)
|
||||
}
|
||||
if maxLen == 0 {
|
||||
continue
|
||||
}
|
||||
distance := levenshtein(contentLine, findLine)
|
||||
similarity += 1 - float64(distance)/float64(maxLen)
|
||||
}
|
||||
return similarity / float64(linesToCheck)
|
||||
}
|
||||
|
||||
// Return EVERY candidate that clears the similarity threshold, best first.
|
||||
// Picking only the best would hide competing blocks from fuzzyEditMatch's
|
||||
// distinct-candidate ambiguity check and could silently edit the wrong
|
||||
// block; with all qualifying spans surfaced, one clear winner still
|
||||
// resolves (single candidate) while two plausible blocks are rejected as
|
||||
// ambiguous. replaceAll consumers take the first (most similar) span.
|
||||
type scoredSpan struct {
|
||||
span span
|
||||
similarity float64
|
||||
}
|
||||
var qualifying []scoredSpan
|
||||
for _, candidate := range candidates {
|
||||
if s := middleSimilarity(candidate); s >= editAnchorSimilarityThreshold {
|
||||
qualifying = append(qualifying, scoredSpan{span: candidate, similarity: s})
|
||||
}
|
||||
}
|
||||
if len(qualifying) == 0 {
|
||||
return nil
|
||||
}
|
||||
sort.SliceStable(qualifying, func(i, j int) bool {
|
||||
return qualifying[i].similarity > qualifying[j].similarity
|
||||
})
|
||||
spans := make([]string, 0, len(qualifying))
|
||||
for _, scored := range qualifying {
|
||||
spans = append(spans, strings.Join(contentLines[scored.span.start:scored.span.end+1], "\n"))
|
||||
}
|
||||
return spans
|
||||
}
|
||||
|
||||
var editWhitespaceRun = regexp.MustCompile(`\s+`)
|
||||
|
||||
func normalizeEditWhitespace(text string) string {
|
||||
return strings.TrimSpace(editWhitespaceRun.ReplaceAllString(text, " "))
|
||||
}
|
||||
|
||||
// whitespaceNormalizedReplacer matches after collapsing all whitespace runs to
|
||||
// a single space: full lines, sub-line spans (via a word-boundary regex), and
|
||||
// multi-line windows.
|
||||
func whitespaceNormalizedReplacer(content, find string) []string {
|
||||
normalizedFind := normalizeEditWhitespace(find)
|
||||
if normalizedFind == "" {
|
||||
return nil
|
||||
}
|
||||
contentLines := strings.Split(content, "\n")
|
||||
var candidates []string
|
||||
var subLinePattern *regexp.Regexp
|
||||
for _, line := range contentLines {
|
||||
normalizedLine := normalizeEditWhitespace(line)
|
||||
if normalizedLine == normalizedFind {
|
||||
candidates = append(candidates, line)
|
||||
continue
|
||||
}
|
||||
if !strings.Contains(normalizedLine, normalizedFind) {
|
||||
continue
|
||||
}
|
||||
if subLinePattern == nil {
|
||||
words := strings.Fields(find)
|
||||
quoted := make([]string, len(words))
|
||||
for i, word := range words {
|
||||
quoted[i] = regexp.QuoteMeta(word)
|
||||
}
|
||||
pattern, err := regexp.Compile(strings.Join(quoted, `\s+`))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
subLinePattern = pattern
|
||||
}
|
||||
if match := subLinePattern.FindString(line); match != "" {
|
||||
candidates = append(candidates, match)
|
||||
}
|
||||
}
|
||||
|
||||
findLines := strings.Split(find, "\n")
|
||||
if len(findLines) > 1 {
|
||||
for i := 0; i+len(findLines) <= len(contentLines); i++ {
|
||||
block := strings.Join(contentLines[i:i+len(findLines)], "\n")
|
||||
if normalizeEditWhitespace(block) == normalizedFind {
|
||||
candidates = append(candidates, block)
|
||||
}
|
||||
}
|
||||
}
|
||||
return candidates
|
||||
}
|
||||
|
||||
// stripCommonIndentation removes the minimum leading-whitespace width shared
|
||||
// by all non-empty lines, so blocks match regardless of their nesting depth.
|
||||
func stripCommonIndentation(text string) string {
|
||||
lines := strings.Split(text, "\n")
|
||||
minIndent := -1
|
||||
for _, line := range lines {
|
||||
if strings.TrimSpace(line) == "" {
|
||||
continue
|
||||
}
|
||||
indent := len(line) - len(strings.TrimLeft(line, " \t"))
|
||||
if minIndent < 0 || indent < minIndent {
|
||||
minIndent = indent
|
||||
}
|
||||
}
|
||||
if minIndent <= 0 {
|
||||
return text
|
||||
}
|
||||
stripped := make([]string, len(lines))
|
||||
for i, line := range lines {
|
||||
if strings.TrimSpace(line) == "" {
|
||||
stripped[i] = line
|
||||
continue
|
||||
}
|
||||
stripped[i] = line[minIndent:]
|
||||
}
|
||||
return strings.Join(stripped, "\n")
|
||||
}
|
||||
|
||||
func indentationFlexibleReplacer(content, find string) []string {
|
||||
normalizedFind := stripCommonIndentation(find)
|
||||
contentLines := strings.Split(content, "\n")
|
||||
findLines := strings.Split(find, "\n")
|
||||
var candidates []string
|
||||
for i := 0; i+len(findLines) <= len(contentLines); i++ {
|
||||
block := strings.Join(contentLines[i:i+len(findLines)], "\n")
|
||||
if stripCommonIndentation(block) == normalizedFind {
|
||||
candidates = append(candidates, block)
|
||||
}
|
||||
}
|
||||
return candidates
|
||||
}
|
||||
|
||||
var editEscapeSequence = regexp.MustCompile("\\\\(n|t|r|'|\"|`|\\\\|\\n|\\$)")
|
||||
|
||||
// unescapeEditString undoes one level of string escaping (\n, \t, \", \\, a
|
||||
// backslash-newline continuation, \$) — the model sometimes reproduces file
|
||||
// content as it appeared inside a quoted string literal.
|
||||
func unescapeEditString(text string) string {
|
||||
return editEscapeSequence.ReplaceAllStringFunc(text, func(match string) string {
|
||||
switch match[1:] {
|
||||
case "n":
|
||||
return "\n"
|
||||
case "t":
|
||||
return "\t"
|
||||
case "r":
|
||||
return "\r"
|
||||
case "\n":
|
||||
return "\n"
|
||||
default:
|
||||
// ', ", `, \, $ all unescape to themselves.
|
||||
return match[1:]
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func escapeNormalizedReplacer(content, find string) []string {
|
||||
unescapedFind := unescapeEditString(find)
|
||||
var candidates []string
|
||||
if strings.Contains(content, unescapedFind) {
|
||||
candidates = append(candidates, unescapedFind)
|
||||
}
|
||||
contentLines := strings.Split(content, "\n")
|
||||
findLines := strings.Split(unescapedFind, "\n")
|
||||
for i := 0; i+len(findLines) <= len(contentLines); i++ {
|
||||
block := strings.Join(contentLines[i:i+len(findLines)], "\n")
|
||||
if unescapeEditString(block) == unescapedFind {
|
||||
candidates = append(candidates, block)
|
||||
}
|
||||
}
|
||||
return candidates
|
||||
}
|
||||
|
||||
// trimmedBoundaryReplacer tolerates stray leading/trailing whitespace (often
|
||||
// blank lines) around an otherwise-exact old_string.
|
||||
func trimmedBoundaryReplacer(content, find string) []string {
|
||||
trimmedFind := strings.TrimSpace(find)
|
||||
if trimmedFind == find || trimmedFind == "" {
|
||||
return nil
|
||||
}
|
||||
var candidates []string
|
||||
if strings.Contains(content, trimmedFind) {
|
||||
candidates = append(candidates, trimmedFind)
|
||||
}
|
||||
contentLines := strings.Split(content, "\n")
|
||||
findLines := strings.Split(find, "\n")
|
||||
for i := 0; i+len(findLines) <= len(contentLines); i++ {
|
||||
block := strings.Join(contentLines[i:i+len(findLines)], "\n")
|
||||
if strings.TrimSpace(block) == trimmedFind {
|
||||
candidates = append(candidates, block)
|
||||
}
|
||||
}
|
||||
return candidates
|
||||
}
|
||||
|
||||
// contextAwareReplacer anchors on the first and last lines and accepts an
|
||||
// equal-length block when at least half of its non-empty middle lines match
|
||||
// after trimming — a cheaper, stricter cousin of blockAnchorReplacer.
|
||||
func contextAwareReplacer(content, find string) []string {
|
||||
findLines := splitFindLines(find)
|
||||
if len(findLines) < 3 {
|
||||
return nil
|
||||
}
|
||||
contentLines := strings.Split(content, "\n")
|
||||
firstAnchor := strings.TrimSpace(findLines[0])
|
||||
lastAnchor := strings.TrimSpace(findLines[len(findLines)-1])
|
||||
var candidates []string
|
||||
for i := 0; i < len(contentLines); i++ {
|
||||
if strings.TrimSpace(contentLines[i]) != firstAnchor {
|
||||
continue
|
||||
}
|
||||
for j := i + 2; j < len(contentLines); j++ {
|
||||
if strings.TrimSpace(contentLines[j]) != lastAnchor {
|
||||
continue
|
||||
}
|
||||
if j-i+1 == len(findLines) {
|
||||
matching, nonEmpty := 0, 0
|
||||
for k := 1; k < len(findLines)-1; k++ {
|
||||
blockLine := strings.TrimSpace(contentLines[i+k])
|
||||
findLine := strings.TrimSpace(findLines[k])
|
||||
if blockLine == "" && findLine == "" {
|
||||
continue
|
||||
}
|
||||
nonEmpty++
|
||||
if blockLine == findLine {
|
||||
matching++
|
||||
}
|
||||
}
|
||||
if nonEmpty == 0 || float64(matching)/float64(nonEmpty) >= 0.5 {
|
||||
candidates = append(candidates, strings.Join(contentLines[i:j+1], "\n"))
|
||||
}
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
return candidates
|
||||
}
|
||||
|
||||
// adaptReplacementToSpan re-shapes the model's replacement to the span a
|
||||
// tolerant matcher resolved. When old_string only matched after normalization,
|
||||
// new_string was written at old_string's (wrong) shape, so applying it raw
|
||||
// would strip the file's indentation or drop a trailing CR:
|
||||
//
|
||||
// 1. Uniform re-indent: when every span line equals delta + the corresponding
|
||||
// find line's indentation (the line-trimmed / indentation-flexible shapes),
|
||||
// the same delta is prepended to every non-blank replacement line. Any
|
||||
// line that breaks the uniform-delta relationship disables the shift —
|
||||
// block-anchor matches with a drifted interior are left untouched.
|
||||
// 2. Trailing CR: a span from a CRLF file ends mid-line at "\r" (candidates
|
||||
// are built by joining lines split on "\n"); the replacement gets the same
|
||||
// trailing "\r" so the file's CRLF pairs stay intact.
|
||||
func adaptReplacementToSpan(span, find, replacement string) string {
|
||||
if delta, ok := uniformIndentDelta(span, find); ok && delta != "" {
|
||||
lines := strings.Split(replacement, "\n")
|
||||
for i, line := range lines {
|
||||
if strings.TrimSpace(line) == "" {
|
||||
continue
|
||||
}
|
||||
lines[i] = delta + line
|
||||
}
|
||||
replacement = strings.Join(lines, "\n")
|
||||
}
|
||||
if strings.HasSuffix(span, "\r") && !strings.HasSuffix(replacement, "\r") {
|
||||
replacement += "\r"
|
||||
}
|
||||
return replacement
|
||||
}
|
||||
|
||||
// uniformIndentDelta returns the indentation prefix that, prepended to every
|
||||
// non-blank find line, yields the corresponding span line's indentation. ok is
|
||||
// false when line counts differ, any line pair disagrees on the delta, or the
|
||||
// span is not simply a uniformly deeper-indented copy of find.
|
||||
func uniformIndentDelta(span, find string) (string, bool) {
|
||||
spanLines := strings.Split(span, "\n")
|
||||
findLines := splitFindLines(find)
|
||||
if len(spanLines) != len(findLines) {
|
||||
return "", false
|
||||
}
|
||||
leadingWhitespace := func(line string) string {
|
||||
return line[:len(line)-len(strings.TrimLeft(line, " \t"))]
|
||||
}
|
||||
delta := ""
|
||||
haveDelta := false
|
||||
for i := range findLines {
|
||||
spanLine := strings.TrimSuffix(spanLines[i], "\r")
|
||||
findLine := strings.TrimSuffix(findLines[i], "\r")
|
||||
if strings.TrimSpace(spanLine) == "" && strings.TrimSpace(findLine) == "" {
|
||||
continue
|
||||
}
|
||||
spanIndent := leadingWhitespace(spanLine)
|
||||
findIndent := leadingWhitespace(findLine)
|
||||
if !strings.HasSuffix(spanIndent, findIndent) {
|
||||
return "", false
|
||||
}
|
||||
lineDelta := spanIndent[:len(spanIndent)-len(findIndent)]
|
||||
if !haveDelta {
|
||||
delta = lineDelta
|
||||
haveDelta = true
|
||||
continue
|
||||
}
|
||||
if lineDelta != delta {
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
return delta, haveDelta
|
||||
}
|
||||
|
||||
// levenshtein computes edit distance with a two-row rolling matrix.
|
||||
func levenshtein(a, b string) int {
|
||||
if a == "" {
|
||||
return len(b)
|
||||
}
|
||||
if b == "" {
|
||||
return len(a)
|
||||
}
|
||||
previous := make([]int, len(b)+1)
|
||||
current := make([]int, len(b)+1)
|
||||
for j := range previous {
|
||||
previous[j] = j
|
||||
}
|
||||
for i := 1; i <= len(a); i++ {
|
||||
current[0] = i
|
||||
for j := 1; j <= len(b); j++ {
|
||||
cost := 1
|
||||
if a[i-1] == b[j-1] {
|
||||
cost = 0
|
||||
}
|
||||
minimum := previous[j] + 1
|
||||
if current[j-1]+1 < minimum {
|
||||
minimum = current[j-1] + 1
|
||||
}
|
||||
if previous[j-1]+cost < minimum {
|
||||
minimum = previous[j-1] + cost
|
||||
}
|
||||
current[j] = minimum
|
||||
}
|
||||
previous, current = current, previous
|
||||
}
|
||||
return previous[len(b)]
|
||||
}
|
||||
@@ -1,369 +0,0 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func runEdit(t *testing.T, dir, initial string, args map[string]any) (Result, string) {
|
||||
t.Helper()
|
||||
path := filepath.Join(dir, "target.txt")
|
||||
if err := os.WriteFile(path, []byte(initial), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if args["path"] == nil {
|
||||
args["path"] = "target.txt"
|
||||
}
|
||||
result := NewEditFileTool(dir).Run(context.Background(), args)
|
||||
after, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return result, string(after)
|
||||
}
|
||||
|
||||
func TestEditExactMatchStillPreferred(t *testing.T) {
|
||||
// The fast path must keep byte-exact semantics: an exact match wins even
|
||||
// when fuzzy strategies would also find candidates elsewhere.
|
||||
result, after := runEdit(t, t.TempDir(), "alpha\nbeta\ngamma\n", map[string]any{
|
||||
"old_string": "beta",
|
||||
"new_string": "delta",
|
||||
})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("expected ok, got %q", result.Output)
|
||||
}
|
||||
if after != "alpha\ndelta\ngamma\n" {
|
||||
t.Fatalf("unexpected content: %q", after)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditFuzzyLineTrimmed(t *testing.T) {
|
||||
// Model reproduced the lines without the file's leading indentation; the
|
||||
// line-trimmed strategy must find the real indented span and preserve the
|
||||
// file's own indentation on the replaced region boundary.
|
||||
initial := "func a() {\n\tx := 1\n\treturn x\n}\n"
|
||||
result, after := runEdit(t, t.TempDir(), initial, map[string]any{
|
||||
"old_string": "x := 1\nreturn x",
|
||||
"new_string": "y := 2\nreturn y",
|
||||
})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("expected ok, got %q", result.Output)
|
||||
}
|
||||
// new_string was written at old_string's outdented shape; the span's
|
||||
// indentation must be re-applied so the file keeps its tabs.
|
||||
if !strings.Contains(after, "\ty := 2\n\treturn y\n}") {
|
||||
t.Fatalf("unexpected content: %q", after)
|
||||
}
|
||||
if strings.Contains(after, "x := 1") {
|
||||
t.Fatalf("old span survived: %q", after)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditFuzzyWhitespaceNormalizedSingleLine(t *testing.T) {
|
||||
// Whitespace runs collapsed: "x := 1" in the file, single spaces from
|
||||
// the model.
|
||||
initial := "start\nx := 1\nend\n"
|
||||
result, after := runEdit(t, t.TempDir(), initial, map[string]any{
|
||||
"old_string": "x := 1",
|
||||
"new_string": "x := 2",
|
||||
})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("expected ok, got %q", result.Output)
|
||||
}
|
||||
if !strings.Contains(after, "x := 2") || strings.Contains(after, ":= 1") {
|
||||
t.Fatalf("unexpected content: %q", after)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditFuzzyBlockAnchorToleratesMiddleDrift(t *testing.T) {
|
||||
// First and last lines anchor the block; one interior line differs slightly
|
||||
// (comment text drifted). Levenshtein similarity keeps it above 0.65.
|
||||
initial := strings.Join([]string{
|
||||
"func handler(w http.ResponseWriter, r *http.Request) {",
|
||||
"\t// write the response body to the client",
|
||||
"\tw.WriteHeader(http.StatusOK)",
|
||||
"\tfmt.Fprint(w, \"done\")",
|
||||
"}",
|
||||
"",
|
||||
}, "\n")
|
||||
find := strings.Join([]string{
|
||||
"func handler(w http.ResponseWriter, r *http.Request) {",
|
||||
"\t// write the response body to client",
|
||||
"\tw.WriteHeader(http.StatusOK)",
|
||||
"\tfmt.Fprint(w, \"done\")",
|
||||
"}",
|
||||
}, "\n")
|
||||
result, after := runEdit(t, t.TempDir(), initial, map[string]any{
|
||||
"old_string": find,
|
||||
"new_string": "func handler(w http.ResponseWriter, r *http.Request) {\n\tw.WriteHeader(http.StatusNoContent)\n}",
|
||||
})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("expected ok, got %q", result.Output)
|
||||
}
|
||||
if !strings.Contains(after, "StatusNoContent") || strings.Contains(after, "StatusOK") {
|
||||
t.Fatalf("unexpected content: %q", after)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditFuzzyIndentationFlexible(t *testing.T) {
|
||||
// The whole block sits one nesting level deeper in the file than in the
|
||||
// model's old_string, with interior relative indentation preserved.
|
||||
initial := "if ok {\n\t\tfor i := range xs {\n\t\t\tsum += xs[i]\n\t\t}\n}\n"
|
||||
result, after := runEdit(t, t.TempDir(), initial, map[string]any{
|
||||
"old_string": "for i := range xs {\n\tsum += xs[i]\n}",
|
||||
"new_string": "sum += total(xs)",
|
||||
})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("expected ok, got %q", result.Output)
|
||||
}
|
||||
// The replacement lands at the span's real depth (two tabs), not at
|
||||
// old_string's outdented depth.
|
||||
if !strings.Contains(after, "\t\tsum += total(xs)") || strings.Contains(after, "range xs") {
|
||||
t.Fatalf("unexpected content: %q", after)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditFuzzyCRLFLineTrimmedPreservesCarriageReturns(t *testing.T) {
|
||||
// CRLF file + a multi-line outdented old_string: the resolved span carries
|
||||
// tabs and trailing CRs; the replacement must inherit both so the file's
|
||||
// CRLF pairs and indentation survive the edit.
|
||||
initial := "func a() {\r\n\tx := 1\r\n\treturn x\r\n}\r\n"
|
||||
result, after := runEdit(t, t.TempDir(), initial, map[string]any{
|
||||
"old_string": "x := 1\nreturn x",
|
||||
"new_string": "y := 2\nreturn y",
|
||||
})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("expected ok, got %q", result.Output)
|
||||
}
|
||||
if !strings.Contains(after, "\ty := 2\r\n\treturn y\r\n}") {
|
||||
t.Fatalf("unexpected content: %q", after)
|
||||
}
|
||||
if strings.Contains(after, "\n\treturn y\n") {
|
||||
t.Fatalf("bare LF introduced into CRLF file: %q", after)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUniformIndentDelta(t *testing.T) {
|
||||
cases := []struct {
|
||||
span, find string
|
||||
want string
|
||||
ok bool
|
||||
}{
|
||||
{"\tx := 1\n\treturn x", "x := 1\nreturn x", "\t", true},
|
||||
{"\t\tfor {\n\t\t\tgo()\n\t\t}", "for {\n\tgo()\n}", "\t\t", true},
|
||||
{"x := 1", "x := 1", "", true}, // no shift needed
|
||||
{"\tx := 1\n return x", "x := 1\nreturn x", "", false}, // mixed delta
|
||||
{"\ta\n\tb\n\tc", "a\nb", "", false}, // line-count mismatch
|
||||
{" a", "\ta", "", false}, // find indent not a suffix
|
||||
}
|
||||
for _, c := range cases {
|
||||
got, ok := uniformIndentDelta(c.span, c.find)
|
||||
if ok != c.ok || (ok && got != c.want) {
|
||||
t.Errorf("uniformIndentDelta(%q, %q) = %q,%v want %q,%v", c.span, c.find, got, ok, c.want, c.ok)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdaptReplacementLeavesNonUniformSpansAlone(t *testing.T) {
|
||||
// Block-anchor style match with a drifted interior: no uniform delta, so
|
||||
// the replacement must pass through untouched.
|
||||
span := "\tfunc h() {\n\t\t// drifted comment\n\t}"
|
||||
find := "func h() {\n// different comment\n}"
|
||||
if got := adaptReplacementToSpan(span, find, "replacement()"); got != "replacement()" {
|
||||
t.Fatalf("non-uniform span must not re-indent: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditFuzzyEscapeNormalized(t *testing.T) {
|
||||
// Model reproduced the line as it would appear inside a quoted string
|
||||
// literal, with escaped quotes.
|
||||
initial := "console.log(\"hello world\")\n"
|
||||
result, after := runEdit(t, t.TempDir(), initial, map[string]any{
|
||||
"old_string": `console.log(\"hello world\")`,
|
||||
"new_string": `console.log("goodbye")`,
|
||||
})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("expected ok, got %q", result.Output)
|
||||
}
|
||||
if !strings.Contains(after, `console.log("goodbye")`) {
|
||||
t.Fatalf("unexpected content: %q", after)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditFuzzyTrimmedBoundary(t *testing.T) {
|
||||
// Stray whitespace around an otherwise-exact old_string; " two " is not a
|
||||
// substring of the file, so only the fuzzy cascade can resolve it.
|
||||
initial := "one\ntwo\nthree\n"
|
||||
result, after := runEdit(t, t.TempDir(), initial, map[string]any{
|
||||
"old_string": " two ",
|
||||
"new_string": "2",
|
||||
})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("expected ok, got %q", result.Output)
|
||||
}
|
||||
if after != "one\n2\nthree\n" {
|
||||
t.Fatalf("unexpected content: %q", after)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTrimmedBoundaryReplacerYieldsTrimmedSpan(t *testing.T) {
|
||||
candidates := trimmedBoundaryReplacer("one\ntwo\nthree\n", "\ntwo\n")
|
||||
if len(candidates) == 0 || candidates[0] != "two" {
|
||||
t.Fatalf("expected trimmed candidate \"two\", got %v", candidates)
|
||||
}
|
||||
if trimmedBoundaryReplacer("anything", "already-trimmed") != nil {
|
||||
t.Fatal("already-trimmed find must yield no candidates")
|
||||
}
|
||||
}
|
||||
|
||||
// Two identical blocks that only match old_string after indentation-tolerant
|
||||
// matching: exact search finds nothing, fuzzy finds both.
|
||||
const ambiguousBlocks = "\tif x {\n\t\tgo()\n\t}\nmid\n\tif x {\n\t\tgo()\n\t}\n"
|
||||
|
||||
func TestEditFuzzyAmbiguousReportsError(t *testing.T) {
|
||||
// No strategy can disambiguate identical blocks, so the tool must refuse
|
||||
// rather than guess, and the file must be untouched.
|
||||
result, after := runEdit(t, t.TempDir(), ambiguousBlocks, map[string]any{
|
||||
"old_string": "if x {\n\tgo()\n}",
|
||||
"new_string": "stop()",
|
||||
})
|
||||
if result.Status != StatusError || !strings.Contains(result.Output, "multiple locations") {
|
||||
t.Fatalf("expected ambiguity error, got %q", result.Output)
|
||||
}
|
||||
if after != ambiguousBlocks {
|
||||
t.Fatalf("file must be unchanged, got %q", after)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditFuzzyReplaceAll(t *testing.T) {
|
||||
// replace_all applies the fuzzy-resolved span at every occurrence.
|
||||
result, after := runEdit(t, t.TempDir(), ambiguousBlocks, map[string]any{
|
||||
"old_string": "if x {\n\tgo()\n}",
|
||||
"new_string": "stop()",
|
||||
"replace_all": true,
|
||||
})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("expected ok, got %q", result.Output)
|
||||
}
|
||||
if strings.Count(after, "stop()") != 2 || strings.Contains(after, "go()") {
|
||||
t.Fatalf("unexpected content: %q", after)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditFuzzyNotFoundKeepsExactError(t *testing.T) {
|
||||
result, after := runEdit(t, t.TempDir(), "alpha\n", map[string]any{
|
||||
"old_string": "omega",
|
||||
"new_string": "x",
|
||||
})
|
||||
if result.Status != StatusError || !strings.Contains(result.Output, "Could not find the exact string") {
|
||||
t.Fatalf("expected not-found error, got %q", result.Output)
|
||||
}
|
||||
if after != "alpha\n" {
|
||||
t.Fatalf("file must be unchanged, got %q", after)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditFuzzyCRLFFile(t *testing.T) {
|
||||
// CRLF file + LF old_string with indentation drift: the CRLF translation
|
||||
// feeds the cascade and the replacement preserves CRLF endings.
|
||||
initial := "func a() {\r\n\tx := 1\r\n}\r\n"
|
||||
result, after := runEdit(t, t.TempDir(), initial, map[string]any{
|
||||
"old_string": "x := 1",
|
||||
"new_string": "x := 2",
|
||||
})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("expected ok, got %q", result.Output)
|
||||
}
|
||||
if !strings.Contains(after, "\tx := 2\r\n") {
|
||||
t.Fatalf("unexpected content: %q", after)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsDisproportionateEditMatch(t *testing.T) {
|
||||
// A candidate span that dwarfs old_string must be refused: anchors bridging
|
||||
// unrelated code would otherwise delete it all.
|
||||
big := strings.Repeat("line\n", 10)
|
||||
if !isDisproportionateEditMatch(big, "a\nb\nc") {
|
||||
t.Fatal("10-line span for 3-line find must be disproportionate")
|
||||
}
|
||||
if isDisproportionateEditMatch("a\nb\nc\nd", "a\nb\nc") {
|
||||
t.Fatal("one extra line within tolerance must be allowed")
|
||||
}
|
||||
if isDisproportionateEditMatch("single line", "single") {
|
||||
t.Fatal("single-line finds are exempt from the byte-length guard")
|
||||
}
|
||||
if !isDisproportionateEditMatch("xx\n"+strings.Repeat("y", 900)+"\nzz", "xx\nab\nzz") {
|
||||
t.Fatal("byte-length blowup must be disproportionate")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLevenshtein(t *testing.T) {
|
||||
cases := []struct {
|
||||
a, b string
|
||||
want int
|
||||
}{
|
||||
{"", "", 0},
|
||||
{"abc", "", 3},
|
||||
{"", "abc", 3},
|
||||
{"kitten", "sitting", 3},
|
||||
{"same", "same", 0},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := levenshtein(c.a, c.b); got != c.want {
|
||||
t.Fatalf("levenshtein(%q,%q) = %d, want %d", c.a, c.b, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditFuzzyDistinctCandidatesAreAmbiguous(t *testing.T) {
|
||||
// Two same-content blocks at DIFFERENT indentation: each resolved span is
|
||||
// distinct and occurs exactly once, so the old literal-uniqueness check
|
||||
// passed and silently edited the first block. Distinct-shaped candidates
|
||||
// from one strategy must instead report ambiguity and leave the file alone.
|
||||
initial := "\tif x {\n\t\tgo()\n\t}\nmid\n\t\tif x {\n\t\t\tgo()\n\t\t}\n"
|
||||
result, after := runEdit(t, t.TempDir(), initial, map[string]any{
|
||||
"old_string": "if x {\n\tgo()\n}",
|
||||
"new_string": "stop()",
|
||||
})
|
||||
if result.Status != StatusError || !strings.Contains(result.Output, "multiple locations") {
|
||||
t.Fatalf("distinct fuzzy candidates must be ambiguous, got %q", result.Output)
|
||||
}
|
||||
if after != initial {
|
||||
t.Fatalf("file must be unchanged, got %q", after)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditFuzzyBlockAnchorTwoPlausibleBlocksAmbiguous(t *testing.T) {
|
||||
// Two blocks share the same first/last anchor lines and BOTH interiors sit
|
||||
// above the similarity threshold. Picking the "best" one would silently
|
||||
// edit a block the model may not have meant; both must surface so the
|
||||
// cascade reports ambiguity and the file stays untouched.
|
||||
initial := strings.Join([]string{
|
||||
"func setup(cfg Config) {",
|
||||
"\tvalue := compute(alpha, beta, gamma)",
|
||||
"}",
|
||||
"",
|
||||
"func setup(cfg Config) {",
|
||||
"\tvalue := compute(alpha, delta, gamma)",
|
||||
"}",
|
||||
"",
|
||||
}, "\n")
|
||||
find := strings.Join([]string{
|
||||
"func setup(cfg Config) {",
|
||||
"\tvalue := compute(alpha, omega, gamma)",
|
||||
"}",
|
||||
}, "\n")
|
||||
result, after := runEdit(t, t.TempDir(), initial, map[string]any{
|
||||
"old_string": find,
|
||||
"new_string": "func setup(cfg Config) {}",
|
||||
})
|
||||
if result.Status != StatusError || !strings.Contains(result.Output, "multiple locations") {
|
||||
t.Fatalf("two plausible anchor blocks must be ambiguous, got %q", result.Output)
|
||||
}
|
||||
if after != initial {
|
||||
t.Fatalf("file must be unchanged, got %q", after)
|
||||
}
|
||||
}
|
||||
@@ -908,15 +908,6 @@ func formatExecCommandOutput(output string, sessionID int, exited bool, exitCode
|
||||
}
|
||||
|
||||
func truncateExecOutput(output string, maxOutputTokens int) (string, bool) {
|
||||
return truncateExecOutputSpill(output, maxOutputTokens, "exec_command")
|
||||
}
|
||||
|
||||
// truncateExecOutputSpill keeps a head/tail window of the output within the
|
||||
// token budget and, on truncation, spills the full output to disk so the model
|
||||
// can grep/read the elided middle instead of re-running the command with a
|
||||
// bigger budget. The spill is best-effort: when it fails the notice simply
|
||||
// omits the file hint.
|
||||
func truncateExecOutputSpill(output string, maxOutputTokens int, toolName string) (string, bool) {
|
||||
if maxOutputTokens <= 0 {
|
||||
maxOutputTokens = defaultMaxOutputTokens
|
||||
}
|
||||
@@ -924,13 +915,9 @@ func truncateExecOutputSpill(output string, maxOutputTokens int, toolName string
|
||||
if len(output) <= maxBytes {
|
||||
return output, false
|
||||
}
|
||||
notice := "\n[zero] output truncated\n"
|
||||
if spillPath := spillTruncatedOutput(toolName, output); spillPath != "" {
|
||||
notice = "\n[zero] output truncated — full output saved to " + spillPath + " (grep or read_file it instead of re-running)\n"
|
||||
}
|
||||
head := maxBytes / 2
|
||||
tail := maxBytes - head
|
||||
return utf8Prefix(output, head) + notice + utf8Suffix(output, tail), true
|
||||
return utf8Prefix(output, head) + "\n[zero] output truncated\n" + utf8Suffix(output, tail), true
|
||||
}
|
||||
|
||||
func utf8Prefix(value string, maxBytes int) string {
|
||||
|
||||
@@ -1,97 +0,0 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Format-on-write for the mutating file tools. When enabled, a successful
|
||||
// edit_file/write_file runs the language's standard formatter on the file it
|
||||
// just wrote, so the model's output always lands in project-canonical style
|
||||
// and never fails a CI format check it cannot see. Off by default (set
|
||||
// ZERO_FORMAT_ON_WRITE=1): auto-reformatting changes bytes the model did not
|
||||
// write, which strict workflows may not want.
|
||||
//
|
||||
// Ordering matters: formatting runs BEFORE the FileTracker re-baseline, and
|
||||
// the caller records the POST-format content. Formatting after the baseline
|
||||
// would make the very next edit look like an external modification and trip
|
||||
// the conflict guard.
|
||||
|
||||
// formatOnWriteTimeout bounds one formatter run; a wedged formatter must never
|
||||
// hang a tool call. On timeout the unformatted write stands.
|
||||
const formatOnWriteTimeout = 10 * time.Second
|
||||
|
||||
// formatterCommands maps a file extension to the formatter argv; the file path
|
||||
// is appended as the final argument. Only in-place, config-respecting,
|
||||
// community-standard formatters — a missing binary silently skips formatting.
|
||||
var formatterCommands = map[string][]string{
|
||||
".go": {"gofmt", "-w"},
|
||||
".rs": {"rustfmt"},
|
||||
".py": {"ruff", "format", "--quiet"},
|
||||
".ts": {"prettier", "--log-level", "silent", "--write"},
|
||||
".tsx": {"prettier", "--log-level", "silent", "--write"},
|
||||
".js": {"prettier", "--log-level", "silent", "--write"},
|
||||
".jsx": {"prettier", "--log-level", "silent", "--write"},
|
||||
".json": {"prettier", "--log-level", "silent", "--write"},
|
||||
".css": {"prettier", "--log-level", "silent", "--write"},
|
||||
".scss": {"prettier", "--log-level", "silent", "--write"},
|
||||
".html": {"prettier", "--log-level", "silent", "--write"},
|
||||
".md": {"prettier", "--log-level", "silent", "--write"},
|
||||
".yaml": {"prettier", "--log-level", "silent", "--write"},
|
||||
".yml": {"prettier", "--log-level", "silent", "--write"},
|
||||
".zig": {"zig", "fmt"},
|
||||
".dart": {"dart", "format"},
|
||||
".tf": {"terraform", "fmt"},
|
||||
".gleam": {"gleam", "format"},
|
||||
".sh": {"shfmt", "-w"},
|
||||
".bash": {"shfmt", "-w"},
|
||||
".c": {"clang-format", "-i"},
|
||||
".h": {"clang-format", "-i"},
|
||||
".cpp": {"clang-format", "-i"},
|
||||
".hpp": {"clang-format", "-i"},
|
||||
".cc": {"clang-format", "-i"},
|
||||
".kt": {"ktlint", "-F"},
|
||||
".swift": {"swiftformat"},
|
||||
".lua": {"stylua"},
|
||||
}
|
||||
|
||||
// formatOnWriteEnabled reports whether the opt-in env toggle is set.
|
||||
func formatOnWriteEnabled() bool {
|
||||
value := strings.TrimSpace(os.Getenv("ZERO_FORMAT_ON_WRITE"))
|
||||
return value != "" && value != "0" && !strings.EqualFold(value, "false")
|
||||
}
|
||||
|
||||
// maybeFormatWrittenFile runs the configured formatter for absolutePath (when
|
||||
// enabled and on PATH) and returns the file's content afterwards. Best-effort
|
||||
// throughout: any failure — no formatter, formatter error, timeout, unreadable
|
||||
// result — returns writtenContent so the caller's state matches the last write
|
||||
// it performed itself.
|
||||
func maybeFormatWrittenFile(ctx context.Context, absolutePath string, writtenContent string) string {
|
||||
if !formatOnWriteEnabled() {
|
||||
return writtenContent
|
||||
}
|
||||
command, ok := formatterCommands[strings.ToLower(filepath.Ext(absolutePath))]
|
||||
if !ok {
|
||||
return writtenContent
|
||||
}
|
||||
if _, err := exec.LookPath(command[0]); err != nil {
|
||||
return writtenContent
|
||||
}
|
||||
formatCtx, cancel := context.WithTimeout(ctx, formatOnWriteTimeout)
|
||||
defer cancel()
|
||||
arguments := append(append([]string(nil), command[1:]...), absolutePath)
|
||||
formatter := exec.CommandContext(formatCtx, command[0], arguments...)
|
||||
formatter.Dir = filepath.Dir(absolutePath)
|
||||
if err := formatter.Run(); err != nil {
|
||||
return writtenContent
|
||||
}
|
||||
formatted, err := os.ReadFile(absolutePath)
|
||||
if err != nil {
|
||||
return writtenContent
|
||||
}
|
||||
return string(formatted)
|
||||
}
|
||||
@@ -1,82 +0,0 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// gofmt ships with the Go toolchain, so it is the one formatter guaranteed to
|
||||
// exist wherever these tests run.
|
||||
func requireGofmt(t *testing.T) {
|
||||
t.Helper()
|
||||
if _, err := exec.LookPath("gofmt"); err != nil {
|
||||
t.Skip("gofmt not on PATH")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatOnWriteDisabledByDefault(t *testing.T) {
|
||||
requireGofmt(t)
|
||||
t.Setenv("ZERO_FORMAT_ON_WRITE", "")
|
||||
dir := t.TempDir()
|
||||
ugly := "package a\n\nfunc A( ) { }\n"
|
||||
|
||||
result := NewWriteFileTool(dir).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{
|
||||
"path": "a.go",
|
||||
"content": ugly,
|
||||
}, RunOptions{})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("write failed: %q", result.Output)
|
||||
}
|
||||
onDisk, err := os.ReadFile(filepath.Join(dir, "a.go"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(onDisk) != ugly {
|
||||
t.Fatalf("formatting must be off by default, got %q", onDisk)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatOnWriteFormatsAndKeepsTrackerConsistent(t *testing.T) {
|
||||
requireGofmt(t)
|
||||
t.Setenv("ZERO_FORMAT_ON_WRITE", "1")
|
||||
dir := t.TempDir()
|
||||
tracker := NewFileTracker()
|
||||
|
||||
write := NewWriteFileTool(dir).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{
|
||||
"path": "a.go",
|
||||
"content": "package a\n\nfunc A( ) { }\n",
|
||||
}, RunOptions{FileTracker: tracker})
|
||||
if write.Status != StatusOK {
|
||||
t.Fatalf("write failed: %q", write.Output)
|
||||
}
|
||||
onDisk, err := os.ReadFile(filepath.Join(dir, "a.go"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(string(onDisk), "func A() {") {
|
||||
t.Fatalf("expected gofmt-formatted content, got %q", onDisk)
|
||||
}
|
||||
|
||||
// The tracker must have been re-baselined to the POST-format content: a
|
||||
// follow-up edit must not trip the external-modification conflict guard.
|
||||
edit := NewEditFileTool(dir).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{
|
||||
"path": "a.go",
|
||||
"old_string": "func A() {",
|
||||
"new_string": "func B() {",
|
||||
}, RunOptions{FileTracker: tracker})
|
||||
if edit.Status != StatusOK {
|
||||
t.Fatalf("follow-up edit must not conflict after formatting: %q", edit.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatOnWriteSkipsUnknownExtensions(t *testing.T) {
|
||||
t.Setenv("ZERO_FORMAT_ON_WRITE", "1")
|
||||
content := maybeFormatWrittenFile(context.Background(), filepath.Join(t.TempDir(), "notes.xyz"), "raw text")
|
||||
if content != "raw text" {
|
||||
t.Fatalf("unknown extension must pass through: %q", content)
|
||||
}
|
||||
}
|
||||
+7
-16
@@ -45,18 +45,18 @@ func NewScopedGlobTool(workspaceRoot string, scope PathScope) Tool {
|
||||
}
|
||||
}
|
||||
|
||||
func (tool globTool) Run(ctx context.Context, args map[string]any) Result {
|
||||
return tool.runWith(ctx, args, readExcluder{})
|
||||
func (tool globTool) Run(_ context.Context, args map[string]any) Result {
|
||||
return tool.runWith(args, readExcluder{})
|
||||
}
|
||||
|
||||
// RunWithSandbox runs glob while skipping subtrees the sandbox policy denies
|
||||
// reads to (DenyRead). With no DenyRead configured the excluder is a no-op and
|
||||
// behavior is unchanged.
|
||||
func (tool globTool) RunWithSandbox(ctx context.Context, args map[string]any, engine *sandbox.Engine) Result {
|
||||
return tool.runWith(ctx, args, sandboxReadExcluder(engine))
|
||||
func (tool globTool) RunWithSandbox(_ context.Context, args map[string]any, engine *sandbox.Engine) Result {
|
||||
return tool.runWith(args, sandboxReadExcluder(engine))
|
||||
}
|
||||
|
||||
func (tool globTool) runWith(ctx context.Context, args map[string]any, exclude readExcluder) Result {
|
||||
func (tool globTool) runWith(args map[string]any, exclude readExcluder) Result {
|
||||
pattern, err := aliasedStringArg(args, []string{"pattern", "glob", "match", "query", "expression"}, "", true, false)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for glob: " + err.Error())
|
||||
@@ -88,11 +88,8 @@ func (tool globTool) runWith(ctx context.Context, args map[string]any, exclude r
|
||||
return errorResult("Error running glob " + fmt.Sprintf("%q", pattern) + ": " + err.Error())
|
||||
}
|
||||
|
||||
matches, err := scanGlob(ctx, root, displayRoot, matcher, includeDirs, exclude)
|
||||
matches, err := scanGlob(root, displayRoot, matcher, includeDirs, exclude)
|
||||
if err != nil {
|
||||
if res, ok := searchCancelledResult("glob", err); ok {
|
||||
return res
|
||||
}
|
||||
return errorResult("Error running glob " + fmt.Sprintf("%q", pattern) + ": " + err.Error())
|
||||
}
|
||||
if len(matches) == 0 {
|
||||
@@ -129,15 +126,9 @@ func (tool globTool) runWith(ctx context.Context, args map[string]any, exclude r
|
||||
}
|
||||
}
|
||||
|
||||
func scanGlob(ctx context.Context, root string, displayRoot string, matcher *regexp.Regexp, includeDirs bool, exclude readExcluder) ([]string, error) {
|
||||
func scanGlob(root string, displayRoot string, matcher *regexp.Regexp, includeDirs bool, exclude readExcluder) ([]string, error) {
|
||||
matches := []string{}
|
||||
err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, walkErr error) error {
|
||||
// Checked first, ahead of walkErr: an unscoped scan over a large tree can
|
||||
// run long enough that cancelling the run must stop the walk promptly
|
||||
// rather than visiting every remaining entry to completion.
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
if walkErr != nil {
|
||||
if path == root {
|
||||
return walkErr
|
||||
|
||||
+10
-33
@@ -55,19 +55,19 @@ func NewScopedGrepTool(workspaceRoot string, scope PathScope) Tool {
|
||||
}
|
||||
}
|
||||
|
||||
func (tool grepTool) Run(ctx context.Context, args map[string]any) Result {
|
||||
return tool.runWith(ctx, args, readExcluder{})
|
||||
func (tool grepTool) Run(_ context.Context, args map[string]any) Result {
|
||||
return tool.runWith(args, readExcluder{})
|
||||
}
|
||||
|
||||
// RunWithSandbox runs the search while skipping subtrees the sandbox policy
|
||||
// denies reads to (DenyRead), so grep never surfaces content from a read-denied
|
||||
// path. With no DenyRead configured the excluder is a no-op and behavior is
|
||||
// unchanged.
|
||||
func (tool grepTool) RunWithSandbox(ctx context.Context, args map[string]any, engine *sandbox.Engine) Result {
|
||||
return tool.runWith(ctx, args, sandboxReadExcluder(engine))
|
||||
func (tool grepTool) RunWithSandbox(_ context.Context, args map[string]any, engine *sandbox.Engine) Result {
|
||||
return tool.runWith(args, sandboxReadExcluder(engine))
|
||||
}
|
||||
|
||||
func (tool grepTool) runWith(ctx context.Context, args map[string]any, exclude readExcluder) Result {
|
||||
func (tool grepTool) runWith(args map[string]any, exclude readExcluder) Result {
|
||||
pattern, err := aliasedStringArg(args, []string{"pattern", "query", "regex", "search", "expression"}, "", true, false)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for grep: " + err.Error())
|
||||
@@ -138,24 +138,15 @@ func (tool grepTool) runWith(ctx context.Context, args map[string]any, exclude r
|
||||
}
|
||||
}
|
||||
|
||||
files, err := grepFiles(ctx, resolvedRoot, target, globMatcher, exclude)
|
||||
files, err := grepFiles(resolvedRoot, target, globMatcher, exclude)
|
||||
if err != nil {
|
||||
if res, ok := searchCancelledResult("grep", err); ok {
|
||||
return res
|
||||
}
|
||||
return errorResult("Error running grep: " + err.Error())
|
||||
}
|
||||
|
||||
// resolveScopedPath returns an absolute displayRoot when the target resolved
|
||||
// to an extra (non-workspace) granted root; emit absolute match paths there so
|
||||
// they survive a round-trip through read_file/edit_file (mirrors glob).
|
||||
matches, err := collectGrepMatches(ctx, resolvedRoot, filepath.IsAbs(displayRoot), files, compiled)
|
||||
if err != nil {
|
||||
if res, ok := searchCancelledResult("grep", err); ok {
|
||||
return res
|
||||
}
|
||||
return errorResult("Error running grep: " + err.Error())
|
||||
}
|
||||
matches := collectGrepMatches(resolvedRoot, filepath.IsAbs(displayRoot), files, compiled)
|
||||
if len(matches) == 0 {
|
||||
if outputMode == "count" {
|
||||
return okResult("0 matches found")
|
||||
@@ -270,7 +261,7 @@ func confineGrepFile(resolvedRoot string, path string) (string, string, bool) {
|
||||
return filepath.ToSlash(relative), resolved, true
|
||||
}
|
||||
|
||||
func grepFiles(ctx context.Context, resolvedRoot string, target string, globMatcher *regexp.Regexp, exclude readExcluder) ([]string, error) {
|
||||
func grepFiles(resolvedRoot string, target string, globMatcher *regexp.Regexp, exclude readExcluder) ([]string, error) {
|
||||
info, err := os.Stat(target)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -297,13 +288,6 @@ func grepFiles(ctx context.Context, resolvedRoot string, target string, globMatc
|
||||
|
||||
files := []string{}
|
||||
err = filepath.WalkDir(target, func(path string, entry os.DirEntry, walkErr error) error {
|
||||
// Checked first, ahead of walkErr: an unscoped search over a large tree
|
||||
// (e.g. a broad parent directory, not just the workspace) can run long
|
||||
// enough that cancelling the run must stop the walk promptly rather than
|
||||
// visiting every remaining entry to completion.
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
if walkErr != nil {
|
||||
if path == target {
|
||||
return walkErr
|
||||
@@ -358,16 +342,9 @@ func grepFiles(ctx context.Context, resolvedRoot string, target string, globMatc
|
||||
return files, nil
|
||||
}
|
||||
|
||||
func collectGrepMatches(ctx context.Context, resolvedRoot string, absolutePaths bool, files []string, compiled *regexp.Regexp) ([]grepMatch, error) {
|
||||
func collectGrepMatches(resolvedRoot string, absolutePaths bool, files []string, compiled *regexp.Regexp) []grepMatch {
|
||||
matches := []grepMatch{}
|
||||
for _, file := range files {
|
||||
// Checked per file, matching the walk's per-entry check in grepFiles: a
|
||||
// broad search can already have collected many files before the run is
|
||||
// cancelled, and reading/regex-scanning each one is real work that must
|
||||
// stop promptly rather than running to completion.
|
||||
if err := ctx.Err(); err != nil {
|
||||
return matches, err
|
||||
}
|
||||
// Re-confine at read time (defense-in-depth) AND to compute the clean
|
||||
// workspace-relative path used in output.
|
||||
relative, resolvedPath, ok := confineGrepFile(resolvedRoot, file)
|
||||
@@ -404,5 +381,5 @@ func collectGrepMatches(ctx context.Context, resolvedRoot string, absolutePaths
|
||||
})
|
||||
}
|
||||
}
|
||||
return matches, nil
|
||||
return matches
|
||||
}
|
||||
@@ -1,18 +0,0 @@
|
||||
package tools
|
||||
|
||||
import "context"
|
||||
|
||||
// inlineDiagnostics renders the post-write diagnostics block a mutating tool
|
||||
// appends to its output. Empty when no Diagnostics callback is wired, the file
|
||||
// is clean, or no language server is available — the tool output is then
|
||||
// byte-identical to the pre-diagnostics behavior.
|
||||
func inlineDiagnostics(ctx context.Context, options RunOptions, absolutePath, relativePath string) string {
|
||||
if options.Diagnostics == nil {
|
||||
return ""
|
||||
}
|
||||
block := options.Diagnostics(ctx, absolutePath)
|
||||
if block == "" {
|
||||
return ""
|
||||
}
|
||||
return "\n\nDiagnostics in " + relativePath + " after this change (fix any errors you introduced):\n" + block
|
||||
}
|
||||
@@ -1,53 +0,0 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// A wired Diagnostics callback must surface its block in the tool output for
|
||||
// both mutating tools, and a clean file (empty block) must leave the output
|
||||
// byte-identical to the pre-diagnostics behavior.
|
||||
func TestMutatingToolsAppendInlineDiagnostics(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(dir, "a.go"), []byte("package a\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
diagnostics := func(_ context.Context, absPath string) string {
|
||||
return absPath + ":1:1: error: undefined: x"
|
||||
}
|
||||
|
||||
editResult := NewEditFileTool(dir).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{
|
||||
"path": "a.go",
|
||||
"old_string": "package a",
|
||||
"new_string": "package b",
|
||||
}, RunOptions{Diagnostics: diagnostics})
|
||||
if editResult.Status != StatusOK {
|
||||
t.Fatalf("edit failed: %q", editResult.Output)
|
||||
}
|
||||
if !strings.Contains(editResult.Output, "Diagnostics in a.go after this change") || !strings.Contains(editResult.Output, "undefined: x") {
|
||||
t.Fatalf("edit output missing diagnostics block: %q", editResult.Output)
|
||||
}
|
||||
|
||||
writeResult := NewWriteFileTool(dir).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{
|
||||
"path": "b.go",
|
||||
"content": "package b\n",
|
||||
}, RunOptions{Diagnostics: diagnostics})
|
||||
if writeResult.Status != StatusOK {
|
||||
t.Fatalf("write failed: %q", writeResult.Output)
|
||||
}
|
||||
if !strings.Contains(writeResult.Output, "Diagnostics in b.go after this change") {
|
||||
t.Fatalf("write output missing diagnostics block: %q", writeResult.Output)
|
||||
}
|
||||
|
||||
clean := NewWriteFileTool(dir).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{
|
||||
"path": "c.go",
|
||||
"content": "package c\n",
|
||||
}, RunOptions{Diagnostics: func(context.Context, string) string { return "" }})
|
||||
if strings.Contains(clean.Output, "Diagnostics") {
|
||||
t.Fatalf("clean file must not gain a diagnostics block: %q", clean.Output)
|
||||
}
|
||||
}
|
||||
@@ -40,12 +40,6 @@ type RunOptions struct {
|
||||
// tool-call progress in the specialist card. nil is a no-op (the default
|
||||
// for every non-Task tool).
|
||||
Progress func(streamjson.Event)
|
||||
// Diagnostics, when set, returns a formatted language-diagnostics block for
|
||||
// a file a mutating tool just wrote ("" when clean or no server available).
|
||||
// edit_file/write_file append it to their output so the model sees an error
|
||||
// it introduced in the same turn instead of waiting for a later verification
|
||||
// pass. nil disables inline diagnostics.
|
||||
Diagnostics func(ctx context.Context, absPath string) string
|
||||
}
|
||||
|
||||
type sandboxAwareTool interface {
|
||||
|
||||
@@ -1,13 +0,0 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
)
|
||||
|
||||
func searchCancelledResult(tool string, err error) (Result, bool) {
|
||||
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
||||
return Result{Status: StatusError, Output: "Error: " + tool + " cancelled."}, true
|
||||
}
|
||||
return Result{}, false
|
||||
}
|
||||
@@ -1,206 +0,0 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// countingCancelContext returns context.Canceled from Err() once it has been
|
||||
// called more than "remaining" times, letting a test deterministically
|
||||
// cancel a loop after a fixed number of iterations instead of relying on
|
||||
// timing. A context.WithTimeout test would be flaky under a loaded or
|
||||
// shared CI runner, exactly the kind of test PR #464 was criticized for
|
||||
// elsewhere (TestCollectRespectsDeadlineUnderContinuousOutput).
|
||||
type countingCancelContext struct {
|
||||
context.Context
|
||||
remaining int
|
||||
}
|
||||
|
||||
func (c *countingCancelContext) Err() error {
|
||||
c.remaining--
|
||||
if c.remaining < 0 {
|
||||
return context.Canceled
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// buildLargeSearchTree creates n files, each containing a grep-matchable
|
||||
// line, so a walk that does NOT respect cancellation would have plenty of
|
||||
// work left to do (and matches to find) after the very first entry.
|
||||
func buildLargeSearchTree(t *testing.T, n int) string {
|
||||
t.Helper()
|
||||
root := t.TempDir()
|
||||
for i := 0; i < n; i++ {
|
||||
path := filepath.Join(root, "file"+strconv.Itoa(i)+".txt")
|
||||
if err := os.WriteFile(path, []byte("needle\n"), 0o644); err != nil {
|
||||
t.Fatalf("WriteFile: %v", err)
|
||||
}
|
||||
}
|
||||
return root
|
||||
}
|
||||
|
||||
// A cancelled run must stop the filesystem walk promptly instead of visiting
|
||||
// every remaining entry to completion. Before this fix, grep's Run/
|
||||
// RunWithSandbox discarded their context entirely, so cancelling the run
|
||||
// (Ctrl+C / /exit) had no effect on an in-flight, unscoped search: the walk
|
||||
// ran to completion regardless, and — because the TUI's exit path waits for
|
||||
// the cancelled run's own response before it can quit — the whole
|
||||
// application was stuck until the walk finished on its own.
|
||||
func TestGrepStopsWalkOnCancelledContext(t *testing.T) {
|
||||
root := buildLargeSearchTree(t, 500)
|
||||
tool := NewGrepTool(root)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel() // already cancelled before the walk starts
|
||||
|
||||
result := tool.Run(ctx, map[string]any{"pattern": "needle"})
|
||||
if result.Status != StatusError {
|
||||
t.Fatalf("status = %q, want error", result.Status)
|
||||
}
|
||||
if result.Output != "Error: grep cancelled." {
|
||||
t.Fatalf("output = %q, want the cancellation message", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrepRunWithSandboxStopsWalkOnCancelledContext(t *testing.T) {
|
||||
root := buildLargeSearchTree(t, 500)
|
||||
tool := NewGrepTool(root)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
sandboxed := tool.(sandboxAwareTool)
|
||||
result := sandboxed.RunWithSandbox(ctx, map[string]any{"pattern": "needle"}, nil)
|
||||
if result.Status != StatusError {
|
||||
t.Fatalf("status = %q, want error", result.Status)
|
||||
}
|
||||
if result.Output != "Error: grep cancelled." {
|
||||
t.Fatalf("output = %q, want the cancellation message", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
// grep with a live (non-cancelled) context must still work normally: this
|
||||
// fix only short-circuits on cancellation, it does not change matching.
|
||||
func TestGrepStillMatchesWithLiveContext(t *testing.T) {
|
||||
root := buildLargeSearchTree(t, 3)
|
||||
tool := NewGrepTool(root)
|
||||
|
||||
result := tool.Run(context.Background(), map[string]any{"pattern": "needle"})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("status = %q, want ok; output=%q", result.Status, result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
// Before this fix, only grepFiles' walk checked ctx: once the walk had
|
||||
// finished discovering candidate files, collectGrepMatches read and
|
||||
// regex-scanned every one of them to completion regardless of cancellation.
|
||||
// This proves the match-collection phase itself stops partway through,
|
||||
// not only when the context was already cancelled before the walk started.
|
||||
// grepFiles must stop mid-walk once ctx is cancelled, not only when the
|
||||
// context was already cancelled before WalkDir starts.
|
||||
func TestGrepFilesStopsMidWalkOnCancelledContext(t *testing.T) {
|
||||
root := buildLargeSearchTree(t, 50)
|
||||
resolvedRoot, err := filepath.EvalSymlinks(root)
|
||||
if err != nil {
|
||||
t.Fatalf("EvalSymlinks: %v", err)
|
||||
}
|
||||
|
||||
const allowed = 5
|
||||
ctx := &countingCancelContext{Context: context.Background(), remaining: allowed}
|
||||
files, err := grepFiles(ctx, resolvedRoot, root, nil, readExcluder{})
|
||||
if !errors.Is(err, context.Canceled) {
|
||||
t.Fatalf("err = %v, want context.Canceled", err)
|
||||
}
|
||||
if len(files) >= 50 {
|
||||
t.Fatalf("files = %d; walk should stop mid-traversal instead of visiting every entry", len(files))
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrepCollectMatchesStopsMidCollectionOnCancelledContext(t *testing.T) {
|
||||
root := buildLargeSearchTree(t, 50)
|
||||
files := make([]string, 50)
|
||||
for i := range files {
|
||||
files[i] = filepath.Join(root, "file"+strconv.Itoa(i)+".txt")
|
||||
}
|
||||
resolvedRoot, err := filepath.EvalSymlinks(root)
|
||||
if err != nil {
|
||||
t.Fatalf("EvalSymlinks: %v", err)
|
||||
}
|
||||
compiled := regexp.MustCompile("needle")
|
||||
|
||||
const allowed = 5
|
||||
ctx := &countingCancelContext{Context: context.Background(), remaining: allowed}
|
||||
matches, err := collectGrepMatches(ctx, resolvedRoot, false, files, compiled)
|
||||
if !errors.Is(err, context.Canceled) {
|
||||
t.Fatalf("err = %v, want context.Canceled", err)
|
||||
}
|
||||
if len(matches) != allowed {
|
||||
t.Fatalf("matches = %d, want exactly %d collected before cancellation stopped the loop", len(matches), allowed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScanGlobStopsMidWalkOnCancelledContext(t *testing.T) {
|
||||
root := buildLargeSearchTree(t, 50)
|
||||
matcher, err := compileGlob("**/*.txt")
|
||||
if err != nil {
|
||||
t.Fatalf("compileGlob: %v", err)
|
||||
}
|
||||
|
||||
const allowed = 5
|
||||
ctx := &countingCancelContext{Context: context.Background(), remaining: allowed}
|
||||
matches, err := scanGlob(ctx, root, root, matcher, false, readExcluder{})
|
||||
if !errors.Is(err, context.Canceled) {
|
||||
t.Fatalf("err = %v, want context.Canceled", err)
|
||||
}
|
||||
if len(matches) >= 50 {
|
||||
t.Fatalf("matches = %d; walk should stop mid-traversal instead of visiting every entry", len(matches))
|
||||
}
|
||||
}
|
||||
|
||||
func TestGlobStopsWalkOnCancelledContext(t *testing.T) {
|
||||
root := buildLargeSearchTree(t, 500)
|
||||
tool := NewGlobTool(root)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
result := tool.Run(ctx, map[string]any{"pattern": "**/*.txt"})
|
||||
if result.Status != StatusError {
|
||||
t.Fatalf("status = %q, want error", result.Status)
|
||||
}
|
||||
if result.Output != "Error: glob cancelled." {
|
||||
t.Fatalf("output = %q, want the cancellation message", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGlobRunWithSandboxStopsWalkOnCancelledContext(t *testing.T) {
|
||||
root := buildLargeSearchTree(t, 500)
|
||||
tool := NewGlobTool(root)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
sandboxed := tool.(sandboxAwareTool)
|
||||
result := sandboxed.RunWithSandbox(ctx, map[string]any{"pattern": "**/*.txt"}, nil)
|
||||
if result.Status != StatusError {
|
||||
t.Fatalf("status = %q, want error", result.Status)
|
||||
}
|
||||
if result.Output != "Error: glob cancelled." {
|
||||
t.Fatalf("output = %q, want the cancellation message", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGlobStillMatchesWithLiveContext(t *testing.T) {
|
||||
root := buildLargeSearchTree(t, 3)
|
||||
tool := NewGlobTool(root)
|
||||
|
||||
result := tool.Run(context.Background(), map[string]any{"pattern": "**/*.txt"})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("status = %q, want ok; output=%q", result.Status, result.Output)
|
||||
}
|
||||
}
|
||||
@@ -1,141 +0,0 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/redaction"
|
||||
)
|
||||
|
||||
// Spill-to-disk for truncated tool output. When a command produces more than
|
||||
// its context budget, the tail was previously simply gone — the model's only
|
||||
// recourse was re-running the (possibly expensive or non-idempotent) command
|
||||
// with a bigger budget. Instead, the full output is written to a cache file
|
||||
// whose path is included in the truncation notice, so the model can grep or
|
||||
// read_file the remainder. Spilling is best-effort: on any error truncation
|
||||
// behaves exactly as before, just without the file hint.
|
||||
|
||||
// spillRetention is how long spilled outputs are kept; files older than this
|
||||
// are opportunistically removed whenever a new spill happens.
|
||||
const spillRetention = 7 * 24 * time.Hour
|
||||
|
||||
// spillRootPath returns the spill directory path without creating or checking
|
||||
// it. Used by the read-path resolver to recognize spill files as readable.
|
||||
func spillRootPath() string {
|
||||
name := "zero-tool-output"
|
||||
if uid := os.Getuid(); uid >= 0 {
|
||||
name = fmt.Sprintf("zero-tool-output-%d", uid)
|
||||
}
|
||||
return filepath.Join(os.TempDir(), name)
|
||||
}
|
||||
|
||||
// resolveSpillReadPath reports whether requestedPath is a file inside the
|
||||
// spill directory and, if so, returns its verified absolute path. Spill files
|
||||
// are zero-created, per-uid owned, and already secret-redacted, so letting the
|
||||
// scoped read tools open them is what makes the truncation notice's
|
||||
// "read_file it" recovery actually work. Symlinks are resolved and the result
|
||||
// must still be inside the spill dir, so a planted link cannot smuggle an
|
||||
// out-of-scope file through this gate.
|
||||
func resolveSpillReadPath(requestedPath string) (string, bool) {
|
||||
if requestedPath == "" || !filepath.IsAbs(requestedPath) {
|
||||
return "", false
|
||||
}
|
||||
root := spillRootPath()
|
||||
cleaned := filepath.Clean(requestedPath)
|
||||
if cleaned == root || !strings.HasPrefix(cleaned, root+string(filepath.Separator)) {
|
||||
return "", false
|
||||
}
|
||||
resolved, err := filepath.EvalSymlinks(cleaned)
|
||||
if err != nil {
|
||||
return "", false
|
||||
}
|
||||
resolvedRoot, err := filepath.EvalSymlinks(root)
|
||||
if err != nil {
|
||||
return "", false
|
||||
}
|
||||
if !strings.HasPrefix(resolved, resolvedRoot+string(filepath.Separator)) {
|
||||
return "", false
|
||||
}
|
||||
return resolved, true
|
||||
}
|
||||
|
||||
// spillDir returns the per-user spill directory, creating it on first use.
|
||||
// Hardened for shared temp dirs (Linux /tmp): the name carries the uid so
|
||||
// users cannot collide, and a pre-existing path is only accepted when it is a
|
||||
// real directory (not a symlink that would redirect spills) owned by the
|
||||
// current user. Any doubt fails the spill — it is best-effort anyway.
|
||||
func spillDir() (string, error) {
|
||||
dir := spillRootPath()
|
||||
if err := os.MkdirAll(dir, 0o700); err != nil {
|
||||
return "", err
|
||||
}
|
||||
// MkdirAll follows symlinks and leaves an existing directory untouched, so
|
||||
// verify what is actually at the path.
|
||||
info, err := os.Lstat(dir)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !info.IsDir() {
|
||||
return "", fmt.Errorf("spill path %s is not a directory", dir)
|
||||
}
|
||||
if err := checkSpillDirOwner(info); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return dir, nil
|
||||
}
|
||||
|
||||
// spillTruncatedOutput writes the full pre-truncation output to the spill
|
||||
// directory and returns the file path, or "" when spilling fails. Output is
|
||||
// scrubbed with the same configured-key redaction the registry applies at the
|
||||
// tool boundary, so a spilled file never holds a secret the transcript would
|
||||
// have hidden.
|
||||
func spillTruncatedOutput(toolName, output string) string {
|
||||
dir, err := spillDir()
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
sweepSpillDir(dir)
|
||||
prefix := strings.Map(func(r rune) rune {
|
||||
if r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '_' || r == '-' {
|
||||
return r
|
||||
}
|
||||
return '-'
|
||||
}, toolName)
|
||||
file, err := os.CreateTemp(dir, prefix+"-*.txt")
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
defer file.Close()
|
||||
scrubbed := redaction.RedactString(output, redaction.Options{})
|
||||
if _, err := file.WriteString(scrubbed); err != nil {
|
||||
_ = os.Remove(file.Name())
|
||||
return ""
|
||||
}
|
||||
return file.Name()
|
||||
}
|
||||
|
||||
// sweepSpillDir removes spill files older than spillRetention. Best-effort:
|
||||
// errors are ignored, the directory is small, and a sweep runs only when a new
|
||||
// spill is about to happen anyway.
|
||||
func sweepSpillDir(dir string) {
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
cutoff := time.Now().Add(-spillRetention)
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
info, err := entry.Info()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if info.ModTime().Before(cutoff) {
|
||||
_ = os.Remove(filepath.Join(dir, entry.Name()))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,91 +0,0 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func spillDirName() string {
|
||||
if uid := os.Getuid(); uid >= 0 {
|
||||
return fmt.Sprintf("zero-tool-output-%d", uid)
|
||||
}
|
||||
return "zero-tool-output"
|
||||
}
|
||||
|
||||
func TestSpillDirRejectsSymlink(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("symlink creation needs privileges on Windows")
|
||||
}
|
||||
tmp := t.TempDir()
|
||||
t.Setenv("TMPDIR", tmp)
|
||||
elsewhere := filepath.Join(tmp, "elsewhere")
|
||||
if err := os.Mkdir(elsewhere, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// An attacker pre-creates the spill path as a symlink to redirect spills.
|
||||
if err := os.Symlink(elsewhere, filepath.Join(tmp, spillDirName())); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if path := spillTruncatedOutput("bash", "sensitive output"); path != "" {
|
||||
t.Fatalf("spill must refuse a symlinked directory, wrote %s", path)
|
||||
}
|
||||
entries, _ := os.ReadDir(elsewhere)
|
||||
if len(entries) != 0 {
|
||||
t.Fatalf("nothing may land behind the symlink: %v", entries)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSpillDirAcceptsOwnedDirectory(t *testing.T) {
|
||||
t.Setenv("TMPDIR", t.TempDir())
|
||||
if path := spillTruncatedOutput("bash", "ok"); path == "" {
|
||||
t.Fatal("spill must work in a clean per-user temp dir")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveSpillReadPathContainment(t *testing.T) {
|
||||
t.Setenv("TMPDIR", t.TempDir())
|
||||
spillPath := spillTruncatedOutput("exec_command", "spilled body")
|
||||
if spillPath == "" {
|
||||
t.Fatal("spill must succeed in a clean temp dir")
|
||||
}
|
||||
resolved, ok := resolveSpillReadPath(spillPath)
|
||||
if !ok || resolved == "" {
|
||||
t.Fatalf("a real spill file must resolve, got ok=%v", ok)
|
||||
}
|
||||
if _, ok := resolveSpillReadPath(filepath.Join(spillRootPath(), "..", "escape.txt")); ok {
|
||||
t.Fatal("path traversal out of the spill dir must be rejected")
|
||||
}
|
||||
if _, ok := resolveSpillReadPath("/etc/hosts"); ok {
|
||||
t.Fatal("paths outside the spill dir must be rejected")
|
||||
}
|
||||
if _, ok := resolveSpillReadPath(spillRootPath()); ok {
|
||||
t.Fatal("the spill dir itself is not a readable file target")
|
||||
}
|
||||
if runtime.GOOS != "windows" {
|
||||
// A symlink planted inside the spill dir pointing outside must not leak.
|
||||
link := filepath.Join(spillRootPath(), "sneaky-link")
|
||||
if err := os.Symlink("/etc/hosts", link); err == nil {
|
||||
if _, ok := resolveSpillReadPath(link); ok {
|
||||
t.Fatal("symlink escaping the spill dir must be rejected")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadFileCanReadSpillFile(t *testing.T) {
|
||||
t.Setenv("TMPDIR", t.TempDir())
|
||||
spillPath := spillTruncatedOutput("exec_command", "line one\nline two\n")
|
||||
if spillPath == "" {
|
||||
t.Fatal("spill must succeed")
|
||||
}
|
||||
workspace := t.TempDir()
|
||||
result := NewReadFileTool(workspace).Run(context.Background(), map[string]any{"path": spillPath})
|
||||
if result.Status != StatusOK || !strings.Contains(result.Output, "line two") {
|
||||
t.Fatalf("read_file must be able to follow the truncation notice: %s %q", result.Status, result.Output)
|
||||
}
|
||||
}
|
||||
@@ -1,24 +0,0 @@
|
||||
//go:build !windows
|
||||
|
||||
package tools
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
// checkSpillDirOwner rejects a spill directory not owned by the current user:
|
||||
// on a shared /tmp another user could have pre-created the path and would then
|
||||
// control its lifetime (deletion/renaming) even though the 0600 spill files
|
||||
// keep their contents private.
|
||||
func checkSpillDirOwner(info os.FileInfo) error {
|
||||
stat, ok := info.Sys().(*syscall.Stat_t)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
if int(stat.Uid) != os.Geteuid() {
|
||||
return fmt.Errorf("spill directory is owned by uid %d, not the current user", stat.Uid)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -1,12 +0,0 @@
|
||||
//go:build windows
|
||||
|
||||
package tools
|
||||
|
||||
import "os"
|
||||
|
||||
// checkSpillDirOwner is a no-op on Windows: %TEMP% is per-user by default and
|
||||
// there is no portable uid to compare. The Lstat symlink/dir check in spillDir
|
||||
// still applies.
|
||||
func checkSpillDirOwner(os.FileInfo) error {
|
||||
return nil
|
||||
}
|
||||
@@ -1,82 +0,0 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestTruncateExecOutputSpillWritesFullOutput(t *testing.T) {
|
||||
t.Setenv("TMPDIR", t.TempDir())
|
||||
long := strings.Repeat("line of build output\n", 5000) // ~100KB > 40KB budget
|
||||
|
||||
truncated, wasTruncated := truncateExecOutputSpill(long, defaultMaxOutputTokens, "bash")
|
||||
if !wasTruncated {
|
||||
t.Fatal("output over budget must truncate")
|
||||
}
|
||||
if !strings.Contains(truncated, "full output saved to ") {
|
||||
t.Fatalf("truncation notice missing spill path: %q", truncated[:200])
|
||||
}
|
||||
// Extract the path and verify the file holds the complete output.
|
||||
start := strings.Index(truncated, "full output saved to ") + len("full output saved to ")
|
||||
end := strings.Index(truncated[start:], " (grep")
|
||||
spillPath := truncated[start : start+end]
|
||||
content, err := os.ReadFile(spillPath)
|
||||
if err != nil {
|
||||
t.Fatalf("spill file unreadable: %v", err)
|
||||
}
|
||||
if string(content) != long {
|
||||
t.Fatalf("spill file must hold the full output: got %d bytes, want %d", len(content), len(long))
|
||||
}
|
||||
}
|
||||
|
||||
func TestTruncateExecOutputUnderBudgetUnchanged(t *testing.T) {
|
||||
output, wasTruncated := truncateExecOutputSpill("short output", defaultMaxOutputTokens, "bash")
|
||||
if wasTruncated || output != "short output" {
|
||||
t.Fatalf("under-budget output must pass through untouched: %q", output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSweepSpillDirRemovesOnlyOldFiles(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
oldFile := filepath.Join(dir, "bash-old.txt")
|
||||
newFile := filepath.Join(dir, "bash-new.txt")
|
||||
for _, path := range []string{oldFile, newFile} {
|
||||
if err := os.WriteFile(path, []byte("x"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
stale := time.Now().Add(-spillRetention - time.Hour)
|
||||
if err := os.Chtimes(oldFile, stale, stale); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
sweepSpillDir(dir)
|
||||
|
||||
if _, err := os.Stat(oldFile); !os.IsNotExist(err) {
|
||||
t.Fatal("stale spill file must be removed")
|
||||
}
|
||||
if _, err := os.Stat(newFile); err != nil {
|
||||
t.Fatal("fresh spill file must survive the sweep")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSpillTruncatedOutputWritesFile(t *testing.T) {
|
||||
t.Setenv("TMPDIR", t.TempDir())
|
||||
path := spillTruncatedOutput("exec_command", "some output body")
|
||||
if path == "" {
|
||||
t.Fatal("spill must return a file path")
|
||||
}
|
||||
if base := filepath.Base(path); !strings.HasPrefix(base, "exec_command-") {
|
||||
t.Fatalf("spill file name must carry the tool prefix: %s", base)
|
||||
}
|
||||
content, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(content) != "some output body" {
|
||||
t.Fatalf("unexpected spill content: %q", content)
|
||||
}
|
||||
}
|
||||
@@ -261,16 +261,10 @@ 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 {
|
||||
switch {
|
||||
case strings.Contains(name, keyword):
|
||||
if 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++
|
||||
@@ -296,19 +290,6 @@ 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 {
|
||||
|
||||
@@ -184,29 +184,6 @@ 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)
|
||||
|
||||
@@ -18,14 +18,9 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
// Raw-body budgets. HTML responses are converted to markdown before they
|
||||
// reach the model (see web_fetch_markdown.go), so a generous raw budget
|
||||
// does not translate into a generous context cost — conversion typically
|
||||
// shrinks a page by an order of magnitude. 64KiB of raw HTML often held
|
||||
// nothing but a page's <head>, which starved research tasks.
|
||||
defaultWebFetchMaxBytes = 256 * 1024
|
||||
maxWebFetchMaxBytes = 2 * 1024 * 1024
|
||||
webFetchTimeout = 30 * time.Second
|
||||
defaultWebFetchMaxBytes = 64 * 1024
|
||||
maxWebFetchMaxBytes = 512 * 1024
|
||||
webFetchTimeout = 10 * time.Second
|
||||
webFetchRedirectLimit = 5
|
||||
webFetchErrorBodyLimit = 4 * 1024
|
||||
webFetchPublicOnlyHint = "web_fetch only supports public remote HTTP/HTTPS URLs. For localhost or private network URLs, use bash with curl so sandbox network permission can apply."
|
||||
@@ -120,17 +115,11 @@ func newWebFetchToolWithClientAndResolver(client *http.Client, resolver webFetch
|
||||
},
|
||||
"max_bytes": {
|
||||
Type: "integer",
|
||||
Description: "Maximum raw response body bytes to download before conversion.",
|
||||
Description: "Maximum response body bytes to return.",
|
||||
Default: defaultWebFetchMaxBytes,
|
||||
Minimum: intPtr(1),
|
||||
Maximum: intPtr(maxWebFetchMaxBytes),
|
||||
},
|
||||
"format": {
|
||||
Type: "string",
|
||||
Description: "auto (default): HTML responses are converted to compact markdown, everything else is returned as-is. raw: never convert. markdown: force conversion.",
|
||||
Enum: []string{"auto", "raw", "markdown"},
|
||||
Default: "auto",
|
||||
},
|
||||
},
|
||||
Required: []string{"url"},
|
||||
AdditionalProperties: false,
|
||||
@@ -181,18 +170,6 @@ func (tool webFetchTool) run(ctx context.Context, args map[string]any) Result {
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for web_fetch: " + err.Error())
|
||||
}
|
||||
format, err := stringArg(args, "format", "auto", false)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for web_fetch: " + err.Error())
|
||||
}
|
||||
format = strings.ToLower(strings.TrimSpace(format))
|
||||
switch format {
|
||||
case "", "auto":
|
||||
format = "auto"
|
||||
case "raw", "markdown":
|
||||
default:
|
||||
return errorResult(`Error: Invalid arguments for web_fetch: format must be "auto", "raw", or "markdown".`)
|
||||
}
|
||||
if err := validateWebFetchURLBeforePermission(rawURL); err != nil {
|
||||
return errorResult("Error: Unsafe URL for web_fetch: " + err.Error())
|
||||
}
|
||||
@@ -247,28 +224,14 @@ func (tool webFetchTool) run(ctx context.Context, args map[string]any) Result {
|
||||
}
|
||||
body = redactWebFetchText(body)
|
||||
|
||||
// HTML responses are converted to compact markdown by default: raw HTML is
|
||||
// mostly markup, so conversion typically shrinks the page by an order of
|
||||
// magnitude and the model reads content instead of boilerplate. format=raw
|
||||
// is the escape hatch for pages the converter mangles.
|
||||
converted := false
|
||||
if format == "markdown" || (format == "auto" && looksLikeHTML(contentType, body)) {
|
||||
if markdown := htmlToMarkdown(body); markdown != "" {
|
||||
body = markdown
|
||||
converted = true
|
||||
}
|
||||
}
|
||||
|
||||
headers := []string{
|
||||
output := strings.Join([]string{
|
||||
"URL: " + finalURL,
|
||||
"Status: " + webFetchStatusLine(response),
|
||||
"Content-Type: " + firstNonEmptyString(contentType, "unknown"),
|
||||
"Bytes: " + strconv.Itoa(len(body)),
|
||||
}
|
||||
if converted {
|
||||
headers = append(headers, "Converted: html -> markdown (pass format: \"raw\" for the original HTML)")
|
||||
}
|
||||
output := strings.Join(append(headers, "", body), "\n")
|
||||
"",
|
||||
body,
|
||||
}, "\n")
|
||||
|
||||
return Result{
|
||||
Status: StatusOK,
|
||||
@@ -280,7 +243,6 @@ func (tool webFetchTool) run(ctx context.Context, args map[string]any) Result {
|
||||
"content_type": contentType,
|
||||
"bytes": strconv.Itoa(len(body)),
|
||||
"truncated": strconv.FormatBool(truncated),
|
||||
"converted": strconv.FormatBool(converted),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,92 +0,0 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"html"
|
||||
"regexp"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// HTML-to-markdown conversion for web_fetch. Raw HTML is mostly boilerplate —
|
||||
// on a typical page well under a tenth of the bytes are readable content — so
|
||||
// returning it raw wastes nearly the whole response budget on markup. The
|
||||
// converter is deliberately dependency-free (regex passes + stdlib entity
|
||||
// unescaping, no HTML parser): the output feeds a language model, not a
|
||||
// browser, so best-effort structure (headings, links, list markers) is enough,
|
||||
// and the format=raw escape hatch covers pages it mangles.
|
||||
|
||||
var (
|
||||
webFetchDropBlockRe = regexp.MustCompile(`(?is)<(script|style|head|noscript|svg|template|iframe)\b[^>]*>.*?</\s*(?:script|style|head|noscript|svg|template|iframe)\s*>`)
|
||||
webFetchCommentRe = regexp.MustCompile(`(?s)<!--.*?-->`)
|
||||
webFetchAnchorRe = regexp.MustCompile(`(?is)<a\b[^>]*?href\s*=\s*["']?([^"'\s>]+)["']?[^>]*>(.*?)</\s*a\s*>`)
|
||||
webFetchHeadingRe = regexp.MustCompile(`(?is)<h([1-6])[^>]*>(.*?)</\s*h[1-6]\s*>`)
|
||||
webFetchListItemRe = regexp.MustCompile(`(?i)<li\b[^>]*>`)
|
||||
webFetchBreakRe = regexp.MustCompile(`(?i)<(?:br|hr)\s*/?>`)
|
||||
webFetchBlockTagRe = regexp.MustCompile(`(?i)</?(?:p|div|section|article|main|header|footer|nav|aside|table|tr|ul|ol|blockquote|pre|figure|form)\b[^>]*>`)
|
||||
webFetchAnyTagRe = regexp.MustCompile(`(?s)<[^>]*>`)
|
||||
webFetchSpaceRunRe = regexp.MustCompile(`[ \t]+`)
|
||||
webFetchBlankRunRe = regexp.MustCompile(`\n{3,}`)
|
||||
)
|
||||
|
||||
// htmlToMarkdown converts an HTML document to compact markdown-flavored text:
|
||||
// headings become #-prefixed lines, anchors become [text](href), list items
|
||||
// become "- " bullets, block boundaries become blank lines, all other markup
|
||||
// is stripped, and entities are unescaped.
|
||||
func htmlToMarkdown(body string) string {
|
||||
text := webFetchCommentRe.ReplaceAllString(body, " ")
|
||||
text = webFetchDropBlockRe.ReplaceAllString(text, " ")
|
||||
text = webFetchHeadingRe.ReplaceAllStringFunc(text, func(match string) string {
|
||||
groups := webFetchHeadingRe.FindStringSubmatch(match)
|
||||
level := int(groups[1][0] - '0')
|
||||
if level < 1 || level > 6 {
|
||||
level = 1
|
||||
}
|
||||
title := strings.TrimSpace(webFetchAnyTagRe.ReplaceAllString(groups[2], " "))
|
||||
return "\n\n" + strings.Repeat("#", level) + " " + title + "\n\n"
|
||||
})
|
||||
text = webFetchAnchorRe.ReplaceAllStringFunc(text, func(match string) string {
|
||||
groups := webFetchAnchorRe.FindStringSubmatch(match)
|
||||
href := groups[1]
|
||||
label := strings.TrimSpace(webFetchAnyTagRe.ReplaceAllString(groups[2], " "))
|
||||
if label == "" {
|
||||
return " "
|
||||
}
|
||||
if strings.HasPrefix(href, "#") || strings.HasPrefix(strings.ToLower(href), "javascript:") {
|
||||
return label
|
||||
}
|
||||
return fmt.Sprintf("[%s](%s)", label, href)
|
||||
})
|
||||
text = webFetchListItemRe.ReplaceAllString(text, "\n- ")
|
||||
text = webFetchBreakRe.ReplaceAllString(text, "\n")
|
||||
text = webFetchBlockTagRe.ReplaceAllString(text, "\n\n")
|
||||
text = webFetchAnyTagRe.ReplaceAllString(text, " ")
|
||||
text = html.UnescapeString(text)
|
||||
|
||||
// Whitespace normalization: collapse space runs, trim line edges, cap blank
|
||||
// runs at one empty line.
|
||||
lines := strings.Split(text, "\n")
|
||||
for i, line := range lines {
|
||||
lines[i] = strings.TrimSpace(webFetchSpaceRunRe.ReplaceAllString(line, " "))
|
||||
}
|
||||
text = strings.Join(lines, "\n")
|
||||
text = webFetchBlankRunRe.ReplaceAllString(text, "\n\n")
|
||||
return strings.TrimSpace(text)
|
||||
}
|
||||
|
||||
// looksLikeHTML reports whether a response should be treated as an HTML
|
||||
// document, by Content-Type first and a body sniff as fallback (servers
|
||||
// mislabel HTML as text/plain often enough to matter).
|
||||
func looksLikeHTML(contentType, body string) bool {
|
||||
lowered := strings.ToLower(contentType)
|
||||
if strings.Contains(lowered, "text/html") || strings.Contains(lowered, "application/xhtml") {
|
||||
return true
|
||||
}
|
||||
if lowered != "" && !strings.Contains(lowered, "text/plain") {
|
||||
return false
|
||||
}
|
||||
head := strings.ToLower(strings.TrimSpace(body))
|
||||
if len(head) > 512 {
|
||||
head = head[:512]
|
||||
}
|
||||
return strings.Contains(head, "<!doctype html") || strings.Contains(head, "<html")
|
||||
}
|
||||
@@ -1,127 +0,0 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
const sampleHTML = `<!DOCTYPE html>
|
||||
<html><head><title>ignored</title><style>body{color:red}</style></head>
|
||||
<body>
|
||||
<script>var tracking = "noise";</script>
|
||||
<h1>Main Title</h1>
|
||||
<p>First paragraph with a <a href="https://example.com/docs">docs link</a> and & entity.</p>
|
||||
<ul><li>alpha</li><li>beta</li></ul>
|
||||
<!-- hidden comment -->
|
||||
<h2>Section</h2>
|
||||
<p>Body with extra spaces.</p>
|
||||
</body></html>`
|
||||
|
||||
func TestHTMLToMarkdown(t *testing.T) {
|
||||
markdown := htmlToMarkdown(sampleHTML)
|
||||
|
||||
for _, want := range []string{
|
||||
"# Main Title",
|
||||
"## Section",
|
||||
"[docs link](https://example.com/docs)",
|
||||
"- alpha",
|
||||
"- beta",
|
||||
"with a", // paragraph text survives
|
||||
"& entity",
|
||||
} {
|
||||
if !strings.Contains(markdown, want) {
|
||||
t.Fatalf("markdown missing %q:\n%s", want, markdown)
|
||||
}
|
||||
}
|
||||
for _, banned := range []string{"tracking", "color:red", "hidden comment", "<p>", "</html>", "ignored"} {
|
||||
if strings.Contains(markdown, banned) {
|
||||
t.Fatalf("markdown must not contain %q:\n%s", banned, markdown)
|
||||
}
|
||||
}
|
||||
if strings.Contains(markdown, "extra spaces") {
|
||||
t.Fatalf("space runs must collapse:\n%s", markdown)
|
||||
}
|
||||
if strings.Contains(markdown, "\n\n\n") {
|
||||
t.Fatalf("blank-line runs must be capped:\n%s", markdown)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLooksLikeHTML(t *testing.T) {
|
||||
cases := []struct {
|
||||
contentType string
|
||||
body string
|
||||
want bool
|
||||
}{
|
||||
{"text/html; charset=utf-8", "anything", true},
|
||||
{"application/xhtml+xml", "anything", true},
|
||||
{"application/json", `{"html":"<html>"}`, false},
|
||||
{"text/plain", "<!doctype html><html>", true}, // mislabeled HTML sniffed
|
||||
{"text/plain", "just text", false},
|
||||
{"", "<html><body>x</body></html>", true},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := looksLikeHTML(c.contentType, c.body); got != c.want {
|
||||
t.Errorf("looksLikeHTML(%q, %.20q) = %v, want %v", c.contentType, c.body, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebFetchConvertsHTMLByDefault(t *testing.T) {
|
||||
tool := newWebFetchToolWithClient(webFetchTestClient(func(request *http.Request) (*http.Response, error) {
|
||||
return webFetchTestResponse(request, http.StatusOK, "text/html; charset=utf-8", sampleHTML), nil
|
||||
}))
|
||||
|
||||
result := tool.Run(context.Background(), map[string]any{"url": "https://example.com/page"})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("expected ok, got %s: %s", result.Status, result.Output)
|
||||
}
|
||||
if !strings.Contains(result.Output, "# Main Title") || strings.Contains(result.Output, "<h1>") {
|
||||
t.Fatalf("HTML must convert to markdown by default:\n%s", result.Output)
|
||||
}
|
||||
if result.Meta["converted"] != "true" {
|
||||
t.Fatalf("converted meta must be true: %#v", result.Meta)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebFetchRawFormatSkipsConversion(t *testing.T) {
|
||||
tool := newWebFetchToolWithClient(webFetchTestClient(func(request *http.Request) (*http.Response, error) {
|
||||
return webFetchTestResponse(request, http.StatusOK, "text/html", sampleHTML), nil
|
||||
}))
|
||||
|
||||
result := tool.Run(context.Background(), map[string]any{"url": "https://example.com/page", "format": "raw"})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("expected ok, got %s: %s", result.Status, result.Output)
|
||||
}
|
||||
if !strings.Contains(result.Output, "<h1>Main Title</h1>") {
|
||||
t.Fatalf("format=raw must keep original HTML:\n%s", result.Output)
|
||||
}
|
||||
if result.Meta["converted"] != "false" {
|
||||
t.Fatalf("converted meta must be false: %#v", result.Meta)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebFetchLeavesNonHTMLUntouched(t *testing.T) {
|
||||
tool := newWebFetchToolWithClient(webFetchTestClient(func(request *http.Request) (*http.Response, error) {
|
||||
return webFetchTestResponse(request, http.StatusOK, "application/json", `{"key":"<b>value</b>"}`), nil
|
||||
}))
|
||||
|
||||
result := tool.Run(context.Background(), map[string]any{"url": "https://example.com/api"})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("expected ok, got %s: %s", result.Status, result.Output)
|
||||
}
|
||||
if !strings.Contains(result.Output, `{"key":"<b>value</b>"}`) {
|
||||
t.Fatalf("non-HTML must pass through untouched:\n%s", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebFetchRejectsInvalidFormat(t *testing.T) {
|
||||
tool := newWebFetchToolWithClient(webFetchTestClient(func(request *http.Request) (*http.Response, error) {
|
||||
return webFetchTestResponse(request, http.StatusOK, "text/plain", "x"), nil
|
||||
}))
|
||||
result := tool.Run(context.Background(), map[string]any{"url": "https://example.com", "format": "xml"})
|
||||
if result.Status != StatusError || !strings.Contains(result.Output, "format must be") {
|
||||
t.Fatalf("invalid format must be rejected: %s %s", result.Status, result.Output)
|
||||
}
|
||||
}
|
||||
@@ -254,14 +254,6 @@ func scopedReadRoots(workspaceRoot string, scope PathScope) ([]string, error) {
|
||||
}
|
||||
|
||||
func resolveScopedReadPath(workspaceRoot string, scope PathScope, requestedPath string) (string, string, error) {
|
||||
// Spill files (truncated tool output saved under the per-uid temp dir) are
|
||||
// readable regardless of scope: the truncation notice tells the model to
|
||||
// read_file/grep them, which must actually work. resolveSpillReadPath
|
||||
// verifies containment after symlink resolution, so this cannot be used to
|
||||
// reach anything outside the spill dir.
|
||||
if spillPath, ok := resolveSpillReadPath(requestedPath); ok {
|
||||
return spillPath, spillPath, nil
|
||||
}
|
||||
if requestedPath == "" || !filepath.IsAbs(requestedPath) || scope == nil {
|
||||
return resolveWorkspacePath(workspaceRoot, requestedPath)
|
||||
}
|
||||
|
||||
@@ -44,7 +44,7 @@ func (tool writeFileTool) Run(ctx context.Context, args map[string]any) Result {
|
||||
return tool.RunWithOptions(ctx, args, RunOptions{})
|
||||
}
|
||||
|
||||
func (tool writeFileTool) RunWithOptions(ctx context.Context, args map[string]any, options RunOptions) Result {
|
||||
func (tool writeFileTool) RunWithOptions(_ context.Context, args map[string]any, options RunOptions) Result {
|
||||
requestedPath, err := aliasedStringArg(args, []string{"path", "file", "file_path", "filename"}, "", true, false)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for write_file: " + err.Error())
|
||||
@@ -110,10 +110,6 @@ func (tool writeFileTool) RunWithOptions(ctx context.Context, args map[string]an
|
||||
if err := os.WriteFile(absolutePath, []byte(content), 0o644); err != nil {
|
||||
return errorResult("Error writing file " + relativePath + ": " + err.Error())
|
||||
}
|
||||
// Optional format-on-write (ZERO_FORMAT_ON_WRITE). Must run BEFORE the
|
||||
// FileTracker baseline: recording pre-format content would make the very
|
||||
// next edit look like an external modification and trip the conflict guard.
|
||||
content = maybeFormatWrittenFile(ctx, absolutePath, content)
|
||||
// Baseline the freshly written content so a later edit/overwrite in this
|
||||
// session compares against what is now on disk.
|
||||
newInfo, _ := os.Stat(absolutePath)
|
||||
@@ -130,7 +126,6 @@ func (tool writeFileTool) RunWithOptions(ctx context.Context, args map[string]an
|
||||
lines++
|
||||
}
|
||||
summary := fmt.Sprintf("%s %s (%d lines).", verb, relativePath, lines)
|
||||
summary += inlineDiagnostics(ctx, options, absolutePath, relativePath)
|
||||
result := okResult(summary)
|
||||
result.ChangedFiles = []string{relativePath}
|
||||
// Card-only preview: a real unified diff (all-green for a create, red/green for
|
||||
|
||||
@@ -1,109 +0,0 @@
|
||||
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)
|
||||
}
|
||||
@@ -1,182 +0,0 @@
|
||||
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])
|
||||
}
|
||||
}
|
||||
@@ -482,9 +482,6 @@ 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")
|
||||
}
|
||||
|
||||
@@ -539,11 +536,7 @@ func (m model) switchProviderModel(providerName, modelID string) (model, string,
|
||||
cmds = append(cmds, cmd)
|
||||
}
|
||||
}
|
||||
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...)
|
||||
return m, fmt.Sprintf("Model\nSwitched to %s · %s", target.Name, target.Model), tea.Batch(cmds...)
|
||||
}
|
||||
|
||||
// profileWithCredential fills a profile's APIKey for provider construction the same
|
||||
|
||||
+10
-34
@@ -41,11 +41,8 @@ const (
|
||||
commandAddDir
|
||||
commandSelfCorrect
|
||||
commandTurns
|
||||
commandRetry
|
||||
commandEdit
|
||||
commandCopy
|
||||
commandExport
|
||||
commandNew
|
||||
commandFast
|
||||
commandUnknown
|
||||
)
|
||||
|
||||
@@ -213,9 +210,9 @@ var commandDefinitions = []commandDefinition{
|
||||
},
|
||||
{
|
||||
name: "/compact",
|
||||
usage: "/compact [status|now]",
|
||||
usage: "/compact [status]",
|
||||
group: commandGroupSession,
|
||||
description: "Compact the transcript now, or show compaction state (/compact status).",
|
||||
description: "Show or request transcript compaction state.",
|
||||
kind: commandCompact,
|
||||
},
|
||||
{
|
||||
@@ -239,6 +236,13 @@ var commandDefinitions = []commandDefinition{
|
||||
description: "Show or set reasoning effort for supported models.",
|
||||
kind: commandEffort,
|
||||
},
|
||||
{
|
||||
name: "/fast",
|
||||
usage: "/fast [on|off]",
|
||||
group: commandGroupModel,
|
||||
description: "Enable fast mode for the current model (/fast off to disable).",
|
||||
kind: commandFast,
|
||||
},
|
||||
{
|
||||
name: "/style",
|
||||
usage: "/style [list|balanced|concise|explanatory|review]",
|
||||
@@ -261,34 +265,6 @@ 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",
|
||||
|
||||
@@ -1,27 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -306,7 +306,7 @@ func (m model) composerPositionAtVisualCell(x int, y int, width int) (int, bool)
|
||||
func (m model) applyComposerKey(msg tea.KeyMsg) (model, bool) {
|
||||
state := m.currentComposerState()
|
||||
switch {
|
||||
case keyIs(msg, tea.KeyEnter) && (keyAlt(msg) || keyShift(msg)):
|
||||
case keyIs(msg, tea.KeyEnter) && keyAlt(msg):
|
||||
m = m.insertComposerTextWithPastePreview(state, "\n", "")
|
||||
case keyCtrl(msg, 'j'):
|
||||
m = m.insertComposerTextWithPastePreview(state, "\n", "")
|
||||
|
||||
@@ -299,8 +299,7 @@ func TestModifiedEnterInsertsNewlineWithoutSubmitting(t *testing.T) {
|
||||
key tea.Msg
|
||||
}{
|
||||
{name: "alt enter", key: testKeyAlt(tea.KeyEnter)},
|
||||
{name: "shift enter", key: testKeyShift(tea.KeyEnter)},
|
||||
{name: "ctrl j", key: testKeyCtrl('j')},
|
||||
{name: "shift enter", key: testKeyCtrl('j')},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
|
||||
@@ -1,27 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
package tui
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/modelregistry"
|
||||
)
|
||||
|
||||
// fastPlan is the pure decision `/fast` makes before touching any state: either
|
||||
// a terminal message to show the user, or a target model to switch to.
|
||||
type fastPlan struct {
|
||||
message string // terminal notice; shown as-is when targetID is empty
|
||||
targetID string // model to switch to ("" means no switch — show message)
|
||||
toFast bool // the target is a fast model (drives the ⚡ prefix)
|
||||
}
|
||||
|
||||
// planFastCommand resolves what `/fast [on|off]` should do for the current model
|
||||
// using only the registry — no side effects. State is DERIVED from the model
|
||||
// identity: the session is "in fast mode" iff the current model has a base
|
||||
// variant. Returns a terminal message, or a target model to switch to.
|
||||
func planFastCommand(registry modelregistry.Registry, currentModel, arg string) fastPlan {
|
||||
// 1. Parse the argument (case-insensitive): no arg / "on" => enable fast,
|
||||
// "off" => return to base, anything else => usage error.
|
||||
enable := true
|
||||
switch strings.ToLower(strings.TrimSpace(arg)) {
|
||||
case "", "on":
|
||||
enable = true
|
||||
case "off":
|
||||
enable = false
|
||||
default:
|
||||
return fastPlan{message: fmt.Sprintf("⚡ /fast takes on or off, not %q", strings.TrimSpace(arg))}
|
||||
}
|
||||
|
||||
// 2. Read the current model.
|
||||
current := strings.TrimSpace(currentModel)
|
||||
if current == "" {
|
||||
return fastPlan{message: "⚡ Pick a model first · /model"}
|
||||
}
|
||||
short := modelShortName(registry, current)
|
||||
|
||||
// 3. Derived state: in fast mode iff the current model has a base variant.
|
||||
_, isFast := registry.BaseVariant(current)
|
||||
|
||||
// 4. No-op cases — already in the requested state.
|
||||
if enable && isFast {
|
||||
return fastPlan{message: "⚡ Already in the fast lane · " + short}
|
||||
}
|
||||
if !enable && !isFast {
|
||||
return fastPlan{message: "Already on base · " + short}
|
||||
}
|
||||
|
||||
// 5. Resolve the target variant.
|
||||
var target modelregistry.ModelEntry
|
||||
var ok bool
|
||||
if enable {
|
||||
target, ok = registry.FastVariant(current)
|
||||
} else {
|
||||
target, ok = registry.BaseVariant(current)
|
||||
}
|
||||
if !ok {
|
||||
return fastPlan{message: "⚡ No fast lane for " + short}
|
||||
}
|
||||
|
||||
// Switching TO a fast model iff the target itself has a base variant.
|
||||
_, toFast := registry.BaseVariant(target.ID)
|
||||
return fastPlan{targetID: target.ID, toFast: toFast}
|
||||
}
|
||||
|
||||
// handleFastCommand implements `/fast [on|off]`: toggle the session model between
|
||||
// its base and fast variant within the same family. The actual switch is
|
||||
// delegated to handleModelCommand — the same path /model uses — so provider
|
||||
// rebuild, the compaction-before-switch guard, and persistence all apply; this
|
||||
// never mutates model state directly.
|
||||
func (m model) handleFastCommand(arg string) (model, string) {
|
||||
registry, err := modelregistry.DefaultRegistry()
|
||||
if err != nil {
|
||||
return m, "Failed to load the model registry: " + err.Error()
|
||||
}
|
||||
|
||||
plan := planFastCommand(registry, m.modelName, arg)
|
||||
if plan.targetID == "" {
|
||||
return m, plan.message
|
||||
}
|
||||
|
||||
var switchText string
|
||||
m, switchText = m.handleModelCommand(plan.targetID)
|
||||
if m.modelName != plan.targetID {
|
||||
// The switch did not complete — either an error, or a "compact first"
|
||||
// guard prompt from handleModelCommand. Surface its own message so the
|
||||
// user sees the real reason instead of a misleading "Switched to ...".
|
||||
return m, switchText
|
||||
}
|
||||
|
||||
short := modelShortName(registry, plan.targetID)
|
||||
if plan.toFast {
|
||||
return m, "⚡ Fast lane on · " + short
|
||||
}
|
||||
return m, "Fast lane off · " + short
|
||||
}
|
||||
|
||||
// modelShortName returns a model's short display name, falling back to the id
|
||||
// when the model is unknown or has no display name.
|
||||
func modelShortName(registry modelregistry.Registry, id string) string {
|
||||
if entry, ok := registry.Get(id); ok {
|
||||
if name := strings.TrimSpace(entry.DisplayName); name != "" {
|
||||
return name
|
||||
}
|
||||
}
|
||||
return strings.TrimSpace(id)
|
||||
}
|
||||
@@ -0,0 +1,200 @@
|
||||
package tui
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/modelregistry"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
func fastTestRegistry(t *testing.T) modelregistry.Registry {
|
||||
t.Helper()
|
||||
registry, err := modelregistry.DefaultRegistry()
|
||||
if err != nil {
|
||||
t.Fatalf("DefaultRegistry() error = %v", err)
|
||||
}
|
||||
return registry
|
||||
}
|
||||
|
||||
// planFastCommand is the pure decision layer — every argument/state case is
|
||||
// exercised here without touching the switch machinery.
|
||||
func TestPlanFastCommand(t *testing.T) {
|
||||
registry := fastTestRegistry(t)
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
model string
|
||||
arg string
|
||||
wantMessage string // substring expected when the plan is terminal
|
||||
wantTarget string // model id expected when the plan switches
|
||||
wantToFast bool
|
||||
}{
|
||||
{name: "invalid arg", model: "gpt-4.1", arg: "turbo", wantMessage: `/fast takes on or off, not "turbo"`},
|
||||
{name: "no model set", model: "", arg: "on", wantMessage: "Pick a model first"},
|
||||
{name: "enable while already fast", model: "gpt-4.1-mini", arg: "on", wantMessage: "Already in the fast lane"},
|
||||
{name: "disable while already base", model: "gpt-4.1", arg: "off", wantMessage: "Already on base"},
|
||||
{name: "enable with no fast variant", model: "claude-opus-4.1", arg: "on", wantMessage: "No fast lane"},
|
||||
{name: "enable (no arg) resolves fast", model: "gpt-4.1", arg: "", wantTarget: "gpt-4.1-mini", wantToFast: true},
|
||||
{name: "enable (on) resolves fast", model: "gpt-4.1", arg: "ON", wantTarget: "gpt-4.1-mini", wantToFast: true},
|
||||
{name: "disable resolves base", model: "gpt-4.1-mini", arg: "off", wantTarget: "gpt-4.1", wantToFast: false},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
plan := planFastCommand(registry, tc.model, tc.arg)
|
||||
if tc.wantTarget != "" {
|
||||
canon, ok := registry.Resolve(tc.wantTarget)
|
||||
if !ok {
|
||||
t.Fatalf("test fixture %q does not resolve", tc.wantTarget)
|
||||
}
|
||||
if plan.targetID != canon.ID {
|
||||
t.Fatalf("targetID = %q, want %q", plan.targetID, canon.ID)
|
||||
}
|
||||
if plan.toFast != tc.wantToFast {
|
||||
t.Fatalf("toFast = %v, want %v", plan.toFast, tc.wantToFast)
|
||||
}
|
||||
if plan.message != "" {
|
||||
t.Fatalf("switching plan should carry no message, got %q", plan.message)
|
||||
}
|
||||
return
|
||||
}
|
||||
if plan.targetID != "" {
|
||||
t.Fatalf("expected a terminal message, got target %q", plan.targetID)
|
||||
}
|
||||
if !strings.Contains(plan.message, tc.wantMessage) {
|
||||
t.Fatalf("message = %q, want it to contain %q", plan.message, tc.wantMessage)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// handleFastCommand's terminal cases return the plan message without switching,
|
||||
// so a bare model exercises them.
|
||||
func TestHandleFastCommandTerminalCases(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
model string
|
||||
arg string
|
||||
wantContains string
|
||||
}{
|
||||
{"invalid arg", "gpt-4.1", "sideways", `/fast takes on or off, not "sideways"`},
|
||||
{"no model", "", "on", "Pick a model first"},
|
||||
{"already fast", "gpt-4.1-mini", "on", "Already in the fast lane"},
|
||||
{"already base", "gpt-4.1", "off", "Already on base"},
|
||||
{"no fast variant", "claude-opus-4.1", "on", "No fast lane"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
m := model{modelName: tc.model}
|
||||
got, notice := m.handleFastCommand(tc.arg)
|
||||
if got.modelName != tc.model {
|
||||
t.Fatalf("model must not change on a terminal case: %q -> %q", tc.model, got.modelName)
|
||||
}
|
||||
if !strings.Contains(notice, tc.wantContains) {
|
||||
t.Fatalf("notice = %q, want it to contain %q", notice, tc.wantContains)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func newFastTestModel(t *testing.T, modelID string) model {
|
||||
t.Helper()
|
||||
configPath := filepath.Join(t.TempDir(), "zero.json")
|
||||
profile := config.ProviderProfile{
|
||||
Name: "openai",
|
||||
ProviderKind: config.ProviderKindOpenAI,
|
||||
BaseURL: config.OpenAIBaseURL,
|
||||
APIKey: "sk-test",
|
||||
Model: modelID,
|
||||
}
|
||||
if _, err := config.UpsertProvider(configPath, profile, true); err != nil {
|
||||
t.Fatalf("write user config: %v", err)
|
||||
}
|
||||
return newModel(context.Background(), Options{
|
||||
UserConfigPath: configPath,
|
||||
ProviderName: "openai",
|
||||
ModelName: modelID,
|
||||
ProviderProfile: profile,
|
||||
Provider: &fakeProvider{},
|
||||
NewProvider: func(config.ProviderProfile) (zeroruntime.Provider, error) {
|
||||
return &fakeProvider{}, nil
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// Enabling switches base -> fast with a ⚡ notice; disabling switches back to the
|
||||
// base with a plain notice. Uses the same switch path /model uses.
|
||||
func TestHandleFastCommandSwitchesToFastAndBack(t *testing.T) {
|
||||
registry := fastTestRegistry(t)
|
||||
m := newFastTestModel(t, "gpt-4.1")
|
||||
|
||||
m, notice := m.handleFastCommand("on")
|
||||
if m.modelName != "gpt-4.1-mini" {
|
||||
t.Fatalf("enable: modelName = %q, want gpt-4.1-mini", m.modelName)
|
||||
}
|
||||
if want := "⚡ Fast lane on · " + modelShortName(registry, "gpt-4.1-mini"); notice != want {
|
||||
t.Fatalf("enable notice = %q, want exactly %q (bolt + fast-lane-on + target short name)", notice, want)
|
||||
}
|
||||
|
||||
m, notice = m.handleFastCommand("off")
|
||||
if m.modelName != "gpt-4.1" {
|
||||
t.Fatalf("disable: modelName = %q, want gpt-4.1", m.modelName)
|
||||
}
|
||||
if want := "Fast lane off · " + modelShortName(registry, "gpt-4.1"); notice != want {
|
||||
t.Fatalf("disable notice = %q, want exactly %q (plain fast-lane-off, no bolt)", notice, want)
|
||||
}
|
||||
}
|
||||
|
||||
// The no-op and no-variant notices must carry their full spec decoration — the ⚡
|
||||
// prefix and the (<short>) name — so a regression dropping either fails loudly
|
||||
// rather than passing a bare substring assertion.
|
||||
func TestFastCommandNoOpAndNoVariantMessagesAreDecorated(t *testing.T) {
|
||||
registry := fastTestRegistry(t)
|
||||
|
||||
plan := planFastCommand(registry, "gpt-4.1-mini", "on")
|
||||
if want := "⚡ Already in the fast lane · " + modelShortName(registry, "gpt-4.1-mini"); plan.message != want {
|
||||
t.Fatalf("already-fast message = %q, want %q", plan.message, want)
|
||||
}
|
||||
|
||||
plan = planFastCommand(registry, "gpt-4.1", "off")
|
||||
if want := "Already on base · " + modelShortName(registry, "gpt-4.1"); plan.message != want {
|
||||
t.Fatalf("already-base message = %q, want %q", plan.message, want)
|
||||
}
|
||||
|
||||
// An unknown / off-catalog model id also reaches the no-variant branch, where
|
||||
// <short> falls back to the raw id.
|
||||
plan = planFastCommand(registry, "some-custom-live-provider-model", "on")
|
||||
if want := "⚡ No fast lane for some-custom-live-provider-model"; plan.message != want {
|
||||
t.Fatalf("no-variant (unknown id) message = %q, want %q", plan.message, want)
|
||||
}
|
||||
}
|
||||
|
||||
// When the underlying switch fails, /fast surfaces the error and leaves the
|
||||
// model unchanged (here: a valid base model but no provider profile to rebuild).
|
||||
func TestHandleFastCommandSurfacesSwitchError(t *testing.T) {
|
||||
m := model{modelName: "gpt-4.1"} // no providerProfile / newProvider wired
|
||||
got, notice := m.handleFastCommand("on")
|
||||
if got.modelName != "gpt-4.1" {
|
||||
t.Fatalf("model must be unchanged on switch failure, got %q", got.modelName)
|
||||
}
|
||||
if strings.HasPrefix(notice, "⚡ Fast lane on ") || strings.HasPrefix(notice, "Fast lane off ") {
|
||||
t.Fatalf("a failed switch must not claim success, got %q", notice)
|
||||
}
|
||||
if strings.TrimSpace(notice) == "" {
|
||||
t.Fatal("a failed switch must surface a reason")
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelShortName(t *testing.T) {
|
||||
registry := fastTestRegistry(t)
|
||||
if got := modelShortName(registry, "gpt-4.1"); strings.TrimSpace(got) == "" {
|
||||
t.Fatal("a known model should resolve to a non-empty short name")
|
||||
}
|
||||
if got := modelShortName(registry, "totally-unknown-model-xyz"); got != "totally-unknown-model-xyz" {
|
||||
t.Fatalf("an unknown model short name = %q, want the raw id", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
package tui
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
tea "charm.land/bubbletea/v2"
|
||||
"charm.land/lipgloss/v2"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/modelregistry"
|
||||
)
|
||||
|
||||
// fastMarkerGlyph is the compact "f" badge shown next to a model that has an
|
||||
// available fast variant — i.e. one where `/fast on` would switch down to a
|
||||
// faster same-family model. It marks the BASE model, never the fast one.
|
||||
const fastMarkerGlyph = "f"
|
||||
|
||||
// fastVariantEntryFor resolves the fast variant of an arbitrary model id, if any.
|
||||
// Like modelContextWindow, it reads the default registry on demand rather than
|
||||
// caching derived state, so a marker can never go stale after a model switch.
|
||||
// DefaultRegistry is memoized, so even the per-frame callers (the active-model
|
||||
// chip renders every frame; the pickers render per keystroke) pay only a map
|
||||
// lookup here, not a registry rebuild.
|
||||
func fastVariantEntryFor(id string) (modelregistry.ModelEntry, bool) {
|
||||
id = strings.TrimSpace(id)
|
||||
if id == "" {
|
||||
return modelregistry.ModelEntry{}, false
|
||||
}
|
||||
registry, err := modelregistry.DefaultRegistry()
|
||||
if err != nil {
|
||||
return modelregistry.ModelEntry{}, false
|
||||
}
|
||||
return registry.FastVariant(id)
|
||||
}
|
||||
|
||||
// currentModelFastVariant resolves the fast variant of the session's active model.
|
||||
func (m model) currentModelFastVariant() (modelregistry.ModelEntry, bool) {
|
||||
return fastVariantEntryFor(m.modelName)
|
||||
}
|
||||
|
||||
// renderFastBadge renders the "f" fast-mode marker in the given (surface-aware)
|
||||
// style, prefixed with a space so it reads as a distinct badge rather than a
|
||||
// trailing letter of the adjacent model name.
|
||||
func renderFastBadge(style lipgloss.Style) string {
|
||||
return " " + style.Render(fastMarkerGlyph)
|
||||
}
|
||||
|
||||
// fastVariantHint is the one-line reveal shown against a focused/selected
|
||||
// fast-capable model. The TUI has no floating-tooltip primitive, so the "hover"
|
||||
// affordance is this inline detail on the active row: it names the target model
|
||||
// and the command that gets there.
|
||||
func fastVariantHint(variantName string) string {
|
||||
variantName = strings.TrimSpace(variantName)
|
||||
if variantName == "" {
|
||||
return ""
|
||||
}
|
||||
return "⚡ fast: " + variantName + " · /fast on"
|
||||
}
|
||||
|
||||
// fastToastDuration is how long a /fast toast stays up before auto-dismissing.
|
||||
const fastToastDuration = 3 * time.Second
|
||||
|
||||
// fastToastExpiredMsg clears a fast-lane toast after fastToastDuration. It is
|
||||
// seq-gated so a newer toast is never dismissed by an older timer.
|
||||
type fastToastExpiredMsg struct{ seq int }
|
||||
|
||||
// showFastToast raises a transient fast-lane notice in the row above the composer
|
||||
// and schedules its dismissal. It replaces the persistent transcript line the
|
||||
// /fast command used to append, so repeated toggles never pile up in history.
|
||||
func (m model) showFastToast(text string) (model, tea.Cmd) {
|
||||
text = strings.TrimSpace(strings.ReplaceAll(text, "\n", " "))
|
||||
if text == "" {
|
||||
return m, nil
|
||||
}
|
||||
m.fastToastSeq++
|
||||
m.fastToast = text
|
||||
seq := m.fastToastSeq
|
||||
return m, tea.Tick(fastToastDuration, func(time.Time) tea.Msg {
|
||||
return fastToastExpiredMsg{seq: seq}
|
||||
})
|
||||
}
|
||||
|
||||
// renderFastToast styles the toast as a filled brand pill so it reads as a
|
||||
// transient popup rather than a permanent status line.
|
||||
func renderFastToast(text string) string {
|
||||
return zeroTheme.badge.Render(" " + strings.TrimSpace(text) + " ")
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
package tui
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/charmbracelet/x/ansi"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/modelregistry"
|
||||
)
|
||||
|
||||
// hasFastBadge reports whether a rendered (styled) line ends with the " f"
|
||||
// fast-mode badge, ignoring trailing layout padding. Labels used in these tests
|
||||
// deliberately contain no trailing " f" so the badge is unambiguous.
|
||||
func hasFastBadge(rendered string) bool {
|
||||
return strings.HasSuffix(strings.TrimRight(ansi.Strip(rendered), " "), " "+fastMarkerGlyph)
|
||||
}
|
||||
|
||||
func TestFastVariantHint(t *testing.T) {
|
||||
if got := fastVariantHint("Claude Haiku 4.5"); got != "⚡ fast: Claude Haiku 4.5 · /fast on" {
|
||||
t.Errorf("fastVariantHint = %q", got)
|
||||
}
|
||||
if got := fastVariantHint(" "); got != "" {
|
||||
t.Errorf("fastVariantHint(blank) = %q, want empty", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderFastBadge(t *testing.T) {
|
||||
if got := ansi.Strip(renderFastBadge(zeroTheme.accent)); got != " "+fastMarkerGlyph {
|
||||
t.Errorf("renderFastBadge visible text = %q, want %q", got, " "+fastMarkerGlyph)
|
||||
}
|
||||
}
|
||||
|
||||
// The active-model chip carries the "f" badge only when the current model has a
|
||||
// fast variant — base models yes, fast models and unpaired top-tier models no.
|
||||
func TestTitleModelSegmentFastBadge(t *testing.T) {
|
||||
cases := []struct {
|
||||
model string
|
||||
want bool
|
||||
}{
|
||||
{"gpt-4.1", true}, // base with a fast variant
|
||||
{"gpt-4.1-mini", false}, // the fast variant itself has none
|
||||
{"claude-opus-4.1", false}, // top-tier, intentionally unpaired
|
||||
{"definitely-not-real", false},
|
||||
{"", false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
seg := model{modelName: tc.model}.titleModelSegment()
|
||||
if got := hasFastBadge(seg); got != tc.want {
|
||||
t.Errorf("titleModelSegment(%q) badge = %v, want %v (visible %q)", tc.model, got, tc.want, ansi.Strip(seg))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A /model picker row shows the badge iff its FastVariant field is populated,
|
||||
// and the field is populated only for base models by decorateFastVariantItems.
|
||||
func TestRenderModelPickerRowFastBadge(t *testing.T) {
|
||||
withFast := renderModelPickerRow(60, false, pickerItem{Label: "gpt-4.1", Value: "gpt-4.1", FastVariant: "GPT-4.1 mini"})
|
||||
if !hasFastBadge(withFast) {
|
||||
t.Errorf("fast-capable row missing badge: %q", ansi.Strip(withFast))
|
||||
}
|
||||
withoutFast := renderModelPickerRow(60, false, pickerItem{Label: "gpt-4.1", Value: "gpt-4.1"})
|
||||
if hasFastBadge(withoutFast) {
|
||||
t.Errorf("row without FastVariant should have no badge: %q", ansi.Strip(withoutFast))
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelPickerItemDetailIncludesFastHint(t *testing.T) {
|
||||
detail := modelPickerItemDetail(pickerItem{Label: "GPT-4.1", Value: "gpt-4.1", FastVariant: "GPT-4.1 mini"})
|
||||
if !strings.Contains(detail, "⚡ fast: GPT-4.1 mini · /fast on") {
|
||||
t.Errorf("detail missing fast hint: %q", detail)
|
||||
}
|
||||
plain := modelPickerItemDetail(pickerItem{Label: "Claude Opus 4.1", Value: "claude-opus-4.1"})
|
||||
if strings.Contains(plain, "fast:") {
|
||||
t.Errorf("non-fast detail should have no hint: %q", plain)
|
||||
}
|
||||
}
|
||||
|
||||
// decorateFastVariantItems stamps the fast variant's display name onto base-model
|
||||
// rows only, using the real registry — the single source of truth for the badge.
|
||||
func TestDecorateFastVariantItems(t *testing.T) {
|
||||
registry, err := modelregistry.DefaultRegistry()
|
||||
if err != nil {
|
||||
t.Fatalf("DefaultRegistry: %v", err)
|
||||
}
|
||||
items := []pickerItem{
|
||||
{Value: "gpt-4.1"}, // base -> GPT-4.1 mini
|
||||
{Value: "gpt-4.1-mini"}, // fast, no variant
|
||||
{Value: "claude-opus-4.1"}, // unpaired
|
||||
{Value: "some-custom-model"}, // not in registry
|
||||
{Value: ""},
|
||||
}
|
||||
decorateFastVariantItems(registry, items)
|
||||
if items[0].FastVariant != "GPT-4.1 mini" {
|
||||
t.Errorf("gpt-4.1 FastVariant = %q, want %q", items[0].FastVariant, "GPT-4.1 mini")
|
||||
}
|
||||
for _, i := range []int{1, 2, 3, 4} {
|
||||
if items[i].FastVariant != "" {
|
||||
t.Errorf("items[%d].FastVariant = %q, want empty", i, items[i].FastVariant)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestShowFastToast(t *testing.T) {
|
||||
// A non-empty message raises the toast, bumps the seq, and schedules dismissal.
|
||||
m, cmd := model{}.showFastToast("⚡ Fast lane on · GPT-4.1 mini")
|
||||
if m.fastToast != "⚡ Fast lane on · GPT-4.1 mini" {
|
||||
t.Fatalf("fastToast = %q", m.fastToast)
|
||||
}
|
||||
if m.fastToastSeq != 1 {
|
||||
t.Fatalf("fastToastSeq = %d, want 1", m.fastToastSeq)
|
||||
}
|
||||
if cmd == nil {
|
||||
t.Fatal("showFastToast must return a dismissal tick cmd")
|
||||
}
|
||||
|
||||
// A multi-line switch-guard message is flattened so the toast stays one line
|
||||
// (the footer slot is fixed-height).
|
||||
m2, _ := model{}.showFastToast("Model\nCannot switch while a run is active.")
|
||||
if strings.Contains(m2.fastToast, "\n") {
|
||||
t.Fatalf("toast must be flattened to one line, got %q", m2.fastToast)
|
||||
}
|
||||
|
||||
// Blank text is a no-op: no toast, no timer.
|
||||
m3, cmd3 := model{}.showFastToast(" ")
|
||||
if m3.fastToast != "" || cmd3 != nil {
|
||||
t.Fatalf("blank text should not raise a toast, got %q / cmd!=nil=%v", m3.fastToast, cmd3 != nil)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderFastToastShowsMessage(t *testing.T) {
|
||||
got := ansi.Strip(renderFastToast("⚡ No fast lane for gpt-5.5"))
|
||||
if !strings.Contains(got, "⚡ No fast lane for gpt-5.5") {
|
||||
t.Fatalf("rendered toast = %q, want it to contain the message", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProviderWizardModelFastBadgeAndHint(t *testing.T) {
|
||||
wizard := &providerWizardState{selectedModel: 0}
|
||||
withFast := wizard.renderSelectableModel(60, 0, providerWizardModel{ID: "gpt-4.1"})
|
||||
if !hasFastBadge(withFast) {
|
||||
t.Errorf("provider wizard base-model row missing badge: %q", ansi.Strip(withFast))
|
||||
}
|
||||
plain := wizard.renderSelectableModel(60, 1, providerWizardModel{ID: "claude-opus-4.1"})
|
||||
if hasFastBadge(plain) {
|
||||
t.Errorf("provider wizard unpaired-model row should have no badge: %q", ansi.Strip(plain))
|
||||
}
|
||||
if detail := providerWizardModelDetail(providerWizardModel{ID: "gpt-4.1"}); !strings.Contains(detail, "⚡ fast:") {
|
||||
t.Errorf("provider wizard detail missing fast hint: %q", detail)
|
||||
}
|
||||
}
|
||||
@@ -121,7 +121,7 @@ func TestEscCancellationLeavesVisibleMarker(t *testing.T) {
|
||||
m.pending = true
|
||||
m.activeRunID = 5
|
||||
m.runCancel = func() {}
|
||||
m.streamingText = []byte("half an answer")
|
||||
m.streamingText = "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 len(next.streamingText) != 0 {
|
||||
if next.streamingText != "" {
|
||||
t.Fatal("streaming text should be cleared after cancel")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -98,9 +98,16 @@ func (m model) modelSupportsVisionTUI() bool {
|
||||
if trimmed == "" {
|
||||
return false
|
||||
}
|
||||
// The curated catalog is authoritative only when it knows the model.
|
||||
if entry, known := m.modelCatalog.Resolve(trimmed); known {
|
||||
return entry.Supports(modelregistry.ModelCapabilityVision)
|
||||
// Check the curated registry first.
|
||||
registry, err := modelregistry.DefaultRegistry()
|
||||
if err == nil {
|
||||
if modelregistry.SupportsVision(registry, trimmed) {
|
||||
return true
|
||||
}
|
||||
// Catalog knows this model and it lacks vision — trust that.
|
||||
if _, known := registry.Get(trimmed); known {
|
||||
return false
|
||||
}
|
||||
}
|
||||
// Check the discovered model list (from models.dev) for InputModalities
|
||||
// containing "image". This covers custom/ollama/cloud models not in the
|
||||
@@ -292,18 +299,6 @@ 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 {
|
||||
|
||||
@@ -483,76 +483,3 @@ 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)
|
||||
}
|
||||
}
|
||||
@@ -36,7 +36,7 @@ func (m model) buildKeybindingGroups() []keybindingGroup {
|
||||
title: "Chat",
|
||||
bindings: []keybinding{
|
||||
{"Enter", "send the message"},
|
||||
{"Shift+Enter / Alt+Enter", "insert a newline (multi-line compose)"},
|
||||
{"Alt+Enter", "insert a newline (multi-line compose)"},
|
||||
{"Esc (\u00d72)", "cancel the run / dismiss a popup / clear the input"},
|
||||
{"Ctrl+C", "cancel the run, then quit"},
|
||||
{"?", "show this help (on an empty input)"},
|
||||
|
||||
@@ -1,31 +0,0 @@
|
||||
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
|
||||
}
|
||||
+59
-173
@@ -19,7 +19,6 @@ 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"
|
||||
@@ -68,7 +67,6 @@ type model struct {
|
||||
gitBranch string
|
||||
providerName string
|
||||
modelName string
|
||||
modelCatalog modelregistry.Registry
|
||||
providerProfile config.ProviderProfile
|
||||
savedProviders []config.ProviderProfile
|
||||
provider zeroruntime.Provider
|
||||
@@ -77,29 +75,24 @@ 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
|
||||
// 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
|
||||
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).
|
||||
@@ -267,30 +260,13 @@ 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 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
|
||||
streamingText string // live assistant text for the current segment
|
||||
streamingReasoning string // live provider reasoning for the current segment
|
||||
streamingReasoningExpanded bool
|
||||
// turnStreamedRunes accumulates every reasoning+answer rune streamed in the
|
||||
@@ -350,9 +326,14 @@ type model struct {
|
||||
// mouse cursor with no button pressed, so it renders in a distinct style —
|
||||
// the visual cue that it's clickable. Requires AllMotion mouse reporting
|
||||
// (see wantsMouseCapture) since idle cursor movement carries no button.
|
||||
hover hoverTarget
|
||||
copyStatus string
|
||||
copyStatusSeq int
|
||||
hover hoverTarget
|
||||
copyStatus string
|
||||
copyStatusSeq int
|
||||
// fastToast is a transient /fast notice shown in the row above the composer for
|
||||
// fastToastDuration, then auto-cleared (seq-gated). It keeps ephemeral fast-lane
|
||||
// feedback out of the persistent transcript.
|
||||
fastToast string
|
||||
fastToastSeq int
|
||||
exitConfirmActive bool
|
||||
exitConfirmSeq int
|
||||
// cancelConfirmActive/cancelConfirmSeq mirror exitConfirmActive/exitConfirmSeq
|
||||
@@ -685,13 +666,9 @@ func newModel(ctx context.Context, options Options) model {
|
||||
sessionStore = sessions.NewStore(sessions.StoreOptions{})
|
||||
}
|
||||
sandboxStore := options.SandboxStore
|
||||
modelCatalog, err := modelregistry.DefaultRegistry()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
usageTracker := options.UsageTracker
|
||||
if usageTracker == nil {
|
||||
usageTracker = usage.NewTracker(usage.TrackerOptions{Registry: &modelCatalog})
|
||||
usageTracker = usage.NewTracker(usage.TrackerOptions{})
|
||||
}
|
||||
prService := options.PrService
|
||||
if prService == nil {
|
||||
@@ -743,7 +720,6 @@ func newModel(ctx context.Context, options Options) model {
|
||||
gitBranch: gitBranch(cwd),
|
||||
providerName: options.ProviderName,
|
||||
modelName: options.ModelName,
|
||||
modelCatalog: modelCatalog,
|
||||
providerProfile: options.ProviderProfile,
|
||||
favoriteModels: favoriteModelSet(options.FavoriteModels),
|
||||
recapsEnabled: options.RecapsEnabled,
|
||||
@@ -796,11 +772,6 @@ 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
|
||||
}
|
||||
@@ -930,22 +901,9 @@ 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()
|
||||
@@ -1062,6 +1020,11 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
m.copyStatus = ""
|
||||
}
|
||||
return m, nil
|
||||
case fastToastExpiredMsg:
|
||||
if msg.seq == m.fastToastSeq {
|
||||
m.fastToast = ""
|
||||
}
|
||||
return m, nil
|
||||
case exitConfirmExpiredMsg:
|
||||
if msg.seq == m.exitConfirmSeq {
|
||||
m.exitConfirmActive = false
|
||||
@@ -1296,7 +1259,7 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
if m.picker != nil {
|
||||
return m.choosePicker()
|
||||
}
|
||||
if keyAlt(msg) || keyShift(msg) {
|
||||
if keyAlt(msg) {
|
||||
if next, ok := m.applyComposerKey(msg); ok {
|
||||
return next, nil
|
||||
}
|
||||
@@ -1617,7 +1580,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 = append(m.streamingText, msg.delta...)
|
||||
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
|
||||
@@ -1900,23 +1863,20 @@ 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.streamingTextString(), "\n"); strings.TrimSpace(text) != "" {
|
||||
if text := strings.TrimRight(m.streamingText, "\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. 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.
|
||||
// metadata a final assistant row would have carried.
|
||||
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 = nil
|
||||
m.streamingText = ""
|
||||
m.streamingReasoning = ""
|
||||
m.streamingReasoningExpanded = false
|
||||
// Roll the completed run's wall-time into the session's rolling average so
|
||||
@@ -2091,14 +2051,14 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
m.streamingReasoning = ""
|
||||
m.streamingReasoningExpanded = false
|
||||
}
|
||||
if text := strings.TrimRight(m.streamingTextString(), "\n"); strings.TrimSpace(text) != "" {
|
||||
if text := strings.TrimRight(m.streamingText, "\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 = nil
|
||||
m.streamingText = ""
|
||||
// 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()
|
||||
@@ -2397,7 +2357,9 @@ func (m model) footerView(width int) string {
|
||||
// a faint idle affordance — discoverable key hints on the left, a jump-to-bottom
|
||||
// cue on the right when scrolled up. Always one line (blank when nothing shows),
|
||||
// so the footer height is unchanged.
|
||||
if copyStatus := strings.TrimSpace(m.copyStatus); copyStatus != "" {
|
||||
if toast := strings.TrimSpace(m.fastToast); toast != "" {
|
||||
footer.WriteString(rightAlignedLine(renderFastToast(toast), width))
|
||||
} else if copyStatus := strings.TrimSpace(m.copyStatus); copyStatus != "" {
|
||||
footer.WriteString(rightAlignedLine(zeroTheme.ink.Render(copyStatus), width))
|
||||
} else if left, right := m.composerIdleHint(), m.jumpToBottomHint(); left != "" || right != "" {
|
||||
footer.WriteString(fitStyledLine(joinHeaderLine(" "+left, right, width), width))
|
||||
@@ -2829,7 +2791,7 @@ func (m model) liveReasoningBodyCap() int {
|
||||
}
|
||||
|
||||
func (m model) interimBlock(width int) string {
|
||||
text := strings.TrimRight(m.streamingTextString(), "\n")
|
||||
text := strings.TrimRight(m.streamingText, "\n")
|
||||
reasoning := strings.TrimRight(m.streamingReasoning, "\n")
|
||||
blocks := []string{}
|
||||
if strings.TrimSpace(reasoning) != "" {
|
||||
@@ -2893,7 +2855,7 @@ func (m model) spinnerGlyph() string {
|
||||
// (reasoning, waiting on the model, or a tool in flight). Cheap and robust — no
|
||||
// transcript scan — so it can't misreport on a long, output-less step.
|
||||
func (m model) workingActivity() string {
|
||||
if strings.TrimSpace(m.streamingTextString()) != "" {
|
||||
if strings.TrimSpace(m.streamingText) != "" {
|
||||
return "writing"
|
||||
}
|
||||
return "thinking"
|
||||
@@ -3864,6 +3826,12 @@ func (m model) handleSubmit() (tea.Model, tea.Cmd) {
|
||||
m, text = m.handleEffortCommand(command.text)
|
||||
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: text})
|
||||
return m, nil
|
||||
case commandFast:
|
||||
var text string
|
||||
m, text = m.handleFastCommand(command.text)
|
||||
// Fast-lane feedback is ephemeral: show it as a transient toast above the
|
||||
// composer (auto-dismissed) instead of a permanent transcript line.
|
||||
return m.showFastToast(text)
|
||||
case commandStyle:
|
||||
text := ""
|
||||
m, text = m.handleStyleCommand(command.text)
|
||||
@@ -3941,62 +3909,6 @@ 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
|
||||
@@ -4011,15 +3923,6 @@ 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{
|
||||
@@ -4207,7 +4110,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.streamingTextString(), "\n"); strings.TrimSpace(text) != "" {
|
||||
if text := strings.TrimRight(m.streamingText, "\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."})
|
||||
@@ -4228,7 +4131,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 = nil
|
||||
m.streamingText = ""
|
||||
m.streamingReasoning = ""
|
||||
m.streamingReasoningExpanded = false
|
||||
// Hard-stop the fade and drop the per-line age map. The next turn's
|
||||
@@ -4305,28 +4208,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 != "" {
|
||||
// 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)
|
||||
}()
|
||||
}
|
||||
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,
|
||||
IncludeLSP: true,
|
||||
Autonomy: selfCorrectAutonomyForMode(options.PermissionMode),
|
||||
})
|
||||
// Inline post-edit diagnostics: edit_file/write_file append error
|
||||
// diagnostics for the file they just wrote to their own output, so the
|
||||
// model sees a break in the same turn. Shares the run's lazy manager.
|
||||
options.FileDiagnostics = agent.NewFileDiagnostics(lspManager, options.Cwd)
|
||||
}
|
||||
|
||||
// Some providers synthesize tool-call ids that repeat within a run (e.g.
|
||||
@@ -4778,13 +4671,6 @@ 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
|
||||
|
||||
@@ -66,8 +66,10 @@ func (m model) modelContextWindow(modelName string) int {
|
||||
if trimmed == "" {
|
||||
return 0
|
||||
}
|
||||
if entry, ok := m.modelCatalog.Resolve(trimmed); ok && entry.ContextLimits.ContextWindow > 0 {
|
||||
return entry.ContextLimits.ContextWindow
|
||||
if registry, err := modelregistry.DefaultRegistry(); err == nil {
|
||||
if entry, ok := registry.Resolve(trimmed); ok && entry.ContextLimits.ContextWindow > 0 {
|
||||
return entry.ContextLimits.ContextWindow
|
||||
}
|
||||
}
|
||||
// Live-discovered window, preferring the ACTIVE provider's models so a model ID
|
||||
// shared across providers resolves to the provider actually in use; only then
|
||||
|
||||
@@ -1,109 +0,0 @@
|
||||
package tui
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/modelregistry"
|
||||
"github.com/Gitlawb/zero/internal/providermodeldiscovery"
|
||||
)
|
||||
|
||||
func TestModelContextWindowUsesCachedCatalog(t *testing.T) {
|
||||
registry := mustTestModelRegistry(t, testModelEntry("custom-long-context", 12345, []modelregistry.ModelCapability{
|
||||
modelregistry.ModelCapabilityChat,
|
||||
}))
|
||||
m := model{modelCatalog: registry}
|
||||
|
||||
if got := m.modelContextWindow("custom-long-context"); got != 12345 {
|
||||
t.Fatalf("modelContextWindow = %d, want 12345", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelSupportsVisionTUITrustsCachedCatalog(t *testing.T) {
|
||||
registry := mustTestModelRegistry(t, testModelEntry("gpt-5-text-only", 12345, []modelregistry.ModelCapability{
|
||||
modelregistry.ModelCapabilityChat,
|
||||
}))
|
||||
m := model{modelName: "gpt-5-text-only", modelCatalog: registry}
|
||||
|
||||
if m.modelSupportsVisionTUI() {
|
||||
t.Fatal("catalog-known text-only model must not fall through to vision name heuristic")
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelSupportsVisionTUIChecksDiscoveredBeforeHeuristic(t *testing.T) {
|
||||
registry := mustTestModelRegistry(t, testModelEntry("custom-known", 12345, []modelregistry.ModelCapability{
|
||||
modelregistry.ModelCapabilityChat,
|
||||
}))
|
||||
m := model{
|
||||
modelName: "gpt-5-text-only",
|
||||
modelCatalog: registry,
|
||||
modelPickerLiveByProvider: map[string][]providermodeldiscovery.Model{
|
||||
"custom": {{ID: "gpt-5-text-only", InputModalities: []string{"text"}}},
|
||||
},
|
||||
}
|
||||
|
||||
if m.modelSupportsVisionTUI() {
|
||||
t.Fatal("discovered text-only model must not fall through to vision name heuristic")
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkModelContextWindowLookup(b *testing.B) {
|
||||
cachedRegistry, err := modelregistry.DefaultRegistry()
|
||||
if err != nil {
|
||||
b.Fatalf("load default registry: %v", err)
|
||||
}
|
||||
|
||||
b.Run("default_registry_each_call", func(b *testing.B) {
|
||||
b.ReportAllocs()
|
||||
for b.Loop() {
|
||||
registry, err := modelregistry.DefaultRegistry()
|
||||
if err != nil {
|
||||
b.Fatalf("load default registry: %v", err)
|
||||
}
|
||||
entry, ok := registry.Resolve("gpt-4.1")
|
||||
if !ok || entry.ContextLimits.ContextWindow <= 0 {
|
||||
b.Fatal("expected gpt-4.1 context window")
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
b.Run("cached_registry", func(b *testing.B) {
|
||||
m := model{modelName: "gpt-4.1", modelCatalog: cachedRegistry}
|
||||
b.ReportAllocs()
|
||||
for b.Loop() {
|
||||
if got := m.modelContextWindow(m.modelName); got <= 0 {
|
||||
b.Fatal("expected gpt-4.1 context window")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func mustTestModelRegistry(t *testing.T, entries ...modelregistry.ModelEntry) modelregistry.Registry {
|
||||
t.Helper()
|
||||
registry, err := modelregistry.NewRegistry(entries)
|
||||
if err != nil {
|
||||
t.Fatalf("create test model registry: %v", err)
|
||||
}
|
||||
return registry
|
||||
}
|
||||
|
||||
func testModelEntry(id string, contextWindow int, capabilities []modelregistry.ModelCapability) modelregistry.ModelEntry {
|
||||
return modelregistry.ModelEntry{
|
||||
ID: id,
|
||||
DisplayName: id,
|
||||
APIModel: id,
|
||||
Provider: modelregistry.ProviderOpenAI,
|
||||
APIProviders: []modelregistry.ProviderKind{modelregistry.ProviderOpenAI},
|
||||
ContextLimits: modelregistry.ContextLimits{ContextWindow: contextWindow, MaxOutputTokens: 1024},
|
||||
Capabilities: capabilities,
|
||||
Cost: modelregistry.ModelCost{
|
||||
Currency: "USD",
|
||||
Unit: "per_1m_tokens",
|
||||
InputPerMillion: 1,
|
||||
OutputPerMillion: 1,
|
||||
Source: "test",
|
||||
SourceLastVerified: "2026-01-01",
|
||||
},
|
||||
Status: modelregistry.ModelStatusActive,
|
||||
Aliases: []string{id},
|
||||
}
|
||||
}
|
||||
@@ -2301,7 +2301,7 @@ func TestStreamingAssistantAfterToolCardGetsRuleSeparator(t *testing.T) {
|
||||
m.width, m.height = 120, 40
|
||||
m.altScreen = true
|
||||
m.pending = true
|
||||
m.streamingText = []byte("Done.")
|
||||
m.streamingText = "Done."
|
||||
m.transcript = append(m.transcript,
|
||||
transcriptRow{kind: rowUser, text: "run it"},
|
||||
transcriptRow{kind: rowAssistant, text: "I'll run it first."},
|
||||
|
||||
@@ -20,10 +20,6 @@ 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()
|
||||
|
||||
@@ -45,11 +41,6 @@ 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) {
|
||||
|
||||
@@ -1151,29 +1151,6 @@ 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
|
||||
@@ -1182,9 +1159,6 @@ 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()
|
||||
|
||||
@@ -1,29 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
+34
-71
@@ -43,6 +43,10 @@ type pickerItem struct {
|
||||
Remote bool
|
||||
Local bool
|
||||
Favorite bool
|
||||
// FastVariant is the display name of this model's faster same-family variant
|
||||
// (empty when it has none). Set only for /model rows; drives the "f" fast-mode
|
||||
// badge and the focused-row "⚡ fast: … · /fast on" hint.
|
||||
FastVariant string
|
||||
}
|
||||
|
||||
// commandPicker is a generic single-select overlay reused by /model and /effort
|
||||
@@ -102,83 +106,16 @@ func (p *commandPicker) applyQuery() {
|
||||
p.selected = clampInt(p.selected, 0, maxInt(0, len(p.items)-1))
|
||||
return
|
||||
}
|
||||
|
||||
// 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
|
||||
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)
|
||||
}
|
||||
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.
|
||||
@@ -222,9 +159,35 @@ func (m model) newModelPicker() *commandPicker {
|
||||
if len(items) == 0 {
|
||||
return nil
|
||||
}
|
||||
decorateFastVariantItems(registry, items)
|
||||
return &commandPicker{kind: pickerModel, title: "Choose a model", items: items, allItems: append([]pickerItem{}, items...), selected: 0}
|
||||
}
|
||||
|
||||
// decorateFastVariantItems stamps each /model row with the display name of its
|
||||
// fast variant (if any), in one pass over the assembled list rather than in every
|
||||
// item builder. Rows whose model id is not in the curated registry (custom or
|
||||
// live-discovered) simply keep an empty FastVariant and show no badge.
|
||||
func decorateFastVariantItems(registry modelregistry.Registry, items []pickerItem) {
|
||||
for index := range items {
|
||||
id := strings.TrimSpace(items[index].Value)
|
||||
if id == "" {
|
||||
continue
|
||||
}
|
||||
if fast, ok := registry.FastVariant(id); ok {
|
||||
items[index].FastVariant = fastVariantDisplayName(fast)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// fastVariantDisplayName is the user-facing name of a fast variant, preferring
|
||||
// its display name and falling back to the id.
|
||||
func fastVariantDisplayName(entry modelregistry.ModelEntry) string {
|
||||
if name := strings.TrimSpace(entry.DisplayName); name != "" {
|
||||
return name
|
||||
}
|
||||
return strings.TrimSpace(entry.ID)
|
||||
}
|
||||
|
||||
// modelPickerProviders returns the providers to list in /model: all saved
|
||||
// providers, falling back to the active profile when none were threaded in.
|
||||
func (m model) modelPickerProviders() []config.ProviderProfile {
|
||||
|
||||
@@ -1,72 +0,0 @@
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
Loaded 100 of 117 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user