Compare commits

..
Author SHA1 Message Date
Kevin Codex 2e2c181d28 fix: surface competing block-anchor spans; make write-barrier test observe cross-probe order
- edit_file: blockAnchorReplacer now returns every span clearing the
  similarity threshold (best first) instead of silently picking the
  single most-similar block. Hiding competitors from the cascade's
  distinct-candidate check could edit the wrong block when two blocks
  share anchors and both interiors are plausible; one clear winner still
  resolves, two plausible blocks now report ambiguity (covered by a
  two-block test that asserts the file is untouched).
- agent tests: the parallel write-barrier test now records both probes
  into one mutex-guarded shared log and asserts the cross-probe ordering
  it always claimed to (write starts after the first read batch ends;
  second read batch starts after the write ends). The previous version
  only compared indices within the read probe's own log, which a
  regression letting the write overlap a batch could still pass; the
  stray no-op mutex assignment is gone.
2026-07-05 17:16:24 +08:00
Kevin Codex 4643dc9348 fix: address review round 2 — permission-event race, fuzzy ambiguity, spill reachability
- agent: serialize OnPermission behind the same mutex as
  OnPermissionRequest in the parallel read batch. buildPermissionEvent
  fires for batched reads with a non-trivial sandbox decision (e.g. a
  grant-matched extra root), and front-end handlers append to shared
  session-recording state without their own locking — two batched reads
  emitting concurrently was a data race the pre-batch serial loop never
  allowed.
- edit_file: reject two or more DISTINCT candidate spans from one fuzzy
  strategy as ambiguous instead of silently editing the first. Duplicate
  blocks at different indentation each resolve to a unique literal, so
  the previous occurrence-uniqueness check passed and wrote the wrong
  span. A single candidate occurring at multiple positions still
  cascades to stricter strategies as before.
- tools: make spill files actually reachable by the recovery path the
  truncation notice advertises — resolveScopedReadPath now recognizes
  files inside the per-uid spill dir (containment verified after symlink
  resolution, so a planted link cannot smuggle an out-of-scope file
  through the gate); read_file/grep can follow the notice instead of
  re-running the command.
2026-07-05 16:57:39 +08:00
Kevin Codex 57e34c790d fix: address PR review feedback on fuzzy edits, diagnostics, spill
- edit_file: adapt the replacement to the fuzzy-resolved span — when the
  span is a uniformly deeper-indented copy of old_string (line-trimmed /
  indentation-flexible matches), the same indentation delta is applied to
  new_string, and a span ending in a bare CR keeps its CR so CRLF files
  never gain a mixed line ending; block-anchor matches with drifted
  interiors are deliberately left untouched
- diagnostics: render workspace-relative paths in inline post-edit
  diagnostics (base name outside the workspace) so absolute paths never
  leak the local username or directory layout into the model prompt or
  session transcript
- spill: harden the spill directory for shared temp dirs — per-uid name,
  reject symlinked or non-directory paths (MkdirAll follows symlinks),
  and on Unix require current-user ownership
- bash: the branch's output cap and truncation-flag threading are
  superseded by main's capture-time bounding (boundedBuffer +
  budgetBashCapture), which already sets Result.Truncated; this rebase
  adopts that design and keeps spill-to-disk on exec_command, where the
  full output is still buffered and a spill is genuinely complete
2026-07-05 16:47:18 +08:00
Kevin Codex d0ae9c98e3 feat: agent quality, caching, retry, and tooling upgrades
Ten self-contained improvements to the core agent engine and tools:

- edit_file: fuzzy fallback matching (7-strategy cascade with uniqueness
  protection and a disproportion guard) when old_string does not match
  byte-for-byte; exact match remains the fast path
- edit_file/write_file: inline error-severity LSP diagnostics in tool
  output so the model sees a break it introduced in the same turn
- lsp: expand server registry from 4 to 28 languages; doctor grades only
  tier-1 servers and lists the long tail informationally
- agent: consecutive read-only tool calls execute concurrently (bounded
  at 8), consumed in original order, never spanning a mutating call
- anthropic: cache_control breakpoints on the last two conversation
  messages so the transcript prefix stays a prompt-cache hit
- providerio: retry schedule raised to 6 attempts with exponential
  2s-30s backoff when the server sends no Retry-After
- tools: truncated exec/bash output spills to a redacted temp file
  referenced in the truncation notice; bash output (previously
  unbounded) now shares exec_command's token budget
- web_fetch: dependency-free HTML-to-markdown conversion (format:
  auto|raw|markdown) and larger limits (256KiB default, 2MiB max, 30s)
- modelregistry: live models.dev overlay refreshes context limits and
  base pricing from a TTL-cached fetch; opt-in at the CLI entrypoint
- tools: opt-in format-on-write (ZERO_FORMAT_ON_WRITE=1) runs standard
  formatters across 27 extensions, ordered before the file-tracker
  re-baseline so conflict detection is preserved
2026-07-05 16:43:50 +08:00
Anandan afb5010a46 polish v2: smoothness & delight (#484)
* feat(tui): actionable hints on provider/model errors

Classify a failed turn's provider error into auth / rate-limit /
connectivity / model-not-found / context-overflow and render a one-line
next step below the red error row, instead of dumping an identical raw
blob for every failure mode.

The classifier lives in a new internal/errhint package so both the TUI
error row and the CLI exec path can share it (TUIHint references slash
commands, CLIHint references zero subcommands). Classification is a
conservative string heuristic keyed off providerio.ClassifiedError's
prefixes plus lower-level DNS/TLS/timeout/context-length signatures,
since the numeric HTTP status is gone by the time the error reaches a UI
surface.

* feat(tools): budget bash output with head+tail truncation

bash was the one tool that returned command output with no byte cap —
'cat large.log' or a verbose test run could dump megabytes into context
and force compaction, while every other read/search tool already applies
a budget.

Cap stdout and stderr at 96KB each, keeping the head and tail of an
oversized stream (build/test failures usually land at the tail) and
dropping the middle behind a marker that points at redirect-to-file +
read_file. Record raw_bytes/emitted_bytes/estimated_tokens/truncated in
Meta like the other tools. Sandbox-denial and shell-issue detection still
run on the full raw output before budgeting.

* fix(gemini): retry 401 with refreshed OAuth token

The Gemini provider used providerio.SendWithRetry (429/503/529 only),
so an expired OAuth token surfaced as a terminal 'API key not valid'
auth error while OpenAI and Anthropic silently recovered via
SendWithAuthRetry's force-refresh-and-replay on 401.

Plumb OAuthResolver through gemini.Options/Provider and the factory
(exactly as the other two providers already do) and switch the stream
send to SendWithAuthRetry. With a nil resolver (API-key users) behavior
is unchanged.

* fix(errhint): gate classification on provider-origin marker

Address CodeRabbit review on #1: agentResponseMsg.err can also carry
local failures (a tool's 'permission denied', a 'file does not exist',
a config error), and broad substrings would have attached a bogus
/provider or /model hint to those.

Classify now returns Unknown unless the message carries a provider marker
the provider layer always attaches (auth error: / rate limit error: /
provider error: / provider request error: / provider stream error:).
Local errors never draw a provider hint; provider errors sub-classify as
before. Added local-failure test cases proving they stay Unknown.

* feat(cli): actionable hint on text-mode exec provider errors

zero exec printed provider failures as a bare '[zero] <raw error>' with
no next step. Append a one-line hint (reusing the shared errhint
classifier) for recognized provider failures — 'run `zero auth`',
'run `zero doctor`', etc. The provider-origin gate means non-provider
codes (sandbox_error, mcp_error) and JSON/stream-json output are
untouched.

* fix(tools): set Result.Truncated for budgeted bash output

Address CodeRabbit review on #3: budgetBashOutput recorded truncation in
meta["truncated"] but never set the Result.Truncated struct field, so
consumers reading the struct (like glob/web_fetch/read_minified do) would
miss truncated bash results. Return the bool from budgetBashOutput and
set Result.Truncated on every return path.

* perf(tui): reuse one LSP manager across prompts

runAgentWithOptions built a fresh lsp.NewManager per run and shut it down
when the run returned, so gopls (and every other language server)
cold-started on the first edit of every turn — 200ms-2s of latency the
manager's own design (long-lived, reused servers) exists to avoid.

Build one manager per session in newModel (cheap: servers start lazily on
first Check) and reuse it across runs; the SelfCorrector still gets a
fresh checker wrapping the shared manager each run. Torn down with a
short deadline in quit(). Runs fall back to a per-run manager when the
session manager is absent (cwd unknown, or a directly-built test model).

* perf(tui): accumulate streamingText as []byte (O(1) append)

m.streamingText += msg.delta was O(len) per delta -> O(n²) across a long
generation (a 10k-token code gen allocates hundreds of MB of intermediate
strings and the UI gets progressively laggier). Accumulate into a []byte
with append (O(1) amortized) and read via streamingTextString().

A []byte rather than strings.Builder because the TUI model is copied by
value on every Update, which would trip strings.Builder's copy check;
[]byte is nil-safe and copies cleanly like the other slice fields.

* perf(tui): coalesce streamed text deltas to one frame

Every OnText delta was its own tea.Msg, so a fast provider (100+ tok/s)
drove 100+ full Update->View cycles per second, each re-parsing the
growing markdown — visible stutter over SSH and wasted CPU.

Batch agentTextMsg deltas at the runtime sink over a ~16ms frame and
forward them as a single message, decoupling render rate from token rate.
Any non-text message flushes pending text first so ordering with
tool-call / reasoning / row messages is preserved; the final
agentResponseMsg (a tea.Cmd return, not a sink message) is safe because
the model already drops deltas for an inactive runID.

* fix(agent): raise tool-failure stop threshold 4->6

guardrails.go halted a run after 4 consecutive same-error tool failures.
A corrective hint fires at 2, and a model iterating on a genuinely tricky
edit can legitimately fail a couple more times after the hint while
converging — stopping at 4 cut those runs short. Raise to 6; the streak
still resets the instant the tool succeeds or hits a different error, so
only true same-error loops are affected.

* fix(agent): more reconnect retries with jitter; narrow to transport errors

#19: bump maxStreamReconnects 2->4 and add up-to-50% jitter on the
exponential backoff (capped at 8s), so a multi-second network blip is
ridden out instead of killing the run on a 2s hiccup, and concurrent runs
(swarms, cron fleets) don't reconnect in lockstep.

#34: drop the '502'/'503' substring matches from shouldReconnect. 503
already exhausted providerio.SendWithRetry (retrying here is a redundant
double-retry) and 502 is non-idempotent by providerio's rule (the POST
may have been processed). Only genuine transport failures — where no
response was received — reconnect now.

* fix(agent): reduce stall retries 2->1 to bound stuck-session time

A no-output stream stall is detected only after the full stream idle
timeout (~5min) elapses, and each retry can idle again — so 2 retries
left an interactive session frozen for ~15min. Drop to 1 retry: keeps the
common single-hiccup recovery while bounding the worst case to ~2x the
idle timeout.

* fix(tui): serialize coalescer forwarding to preserve message order

Address CodeRabbit review on #7: flush() and the run-switch path drained
the text buffer under c.mu but called forward() after releasing it, so a
timer-fired text flush could race a concurrent non-text send() and land
after the tool-call/reasoning message it should precede.

Hold c.mu across drain AND forward (drainAndForwardLocked), so whoever
holds the lock delivers atomically and the other caller blocks until it
is done — text can never overtake a following non-text message. Added a
concurrency stress test (timer racing an inline send) that passes under
-race.

* feat(tui): accept /compact now

Bare /compact already triggers compaction, but /compact now — what users
reach for when the context gauge climbs — hit the usage-error branch.
Accept the 'now' keyword and advertise it in the command usage.

* feat(tui): add /retry, /edit, /copy, /export commands

/retry resends the last prompt; /edit recalls it into the composer to
tweak and resend — both read a new lastPrompt field captured verbatim
(pre-expansion) in launchPrompt. /copy puts the last answer on the
clipboard via the existing copy machinery; /export writes a plain
role-prefixed transcript to a file (timestamped default, or a given
path). Real gaps for SSH users with no mouse select.

* feat(tools): separator-insensitive tool_search matching

tool_search ranked deferred tools by exact substring, so a model that
typed 'webfetch' (dropping the underscore) matched nothing and looped.
Add a separator-squashing fallback ('web_fetch' -> 'webfetch') scored
below an exact substring match, so precise queries still rank first.

* feat(tools): signal ask_user dismissal distinctly from a blank field

Empty answers all rendered as '(no answer provided)', so the model
couldn't tell a wholesale dismissal (user closed the prompt without
answering anything) from one field left blank amid real answers — and
might invent a default. FormatAskUserAnswers now flags a full dismissal
up front as a skip and marks individual empties '(left blank)' vs
'(skipped)'.

* feat(tui): fuzzy ranking in model/session pickers

The pickers filtered by flat strings.Contains with no ranking, so
finding 'sonnet 4.5' among 50+ models meant an exact prefix or scrolling.
Rank matches (exact < prefix < contains < subsequence, reusing
fuzzySubsequenceGap) so the closest match lands on top, and 'snt45' now
matches 'Sonnet 4.5'. Groups stay contiguous — ordered by their best
match, never split — so the grouped model picker still renders one header
per provider. Covers both /model and /resume (shared commandPicker).

* perf(agent): cache per-tool schema render across turns

partitionTools re-ran the recursive schema->map conversion
(schemaToRuntimeMap) for every tool on every turn — redundant, since a
tool's advertised name/description/schema is stable for the run. Add a
per-run definition cache keyed by tool name (nil-safe; the plain
partitionTools entrypoint and tests pass nil for fresh renders).

The partitioning itself (visibility, deferral, ordering) still recomputes
every turn — it must, because a tool's deferred state can flip mid-run
(swarm tools un-defer once a swarm is active). Only the expensive schema
render is memoized. tool_search is excluded by its callers (dynamic
description) so it never poisons the cache.

* perf(agent): calibrate compaction estimate against real prompt tokens

The byte/4 heuristic (ApproxTextTokens) over-counts code-heavy content
~15-20%, so with triggerRatio=0.7 compaction fired at ~60% of true
capacity — premature summarizer calls that degrade quality. Each turn now
folds (rawEstimate, provider InputTokens) into an EMA calibration ratio,
clamped per-sample to a sane band; maybeCompact and the reactive path
scale their estimates by it. Turn 1 uses the raw estimate (no data yet);
later turns compact near real capacity.

* style: gofmt the new agent test files

The CI gofmt check flagged the #22/#25 test files (struct-field and
switch alignment). No behavior change.

* feat(agent): turn-based plan staleness reminder

The stale-plan nudge only fired after 10 tool calls since the last
update_plan, so a plan that drifts stale across many low-tool-call turns
(one call per turn) took many turns to catch — the 'plan set turn 1 stays
pending forever' case. Add a turn-count complement: after
stalePlanTurnThreshold (8) turns without an update AND with items still
pending, fire the same one-shot reminder. Gated on pending items so a
fully-completed plan is never nagged; the existing tool-call trigger is
unchanged. (The roadmap's fuzzy 'tool calls unrelated to plan items'
relevance match is deliberately omitted — too false-positive-prone.)

* feat(agent): preserve recent-edits summary across compaction

Compaction preserved the plan, loaded skills, tool schemas, and project
instructions, but not file edits — so after the editing turns were elided
the model knew it had changed a file only vaguely, and would re-read to
rediscover its own footprint. Extract write_file/edit_file targets from
the elided middle with a one-line note from each tool result, merge them
across repeated compactions (newer note per path wins), and carry them in
the preserved-state JSON (capped at 20 paths / 160 bytes per note).

* feat(tui): workspace-scoped /resume + fewer event reads

/resume listed sessions from every project globally, and read the full
event file for each of them to test for resumable content. Filter the
picker (and /resume latest) by the current workspace Cwd, checked BEFORE
the per-session read so a large global history no longer pays N full file
reads to build one workspace's list; also skip zero-event sessions via
metadata without any read. Sessions with no recorded Cwd (older runs) and
an unknown current workspace stay visible so nothing is hidden that can't
be confidently placed. Explicit /resume <id> still resolves any session.

* feat(tui): warn on model switch that drops staged images

Staged images were silently dropped at submit when the active model had
no vision support — the user only learned after sending. Chips already
render above the composer (renderAttachmentChips); add an immediate amber
warning line to the /model switch status (both same-provider and
cross-provider paths) when images are staged and the newly selected model
can't accept them, so the drop is surfaced at switch time.

* feat(tui): inline recovery affordance on setup errors

A first-run setup error rendered as a bare red 'error: <text>' with no
pointer to the fix. Append a faint, stage-tailored recovery line (edit
the endpoint/name/key above then Enter, or pick another model) that names
only keys the stage's footer confirms, so the guidance is always accurate.
Empty for stages whose footer already makes recovery obvious.

* fix: address CodeRabbit review (5 findings)

- loop.go: recompute the compaction calibration estimate at calibration
  time, not request-build time, so a reactive-compaction reissue that
  shrinks the request feeds the EMA the sample that actually completed.
- errhint: match bare HTTP status codes (401/403/429/529) only at digit
  boundaries, so an incidental number ('4290ms', 'id 14015') isn't
  mis-bucketed as auth/rate-limit.
- bash: drop the stray unmatched ']' in the output-truncation marker.
- tui /retry: apply the same exiting/compactInFlight guards commandPrompt
  has, so a retry can't race compactResultMsg's rewrite of session state.
- tui /export: write the transcript 0o600 (it may contain echoed
  secrets), not world/group-readable 0o644.

* test(tui): skip export file-mode assertion on Windows

os.WriteFile ignores Unix permission bits on Windows and Stat reports
0666, so the 0o600 assertion failed there. The production 0o600 is
honored on Linux/macOS (where it matters); guard the check with
runtime.GOOS so Windows smoke passes.

* fix(tui): align /resume latest filters + case-insensitive workspace match on Windows

Address CodeRabbit review on #21:
- latestResumableInWorkspace now applies the SAME filters as the picker
  (skip zero-event metadata and empty/failed runs), so /resume latest
  can't land on a session the picker intentionally hides.
- sessionMatchesWorkspace compares paths case-insensitively on Windows
  (its filesystem is), so the same workspace spelled with different casing
  no longer hides resumable sessions; other platforms stay case-sensitive.

* fix(tui): /retry resends the last prompt's image/PDF attachments

/retry relaunched only lastPrompt's text, but launchPrompt clears the
pending image/document queues once a turn is sent. A vision- or PDF-backed
prompt that failed would silently retry as text-only and answer a different
task. Snapshot the consumed attachments at launch and re-stage them on
/retry so launchPrompt rebuilds an identical request (document preamble,
images, and the submit-time vision re-check). startNewSession drops the
snapshot so a post-/new /retry cannot leak the previous session's
attachments. Also correct the OAuthResolver factory comment, which omitted
the Gemini provider it is now passed to.

* fix: address jatmn review (reconnect 5xx, /edit attachments, edit-cap, coalescer test)

- reconnect: exclude HTTP 5xx (500/502/503/504) before the transport
  substring match so "504 Gateway Timeout" no longer slips through on the
  generic "timeout" needle. A gateway timeout is non-idempotent — the
  completion POST may already have reached the model — so replaying the
  connect risked duplicate billable work. Reuses a digit-boundary status
  matcher exported from errhint (HasStatusCode). Test asserts 504/500 stay
  non-reconnectable.

- tui /edit: re-stage the remembered image/PDF snapshot alongside the
  recalled prompt text, mirroring the /retry fix — editing a vision- or
  document-backed prompt no longer silently resends a text-only version.

- compaction recent-edits: order edits by LAST edit, not first, and move a
  re-touched path to the newest position on merge, so the maxRecentEdits
  tail cap keeps the file the model most recently edited instead of dropping
  it. Adds mergeRecentEdits (edit-specific) and a capped-list regression.

- coalescer test: inject the frame timer via afterFunc so
  TestCoalescerBatchesDeltas proves buffering deterministically instead of
  racing the real 16ms wall-clock timer.

* fix(tools): bound bash output capture in memory, not just model-visible text

Previously stdout/stderr streamed into unbounded bytes.Buffers and were only
truncated to the 96 KiB budget AFTER the command finished, so a runaway command
(cat huge.log, yes) could grow Zero's memory until it stalled or OOMed before
budgetBashOutput ever ran. Replace the capture with boundedBuffer, an io.Writer
that keeps only the head+tail each stream will surface and discards the middle as
it arrives, while counting the true total for the truncation marker. Memory is now
bounded to ~head+2×tail per stream regardless of output size.

truncateHeadTail is refactored onto a total-aware core so the string path (and its
tests) are unchanged, and budgetBashCapture reports the true raw_bytes even though
only a bounded slice was held. Adds boundedBuffer/capture unit tests and an
end-to-end test that streams ~5× the budget and asserts bounded emission with an
accurate raw_bytes.
2026-07-05 13:15:14 +05:30
68 changed files with 3341 additions and 207 deletions

No files matched your search

+41 -4
View File
@@ -325,6 +325,43 @@ type compactionState struct {
// OnText is deliberately NOT forwarded (compaction stays invisible to the user),
// but its token COST must still be counted so usage reports and budgets include it.
onUsage func(Usage)
// calibrationRatio scales the raw byte/4 token estimate toward the provider's
// real prompt-token count. ApproxTextTokens over-counts code-heavy content by
// ~15-20%, which would trip compaction early (at ~60% of true capacity). It
// starts at 1.0 and converges via an EMA as each turn reports actual usage, so
// later turns compact nearer to real capacity. Zero is treated as 1.0.
calibrationRatio float64
}
// calibrate folds one turn's (rawEstimate, actualPromptTokens) sample into the
// running calibration ratio. A single sample is clamped to a sane band so an
// outlier (a huge cache-read turn, a provider-overhead spike) can't skew the
// estimate enough to disable or thrash compaction.
func (state *compactionState) calibrate(rawEstimate int, actualPromptTokens int) {
if !state.enabled || rawEstimate <= 0 || actualPromptTokens <= 0 {
return
}
sample := float64(actualPromptTokens) / float64(rawEstimate)
if sample < 0.5 {
sample = 0.5
} else if sample > 2.0 {
sample = 2.0
}
if state.calibrationRatio <= 0 {
state.calibrationRatio = 1.0
}
const alpha = 0.3 // weight on the newest sample; smooths jitter across turns
state.calibrationRatio = state.calibrationRatio*(1-alpha) + sample*alpha
}
// calibratedTokens applies the learned ratio to a raw estimate. Before any sample
// arrives (ratio unset) it returns the raw estimate unchanged.
func (state *compactionState) calibratedTokens(raw int) int {
if state.calibrationRatio <= 0 {
return raw
}
return int(float64(raw) * state.calibrationRatio)
}
func newCompactionState(options Options) *compactionState {
@@ -353,7 +390,7 @@ func (state *compactionState) maybeCompact(
// the messages; both the threshold check and the shrink check below use the
// same term so they stay consistent.
toolTokens := estimateToolDefTokens(tools)
size := estimateTokens(messages) + toolTokens
size := state.calibratedTokens(estimateTokens(messages) + toolTokens)
if size <= state.threshold {
return messages
}
@@ -370,7 +407,7 @@ func (state *compactionState) maybeCompact(
// entirely and preserve recent turns verbatim.
if pruned, reclaimed := pruneStaleToolOutput(messages, state.preserveLast); reclaimed > 0 {
messages = pruned
size = estimateTokens(messages) + toolTokens
size = state.calibratedTokens(estimateTokens(messages) + toolTokens)
if size <= state.threshold {
state.lowWaterMark = size
return messages
@@ -386,7 +423,7 @@ func (state *compactionState) maybeCompact(
// later turn) can try again; we never drop messages on failure here.
return messages
}
newSize := estimateTokens(compacted) + toolTokens
newSize := state.calibratedTokens(estimateTokens(compacted) + toolTokens)
if newSize >= size {
// Compaction did not actually shrink anything (e.g. nothing to
// summarize). Leave the history untouched and don't churn next turn.
@@ -444,7 +481,7 @@ func (state *compactionState) recover(
// the SAME combined (messages + tool-defs) domain maybeCompact uses, so the
// proactive shrink-guard compares like with like.
state.reactiveAttempted = true
state.lowWaterMark = estimateTokens(result) + estimateToolDefTokens(tools)
state.lowWaterMark = state.calibratedTokens(estimateTokens(result) + estimateToolDefTokens(tools))
return result, true, nil
}
@@ -0,0 +1,50 @@
package agent
import "testing"
func TestCalibrationConvergesTowardActual(t *testing.T) {
state := &compactionState{enabled: true}
// Before any sample, the estimate is unchanged.
if got := state.calibratedTokens(1000); got != 1000 {
t.Fatalf("uncalibrated tokens = %d, want 1000", got)
}
// The estimator over-counts: our raw estimate is 1000 but the provider reports
// 850 (a 0.85 ratio). Feeding that sample repeatedly should pull the ratio
// toward 0.85, so a 1000-token raw estimate calibrates downward.
for range 20 {
state.calibrate(1000, 850)
}
got := state.calibratedTokens(1000)
if got >= 1000 || got < 820 || got > 880 {
t.Fatalf("calibrated tokens = %d, want ~850 after convergence", got)
}
}
func TestCalibrateIgnoresDegenerateSamples(t *testing.T) {
state := &compactionState{enabled: true}
state.calibrate(0, 500) // zero estimate
state.calibrate(500, 0) // zero actual
if state.calibrationRatio != 0 {
t.Fatalf("degenerate samples must not move the ratio, got %v", state.calibrationRatio)
}
// Disabled state never calibrates.
disabled := &compactionState{enabled: false}
disabled.calibrate(1000, 850)
if disabled.calibrationRatio != 0 {
t.Fatalf("disabled compaction must not calibrate, got %v", disabled.calibrationRatio)
}
}
func TestCalibrateClampsOutliers(t *testing.T) {
state := &compactionState{enabled: true}
// A wild outlier (10x) is clamped to 2.0, so one bad sample can't blow up the
// ratio and disable compaction.
for range 50 {
state.calibrate(1000, 10000)
}
if state.calibrationRatio > 2.0 {
t.Fatalf("ratio should be clamped at 2.0, got %v", state.calibrationRatio)
}
}
+194 -4
View File
@@ -20,6 +20,16 @@ const (
toolNameUpdatePlan = "update_plan"
toolNameToolSearch = "tool_search"
toolNameSkill = "skill"
toolNameWriteFile = "write_file"
toolNameEditFile = "edit_file"
)
// maxRecentEdits caps how many edited files are carried across a compaction, and
// maxEditNoteBytes caps each edit's one-line note, so a long editing run can't
// bloat the preserved-state block.
const (
maxRecentEdits = 20
maxEditNoteBytes = 160
)
const (
@@ -100,12 +110,123 @@ func formatPlanArguments(arguments string) string {
}
// skillEntry is a named preserved body. It began as loaded-skill state and is
// reused for loaded tools and project instruction blocks.
// reused for loaded tools, project instruction blocks, and recent file edits.
type skillEntry struct {
name string
body string
}
// recentEdits returns the files mutated by write_file/edit_file calls in messages
// — latest note per path, in last-seen order — as skillEntry{name: path, body:
// note}. After compaction elides the editing turns, this tells the model WHAT it
// changed in each file (from the tool's result) so it need not re-read to
// rediscover its own footprint. Capped at maxRecentEdits paths.
//
// Ordering is by LAST edit, not first: re-editing a file moves it to the newest
// position so the tail cap below keeps the file the model most recently touched
// rather than an earlier, now-stale entry.
func recentEdits(messages []zeroruntime.Message) []skillEntry {
pathByID := map[string]string{}
sequence := make([]string, 0)
for _, message := range messages {
for _, call := range message.ToolCalls {
if call.Name != toolNameWriteFile && call.Name != toolNameEditFile {
continue
}
path := editPathFromArguments(call.Arguments)
if path == "" {
continue
}
if call.ID != "" {
pathByID[call.ID] = path
}
sequence = append(sequence, path)
}
}
order := lastSeenOrder(sequence)
if len(order) == 0 {
return nil
}
noteByPath := map[string]string{}
for _, message := range messages {
if message.Role != zeroruntime.MessageRoleTool || message.ToolCallID == "" {
continue
}
if path, ok := pathByID[message.ToolCallID]; ok {
noteByPath[path] = editNote(message.Content)
}
}
// Keep the most recent maxRecentEdits paths (the tail of last-seen order).
if len(order) > maxRecentEdits {
order = order[len(order)-maxRecentEdits:]
}
entries := make([]skillEntry, 0, len(order))
for _, path := range order {
entries = append(entries, skillEntry{name: path, body: noteByPath[path]})
}
return entries
}
// lastSeenOrder dedupes paths keeping each at its LAST occurrence, preserving
// chronological order otherwise. A re-edited path therefore lands at the newest
// position rather than staying pinned to where it first appeared.
func lastSeenOrder(paths []string) []string {
lastIdx := make(map[string]int, len(paths))
for i, p := range paths {
lastIdx[p] = i
}
order := make([]string, 0, len(lastIdx))
for i, p := range paths {
if lastIdx[p] == i {
order = append(order, p)
}
}
return order
}
// editPathFromArguments pulls the target file path from a write_file/edit_file
// call's JSON arguments (path/file/file_path/filename aliases). Returns "" on
// malformed arguments or no path.
func editPathFromArguments(arguments string) string {
var parsed struct {
Path string `json:"path"`
File string `json:"file"`
FilePath string `json:"file_path"`
Filename string `json:"filename"`
}
if err := json.Unmarshal([]byte(strings.TrimSpace(arguments)), &parsed); err != nil {
return ""
}
for _, candidate := range []string{parsed.Path, parsed.File, parsed.FilePath, parsed.Filename} {
if s := strings.TrimSpace(candidate); s != "" {
return s
}
}
return ""
}
// editNote reduces an edit tool's result to a single short line (its first
// non-empty line, byte-capped on a rune boundary) for the preserved summary.
func editNote(content string) string {
for _, line := range strings.Split(content, "\n") {
s := strings.TrimSpace(line)
if s == "" {
continue
}
if len(s) > maxEditNoteBytes {
limit := maxEditNoteBytes
for limit > 0 && !utf8.RuneStart(s[limit]) {
limit--
}
s = s[:limit] + "…"
}
return s
}
return ""
}
// loadedSkills returns the skills loaded via the skill tool in messages — the
// latest body per name, in first-seen order — matching each skill tool call to
// its tool result by id.
@@ -310,11 +431,17 @@ func capBody(body string) string {
// preservedState is the JSON shape of the carried-across-compaction block.
type preservedState struct {
Plan string `json:"plan,omitempty"`
RecentEdits []preservedEdit `json:"recent_edits,omitempty"`
Tools []preservedTool `json:"tools,omitempty"`
Skills []preservedSkill `json:"skills,omitempty"`
ProjectInstructions []preservedInstruction `json:"project_instructions,omitempty"`
}
type preservedEdit struct {
Path string `json:"path"`
Note string `json:"note,omitempty"`
}
type preservedTool struct {
Name string `json:"name"`
Body string `json:"body"`
@@ -347,6 +474,12 @@ func appendPreservedState(summary string, middle []zeroruntime.Message) string {
plan = priorState.Plan
}
// Recent edits: merge edits preserved earlier with fresh write_file/edit_file
// results in middle (newer note per path wins AND moves to the newest
// position), so the tail cap keeps the files the model most recently touched
// across repeated compactions rather than dropping a just-re-edited file.
edits := capRecentEdits(mergeRecentEdits(preservedEditsToEntries(priorState.RecentEdits), recentEdits(middle)))
// Tools: preserve deferred tool_search schemas from the transcript. Fresh
// loads override older carried copies by name.
tools := mergeSkillEntries(preservedToolsToEntries(priorState.Tools), loadedToolSchemas(middle))
@@ -360,19 +493,76 @@ func appendPreservedState(summary string, middle []zeroruntime.Message) string {
projectInstructionEntries(middle),
)
if block := formatPreservedState(plan, tools, skills, instructions); block != "" {
if block := formatPreservedState(plan, edits, tools, skills, instructions); block != "" {
summary += "\n\n" + block
}
return summary
}
// mergeRecentEdits overlays fresh edits onto edits preserved by an earlier
// compaction. Unlike mergeSkillEntries (which keeps refreshed entries in their
// original slot), a path touched again by a fresh edit MOVES to the newest
// position and takes the fresh note, so capRecentEdits — which keeps only the
// most-recent tail — never drops a file the model just re-edited. Paths present
// only in the older set keep their relative order, ahead of the fresh ones.
func mergeRecentEdits(older, newer []skillEntry) []skillEntry {
if len(newer) == 0 {
return older
}
freshBody := make(map[string]string, len(newer))
freshOrder := make([]string, 0, len(newer))
for _, e := range newer {
if _, ok := freshBody[e.name]; !ok {
freshOrder = append(freshOrder, e.name)
}
freshBody[e.name] = e.body
}
merged := make([]skillEntry, 0, len(older)+len(freshOrder))
// Older entries that were NOT re-edited keep their position, oldest first.
for _, e := range older {
if _, refreshed := freshBody[e.name]; refreshed {
continue // re-appended at its newest position below
}
merged = append(merged, e)
}
// Fresh edits follow, in last-seen order, each carrying the fresh note — so
// the most recently touched files sit at the tail the cap preserves.
for _, name := range freshOrder {
merged = append(merged, skillEntry{name: name, body: freshBody[name]})
}
return merged
}
// capRecentEdits bounds the preserved edit list to the most recent maxRecentEdits
// entries after a merge, so repeated compactions can't grow it without limit.
func capRecentEdits(entries []skillEntry) []skillEntry {
if len(entries) > maxRecentEdits {
return entries[len(entries)-maxRecentEdits:]
}
return entries
}
func preservedEditsToEntries(edits []preservedEdit) []skillEntry {
entries := make([]skillEntry, 0, len(edits))
for _, e := range edits {
if e.Path == "" {
continue
}
entries = append(entries, skillEntry{name: e.Path, body: e.Note})
}
return entries
}
// formatPreservedState renders state as the labelled, single-line
// JSON block. Returns "" when there is nothing to preserve.
func formatPreservedState(plan string, tools, skills, instructions []skillEntry) string {
if plan == "" && len(tools) == 0 && len(skills) == 0 && len(instructions) == 0 {
func formatPreservedState(plan string, edits, tools, skills, instructions []skillEntry) string {
if plan == "" && len(edits) == 0 && len(tools) == 0 && len(skills) == 0 && len(instructions) == 0 {
return ""
}
state := preservedState{Plan: plan}
for _, e := range edits {
state.RecentEdits = append(state.RecentEdits, preservedEdit{Path: e.name, Note: e.body})
}
for _, t := range tools {
state.Tools = append(state.Tools, preservedTool{Name: t.name, Body: t.body})
}
@@ -0,0 +1,121 @@
package agent
import (
"fmt"
"strings"
"testing"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
// recentEdits extracts each mutated file's path and a one-line note from the
// matching tool result, latest note per path in last-seen order.
func TestRecentEditsExtractsPathsAndNotes(t *testing.T) {
messages := []zeroruntime.Message{
{Role: zeroruntime.MessageRoleAssistant, ToolCalls: []zeroruntime.ToolCall{
{ID: "e1", Name: "write_file", Arguments: `{"path":"internal/foo.go","content":"package foo"}`},
}},
{Role: zeroruntime.MessageRoleTool, ToolCallID: "e1", Content: "Wrote internal/foo.go (12 lines)"},
{Role: zeroruntime.MessageRoleAssistant, ToolCalls: []zeroruntime.ToolCall{
{ID: "e2", Name: "edit_file", Arguments: `{"path":"internal/bar.go","old_string":"a","new_string":"b"}`},
}},
{Role: zeroruntime.MessageRoleTool, ToolCallID: "e2", Content: "Applied edit to internal/bar.go"},
}
edits := recentEdits(messages)
if len(edits) != 2 {
t.Fatalf("expected 2 edited files, got %d: %#v", len(edits), edits)
}
if edits[0].name != "internal/foo.go" || !strings.Contains(edits[0].body, "12 lines") {
t.Fatalf("first edit = %#v, want foo.go with its note", edits[0])
}
if edits[1].name != "internal/bar.go" || !strings.Contains(edits[1].body, "Applied edit") {
t.Fatalf("second edit = %#v, want bar.go with its note", edits[1])
}
}
// After compaction elides the editing turns, the preserved-state block still
// names the edited files and what changed, so the model needn't re-read them.
func TestCompactionPreservesRecentEdits(t *testing.T) {
messages := []zeroruntime.Message{
{Role: zeroruntime.MessageRoleSystem, Content: "system"},
{Role: zeroruntime.MessageRoleUser, Content: "add a flag"},
{Role: zeroruntime.MessageRoleAssistant, Content: "editing", ToolCalls: []zeroruntime.ToolCall{
{ID: "e1", Name: "write_file", Arguments: `{"path":"cmd/main.go","content":"..."}`},
}},
{Role: zeroruntime.MessageRoleTool, ToolCallID: "e1", Content: "Wrote cmd/main.go (adds --version flag)"},
{Role: zeroruntime.MessageRoleAssistant, Content: "done"},
{Role: zeroruntime.MessageRoleUser, Content: "continue"},
{Role: zeroruntime.MessageRoleAssistant, Content: "continuing"},
}
summary := compactStateConversation(t, messages)
state := parsePreservedStateBlock(summary)
if len(state.RecentEdits) != 1 {
t.Fatalf("expected 1 preserved edit, got %#v", state.RecentEdits)
}
if state.RecentEdits[0].Path != "cmd/main.go" || !strings.Contains(state.RecentEdits[0].Note, "--version") {
t.Fatalf("preserved edit = %#v, want cmd/main.go + its note", state.RecentEdits[0])
}
}
// A fresh edit note for a path overrides the one carried from an earlier
// compaction (newer wins) and moves the path to the newest position rather than
// duplicating it or leaving it pinned to its old slot.
func TestRecentEditsMergeNewerWins(t *testing.T) {
prior := preservedState{RecentEdits: []preservedEdit{{Path: "a.go", Note: "old note"}, {Path: "z.go", Note: "z"}}}
older := preservedEditsToEntries(prior.RecentEdits)
newer := []skillEntry{{name: "a.go", body: "new note"}, {name: "b.go", body: "added"}}
merged := mergeRecentEdits(older, newer)
// z.go (untouched) keeps its lead; a.go is re-edited so it takes the new note
// and moves behind z.go, ahead of the brand-new b.go: [z.go, a.go, b.go].
if len(merged) != 3 {
t.Fatalf("expected z.go + refreshed a.go + b.go, got %#v", merged)
}
if merged[0].name != "z.go" {
t.Fatalf("untouched z.go should stay first, got %#v", merged)
}
if merged[1].name != "a.go" || merged[1].body != "new note" {
t.Fatalf("re-edited a.go should move to the newest edits with the new note, got %#v", merged[1])
}
if merged[2].name != "b.go" {
t.Fatalf("brand-new b.go should be last, got %#v", merged[2])
}
}
// Regression: once more than maxRecentEdits distinct files are tracked, re-editing
// an early file must move it into the capped tail rather than leaving it to be
// dropped — otherwise the next compaction omits exactly the file the model most
// recently touched.
func TestRecentEditsCapKeepsReeditedFile(t *testing.T) {
// Earlier compaction preserved exactly maxRecentEdits files, f0 the oldest.
older := make([]skillEntry, 0, maxRecentEdits)
for i := 0; i < maxRecentEdits; i++ {
older = append(older, skillEntry{name: fmt.Sprintf("f%d.go", i), body: "old"})
}
// Fresh window re-edits the OLDEST file (f0) and adds a brand-new file, so the
// merged list exceeds the cap and something must be dropped.
newer := []skillEntry{{name: "f0.go", body: "re-edited"}, {name: "fresh.go", body: "new"}}
merged := capRecentEdits(mergeRecentEdits(older, newer))
if len(merged) != maxRecentEdits {
t.Fatalf("cap should keep %d entries, got %d", maxRecentEdits, len(merged))
}
var f0 *skillEntry
for i := range merged {
if merged[i].name == "f0.go" {
f0 = &merged[i]
}
if merged[i].name == "f1.go" {
t.Fatalf("f1.go (now the least-recently touched) should have been dropped, got %#v", merged)
}
}
if f0 == nil {
t.Fatalf("re-edited f0.go must survive the cap, got %#v", merged)
}
if f0.body != "re-edited" {
t.Fatalf("surviving f0.go should carry the fresh note, got %q", f0.body)
}
}
+22 -4
View File
@@ -3,6 +3,8 @@ package agent
import (
"context"
"os"
"path/filepath"
"strings"
"time"
"github.com/Gitlawb/zero/internal/lsp"
@@ -18,9 +20,12 @@ const fileDiagnosticsTimeout = 10 * time.Second
// 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. Returns nil when
// manager is nil, disabling inline diagnostics entirely.
func NewFileDiagnostics(manager *lsp.Manager) func(context.Context, string) string {
// 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
}
@@ -39,6 +44,19 @@ func NewFileDiagnostics(manager *lsp.Manager) func(context.Context, string) stri
if len(errors) == 0 {
return ""
}
return lsp.FormatDiagnostics(absPath, errors)
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)
}
+24
View File
@@ -0,0 +1,24 @@
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)
}
}
}
+30 -8
View File
@@ -22,6 +22,12 @@ const (
// calls have executed since the last update_plan call.
staleToolCallThreshold = 10
// stalePlanTurnThreshold injects the same one-shot reminder once this many
// turns have passed since the last update_plan while plan items are still
// pending — the turn-based complement to staleToolCallThreshold, catching a
// plan that drifts stale across many turns that each make few tool calls.
stalePlanTurnThreshold = 8
// toolOnlyProgressReminderAt injects a one-shot progress nudge after this
// many consecutive turns contain tool calls but no visible assistant text.
// It does not stop the run; it tells the model to synthesize what it already
@@ -43,8 +49,12 @@ const (
toolFailureHintAt = 2
// toolFailureStopAt halts the run after a tool fails this many times in a row
// with the same error, so NO model (weak or strong) burns turns looping on a
// bad call.
toolFailureStopAt = 4
// bad call. Set to 6 (not 4): a corrective hint fires at toolFailureHintAt (2),
// and a model iterating on a genuinely tricky edit can legitimately fail a few
// times after the hint while converging — stopping at 4 cut those runs short.
// The streak still resets the moment the tool succeeds or hits a different
// error, so this only affects true same-error loops.
toolFailureStopAt = 6
// maxContinueNudges bounds how many times the headless completion gate
// (Options.RequireCompletionSignal) re-prompts a model that stopped without a
@@ -354,8 +364,13 @@ type guardState struct {
emptyTurns int
totalToolCalls int
toolCallsSincePlanUpdate int
planEverCalled bool
notCalledReminderSent bool
// turnsSincePlanUpdate counts turns (not individual tool calls) since the last
// update_plan, so a plan that goes stale across many low-tool-call turns is
// still caught — the tool-call counter alone can take many turns to trip when
// the model makes only one call per turn.
turnsSincePlanUpdate int
planEverCalled bool
notCalledReminderSent bool
// staleReminderSent records whether the stale reminder has already fired for
// the current stale interval. It is cleared when a plan update opens a new
// interval, making the reminder one-shot per interval rather than per turn.
@@ -427,11 +442,16 @@ func (state *guardState) observeTurn(collected zeroruntime.CollectedStream) (sto
state.toolOnlyReminderSent = false
}
// One turn has passed; the plan-update below resets this to 0 when the model
// refreshes the plan this turn.
state.turnsSincePlanUpdate++
for _, call := range collected.ToolCalls {
state.totalToolCalls++
if call.Name == planToolName {
state.planEverCalled = true
state.toolCallsSincePlanUpdate = 0
state.turnsSincePlanUpdate = 0
// A fresh plan update opens a new stale interval.
state.staleReminderSent = false
// Record how many items remain so the completion gate knows whether
@@ -485,10 +505,12 @@ func (state *guardState) progressReminder() string {
// number of turns completed so far).
func (state *guardState) planReminder(turn int) string {
// STALE reminder takes priority: a long run without a plan update is the
// stronger signal. One-shot per stale interval.
if state.planEverCalled &&
!state.staleReminderSent &&
state.toolCallsSincePlanUpdate >= staleToolCallThreshold {
// stronger signal. Fires on either the tool-call streak OR a turn streak with
// pending items (so a plan drifting stale across many low-call turns is caught,
// while a fully-completed plan is left alone). One-shot per stale interval.
if state.planEverCalled && !state.staleReminderSent &&
(state.toolCallsSincePlanUpdate >= staleToolCallThreshold ||
(state.turnsSincePlanUpdate >= stalePlanTurnThreshold && state.planItemsPending > 0)) {
state.staleReminderSent = true
return planStaleReminder(state.toolCallsSincePlanUpdate)
}
+54 -5
View File
@@ -25,7 +25,12 @@ const maxTurnsFinalAnswerPrompt = "You have reached the tool-turn limit. Do not
// WITH NO OUTPUT yet is re-issued on a fresh connection before giving up. Only
// the no-output case is retried (a partial turn would duplicate), so this is a
// safe recovery for a stalled/dead pooled connection.
const maxStreamStallRetries = 2
//
// Set to 1 (not 2): each attempt can itself idle for the full stream timeout
// (~5min) before the stall is even detected, so 2 retries left an interactive
// session frozen for ~15min. One retry keeps the common single-hiccup recovery
// while bounding the worst case to ~2× the idle timeout.
const maxStreamStallRetries = 1
const (
toolResultMetaControl = "control"
@@ -145,6 +150,11 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
// has already been demanded this run, so it fires at most once.
acceptanceRequested := false
// toolDefCache memoizes each tool's rendered JSON-schema definition across
// turns (a tool's advertised schema is stable for the run), so partitionTools
// doesn't re-run the recursive schema→map conversion for every tool every turn.
toolDefCache := map[string]zeroruntime.ToolDefinition{}
result := Result{Messages: copyMessages(messages)}
for turn := 0; turn < maxTurns; turn++ {
result.Turns = turn + 1
@@ -153,7 +163,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
// the tool-definition tokens (they ride on every request) in its estimate.
// partitionTools depends only on registry/permissions/options/loaded, not on
// the messages, so computing it before compaction is safe.
exposed, _ := partitionTools(registry, permissionMode, options, loaded)
exposed, _ := partitionToolsCached(registry, permissionMode, options, loaded, toolDefCache)
// PROACTIVE compaction: if the history is approaching the model's
// context window, summarize the oldest middle before building the
@@ -340,6 +350,16 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
}
}
// Calibrate the compaction token estimator against the provider's real
// prompt-token count for the request we just sent, so later turns trigger
// compaction near true capacity instead of ~15% early on code-heavy history.
// Recompute the estimate HERE (not at request-build time): a reactive
// compaction may have replaced `messages` with a smaller set and re-sent, so
// this reflects the request that actually produced collected.Usage. The
// assistant reply is appended below, after this, so `messages` is still the
// sent request.
compactor.calibrate(estimateTokens(messages)+estimateToolDefTokens(exposed), collected.Usage.InputTokens)
// Carry the turn's terminal stop reason so a final answer cut off at the
// output token cap (or by a content filter) is reported as truncated. A
// tool-call turn normalizes to "" and clears any prior reason.
@@ -2488,6 +2508,17 @@ func permissionActionFromSandbox(action sandbox.Action) PermissionAction {
// exposed. The exposed slice is alpha-sorted by name, matching the legacy order
// so the inactive path is stable.
func partitionTools(registry *tools.Registry, permissionMode PermissionMode, options Options, loaded map[string]bool) ([]zeroruntime.ToolDefinition, string) {
return partitionToolsCached(registry, permissionMode, options, loaded, nil)
}
// partitionToolsCached is partitionTools with an optional per-tool definition
// cache. The partitioning itself (visibility, deferral, ordering) is recomputed
// every call — it must be, because a tool's deferred state can flip mid-run (e.g.
// swarm tools un-defer once a swarm is active). Only the expensive part —
// rendering each tool's JSON-schema parameters — is memoized by tool name, since a
// tool's advertised name/description/schema is stable for the run. defCache nil
// disables caching (used by tests and the plain partitionTools entrypoint).
func partitionToolsCached(registry *tools.Registry, permissionMode PermissionMode, options Options, loaded map[string]bool, defCache map[string]zeroruntime.ToolDefinition) ([]zeroruntime.ToolDefinition, string) {
registeredTools := registry.All()
visible := make([]tools.Tool, 0, len(registeredTools))
@@ -2531,7 +2562,7 @@ func partitionTools(registry *tools.Registry, permissionMode PermissionMode, opt
if tool.Name() == tools.ToolSearchToolName {
continue
}
definitions = append(definitions, runtimeToolDefinition(tool))
definitions = append(definitions, cachedRuntimeToolDefinition(defCache, tool))
}
sort.Slice(definitions, func(left int, right int) bool {
return definitions[left].Name < definitions[right].Name
@@ -2563,13 +2594,13 @@ func partitionTools(registry *tools.Registry, permissionMode PermissionMode, opt
}
if tools.IsDeferred(tool) {
if loaded[name] {
loadedTail = append(loadedTail, runtimeToolDefinition(tool))
loadedTail = append(loadedTail, cachedRuntimeToolDefinition(defCache, tool))
} else {
hiddenTools = append(hiddenTools, tool)
}
continue
}
eager = append(eager, runtimeToolDefinition(tool))
eager = append(eager, cachedRuntimeToolDefinition(defCache, tool))
}
sort.Slice(eager, func(left int, right int) bool {
return eager[left].Name < eager[right].Name
@@ -2603,6 +2634,24 @@ func partitionTools(registry *tools.Registry, permissionMode PermissionMode, opt
return definitions, discovery
}
// cachedRuntimeToolDefinition returns the tool's rendered definition, reusing a
// cached render when defCache holds one for this tool name. A tool's advertised
// definition is stable across a run, so caching skips the recursive schema→map
// conversion (schemaToRuntimeMap) that would otherwise run for every tool on every
// turn. tool_search is excluded by its callers (its description is dynamic), so it
// never poisons the cache. A nil cache computes fresh.
func cachedRuntimeToolDefinition(defCache map[string]zeroruntime.ToolDefinition, tool tools.Tool) zeroruntime.ToolDefinition {
if defCache == nil {
return runtimeToolDefinition(tool)
}
if def, ok := defCache[tool.Name()]; ok {
return def
}
def := runtimeToolDefinition(tool)
defCache[tool.Name()] = def
return def
}
// runtimeToolDefinition renders a tool's advertised definition (name, description,
// JSON-schema parameters) as sent to the provider.
func runtimeToolDefinition(tool tools.Tool) zeroruntime.ToolDefinition {
+19 -7
View File
@@ -51,21 +51,33 @@ func parallelSafeToolCall(registry *tools.Registry, call ToolCall, options Optio
}
// executeParallelReadBatch runs calls[start:end] concurrently (bounded by
// maxParallelReadTools) and returns results indexed relative to start. Any
// permission prompt that still occurs (a sandbox preflight can demand one even
// for an auto-allowed read) is serialized so an interactive front-end never
// sees two prompts at once.
// 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 {
var promptMutex sync.Mutex
inner := options.OnPermissionRequest
batchOptions.OnPermissionRequest = func(ctx context.Context, request PermissionRequest) (PermissionDecision, error) {
promptMutex.Lock()
defer promptMutex.Unlock()
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)
+42 -22
View File
@@ -17,6 +17,9 @@ 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
@@ -24,6 +27,24 @@ type probeTool struct {
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 {
@@ -45,7 +66,13 @@ func (tool *probeTool) Run(_ context.Context, args map[string]any) tools.Result
}
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)
@@ -127,10 +154,9 @@ func TestRunExecutesConsecutiveReadsConcurrently(t *testing.T) {
}
func TestRunParallelReadsNeverSpanMutatingCall(t *testing.T) {
read := &probeTool{name: "probe_read", sideEffect: tools.SideEffectRead, delay: 30 * time.Millisecond}
write := &probeTool{name: "probe_write", sideEffect: tools.SideEffectWrite}
// Shared log across both probes so relative ordering is observable.
write.mu = sync.Mutex{}
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)
@@ -151,32 +177,26 @@ func TestRunParallelReadsNeverSpanMutatingCall(t *testing.T) {
if err != nil {
t.Fatal(err)
}
// The write must start only after r1+r2 finished, and r3/r4 must start only
// after the write finished: batches never cross a mutating call.
index := func(log []string, entry string) int {
// 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
}
readLog := func() []string { read.mu.Lock(); defer read.mu.Unlock(); return append([]string(nil), read.log...) }()
writeLog := func() []string { write.mu.Lock(); defer write.mu.Unlock(); return append([]string(nil), write.log...) }()
if index(writeLog, "start:w") == -1 {
t.Fatalf("write probe never ran: %v", writeLog)
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)
}
for _, entry := range []string{"end:r1", "end:r2"} {
if index(readLog, entry) == -1 {
t.Fatalf("first read batch incomplete: %v", readLog)
}
}
// r3/r4 must appear strictly after r1/r2 in the read log (the write barrier
// between the two batches forces full separation).
firstBatchMaxEnd := max(index(readLog, "end:r1"), index(readLog, "end:r2"))
secondBatchMinStart := min(index(readLog, "start:r3"), index(readLog, "start:r4"))
if secondBatchMinStart < firstBatchMaxEnd {
t.Fatalf("second read batch started before first batch ended: %v", readLog)
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)
+63
View File
@@ -0,0 +1,63 @@
package agent
import (
"context"
"reflect"
"testing"
"github.com/Gitlawb/zero/internal/tools"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
// countingSchemaTool records how many times its Parameters() schema is read, so a
// test can prove the definition cache avoids re-rendering across turns.
type countingSchemaTool struct {
name string
calls *int
}
func (t countingSchemaTool) Name() string { return t.name }
func (t countingSchemaTool) Description() string { return "counts schema reads" }
func (t countingSchemaTool) Parameters() tools.Schema {
*t.calls++
return tools.Schema{Type: "object", AdditionalProperties: false, Properties: map[string]tools.PropertySchema{
"x": {Type: "string"},
}}
}
func (t countingSchemaTool) Safety() tools.Safety {
return tools.Safety{SideEffect: tools.SideEffectRead, Permission: tools.PermissionAllow}
}
func (t countingSchemaTool) Run(_ context.Context, _ map[string]any) tools.Result {
return tools.Result{Status: tools.StatusOK}
}
// The cache renders a tool's schema once and reuses it across calls, and its
// output is identical to the uncached path.
func TestPartitionToolsCacheRendersOnceAndMatchesUncached(t *testing.T) {
calls := 0
registry := tools.NewRegistry()
registry.Register(countingSchemaTool{name: "alpha", calls: &calls})
registry.Register(countingSchemaTool{name: "beta", calls: &calls})
uncached, _ := partitionTools(registry, PermissionModeAuto, Options{}, map[string]bool{})
base := calls
cache := map[string]zeroruntime.ToolDefinition{}
first, _ := partitionToolsCached(registry, PermissionModeAuto, Options{}, map[string]bool{}, cache)
rendersAfterFirst := calls - base
second, _ := partitionToolsCached(registry, PermissionModeAuto, Options{}, map[string]bool{}, cache)
rendersAfterSecond := calls - base
if rendersAfterFirst != 2 {
t.Fatalf("first cached call should render 2 tools once each, rendered %d times", rendersAfterFirst)
}
if rendersAfterSecond != 2 {
t.Fatalf("second cached call must reuse the cache (no new schema reads), total renders = %d", rendersAfterSecond)
}
if !reflect.DeepEqual(first, second) {
t.Fatal("cached calls must return identical definitions")
}
if !reflect.DeepEqual(first, uncached) {
t.Fatal("cached output must match the uncached partitionTools output")
}
}
+64
View File
@@ -0,0 +1,64 @@
package agent
import (
"strings"
"testing"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
func planUpdateTurn(status string) zeroruntime.CollectedStream {
return zeroruntime.CollectedStream{ToolCalls: []zeroruntime.ToolCall{
{Name: "update_plan", Arguments: `{"plan":[{"content":"step one","status":"` + status + `"}]}`},
}}
}
// Turn-based staleness: a plan that goes many turns without an update — while
// items are still pending — draws the stale reminder even though few tool calls
// have run (the tool-call streak alone would not have tripped).
func TestPlanReminderFiresOnTurnStaleness(t *testing.T) {
state := newGuardState()
state.observeTurn(planUpdateTurn("in_progress"))
// Text-only turns advance the turn counter without adding tool calls.
for range stalePlanTurnThreshold {
state.observeTurn(zeroruntime.CollectedStream{Text: "still working…"})
}
if state.toolCallsSincePlanUpdate >= staleToolCallThreshold {
t.Fatalf("precondition: the tool-call trigger must not fire (%d calls)", state.toolCallsSincePlanUpdate)
}
if got := state.planReminder(stalePlanTurnThreshold + 2); !strings.Contains(got, planStaleReminderMarker) {
t.Fatalf("expected a turn-based stale reminder, got %q", got)
}
}
// A fully-completed plan is NOT stale: no pending items means the turn-based
// trigger stays quiet even after many turns, so a finished plan isn't nagged.
func TestPlanReminderSkipsTurnStalenessWhenComplete(t *testing.T) {
state := newGuardState()
state.observeTurn(planUpdateTurn("completed"))
for range stalePlanTurnThreshold + 3 {
state.observeTurn(zeroruntime.CollectedStream{Text: "wrapping up…"})
}
if got := state.planReminder(stalePlanTurnThreshold + 5); got != "" {
t.Fatalf("a completed plan must not draw a stale reminder, got %q", got)
}
}
// A plan update resets the turn counter, so the stale reminder does not fire
// right after a fresh update.
func TestPlanUpdateResetsTurnStaleness(t *testing.T) {
state := newGuardState()
state.observeTurn(planUpdateTurn("in_progress"))
for range stalePlanTurnThreshold - 1 {
state.observeTurn(zeroruntime.CollectedStream{Text: "working…"})
}
// Refresh the plan — resets the turn streak.
state.observeTurn(planUpdateTurn("in_progress"))
if got := state.planReminder(stalePlanTurnThreshold + 1); got != "" {
t.Fatalf("a just-refreshed plan must not be stale, got %q", got)
}
}
+45 -5
View File
@@ -3,9 +3,11 @@ package agent
import (
"context"
"fmt"
"math/rand"
"strings"
"time"
"github.com/Gitlawb/zero/internal/errhint"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
@@ -17,10 +19,20 @@ import (
// times. We retry ONLY the connect (not a partially-consumed stream), so no
// already-forwarded OnText is ever duplicated.
const (
maxStreamReconnects = 2
streamReconnectBase = 500 * time.Millisecond
// maxStreamReconnects is how many times the connect is re-issued after a
// transient disconnect. 4 (not 2): with jittered exponential backoff this rides
// out a multi-second network blip (~0.5s + 1s + 2s + 4s worst case) instead of
// dying on a 2s hiccup and re-burning every token on a restart.
maxStreamReconnects = 4
// streamReconnectMax caps a single backoff so the tail attempts don't wait
// minutes on a long outage.
streamReconnectMax = 8 * time.Second
)
// streamReconnectBase is the first-retry delay (doubled each subsequent attempt).
// A var, not a const, so tests can shrink it to keep the exhaustion path fast.
var streamReconnectBase = 500 * time.Millisecond
// reconnectNotifier is called before each retry with the 1-based attempt number
// and the max, so the caller can surface a "Reconnecting N/max" notice. Nil is
// fine.
@@ -73,7 +85,7 @@ func streamWithReconnect(ctx context.Context, provider Provider, request zerorun
if notify != nil {
notify(attempt, maxStreamReconnects)
}
if waitErr := sleepWithContext(ctx, backoffFor(attempt)); waitErr != nil {
if waitErr := sleepWithContext(ctx, jitteredBackoff(attempt)); waitErr != nil {
return nil, err // ctx cancelled while waiting; surface the original error
}
stream, err = provider.StreamCompletion(ctx, request)
@@ -96,6 +108,19 @@ func shouldReconnect(ctx context.Context, err error) bool {
if isContextLimitError(msg) || isImageRejectionError(err) {
return false
}
// HTTP 5xx statuses are non-reconnectable and must be excluded BEFORE the
// transport substring match below — otherwise "504 Gateway Timeout" would slip
// through on the generic "timeout" needle. 503 already exhausted
// providerio.SendWithRetry (retrying is a redundant double-retry); 500/502/504
// are non-idempotent by providerio's rule — the completion POST may already
// have reached the model before the upstream/gateway gave up, so replaying the
// connect risks duplicate billable work. Digit-boundary matched so an
// incidental "504" inside a latency/id number is not mistaken for a status.
if errhint.HasStatusCode(msg, "500", "502", "503", "504") {
return false
}
// Transport-level disconnects only. A genuine transport failure (EOF, reset,
// refused, timeout) means no response was received, which is safe to reconnect.
for _, needle := range []string{
"eof",
"connection reset",
@@ -106,8 +131,6 @@ func shouldReconnect(ctx context.Context, err error) bool {
"timed out",
"temporarily unavailable",
"i/o timeout",
"503",
"502",
"server closed",
"unexpected end",
} {
@@ -118,14 +141,31 @@ func shouldReconnect(ctx context.Context, err error) bool {
return false
}
// backoffFor is the deterministic exponential base delay for a 1-based attempt,
// capped at streamReconnectMax. Jitter is layered on separately (jitteredBackoff).
func backoffFor(attempt int) time.Duration {
d := streamReconnectBase
for i := 1; i < attempt; i++ {
if d >= streamReconnectMax {
return streamReconnectMax
}
d *= 2
}
if d > streamReconnectMax {
d = streamReconnectMax
}
return d
}
// jitteredBackoff adds up to 50% random jitter on top of backoffFor so concurrent
// runs (swarm members, a cron fleet) that all trip on the same outage don't
// reconnect in lockstep and hammer a recovering endpoint. Never shorter than the
// deterministic base, so backoff still grows attempt over attempt.
func jitteredBackoff(attempt int) time.Duration {
base := backoffFor(attempt)
return base + time.Duration(rand.Int63n(int64(base/2)+1))
}
func sleepWithContext(ctx context.Context, d time.Duration) error {
timer := time.NewTimer(d)
defer timer.Stop()
+30 -1
View File
@@ -44,6 +44,10 @@ func TestStreamWithReconnectRecoversFromTransientDisconnect(t *testing.T) {
}
func TestStreamWithReconnectGivesUpAfterMax(t *testing.T) {
// Shrink the backoff so exhausting all retries stays fast (real base would
// sleep ~7.5s across 4 attempts).
defer func(orig time.Duration) { streamReconnectBase = orig }(streamReconnectBase)
streamReconnectBase = time.Millisecond
// Always fails with a disconnect error → exhausts retries and returns it.
p := &flakyProvider{failBefore: 99, failErr: errors.New("connection reset by peer")}
_, err := streamWithReconnect(context.Background(), p, zeroruntime.CompletionRequest{}, nil)
@@ -87,7 +91,7 @@ func TestShouldReconnectClassification(t *testing.T) {
ctx := context.Background()
disconnects := []string{
"unexpected EOF", "connection reset by peer", "broken pipe",
"i/o timeout", "503 Service Unavailable", "server closed the connection",
"i/o timeout", "server closed the connection", "connection refused",
}
for _, m := range disconnects {
if !shouldReconnect(ctx, errors.New(m)) {
@@ -97,6 +101,13 @@ func TestShouldReconnectClassification(t *testing.T) {
notDisconnects := []string{
"context length exceeded", "invalid api key", "model not found",
"400 bad request: unsupported parameter",
// HTTP 5xx statuses are handled by providerio.SendWithRetry (503) or are
// non-idempotent (500/502/504); the reconnect path must not double-retry
// them. "504 Gateway Timeout" must NOT slip through on the generic "timeout"
// transport signal — a gateway timeout is non-idempotent, since the upstream
// may have already processed the completion POST before giving up.
"503 Service Unavailable", "502 Bad Gateway",
"504 Gateway Timeout", "500 Internal Server Error",
}
for _, m := range notDisconnects {
if shouldReconnect(ctx, errors.New(m)) {
@@ -112,6 +123,24 @@ func TestBackoffGrows(t *testing.T) {
if backoffFor(2) != 2*streamReconnectBase {
t.Fatalf("attempt 2 backoff = %v, want %v", backoffFor(2), 2*streamReconnectBase)
}
// The exponential base is capped so late attempts don't wait minutes.
if got := backoffFor(20); got != streamReconnectMax {
t.Fatalf("attempt 20 backoff = %v, want cap %v", got, streamReconnectMax)
}
}
func TestJitteredBackoffStaysInBounds(t *testing.T) {
// Jitter never drops below the deterministic base and never exceeds base*1.5,
// so backoff still grows attempt over attempt while decorrelating retries.
for attempt := 1; attempt <= 5; attempt++ {
base := backoffFor(attempt)
for i := 0; i < 200; i++ {
got := jitteredBackoff(attempt)
if got < base || got > base+base/2 {
t.Fatalf("attempt %d jittered backoff %v out of [%v, %v]", attempt, got, base, base+base/2)
}
}
}
}
func TestReconnectNoticeRoutesThroughReasoning(t *testing.T) {
+10 -1
View File
@@ -12,6 +12,7 @@ import (
"github.com/Gitlawb/zero/internal/agent"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/errhint"
"github.com/Gitlawb/zero/internal/imageinput"
"github.com/Gitlawb/zero/internal/lsp"
"github.com/Gitlawb/zero/internal/modelregistry"
@@ -718,7 +719,7 @@ func newExecSelfCorrector(enabled bool, workspaceRoot string, autonomy string) (
defer cancel()
_ = manager.Shutdown(ctx)
}
fileDiagnostics := agent.NewFileDiagnostics(manager)
fileDiagnostics := agent.NewFileDiagnostics(manager, workspaceRoot)
if !enabled {
return nil, fileDiagnostics, cleanup
}
@@ -961,6 +962,14 @@ func writeExecProviderError(stdout io.Writer, stderr io.Writer, format execOutpu
if _, err := fmt.Fprintf(stderr, "[zero] %s\n", message); err != nil {
return exitCrash
}
// Append a one-line next step for recognized provider failures (auth /
// rate-limit / connectivity / …). The classifier gates on a provider-origin
// marker, so non-provider codes (sandbox_error, mcp_error) never draw a hint.
if hint := errhint.CLIHint(errors.New(message)); hint != "" {
if _, err := fmt.Fprintf(stderr, "[zero] %s\n", hint); err != nil {
return exitCrash
}
}
return exitProvider
}
+50
View File
@@ -0,0 +1,50 @@
package cli
import (
"bytes"
"strings"
"testing"
)
// A recognized provider error in text mode prints the raw message plus a one-line
// actionable hint; a non-provider error (or JSON mode) prints no hint.
func TestWriteExecProviderErrorAppendsHint(t *testing.T) {
t.Run("provider auth error gets a hint", func(t *testing.T) {
var stdout, stderr bytes.Buffer
code := writeExecProviderError(&stdout, &stderr, execOutputText, "provider_error",
"auth error: your API key is missing or invalid")
if code != exitProvider {
t.Fatalf("exit code = %d, want exitProvider", code)
}
out := stderr.String()
if !strings.Contains(out, "auth error:") {
t.Fatalf("expected raw message, got %q", out)
}
if !strings.Contains(out, "zero auth") {
t.Fatalf("expected an actionable hint referencing `zero auth`, got %q", out)
}
})
t.Run("non-provider error gets no hint", func(t *testing.T) {
var stdout, stderr bytes.Buffer
writeExecProviderError(&stdout, &stderr, execOutputText, "sandbox_error",
"sandbox setup failed: permission denied")
out := stderr.String()
// Exactly one "[zero]" line — no spurious hint attached to a local error.
if n := strings.Count(out, "[zero]"); n != 1 {
t.Fatalf("expected exactly one [zero] line for a non-provider error, got %d:\n%s", n, out)
}
})
t.Run("json mode never appends a hint line", func(t *testing.T) {
var stdout, stderr bytes.Buffer
writeExecProviderError(&stdout, &stderr, execOutputJSON, "provider_error",
"auth error: bad key")
if stderr.Len() != 0 {
t.Fatalf("json mode must not write to stderr, got %q", stderr.String())
}
if strings.Contains(stdout.String(), "zero auth") {
t.Fatalf("json mode must not inject a hint into the structured payload, got %q", stdout.String())
}
})
}
+167
View File
@@ -0,0 +1,167 @@
// Package errhint classifies provider/model failures into a few user-actionable
// categories and turns them into a one-line "next step" hint.
//
// Provider errors already arrive with a classified string prefix from
// providerio.ClassifiedError ("auth error:", "rate limit error:", …); lower-level
// failures (DNS, TLS, timeouts, context-length) arrive as raw driver or library
// messages. Classify matches both so the interactive error row (TUI) and the
// `zero exec` provider-error path (CLI) can append one concrete next step instead
// of dumping an identical red blob for every failure mode.
package errhint
import "strings"
// Category buckets a provider/model failure into a small set of classes that each
// map to a distinct recovery action.
type Category int
const (
// Unknown means the error didn't match any known signature; callers should
// emit no hint rather than guess.
Unknown Category = iota
Auth
RateLimit
Connectivity
ModelNotFound
ContextOverflow
)
// providerMarkers are the prefixes the provider layer attaches to every
// provider-originated failure: providerio.ClassifiedError emits "auth error:",
// "rate limit error:", "provider error:", and "provider request error:", and the
// streaming paths wrap transport/decoding failures as "provider stream error:".
// A UI surface's error can also be a *local* failure (a tool's "permission
// denied", a "file does not exist", a config error), so Classify only proceeds
// past this gate for messages that are recognizably from the provider — otherwise
// a broad substring like "does not exist" would attach a bogus /model hint to an
// unrelated local error.
var providerMarkers = []string{
"auth error:",
"rate limit error:",
"provider error:",
"provider request error:",
"provider stream error:",
}
// Classify buckets err by scanning its message for known signatures. It is a
// deliberately conservative string heuristic — the provider layer has already
// discarded the numeric HTTP status by the time the error reaches a UI surface,
// so the message is all we have. It first gates on a provider-origin marker (see
// providerMarkers) so local failures never draw a provider hint, then sub-classifies.
// Order matters: more specific signatures are tested before broader ones (e.g.
// "context length" as overflow before the generic "timeout" as connectivity).
func Classify(err error) Category {
if err == nil {
return Unknown
}
m := strings.ToLower(err.Error())
if !containsAny(m, providerMarkers...) {
return Unknown
}
switch {
case containsAny(m, "auth error:", "unauthorized", "api key", "api_key", "invalid_api_key",
"authentication", "permission denied", "forbidden") || containsStatusCode(m, "401", "403"):
return Auth
case containsAny(m, "rate limit", "rate_limit", "too many requests", "quota",
"resource_exhausted", "overloaded") || containsStatusCode(m, "429", "529"):
return RateLimit
case containsAny(m, "context length", "context window", "maximum context", "context_length_exceeded",
"too many tokens", "prompt is too long", "reduce the length", "maximum context length"):
return ContextOverflow
case containsAny(m, "model not found", "model_not_found", "does not exist", "unknown model",
"no such model", "unsupported model", "invalid model", "model is not"):
return ModelNotFound
case containsAny(m, "dial tcp", "no such host", "connection refused", "network is unreachable",
"i/o timeout", "context deadline exceeded", "tls handshake", "connection reset",
"unexpected eof", "lookup ", "timeout"):
return Connectivity
default:
return Unknown
}
}
// TUIHint returns a one-line hint referencing interactive slash commands, or ""
// when the category is Unknown. Meant to sit under the raw error in the live
// error row.
func TUIHint(err error) string {
switch Classify(err) {
case Auth:
return "API key rejected — run /provider to re-check your credentials"
case RateLimit:
return "Rate limited — wait a moment, or switch model with /model"
case Connectivity:
return "Can't reach the provider — run /doctor --connectivity"
case ModelNotFound:
return "Model unavailable — pick another with /model"
case ContextOverflow:
return "Context window full — run /compact to free space"
default:
return ""
}
}
// CLIHint returns a one-line hint referencing `zero …` subcommands, or "" when the
// category is Unknown. Meant for the non-interactive `zero exec` error path, where
// slash commands don't apply.
func CLIHint(err error) string {
switch Classify(err) {
case Auth:
return "API key rejected — run `zero auth` or set the provider's API key, then retry"
case RateLimit:
return "Rate limited — wait a moment, or switch model with --model"
case Connectivity:
return "Can't reach the provider — run `zero doctor`"
case ModelNotFound:
return "Model unavailable — run `zero doctor` or pick another with --model"
case ContextOverflow:
return "Context window full — shorten the prompt or start a fresh session"
default:
return ""
}
}
func containsAny(haystack string, needles ...string) bool {
for _, n := range needles {
if strings.Contains(haystack, n) {
return true
}
}
return false
}
// containsStatusCode reports whether haystack contains any of the given HTTP
// status codes as a standalone number — not embedded in a longer digit run like
// "completed in 4290ms" or "request id 14015" — so an incidental number can't be
// mis-bucketed as an auth/rate-limit failure.
func containsStatusCode(haystack string, codes ...string) bool {
return HasStatusCode(haystack, codes...)
}
// HasStatusCode reports whether haystack contains any of the given HTTP status
// codes as a standalone number — not embedded in a longer digit run like
// "completed in 4290ms" or "request id 14015". Exported so other packages (e.g.
// the reconnect classifier) can gate on a status code without re-implementing
// the digit-boundary check.
func HasStatusCode(haystack string, codes ...string) bool {
for _, code := range codes {
for from := 0; ; {
rel := strings.Index(haystack[from:], code)
if rel < 0 {
break
}
pos := from + rel
beforeOK := pos == 0 || !isASCIIDigit(haystack[pos-1])
end := pos + len(code)
afterOK := end >= len(haystack) || !isASCIIDigit(haystack[end])
if beforeOK && afterOK {
return true
}
from = pos + 1
}
}
return false
}
func isASCIIDigit(b byte) bool {
return b >= '0' && b <= '9'
}
+94
View File
@@ -0,0 +1,94 @@
package errhint
import (
"errors"
"strings"
"testing"
)
func TestClassify(t *testing.T) {
cases := []struct {
name string
msg string
want Category
}{
{"nil-ish empty", "", Unknown},
{"providerio auth prefix", "auth error: your API key is missing or invalid — run `zero auth`", Auth},
{"raw 401", "provider request error: 401 Unauthorized", Auth},
{"invalid api key", "provider request error: invalid_api_key: incorrect key provided", Auth},
{"rate limit prefix", "rate limit error: 429 too many requests", RateLimit},
{"overloaded", "provider error: model is overloaded, please retry", RateLimit},
{"resource exhausted gemini", "provider stream error: rpc error: code = ResourceExhausted desc = quota exceeded", RateLimit},
{"context length openai", "provider request error: this model's maximum context length is 128000 tokens", ContextOverflow},
{"context_length_exceeded", "provider request error: context_length_exceeded", ContextOverflow},
{"model not found", "provider request error: 404 the model `gpt-9` does not exist", ModelNotFound},
{"unknown model", "provider request error: unknown model: sonnet-99", ModelNotFound},
{"dns failure", "provider stream error: Post \"https://api...\": dial tcp: lookup api.foo.com: no such host", Connectivity},
{"connection refused", "provider stream error: dial tcp 127.0.0.1:443: connection refused", Connectivity},
{"deadline", "provider stream error: context deadline exceeded (Client.Timeout exceeded)", Connectivity},
{"provider-marked but no sub-signature", "provider error: something totally unexpected happened", Unknown},
// An incidental number that merely contains a status-code digit run must NOT
// be mis-bucketed (digit-boundary check), even behind a provider marker.
{"incidental latency number", "provider error: request completed in 4290ms then failed", Unknown},
{"incidental request id", "provider error: request id 14015 failed unexpectedly", Unknown},
{"standalone 401 still auth", "provider request error: got 401 from upstream", Auth},
{"standalone 429 still rate limit", "provider error: 429 returned", RateLimit},
// Local (non-provider) failures must NOT be classified — no provider marker,
// so no bogus recovery hint attaches to them.
{"local fs permission denied", "open /etc/shadow: permission denied", Unknown},
{"local file missing", "stat foo.go: no such file or directory", Unknown},
{"local model config typo", "unknown model: sonnet-99", Unknown},
{"tool timeout local", "tool bash timed out after 600s", Unknown},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
if got := Classify(errors.New(tc.msg)); got != tc.want {
t.Fatalf("Classify(%q) = %v, want %v", tc.msg, got, tc.want)
}
})
}
}
func TestClassifyNil(t *testing.T) {
if got := Classify(nil); got != Unknown {
t.Fatalf("Classify(nil) = %v, want Unknown", got)
}
}
// Context-overflow must win over the connectivity "timeout" catch-all: a message
// mentioning both "context length" and a timeout is an overflow, not a network
// problem.
func TestContextOverflowBeatsConnectivity(t *testing.T) {
err := errors.New("provider request error: maximum context length exceeded; request timeout")
if got := Classify(err); got != ContextOverflow {
t.Fatalf("Classify = %v, want ContextOverflow", got)
}
}
func TestHintsPresentForKnownCategories(t *testing.T) {
known := []error{
errors.New("auth error: bad key"),
errors.New("rate limit error: 429"),
errors.New("provider stream error: dial tcp: no such host"),
errors.New("provider request error: model does not exist"),
errors.New("provider request error: maximum context length is 128000 tokens"),
}
for _, err := range known {
if h := TUIHint(err); strings.TrimSpace(h) == "" {
t.Fatalf("TUIHint(%q) empty, want a hint", err)
}
if h := CLIHint(err); strings.TrimSpace(h) == "" {
t.Fatalf("CLIHint(%q) empty, want a hint", err)
}
}
// Unknown yields no hint on either surface.
unknown := errors.New("provider error: mystery")
if h := TUIHint(unknown); h != "" {
t.Fatalf("TUIHint(unknown) = %q, want empty", h)
}
if h := CLIHint(unknown); h != "" {
t.Fatalf("CLIHint(unknown) = %q, want empty", h)
}
}
+2 -1
View File
@@ -24,7 +24,7 @@ type Options struct {
ModelRegistry *modelregistry.Registry
// OAuthResolver, when set, lets the provider authenticate model calls with an
// OAuth bearer token (preferred over the API key). nil => API-key auth only.
// Applied to the OpenAI and Anthropic providers.
// Applied to the OpenAI, Anthropic, and Google (Gemini) providers.
OAuthResolver providerio.TokenResolver
// OAuthLoginKey is the credential-store key OAuthResolver bound to (empty when
// there is no OAuth login). It is the SAME selection the bearer resolver made,
@@ -91,6 +91,7 @@ func New(profile config.ProviderProfile, options Options) (zeroruntime.Provider,
AuthScheme: profile.AuthScheme,
AuthHeaderValue: profile.AuthHeaderValue,
CustomHeaders: profile.CustomHeaders,
OAuthResolver: options.OAuthResolver,
MaxTokens: resolved.maxOutputTokens,
HTTPClient: options.HTTPClient,
UserAgent: options.UserAgent,
+16 -8
View File
@@ -49,6 +49,10 @@ type Options struct {
CustomHeaders map[string]string
HTTPClient *http.Client
UserAgent string
// OAuthResolver, when set, supplies an OAuth bearer credential per request and
// is retried once with a forced token refresh after an upstream 401 (matching
// the OpenAI and Anthropic providers). Nil falls back to plain API-key auth.
OAuthResolver providerio.TokenResolver
// StreamIdleTimeout aborts the stream if no data arrives for this long.
// When unset, Zero uses providerio.ResolveStreamIdleTimeout — the
// ZERO_STREAM_IDLE_TIMEOUT override or providerio.DefaultStreamIdleTimeout.
@@ -67,6 +71,7 @@ type Provider struct {
customHeaders map[string]string
httpClient *http.Client
userAgent string
oauthResolver providerio.TokenResolver
streamIdleTimeout time.Duration
}
@@ -95,6 +100,7 @@ func New(options Options) (*Provider, error) {
customHeaders: providerio.CopyHeaders(options.CustomHeaders),
httpClient: providerio.HTTPClient(options.HTTPClient),
userAgent: options.UserAgent,
oauthResolver: options.OAuthResolver,
streamIdleTimeout: providerio.ResolveStreamIdleTimeout(options.StreamIdleTimeout),
}, nil
}
@@ -127,20 +133,22 @@ func (provider *Provider) stream(ctx context.Context, body []byte, events chan<-
streamCtx, cancelStream := context.WithCancel(ctx)
defer cancelStream()
response, err := providerio.SendWithRetry(streamCtx, provider.httpClient, http.MethodPost, provider.streamURL(), body, func(request *http.Request) {
request.Header.Set("Content-Type", "application/json")
if provider.userAgent != "" {
request.Header.Set("User-Agent", provider.userAgent)
}
providerio.ApplyAuthHeaders(request, providerio.AuthHeaders{
response, err := providerio.SendWithAuthRetry(streamCtx, provider.httpClient, http.MethodPost, provider.streamURL(), body,
providerio.AuthHeaders{
APIKey: provider.apiKey,
DefaultAuthHeader: "x-goog-api-key",
AuthHeader: provider.authHeader,
AuthScheme: provider.authScheme,
AuthHeaderValue: provider.authHeaderValue,
CustomHeaders: provider.customHeaders,
})
}, 0)
},
provider.oauthResolver,
func(request *http.Request) {
request.Header.Set("Content-Type", "application/json")
if provider.userAgent != "" {
request.Header.Set("User-Agent", provider.userAgent)
}
}, 0)
if err != nil {
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventError, Error: provider.redact("provider stream error: " + err.Error())})
return
@@ -370,6 +370,57 @@ func TestStreamCompletionClassifiesHTTPAndPromptBlockErrors(t *testing.T) {
}
}
// A 401 with an OAuth resolver is retried once with a force-refreshed token; the
// replayed request carries the refreshed bearer and succeeds.
func TestStreamCompletionRetries401WithRefreshedToken(t *testing.T) {
var attempts int
var secondAuth string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
attempts++
if attempts == 1 {
http.Error(w, `{"error":{"message":"token expired"}}`, http.StatusUnauthorized)
return
}
secondAuth = r.Header.Get("Authorization")
writeSSE(w, `{}`)
}))
defer server.Close()
var forceRefreshOnRetry bool
resolver := func(ctx context.Context, forceRefresh bool) (string, string, bool, error) {
if forceRefresh {
forceRefreshOnRetry = true
return "Authorization", "Bearer refreshed", true, nil
}
return "Authorization", "Bearer stale", true, nil
}
provider, err := New(Options{
APIKey: "sk-google",
BaseURL: server.URL,
Model: "gemini-test",
OAuthResolver: resolver,
})
if err != nil {
t.Fatalf("New returned error: %v", err)
}
stream, err := provider.StreamCompletion(context.Background(), validRequest())
if err != nil {
t.Fatalf("StreamCompletion returned error: %v", err)
}
drain(stream)
if attempts != 2 {
t.Fatalf("attempts = %d, want 2 (initial 401 + one refreshed retry)", attempts)
}
if !forceRefreshOnRetry {
t.Fatalf("resolver was not called with forceRefresh on the retry")
}
if secondAuth != "Bearer refreshed" {
t.Fatalf("retry Authorization = %q, want refreshed bearer", secondAuth)
}
}
func TestStreamCompletionEmitsStreamErrorObject(t *testing.T) {
provider := newTestProviderWithKey(t, "sk-google", func(w http.ResponseWriter, r *http.Request) {
writeSSE(w, `{"error":{"code":429,"message":"stream failed sk-google","status":"RESOURCE_EXHAUSTED"}}`)
+23 -2
View File
@@ -272,15 +272,36 @@ func coerceAskUserOptions(optionsValue, descsValue any) (labels []string, descri
// FormatAskUserAnswers renders question/answer pairs into a clear, model-readable
// block. Missing answers are surfaced explicitly so the model never silently
// treats an unanswered question as answered.
//
// It distinguishes two shapes of "empty" the model would otherwise conflate: a
// wholesale dismissal (the user closed the prompt without answering ANYTHING) is
// flagged up front as a skip, so the model doesn't invent a default; a single
// blank field amid other answers is marked "(left blank)".
func FormatAskUserAnswers(questions []AskUserQuestion, answers []string) string {
lines := make([]string, 0, len(questions)*3)
anyAnswered := false
for index := range questions {
if index < len(answers) && strings.TrimSpace(answers[index]) != "" {
anyAnswered = true
break
}
}
lines := make([]string, 0, len(questions)*3+2)
if len(questions) > 0 && !anyAnswered {
lines = append(lines, "[note] The user dismissed this prompt without answering. Treat this as a skip: do not assume a default — ask again more specifically, or proceed only if the intent is already unambiguous.")
lines = append(lines, "")
}
for index, question := range questions {
answer := ""
if index < len(answers) {
answer = strings.TrimSpace(answers[index])
}
if answer == "" {
answer = "(no answer provided)"
if anyAnswered {
answer = "(left blank)"
} else {
answer = "(skipped)"
}
}
lines = append(lines, fmt.Sprintf("%d. [question] %s", index+1, question.Question))
lines = append(lines, "[answer] "+answer)
+34
View File
@@ -0,0 +1,34 @@
package tools
import (
"strings"
"testing"
)
func TestFormatAskUserAnswersDistinguishesSkipFromBlank(t *testing.T) {
questions := []AskUserQuestion{
{Question: "Which database?"},
{Question: "Migrate existing data?"},
}
// Wholesale dismissal (nothing answered): flagged as a skip up front.
dismissed := FormatAskUserAnswers(questions, []string{"", ""})
if !strings.Contains(dismissed, "dismissed this prompt") || !strings.Contains(dismissed, "Treat this as a skip") {
t.Fatalf("all-empty answers must be flagged as a skip, got:\n%s", dismissed)
}
if strings.Contains(dismissed, "left blank") {
t.Fatalf("a full dismissal should read as skipped, not left-blank:\n%s", dismissed)
}
if !strings.Contains(dismissed, "(skipped)") {
t.Fatalf("dismissed questions should render (skipped):\n%s", dismissed)
}
// Partial answer: the blank one is "left blank", no wholesale-skip note.
partial := FormatAskUserAnswers(questions, []string{"postgres", ""})
if strings.Contains(partial, "dismissed this prompt") {
t.Fatalf("a partial answer is not a dismissal:\n%s", partial)
}
if !strings.Contains(partial, "postgres") || !strings.Contains(partial, "(left blank)") {
t.Fatalf("partial answers should show the answer and mark the blank one left-blank:\n%s", partial)
}
}
+149 -23
View File
@@ -1,7 +1,6 @@
package tools
import (
"bytes"
"context"
"errors"
"fmt"
@@ -123,10 +122,14 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS
defer plan.Cleanup()
addSandboxMeta(meta, plan)
var stdout bytes.Buffer
var stderr bytes.Buffer
command.Stdout = &stdout
command.Stderr = &stderr
// Bound the capture so a command with runaway output (`cat huge.log`, `yes`)
// can't grow Zero's memory before truncation: only the head+tail each stream
// will ever surface to the model are retained, the middle is discarded as it
// streams, and the true size is counted for the truncation marker.
stdout := newBoundedBuffer(bashOutputBudgetBytes, bashOutputBudgetBytes)
stderr := newBoundedBuffer(bashOutputBudgetBytes, bashOutputBudgetBytes)
command.Stdout = stdout
command.Stderr = stderr
// Kill the shell as a process group on timeout and bound the post-kill I/O
// wait, so a backgrounded child cannot outlive the command or hang Run().
@@ -138,7 +141,12 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS
err = command.Run()
exitCode := commandExitCode(err)
meta["exit_code"] = strconv.Itoa(exitCode)
stderrText := appendSandboxBlocks(stderr.String(), monitor.Stop())
stdoutText := stdout.retained()
stderrRetained := stderr.retained()
stderrText := appendSandboxBlocks(stderrRetained, monitor.Stop())
// Sandbox blocks are extra model-visible stderr bytes appended after capture;
// count them toward the true total so the budget/marker stay accurate.
stderrTotal := stderr.total + (len(stderrText) - len(stderrRetained))
if errors.Is(commandCtx.Err(), context.DeadlineExceeded) {
return Result{
@@ -155,26 +163,31 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS
Meta: meta,
}
}
markLikelySandboxDenial(meta, plan, exitCode, stdout.String(), stderrText)
markLikelySandboxDenial(meta, plan, exitCode, stdoutText, stderrText)
outText, errText, truncated := budgetBashCapture(stdoutText, stdout.total, stderrText, stderrTotal, meta)
return Result{
Status: StatusError,
Output: formatBashOutputWithShellHint(commandText, stdout.String(), stderrText, exitCode, meta),
Meta: meta,
Status: StatusError,
Output: formatBashOutputWithShellHint(commandText, outText, errText, exitCode, meta),
Truncated: truncated,
Meta: meta,
}
}
markLikelySandboxDenial(meta, plan, exitCode, stdout.String(), stderrText)
markLikelySandboxDenial(meta, plan, exitCode, stdoutText, stderrText)
outText, errText, truncated := budgetBashCapture(stdoutText, stdout.total, stderrText, stderrTotal, meta)
if meta[SandboxLikelyDeniedMeta] == "true" {
return Result{
Status: StatusError,
Output: formatBashOutputWithShellHint(commandText, stdout.String(), stderrText, exitCode, meta),
Meta: meta,
Status: StatusError,
Output: formatBashOutputWithShellHint(commandText, outText, errText, exitCode, meta),
Truncated: truncated,
Meta: meta,
}
}
return Result{
Status: StatusOK,
Output: formatBashOutput(stdout.String(), stderrText, exitCode),
Meta: meta,
Status: StatusOK,
Output: formatBashOutput(outText, errText, exitCode),
Truncated: truncated,
Meta: meta,
}
}
@@ -354,12 +367,125 @@ func formatBashOutput(stdout string, stderr string, exitCode int) string {
if len(parts) == 0 {
return "Command completed with no output."
}
// Cap the context cost of a verbose command at the same default budget as
// exec_command (head/tail window); the full output is spilled to disk and
// referenced in the truncation notice, so nothing is lost. Previously bash
// output was unbounded — one chatty build log could flood the context.
output, _ := truncateExecOutputSpill(strings.Join(parts, "\n"), defaultMaxOutputTokens, "bash")
return output
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 {
+127
View File
@@ -0,0 +1,127 @@
package tools
import (
"fmt"
"strconv"
"strings"
"testing"
)
// Small output passes through untouched and records raw==emitted, no truncated flag.
func TestBudgetBashOutputSmallPassesThrough(t *testing.T) {
meta := map[string]string{}
out, errStr, truncated := budgetBashOutput("hello\n", "warn\n", meta)
if out != "hello\n" || errStr != "warn\n" {
t.Fatalf("small output altered: out=%q err=%q", out, errStr)
}
if truncated {
t.Fatalf("small output must report truncated=false")
}
if meta["truncated"] == "true" {
t.Fatalf("small output must not be flagged truncated: %v", meta)
}
if meta["raw_bytes"] != strconv.Itoa(len("hello\n")+len("warn\n")) {
t.Fatalf("raw_bytes wrong: %v", meta)
}
}
// Oversized stdout is truncated head+tail: both the first and last lines survive,
// the middle is dropped behind a marker, and meta is flagged.
func TestBudgetBashOutputTruncatesHeadAndTail(t *testing.T) {
head := "FIRST_LINE_MARKER\n"
tail := "\nLAST_LINE_MARKER"
big := head + strings.Repeat("x", bashOutputBudgetBytes) + tail
meta := map[string]string{}
out, _, truncated := budgetBashOutput(big, "", meta)
if !truncated {
t.Fatalf("oversized output must report truncated=true")
}
if !strings.Contains(out, "FIRST_LINE_MARKER") {
t.Fatalf("head lost after truncation")
}
if !strings.Contains(out, "LAST_LINE_MARKER") {
t.Fatalf("tail lost after truncation (failures live at the tail)")
}
if !strings.Contains(out, "output truncated") {
t.Fatalf("expected a truncation marker, got:\n%s", out[:200])
}
if len(out) > bashOutputBudgetBytes {
t.Fatalf("emitted %d bytes exceeds budget %d", len(out), bashOutputBudgetBytes)
}
if meta["truncated"] != "true" {
t.Fatalf("expected truncated=true, got %v", meta)
}
if meta["raw_bytes"] != strconv.Itoa(len(big)) {
t.Fatalf("raw_bytes = %s, want %d", meta["raw_bytes"], len(big))
}
if got, _ := strconv.Atoi(meta["emitted_bytes"]); got != len(out) {
t.Fatalf("emitted_bytes = %s, want %d", meta["emitted_bytes"], len(out))
}
}
// boundedBuffer must keep memory bounded to head+tail while counting every byte,
// so a runaway command can't OOM Zero before truncation runs. It keeps the very
// first bytes (head) and the very last bytes (tail), discarding the middle.
func TestBoundedBufferKeepsHeadAndTailBounded(t *testing.T) {
b := newBoundedBuffer(8, 8)
total := 0
for i := 0; i < 1000; i++ {
chunk := []byte(fmt.Sprintf("%05d.", i)) // 6 bytes each
n, err := b.Write(chunk)
if err != nil || n != len(chunk) {
t.Fatalf("Write = (%d, %v), want (%d, nil)", n, err, len(chunk))
}
total += len(chunk)
}
if b.total != total {
t.Fatalf("total = %d, want %d (every byte must be counted)", b.total, total)
}
// Retained memory is bounded regardless of how much streamed through.
if len(b.head) > 8 {
t.Fatalf("head grew past its cap: %d", len(b.head))
}
if len(b.tail) > 2*8 {
t.Fatalf("tail grew past 2×cap (compaction failed): %d", len(b.tail))
}
r := b.retained()
if len(r) > 8+8 {
t.Fatalf("retained %d bytes, want <= headCap+tailCap", len(r))
}
if !strings.HasPrefix(r, "00000.") {
t.Fatalf("retained head lost, got prefix %q", r[:min(6, len(r))])
}
if !strings.HasSuffix(r, "00999.") {
t.Fatalf("retained tail lost, got %q", r)
}
}
// budgetBashCapture reports the TRUE total (not the retained size) in the marker
// and raw_bytes, even though only a bounded head+tail was ever held in memory.
func TestBudgetBashCaptureReportsTrueTotal(t *testing.T) {
// Retained head+tail as boundedBuffer would hand over; the real command produced
// far more than was kept.
retained := "HEAD_START" + strings.Repeat("y", bashOutputBudgetBytes) + "TAIL_END"
total := 10 * bashOutputBudgetBytes
meta := map[string]string{}
out, _, truncated := budgetBashCapture(retained, total, "", 0, meta)
if !truncated {
t.Fatal("oversized capture must report truncated=true")
}
if !strings.Contains(out, "HEAD_START") || !strings.Contains(out, "TAIL_END") {
t.Fatalf("head/tail lost after budgeting:\n%s", out[:min(120, len(out))])
}
if len(out) > bashOutputBudgetBytes {
t.Fatalf("emitted %d bytes exceeds budget %d", len(out), bashOutputBudgetBytes)
}
if meta["raw_bytes"] != strconv.Itoa(total) {
t.Fatalf("raw_bytes = %s, want the true total %d", meta["raw_bytes"], total)
}
// The marker must cite the true omitted count (total-budget), not retained-budget.
if !strings.Contains(out, strconv.Itoa(total-bashOutputBudgetBytes)) {
t.Fatalf("marker should cite the true omitted byte count:\n%s", out[:min(200, len(out))])
}
}
+38
View File
@@ -8,6 +8,7 @@ import (
"os"
"path/filepath"
"runtime"
"strconv"
"strings"
"testing"
"time"
@@ -252,6 +253,43 @@ func TestBashToolRunsCommandInWorkspace(t *testing.T) {
}
}
// A command with runaway output must not be buffered whole in memory: the capture
// is bounded to head+tail, yet raw_bytes still reports the true (much larger) size.
// If capture were unbounded this would balloon Zero's memory before truncation.
func TestBashToolBoundsRunawayOutputCapture(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses a POSIX shell pipeline (yes | head)")
}
const produced = 500000 // ~5× the 96 KiB budget
result := NewBashTool(t.TempDir()).Run(context.Background(), map[string]any{
"command": fmt.Sprintf("yes ABCDEFGH | head -c %d", produced),
})
if result.Status != StatusOK {
t.Fatalf("expected ok status, got %s: %s", result.Status, result.Output)
}
if !result.Truncated || result.Meta["truncated"] != "true" {
t.Fatalf("runaway output should be flagged truncated, meta=%v", result.Meta)
}
if !strings.Contains(result.Output, "output truncated") {
t.Fatalf("expected a truncation marker in output")
}
// The emitted (model-visible) text is capped near the budget...
if len(result.Output) > bashOutputBudgetBytes+1024 {
t.Fatalf("emitted %d bytes far exceeds budget %d", len(result.Output), bashOutputBudgetBytes)
}
// ...but raw_bytes reflects the full stream, proving the total was counted while
// only a bounded slice was ever held (raw >> the ~2×budget retention cap).
raw, err := strconv.Atoi(result.Meta["raw_bytes"])
if err != nil {
t.Fatalf("raw_bytes not an int: %q", result.Meta["raw_bytes"])
}
if raw < produced {
t.Fatalf("raw_bytes = %d, want >= %d (full stream counted)", raw, produced)
}
}
func TestBashToolUsesRequestedCwd(t *testing.T) {
root := t.TempDir()
nested := filepath.Join(root, "nested")
+4 -1
View File
@@ -114,7 +114,10 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
switch {
case ferr == nil:
oldString = search
newString = findNew
// 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.")
+123 -26
View File
@@ -13,6 +13,7 @@ package tools
import (
"errors"
"regexp"
"sort"
"strings"
)
@@ -45,26 +46,42 @@ func fuzzyEditMatch(content, find string, replaceAll bool) (string, error) {
}
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 == "" {
if search == "" || seen[search] {
continue
}
index := strings.Index(content, search)
if index < 0 {
continue
}
found = true
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.LastIndex(content, search) != index {
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
@@ -202,22 +219,33 @@ func blockAnchorReplacer(content, find string) []string {
return similarity / float64(linesToCheck)
}
best := candidates[0]
if len(candidates) > 1 {
bestSimilarity := -1.0
for _, candidate := range candidates {
if s := middleSimilarity(candidate); s > bestSimilarity {
bestSimilarity = s
best = candidate
}
// 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 bestSimilarity < editAnchorSimilarityThreshold {
return nil
}
} else if middleSimilarity(best) < editAnchorSimilarityThreshold {
}
if len(qualifying) == 0 {
return nil
}
return []string{strings.Join(contentLines[best.start:best.end+1], "\n")}
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+`)
@@ -422,6 +450,75 @@ func contextAwareReplacer(content, find string) []string {
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 == "" {
+108 -3
View File
@@ -47,11 +47,13 @@ func TestEditFuzzyLineTrimmed(t *testing.T) {
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": "\ty := 2\n\treturn y",
"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)
}
@@ -112,16 +114,69 @@ func TestEditFuzzyIndentationFlexible(t *testing.T) {
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": "\t\tsum += total(xs)",
"new_string": "sum += total(xs)",
})
if result.Status != StatusOK {
t.Fatalf("expected ok, got %q", result.Output)
}
if !strings.Contains(after, "total(xs)") || strings.Contains(after, "range xs") {
// 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.
@@ -262,3 +317,53 @@ func TestLevenshtein(t *testing.T) {
}
}
}
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)
}
}
+59 -2
View File
@@ -1,6 +1,7 @@
package tools
import (
"fmt"
"os"
"path/filepath"
"strings"
@@ -21,12 +22,68 @@ import (
// are opportunistically removed whenever a new spill happens.
const spillRetention = 7 * 24 * time.Hour
// spillDir returns the shared spill directory, creating it on first use.
// 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 := filepath.Join(os.TempDir(), "zero-tool-output")
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
}
+91
View File
@@ -0,0 +1,91 @@
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)
}
}
+24
View File
@@ -0,0 +1,24 @@
//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
}
+12
View File
@@ -0,0 +1,12 @@
//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
}
-12
View File
@@ -39,18 +39,6 @@ func TestTruncateExecOutputUnderBudgetUnchanged(t *testing.T) {
}
}
func TestFormatBashOutputCapsVerboseOutput(t *testing.T) {
t.Setenv("TMPDIR", t.TempDir())
long := strings.Repeat("x", 200*1024)
formatted := formatBashOutput(long, "", 0)
if len(formatted) > defaultMaxOutputTokens*4+1024 {
t.Fatalf("bash output must be capped near the token budget, got %d bytes", len(formatted))
}
if !strings.Contains(formatted, "output truncated") {
t.Fatal("capped bash output must carry the truncation notice")
}
}
func TestSweepSpillDirRemovesOnlyOldFiles(t *testing.T) {
dir := t.TempDir()
oldFile := filepath.Join(dir, "bash-old.txt")
+20 -1
View File
@@ -261,10 +261,16 @@ func (tool toolSearchTool) rankByKeyword(query string, deferred []Tool) []Tool {
for index, candidate := range deferred {
name := strings.ToLower(candidate.Name())
desc := strings.ToLower(candidate.Description())
nameSquashed := squashSeparators(name)
score := 0
for _, keyword := range keywords {
if strings.Contains(name, keyword) {
switch {
case strings.Contains(name, keyword):
score += 2
case strings.Contains(nameSquashed, squashSeparators(keyword)):
// Separator-insensitive fallback: "webfetch" matches "web_fetch".
// Ranked below an exact substring so precise queries still win.
score++
}
if strings.Contains(desc, keyword) {
score++
@@ -290,6 +296,19 @@ func (tool toolSearchTool) rankByKeyword(query string, deferred []Tool) []Tool {
return matches
}
// squashSeparators drops separators so "web_fetch", "web-fetch", and "web fetch"
// all normalize to "webfetch". It lets a query missing the separators still match
// a tool name, the common way a model mistypes a tool ("webfetch").
func squashSeparators(s string) string {
return strings.Map(func(r rune) rune {
switch r {
case '_', '-', ' ', '.':
return -1
}
return r
}, s)
}
// noMatchMessage reports that nothing loaded and names the available deferred
// tools so the model can retry with a valid select: query.
func (tool toolSearchTool) noMatchMessage(query string, deferred []Tool) string {
+23
View File
@@ -184,6 +184,29 @@ func TestToolSearchKeywordRanksByNameThenDescription(t *testing.T) {
}
}
// A query missing the name's separators still matches: "webfetch" -> web_fetch.
func TestToolSearchKeywordMatchesSeparatorInsensitive(t *testing.T) {
reg := NewRegistry()
reg.Register(searchFakeTool{
name: "web_fetch",
description: "Fetch a URL.",
parameters: Schema{Type: "object", AdditionalProperties: false},
})
reg.Register(searchFakeTool{
name: "stock_quote",
description: "Get a stock price.",
parameters: Schema{Type: "object", AdditionalProperties: false},
})
tool := NewToolSearchTool(reg).(optionsAwareTool)
result := tool.RunWithOptions(context.Background(),
map[string]any{"query": "webfetch"}, RunOptions{})
if got := result.Meta["load_tools"]; got != "web_fetch" {
t.Fatalf("load_tools = %q, want web_fetch (separator-insensitive match)", got)
}
}
func TestToolSearchKeywordExcludesNonMatches(t *testing.T) {
reg := newDeferredFixtureRegistry()
tool := NewToolSearchTool(reg).(optionsAwareTool)
+8
View File
@@ -254,6 +254,14 @@ 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)
}
+109
View File
@@ -0,0 +1,109 @@
package tui
import (
"sync"
"time"
tea "charm.land/bubbletea/v2"
)
// streamCoalesceInterval is roughly one 60fps frame. Assistant-text deltas that
// arrive within this window are merged into a single agentTextMsg, so the render
// rate decouples from the token rate: a fast local model (100+ tok/s) no longer
// forces 100+ full Update→View cycles (each re-parsing the growing markdown) per
// second. Rendering stays smooth regardless of provider speed.
const streamCoalesceInterval = 16 * time.Millisecond
// textCoalescer batches agentTextMsg deltas before forwarding them to the Bubble
// Tea program. Any OTHER message flushes the pending text first, so ordering
// between streamed prose and tool-call / reasoning / row / usage messages is
// preserved. The turn's final agentResponseMsg does not pass through here (it is
// a tea.Cmd return, not a sink message), but the model drops deltas whose runID
// is no longer active, so a flush that races just past end-of-turn is harmless.
//
// Sink messages originate from the single agent goroutine and so arrive
// serially; the only concurrent caller is the flush timer. The mutex guards the
// buffer/timer AND is held across the downstream forward, so a timer-fired text
// flush can never overtake a concurrent non-text message: whoever holds the lock
// drains and forwards atomically, and the other caller blocks until it is done.
type textCoalescer struct {
forward func(tea.Msg) // downstream sink (external sink + program.Send)
// afterFunc schedules fn to run after one frame interval and returns a
// stoppable timer. Defaults to a real time.AfterFunc(streamCoalesceInterval, …);
// tests swap in a controllable timer so flush timing is deterministic instead of
// racing the 16ms wall clock.
afterFunc func(fn func()) coalesceTimer
mu sync.Mutex
buf []byte
runID int
timer coalesceTimer
}
// coalesceTimer is the subset of *time.Timer the coalescer needs. Abstracted
// behind afterFunc so a test can substitute a timer it controls.
type coalesceTimer interface {
Stop() bool
}
func newTextCoalescer(forward func(tea.Msg)) *textCoalescer {
return &textCoalescer{
forward: forward,
afterFunc: func(fn func()) coalesceTimer {
return time.AfterFunc(streamCoalesceInterval, fn)
},
}
}
// send is the coalescing entry point installed as the RuntimeMessageSink.
func (c *textCoalescer) send(msg tea.Msg) {
c.mu.Lock()
defer c.mu.Unlock()
text, ok := msg.(agentTextMsg)
if !ok {
// Non-text message: flush buffered text first (preserving order), then
// forward it — both under the lock so nothing can interleave between them.
c.drainAndForwardLocked()
c.forward(msg)
return
}
// A delta for a different run than the one buffered: flush the old run's text
// before buffering the new run's. In practice runs are sequential (the prior
// run's end already flushed via a non-text message), so this is belt-and-braces.
if len(c.buf) > 0 && text.runID != c.runID {
c.drainAndForwardLocked()
}
c.runID = text.runID
c.buf = append(c.buf, text.delta...)
if c.timer == nil {
c.timer = c.afterFunc(c.flush)
}
}
// flush forwards any buffered text as one agentTextMsg. Runs on the timer
// goroutine; the lock it takes serializes it against send so its output can't be
// reordered around a concurrent non-text message.
func (c *textCoalescer) flush() {
c.mu.Lock()
defer c.mu.Unlock()
c.drainAndForwardLocked()
}
// drainAndForwardLocked forwards any buffered text as one agentTextMsg and stops
// the timer, all while the caller holds c.mu — so a text flush and any non-text
// forward are strictly ordered and never interleave. A no-op when nothing is
// buffered. string(c.buf) copies, so reusing the backing array via [:0] is safe.
func (c *textCoalescer) drainAndForwardLocked() {
if len(c.buf) == 0 {
return
}
if c.timer != nil {
c.timer.Stop()
c.timer = nil
}
msg := agentTextMsg{runID: c.runID, delta: string(c.buf)}
c.buf = c.buf[:0]
c.forward(msg)
}
+182
View File
@@ -0,0 +1,182 @@
package tui
import (
"sync"
"testing"
"time"
tea "charm.land/bubbletea/v2"
)
// recorder is a thread-safe sink that captures forwarded messages in order.
type recorder struct {
mu sync.Mutex
msgs []tea.Msg
}
func (r *recorder) forward(msg tea.Msg) {
r.mu.Lock()
defer r.mu.Unlock()
r.msgs = append(r.msgs, msg)
}
func (r *recorder) snapshot() []tea.Msg {
r.mu.Lock()
defer r.mu.Unlock()
out := make([]tea.Msg, len(r.msgs))
copy(out, r.msgs)
return out
}
// manualTimer is a coalesceTimer that never fires on its own, so a test can prove
// buffering deterministically without racing the real 16ms frame timer: the flush
// happens only when the test calls c.flush() explicitly.
type manualTimer struct{ stopped bool }
func (m *manualTimer) Stop() bool {
prev := m.stopped
m.stopped = true
return !prev
}
// Rapid deltas within one frame collapse into a single agentTextMsg carrying the
// concatenated text.
func TestCoalescerBatchesDeltas(t *testing.T) {
rec := &recorder{}
c := newTextCoalescer(rec.forward)
// Swap in a timer that never auto-fires so the "still buffered" assertion below
// cannot flake if the scheduler pauses this goroutine past the frame interval.
c.afterFunc = func(func()) coalesceTimer { return &manualTimer{} }
c.send(agentTextMsg{runID: 1, delta: "Hel"})
c.send(agentTextMsg{runID: 1, delta: "lo, "})
c.send(agentTextMsg{runID: 1, delta: "world"})
// Nothing forwarded yet — still buffered within the frame (the timer we injected
// never fires; only the explicit flush below delivers the text).
if got := rec.snapshot(); len(got) != 0 {
t.Fatalf("deltas should buffer, forwarded %d early: %#v", len(got), got)
}
c.flush()
got := rec.snapshot()
if len(got) != 1 {
t.Fatalf("expected 1 coalesced message, got %d: %#v", len(got), got)
}
text, ok := got[0].(agentTextMsg)
if !ok || text.delta != "Hello, world" || text.runID != 1 {
t.Fatalf("coalesced message = %#v, want agentTextMsg{1, \"Hello, world\"}", got[0])
}
}
// A non-text message flushes buffered text first, so ordering (text before the
// tool call it precedes) is preserved.
func TestCoalescerFlushesTextBeforeOtherMessages(t *testing.T) {
rec := &recorder{}
c := newTextCoalescer(rec.forward)
c.send(agentTextMsg{runID: 1, delta: "about to run"})
c.send(toolCallStreamStartMsg{runID: 1, id: "t1", name: "bash"})
got := rec.snapshot()
if len(got) != 2 {
t.Fatalf("expected text then tool-call, got %d: %#v", len(got), got)
}
if text, ok := got[0].(agentTextMsg); !ok || text.delta != "about to run" {
t.Fatalf("first forwarded message must be the flushed text, got %#v", got[0])
}
if _, ok := got[1].(toolCallStreamStartMsg); !ok {
t.Fatalf("second forwarded message must be the tool-call start, got %#v", got[1])
}
}
// A delta for a new run flushes the previous run's buffered text before buffering
// the new run's, so text is never mis-attributed across runs.
func TestCoalescerFlushesOnRunSwitch(t *testing.T) {
rec := &recorder{}
c := newTextCoalescer(rec.forward)
c.send(agentTextMsg{runID: 1, delta: "run one text"})
c.send(agentTextMsg{runID: 2, delta: "run two text"})
c.flush()
got := rec.snapshot()
if len(got) != 2 {
t.Fatalf("expected one message per run, got %d: %#v", len(got), got)
}
if first := got[0].(agentTextMsg); first.runID != 1 || first.delta != "run one text" {
t.Fatalf("first message = %#v, want run 1 text flushed on switch", got[0])
}
if second := got[1].(agentTextMsg); second.runID != 2 || second.delta != "run two text" {
t.Fatalf("second message = %#v, want run 2 text", got[1])
}
}
// Under concurrency (the frame timer firing while a non-text message is sent),
// forwarding must stay serialized and ordered: buffered text is always delivered
// before a following non-text message, never after it. Run with -race to catch
// the interleaving CodeRabbit flagged.
func TestCoalescerOrdersTextBeforeNonTextUnderConcurrency(t *testing.T) {
for iter := 0; iter < 200; iter++ {
rec := &recorder{}
c := newTextCoalescer(rec.forward)
// Buffer text, then let the timer race against an inline non-text send.
c.send(agentTextMsg{runID: 1, delta: "before-tool"})
done := make(chan struct{})
go func() {
c.send(toolCallStreamStartMsg{runID: 1, id: "t1", name: "bash"})
close(done)
}()
<-done
// Give the timer a chance to also fire, then settle.
time.Sleep(2 * streamCoalesceInterval)
c.flush()
got := rec.snapshot()
textIdx, toolIdx := -1, -1
for i, m := range got {
switch m.(type) {
case agentTextMsg:
if textIdx == -1 {
textIdx = i
}
case toolCallStreamStartMsg:
toolIdx = i
}
}
if textIdx == -1 || toolIdx == -1 {
t.Fatalf("iter %d: missing messages: %#v", iter, got)
}
if textIdx > toolIdx {
t.Fatalf("iter %d: text (%d) forwarded after tool-call (%d): %#v", iter, textIdx, toolIdx, got)
}
}
}
// The frame timer flushes buffered text on its own without an explicit flush or a
// following message.
func TestCoalescerTimerFlushes(t *testing.T) {
rec := &recorder{}
c := newTextCoalescer(rec.forward)
c.send(agentTextMsg{runID: 1, delta: "timer-driven"})
deadline := time.Now().Add(time.Second)
for time.Now().Before(deadline) {
if len(rec.snapshot()) > 0 {
break
}
time.Sleep(streamCoalesceInterval)
}
got := rec.snapshot()
if len(got) != 1 {
t.Fatalf("timer should have flushed exactly one message, got %d: %#v", len(got), got)
}
if text := got[0].(agentTextMsg); text.delta != "timer-driven" {
t.Fatalf("timer-flushed message = %#v", got[0])
}
}
+8 -1
View File
@@ -482,6 +482,9 @@ func (m model) handleModelCommand(args string) (model, string) {
lines = append(lines, target.notice)
}
lines = append(lines, summary)
if warn := m.visionDropWarning(); warn != "" {
lines = append(lines, warn)
}
return m, strings.Join(lines, "\n")
}
@@ -536,7 +539,11 @@ func (m model) switchProviderModel(providerName, modelID string) (model, string,
cmds = append(cmds, cmd)
}
}
return m, fmt.Sprintf("Model\nSwitched to %s · %s", target.Name, target.Model), tea.Batch(cmds...)
status := fmt.Sprintf("Model\nSwitched to %s · %s", target.Name, target.Model)
if warn := m.visionDropWarning(); warn != "" {
status += "\n" + warn
}
return m, status, tea.Batch(cmds...)
}
// profileWithCredential fills a profile's APIKey for provider construction the same
+34 -2
View File
@@ -41,6 +41,10 @@ const (
commandAddDir
commandSelfCorrect
commandTurns
commandRetry
commandEdit
commandCopy
commandExport
commandNew
commandUnknown
)
@@ -209,9 +213,9 @@ var commandDefinitions = []commandDefinition{
},
{
name: "/compact",
usage: "/compact [status]",
usage: "/compact [status|now]",
group: commandGroupSession,
description: "Show or request transcript compaction state.",
description: "Compact the transcript now, or show compaction state (/compact status).",
kind: commandCompact,
},
{
@@ -257,6 +261,34 @@ var commandDefinitions = []commandDefinition{
description: "Show or set the per-run tool-turn budget for this session (raise it for long multi-step tasks).",
kind: commandTurns,
},
{
name: "/retry",
usage: "/retry",
group: commandGroupSession,
description: "Resend your last prompt.",
kind: commandRetry,
},
{
name: "/edit",
usage: "/edit",
group: commandGroupSession,
description: "Recall your last prompt into the composer to edit and resend.",
kind: commandEdit,
},
{
name: "/copy",
usage: "/copy",
group: commandGroupSession,
description: "Copy the last answer to the clipboard.",
kind: commandCopy,
},
{
name: "/export",
usage: "/export [path]",
group: commandGroupSession,
description: "Write the conversation transcript to a file.",
kind: commandExport,
},
{
name: "/help",
usage: "/help",
+27
View File
@@ -0,0 +1,27 @@
package tui
import (
"strings"
"testing"
)
// "/compact now" is accepted as a manual trigger (not a usage error), while a
// bogus argument still reports usage.
func TestCompactNowAccepted(t *testing.T) {
m := model{}
_, nowText, _ := m.handleCompactCommand("now")
if strings.Contains(nowText, "usage:") {
t.Fatalf("/compact now must be accepted as a trigger, got usage error: %q", nowText)
}
_, statusText, _ := m.handleCompactCommand("status")
if strings.Contains(statusText, "usage:") {
t.Fatalf("/compact status must not error, got: %q", statusText)
}
_, bogusText, _ := m.handleCompactCommand("frobnicate")
if !strings.Contains(bogusText, "usage:") {
t.Fatalf("/compact frobnicate must report usage, got: %q", bogusText)
}
}
+27
View File
@@ -0,0 +1,27 @@
package tui
import (
"strings"
"testing"
)
// A hinted error row renders the raw error and the faint one-line next step
// below it; an unhinted row renders just the error.
func TestRenderErrorRowShowsHint(t *testing.T) {
withHint := renderErrorRow(transcriptRow{
kind: rowError,
text: "auth error: your API key is missing or invalid",
hint: "API key rejected — run /provider to re-check your credentials",
}, 80)
if !strings.Contains(withHint, "auth error:") {
t.Fatalf("expected raw error text in output, got:\n%s", withHint)
}
if !strings.Contains(withHint, "/provider") {
t.Fatalf("expected hint referencing /provider, got:\n%s", withHint)
}
noHint := renderErrorRow(transcriptRow{kind: rowError, text: "provider error: mystery"}, 80)
if strings.Contains(noHint, "→") {
t.Fatalf("unhinted error must not render a hint arrow, got:\n%s", noHint)
}
}
+2 -2
View File
@@ -121,7 +121,7 @@ func TestEscCancellationLeavesVisibleMarker(t *testing.T) {
m.pending = true
m.activeRunID = 5
m.runCancel = func() {}
m.streamingText = "half an answer"
m.streamingText = []byte("half an answer")
updated, _ := m.Update(testKey(tea.KeyEsc))
next := updated.(model)
@@ -133,7 +133,7 @@ func TestEscCancellationLeavesVisibleMarker(t *testing.T) {
if !transcriptContains(next.transcript, "half an answer") {
t.Fatalf("expected the partial streamed answer to be preserved, got %#v", next.transcript)
}
if next.streamingText != "" {
if len(next.streamingText) != 0 {
t.Fatal("streaming text should be cleared after cancel")
}
}
+12
View File
@@ -292,6 +292,18 @@ func renderImageChips(labels []string) string {
// and staged documents, e.g. "[Image #1] [Image #2] [Doc #1]". Returns "" when
// nothing is staged. Numbered (not named) so a long screenshot path never shows
// in the composer.
// visionDropWarning returns a one-line notice when images are staged but the
// (now active) model can't accept them, so switching to a non-vision model warns
// the user immediately at switch time instead of silently dropping the images at
// submit. Empty when there is nothing staged or the model supports vision.
func (m model) visionDropWarning() string {
if len(m.pendingImages) == 0 || m.modelSupportsVisionTUI() {
return ""
}
return fmt.Sprintf("⚠ %d staged image(s) will be dropped — %s has no vision support.",
len(m.pendingImages), displayValue(m.modelName, "the active model"))
}
func renderAttachmentChips(imageLabels []string, docs []pendingDocument) string {
chips := make([]string, 0, len(imageLabels)+len(docs))
for i := range imageLabels {
+73
View File
@@ -483,3 +483,76 @@ func TestRenderAttachmentChips(t *testing.T) {
t.Fatalf("chip row %q should not show file names", got)
}
}
// TestRetryResendsAttachments guards the /retry contract: it must resend the exact
// request the last prompt carried — including its staged image and PDF document —
// not a degraded text-only prompt. launchPrompt clears the pending queues once a
// turn is sent, so without re-staging the remembered snapshot a vision/document-
// backed turn that failed would silently retry as a different task.
func TestRetryResendsAttachments(t *testing.T) {
root := t.TempDir()
captured := make(chan []zeroruntime.ImageBlock, 1)
provider := &fakeProvider{events: []zeroruntime.StreamEvent{
{Type: zeroruntime.StreamEventText, Content: "ok"},
{Type: zeroruntime.StreamEventDone},
}}
m := newModel(context.Background(), Options{
Cwd: root,
ProviderName: "openai",
ModelName: "gpt-4.1",
Provider: provider,
Registry: tools.NewRegistry(),
SessionStore: testSessionStore(t),
})
m.agentOptions.OnText = func(string) {}
m.captureRunImages = func(imgs []zeroruntime.ImageBlock) { captured <- imgs }
// State left after a prior vision+PDF prompt was submitted: the pending queues
// are cleared, but the remembered snapshot survives so /retry can reproduce it.
m.lastPrompt = "describe both"
m.lastImages = []zeroruntime.ImageBlock{{MediaType: "image/png", Data: []byte{0x89, 'P', 'N', 'G'}}}
m.lastImageLabels = []string{"photo.png"}
m.lastDocuments = []pendingDocument{{label: "spec.pdf", text: "Top secret design notes"}}
m.input.SetValue("/retry")
updated, cmd := m.handleSubmit()
next := updated.(model)
if cmd == nil {
t.Fatal("/retry with a remembered prompt should launch a run")
}
// The retried turn re-consumes and then clears the queues, exactly like a fresh
// submit; the snapshot must survive so a second /retry stays reproducible.
if len(next.pendingImages) != 0 || len(next.pendingImageLabels) != 0 || len(next.pendingDocuments) != 0 {
t.Fatalf("retry must clear pending queues after resend, got imgs=%d labels=%d docs=%d",
len(next.pendingImages), len(next.pendingImageLabels), len(next.pendingDocuments))
}
if len(next.lastImages) != 1 || len(next.lastDocuments) != 1 {
t.Fatalf("retry must keep the snapshot for a subsequent retry, got imgs=%d docs=%d",
len(next.lastImages), len(next.lastDocuments))
}
updated, _ = next.Update(execCmd(cmd))
_ = updated.(model)
select {
case imgs := <-captured:
if len(imgs) != 1 || imgs[0].MediaType != "image/png" {
t.Fatalf("retried run must resend the remembered image, got %#v", imgs)
}
default:
t.Fatal("expected the retried run to receive the remembered image")
}
if len(provider.requests) != 1 {
t.Fatalf("expected one provider request, got %d", len(provider.requests))
}
msgs := provider.requests[0].Messages
last := msgs[len(msgs)-1]
if !strings.Contains(last.Content, "Top secret design notes") {
t.Fatalf("retried prompt should re-prepend the document text, got:\n%s", last.Content)
}
if !strings.Contains(last.Content, "describe both") {
t.Fatalf("retried prompt should include the remembered user text, got:\n%s", last.Content)
}
}
+31
View File
@@ -0,0 +1,31 @@
package tui
import (
"context"
"testing"
)
// The LSP manager is built once at construction (when cwd is known) and reused —
// a fresh manager per run would cold-start gopls on the first edit of every turn.
func TestNewModelBuildsSharedLSPManager(t *testing.T) {
m := newModel(context.Background(), Options{Cwd: t.TempDir()})
if m.lspManager == nil {
t.Fatal("expected a session-long lspManager when cwd is set")
}
// The pointer is stable across model copies (Bubble Tea passes the model by
// value every Update), so every run shares the same warm servers.
copied := m
if copied.lspManager != m.lspManager {
t.Fatal("lspManager pointer must survive a model copy")
}
}
// shutdownLSPManager must be nil-safe: a model built without a cwd has no manager,
// and quitting it must not panic.
func TestShutdownLSPManagerNilSafe(t *testing.T) {
m := model{}
if m.lspManager != nil {
t.Fatal("expected no manager on a zero-value model")
}
m.shutdownLSPManager() // must not panic
}
+158 -36
View File
@@ -19,6 +19,7 @@ import (
"github.com/Gitlawb/zero/internal/agent"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/doctor"
"github.com/Gitlawb/zero/internal/errhint"
"github.com/Gitlawb/zero/internal/lsp"
internalmcp "github.com/Gitlawb/zero/internal/mcp"
"github.com/Gitlawb/zero/internal/modelregistry"
@@ -76,24 +77,29 @@ type model struct {
discoverProviderModels func(context.Context, config.ProviderProfile) ([]providermodeldiscovery.Model, error)
discoverOllamaContextWindow func(ctx context.Context, baseURL string, model string) (int, error)
registry *tools.Registry
sessionStore *sessions.Store
sandboxStore *sandbox.GrantStore
mcpConfig config.MCPConfig
mcpPermissionStore *internalmcp.PermissionStore
mcpTokenStore *internalmcp.TokenStore
mcpCommand func(context.Context, []string) MCPCommandResult
sandboxSetupCommand func(context.Context) SandboxSetupCommandResult
mcpViewStateCache MCPViewState
mcpViewStateReady bool
mcpCommandSeq int
mcpCommandCancel context.CancelFunc
sandboxSetupSeq int
sandboxSetupInFlight bool
doctorCommandSeq int
doctorInFlight bool
doctorFrame int
activeSession sessions.Metadata
sessionEvents []sessions.Event
// lspManager is created once per session and reused across prompts so gopls (and
// other language servers) stay warm — a fresh manager per run would cold-start
// the server on the first edit of every turn. Nil when cwd is unknown; runs then
// fall back to a per-run manager. Torn down in quit().
lspManager *lsp.Manager
sessionStore *sessions.Store
sandboxStore *sandbox.GrantStore
mcpConfig config.MCPConfig
mcpPermissionStore *internalmcp.PermissionStore
mcpTokenStore *internalmcp.TokenStore
mcpCommand func(context.Context, []string) MCPCommandResult
sandboxSetupCommand func(context.Context) SandboxSetupCommandResult
mcpViewStateCache MCPViewState
mcpViewStateReady bool
mcpCommandSeq int
mcpCommandCancel context.CancelFunc
sandboxSetupSeq int
sandboxSetupInFlight bool
doctorCommandSeq int
doctorInFlight bool
doctorFrame int
activeSession sessions.Metadata
sessionEvents []sessions.Event
// titledSessions records session ids for which a model-generated title has
// already been attempted this process, so a finished turn re-fires the title
// generator at most once per session (even before its async result lands).
@@ -261,13 +267,30 @@ type model struct {
headerPrinted bool
// Composer input history (shell-style ↑/↓ recall of submitted inputs).
// lastPrompt is the verbatim text of the most recent submitted prompt, so
// /retry can resend it and /edit can recall it into the composer.
lastPrompt string
// lastImages/lastImageLabels/lastDocuments remember the attachments consumed
// by the most recent submitted prompt. launchPrompt clears the pending queues
// once a turn is sent, so /retry re-stages these to reproduce the exact same
// request — otherwise a vision/PDF-backed prompt would silently retry as
// text-only and answer a different task. They share the underlying image bytes
// with the sent turn (never mutated in place), so no deep copy is needed.
lastImages []zeroruntime.ImageBlock
lastImageLabels []string
lastDocuments []pendingDocument
// historyIdx == len(inputHistory) means "not navigating"; historyDraft
// preserves whatever was typed before recall started.
inputHistory []string
historyIdx int
historyDraft string
streamingText string // live assistant text for the current segment
// streamingText is the live assistant text for the current segment, accumulated
// as []byte so each delta is an O(1) amortized append instead of the O(n²) that
// string += delta incurs across a long generation. Read via streamingTextString().
// A []byte (not strings.Builder) because the model is copied by value on every
// Update, which would trip strings.Builder's copy check.
streamingText []byte
streamingReasoning string // live provider reasoning for the current segment
streamingReasoningExpanded bool
// turnStreamedRunes accumulates every reasoning+answer rune streamed in the
@@ -773,6 +796,11 @@ func newModel(ctx context.Context, options Options) model {
// Streaming text always renders statically at base ink (the disabled path in
// styleStreamingLine), so no accent glow and no per-line fade ticks.
m.fadeDisabled = true
// One session-long LSP manager (cheap to build — servers start lazily on the
// first Check), reused across prompts so gopls stays warm between turns.
if cwd != "" {
m.lspManager = lsp.NewManager(cwd)
}
m.refreshMCPViewState()
return m
}
@@ -902,9 +930,22 @@ func (m model) noBlockingModal() bool {
func (m model) quit() (tea.Model, tea.Cmd) {
m.stopPRWatcher()
m.stopAllBackgroundTerminalSessions()
m.shutdownLSPManager()
return m, tea.Quit
}
// shutdownLSPManager gracefully stops the session-long language servers on exit.
// Best-effort with a short deadline so a slow server can't hang the quit; the
// servers are our child processes and would be reaped on exit regardless.
func (m model) shutdownLSPManager() {
if m.lspManager == nil {
return
}
shutdownCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
_ = m.lspManager.Shutdown(shutdownCtx)
}
func (m model) handleCtrlC() (tea.Model, tea.Cmd) {
if !m.pending && m.composerValue() != "" && m.noBlockingModal() && !m.transcriptDetailed && !m.subchat.active {
m.clearComposer()
@@ -1576,7 +1617,7 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
// Streaming text means any in-progress tool call has finished — clear the
// live "writing" block so it doesn't linger over new prose.
m.clearStreamingToolCall()
m.streamingText += msg.delta
m.streamingText = append(m.streamingText, msg.delta...)
m.turnStreamedRunes += utf8.RuneCountInString(msg.delta)
// recordStreamingDelta appends a time.Time to lineAges for every
// newline in the delta and bumps lastStreamActivity. It also
@@ -1859,20 +1900,23 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
if row, ok := reasoningTranscriptRow("", msg.runID, m.streamingReasoning); ok {
m.transcript = appendTranscriptRow(m.transcript, row)
}
if text := strings.TrimRight(m.streamingText, "\n"); strings.TrimSpace(text) != "" {
if text := strings.TrimRight(m.streamingTextString(), "\n"); strings.TrimSpace(text) != "" {
m.transcript = appendTranscriptRow(m.transcript, transcriptRow{kind: rowAssistant, text: text})
}
// The error row terminates the turn, so it carries the done-line
// metadata a final assistant row would have carried.
// metadata a final assistant row would have carried. A recognized
// provider failure (auth/rate-limit/connectivity/…) also carries a
// one-line next step so the user isn't left staring at a raw blob.
m.transcript = appendTranscriptRow(m.transcript, transcriptRow{
kind: rowError,
text: msg.err.Error(),
hint: errhint.TUIHint(msg.err),
final: true,
turnTools: msg.turnTools,
turnElapsed: msg.turnElapsed,
})
}
m.streamingText = ""
m.streamingText = nil
m.streamingReasoning = ""
m.streamingReasoningExpanded = false
// Roll the completed run's wall-time into the session's rolling average so
@@ -2047,14 +2091,14 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
m.streamingReasoning = ""
m.streamingReasoningExpanded = false
}
if text := strings.TrimRight(m.streamingText, "\n"); strings.TrimSpace(text) != "" {
if text := strings.TrimRight(m.streamingTextString(), "\n"); strings.TrimSpace(text) != "" {
m.transcript = appendTranscriptRow(m.transcript, transcriptRow{kind: rowAssistant, text: text})
// This interim narration is the agent explaining what it's about to
// do — attribute it to the active plan step so the step-detail card
// can replay the agent's own account of the work.
m = m.captureStepNarration(text)
}
m.streamingText = ""
m.streamingText = nil
// The tool call has finalized into its card — drop the live "writing"
// preview so it doesn't linger or duplicate beneath the card.
m.clearStreamingToolCall()
@@ -2785,7 +2829,7 @@ func (m model) liveReasoningBodyCap() int {
}
func (m model) interimBlock(width int) string {
text := strings.TrimRight(m.streamingText, "\n")
text := strings.TrimRight(m.streamingTextString(), "\n")
reasoning := strings.TrimRight(m.streamingReasoning, "\n")
blocks := []string{}
if strings.TrimSpace(reasoning) != "" {
@@ -2849,7 +2893,7 @@ func (m model) spinnerGlyph() string {
// (reasoning, waiting on the model, or a tool in flight). Cheap and robust — no
// transcript scan — so it can't misreport on a long, output-less step.
func (m model) workingActivity() string {
if strings.TrimSpace(m.streamingText) != "" {
if strings.TrimSpace(m.streamingTextString()) != "" {
return "writing"
}
return "thinking"
@@ -3897,6 +3941,62 @@ func (m model) handleSubmit() (tea.Model, tea.Cmd) {
}
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: "$ " + cmdText})
return m, runBashEscape(m.cwd, cmdText)
case commandRetry:
// /retry launches a run, so it needs the same guards a normal prompt gets:
// never start one while exiting (would strand the shutdown flush) or during
// compaction (would race compactResultMsg's wholesale rewrite of
// transcript/sessionEvents and silently drop events).
if m.exiting {
return m, nil
}
if m.pending {
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: "Retry\ncannot retry while a run is in progress."})
return m, nil
}
if m.compactInFlight {
m.transcript = reduceTranscript(m.transcript, transcriptAction{
kind: actionAppendSystem,
text: "Retry\nstatus: warning\nCompaction is running. Retry once it finishes.",
})
return m, nil
}
if strings.TrimSpace(m.lastPrompt) == "" {
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: "Retry\nno previous prompt to resend."})
return m, nil
}
// Re-stage the attachments the last prompt carried so launchPrompt rebuilds
// an identical request (document preamble + images + vision re-check). Without
// this the queues are empty and /retry would resend a text-only prompt,
// silently dropping the image/PDF context and answering a different task.
m.pendingImages = m.lastImages
m.pendingImageLabels = m.lastImageLabels
m.pendingDocuments = m.lastDocuments
return m.launchPrompt(m.lastPrompt)
case commandEdit:
if strings.TrimSpace(m.lastPrompt) == "" {
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: "Edit\nno previous prompt to recall."})
return m, nil
}
// Re-stage the remembered attachments alongside the recalled text so an
// edited resend carries the same image/PDF context — the reappearing chip
// row is the visible confirmation. Without this, editing a vision- or
// document-backed prompt would silently submit a text-only version and
// answer a different task (the same gap /retry guards against).
m.pendingImages = m.lastImages
m.pendingImageLabels = m.lastImageLabels
m.pendingDocuments = m.lastDocuments
m.input.SetValue(m.lastPrompt)
return m, nil
case commandCopy:
text := m.lastAssistantAnswer()
if strings.TrimSpace(text) == "" {
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: "Copy\nno answer to copy yet."})
return m, nil
}
return m, copyTranscriptSelectionCmd(text)
case commandExport:
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: m.handleExportCommand(command.text)})
return m, nil
case commandPrompt:
if intent, ok := detectMCPSetupIntent(command.text); ok {
return m.openMCPAddWizardFromIntent(intent), nil
@@ -3911,6 +4011,15 @@ func (m model) handleSubmit() (tea.Model, tea.Cmd) {
// composer. Queued prompts use this path too, so session and image behavior
// stays identical to immediate submissions.
func (m model) launchPrompt(prompt string) (model, tea.Cmd) {
// Remember the verbatim prompt (before specialist/document expansion) so /retry
// and /edit can act on exactly what the user submitted. Snapshot the staged
// attachments too: launchPrompt clears the pending queues below, so /retry
// re-stages these to resend an identical vision/PDF-backed request rather than
// a degraded text-only one.
m.lastPrompt = prompt
m.lastImages = m.pendingImages
m.lastImageLabels = m.pendingImageLabels
m.lastDocuments = m.pendingDocuments
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendUser, text: prompt})
if m.provider == nil {
m.transcript = reduceTranscript(m.transcript, transcriptAction{
@@ -4098,7 +4207,7 @@ func (m *model) cancelRun() {
if row, ok := reasoningTranscriptRow("", m.activeRunID, m.streamingReasoning); ok {
m.transcript = appendTranscriptRow(m.transcript, row)
}
if text := strings.TrimRight(m.streamingText, "\n"); strings.TrimSpace(text) != "" {
if text := strings.TrimRight(m.streamingTextString(), "\n"); strings.TrimSpace(text) != "" {
m.transcript = appendTranscriptRow(m.transcript, transcriptRow{kind: rowAssistant, text: text})
}
m.transcript = appendTranscriptRow(m.transcript, transcriptRow{kind: rowSystem, text: "Run cancelled."})
@@ -4119,7 +4228,7 @@ func (m *model) cancelRun() {
m.pendingAskUser = nil
// The interim block renders streamingText live; a cancelled run's partial
// answer must not leak into (and concatenate with) the next turn's stream.
m.streamingText = ""
m.streamingText = nil
m.streamingReasoning = ""
m.streamingReasoningExpanded = false
// Hard-stop the fade and drop the per-line age map. The next turn's
@@ -4196,12 +4305,18 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
// matching exec; the per-turn lsp.Manager is torn down when this run
// returns; auto-fix vs report-only follows the active permission mode.
if !runOptions.specDraft && options.Cwd != "" {
lspManager := lsp.NewManager(options.Cwd)
defer func() {
shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
_ = lspManager.Shutdown(shutdownCtx)
}()
// Prefer the session-long manager (kept warm across prompts). Only when it
// is absent — e.g. cwd was unknown at construction, or a test built the
// model directly — fall back to a per-run manager that is shut down here.
lspManager := m.lspManager
if lspManager == nil {
lspManager = lsp.NewManager(options.Cwd)
defer func() {
shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
_ = lspManager.Shutdown(shutdownCtx)
}()
}
options.SelfCorrect = agent.NewSelfCorrector(options.Cwd, agent.NewLSPDiagnosticsChecker(lspManager), agent.NewProjectVerifier(options.Cwd), agent.SelfCorrectConfig{
Enabled: true,
IncludeTests: m.selfCorrectTests,
@@ -4211,7 +4326,7 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
// 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.FileDiagnostics = agent.NewFileDiagnostics(lspManager, options.Cwd)
}
// Some providers synthesize tool-call ids that repeat within a run (e.g.
@@ -4663,6 +4778,13 @@ func (m model) sendAgentText(runID int, delta string) {
m.runtimeMessageSink(agentTextMsg{runID: runID, delta: delta})
}
// streamingTextString returns the accumulated live assistant text. streamingText
// is stored as []byte for O(1) amortized appends; the conversion here is bounded
// by the segment length, the same cost the renderer already pays.
func (m model) streamingTextString() string {
return string(m.streamingText)
}
func (m model) sendToolCallStreamStart(runID int, id, name string) {
if m.runtimeMessageSink == nil {
return
+1 -1
View File
@@ -2301,7 +2301,7 @@ func TestStreamingAssistantAfterToolCardGetsRuleSeparator(t *testing.T) {
m.width, m.height = 120, 40
m.altScreen = true
m.pending = true
m.streamingText = "Done."
m.streamingText = []byte("Done.")
m.transcript = append(m.transcript,
transcriptRow{kind: rowUser, text: "run it"},
transcriptRow{kind: rowAssistant, text: "I'll run it first."},
+9
View File
@@ -20,6 +20,10 @@ func TestStartNewSessionResetsState(t *testing.T) {
m.pendingImageLabels = []string{"pic.png"}
m.pendingDocuments = []pendingDocument{{label: "doc.pdf"}}
m.queuedMessage = "queued"
// The /retry attachment snapshot is prior-session state too and must not survive.
m.lastImages = make([]zeroruntime.ImageBlock, 1)
m.lastImageLabels = []string{"pic.png"}
m.lastDocuments = []pendingDocument{{label: "doc.pdf"}}
next := m.startNewSession()
@@ -41,6 +45,11 @@ func TestStartNewSessionResetsState(t *testing.T) {
t.Fatalf("startNewSession must clear staged input, got images=%d labels=%d docs=%d queued=%q",
len(next.pendingImages), len(next.pendingImageLabels), len(next.pendingDocuments), next.queuedMessage)
}
// The /retry snapshot must not leak the previous session's attachments.
if len(next.lastImages) != 0 || len(next.lastImageLabels) != 0 || len(next.lastDocuments) != 0 {
t.Fatalf("startNewSession must clear the retry snapshot, got images=%d labels=%d docs=%d",
len(next.lastImages), len(next.lastImageLabels), len(next.lastDocuments))
}
}
func TestNewCommandStartsFreshSession(t *testing.T) {
+26
View File
@@ -1151,6 +1151,29 @@ func (m model) setupProvider() SetupProviderOption {
return m.setup.providers[index]
}
// setupErrorAffordance returns a faint one-line recovery hint tailored to the
// current setup stage, so a first-run error points the user at the inline fix
// instead of leaving them staring at a red line. Empty for stages whose footer
// already makes recovery obvious. Each hint names only keys the stage's footer
// confirms, so the guidance is always accurate.
func (m model) setupErrorAffordance() string {
switch m.setup.stage {
case setupStageEndpoint:
return "→ edit the endpoint above, then Enter to retry (← to go back)"
case setupStageCredentials:
if m.setupCredentialInputActive() {
return "→ re-enter the key above, then Enter to retry (← to go back)"
}
case setupStageName:
return "→ edit the name above, then Enter to continue (← to go back)"
case setupStageModel:
if !m.setup.modelLoad {
return "→ pick another model with ↑/↓, or type to search"
}
}
return ""
}
func (m model) setupView(width int) string {
if width <= 0 {
width = defaultStartupWidth
@@ -1159,6 +1182,9 @@ func (m model) setupView(width int) string {
content := m.setupStageLines(width, height)
if m.setup.err != "" {
content = append(content, "", zeroTheme.red.Render("error: "+m.setup.err))
if hint := m.setupErrorAffordance(); hint != "" {
content = append(content, zeroTheme.faint.Render(hint))
}
}
progress := m.setupProgressText()
footer := m.setupFooter()
@@ -0,0 +1,29 @@
package tui
import (
"strings"
"testing"
)
// Each input/choice stage's error affordance names an accurate recovery step;
// stages without a tailored hint return empty.
func TestSetupErrorAffordance(t *testing.T) {
endpoint := model{}
endpoint.setup.stage = setupStageEndpoint
if got := endpoint.setupErrorAffordance(); !strings.Contains(got, "endpoint") || !strings.Contains(got, "Enter") {
t.Fatalf("endpoint affordance = %q", got)
}
name := model{}
name.setup.stage = setupStageName
if got := name.setupErrorAffordance(); !strings.Contains(got, "name") {
t.Fatalf("name affordance = %q", got)
}
// A stage with no tailored hint (method chooser) returns empty.
method := model{}
method.setup.stage = setupStageMethod
if got := method.setupErrorAffordance(); got != "" {
t.Fatalf("method stage should have no affordance, got %q", got)
}
}
+71 -4
View File
@@ -102,16 +102,83 @@ func (p *commandPicker) applyQuery() {
p.selected = clampInt(p.selected, 0, maxInt(0, len(p.items)-1))
return
}
filtered := make([]pickerItem, 0, len(source))
for _, item := range source {
if strings.Contains(strings.ToLower(strings.Join([]string{item.Group, item.Label, item.Value, item.Meta}, " ")), query) {
filtered = append(filtered, item)
// Rank matches (exact < prefix < contains < subsequence) instead of a flat
// substring filter, so the closest match to "sonnet 4.5" lands at the top
// rather than buried behind scrolling. Groups stay contiguous — a group is
// ordered by its best-matching item and never split into two header blocks —
// so the grouped model picker still renders one header per provider.
type entry struct {
item pickerItem
score int
order int
}
groupFirst := map[string]int{}
groupBest := map[string]int{}
entries := make([]entry, 0, len(source))
for index, item := range source {
score, ok := scorePickerItem(item, query)
if !ok {
continue
}
if _, seen := groupFirst[item.Group]; !seen {
groupFirst[item.Group] = index
}
if best, seen := groupBest[item.Group]; !seen || score < best {
groupBest[item.Group] = score
}
entries = append(entries, entry{item: item, score: score, order: index})
}
sort.SliceStable(entries, func(a, b int) bool {
ga, gb := entries[a].item.Group, entries[b].item.Group
if ga != gb {
// Most-relevant group first; ties keep original group appearance order so
// a group is never scattered across the list.
if groupBest[ga] != groupBest[gb] {
return groupBest[ga] < groupBest[gb]
}
return groupFirst[ga] < groupFirst[gb]
}
// Within a group: best match first, then original order.
if entries[a].score != entries[b].score {
return entries[a].score < entries[b].score
}
return entries[a].order < entries[b].order
})
filtered := make([]pickerItem, 0, len(entries))
for _, e := range entries {
filtered = append(filtered, e.item)
}
p.items = filtered
p.selected = 0
}
// scorePickerItem ranks an item against a lowercased query; lower is better, and
// ok is false when it doesn't match at all. Tiers mirror scoreFileSuggestion:
// exact/prefix/contains on the label (what the user reads) beat matches deeper in
// the joined haystack, and a fuzzy subsequence is the last-resort match.
func scorePickerItem(item pickerItem, query string) (int, bool) {
label := strings.ToLower(item.Label)
hay := strings.ToLower(strings.Join([]string{item.Group, item.Label, item.Value, item.Meta}, " "))
switch {
case label == query:
return 0, true
case strings.HasPrefix(label, query):
return 20, true
case strings.Contains(label, query):
return 40, true
case strings.HasPrefix(hay, query):
return 60, true
case strings.Contains(hay, query):
return 80, true
default:
if gap, ok := fuzzySubsequenceGap(hay, query); ok {
return 120 + gap, true
}
return 0, false
}
}
// newModelPicker lists active (non-deprecated) models, preselecting the active
// one. Returns nil when the catalog is unavailable so the caller falls back to
// the plain status text.
+72
View File
@@ -0,0 +1,72 @@
package tui
import (
"testing"
)
func newFuzzyTestPicker(items []pickerItem) *commandPicker {
return &commandPicker{
kind: pickerModel,
items: append([]pickerItem{}, items...),
allItems: append([]pickerItem{}, items...),
}
}
// A closer match ranks above a merely-containing one, so the best result lands at
// the top instead of in original order.
func TestPickerRanksBestMatchFirst(t *testing.T) {
p := newFuzzyTestPicker([]pickerItem{
{Label: "Claude Opus 4.5", Value: "opus-4.5"},
{Label: "Claude Sonnet 4.5", Value: "sonnet-4.5"},
{Label: "Sonnet", Value: "sonnet"}, // exact-ish prefix match should win
})
p.query = "sonnet"
p.applyQuery()
if len(p.items) != 2 {
t.Fatalf("expected 2 matches for 'sonnet', got %d: %#v", len(p.items), p.items)
}
if p.items[0].Value != "sonnet" {
t.Fatalf("exact/prefix match should rank first, got %q", p.items[0].Value)
}
}
// A subsequence query (non-contiguous characters) still matches when no substring does.
func TestPickerFuzzySubsequenceMatches(t *testing.T) {
p := newFuzzyTestPicker([]pickerItem{
{Label: "Claude Sonnet 4.5", Value: "sonnet-4.5"},
{Label: "GPT-5", Value: "gpt-5"},
})
p.query = "snt45" // subsequence of "sonnet 4.5", not a substring
p.applyQuery()
if len(p.items) != 1 || p.items[0].Value != "sonnet-4.5" {
t.Fatalf("subsequence query should match sonnet-4.5, got %#v", p.items)
}
}
// Ranking keeps each group contiguous: a filtered result never splits one group
// into two separate blocks, so group headers still render once per group.
func TestPickerRankingKeepsGroupsContiguous(t *testing.T) {
p := newFuzzyTestPicker([]pickerItem{
{Group: "openai", Label: "gpt-5 mini", Value: "a"},
{Group: "anthropic", Label: "sonnet 4.5", Value: "b"},
{Group: "openai", Label: "gpt-5 codex", Value: "c"},
{Group: "anthropic", Label: "opus 4.5", Value: "d"},
})
p.query = "5" // matches both openai models and the "4.5" anthropic ones
p.applyQuery()
// Every item's group must appear as one contiguous run.
seen := map[string]bool{}
last := ""
for _, item := range p.items {
if item.Group != last {
if seen[item.Group] {
t.Fatalf("group %q split into non-contiguous blocks: %#v", item.Group, p.items)
}
seen[item.Group] = true
last = item.Group
}
}
}
+6
View File
@@ -987,6 +987,12 @@ func isDoctorResultHeading(value string) bool {
func renderErrorRow(row transcriptRow, width int) string {
note := noteBox(row.text, width, zeroTheme.cardErr, zeroTheme.red)
// A recognized failure carries a one-line next step. Render it just below the
// red box in the faint metadata style so it reads as guidance, not more error
// text (and to avoid nesting ANSI styles inside noteBox's per-line red wrap).
if hint := strings.TrimSpace(row.hint); hint != "" {
note += "\n" + fitStyledLine(zeroTheme.faint.Render("→ "+hint), width)
}
return note
}
+7 -7
View File
@@ -224,14 +224,14 @@ func TestCommandCardRowTrimsIndentedActionsLabel(t *testing.T) {
func TestInterimBlockShowsStreamingTextWithCursor(t *testing.T) {
m := limeTestModel()
m.pending = true
m.streamingText = "I'll add a --version flag"
m.streamingText = []byte("I'll add a --version flag")
got := plainRender(t, m.interimBlock(96))
if !strings.Contains(got, "I'll add a --version flag") || !strings.Contains(got, "▌") {
t.Fatalf("interim block = %q, want streamed text with trailing cursor", got)
}
// Before the first delta the block falls back to the liveness spinner.
m.streamingText = ""
m.streamingText = nil
if got := plainRender(t, m.interimBlock(96)); !strings.Contains(got, "Working") {
t.Fatalf("empty interim block = %q, want the liveness label", got)
}
@@ -240,13 +240,13 @@ func TestInterimBlockShowsStreamingTextWithCursor(t *testing.T) {
func TestInterimBlockRendersStreamingMarkdownTable(t *testing.T) {
m := limeTestModel()
m.pending = true
m.streamingText = strings.Join([]string{
m.streamingText = []byte(strings.Join([]string{
"Here's the comparison:",
"",
"| Category | System A | System B |",
"|---|---|---|",
"| **Label** | Alpha | Beta |",
}, "\n")
}, "\n"))
rendered := m.interimBlock(72)
got := plainRender(t, rendered)
@@ -1528,10 +1528,10 @@ func TestCancelRunClearsStreamingText(t *testing.T) {
m := limeTestModel()
m.pending = true
m.activeRunID = 3
m.streamingText = "partial answer from a doomed run"
m.streamingText = []byte("partial answer from a doomed run")
m.cancelRun()
if m.streamingText != "" {
t.Fatalf("cancelRun must clear streamingText, got %q", m.streamingText)
if len(m.streamingText) != 0 {
t.Fatalf("cancelRun must clear streamingText, got %q", string(m.streamingText))
}
}
+31
View File
@@ -0,0 +1,31 @@
package tui
import (
"runtime"
"testing"
)
func TestSessionMatchesWorkspace(t *testing.T) {
cases := []struct {
name string
sessionCwd string
workspace string
wantVisible bool
}{
{"same workspace", "/home/u/proj", "/home/u/proj", true},
{"trailing slash normalizes", "/home/u/proj/", "/home/u/proj", true},
{"different workspace hidden", "/home/u/other", "/home/u/proj", false},
{"session with no cwd stays visible", "", "/home/u/proj", true},
{"unknown current workspace keeps all", "/home/u/other", "", true},
// Casing: matched case-insensitively on Windows (its FS is), case-sensitively
// elsewhere. This exercises the platform-specific branch on each OS's CI.
{"case differs only", "/home/U/Proj", "/home/u/proj", runtime.GOOS == "windows"},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
if got := sessionMatchesWorkspace(tc.sessionCwd, tc.workspace); got != tc.wantVisible {
t.Fatalf("sessionMatchesWorkspace(%q, %q) = %v, want %v", tc.sessionCwd, tc.workspace, got, tc.wantVisible)
}
})
}
}
+96
View File
@@ -0,0 +1,96 @@
package tui
import (
"context"
"testing"
tea "charm.land/bubbletea/v2"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
func TestCommandsRegistered(t *testing.T) {
for _, name := range []string{"/retry", "/edit", "/copy", "/export"} {
if cmd, ok := resolveCommand(name); !ok {
t.Fatalf("%s should be a registered command", name)
} else if cmd.name != name {
t.Fatalf("resolveCommand(%q) = %q", name, cmd.name)
}
}
}
// /edit recalls the last prompt into the composer for editing.
func TestEditRecallsLastPrompt(t *testing.T) {
m := newModel(context.Background(), Options{})
m.lastPrompt = "refactor the parser"
m.input.SetValue("/edit")
updated, _ := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if got := next.composerValue(); got != "refactor the parser" {
t.Fatalf("/edit should recall last prompt into composer, got %q", got)
}
}
// /edit must re-stage the remembered attachments alongside the recalled text, so an
// edited resend of a vision/PDF-backed prompt carries the same image/document
// context instead of silently sending a text-only version.
func TestEditRestagesAttachments(t *testing.T) {
m := newModel(context.Background(), Options{ModelName: "gpt-4.1"})
m.lastPrompt = "describe the diagram"
m.lastImages = []zeroruntime.ImageBlock{{MediaType: "image/png"}}
m.lastImageLabels = []string{"diagram.png"}
m.lastDocuments = []pendingDocument{{label: "spec.pdf", text: "notes"}}
m.input.SetValue("/edit")
updated, _ := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if got := next.composerValue(); got != "describe the diagram" {
t.Fatalf("/edit should recall the prompt text, got %q", got)
}
if len(next.pendingImages) != 1 || len(next.pendingImageLabels) != 1 || next.pendingImageLabels[0] != "diagram.png" {
t.Fatalf("/edit should re-stage the remembered image, got imgs=%d labels=%v",
len(next.pendingImages), next.pendingImageLabels)
}
if len(next.pendingDocuments) != 1 || next.pendingDocuments[0].label != "spec.pdf" {
t.Fatalf("/edit should re-stage the remembered document, got %#v", next.pendingDocuments)
}
}
// /retry with no prior prompt reports that there's nothing to resend rather than
// launching an empty run.
func TestRetryWithoutPriorPromptIsNoOp(t *testing.T) {
m := newModel(context.Background(), Options{})
m.input.SetValue("/retry")
updated, _ := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if !transcriptContains(next.transcript, "no previous prompt") {
t.Fatalf("/retry with no history should note there's nothing to resend, got %#v", next.transcript)
}
}
// /retry must not launch a run while compaction is rewriting session state — the
// same guard a normal prompt has.
func TestRetryBlockedDuringCompaction(t *testing.T) {
m := newModel(context.Background(), Options{})
m.lastPrompt = "do the thing"
m.compactInFlight = true
m.input.SetValue("/retry")
updated, cmd := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if next.pending {
t.Fatal("/retry must not start a run during compaction")
}
if cmd != nil {
t.Fatal("/retry during compaction must not return a run command")
}
if !transcriptContains(next.transcript, "Compaction is running") {
t.Fatalf("/retry during compaction should warn, got %#v", next.transcript)
}
}
+5 -1
View File
@@ -25,7 +25,7 @@ func Run(ctx context.Context, options Options) int {
externalSink := options.RuntimeMessageSink
var program *tea.Program
options.RuntimeMessageSink = func(msg tea.Msg) {
forward := func(msg tea.Msg) {
if externalSink != nil {
externalSink(msg)
}
@@ -33,6 +33,10 @@ func Run(ctx context.Context, options Options) int {
program.Send(msg)
}
}
// Coalesce streamed assistant-text deltas to ~one frame each so a fast provider
// can't drive a full Update→View per token; every other message flushes pending
// text first, keeping order intact.
options.RuntimeMessageSink = newTextCoalescer(forward).send
options.AltScreen = useAltScreen(options)
programOpts := []tea.ProgramOption{
+71 -3
View File
@@ -4,6 +4,8 @@ import (
"encoding/json"
"errors"
"fmt"
"path/filepath"
"runtime"
"strings"
"time"
@@ -80,6 +82,12 @@ func (m model) startNewSession() model {
m.pendingImageLabels = nil
m.pendingDocuments = nil
m.queuedMessage = ""
// The remembered /retry attachment snapshot belongs to the previous session
// too — dropping it keeps a post-/new /retry from re-staging old images or
// documents. lastPrompt is composer history (like inputHistory) and stays.
m.lastImages = nil
m.lastImageLabels = nil
m.lastDocuments = nil
note := "Started a new session."
if previousID != "" {
@@ -220,9 +228,11 @@ func (m model) sessionPrompt(prompt string) string {
func (m model) resolveResumeSession(args string) (*sessions.Metadata, error) {
if strings.EqualFold(args, "latest") {
// Latest *resumable* conversation, so "latest" never lands on a child or
// spec sub-run. An explicit `/resume <id>` below still resolves any session.
latest, err := m.sessionStore.LatestResumable()
// Latest *resumable* conversation IN THIS WORKSPACE, so "latest" never lands
// on a child, a spec sub-run, or a session from another project (matching the
// workspace-scoped picker). An explicit `/resume <id>` below still resolves
// any session regardless of workspace.
latest, err := m.latestResumableInWorkspace()
if err != nil {
return nil, err
}
@@ -321,6 +331,18 @@ func (m model) newSessionPicker() *commandPicker {
now := m.now()
items := make([]pickerItem, 0, len(metas))
for _, meta := range metas {
// Workspace-scoped: hide sessions from other project directories so /resume
// lists this workspace's history, not every project's. Checked BEFORE the
// per-session event read below, so a large global history doesn't pay 50
// full file reads to build one workspace's list. Sessions with no recorded
// Cwd (older runs) stay visible rather than vanishing.
if !sessionMatchesWorkspace(meta.Cwd, m.cwd) {
continue
}
// A zero-event session has nothing to resume — skip it without a file read.
if meta.EventCount == 0 {
continue
}
// Skip empty/failed runs (no assistant output, no tool calls) — e.g. the
// same prompt retried while the model wasn't responding. They have nothing
// to resume and otherwise flood the list with identical rows. Still on disk.
@@ -356,6 +378,52 @@ func (m model) newSessionPicker() *commandPicker {
// resuming: a tool call/result, or a non-user message with real content (not the
// no-output guardrail stop). Empty/failed runs return false and are hidden from
// the picker (they stay on disk). Errors fail open (the session is kept).
// latestResumableInWorkspace returns the most-recently-updated resumable session
// belonging to the current workspace, or nil when none exist. ListResumable is
// ordered latest-first, so the first qualifying match is the latest. It applies
// the SAME filters as newSessionPicker — workspace membership, non-empty
// metadata, and real resumable content — so `/resume latest` never lands on a
// zero-event or empty/failed run the picker would have hidden.
func (m model) latestResumableInWorkspace() (*sessions.Metadata, error) {
metas, err := m.sessionStore.ListResumable()
if err != nil {
return nil, err
}
for i := range metas {
if !sessionMatchesWorkspace(metas[i].Cwd, m.cwd) {
continue
}
if metas[i].EventCount == 0 {
continue
}
if !m.sessionHasResumableContent(metas[i].SessionID) {
continue
}
return &metas[i], nil
}
return nil, nil
}
// sessionMatchesWorkspace reports whether a session recorded in sessionCwd
// belongs to the current workspaceCwd. A session with no recorded Cwd (older
// runs) or an unknown current workspace is kept visible rather than filtered out,
// so the scoping never hides history it can't confidently place elsewhere. On
// Windows the comparison is case-insensitive, since the filesystem is and the
// same workspace can be spelled with different casing (C:\Proj vs c:\proj).
func sessionMatchesWorkspace(sessionCwd, workspaceCwd string) bool {
sessionCwd = strings.TrimSpace(sessionCwd)
workspaceCwd = strings.TrimSpace(workspaceCwd)
if sessionCwd == "" || workspaceCwd == "" {
return true
}
a := filepath.Clean(sessionCwd)
b := filepath.Clean(workspaceCwd)
if runtime.GOOS == "windows" {
return strings.EqualFold(a, b)
}
return a == b
}
func (m model) sessionHasResumableContent(sessionID string) bool {
events, err := m.sessionStore.ReadEvents(sessionID)
if err != nil {
+4 -2
View File
@@ -326,8 +326,10 @@ func (m model) handleCompactCommand(args string) (model, string, tea.Cmd) {
if args == "status" {
return m, m.compactText(false), nil
}
if args != "" {
return m, "Compact\nusage: /compact [status]", nil
// Bare "/compact" already triggers compaction; accept "now" too since that is
// what users reach for when they see the context gauge climbing.
if args != "" && args != "now" {
return m, "Compact\nusage: /compact [status|now]", nil
}
if m.compactInFlight {
return m, m.compactText(true), nil
+3 -3
View File
@@ -7,7 +7,7 @@ import (
func TestStreamingCodeRendersHighlighted(t *testing.T) {
m := model{
streamingText: "```go\nfunc main() {}\n```",
streamingText: []byte("```go\nfunc main() {}\n```"),
pending: true,
}
out := m.interimBlock(80)
@@ -86,7 +86,7 @@ func TestStreamingMarkdownStablePrefixUsesRenderCache(t *testing.T) {
func TestStreamingBuffersOpenFencedCodeBlock(t *testing.T) {
open := model{
streamingText: "Here is the script:\n```python\nfrom datetime import datetime\nprint(datetime.now())",
streamingText: []byte("Here is the script:\n```python\nfrom datetime import datetime\nprint(datetime.now())"),
pending: true,
}
openOut := plainRender(t, open.interimBlock(90))
@@ -98,7 +98,7 @@ func TestStreamingBuffersOpenFencedCodeBlock(t *testing.T) {
}
closed := model{
streamingText: open.streamingText + "\n```",
streamingText: []byte(string(open.streamingText) + "\n```"),
pending: true,
}
closedOut := closed.interimBlock(90)
+1
View File
@@ -35,6 +35,7 @@ type transcriptRow struct {
tool string // tool name, for tool call/result rows
status tools.Status // result status, for tool result rows
detail string // raw multi-line output (e.g. a diff to render as a card)
hint string // one-line actionable hint, rendered faintly below error rows
arg string // secondary argument hint (pattern/command), for tool call rows
runID int // owning run, for tool call rows (0 = rehydrated/unknown)
permission *agent.PermissionEvent
+83
View File
@@ -0,0 +1,83 @@
package tui
import (
"fmt"
"os"
"path/filepath"
"strings"
)
// lastAssistantAnswer returns the text of the most recent assistant row — the
// last final answer if one exists, else the last assistant row of any kind. Empty
// when the conversation has no assistant text yet.
func (m model) lastAssistantAnswer() string {
lastAny := ""
for i := len(m.transcript) - 1; i >= 0; i-- {
row := m.transcript[i]
if row.kind != rowAssistant {
continue
}
if row.final {
return row.text
}
if lastAny == "" {
lastAny = row.text
}
}
return lastAny
}
// plainTranscriptText renders the conversation as a plain, role-prefixed text
// document for /export — the readable content (user prompts, assistant answers,
// system notes, errors), skipping tool-call/permission UI noise.
func (m model) plainTranscriptText() string {
var b strings.Builder
for _, row := range m.transcript {
var prefix string
switch row.kind {
case rowUser:
prefix = "you: "
case rowAssistant:
prefix = "zero: "
case rowSystem:
prefix = "· "
case rowError:
prefix = "error: "
default:
continue
}
text := strings.TrimRight(row.text, "\n")
if strings.TrimSpace(text) == "" {
continue
}
b.WriteString(prefix)
b.WriteString(text)
b.WriteString("\n\n")
}
return b.String()
}
// handleExportCommand writes the transcript to a file and returns a status
// message. With no argument it derives a timestamped filename in the workspace;
// a relative path is resolved against the workspace root.
func (m model) handleExportCommand(args string) string {
body := m.plainTranscriptText()
if strings.TrimSpace(body) == "" {
return "Export\nnothing to export yet."
}
path := strings.TrimSpace(args)
if path == "" {
stamp := m.now().Format("20060102-150405")
path = fmt.Sprintf("zero-transcript-%s.txt", stamp)
}
if !filepath.IsAbs(path) && m.cwd != "" {
path = filepath.Join(m.cwd, path)
}
// 0o600: a transcript can include tool/bash output that echoed secrets, so it
// must not be world/group-readable on a shared machine.
if err := os.WriteFile(path, []byte(body), 0o600); err != nil {
return "Export\nfailed to write " + path + ": " + err.Error()
}
return "Export\nwrote transcript to " + path
}
+103
View File
@@ -0,0 +1,103 @@
package tui
import (
"os"
"path/filepath"
"runtime"
"strings"
"testing"
"time"
)
func TestLastAssistantAnswerPrefersFinal(t *testing.T) {
m := model{transcript: []transcriptRow{
{kind: rowUser, text: "hi"},
{kind: rowAssistant, text: "interim narration"},
{kind: rowAssistant, text: "the final answer", final: true},
{kind: rowSystem, text: "worked for 3s"},
}}
if got := m.lastAssistantAnswer(); got != "the final answer" {
t.Fatalf("lastAssistantAnswer = %q, want the final answer", got)
}
// With no final row, falls back to the most recent assistant row.
m2 := model{transcript: []transcriptRow{
{kind: rowAssistant, text: "first"},
{kind: rowAssistant, text: "second"},
}}
if got := m2.lastAssistantAnswer(); got != "second" {
t.Fatalf("fallback lastAssistantAnswer = %q, want second", got)
}
// Empty when there's no assistant text.
if got := (model{transcript: []transcriptRow{{kind: rowUser, text: "hi"}}}).lastAssistantAnswer(); got != "" {
t.Fatalf("expected empty answer, got %q", got)
}
}
func TestPlainTranscriptTextSkipsNoise(t *testing.T) {
m := model{transcript: []transcriptRow{
{kind: rowUser, text: "add a flag"},
{kind: rowToolCall, tool: "bash", text: "go build"},
{kind: rowToolResult, tool: "bash", text: "ok"},
{kind: rowAssistant, text: "Done — added --version.", final: true},
{kind: rowError, text: "provider error: boom"},
}}
out := m.plainTranscriptText()
if !strings.Contains(out, "you: add a flag") || !strings.Contains(out, "zero: Done — added --version.") {
t.Fatalf("export missing conversation text:\n%s", out)
}
if !strings.Contains(out, "error: provider error: boom") {
t.Fatalf("export should include error rows:\n%s", out)
}
if strings.Contains(out, "go build") || strings.Contains(out, "bash") {
t.Fatalf("export should skip tool-call noise:\n%s", out)
}
}
func TestHandleExportCommandWritesFile(t *testing.T) {
dir := t.TempDir()
m := model{
cwd: dir,
now: func() time.Time { return time.Date(2026, 7, 4, 9, 30, 0, 0, time.UTC) },
transcript: []transcriptRow{
{kind: rowUser, text: "hello"},
{kind: rowAssistant, text: "hi there", final: true},
},
}
// Explicit relative path resolves against cwd.
msg := m.handleExportCommand("out.txt")
if !strings.Contains(msg, "wrote transcript to") {
t.Fatalf("export status = %q", msg)
}
outPath := filepath.Join(dir, "out.txt")
data, err := os.ReadFile(outPath)
if err != nil {
t.Fatalf("reading exported file: %v", err)
}
if !strings.Contains(string(data), "you: hello") || !strings.Contains(string(data), "zero: hi there") {
t.Fatalf("exported content = %q", string(data))
}
// The transcript may contain secrets echoed in tool output; it must not be
// world/group-readable. Unix perm bits are meaningless on Windows (os.WriteFile
// ignores them and Stat reports 0666), so only assert the mode where it applies.
if runtime.GOOS != "windows" {
if info, err := os.Stat(outPath); err != nil {
t.Fatalf("stat exported file: %v", err)
} else if perm := info.Mode().Perm(); perm != 0o600 {
t.Fatalf("exported file mode = %o, want 0600", perm)
}
}
// No-arg export derives a timestamped filename in cwd.
if msg := m.handleExportCommand(""); !strings.Contains(msg, "zero-transcript-20260704-093000.txt") {
t.Fatalf("default export status = %q", msg)
}
// Nothing to export on an empty conversation.
empty := model{cwd: dir, now: m.now}
if msg := empty.handleExportCommand(""); !strings.Contains(msg, "nothing to export") {
t.Fatalf("empty export status = %q", msg)
}
}
+1 -1
View File
@@ -366,7 +366,7 @@ func (m model) transcriptBodyItems(width int, emptyOverlay string) []transcriptB
}
if m.pending {
pendingShowsAssistantText := m.pendingPermission == nil && m.pendingAskUser == nil && strings.TrimSpace(m.streamingText) != ""
pendingShowsAssistantText := m.pendingPermission == nil && m.pendingAskUser == nil && strings.TrimSpace(m.streamingTextString()) != ""
if pendingShowsAssistantText && havePreviousKind && shouldRuleBeforeTurn(previousKind, rowAssistant) {
items = append(items, transcriptRuleBodyItem(contentWidth, gutter))
} else {
+23
View File
@@ -0,0 +1,23 @@
package tui
import (
"strings"
"testing"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
func TestVisionDropWarning(t *testing.T) {
// No staged images: never warns, regardless of model.
if got := (model{}).visionDropWarning(); got != "" {
t.Fatalf("no images should give no warning, got %q", got)
}
// Staged images + a model with no vision support (empty model name qualifies):
// warn immediately, naming the count.
withImg := model{pendingImages: make([]zeroruntime.ImageBlock, 2)}
warn := withImg.visionDropWarning()
if !strings.Contains(warn, "will be dropped") || !strings.Contains(warn, "2 staged") {
t.Fatalf("expected a drop warning naming the count, got %q", warn)
}
}
+4 -4
View File
@@ -191,7 +191,7 @@ func TestInterimBlockShowsWorkingLineWithStreamedText(t *testing.T) {
base := time.Date(2026, 6, 14, 10, 0, 0, 0, time.UTC)
m.now = func() time.Time { return base.Add(12 * time.Second) }
m.turnStartedAt = base
m.streamingText = "partial answer so far"
m.streamingText = []byte("partial answer so far")
got := plainRender(t, m.interimBlock(96))
if !strings.Contains(got, "partial answer so far") {
@@ -241,7 +241,7 @@ func TestWorkingTokenIndicatorAccumulatesAcrossSegmentClears(t *testing.T) {
// Simulate the segment boundary that clears the live buffers, then stream
// answer text in the next segment.
m.streamingReasoning = ""
m.streamingText = ""
m.streamingText = nil
updated, _ = m.Update(agentTextMsg{runID: rid, delta: strings.Repeat("b", 40)})
m = updated.(model)
@@ -288,7 +288,7 @@ func TestInterimBlockShowsReasoningPreviewWhileThinking(t *testing.T) {
m.now = func() time.Time { return base.Add(90 * time.Second) }
m.turnStartedAt = base
m.streamingReasoning = "analyzing the layout\nthe patch was corrupt so re-planning the css edits"
m.streamingText = "" // thinking phase: no answer yet
m.streamingText = nil // thinking phase: no answer yet
got := plainRender(t, m.interimBlock(96))
if !strings.Contains(got, "re-planning the css edits") {
@@ -307,7 +307,7 @@ func TestInterimBlockNoPreviewWhenReasoningExpanded(t *testing.T) {
m.now = func() time.Time { return time.Date(2026, 6, 18, 23, 0, 0, 0, time.UTC) }
m.streamingReasoningExpanded = true
m.streamingReasoning = "only line of reasoning here"
m.streamingText = ""
m.streamingText = nil
got := plainRender(t, m.interimBlock(96))
if strings.Count(got, "only line of reasoning here") != 1 {
t.Fatalf("reasoning should appear exactly once when expanded (no preview dup):\n%s", got)