Compare commits

..
Author SHA1 Message Date
anandh8x daabe11164 fix(providers): send stable session headers to OpenCode 2026-09-15 18:02:59 +05:30
112 changed files with 686 additions and 9728 deletions

No files matched your search

+2 -9
View File
@@ -6,7 +6,7 @@
<p align="center">
<a href="LICENSE"><img alt="license" src="https://img.shields.io/badge/license-MIT-blue"></a>
<img alt="Go 1.26.9+" src="https://img.shields.io/badge/Go-1.26.9+-00ADD8?logo=go&logoColor=white">
<img alt="Go 1.26.6+" src="https://img.shields.io/badge/Go-1.26.6+-00ADD8?logo=go&logoColor=white">
<img alt="25+ providers" src="https://img.shields.io/badge/providers-25+-34E2EA">
<a href="https://discord.gg/CaQDS6wdFn"><img alt="Discord" src="https://img.shields.io/badge/Discord-join-5865F2?logo=discord&logoColor=white"></a>
<br>
@@ -76,7 +76,7 @@ irm https://raw.githubusercontent.com/Gitlawb/zero/main/scripts/install.ps1 | ie
### From source
Source builds require Go 1.26.9+.
Source builds require Go 1.26.6+.
```bash
git clone https://github.com/Gitlawb/zero.git
@@ -331,13 +331,6 @@ zero update --check check for newer releases
zero upgrade download, verify, and install the latest release
```
Cron keeps the newest 1,000 run outcomes per job in `runs.jsonl`. Each new run
atomically replaces the history with the retained tail; existing larger histories
are trimmed on their next run. Archive the file before that run if you need older
outcomes. This retention does not delete session directories or reset fire counts.
Malformed records and legacy lines of 1 MiB or more are skipped when reading or
compacting history; new records of that size are rejected.
## Extending Zero
### Project and personal instructions
+2 -2
View File
@@ -6,7 +6,7 @@
<p align="center">
<a href="LICENSE"><img alt="license" src="https://img.shields.io/badge/license-MIT-blue"></a>
<img alt="Go 1.26.9+" src="https://img.shields.io/badge/Go-1.26.9+-00ADD8?logo=go&logoColor=white">
<img alt="Go 1.26.6+" src="https://img.shields.io/badge/Go-1.26.6+-00ADD8?logo=go&logoColor=white">
<img alt="25+ providers" src="https://img.shields.io/badge/providers-25+-34E2EA">
<a href="https://discord.gg/CaQDS6wdFn"><img alt="Discord" src="https://img.shields.io/badge/Discord-join-5865F2?logo=discord&logoColor=white"></a>
<br>
@@ -57,7 +57,7 @@ irm https://raw.githubusercontent.com/Gitlawb/zero/main/scripts/install.ps1 | ie
### 从源码构建
源码构建需要 Go 1.26.9+。
源码构建需要 Go 1.26.6+。
```bash
git clone https://github.com/Gitlawb/zero.git
+2 -2
View File
@@ -153,7 +153,7 @@ describe a version that is no longer on disk. Prefer `mise upgrade`.
go install github.com/Gitlawb/zero/cmd/zero@latest
```
This builds from source, so it needs Go 1.26.9+ and it does not go through the
This builds from source, so it needs Go 1.26.6+ and it does not go through the
release archives. Two consequences worth knowing before you pick it:
- On Linux you also need the sandbox helper, which is a separate binary and is
@@ -180,7 +180,7 @@ Build a local binary:
go build -o zero ./cmd/zero
```
Source builds require Go 1.26.9+.
Source builds require Go 1.26.6+.
### Sandbox Helpers For Source Builds
+1 -1
View File
@@ -1,6 +1,6 @@
module github.com/Gitlawb/zero
go 1.26.9
go 1.26.6
require (
charm.land/bubbles/v2 v2.1.1
+5 -74
View File
@@ -15,7 +15,6 @@ import (
"github.com/Gitlawb/zero/internal/agent"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/modelregistry"
"github.com/Gitlawb/zero/internal/providercatalog"
"github.com/Gitlawb/zero/internal/providermodelcatalog"
"github.com/Gitlawb/zero/internal/providermodeldiscovery"
@@ -107,14 +106,6 @@ type acpSession struct {
restrictModels bool
cancel context.CancelFunc
history []turnRecord
// readOnly marks a session published by session/load for a kind that is
// not resumable (a child, side or spec sub-run). Load still replays its
// transcript, which is the render-only use a client needs, but a prompt
// would continue the sub-run as a standalone conversation, which every
// other resume path refuses (sessions.IsResumableKind). Set once when the
// session is first registered and never changed, so it is read without mu.
readOnly bool
}
// NewAgent builds the ACP server and registers its method handlers on conn.
@@ -213,7 +204,7 @@ func (a *Agent) handleSessionNew(ctx context.Context, params json.RawMessage) (a
if err != nil {
return nil, RPCError(codeInternalError, "create session: "+err.Error())
}
sess, _ := a.registerSession(meta.SessionID, root, nil, model, models, restrictModels, false)
sess, _ := a.registerSession(meta.SessionID, root, nil, model, models, restrictModels)
return NewSessionResult{
SessionID: sess.id,
ConfigOptions: a.configOptions(sess),
@@ -320,10 +311,7 @@ func (a *Agent) activatePersistedSession(ctx context.Context, p LoadSessionParam
models = append(models, SessionConfigOptionValue{Value: persistedModel, Name: persistedModel})
}
}
// Resume has already refused a non-resumable kind above, so only load can
// reach here with one, and it publishes that session for viewing only.
readOnly := !sessions.IsResumableKind(meta.SessionKind)
sess, existed := a.registerSession(meta.SessionID, root, history, model, models, restrictModels, readOnly)
sess, existed := a.registerSession(meta.SessionID, root, history, model, models, restrictModels)
if existed {
_, messages, historyWarning, historyErr = a.refreshSessionHistory(sess, requireHistoryLog)
if operation == persistedSessionResume && historyErr != nil {
@@ -445,9 +433,6 @@ func (a *Agent) handleSessionPrompt(ctx context.Context, params json.RawMessage)
if sess == nil {
return nil, RPCError(codeInvalidParams, "unknown session: "+p.SessionID)
}
if sess.readOnly {
return nil, RPCError(codeInvalidParams, "session is not resumable, so it was loaded for viewing only and cannot be prompted: "+p.SessionID)
}
// Serialize turns for this session so two prompts can't interleave history or
// fight over the single cancel slot. session/cancel still works concurrently
@@ -523,30 +508,6 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string,
}
queue(messageEvent("user", userText))
visionCache := make(map[string]bool)
var visionCacheMu sync.Mutex
supportsVision := func(modelID string) bool {
modelID = strings.TrimSpace(modelID)
if modelID == "" {
return false
}
visionCacheMu.Lock()
defer visionCacheMu.Unlock()
if cached, ok := visionCache[modelID]; ok {
return cached
}
supported := a.modelSupportsVision(ctx, resolved.Provider, modelID)
visionCache[modelID] = supported
return supported
}
effectivePrompt := userText
if len(images) > 0 && !supportsVision(resolved.Provider.Model) {
msg := fmt.Sprintf("Model %s does not support image input; ignoring %d prompt image(s).", resolved.Provider.Model, len(images))
note.text("[zero] " + msg + "\n\n")
effectivePrompt = fmt.Sprintf("[Note: %s]\n\n%s", msg, userText)
images = nil
}
opts := agent.Options{
Cwd: sess.cwd,
SessionID: sess.id,
@@ -558,7 +519,6 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string,
DeferThreshold: workspace.DeferThreshold,
MaxTurns: resolved.MaxTurns,
Images: images,
SupportsVision: supportsVision,
OnText: note.text,
OnReasoning: note.thought,
OnToolCall: func(call agent.ToolCall) {
@@ -577,7 +537,7 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string,
},
}
agentPrompt := buildPrompt(sess.snapshotHistory(), effectivePrompt)
agentPrompt := buildPrompt(sess.snapshotHistory(), userText)
result, runErr := a.deps.RunAgent(ctx, agentPrompt, provider, opts)
if result.FinalAnswer != "" {
queue(messageEvent("assistant", result.FinalAnswer))
@@ -1194,13 +1154,13 @@ func promptImages(blocks []ContentBlock) []zeroruntime.ImageBlock {
// The bool reports that case so persisted activation can re-read and apply disk
// history under turnMu. history is set BEFORE first publication so no concurrent
// prompt can read a half-initialized session.
func (a *Agent) registerSession(id, cwd string, history []turnRecord, model string, models []SessionConfigOptionValue, restrictModels bool, readOnly bool) (*acpSession, bool) {
func (a *Agent) registerSession(id, cwd string, history []turnRecord, model string, models []SessionConfigOptionValue, restrictModels bool) (*acpSession, bool) {
a.mu.Lock()
defer a.mu.Unlock()
if existing := a.sessions[id]; existing != nil {
return existing, true
}
sess := &acpSession{id: id, cwd: cwd, mode: agent.PermissionModeAuto, model: model, models: models, restrictModels: restrictModels, history: history, readOnly: readOnly}
sess := &acpSession{id: id, cwd: cwd, mode: agent.PermissionModeAuto, model: model, models: models, restrictModels: restrictModels, history: history}
a.sessions[id] = sess
return sess, false
}
@@ -1345,32 +1305,3 @@ func sameWorkspace(left, right string) bool {
}
return os.SameFile(leftInfo, rightInfo)
}
func (a *Agent) modelSupportsVision(ctx context.Context, profile config.ProviderProfile, modelID string) bool {
trimmed := strings.TrimSpace(modelID)
if trimmed == "" {
return false
}
if a.deps.DiscoverModels != nil {
if discovered, err := a.deps.DiscoverModels(ctx, profile); err == nil {
for _, dm := range discovered {
if strings.EqualFold(strings.TrimSpace(dm.ID), trimmed) {
if len(dm.InputModalities) > 0 {
for _, mod := range dm.InputModalities {
if strings.EqualFold(strings.TrimSpace(mod), "image") {
return true
}
}
return false
}
break
}
}
}
}
reg, _ := modelregistry.DefaultRegistry()
if entry, known := reg.Resolve(trimmed); known {
return entry.Supports(modelregistry.ModelCapabilityVision)
}
return modelregistry.SupportsVision(reg, trimmed)
}
-320
View File
@@ -1,9 +1,7 @@
package acp
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
@@ -849,48 +847,6 @@ func TestACPRunTurnClosesWorkspaceAfterSuccessAndFailure(t *testing.T) {
}
}
func TestACPWiresSupportsVision(t *testing.T) {
deps := testDeps(t)
var captured agent.Options
deps.RunAgent = func(_ context.Context, _ string, _ zeroruntime.Provider, opts agent.Options) (agent.Result, error) {
captured = opts
return agent.Result{FinalAnswer: "done"}, nil
}
deps.DiscoverModels = func(_ context.Context, _ config.ProviderProfile) ([]providermodeldiscovery.Model, error) {
return []providermodeldiscovery.Model{
{ID: "custom-vision", InputModalities: []string{"text", "image"}},
{ID: "custom-text", InputModalities: []string{"text"}},
}, nil
}
h := newHarness(t, deps)
defer h.stop()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
var newRes NewSessionResult
if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir(), McpServers: []McpServer{}}, &newRes); err != nil {
t.Fatalf("session/new: %v", err)
}
var promptRes PromptResult
if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{
SessionID: newRes.SessionID,
Prompt: []ContentBlock{
{Type: "text", Text: "hello"},
},
}, &promptRes); err != nil {
t.Fatalf("session/prompt: %v", err)
}
if captured.SupportsVision == nil {
t.Fatal("SupportsVision was not wired into agent.Options")
}
if !captured.SupportsVision("custom-vision") {
t.Fatal("SupportsVision(custom-vision) = false, want true")
}
if captured.SupportsVision("custom-text") {
t.Fatal("SupportsVision(custom-text) = true, want false")
}
}
// TestACPRejectsInvalidCwd confirms session/new fails when the workspace root
// resolver rejects the client cwd (e.g. filesystem root).
func TestACPRejectsInvalidCwd(t *testing.T) {
@@ -1295,85 +1251,6 @@ func TestACPResumeAppliesStandaloneSessionKindPolicy(t *testing.T) {
}, &LoadSessionResult{}); err != nil {
t.Fatalf("session/load should retain render-only access: %v", err)
}
// The check above only showed the refused resume never registered the
// session. Load does register it, so the prompt that matters is the one
// AFTER the load: render-only means the transcript can be read and the
// sub-run cannot be continued.
err = h.client.Call(ctx, MethodSessionPrompt, PromptParams{
SessionID: created.SessionID, Prompt: []ContentBlock{TextBlock("must stay closed after load")},
}, &PromptResult{})
if !errors.As(err, &rpcErr) || rpcErr.Code != codeInvalidParams {
t.Fatalf("session/prompt after load of a %s session = %v, want invalid params", tc.name, err)
}
})
}
}
// LOAD PUBLISHES A SUB-RUN FOR VIEWING, NOT FOR CONTINUING. Every other resume
// path (the TUI, exec sessions, the exec CLI, and ACP's own session/resume)
// refuses a kind sessions.IsResumableKind rejects. session/load has a real
// render-only use, a client rebuilding a transcript it can scroll, so it keeps
// loading those sessions; what it must not do is publish one that accepts a
// prompt, because the prompt would carry on a child, side or spec sub-run as a
// standalone conversation. Checked in both directions, so a gate that refused
// everything would fail on the resumable kinds.
func TestACPLoadPublishesSubRunsReadOnly(t *testing.T) {
for _, tc := range []struct {
name string
kind sessions.SessionKind
readOnly bool
}{
{name: "regular", kind: ""},
{name: "fork", kind: sessions.SessionKindFork},
{name: "child", kind: sessions.SessionKindChild, readOnly: true},
{name: "side", kind: sessions.SessionKindSide, readOnly: true},
{name: "spec draft", kind: sessions.SessionKindSpecDraft, readOnly: true},
{name: "spec impl", kind: sessions.SessionKindSpecImpl, readOnly: true},
} {
t.Run(tc.name, func(t *testing.T) {
if got := !sessions.IsResumableKind(tc.kind); got != tc.readOnly {
t.Fatalf("SETUP INVALID: IsResumableKind(%q) disagrees with this table (readOnly=%v), so the cases no longer mean what they say", tc.kind, tc.readOnly)
}
deps := testDeps(t)
workspace := t.TempDir()
created, err := deps.Store.Create(sessions.CreateInput{
SessionID: "load-" + strings.ReplaceAll(tc.name, " ", "-"), Cwd: workspace, SessionKind: tc.kind,
})
if err != nil {
t.Fatal(err)
}
h := newHarness(t, deps)
defer h.stop()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := h.client.Call(ctx, MethodSessionLoad, LoadSessionParams{
SessionID: created.SessionID, Cwd: workspace,
}, &LoadSessionResult{}); err != nil {
t.Fatalf("session/load of a %s session: %v", tc.name, err)
}
err = h.client.Call(ctx, MethodSessionPrompt, PromptParams{
SessionID: created.SessionID, Prompt: []ContentBlock{TextBlock("continue")},
}, &PromptResult{})
// A resumable session has to be genuinely promptable, not merely "not
// refused as view-only": any other failure would hide a broken
// writable path behind a passing test.
if !tc.readOnly {
if err != nil {
t.Fatalf("session/prompt after load of a %s session = %v, want success", tc.name, err)
}
return
}
// The decision is the RPC code; the message only distinguishes the
// view-only refusal from other invalid-params errors such as an
// unknown session id.
var rpcErr *rpcError
if !errors.As(err, &rpcErr) || rpcErr.Code != codeInvalidParams {
t.Fatalf("session/prompt after load of a %s session = %v, want invalid params", tc.name, err)
}
if !strings.Contains(rpcErr.Message, "loaded for viewing only") {
t.Fatalf("session/prompt after load of a %s session was refused for another reason: %v", tc.name, err)
}
})
}
}
@@ -2541,200 +2418,3 @@ func payloadString(payload any, key string) string {
value, _ := decoded[key].(string)
return value
}
func TestACPPromptUnsupportedModelDropsImagesAndNotifies(t *testing.T) {
deps := testDeps(t)
deps.DiscoverModels = func(ctx context.Context, p config.ProviderProfile) ([]providermodeldiscovery.Model, error) {
return []providermodeldiscovery.Model{
{ID: "fake-model", InputModalities: []string{"text"}},
}, nil
}
var capturedOpts agent.Options
var capturedPrompt string
deps.RunAgent = func(ctx context.Context, prompt string, provider zeroruntime.Provider, opts agent.Options) (agent.Result, error) {
capturedPrompt = prompt
capturedOpts = opts
return agent.Result{FinalAnswer: "Answered without images"}, nil
}
h := newHarness(t, deps)
defer h.stop()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
var initRes InitializeResult
if err := h.client.Call(ctx, MethodInitialize, InitializeParams{ProtocolVersion: ProtocolVersion}, &initRes); err != nil {
t.Fatalf("initialize: %v", err)
}
var newRes NewSessionResult
if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir(), McpServers: []McpServer{}}, &newRes); err != nil {
t.Fatalf("session/new: %v", err)
}
rawPng := "\x89PNG\r\n\x1a\nfake-image-bytes"
b64Png := base64.StdEncoding.EncodeToString([]byte(rawPng))
var promptRes PromptResult
if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{
SessionID: newRes.SessionID,
Prompt: []ContentBlock{
TextBlock("look at this image"),
ImageBlock(b64Png, "image/png"),
},
}, &promptRes); err != nil {
t.Fatalf("session/prompt: %v", err)
}
if promptRes.StopReason != StopEndTurn {
t.Fatalf("stopReason = %q, want %q", promptRes.StopReason, StopEndTurn)
}
got := drainTextUntil(t, h.updates, func(text string) bool {
return strings.Contains(text, "does not support image input")
})
if !strings.Contains(got, "Model fake-model does not support image input; ignoring 1 prompt image(s).") {
t.Fatalf("streamed text = %q, want drop notice", got)
}
if len(capturedOpts.Images) != 0 {
t.Fatalf("captured %d images, want 0 (images should be withheld)", len(capturedOpts.Images))
}
if !strings.Contains(capturedPrompt, "[Note: Model fake-model does not support image input; ignoring 1 prompt image(s).]") {
t.Fatalf("captured prompt = %q, want prompt note", capturedPrompt)
}
}
func TestACPSessionPrompt_SupportsVisionMemoizedPerRun(t *testing.T) {
deps := testDeps(t)
var discoverCount int
var discoverMu sync.Mutex
deps.DiscoverModels = func(ctx context.Context, p config.ProviderProfile) ([]providermodeldiscovery.Model, error) {
discoverMu.Lock()
discoverCount++
discoverMu.Unlock()
return []providermodeldiscovery.Model{
{ID: "fake-model", InputModalities: []string{"text", "image"}},
}, nil
}
deps.RunAgent = func(ctx context.Context, prompt string, provider zeroruntime.Provider, opts agent.Options) (agent.Result, error) {
for i := 0; i < 5; i++ {
if !opts.SupportsVision(opts.Model) {
t.Error("expected SupportsVision to report true")
}
}
return agent.Result{FinalAnswer: "Done"}, nil
}
h := newHarness(t, deps)
defer h.stop()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
var initRes InitializeResult
if err := h.client.Call(ctx, MethodInitialize, InitializeParams{ProtocolVersion: ProtocolVersion}, &initRes); err != nil {
t.Fatalf("initialize: %v", err)
}
var newRes NewSessionResult
if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir(), McpServers: []McpServer{}}, &newRes); err != nil {
t.Fatalf("session/new: %v", err)
}
discoverMu.Lock()
discoverCount = 0
discoverMu.Unlock()
rawPng := "\x89PNG\r\n\x1a\nfake-image-bytes"
b64Png := base64.StdEncoding.EncodeToString([]byte(rawPng))
var promptRes PromptResult
if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{
SessionID: newRes.SessionID,
Prompt: []ContentBlock{
TextBlock("look at this"),
ImageBlock(b64Png, "image/png"),
},
}, &promptRes); err != nil {
t.Fatalf("session/prompt: %v", err)
}
if promptRes.StopReason != StopEndTurn {
t.Fatalf("stopReason = %q, want %q", promptRes.StopReason, StopEndTurn)
}
discoverMu.Lock()
count := discoverCount
discoverMu.Unlock()
if count != 1 {
t.Fatalf("DiscoverModels was called %d times, want exactly 1 call (memoized per run)", count)
}
}
func TestACPEndToEndImagePromptSupportedModel(t *testing.T) {
deps := testDeps(t)
deps.DiscoverModels = func(ctx context.Context, p config.ProviderProfile) ([]providermodeldiscovery.Model, error) {
return []providermodeldiscovery.Model{
{ID: "discovered-vision-model", InputModalities: []string{"text", "image"}},
}, nil
}
deps.ResolveConfig = func(_ string, o config.Overrides) (config.ResolvedConfig, error) {
model := "discovered-vision-model"
if o.Provider.Model != "" {
model = o.Provider.Model
}
return config.ResolvedConfig{
Provider: config.ProviderProfile{Name: "fake", Model: model},
MaxTurns: 4,
}, nil
}
var capturedOpts agent.Options
var capturedPrompt string
deps.RunAgent = func(ctx context.Context, prompt string, provider zeroruntime.Provider, opts agent.Options) (agent.Result, error) {
capturedPrompt = prompt
capturedOpts = opts
return agent.Result{FinalAnswer: "I can see the image!"}, nil
}
h := newHarness(t, deps)
defer h.stop()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
var initRes InitializeResult
if err := h.client.Call(ctx, MethodInitialize, InitializeParams{ProtocolVersion: ProtocolVersion}, &initRes); err != nil {
t.Fatalf("initialize: %v", err)
}
var newRes NewSessionResult
if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir(), McpServers: []McpServer{}}, &newRes); err != nil {
t.Fatalf("session/new: %v", err)
}
rawPng := []byte("\x89PNG\r\n\x1a\nreal-png-bytes")
b64Png := base64.StdEncoding.EncodeToString(rawPng)
var promptRes PromptResult
if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{
SessionID: newRes.SessionID,
Prompt: []ContentBlock{
TextBlock("analyze this image"),
ImageBlock(b64Png, "image/png"),
},
}, &promptRes); err != nil {
t.Fatalf("session/prompt: %v", err)
}
if promptRes.StopReason != StopEndTurn {
t.Fatalf("stopReason = %q, want %q", promptRes.StopReason, StopEndTurn)
}
if len(capturedOpts.Images) != 1 {
t.Fatalf("captured %d images, want 1", len(capturedOpts.Images))
}
if capturedOpts.Images[0].MediaType != "image/png" {
t.Fatalf("mediaType = %q, want image/png", capturedOpts.Images[0].MediaType)
}
if !bytes.Equal(capturedOpts.Images[0].Data, rawPng) {
t.Fatalf("image bytes mismatch: got %v, want %v", capturedOpts.Images[0].Data, rawPng)
}
if strings.Contains(capturedPrompt, "does not support image input") {
t.Fatalf("prompt unexpectedly contains refusal note: %q", capturedPrompt)
}
}
+31 -103
View File
@@ -13,7 +13,6 @@ import (
"github.com/Gitlawb/zero/internal/execution"
"github.com/Gitlawb/zero/internal/hooks"
"github.com/Gitlawb/zero/internal/modelregistry"
"github.com/Gitlawb/zero/internal/redaction"
"github.com/Gitlawb/zero/internal/sandbox"
"github.com/Gitlawb/zero/internal/streamjson"
@@ -26,18 +25,14 @@ const maxTurnsAnswer = "Agent reached maximum number of turns without a final an
const maxTurnsFinalAnswerPrompt = "You have reached the tool-turn limit. Do not call tools. Give a concise final answer now: summarize what you completed, what you found, and any remaining blockers."
// maxStreamStallRetries bounds how many times a turn that timed out (idle/stall)
// OR hit a mid-stream transport abort (#973) with no answer text yet is re-issued
// on a fresh connection before giving up. Incomplete tool-call previews are
// allowed through the gate (they are never dispatched before the error return).
// The safety claim is only that visible answer prose is not duplicated on retry —
// not that every partial stream shape is excluded. Covers a stalled/dead pooled
// connection or a socket abort before tool dispatch.
// WITH NO OUTPUT yet is re-issued on a fresh connection before giving up. Only
// the no-output case is retried (a partial turn would duplicate), so this is a
// safe recovery for a stalled/dead pooled connection.
//
// Set to 1 (not 2): each attempt can itself idle for the full stream timeout
// (~5min) before the stall is even detected, so 2 retries left an interactive
// session frozen for ~15min. One retry keeps the common single-hiccup recovery
// while bounding the worst case to ~2× the idle timeout. Mid-stream transport
// aborts share this bound (do not raise it for #973).
// while bounding the worst case to ~2× the idle timeout.
const maxStreamStallRetries = 1
const (
@@ -379,11 +374,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
rec.StampFirstToken()
}
if onText != nil {
// Zero-length chunks are not committed answer prose; counting
// them would block an eligible mid-stream abort retry.
if s != "" {
forwardedVisibleText = true
}
forwardedVisibleText = true
onText(s)
}
}
@@ -459,14 +450,9 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
result.Messages = copyMessages(messages)
return result, ctx.Err()
}
// A stream idle/stall timeout OR mid-stream transport abort (#973 — wsarecv /
// connection reset / unexpected EOF / connection closed / broken pipe) is
// safely re-issued when the turn committed NO answer text — no forwarded
// visible prose (forwardedVisibleText) and no collected final text
// (collected.Text). The abort gate matches midStreamAbortNeedles only, not
// the full shouldReconnect list: connect-phase timeouts and connection
// refused are NOT retried here (a slow healthy server must not cost a
// second prefill). This covers three cases:
// A stream idle/stall timeout is safely re-issued when the turn committed NO
// answer text — no forwarded visible prose (forwardedVisibleText) and no
// collected final text (collected.Text). This covers two cases:
// 1. Nothing streamed at all before the connection died (the original
// macOS stale-pooled-connection hang past the response-header timeout).
// 2. The model streamed transient reasoning and began a tool call (e.g. a
@@ -477,27 +463,13 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
// re-render is transient previews, not duplicated answer text. This is
// why the gate no longer excludes collected.ToolCalls: an incomplete
// tool call from a timed-out stream is discard-and-retry, not output.
// 3. Mid-stream socket abort after connect succeeded (Windows WSAECONNABORTED /
// wsarecv, connection reset by peer). Same no-tool-executed safety: the
// error returns before tool dispatch, so a bounded retry is safe.
// A turn that forwarded real prose is NOT retried (it would duplicate visible
// answer text) and falls through to the error return below. Capped +
// exponential backoff, with a user-visible notice per attempt.
for attempt := 1; attempt <= maxStreamStallRetries &&
(isStreamTimeoutError(collected.Error) || isMidStreamTransportAbort(collected.Error)) &&
!forwardedVisibleText &&
isStreamTimeoutError(collected.Error) && !forwardedVisibleText &&
collected.Text == ""; attempt++ {
var notify reconnectNotifier
if isStreamTimeoutError(collected.Error) {
notify = stallRetryNoticeFor(options)
} else {
notify = reconnectNoticeFor(options)
// Count this post-connect reissue even when the replacement
// connect succeeds immediately. streamWithReconnect only
// increments reconnect_count after a failed connect attempt.
options.Trace.Counter(trace.CounterReconnectCount, 1)
}
if notify != nil {
if notify := stallRetryNoticeFor(options); notify != nil {
notify(attempt, maxStreamStallRetries)
}
if err := sleepWithContext(ctx, backoffFor(attempt)); err != nil {
@@ -510,22 +482,12 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
retryStream, retryErr := streamWithReconnect(ctx, provider, retryRequest, reconnectNoticeFor(options))
if retryErr != nil {
result.Messages = copyMessages(messages)
if ctx.Err() != nil {
return result, ctx.Err()
}
return result, retryErr
}
stallGenSpan := options.Trace.Span(trace.SpanGeneration)
collected = zeroruntime.CollectStreamWithOptions(ctx, retryStream, forwardingOpts)
stallGenSpan.End()
}
// Recheck ctx after a retried CollectStream: helpers.go stores
// ctx.Err().Error() in collected.Error, and errors.New of that string
// drops the context.Canceled sentinel ACP branches on.
if ctx.Err() != nil {
result.Messages = copyMessages(messages)
return result, ctx.Err()
}
if collected.Error != "" {
// Route a reissued stream's non-stall error through the SAME recovery as
// the initial stream (image-rejection wrapping / context-limit compaction)
@@ -536,10 +498,6 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
result.Messages = copyMessages(messages)
return result, stop
}
if ctx.Err() != nil {
result.Messages = copyMessages(messages)
return result, ctx.Err()
}
if collected.Error != "" {
result.Messages = copyMessages(messages)
return result, errors.New(collected.Error)
@@ -688,25 +646,9 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
// recorded. They travel as user messages, and a user message between two
// tool_results breaks strict provider replay — Anthropic coalesces them
// into one user block list and requires the tool_result blocks first, so
// Images ride a following USER message rather than the tool result
// above. Every provider drops images on a tool-role message —
// Anthropic's tool_result content is a string, Gemini's is a
// functionResponse, and OpenAI guards its image parts to the user role
// — so attaching them there would silently deliver nothing. A separate
// message also keeps the one-tool-result-per-tool-call pairing intact,
// which the providers validate. Vision-gate evaluation is deferred until
// after any model switch this turn resolves, so images can reach an
// escalated vision-capable model.
var toolResultsWithImages []ToolResult
buildToolImageMessages := func() []zeroruntime.Message {
var out []zeroruntime.Message
for _, tr := range toolResultsWithImages {
if imageMessage, ok := toolResultImageMessage(tr, options); ok {
out = append(out, imageMessage)
}
}
return out
}
// interleaving yields [tool_result, text, image, tool_result] and a 400.
// Same reason the self-correction feedback below is deferred.
var toolImageMessages []zeroruntime.Message
// Parallel read-ahead state: results for calls[precomputedStart:precomputedEnd]
// executed concurrently, consumed strictly in order below.
var precomputed []precomputedToolResult
@@ -763,8 +705,15 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
IsError: toolResult.Status == tools.StatusError,
ChangedFiles: append([]string(nil), toolResult.ChangedFiles...),
})
if len(toolResult.Images) > 0 {
toolResultsWithImages = append(toolResultsWithImages, toolResult)
// Images ride a following USER message rather than the tool result
// above. Every provider drops images on a tool-role message —
// Anthropic's tool_result content is a string, Gemini's is a
// functionResponse, and OpenAI guards its image parts to the user role
// — so attaching them there would silently deliver nothing. A separate
// message also keeps the one-tool-result-per-tool-call pairing intact,
// which the providers validate.
if imageMessage, ok := toolResultImageMessage(toolResult); ok {
toolImageMessages = append(toolImageMessages, imageMessage)
}
// A tool may demand the run ABORT — a canceled/timed-out ask_user prompt
@@ -776,13 +725,13 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
}
if abortErr != nil {
messages = appendAbortedToolResults(messages, collected.ToolCalls[index+1:])
messages = append(messages, buildToolImageMessages()...)
messages = append(messages, toolImageMessages...)
result.Messages = copyMessages(messages)
return result, abortErr
}
if stopReason := stopReasonFromToolResult(toolResult); stopReason != "" {
messages = appendAbortedToolResults(messages, collected.ToolCalls[index+1:])
messages = append(messages, buildToolImageMessages()...)
messages = append(messages, toolImageMessages...)
result.FinalAnswer = toolResult.ModelOutput()
result.StopReason = stopReason
result.Messages = copyMessages(messages)
@@ -806,7 +755,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
// messages stay valid for a strict provider replay (Anthropic
// rejects a tool_use with no answering tool_result).
messages = appendAbortedToolResults(messages, collected.ToolCalls[index+1:])
messages = append(messages, buildToolImageMessages()...)
messages = append(messages, toolImageMessages...)
result.FinalAnswer = toolFailureStopAnswer(call.Name, outcome.Count)
result.Messages = copyMessages(messages)
return result, nil
@@ -825,6 +774,10 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
postEditDiagnostics.enqueue(ctx, toolResult.ChangedFiles)
}
}
// Every tool_result for this turn is now recorded, including aborted
// placeholders, so the images can follow without splitting them.
messages = append(messages, toolImageMessages...)
toolImageMessages = nil
// Run post-edit self-correction once over the union of files this turn
// changed, then append any feedback after every tool_result is recorded so
@@ -912,12 +865,6 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
}
}
// Every tool_result, self-correct notice, and model-switch update for this
// turn is now in place; evaluate and append image messages against the
// effective (potentially escalated) model.
messages = append(messages, buildToolImageMessages()...)
toolResultsWithImages = nil
// A turn can mix valid tool calls with a dropped (nameless) one. The valid
// calls executed above; surface the dropped call too so it is never
// silently ignored just because the turn also did real work. This is
@@ -3517,10 +3464,8 @@ func copyMessages(messages []Message) []Message {
//
// The text names the tool so the model can tell which call an image came from
// when several ran in one turn — the images arrive detached from their tool
// result, so nothing else associates them. If the effective model cannot
// accept images, the tool's text result is left unchanged and this follow-up
// is a notice with no image parts, matching the CLI/TUI vision gate.
func toolResultImageMessage(result ToolResult, options Options) (zeroruntime.Message, bool) {
// result, so nothing else associates them.
func toolResultImageMessage(result ToolResult) (zeroruntime.Message, bool) {
images := make([]zeroruntime.ImageBlock, 0, len(result.Images))
for _, image := range result.Images {
if len(image.Data) == 0 {
@@ -3541,26 +3486,9 @@ func toolResultImageMessage(result ToolResult, options Options) (zeroruntime.Mes
if label == "" {
label = "tool"
}
if !modelAcceptsToolImages(options) {
return zeroruntime.Message{
Role: zeroruntime.MessageRoleUser,
Content: "Image output from " + label + " was not sent because the current model does not support image input. Switch to a vision-capable model before retrying image delivery, including any images skipped by the tool's image budget.",
}, true
}
return zeroruntime.Message{
Role: zeroruntime.MessageRoleUser,
Content: "Image output from " + label + ":",
Images: images,
}, true
}
func modelAcceptsToolImages(options Options) bool {
if options.SupportsVision != nil {
return options.SupportsVision(options.Model)
}
registry, err := modelregistry.DefaultRegistry()
if err != nil {
return modelregistry.VisionCapableByName(options.Model)
}
return modelregistry.SupportsVision(registry, options.Model)
}
-487
View File
@@ -1,487 +0,0 @@
package agent
import (
"context"
"errors"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/Gitlawb/zero/internal/tools"
"github.com/Gitlawb/zero/internal/trace"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
// midStreamAbortProvider connects successfully but emits a transport-abort
// StreamEventError on the first abortBefore calls, then succeeds with "done".
// hangOnCall, if > 0, makes that 1-based call block until ctx is done so a
// cancel during the retried CollectStream can be reproduced.
// connectFailFrom, if > 0, makes that 1-based call and later return a
// reconnectable connect error so streamWithReconnect enters its backoff.
type midStreamAbortProvider struct {
calls int32
abortBefore int32
abortError string
partialText string
emptyTextEvent bool
partialToolCall string
hangOnCall int32
connectFailFrom int32
connectErr string
started chan struct{}
}
func (p *midStreamAbortProvider) StreamCompletion(ctx context.Context, _ zeroruntime.CompletionRequest) (<-chan zeroruntime.StreamEvent, error) {
n := atomic.AddInt32(&p.calls, 1)
if p.connectFailFrom > 0 && n >= p.connectFailFrom {
if p.started != nil {
select {
case <-p.started:
default:
close(p.started)
}
}
errMsg := p.connectErr
if errMsg == "" {
errMsg = "connection reset by peer"
}
return nil, errors.New(errMsg)
}
if p.hangOnCall > 0 && n == p.hangOnCall {
if p.started != nil {
close(p.started)
}
ch := make(chan zeroruntime.StreamEvent)
go func() {
<-ctx.Done()
close(ch)
}()
return ch, nil
}
ch := make(chan zeroruntime.StreamEvent, 5)
if n <= p.abortBefore {
if p.emptyTextEvent {
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventText, Content: ""}
}
if p.partialText != "" {
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventText, Content: p.partialText}
}
if p.partialToolCall != "" {
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "tc_1", ToolName: p.partialToolCall}
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "tc_1", ArgumentsFragment: `{"path":"x.html","content":"<!doctype`}
}
errMsg := p.abortError
if errMsg == "" {
errMsg = "provider stream error: read: connection reset by peer"
}
ch <- zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventError,
Error: errMsg,
}
close(ch)
return ch, nil
}
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventText, Content: "done"}
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}
close(ch)
return ch, nil
}
func TestRunRetriesMidStreamConnectionReset(t *testing.T) {
p := &midStreamAbortProvider{abortBefore: 1}
result, err := Run(context.Background(), "go", p, Options{Registry: tools.NewRegistry()})
if err != nil {
t.Fatalf("a mid-stream connection reset should retry to success, got %v", err)
}
if result.FinalAnswer != "done" {
t.Fatalf("final answer = %q, want %q", result.FinalAnswer, "done")
}
if got := atomic.LoadInt32(&p.calls); got != 2 {
t.Fatalf("want 2 calls (1 abort + 1 retry), got %d", got)
}
}
func TestRunRetriesMidStreamWindowsAbort(t *testing.T) {
p := &midStreamAbortProvider{
abortBefore: 1,
abortError: "provider stream error: read: wsarecv: An established connection was aborted by the software in your host machine",
}
result, err := Run(context.Background(), "go", p, Options{Registry: tools.NewRegistry()})
if err != nil {
t.Fatalf("a Windows mid-stream abort should retry to success, got %v", err)
}
if result.FinalAnswer != "done" {
t.Fatalf("final answer = %q, want %q", result.FinalAnswer, "done")
}
if got := atomic.LoadInt32(&p.calls); got != 2 {
t.Fatalf("want 2 calls (1 abort + 1 retry), got %d", got)
}
}
// wsarecv alone (no "connection was aborted" substring) must still retry, so
// deleting the wsarecv needle cannot go green.
func TestRunRetriesMidStreamWsarecvNeedle(t *testing.T) {
p := &midStreamAbortProvider{
abortBefore: 1,
abortError: "provider stream error: read: wsarecv: 10053",
}
result, err := Run(context.Background(), "go", p, Options{Registry: tools.NewRegistry()})
if err != nil {
t.Fatalf("wsarecv without 'connection was aborted' should retry to success, got %v", err)
}
if result.FinalAnswer != "done" {
t.Fatalf("final answer = %q, want %q", result.FinalAnswer, "done")
}
if got := atomic.LoadInt32(&p.calls); got != 2 {
t.Fatalf("want 2 calls (1 abort + 1 retry), got %d", got)
}
}
func TestRunDoesNotRetryMidStreamAbortAfterPartialOutput(t *testing.T) {
sawText := 0
p := &midStreamAbortProvider{abortBefore: 1, partialText: "partial"}
_, err := Run(context.Background(), "go", p, Options{
Registry: tools.NewRegistry(),
OnText: func(string) { sawText++ },
})
if err == nil {
t.Fatal("a mid-stream abort after partial output must NOT be retried; want an error")
}
if got := atomic.LoadInt32(&p.calls); got != 1 {
t.Fatalf("partial-then-abort must not retry, got %d calls", got)
}
if sawText == 0 {
t.Fatal("OnText must have fired so forwardedVisibleText is the named guard")
}
}
// A zero-length OnText chunk is not committed answer prose, so an eligible
// transport abort still retries.
func TestRunRetriesMidStreamAbortAfterEmptyTextEvent(t *testing.T) {
p := &midStreamAbortProvider{abortBefore: 1, emptyTextEvent: true}
result, err := Run(context.Background(), "go", p, Options{
Registry: tools.NewRegistry(),
OnText: func(string) {},
})
if err != nil {
t.Fatalf("empty OnText must not block mid-stream abort retry, got %v", err)
}
if result.FinalAnswer != "done" {
t.Fatalf("final answer = %q, want %q", result.FinalAnswer, "done")
}
if got := atomic.LoadInt32(&p.calls); got != 2 {
t.Fatalf("want 2 calls (1 abort + 1 retry), got %d", got)
}
}
func TestRunRetriesMidStreamAbortAfterIncompleteToolCall(t *testing.T) {
starts := 0
dispatched := 0
p := &midStreamAbortProvider{abortBefore: 1, partialToolCall: "write_file"}
result, err := Run(context.Background(), "go", p, Options{
Registry: tools.NewRegistry(),
OnToolCallStart: func(string, string) { starts++ },
OnToolCallDelta: func(string, string) {},
OnToolCall: func(ToolCall) { dispatched++ },
})
if err != nil {
t.Fatalf("a mid-stream abort mid-incomplete-tool-call should retry to success, got %v", err)
}
if result.FinalAnswer != "done" {
t.Fatalf("final answer = %q, want %q", result.FinalAnswer, "done")
}
if got := atomic.LoadInt32(&p.calls); got != 2 {
t.Fatalf("want 2 calls (1 abort + 1 retry), got %d", got)
}
if starts != 1 {
t.Fatalf("the incomplete tool call should have been forwarded once before the abort, got %d starts", starts)
}
if dispatched != 0 {
t.Fatalf("OnToolCall must never fire for an incomplete aborted tool call (no dispatch), got %d", dispatched)
}
}
func TestRunGivesUpAfterMaxMidStreamAbortRetries(t *testing.T) {
if maxStreamStallRetries != 1 {
t.Fatalf("maxStreamStallRetries = %d, want 1 (#973 must not raise this bound)", maxStreamStallRetries)
}
defer func(orig time.Duration) { streamReconnectBase = orig }(streamReconnectBase)
streamReconnectBase = time.Millisecond
p := &midStreamAbortProvider{abortBefore: 99}
_, err := Run(context.Background(), "go", p, Options{Registry: tools.NewRegistry()})
if err == nil {
t.Fatal("a persistent mid-stream abort must surface an error after exhausting retries")
}
if got := atomic.LoadInt32(&p.calls); got != 2 {
t.Fatalf("want 2 calls (1 abort + 1 retry); pin the bound at 1, got %d", got)
}
}
func TestRunMidStreamAbortUsesReconnectNotice(t *testing.T) {
defer func(orig time.Duration) { streamReconnectBase = orig }(streamReconnectBase)
streamReconnectBase = time.Millisecond
var notices string
p := &midStreamAbortProvider{abortBefore: 1}
result, err := Run(context.Background(), "go", p, Options{
Registry: tools.NewRegistry(),
OnReasoning: func(s string) { notices += s },
})
if err != nil {
t.Fatalf("retry should succeed, got %v", err)
}
if result.FinalAnswer != "done" {
t.Fatalf("final answer = %q, want %q", result.FinalAnswer, "done")
}
lower := strings.ToLower(notices)
if !strings.Contains(lower, "connection lost") || !strings.Contains(lower, "reconnecting") {
t.Fatalf("transport abort must use reconnect wording, notices = %q", notices)
}
if strings.Contains(lower, "model stalled") {
t.Fatalf("transport abort must not use stall wording, notices = %q", notices)
}
}
func TestRunCancelDuringMidStreamRetryPreservesContextCanceled(t *testing.T) {
defer func(orig time.Duration) { streamReconnectBase = orig }(streamReconnectBase)
streamReconnectBase = time.Millisecond
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
started := make(chan struct{})
p := &midStreamAbortProvider{abortBefore: 1, hangOnCall: 2, started: started}
errCh := make(chan error, 1)
go func() {
_, err := Run(ctx, "go", p, Options{Registry: tools.NewRegistry()})
errCh <- err
}()
select {
case <-started:
case <-time.After(5 * time.Second):
t.Fatal("timed out waiting for retried stream")
}
cancel()
err := <-errCh
if err == nil {
t.Fatal("want context.Canceled after cancel during retried stream")
}
if !errors.Is(err, context.Canceled) {
t.Fatalf("cancel during retried stream must keep the sentinel, got %q (errors.Is Canceled = false)", err)
}
}
func TestRunCancelDuringReplacementConnectBackoffPreservesContextCanceled(t *testing.T) {
defer func(orig time.Duration) { streamReconnectBase = orig }(streamReconnectBase)
streamReconnectBase = 200 * time.Millisecond
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
started := make(chan struct{})
p := &midStreamAbortProvider{abortBefore: 1, connectFailFrom: 2, started: started}
errCh := make(chan error, 1)
go func() {
_, err := Run(ctx, "go", p, Options{Registry: tools.NewRegistry()})
errCh <- err
}()
select {
case <-started:
case <-time.After(5 * time.Second):
t.Fatal("timed out waiting for replacement connect failure")
}
cancel()
err := <-errCh
if err == nil {
t.Fatal("want context.Canceled after cancel during replacement connect backoff")
}
if !errors.Is(err, context.Canceled) {
t.Fatalf("cancel during replacement connect backoff must keep the sentinel, got %q (errors.Is Canceled = false)", err)
}
}
func TestRunMidStreamAbortRecordsReconnectCountOnImmediateReplacement(t *testing.T) {
defer func(orig time.Duration) { streamReconnectBase = orig }(streamReconnectBase)
streamReconnectBase = time.Millisecond
rec := trace.NewRecorder("midstream-abort", "run-1", "test")
p := &midStreamAbortProvider{abortBefore: 1}
result, err := Run(context.Background(), "go", p, Options{
Registry: tools.NewRegistry(),
Trace: rec,
})
if err != nil {
t.Fatalf("retry should succeed, got %v", err)
}
if result.FinalAnswer != "done" {
t.Fatalf("final answer = %q, want %q", result.FinalAnswer, "done")
}
if got := atomic.LoadInt32(&p.calls); got != 2 {
t.Fatalf("want 2 calls (1 abort + 1 immediate replacement), got %d", got)
}
gotTrace := rec.Finish()
if got := gotTrace.Counter(trace.CounterReconnectCount); got != 1 {
t.Fatalf("reconnect_count = %d, want 1 for a successful replacement connection", got)
}
if got := gotTrace.Counter(trace.CounterModelRequests); got != 2 {
t.Fatalf("model_requests = %d, want 2 (initial + replacement)", got)
}
}
func TestRunDoesNotRetryResponseHeaderTimeout(t *testing.T) {
p := &midStreamAbortProvider{
abortBefore: 99,
abortError: "net/http: timeout awaiting response headers",
}
_, err := Run(context.Background(), "go", p, Options{Registry: tools.NewRegistry()})
if err == nil {
t.Fatal("a response-header timeout must not be retried as a mid-stream abort")
}
if got := atomic.LoadInt32(&p.calls); got != 1 {
t.Fatalf("header timeout must not retry, got %d calls", got)
}
}
func TestIsMidStreamTransportAbort(t *testing.T) {
aborts := []string{
"provider stream error: read: connection reset by peer",
"provider stream error: unexpected EOF",
"provider stream error: read: wsarecv: 10053",
"provider stream error: read: wsarecv: An established connection was aborted by the software in your host machine",
"An existing connection was forcibly closed by the remote host",
"provider stream error: read: connection closed",
"write: broken pipe",
"provider stream error: server closed the connection",
"provider stream error: software caused connection abort",
}
for _, m := range aborts {
if !isMidStreamTransportAbort(m) {
t.Fatalf("want mid-stream transport abort: %q", m)
}
}
notAborts := []string{
"",
"context length exceeded",
"rate limit error: slow down",
"model not found",
"net/http: timeout awaiting response headers",
"dial tcp 10.0.0.1:443: connect: connection refused",
"i/o timeout",
"500 Internal Server Error",
"502 Bad Gateway",
"503 Service Unavailable",
"504 Gateway Timeout",
"provider request error: request does not satisfy oneOf schema",
"provider request error: connection closed is not a supported finish reason",
"provider error: connection closed is not a supported finish reason",
"provider error: upstream connection reset by peer",
"provider error: unexpected EOF while reading upstream response",
"provider error: 502 Bad Gateway",
"auth error: connection reset by peer",
"rate limit error: server closed",
"schema error: property oneOf is invalid",
}
for _, m := range notAborts {
if isMidStreamTransportAbort(m) {
t.Fatalf("must NOT classify as mid-stream transport abort: %q", m)
}
}
}
func TestRunDoesNotRetryClassifiedProviderErrorWithSocketPhrase(t *testing.T) {
p := &midStreamAbortProvider{
abortBefore: 1,
abortError: "provider request error: connection closed is not a supported finish reason",
}
var notices string
opts := Options{
Registry: tools.NewRegistry(),
OnReasoning: func(s string) { notices += s },
}
_, err := Run(context.Background(), "go", p, opts)
if err == nil {
t.Fatal("expected application error, got nil")
}
if got := atomic.LoadInt32(&p.calls); got != 1 {
t.Fatalf("classified provider error with 'connection closed' must not retry (got %d calls, want 1)", got)
}
if notices != "" {
t.Fatalf("expected no reconnect notices, got %q", notices)
}
if !strings.Contains(err.Error(), "connection closed") {
t.Fatalf("expected original error preserved, got %v", err)
}
}
func TestRunDoesNotRetryProviderErrorWithSocketPhrase(t *testing.T) {
p := &midStreamAbortProvider{
abortBefore: 1,
abortError: "provider error: unexpected EOF while reading upstream response",
}
var notices string
opts := Options{
Registry: tools.NewRegistry(),
OnReasoning: func(s string) { notices += s },
}
_, err := Run(context.Background(), "go", p, opts)
if err == nil {
t.Fatal("expected application error, got nil")
}
if got := atomic.LoadInt32(&p.calls); got != 1 {
t.Fatalf("provider error with 'unexpected EOF' must not retry (got %d calls, want 1)", got)
}
if notices != "" {
t.Fatalf("expected no reconnect notices, got %q", notices)
}
if got := err.Error(); got != p.abortError {
t.Fatalf("error = %q, want original provider error %q", got, p.abortError)
}
}
func TestRunDoesNotRetryApplicationErrorContainingEOFSubstring(t *testing.T) {
p := &midStreamAbortProvider{
abortBefore: 1,
abortError: "provider request error: request does not satisfy oneOf schema",
}
var notices string
opts := Options{
Registry: tools.NewRegistry(),
OnReasoning: func(s string) { notices += s },
}
_, err := Run(context.Background(), "go", p, opts)
if err == nil {
t.Fatal("expected application error, got nil")
}
if got := atomic.LoadInt32(&p.calls); got != 1 {
t.Fatalf("application error with 'oneOf' must not retry (got %d calls, want 1)", got)
}
if notices != "" {
t.Fatalf("expected no reconnect notices, got %q", notices)
}
if !strings.Contains(err.Error(), "oneOf") {
t.Fatalf("expected original error preserved, got %v", err)
}
}
func TestRunRetriesMidStreamUnexpectedEOF(t *testing.T) {
defer func(orig time.Duration) { streamReconnectBase = orig }(streamReconnectBase)
streamReconnectBase = time.Millisecond
p := &midStreamAbortProvider{
abortBefore: 1,
abortError: "provider stream error: unexpected EOF",
}
var notices string
opts := Options{
Registry: tools.NewRegistry(),
OnReasoning: func(s string) { notices += s },
}
result, err := Run(context.Background(), "go", p, opts)
if err != nil {
t.Fatalf("unexpected EOF should retry to success, got %v", err)
}
if result.FinalAnswer != "done" {
t.Fatalf("final answer = %q, want %q", result.FinalAnswer, "done")
}
if got := atomic.LoadInt32(&p.calls); got != 2 {
t.Fatalf("want 2 calls (1 abort + 1 retry), got %d", got)
}
lower := strings.ToLower(notices)
if !strings.Contains(lower, "connection lost") || !strings.Contains(lower, "reconnecting") {
t.Fatalf("expected reconnect notice, got %q", notices)
}
}
-89
View File
@@ -144,15 +144,6 @@ func shouldReconnect(ctx context.Context, err error) bool {
"i/o timeout",
"server closed",
"unexpected end",
// Windows mid-stream aborts (WSAECONNABORTED / wsarecv) — host or AV
// forcibly closes an established socket during CollectStream.
// "software caused connection abort" is the Unix/macOS ECONNABORTED
// text; it also appears in midStreamAbortNeedles, which must stay a
// subset of this list.
"wsarecv",
"connection was aborted",
"software caused connection abort",
"forcibly closed",
} {
if strings.Contains(msg, needle) {
return true
@@ -161,86 +152,6 @@ func shouldReconnect(ctx context.Context, err error) bool {
return false
}
// midStreamAbortNeedles are the transport-abort classes retried AFTER connect
// succeeded and the stream body has started. This is a subset of shouldReconnect:
// connect-phase signals (timeout, connection refused, "temporarily unavailable")
// stay on the connect retry path, because matching them here would re-prefill a
// healthy-but-slow server (ollama cloud header timeouts) and tell the user the
// connection was lost. Issue #973 names WSAECONNABORTED / connection reset /
// unexpected stream EOF.
var midStreamAbortNeedles = []string{
"unexpected end",
"connection reset",
"broken pipe",
"connection closed",
"server closed",
"wsarecv",
"connection was aborted",
"software caused connection abort",
"forcibly closed",
}
var classifiedNonTransportPrefixes = []string{
"provider request error:",
"provider error:",
"auth error:",
"rate limit error:",
}
// isMidStreamTransportAbort reports whether a collected stream error string is a
// retryable mid-stream transport abort. It does NOT delegate to shouldReconnect:
// that list is a connect-phase argument ("no response was received"). The stream
// body has already started here, so only abort/reset/EOF/close classes retry.
//
// Used for failures DURING the stream body AFTER a successful connect and BEFORE
// tool dispatch — the incomplete turn committed no answer text and executed no
// tools, so a bounded re-issue is safe (same safety rules as the stall path).
func isMidStreamTransportAbort(message string) bool {
lowered := strings.ToLower(strings.TrimSpace(message))
if lowered == "" || isContextLimitError(lowered) {
return false
}
for _, prefix := range classifiedNonTransportPrefixes {
if strings.Contains(lowered, prefix) {
return false
}
}
if errhint.HasStatusCode(lowered, "500", "502", "503", "504") {
return false
}
if containsWordBoundary(lowered, "eof") {
return true
}
for _, needle := range midStreamAbortNeedles {
if strings.Contains(lowered, needle) {
return true
}
}
return false
}
func containsWordBoundary(text, word string) bool {
start := 0
for {
idx := strings.Index(text[start:], word)
if idx < 0 {
return false
}
pos := start + idx
endPos := pos + len(word)
leftBoundary := pos == 0 || !isAlphaNumByte(text[pos-1])
rightBoundary := endPos == len(text) || !isAlphaNumByte(text[endPos])
if leftBoundary && rightBoundary {
return true
}
start = pos + 1
}
}
func isAlphaNumByte(c byte) bool {
return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '_'
}
// 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 {
-6
View File
@@ -92,12 +92,6 @@ func TestShouldReconnectClassification(t *testing.T) {
disconnects := []string{
"unexpected EOF", "connection reset by peer", "broken pipe",
"i/o timeout", "server closed the connection", "connection refused",
// Windows socket aborts also remain reconnectable at connect time.
"wsarecv: An established connection was aborted by the software in your host machine",
"read: connection was aborted",
// Unix/macOS ECONNABORTED text.
"read: software caused connection abort",
"An existing connection was forcibly closed by the remote host",
}
for _, m := range disconnects {
if !shouldReconnect(ctx, errors.New(m)) {
+1 -198
View File
@@ -10,24 +10,6 @@ import (
"github.com/Gitlawb/zero/internal/zeroruntime"
)
func TestNonVisionImageNoticeQualifiesBudgetRecovery(t *testing.T) {
result := ToolResult{
Name: "mcp_screenshot",
Output: "One image exceeded this result's remaining image budget. Retrying with fewer images can recover this payload.",
Images: []zeroruntime.ImageBlock{{MediaType: "image/png", Data: []byte("image")}},
}
notice, ok := toolResultImageMessage(result, Options{
SupportsVision: func(string) bool { return false },
})
if !ok || len(notice.Images) != 0 {
t.Fatal("non-vision delivery must produce a text-only notice")
}
if !strings.Contains(notice.Content, "Switch to a vision-capable model before retrying") ||
!strings.Contains(notice.Content, "budget") {
t.Fatalf("notice does not qualify budget recovery: %s", notice.Content)
}
}
// imageTool returns a tool result carrying an image, the way a screenshot tool
// should be able to.
type imageTool struct {
@@ -80,7 +62,6 @@ func TestRunDeliversToolResultImagesToTheModel(t *testing.T) {
result, err := Run(context.Background(), "screenshot please", provider, Options{
Registry: registry,
MaxTurns: 2,
Model: "gpt-4o",
})
if err != nil {
t.Fatalf("Run: %v", err)
@@ -105,16 +86,11 @@ func TestRunDeliversToolResultImagesToTheModel(t *testing.T) {
// The tool-result pairing must be untouched: one tool result per tool call.
toolResults := 0
var toolText string
for _, message := range result.Messages {
if message.Role == zeroruntime.MessageRoleTool {
toolResults++
toolText = message.Content
}
}
if toolText != "Captured a screenshot." {
t.Fatalf("tool text = %q, want preserved", toolText)
}
if toolResults != 1 {
t.Errorf("got %d tool-result messages for 1 tool call", toolResults)
}
@@ -160,7 +136,7 @@ func TestRunKeepsToolResultsContiguousWhenAToolReturnsAnImage(t *testing.T) {
{{Type: zeroruntime.StreamEventText, Content: "done"}, {Type: zeroruntime.StreamEventDone}},
}}
result, err := Run(context.Background(), "two calls", provider, Options{Registry: registry, MaxTurns: 2, Model: "gpt-4o"})
result, err := Run(context.Background(), "two calls", provider, Options{Registry: registry, MaxTurns: 2})
if err != nil {
t.Fatalf("Run: %v", err)
}
@@ -229,176 +205,3 @@ func messageShape(messages []zeroruntime.Message) string {
}
return "[" + strings.Join(parts, " ") + "]"
}
func TestRunDropsToolResultImagesForANonVisionModel(t *testing.T) {
registry := tools.NewRegistry()
registry.Register(imageTool{media: "image/png", data: []byte("\x89PNG\r\n\x1a\nfake")})
provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{
{
{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call_1", ToolName: "capture"},
{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call_1", ArgumentsFragment: `{}`},
{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call_1"},
{Type: zeroruntime.StreamEventDone},
},
{{Type: zeroruntime.StreamEventText, Content: "ok"}, {Type: zeroruntime.StreamEventDone}},
}}
result, err := Run(context.Background(), "screenshot please", provider, Options{
Registry: registry,
MaxTurns: 2,
Model: "totally-made-up-model",
})
if err != nil {
t.Fatalf("Run: %v", err)
}
var toolText string
var noticed bool
for _, message := range result.Messages {
if message.Role == zeroruntime.MessageRoleTool {
toolText = message.Content
}
if len(message.Images) > 0 {
t.Fatalf("non-vision model received image bytes on %q: %s", message.Role, messageShape(result.Messages))
}
if strings.Contains(message.Content, "does not support image input") {
noticed = true
}
}
if toolText != "Captured a screenshot." {
t.Fatalf("tool text = %q, want preserved on the non-vision path", toolText)
}
if !noticed {
t.Fatalf("non-vision model was not told the image was dropped; recorded %s", messageShape(result.Messages))
}
}
type switchAndCaptureTool struct {
targetModel string
}
func (switchAndCaptureTool) Name() string { return "switch_and_capture" }
func (switchAndCaptureTool) Description() string { return "Switches model and captures image" }
func (switchAndCaptureTool) Parameters() tools.Schema {
return tools.Schema{Type: "object", Properties: map[string]tools.PropertySchema{}}
}
func (switchAndCaptureTool) Safety() tools.Safety {
return tools.Safety{Permission: tools.PermissionAllow}
}
func (t switchAndCaptureTool) Run(context.Context, map[string]any) tools.Result {
return tools.Result{
Status: tools.StatusOK,
Output: "[image returned by tool]",
Images: []zeroruntime.ImageBlock{{MediaType: "image/png", Data: []byte("\x89PNG\r\n\x1a\nfake")}},
Meta: map[string]string{"escalate_to_model": t.targetModel},
}
}
func TestRunToolImagesRespectsModelSwitch(t *testing.T) {
t.Run("Non-vision to vision model escalation forwards images", func(t *testing.T) {
registry := tools.NewRegistry()
registry.Register(switchAndCaptureTool{targetModel: "gpt-4o"})
provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{
{
{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call_1", ToolName: "switch_and_capture"},
{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call_1", ArgumentsFragment: `{}`},
{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call_1"},
{Type: zeroruntime.StreamEventDone},
},
{
{Type: zeroruntime.StreamEventText, Content: "I see it on gpt-4o."},
{Type: zeroruntime.StreamEventDone},
},
}}
switched := false
switcher := func(_ context.Context, modelID string) (Provider, error) {
switched = true
return provider, nil
}
result, err := Run(context.Background(), "test", provider, Options{
Registry: registry,
MaxTurns: 2,
Model: "non-vision-initial",
ModelSwitcher: switcher,
SupportsVision: func(modelID string) bool {
return modelID == "gpt-4o"
},
})
if err != nil {
t.Fatalf("Run: %v", err)
}
if !switched {
t.Fatal("expected model switch to occur")
}
var carrier *zeroruntime.Message
for index := range result.Messages {
if len(result.Messages[index].Images) > 0 {
carrier = &result.Messages[index]
}
if strings.Contains(result.Messages[index].Content, "does not support image input") {
t.Fatalf("images were unexpectedly dropped after switch to vision model: %v", result.Messages[index])
}
}
if carrier == nil {
t.Fatal("expected image carrier message after escalation to vision model")
}
})
t.Run("Vision to non-vision model switch drops images with notice", func(t *testing.T) {
registry := tools.NewRegistry()
registry.Register(switchAndCaptureTool{targetModel: "text-only-dest"})
provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{
{
{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call_1", ToolName: "switch_and_capture"},
{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call_1", ArgumentsFragment: `{}`},
{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call_1"},
{Type: zeroruntime.StreamEventDone},
},
{
{Type: zeroruntime.StreamEventText, Content: "ok."},
{Type: zeroruntime.StreamEventDone},
},
}}
switcher := func(_ context.Context, modelID string) (Provider, error) {
return provider, nil
}
result, err := Run(context.Background(), "test", provider, Options{
Registry: registry,
MaxTurns: 2,
Model: "gpt-4o",
ModelSwitcher: switcher,
SupportsVision: func(modelID string) bool {
return modelID == "gpt-4o"
},
})
if err != nil {
t.Fatalf("Run: %v", err)
}
var noticed bool
var toolOutputSeen bool
for _, m := range result.Messages {
if len(m.Images) > 0 {
t.Fatalf("image delivered to non-vision destination model: %v", m)
}
if strings.Contains(m.Content, "does not support image input") {
noticed = true
}
if strings.Contains(m.Content, "[image returned by tool]") {
toolOutputSeen = true
}
}
if !noticed {
t.Fatal("expected drop notice when switching to non-vision model")
}
if !toolOutputSeen {
t.Fatal("expected tool output preserved in message history")
}
})
}
-5
View File
@@ -341,11 +341,6 @@ type Options struct {
// nil for text-only runs (the seeded message then carries no images, exactly
// as before).
Images []zeroruntime.ImageBlock
// SupportsVision, when set, reports whether the effective model accepts
// image input. Tool-produced images are dropped at the shared delivery
// boundary when this is false. nil uses the curated catalog plus the
// name heuristic via modelregistry.SupportsVision.
SupportsVision func(modelID string) bool
// ContextWindow is the model's maximum input token budget. When > 0 the agent
// loop compacts long conversations once the estimated size crosses a fraction
// of this window. 0 DISABLES compaction entirely (every existing caller/test
+1 -8
View File
@@ -4,7 +4,6 @@ package background
import (
"fmt"
"os"
"os/exec"
"github.com/Gitlawb/zero/internal/execution"
@@ -47,13 +46,7 @@ func terminateOwnedProcess(cmd *exec.Cmd) (bool, error) {
terminateErr = terminateProcessForTest(cmd.Process.Pid)
})
if err != nil {
// WithHandle refuses only a process that has already been waited or
// released, so it has finished. Say so in the form exec.Cmd understands:
// a Cancel error wrapping os.ErrProcessDone is "nothing to cancel", where
// the unexported "already released" error it would otherwise carry made
// Wait report a spurious cancel failure. A dead root has no tree left to
// find either way.
return true, fmt.Errorf("pin process %d for termination: %w: %w", cmd.Process.Pid, os.ErrProcessDone, err)
return false, fmt.Errorf("pin process %d for termination: %w", cmd.Process.Pid, err)
}
return alreadyExited, terminateErr
}
-23
View File
@@ -73,29 +73,6 @@ func TerminateCommand(cmd *exec.Cmd) error {
return nil
}
// TerminateCommandGroup stops a started command's process group on POSIX, or its
// process tree on Windows, WITHOUT reaping it.
//
// It is TerminateCommand without the Wait, for a caller whose Wait already runs
// somewhere else: a worker pool waiting on the command in its own goroutine, or
// the command's own Cancel hook, which exec.Cmd invokes while its Wait is in
// progress. TerminateCommand would be a second Wait there. Both used to pass
// cmd.Process.Pid to TerminateProcess, which rediscovers the group at signal time
// with Getpgid, and once the leader has exited that lookup fails, so only the dead
// leader was signalled and its descendants kept running. When cmd was configured
// by ConfigureChildProcessGroup the group now comes from that launch-time
// configuration instead. On Windows it is the tree kill TerminateProcess performs,
// with the process handle pinned while it runs so the PID cannot be recycled
// underneath it; once the command has been waited that pin fails with an error
// wrapping os.ErrProcessDone.
func TerminateCommandGroup(cmd *exec.Cmd) error {
if cmd == nil || cmd.Process == nil {
return nil
}
_, err := terminateOwnedProcess(cmd)
return err
}
// classifyWaitError treats a non-zero exit status as success: a process being
// terminated is expected to report one (or a signal), so the only interesting
// failure is Wait itself not working.
@@ -1,71 +0,0 @@
package background
import (
"errors"
"os"
"os/exec"
"runtime"
"testing"
"time"
)
// TerminateCommandGroup LEAVES THE WAIT TO ITS CALLER. Its callers already have a
// Wait running or about to run: the daemon pool waits on its workers, and
// exec.Cmd calls the Cancel hook from inside Wait. Reaping here, the way
// TerminateCommand does, would make their Wait fail with "Wait was already
// called" instead of reporting how the command ended.
func TestTerminateCommandGroupLeavesTheWaitToTheCaller(t *testing.T) {
cmd := exec.Command("sleep", "30")
if runtime.GOOS == "windows" {
cmd = exec.Command("powershell.exe", "-NoProfile", "-NonInteractive", "-Command", "Start-Sleep -Seconds 30")
}
ConfigureChildProcessGroup(cmd)
if err := cmd.Start(); err != nil {
t.Fatalf("start: %v", err)
}
if err := TerminateCommandGroup(cmd); err != nil {
_ = cmd.Process.Kill()
_ = cmd.Wait()
t.Fatalf("TerminateCommandGroup: %v", err)
}
waited := make(chan error, 1)
go func() { waited <- cmd.Wait() }()
select {
case err := <-waited:
var exitErr *exec.ExitError
if err != nil && !errors.As(err, &exitErr) {
t.Fatalf("the caller's own Wait failed after TerminateCommandGroup: %v", err)
}
case <-time.After(10 * time.Second):
_ = cmd.Process.Kill()
t.Fatal("the command was still running 10s after TerminateCommandGroup")
}
}
// A command that was never started, or a nil one, is nothing to stop.
func TestTerminateCommandGroupWithNoProcessIsANoOp(t *testing.T) {
if err := TerminateCommandGroup(nil); err != nil {
t.Fatalf("nil command: %v", err)
}
if err := TerminateCommandGroup(exec.Command("sleep", "1")); err != nil {
t.Fatalf("unstarted command: %v", err)
}
}
// After the caller has waited, Windows cannot pin the process any more. The error
// must still read as "already finished", which is what exec.Cmd's Cancel contract
// and the pool both treat as nothing wrong.
func TestTerminateCommandGroupAfterWaitReportsProcessDone(t *testing.T) {
if runtime.GOOS != "windows" {
t.Skip("POSIX signals the group, which is simply gone after the wait")
}
cmd := exec.Command("cmd.exe", "/c", "exit", "0")
ConfigureChildProcessGroup(cmd)
if err := cmd.Run(); err != nil {
t.Fatalf("run: %v", err)
}
if err := TerminateCommandGroup(cmd); err != nil && !errors.Is(err, os.ErrProcessDone) {
t.Fatalf("after Wait the error is %v, want nil or one wrapping os.ErrProcessDone", err)
}
}
+2 -12
View File
@@ -695,7 +695,7 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
// Refresh the models.dev pricing/limits cache in the background when stale;
// the overlay is read at registry construction from the cache file, so this
// benefits the next run and never blocks or fails this one.
modelregistry.StartModelsDevRefresh()
go func() { _ = modelregistry.RefreshModelsDevCache(context.Background()) }()
workspaceRoot, err := deps.getwd()
if err != nil {
@@ -727,13 +727,7 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
resolved.Provider = usable
resolved.ActiveProvider = usable.Name
} else {
// Fresh-onboarding reset, but the user's notification preference is
// NOT config-to-redo: clearing it here would surface as an empty
// Options.Notify, and the TUI's unconfigured default (both/unfocused)
// would then resurrect alerts a user explicitly turned off while the
// setup wizard runs (maintainer review, PR #1001). Carry the stored
// block through; other wizard-visible state starts clean as before.
resolved = config.ResolvedConfig{Notify: resolved.Notify}
resolved = config.ResolvedConfig{}
forceSetup = true
}
}
@@ -965,9 +959,6 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
// notice when project hooks/plugins were dropped for an untrusted workspace.
hookDispatcher, hookSkip := newHookDispatcherWithExtra(workspaceRoot, pluginActivation.hooks, trustRoot, executionRunner)
emitTrustNotice(stderr, hookSkip, pluginActivation.trustSkip, mcpSkip)
// From the engine the session's commands run through, so the notice
// describes the sandbox they actually get.
startupNotices := []string{sandboxEngine.DegradedNotice()}
return deps.runTUI(context.Background(), tui.Options{
Cwd: workspaceRoot,
Version: version,
@@ -1055,7 +1046,6 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
PermissionMode: permissionMode,
Notify: resolved.Notify,
KeyBindings: resolved.KeyBindings,
StartupNotices: startupNotices,
STT: resolved.STT,
BuildDictationTranscriber: newDictationTranscriberFactory(resolved, userConfigPath, sttServerManager),
ShutdownDictationServer: sttServerManager.Shutdown,
-67
View File
@@ -377,73 +377,6 @@ func TestRunNoArgsFallsBackToUsableProviderWhenNoneMarkedActive(t *testing.T) {
}
}
// Maintainer regression (PR #1001): a stored explicit notify opt-out must
// survive the provider-recovery startup path. With a stale activeProvider and
// another usable saved provider, Resolve returns the partial config alongside
// ErrNoActiveProvider and the recovery branch forwards it — dropping the
// notify policy there would surface an empty Options.Notify, and the TUI's
// unconfigured default (both) would resurrect alerts the user explicitly
// turned off. The same recovery branch that clears the resolved config for
// the setup wizard carries the stored notify block through.
func TestRunNoArgsPreservesNotifyOptOutThroughProviderRecovery(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer
cwd := t.TempDir()
setCLIUserConfigRoot(t)
userConfigPath := filepath.Join(t.TempDir(), "zero", "config.json")
var launchedOptions tui.Options
launched := false
usable := config.ProviderProfile{
Name: "work",
ProviderKind: config.ProviderKindOpenAI,
BaseURL: config.OpenAIBaseURL,
APIKey: "sk-test",
Model: "gpt-test",
}
exitCode := runWithDeps([]string{}, &stdout, &stderr, appDeps{
getwd: func() (string, error) {
return cwd, nil
},
resolveConfig: func(workspaceRoot string, overrides config.Overrides) (config.ResolvedConfig, error) {
// Stale activeProvider + usable saved provider, and the user's
// config.json explicitly opts out of notifications. The fixed
// resolver carries the parsed notify block through the error path.
return config.ResolvedConfig{
Providers: []config.ProviderProfile{usable},
Notify: config.NotifyConfig{Mode: "off", FocusMode: "always"},
}, fmt.Errorf("%w: active provider %q not found", config.ErrNoActiveProvider, "ghost")
},
newProvider: func(profile config.ProviderProfile) (zeroruntime.Provider, error) {
return &cliFakeProvider{}, nil
},
userConfigPath: func() (string, error) {
return userConfigPath, nil
},
registerMCPTools: func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
return noopMCPRuntime{}, nil
},
runTUI: func(ctx context.Context, options tui.Options) int {
launched = true
launchedOptions = options
return 0
},
})
if exitCode != 0 {
t.Fatalf("exit code = %d, want 0, stderr=%q", exitCode, stderr.String())
}
if !launched {
t.Fatal("TUI was not launched")
}
// The explicit opt-out reaches the TUI intact — NOT the both/unfocused
// unconfigured default.
if launchedOptions.Notify.Mode != "off" || launchedOptions.Notify.FocusMode != "always" {
t.Fatalf("Options.Notify = %+v, want the stored off/always opt-out preserved through recovery", launchedOptions.Notify)
}
}
func TestRunNoArgsFailsWhenResolveErrorIsNotProviderRelated(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer
+1 -29
View File
@@ -26,30 +26,6 @@ type modelSummary = zerocommands.ModelSnapshot
type providerCatalogSummary = zerocommands.ProviderCatalogSnapshot
func runConfig(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int {
// The first non-flag argument is a subcommand. With no positional argument
// (or only flag arguments), the read-only summary path runs — this keeps
// `zero config` and `zero config --json` working unchanged.
command := "summary"
rest := args
if len(args) > 0 && !strings.HasPrefix(args[0], "-") {
command = strings.ToLower(strings.TrimSpace(args[0]))
rest = args[1:]
}
switch command {
case "summary":
return runConfigSummary(rest, stdout, stderr, deps)
case "notify":
return runConfigNotify(rest, stdout, stderr, deps)
case "help":
if err := writeConfigHelp(stdout); err != nil {
return exitCrash
}
return exitSuccess
}
return writeExecUsageError(stderr, fmt.Sprintf("unknown config command %q", command))
}
func runConfigSummary(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int {
options, help, err := parseCommandCenterArgs(args, false, false)
if err != nil {
return writeExecUsageError(stderr, err.Error())
@@ -490,12 +466,8 @@ func formatProviderCatalogValue(value string, fallback string) string {
func writeConfigHelp(w io.Writer) error {
_, err := fmt.Fprint(w, `Usage:
zero config [flags]
zero config notify [flags]
Inspects resolved Go configuration without printing secrets. The notify
subcommand reads or updates the stored global notification preference,
which controls both the completion and needs-input alerts —
run "zero config notify --help" for details.
Inspects resolved Go configuration without printing secrets.
Flags:
--json Print JSON summary
+1 -5
View File
@@ -40,11 +40,7 @@ var completionRoot = completionNode{
},
{names: []string{"daemon"}, children: leafNodes("start", "stop", "status", "run", "attach", "serve-remote", "link")},
{names: []string{"setup"}},
{names: []string{"config"}, flags: []string{"-h", "--help", "--json"}, children: []completionNode{
{names: []string{"notify"}, flags: []string{
"-h", "--help", "--mode", "--focus", "--reset", "--json",
}},
}},
{names: []string{"config"}},
{names: []string{"models"}, children: []completionNode{{names: []string{"list", "ls"}}}},
{names: []string{"providers"}, children: []completionNode{
{names: []string{"current"}}, {names: []string{"list"}}, {names: []string{"catalog"}},
-5
View File
@@ -165,11 +165,6 @@ func TestCompletionTreeCoversAliasesNestingAndCommonFlags(t *testing.T) {
assertCandidates(t, byPath["completions"], "bash", "zsh", "fish", "powershell", "elvish")
assertCandidates(t, byPath["plugins"], "list", "add", "info", "remove", "rm")
assertCandidates(t, byPath["plugin"], "list", "add", "info", "remove", "rm")
// `config` gained the notify subcommand; the completion tree must expose
// it (and its flags) so shell completion cannot go stale for new CLI
// surfaces (maintainer review, PR #1001).
assertCandidates(t, byPath["config"], "notify", "--json", "--help")
assertCandidates(t, byPath["config notify"], "--mode", "--focus", "--reset", "--json", "--help")
}
func assertCandidates(t *testing.T, got []string, wants ...string) {
-177
View File
@@ -1,177 +0,0 @@
package cli
import (
"fmt"
"io"
"strings"
"github.com/Gitlawb/zero/internal/config"
)
// runConfigNotify implements `zero config notify`: with no flags it prints the
// stored mode/focusMode; --mode/--focus update them via the same
// config.SetNotify writer the TUI /notify command uses, so all surfaces stay
// in lockstep; --reset blanks both fields so the TUI's effective default
// applies again (an unconfigured headless run stays silent).
//
// The command manages a user preference, so it talks ONLY to the user's own
// config file (config.UserNotify / config.SetNotify) and never runs the full
// config resolution: resolving providers would fail with ErrNoActiveProvider
// for a brand-new user, locking them out of setting notifications before they
// have even configured a provider (CodeRabbit review, PR #1001). The display
// therefore reports the USER'S stored values — a project config that
// overrides notify for one repo is not shown here, by the same logic the
// maintainer applied to the write path.
func runConfigNotify(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int {
options, help, err := parseConfigNotifyArgs(args)
if err != nil {
return writeExecUsageError(stderr, err.Error())
}
if help {
if err := writeConfigNotifyHelp(stdout); err != nil {
return exitCrash
}
return exitSuccess
}
configPath, err := deps.userConfigPath()
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
if options.mode != "" || options.focus != "" || options.reset {
// One serialized read-merge-write transaction: the lock covers
// reading the stored block, applying only the explicit fields, and
// replacing the file, so two concurrent partial updates (e.g.
// --mode off and --focus always from two terminals) cannot lose
// each other's change (maintainer review, PR #1001). Omitted
// fields preserve the values stored in the user's OWN file — blank
// stays blank; --reset is the only path that clears both fields.
_, err := config.UpdateNotify(configPath, func(current config.NotifyConfig) config.NotifyConfig {
if options.reset {
return config.NotifyConfig{}
}
if options.mode != "" {
current.Mode = options.mode
}
if options.focus != "" {
current.FocusMode = options.focus
}
return current
})
if err != nil {
return writeAppError(stderr, err.Error(), exitUsage)
}
}
// Report the stored values (blank renders as "(default)").
current, err := config.UserNotify(configPath)
if err != nil {
return writeAppError(stderr, err.Error(), exitUsage)
}
if options.json {
if err := writePrettyJSON(stdout, map[string]any{
"mode": current.Mode,
"focusMode": current.FocusMode,
}); err != nil {
return exitCrash
}
return exitSuccess
}
lines := []string{
"Notify",
"mode: " + displayCLIValue(current.Mode, "(default)"),
"focusMode: " + displayCLIValue(current.FocusMode, "(default)"),
}
if _, err := fmt.Fprintln(stdout, strings.Join(lines, "\n")); err != nil {
return exitCrash
}
return exitSuccess
}
type configNotifyOptions struct {
mode string
focus string
reset bool
json bool
}
func parseConfigNotifyArgs(args []string) (configNotifyOptions, bool, error) {
options := configNotifyOptions{}
for index := 0; index < len(args); index++ {
arg := args[index]
switch {
case arg == "-h" || arg == "--help" || arg == "help":
return options, true, nil
case arg == "--json":
options.json = true
case arg == "--reset":
options.reset = true
case arg == "--mode":
value, next, err := nextFlagValue(args, index, arg)
if err != nil {
return options, false, err
}
options.mode = value
index = next
case strings.HasPrefix(arg, "--mode="):
value, err := requiredInlineFlagValue(arg, "--mode")
if err != nil {
return options, false, err
}
options.mode = value
case arg == "--focus":
value, next, err := nextFlagValue(args, index, arg)
if err != nil {
return options, false, err
}
options.focus = value
index = next
case strings.HasPrefix(arg, "--focus="):
value, err := requiredInlineFlagValue(arg, "--focus")
if err != nil {
return options, false, err
}
options.focus = value
case strings.HasPrefix(arg, "-"):
return options, false, execUsageError{fmt.Sprintf("unknown flag %q", arg)}
default:
return options, false, execUsageError{fmt.Sprintf("unexpected argument %q", arg)}
}
}
return options, false, nil
}
func writeConfigNotifyHelp(w io.Writer) error {
_, err := fmt.Fprint(w, "Usage:\n"+
" zero config notify [flags]\n"+
"\n"+
"Print or update the stored global notification preference.\n"+
"\n"+
"The preference controls BOTH notification kinds: the turn-completion\n"+
"(\"Zero: ready\") alert and the needs-input alert. mode off silences both;\n"+
"the focus mode (unfocused, always, focused) applies to both.\n"+
"\n"+
"When run with no flag, prints the stored mode and focusMode; a field you\n"+
"never set shows as (default) — the TUI alerts with bell + notification,\n"+
"firing only when the terminal is unfocused, while an unconfigured\n"+
"headless run stays silent. Omitted flags preserve the values stored in\n"+
"YOUR config file; --reset clears both so the TUI's effective default\n"+
"applies again. If --reset is combined with --mode/--focus, --reset wins\n"+
"and the explicit values are ignored.\n"+
"\n"+
"Examples:\n"+
" zero config notify\n"+
" zero config notify --json\n"+
" zero config notify --mode both --focus unfocused\n"+
" zero config notify --mode off\n"+
" zero config notify --reset # clear the stored preference\n"+
"\n"+
"Flags:\n"+
" --mode <off|bell|notify|both> Notification mechanism (both kinds)\n"+
" --focus <unfocused|always|focused> When the alert fires\n"+
" --reset Clear the stored preference so the TUI effective default applies\n"+
" --json Machine-readable output\n"+
" -h, --help Show this help\n")
return err
}
-481
View File
@@ -1,481 +0,0 @@
package cli
import (
"bytes"
"encoding/json"
"os"
"path/filepath"
"strings"
"testing"
"github.com/Gitlawb/zero/internal/config"
)
// `zero config notify` with no flag and a fresh config: nothing is configured,
// the resolver leaves the fields blank (defaults deliberately live in the TUI,
// not the resolver — maintainer review, PR #1001), and the command reports
// "(default)" for both.
func TestRunConfigNotifyPrintsUnconfiguredAsDefault(t *testing.T) {
configPath := filepath.Join(t.TempDir(), "config.json")
// A valid openai profile so the resolver does not error with
// ErrNoActiveProvider. The notify defaults are applied independently of
// the provider resolution path.
seed := `{
"activeProvider": "openai",
"providers": [{
"name": "openai",
"providerKind": "openai",
"baseUrl": "https://api.openai.com/v1",
"model": "gpt-4.1",
"apiKeyEnv": "OPENAI_API_KEY"
}]
}`
if err := os.WriteFile(configPath, []byte(seed), 0o600); err != nil {
t.Fatalf("seed config: %v", err)
}
deps := commandCenterDeps(t)
deps.userConfigPath = func() (string, error) { return configPath, nil }
deps.resolveConfig = func(workspaceRoot string, overrides config.Overrides) (config.ResolvedConfig, error) {
return config.Resolve(config.ResolveOptions{UserConfigPath: configPath, Env: map[string]string{"OPENAI_API_KEY": "sk-test"}})
}
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"config", "notify"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("exit = %d, want %d: %s", exitCode, exitSuccess, stderr.String())
}
if !strings.Contains(stdout.String(), "mode: (default)") {
t.Errorf("stdout should show unconfigured mode as (default), got: %s", stdout.String())
}
if !strings.Contains(stdout.String(), "focusMode: (default)") {
t.Errorf("stdout should show unconfigured focus as (default), got: %s", stdout.String())
}
}
// `zero config notify --json` emits a machine-readable payload. A configured
// pair round-trips; unconfigured fields are empty strings (never defaults
// filled in), so scripts can distinguish "user chose" from "default".
func TestRunConfigNotifyPrintsJSON(t *testing.T) {
configPath := filepath.Join(t.TempDir(), "config.json")
seed := `{
"activeProvider": "openai",
"providers": [{
"name": "openai",
"providerKind": "openai",
"baseUrl": "https://api.openai.com/v1",
"model": "gpt-4.1",
"apiKeyEnv": "OPENAI_API_KEY"
}],
"notify": {"mode": "bell", "focusMode": "always"}
}`
if err := os.WriteFile(configPath, []byte(seed), 0o600); err != nil {
t.Fatalf("seed config: %v", err)
}
deps := commandCenterDeps(t)
deps.userConfigPath = func() (string, error) { return configPath, nil }
deps.resolveConfig = func(workspaceRoot string, overrides config.Overrides) (config.ResolvedConfig, error) {
return config.Resolve(config.ResolveOptions{UserConfigPath: configPath, Env: map[string]string{"OPENAI_API_KEY": "sk-test"}})
}
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"config", "notify", "--json"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("exit = %d, want %d: %s", exitCode, exitSuccess, stderr.String())
}
var payload map[string]any
if err := json.Unmarshal(stdout.Bytes(), &payload); err != nil {
t.Fatalf("output is not valid JSON: %v\n%s", err, stdout.String())
}
if payload["mode"] != "bell" {
t.Errorf("mode = %v, want bell", payload["mode"])
}
if payload["focusMode"] != "always" {
t.Errorf("focusMode = %v, want always", payload["focusMode"])
}
}
// The unconfigured JSON shape: fields are empty strings, never defaults
// filled in, so scripts can tell "user chose" from "default".
func TestRunConfigNotifyPrintsJSONUnconfigured(t *testing.T) {
configPath := filepath.Join(t.TempDir(), "config.json")
seed := `{
"activeProvider": "openai",
"providers": [{
"name": "openai",
"providerKind": "openai",
"baseUrl": "https://api.openai.com/v1",
"model": "gpt-4.1",
"apiKeyEnv": "OPENAI_API_KEY"
}]
}`
if err := os.WriteFile(configPath, []byte(seed), 0o600); err != nil {
t.Fatalf("seed config: %v", err)
}
deps := commandCenterDeps(t)
deps.userConfigPath = func() (string, error) { return configPath, nil }
deps.resolveConfig = func(workspaceRoot string, overrides config.Overrides) (config.ResolvedConfig, error) {
return config.Resolve(config.ResolveOptions{UserConfigPath: configPath, Env: map[string]string{"OPENAI_API_KEY": "sk-test"}})
}
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"config", "notify", "--json"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("exit = %d, want %d: %s", exitCode, exitSuccess, stderr.String())
}
var payload map[string]any
if err := json.Unmarshal(stdout.Bytes(), &payload); err != nil {
t.Fatalf("output is not valid JSON: %v\n%s", err, stdout.String())
}
if payload["mode"] != "" {
t.Errorf("mode = %v, want empty (unconfigured)", payload["mode"])
}
if payload["focusMode"] != "" {
t.Errorf("focusMode = %v, want empty (unconfigured)", payload["focusMode"])
}
}
// `zero config notify --mode off` writes the new value to disk and prints
// confirmation. The user can read it back by running the command again.
func TestRunConfigNotifyWritesModeChange(t *testing.T) {
configPath := filepath.Join(t.TempDir(), "config.json")
seed := `{
"activeProvider": "openai",
"providers": [{
"name": "openai",
"providerKind": "openai",
"baseUrl": "https://api.openai.com/v1",
"model": "gpt-4.1",
"apiKeyEnv": "OPENAI_API_KEY"
}],
"notify": {"mode": "both", "focusMode": "always"}
}`
if err := os.WriteFile(configPath, []byte(seed), 0o600); err != nil {
t.Fatalf("seed config: %v", err)
}
deps := commandCenterDeps(t)
deps.userConfigPath = func() (string, error) { return configPath, nil }
deps.resolveConfig = func(workspaceRoot string, overrides config.Overrides) (config.ResolvedConfig, error) {
return config.Resolve(config.ResolveOptions{UserConfigPath: configPath, Env: map[string]string{"OPENAI_API_KEY": "sk-test"}})
}
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"config", "notify", "--mode", "off"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("exit = %d, want %d: %s", exitCode, exitSuccess, stderr.String())
}
cfg := readFileConfig(t, configPath)
if cfg.Notify.Mode != "off" {
t.Errorf("Notify.Mode = %q, want off", cfg.Notify.Mode)
}
// A mode-only update must preserve the configured focusMode, not wipe it.
if cfg.Notify.FocusMode != "always" {
t.Errorf("Notify.FocusMode = %q, want preserved %q", cfg.Notify.FocusMode, "always")
}
if !strings.Contains(stdout.String(), "mode: off") {
t.Errorf("stdout should confirm the change, got: %s", stdout.String())
}
}
// A focus-only update preserves the configured mode.
func TestRunConfigNotifyFocusOnlyPreservesMode(t *testing.T) {
configPath := filepath.Join(t.TempDir(), "config.json")
seed := `{
"activeProvider": "openai",
"providers": [{
"name": "openai",
"providerKind": "openai",
"baseUrl": "https://api.openai.com/v1",
"model": "gpt-4.1",
"apiKeyEnv": "OPENAI_API_KEY"
}],
"notify": {"mode": "bell", "focusMode": "unfocused"}
}`
if err := os.WriteFile(configPath, []byte(seed), 0o600); err != nil {
t.Fatalf("seed config: %v", err)
}
deps := commandCenterDeps(t)
deps.userConfigPath = func() (string, error) { return configPath, nil }
deps.resolveConfig = func(workspaceRoot string, overrides config.Overrides) (config.ResolvedConfig, error) {
return config.Resolve(config.ResolveOptions{UserConfigPath: configPath, Env: map[string]string{"OPENAI_API_KEY": "sk-test"}})
}
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"config", "notify", "--focus", "always"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("exit = %d, want %d: %s", exitCode, exitSuccess, stderr.String())
}
cfg := readFileConfig(t, configPath)
if cfg.Notify.Mode != "bell" {
t.Errorf("Notify.Mode = %q, want preserved %q", cfg.Notify.Mode, "bell")
}
if cfg.Notify.FocusMode != "always" {
t.Errorf("Notify.FocusMode = %q, want always", cfg.Notify.FocusMode)
}
}
// Maintainer regression (PR #1001): a partial update must seed from the USER'S
// OWN file, never from the resolved view. A project .zero/config.json setting
// mode=off resolves into the session, but `--focus always` inside that repo
// must NOT copy the project's off into the user's global config.
func TestRunConfigNotifyDoesNotCopyProjectNotifyIntoUserConfig(t *testing.T) {
configPath := filepath.Join(t.TempDir(), "user.json")
projectPath := filepath.Join(t.TempDir(), "project.json")
if err := os.WriteFile(configPath, []byte(`{"activeProvider": "openai"}`), 0o600); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(projectPath, []byte(`{"notify": {"mode": "off"}}`), 0o600); err != nil {
t.Fatal(err)
}
deps := commandCenterDeps(t)
deps.userConfigPath = func() (string, error) { return configPath, nil }
// Resolve EXACTLY the way production does: user config + project config
// merged, so resolved.Notify.Mode is the project's "off".
deps.resolveConfig = func(workspaceRoot string, overrides config.Overrides) (config.ResolvedConfig, error) {
return config.Resolve(config.ResolveOptions{
UserConfigPath: configPath,
ProjectConfigPath: projectPath,
Env: map[string]string{"OPENAI_API_KEY": "sk-test"},
})
}
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"config", "notify", "--focus", "always"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("exit = %d, want %d: %s", exitCode, exitSuccess, stderr.String())
}
cfg := readFileConfig(t, configPath)
// The user's file must keep an unspecified mode unspecified — the
// project's "off" stays in the project file where it belongs.
if cfg.Notify.Mode != "" {
t.Errorf("user Notify.Mode = %q, want blank (project's off must not leak into the user file)", cfg.Notify.Mode)
}
if cfg.Notify.FocusMode != "always" {
t.Errorf("Notify.FocusMode = %q, want always", cfg.Notify.FocusMode)
}
// The project file is untouched.
project := readFileConfig(t, projectPath)
if project.Notify.Mode != "off" {
t.Errorf("project Notify.Mode = %q, want untouched off", project.Notify.Mode)
}
}
// Maintainer regression (PR #1001): with a clean config, `--mode off` must not
// also pin focusMode as an explicit choice — blank means "use the built-in
// default", and a partial update keeps it that way.
func TestRunConfigNotifyDoesNotPinDefaultsAsExplicitChoices(t *testing.T) {
configPath := filepath.Join(t.TempDir(), "config.json")
if err := os.WriteFile(configPath, []byte(`{"activeProvider": "openai"}`), 0o600); err != nil {
t.Fatal(err)
}
deps := commandCenterDeps(t)
deps.userConfigPath = func() (string, error) { return configPath, nil }
deps.resolveConfig = func(workspaceRoot string, overrides config.Overrides) (config.ResolvedConfig, error) {
return config.Resolve(config.ResolveOptions{UserConfigPath: configPath, Env: map[string]string{"OPENAI_API_KEY": "sk-test"}})
}
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"config", "notify", "--mode", "off"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("exit = %d, want %d: %s", exitCode, exitSuccess, stderr.String())
}
cfg := readFileConfig(t, configPath)
if cfg.Notify.Mode != "off" {
t.Errorf("Notify.Mode = %q, want off", cfg.Notify.Mode)
}
if cfg.Notify.FocusMode != "" {
t.Errorf("Notify.FocusMode = %q, want blank (unspecified stays unspecified; the default must not be pinned)", cfg.Notify.FocusMode)
}
}
// `--mode` and `--focus` together update both fields in one call.
func TestRunConfigNotifyWritesModeAndFocus(t *testing.T) {
configPath := filepath.Join(t.TempDir(), "config.json")
seed := `{
"activeProvider": "openai",
"providers": [{
"name": "openai",
"providerKind": "openai",
"baseUrl": "https://api.openai.com/v1",
"model": "gpt-4.1",
"apiKeyEnv": "OPENAI_API_KEY"
}]
}`
if err := os.WriteFile(configPath, []byte(seed), 0o600); err != nil {
t.Fatalf("seed config: %v", err)
}
deps := commandCenterDeps(t)
deps.userConfigPath = func() (string, error) { return configPath, nil }
deps.resolveConfig = func(workspaceRoot string, overrides config.Overrides) (config.ResolvedConfig, error) {
return config.Resolve(config.ResolveOptions{UserConfigPath: configPath, Env: map[string]string{"OPENAI_API_KEY": "sk-test"}})
}
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"config", "notify", "--mode", "both", "--focus", "unfocused"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("exit = %d, want %d: %s", exitCode, exitSuccess, stderr.String())
}
cfg := readFileConfig(t, configPath)
if cfg.Notify.Mode != "both" || cfg.Notify.FocusMode != "unfocused" {
t.Errorf("Notify = %+v, want mode=both focusMode=unfocused", cfg.Notify)
}
}
// `--mode loud` is a usage error. The config must not be mutated.
func TestRunConfigNotifyRejectsInvalidMode(t *testing.T) {
configPath := filepath.Join(t.TempDir(), "config.json")
seed := `{
"activeProvider": "openai",
"providers": [{
"name": "openai",
"providerKind": "openai",
"baseUrl": "https://api.openai.com/v1",
"model": "gpt-4.1",
"apiKeyEnv": "OPENAI_API_KEY"
}],
"notify": {"mode": "off"}
}`
if err := os.WriteFile(configPath, []byte(seed), 0o600); err != nil {
t.Fatalf("seed config: %v", err)
}
deps := commandCenterDeps(t)
deps.userConfigPath = func() (string, error) { return configPath, nil }
deps.resolveConfig = func(workspaceRoot string, overrides config.Overrides) (config.ResolvedConfig, error) {
return config.Resolve(config.ResolveOptions{UserConfigPath: configPath, Env: map[string]string{"OPENAI_API_KEY": "sk-test"}})
}
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"config", "notify", "--mode", "loud"}, &stdout, &stderr, deps)
if exitCode == exitSuccess {
t.Fatalf("expected failure for invalid mode, got success; stdout=%s", stdout.String())
}
cfg := readFileConfig(t, configPath)
if cfg.Notify.Mode != "off" {
t.Errorf("Notify.Mode = %q, want off (unchanged after failed write)", cfg.Notify.Mode)
}
}
// `--reset` blanks both fields so the TUI's effective default applies again
// (an unconfigured headless run stays silent). Useful for "go back to the
// recommended setup" after a custom value.
func TestRunConfigNotifyResetClearsStoredValues(t *testing.T) {
configPath := filepath.Join(t.TempDir(), "config.json")
seed := `{
"activeProvider": "openai",
"providers": [{
"name": "openai",
"providerKind": "openai",
"baseUrl": "https://api.openai.com/v1",
"model": "gpt-4.1",
"apiKeyEnv": "OPENAI_API_KEY"
}],
"notify": {"mode": "off", "focusMode": "always"}
}`
if err := os.WriteFile(configPath, []byte(seed), 0o600); err != nil {
t.Fatalf("seed config: %v", err)
}
deps := commandCenterDeps(t)
deps.userConfigPath = func() (string, error) { return configPath, nil }
deps.resolveConfig = func(workspaceRoot string, overrides config.Overrides) (config.ResolvedConfig, error) {
return config.Resolve(config.ResolveOptions{UserConfigPath: configPath, Env: map[string]string{"OPENAI_API_KEY": "sk-test"}})
}
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"config", "notify", "--reset"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("exit = %d, want %d: %s", exitCode, exitSuccess, stderr.String())
}
cfg := readFileConfig(t, configPath)
if cfg.Notify.Mode != "" || cfg.Notify.FocusMode != "" {
t.Errorf("Notify after reset = %+v, want empty (defaults apply)", cfg.Notify)
}
}
// CodeRabbit regression (PR #1001): the command manages a user preference and
// must not require provider resolution. A brand-new user with NO provider
// configured (the resolver would return ErrNoActiveProvider) can still read,
// set, and reset their notification preference.
func TestRunConfigNotifyWorksWithoutAnyProviderConfigured(t *testing.T) {
configPath := filepath.Join(t.TempDir(), "config.json")
if err := os.WriteFile(configPath, []byte(`{}`), 0o600); err != nil {
t.Fatal(err)
}
// resolveConfig fails the way the real resolver does for a fresh user —
// the command must never call it, so a panic-free stub is enough to prove
// the point; use the failing resolver to catch any regression to the
// resolve-first shape.
deps := commandCenterDeps(t)
deps.userConfigPath = func() (string, error) { return configPath, nil }
deps.resolveConfig = func(workspaceRoot string, overrides config.Overrides) (config.ResolvedConfig, error) {
return config.Resolve(config.ResolveOptions{UserConfigPath: configPath, Env: map[string]string{}})
}
// Read works and shows defaults.
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"config", "notify"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("read: exit = %d, want %d: %s", exitCode, exitSuccess, stderr.String())
}
if !strings.Contains(stdout.String(), "mode: (default)") {
t.Errorf("read should show (default), got: %s", stdout.String())
}
// Write works.
stdout.Reset()
exitCode = runWithDeps([]string{"config", "notify", "--mode", "bell", "--focus", "always"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("write: exit = %d, want %d: %s", exitCode, exitSuccess, stderr.String())
}
cfg := readFileConfig(t, configPath)
if cfg.Notify.Mode != "bell" || cfg.Notify.FocusMode != "always" {
t.Fatalf("write: Notify = %+v, want bell/always", cfg.Notify)
}
// JSON read reflects the stored pair.
stdout.Reset()
exitCode = runWithDeps([]string{"config", "notify", "--json"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("json: exit = %d, want %d: %s", exitCode, exitSuccess, stderr.String())
}
var payload map[string]any
if err := json.Unmarshal(stdout.Bytes(), &payload); err != nil {
t.Fatalf("json output invalid: %v\n%s", err, stdout.String())
}
if payload["mode"] != "bell" || payload["focusMode"] != "always" {
t.Errorf("json = mode %v focus %v, want bell/always", payload["mode"], payload["focusMode"])
}
// Reset works.
stdout.Reset()
exitCode = runWithDeps([]string{"config", "notify", "--reset"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("reset: exit = %d, want %d: %s", exitCode, exitSuccess, stderr.String())
}
cfg = readFileConfig(t, configPath)
if cfg.Notify.Mode != "" || cfg.Notify.FocusMode != "" {
t.Errorf("reset: Notify = %+v, want empty", cfg.Notify)
}
}
// `zero config` (no subcommand) still works after the dispatch change.
func TestRunConfigSummaryStillWorks(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"config"}, &stdout, &stderr, commandCenterDeps(t))
if exitCode != exitSuccess {
t.Fatalf("exit = %d, want %d: %s", exitCode, exitSuccess, stderr.String())
}
if !strings.Contains(stdout.String(), "Config") {
t.Errorf("stdout should show the config summary, got: %s", stdout.String())
}
}
-116
View File
@@ -2,10 +2,6 @@ package cli
import (
"bytes"
"io"
"os"
"path/filepath"
"strings"
"sync"
"testing"
"time"
@@ -43,115 +39,3 @@ func TestFireJobClaimPreventsDoubleFire(t *testing.T) {
t.Fatalf("NextRunAt must advance past now, got %s", d.NextRunAt)
}
}
func TestFireJobOverlappingAdjacentSlots(t *testing.T) {
for _, olderFirst := range []bool{false, true} {
name := "newer completes first"
if olderFirst {
name = "older completes first"
}
t.Run(name, func(t *testing.T) {
root := t.TempDir()
stores := []*cron.Store{
cron.NewStore(cron.StoreOptions{RootDir: root}),
cron.NewStore(cron.StoreOptions{RootDir: root}),
}
now := time.Date(2026, 6, 9, 9, 0, 0, 0, time.UTC)
job, err := stores[0].Add(cron.Job{Expr: "* * * * *", Prompt: "fire me", NextRunAt: now, FireCount: 7})
if err != nil {
t.Fatal(err)
}
var release, done [2]chan struct{}
for i := range stores {
release[i], done[i] = make(chan struct{}), make(chan struct{})
// Always release blocked runners, including on assertion failures.
defer func() {
select {
case <-release[i]:
default:
close(release[i])
}
<-done[i]
}()
started := make(chan struct{})
go func() {
defer close(done[i])
fireJob(stores[i], func() time.Time { return now.Add(time.Duration(i) * time.Minute) }, job, io.Discard, io.Discard,
func(_ []string, _, _ io.Writer) int {
close(started)
<-release[i]
return i // Include a failed execution in the accounting.
})
}()
select {
case <-started:
case <-done[i]:
t.Fatal("runner did not claim its slot")
}
}
order := []int{1, 0}
if olderFirst {
order = []int{0, 1}
}
for n, i := range order {
close(release[i])
<-done[i]
current, err := stores[0].Get(job.ID)
if err != nil {
t.Fatal(err)
}
if current.FireCount != 8+n || !current.NextRunAt.Equal(now.Add(2*time.Minute)) {
t.Fatalf("after completion %d: FireCount=%d NextRunAt=%s; want count %d and next %s", n+1, current.FireCount, current.NextRunAt, 8+n, now.Add(2*time.Minute))
}
claimed, err := claimFire(stores[1], now.Add(time.Minute), job.ID)
if err != nil || claimed {
t.Fatalf("already claimed slot became available: claimed=%v err=%v", claimed, err)
}
}
runs, err := stores[0].Runs(job.ID, 0)
if err != nil || len(runs) != 2 {
t.Fatalf("runs=%v err=%v", runs, err)
}
for n, i := range order {
if !runs[n].At.Equal(now.Add(time.Duration(i)*time.Minute)) || runs[n].ExitCode != i {
t.Fatalf("unexpected run record: %+v", runs[n])
}
}
})
}
}
func TestFireJobCompletionReadErrorDoesNotOverwriteState(t *testing.T) {
root := t.TempDir()
store := cron.NewStore(cron.StoreOptions{RootDir: root})
now := time.Date(2026, 6, 9, 9, 0, 0, 0, time.UTC)
job, err := store.Add(cron.Job{Expr: "* * * * *", Prompt: "fire me", NextRunAt: now})
if err != nil {
t.Fatal(err)
}
metadata := filepath.Join(root, job.ID, "metadata.json")
var stdout, stderr bytes.Buffer
fireJob(store, func() time.Time { return now }, job, &stdout, &stderr,
func(_ []string, _, _ io.Writer) int {
// A corrupt read must not cause completion to replace unknown newer state
// with the stale pre-execution job. This works without permission tricks.
if err := os.WriteFile(metadata, []byte("unreadable job"), 0o600); err != nil {
t.Fatal(err)
}
return 0
})
data, err := os.ReadFile(metadata)
if err != nil || string(data) != "unreadable job" {
t.Fatalf("completion overwrote unreadable state: %q, err=%v", data, err)
}
if !strings.Contains(stderr.String(), "failed to persist job state") {
t.Fatalf("missing persistence warning: %s", stderr.String())
}
if want := "fired " + job.ID + " -> exit 0 (next: unknown)\n"; stdout.String() != want {
t.Fatalf("completion summary = %q; want %q", stdout.String(), want)
}
runs, err := store.Runs(job.ID, 0)
if err != nil || len(runs) != 1 || runs[0].ExitCode != 0 {
t.Fatalf("execution should still be recorded: runs=%v err=%v", runs, err)
}
}
+38 -19
View File
@@ -129,15 +129,17 @@ func reconcileOverdue(store *cron.Store, now func() time.Time, ids []string, std
}
}
// fireJob claims one job, runs it via the exec runner, and records the outcome.
// Separate scheduler processes can have adjacent slots executing concurrently.
// fireJob runs one job via the exec runner, records the outcome, advances the
// schedule, and persists. The foreground loop is single-goroutine, so the
// previous fire has already returned before the next tick — no overlap.
func fireJob(store *cron.Store, now func() time.Time, job cron.Job, stdout io.Writer, stderr io.Writer, exec execRunner) {
fired := now()
// Atomically claim this fire before running it so concurrent schedulers (or
// --once overlapping the forever loop) can't both execute the same due job: the
// winner advances NextRunAt under the per-job lock, and a loser sees the
// advanced time and skips. Only the claim writes schedule state; completion
// must not overwrite a later slot's claim with this run's stale job copy.
// advanced time and skips. The post-exec advance recomputes the same
// NextRunAt (sched.Next(fired) is deterministic), so the claim serializes the
// fire decision without changing any persisted schedule state (M9).
claimed, claimErr := claimFire(store, fired, job.ID)
if claimErr != nil {
if errors.Is(claimErr, cron.ErrJobNotFound) {
@@ -176,25 +178,44 @@ func fireJob(store *cron.Store, now func() time.Time, job cron.Job, stdout io.Wr
rec.Error = cronTruncate(detail, 500)
}
// Explain schedules that claimFire paused because they cannot advance.
job.FireCount++
// Advance the schedule. If the expression can no longer produce a future run
// (became invalid, or is an impossible spec whose Next is zero), pause the job
// so it cannot re-fire on every tick.
if sched, perr := cron.Parse(job.Expr); perr != nil {
job.Status = cron.StatusPaused
if rec.Error == "" {
rec.Error = "invalid schedule; job paused: " + perr.Error()
}
} else if nxt := sched.Next(fired.Truncate(time.Minute)); nxt.IsZero() {
job.Status = cron.StatusPaused
if rec.Error == "" {
rec.Error = "schedule no longer fires; job paused"
}
} else {
// Minute-aligned input keeps this advance identical to claimFire's and lets
// the DST fall-back collapse guard engage (AUDIT-M4).
job.NextRunAt = nxt
}
// Count this completion against current state under the cross-process lock.
// The claim already advanced or paused the schedule. Leave that state alone:
// another scheduler may have claimed a later slot or a user may have edited it.
// Re-read and persist the advanced state ATOMICALLY under the per-job lock. The
// job may have been paused or removed while it executed, and this in-memory copy
// is stale from tick start. Mutate closes the read-modify-write window a single
// scheduler alone could only narrow (Store.Mutate took a cross-process lock): a
// concurrent scheduler or an external pause/remove landing here can no longer be
// clobbered.
persisted, err := store.Mutate(job.ID, func(current cron.Job, readErr error) (cron.Job, error) {
if readErr != nil {
// Never replace unknown newer state with the pre-execution snapshot.
return current, readErr
// A transient read failure (IO/permission) is NOT removal — persist the
// computed next state anyway so the schedule advances and the job does not
// re-fire next tick.
fmt.Fprintf(stderr, "warning: could not re-read job %s before persist: %v\n", job.ID, readErr)
return job, nil
}
current.FireCount = job.FireCount
current.NextRunAt = job.NextRunAt
if current.Status != cron.StatusPaused {
current.Status = job.Status
}
current.FireCount++
return current, nil
})
if errors.Is(err, cron.ErrJobNotFound) {
@@ -208,21 +229,19 @@ func fireJob(store *cron.Store, now func() time.Time, job cron.Job, stdout io.Wr
if aerr := store.AppendRun(job.ID, rec); aerr != nil {
fmt.Fprintf(stderr, "warning: failed to record run for %s: %v\n", job.ID, aerr)
}
next := "unknown"
if err != nil {
fmt.Fprintf(stderr, "warning: failed to persist job state for %s: %v\n", job.ID, err)
} else {
next = formatCronTime(persisted.NextRunAt)
persisted = job
}
fmt.Fprintf(stdout, "fired %s -> exit %d (next: %s)\n", job.ID, code, next)
fmt.Fprintf(stdout, "fired %s -> exit %d (next: %s)\n", job.ID, code, formatCronTime(persisted.NextRunAt))
}
// claimFire atomically claims a due job's fire under the per-job lock and reports
// whether THIS caller should run it. It returns false only when the job was paused
// externally or another scheduler already advanced NextRunAt past `fired` (so a
// valid recurring job never double-fires). A valid due job has its NextRunAt
// advanced here to claim the slot; an invalid or exhausted schedule is paused
// here and the caller still fires once (preserving single-scheduler behavior).
// advanced here to claim the slot; an invalid or exhausted schedule is left for the
// caller to fire-once and pause (preserving existing single-scheduler behavior).
func claimFire(store *cron.Store, fired time.Time, id string) (bool, error) {
claimed := true
_, err := store.Mutate(id, func(current cron.Job, readErr error) (cron.Job, error) {
@@ -249,8 +268,8 @@ func claimFire(store *cron.Store, fired time.Time, id string) (bool, error) {
}
// Unparseable or unadvanceable schedule (an impossible spec whose Next is
// zero): pause the job inside this same locked claim so a concurrent scheduler
// sees a non-active job and does NOT also fire it. The winner still fires once.
// Without this, both callers leave
// sees a non-active job and does NOT also fire it. The winner still fires once
// and fireJob's post-exec keeps it paused. Without this, both callers leave
// NextRunAt unchanged, both see the job due, and both fire the same slot
// (AUDIT-M5).
current.Status = cron.StatusPaused
+2 -2
View File
@@ -53,7 +53,7 @@ func TestCronRunOnceFiresDueJobs(t *testing.T) {
if d.FireCount != 1 || !d.NextRunAt.After(now) {
t.Fatalf("due job not advanced: %+v", d)
}
runs, _ := store.Runs(due.ID, 0)
runs, _ := store.Runs(due.ID)
if len(runs) != 1 {
t.Fatalf("expected 1 run record, got %d", len(runs))
}
@@ -146,7 +146,7 @@ func TestCronRunPausesUnadvanceableJob(t *testing.T) {
if d.Status != cron.StatusPaused {
t.Fatalf("unadvanceable job must be paused, got status=%q", d.Status)
}
runs, _ := store.Runs(job.ID, 0)
runs, _ := store.Runs(job.ID)
if len(runs) != 1 || runs[0].Error == "" {
t.Fatalf("expected one run record with an error, got %+v", runs)
}
+1 -15
View File
@@ -181,7 +181,7 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
// Refresh the models.dev pricing/limits cache in the background when stale;
// the overlay is read at registry construction from the cache file, so this
// benefits the next run and never blocks or fails this one.
modelregistry.StartModelsDevRefresh()
go func() { _ = modelregistry.RefreshModelsDevCache(context.Background()) }()
// A mode seeds model/effort/max-turns/tool filters as a preset. Expand it up
// front — before tool-filter validation and the --list-tools branch — so a
@@ -347,17 +347,6 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
} else if notice != "" {
_, _ = fmt.Fprintln(stderr, "[zero] "+notice)
}
// Before any tool runs, and before the MCP servers and plugins below start
// through the same sandbox, so a run that goes ahead with reduced isolation
// says so where its output starts, not only when someone runs `zero doctor`.
// A notice that cannot be written stops the run, as the image and effort
// notices further down do: going ahead would mean reduced isolation and no
// warning delivered.
if notice := sandboxEngine.DegradedNotice(); notice != "" {
if _, err := fmt.Fprintln(stderr, "[zero] "+notice); err != nil {
return exitCrash
}
}
executionRunner.SetPreparer(sandboxEngine)
if permissionMode != agent.PermissionModePlan {
mcpRuntime, mcpSkip, err = registerMCPToolsForWorkspace(context.Background(), workspaceRoot, registry, deps, execMCPAutonomy(options), trustRoot, executionRunner)
@@ -671,9 +660,6 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
ContextWindowFor: func(modelID string) int {
return modelregistry.AgentContextWindow(modelContextWindow(modelRegistry, modelID))
},
SupportsVision: func(modelID string) bool {
return modelregistry.SupportsVision(modelRegistry, modelID)
},
ReasoningEffort: forwardEffort,
Trace: traceRecorder,
Cwd: workspaceRoot,
+11 -11
View File
@@ -70,8 +70,8 @@ func TestRunExecListsFilteredToolsWithoutPromptOrProvider(t *testing.T) {
if exitCode != exitSuccess {
t.Fatalf("expected exit code %d, got %d: %s", exitSuccess, exitCode, stderr.String())
}
if stderr.String() != hostSandboxNotice(t) {
t.Fatalf("expected no stderr beyond this host's sandbox notice, got %q", stderr.String())
if stderr.Len() != 0 {
t.Fatalf("expected empty stderr, got %q", stderr.String())
}
if providerBuilt {
t.Fatal("provider should not be constructed for --list-tools")
@@ -108,8 +108,8 @@ func TestRunExecListsToolsAsStreamJSONWhenRequested(t *testing.T) {
if exitCode != exitSuccess {
t.Fatalf("expected exit code %d, got %d: %s", exitSuccess, exitCode, stderr.String())
}
if stderr.String() != hostSandboxNotice(t) {
t.Fatalf("expected no stderr beyond this host's sandbox notice, got %q", stderr.String())
if stderr.Len() != 0 {
t.Fatalf("expected empty stderr, got %q", stderr.String())
}
if providerBuilt {
t.Fatal("provider should not be constructed for --list-tools")
@@ -176,8 +176,8 @@ func TestRunExecListsMCPToolsWithoutProviderConstruction(t *testing.T) {
if !closed {
t.Fatal("MCP runtime was not closed after --list-tools")
}
if stderr.String() != hostSandboxNotice(t) {
t.Fatalf("expected no stderr beyond this host's sandbox notice, got %q", stderr.String())
if stderr.Len() != 0 {
t.Fatalf("expected empty stderr, got %q", stderr.String())
}
if !strings.Contains(stdout.String(), "mcp_docs_lookup") || !strings.Contains(stdout.String(), "Lookup documentation") {
t.Fatalf("expected MCP tool in list output, got %q", stdout.String())
@@ -323,8 +323,8 @@ func TestRunExecStreamJSONOutputsRunEndAndRecordsSession(t *testing.T) {
if exitCode != exitSuccess {
t.Fatalf("expected exit code %d, got %d: %s", exitSuccess, exitCode, stderr.String())
}
if stderr.String() != hostSandboxNotice(t) {
t.Fatalf("expected no stderr beyond this host's sandbox notice, got %q", stderr.String())
if stderr.Len() != 0 {
t.Fatalf("expected empty stderr, got %q", stderr.String())
}
events := decodeJSONLines(t, stdout.String())
@@ -388,8 +388,8 @@ func TestRunExecStreamJSONEmitsAndRecordsPermissionEvents(t *testing.T) {
if exitCode != exitSuccess {
t.Fatalf("expected exit code %d, got %d: %s", exitSuccess, exitCode, stderr.String())
}
if stderr.String() != hostSandboxNotice(t) {
t.Fatalf("expected no stderr beyond this host's sandbox notice, got %q", stderr.String())
if stderr.Len() != 0 {
t.Fatalf("expected empty stderr, got %q", stderr.String())
}
events := decodeJSONLines(t, stdout.String())
@@ -460,7 +460,7 @@ func TestRunExecStreamJSONEmitsAndRecordsPermissionEvents(t *testing.T) {
if exitCode != exitSuccess {
t.Fatalf("approved exec exit = %d, stderr = %q", exitCode, stderr.String())
}
if stderr.String() != hostSandboxNotice(t) {
if stderr.String() != "" {
t.Fatalf("approved exec wrote stderr = %q", stderr.String())
}
approvedEvents := decodeJSONLines(t, stdout.String())
+6 -9
View File
@@ -118,15 +118,12 @@ func runExecSpecDraft(run execSpecDraftRun) int {
hookDispatcher, hookSkip := newHookDispatcher(run.workspaceRoot, run.trustRoot, execution.NewRunner(run.sandboxEngine))
emitTrustNotice(run.stderr, hookSkip, run.mcpSkip)
result, err := agent.Run(runCtx, run.prompt, run.provider, agent.Options{
MaxTurns: run.resolved.MaxTurns,
ContextWindow: resolveAgentContextWindow(runCtx, run.modelRegistry, run.resolved.Provider),
SessionID: draftSession.SessionID,
SessionTitle: run.sessionTitle,
ProviderName: run.resolved.Provider.Name,
Model: run.resolved.Provider.Model,
SupportsVision: func(modelID string) bool {
return modelregistry.SupportsVision(run.modelRegistry, modelID)
},
MaxTurns: run.resolved.MaxTurns,
ContextWindow: resolveAgentContextWindow(runCtx, run.modelRegistry, run.resolved.Provider),
SessionID: draftSession.SessionID,
SessionTitle: run.sessionTitle,
ProviderName: run.resolved.Provider.Name,
Model: run.resolved.Provider.Model,
ReasoningEffort: run.reasoningEffort,
Profile: run.profilePolicy,
Cwd: run.workspaceRoot,
+20 -20
View File
@@ -568,8 +568,8 @@ func TestRunExecAcceptsLegacyModelProfileFlags(t *testing.T) {
if !strings.Contains(stdout, "hello") {
t.Fatalf("expected prompt output, got %q", stdout)
}
if stderr != hostSandboxNotice(t) {
t.Fatalf("expected no stderr beyond this host's sandbox notice, got %q", stderr)
if stderr != "" {
t.Fatalf("expected empty stderr, got %q", stderr)
}
}
@@ -696,8 +696,8 @@ func TestRunExecAssemblesInlineAndFilePromptRelativeToCwd(t *testing.T) {
if !strings.Contains(stdout, "inline prompt\n\nfile prompt") {
t.Fatalf("expected inline and file prompt joined by blank line, got %q", stdout)
}
if stderr != hostSandboxNotice(t) {
t.Fatalf("expected no stderr beyond this host's sandbox notice, got %q", stderr)
if stderr != "" {
t.Fatalf("expected empty stderr, got %q", stderr)
}
}
@@ -715,8 +715,8 @@ func TestRunExecAcceptsFileOnlyPrompt(t *testing.T) {
if !strings.Contains(stdout, "file only prompt") {
t.Fatalf("expected file prompt output, got %q", stdout)
}
if stderr != hostSandboxNotice(t) {
t.Fatalf("expected no stderr beyond this host's sandbox notice, got %q", stderr)
if stderr != "" {
t.Fatalf("expected empty stderr, got %q", stderr)
}
}
@@ -798,8 +798,8 @@ func TestRunExecJSONOutputsNDJSONEvents(t *testing.T) {
if exitCode != 0 {
t.Fatalf("expected exit code 0, got %d: %s", exitCode, stderr)
}
if stderr != hostSandboxNotice(t) {
t.Fatalf("expected no stderr beyond this host's sandbox notice, got %q", stderr)
if stderr != "" {
t.Fatalf("expected empty stderr, got %q", stderr)
}
events := decodeJSONLines(t, stdout)
@@ -833,8 +833,8 @@ func TestRunExecResolvesCanonicalModelAlias(t *testing.T) {
if exitCode != 0 {
t.Fatalf("expected exit code 0, got %d: %s", exitCode, stderr)
}
if stderr != hostSandboxNotice(t) {
t.Fatalf("expected no stderr beyond this host's sandbox notice for active model, got %q", stderr)
if stderr != "" {
t.Fatalf("expected empty stderr for active model, got %q", stderr)
}
events := decodeJSONLines(t, stdout)
if got := events[0]["model"]; got != "gpt-4.1" {
@@ -926,8 +926,8 @@ func TestRunExecJSONUnsafeOutputsWarningEvent(t *testing.T) {
if exitCode != 0 {
t.Fatalf("expected exit code 0, got %d: %s", exitCode, stderr)
}
if stderr != hostSandboxNotice(t) {
t.Fatalf("expected no stderr beyond this host's sandbox notice, got %q", stderr)
if stderr != "" {
t.Fatalf("expected empty stderr, got %q", stderr)
}
events := decodeJSONLines(t, stdout)
@@ -970,10 +970,10 @@ func TestRunExecUsesProjectConfigAndOpenAICompatibleProvider(t *testing.T) {
"name": "local",
"provider_kind": "openai-compatible",
"base_url": "` + server.URL + `",
"api_key": "sk-local",
"model": "local-model"
}]
}`
"api_key": "sk-local",
"model": "local-model"
}]
}`
if err := os.WriteFile(filepath.Join(configDir, "config.json"), []byte(writeConfig), 0o600); err != nil {
t.Fatal(err)
}
@@ -988,8 +988,8 @@ func TestRunExecUsesProjectConfigAndOpenAICompatibleProvider(t *testing.T) {
if strings.TrimSpace(stdout.String()) != "provider ok" {
t.Fatalf("stdout = %q, want provider response", stdout.String())
}
if stderr.String() != hostSandboxNotice(t) {
t.Fatalf("expected no stderr beyond this host's sandbox notice, got %q", stderr.String())
if stderr.Len() != 0 {
t.Fatalf("expected empty stderr, got %q", stderr.String())
}
if gotAuth != "Bearer sk-local" {
t.Fatalf("Authorization = %q, want project config token", gotAuth)
@@ -1204,8 +1204,8 @@ func TestRunExecAutoHighEmitsUnsafeWarning(t *testing.T) {
if exitCode != exitSuccess {
t.Fatalf("expected exit code %d, got %d: %s", exitSuccess, exitCode, stderr)
}
if stderr != hostSandboxNotice(t) {
t.Fatalf("expected no stderr beyond this host's sandbox notice, got %q", stderr)
if stderr != "" {
t.Fatalf("expected empty stderr, got %q", stderr)
}
events := decodeJSONLines(t, stdout)
+1 -1
View File
@@ -69,7 +69,7 @@ func TestRunExecListToolsHonorsDisabledLocalControlConfig(t *testing.T) {
if exitCode != exitSuccess {
t.Fatalf("exitCode = %d stdout=%s stderr=%s", exitCode, stdout.String(), stderr.String())
}
if stderr.String() != hostSandboxNotice(t) {
if stderr.Len() != 0 {
t.Fatalf("stderr = %q, want empty", stderr.String())
}
var payload struct {
@@ -1,237 +0,0 @@
package cli
import (
"bytes"
"context"
"encoding/json"
"strings"
"testing"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/mcp"
"github.com/Gitlawb/zero/internal/sandbox"
"github.com/Gitlawb/zero/internal/tools"
"github.com/Gitlawb/zero/internal/tui"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
// hostSandboxNotice is the degraded-sandbox line exec prints on THIS machine for
// a run under the default sandbox policy, or "" where the sandbox is not
// degraded. Exec tests that expect a quiet stderr compare against it rather than
// against "": unless a test pins a backend, the run gets the host's own, and a
// Linux runner without the sandbox helper is degraded where macOS and Windows
// are not.
func hostSandboxNotice(t *testing.T) string {
t.Helper()
engine := sandbox.NewEngine(sandbox.EngineOptions{
WorkspaceRoot: t.TempDir(),
Policy: sandbox.DefaultPolicy(),
Backend: sandbox.SelectBackend(sandbox.BackendOptions{}),
})
notice := engine.DegradedNotice()
if notice == "" {
return ""
}
return "[zero] " + notice + "\n"
}
// clearSandboxNestingMarkers keeps a test that pins the notice from inheriting
// the nesting markers when the suite itself runs inside a Zero sandbox.
func clearSandboxNestingMarkers(t *testing.T) {
t.Helper()
t.Setenv(sandbox.EnvSandboxed, "")
t.Setenv(sandbox.EnvSandboxBackend, "")
}
// unavailableTestSandbox is the backend a Linux box without the helper gets,
// pinned so a test sees the degraded case on every platform.
func unavailableTestSandbox(sandbox.BackendOptions) sandbox.Backend {
return sandbox.Backend{Name: sandbox.BackendUnavailable, Platform: "linux", Fallback: true, Message: "Linux sandbox helper is not available"}
}
const unavailableTestSandboxNotice = "[zero] Sandbox enforcement is degraded: Linux sandbox helper is not available. See `zero sandbox policy --effective`.\n"
func runExecWithSandbox(t *testing.T, args []string, backend func(sandbox.BackendOptions) sandbox.Backend, sandboxConfig config.SandboxConfig) (int, string, string) {
t.Helper()
isolateCLIUserState(t)
var stdout, stderr bytes.Buffer
cwd := t.TempDir()
exitCode := runWithDeps(args, &stdout, &stderr, appDeps{
getwd: func() (string, error) { return cwd, nil },
resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) {
return config.ResolvedConfig{
ActiveProvider: "echo",
Provider: config.ProviderProfile{
Name: "echo",
ProviderKind: config.ProviderKindOpenAICompatible,
BaseURL: "http://127.0.0.1/v1",
Model: "echo-model",
},
MaxTurns: 3,
Sandbox: sandboxConfig,
}, nil
},
newProvider: func(config.ProviderProfile) (zeroruntime.Provider, error) {
return echoExecProvider{}, nil
},
selectSandboxBackend: backend,
})
return exitCode, stdout.String(), stderr.String()
}
// A RUN THAT GOES AHEAD WITH REDUCED ISOLATION SAYS SO. Before this only `zero
// doctor` and `zero sandbox policy` did, so `zero exec` ran every tool with the
// native sandbox effectively off and printed nothing about it. #1041.
func TestExecSaysOnStderrWhenTheSandboxIsDegraded(t *testing.T) {
clearSandboxNestingMarkers(t)
exitCode, _, stderr := runExecWithSandbox(t, []string{"exec", "hello"}, unavailableTestSandbox, config.SandboxConfig{})
if exitCode != 0 {
t.Fatalf("exit code = %d, stderr %q", exitCode, stderr)
}
if stderr != unavailableTestSandboxNotice {
t.Fatalf("stderr = %q, want exactly the degraded notice %q", stderr, unavailableTestSandboxNotice)
}
}
// The notice goes to stderr in the JSON modes as well, so stdout stays one JSON
// object per line for whatever reads it.
func TestExecKeepsTheDegradedNoticeOffStdoutInJSONMode(t *testing.T) {
clearSandboxNestingMarkers(t)
exitCode, stdout, stderr := runExecWithSandbox(t, []string{"exec", "--output-format", "json", "hello"}, unavailableTestSandbox, config.SandboxConfig{})
if exitCode != 0 {
t.Fatalf("exit code = %d, stderr %q", exitCode, stderr)
}
if stderr != unavailableTestSandboxNotice {
t.Fatalf("stderr = %q, want exactly the degraded notice", stderr)
}
for _, line := range strings.Split(strings.TrimSpace(stdout), "\n") {
var event map[string]any
if err := json.Unmarshal([]byte(line), &event); err != nil {
t.Fatalf("stdout carries a line that is not JSON: %q (%v)", line, err)
}
}
}
// A sandbox the user turned off is their choice, not a degradation.
func TestExecSaysNothingWhenTheSandboxIsTurnedOff(t *testing.T) {
clearSandboxNestingMarkers(t)
off := false
exitCode, _, stderr := runExecWithSandbox(t, []string{"exec", "hello"}, unavailableTestSandbox, config.SandboxConfig{Enabled: &off})
if exitCode != 0 {
t.Fatalf("exit code = %d, stderr %q", exitCode, stderr)
}
if stderr != "" {
t.Fatalf("stderr = %q, want nothing for a disabled sandbox", stderr)
}
}
// countingExecProvider is echoExecProvider, counting the completions it streams.
type countingExecProvider struct {
echoExecProvider
streamed *int
}
func (provider countingExecProvider) StreamCompletion(ctx context.Context, request zeroruntime.CompletionRequest) (<-chan zeroruntime.StreamEvent, error) {
*provider.streamed++
return provider.echoExecProvider.StreamCompletion(ctx, request)
}
// A degraded notice that cannot be written stops the run before the model is
// asked for anything, as the image and effort notices do, instead of going ahead
// with reduced isolation and no warning delivered.
func TestExecStopsWhenTheDegradedNoticeCannotBeWritten(t *testing.T) {
clearSandboxNestingMarkers(t)
isolateCLIUserState(t)
var stdout bytes.Buffer
cwd := t.TempDir()
streamed := 0
exitCode := runWithDeps([]string{"exec", "hello"}, &stdout, failingWriter{}, appDeps{
getwd: func() (string, error) { return cwd, nil },
resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) {
return config.ResolvedConfig{
ActiveProvider: "echo",
Provider: config.ProviderProfile{
Name: "echo",
ProviderKind: config.ProviderKindOpenAICompatible,
BaseURL: "http://127.0.0.1/v1",
Model: "echo-model",
},
MaxTurns: 3,
}, nil
},
newProvider: func(config.ProviderProfile) (zeroruntime.Provider, error) {
return countingExecProvider{streamed: &streamed}, nil
},
selectSandboxBackend: unavailableTestSandbox,
})
if exitCode != exitCrash {
t.Fatalf("exit code = %d, want %d when the degraded notice cannot be written", exitCode, exitCrash)
}
if streamed != 0 {
t.Fatalf("the model was asked %d times: the run went ahead without its warning", streamed)
}
}
// launchTUIWithSandbox runs the root command the way a bare `zero` does, with
// the sandbox backend pinned to the unavailable one, and returns the options the
// TUI would have started with.
func launchTUIWithSandbox(t *testing.T, sandboxConfig config.SandboxConfig) tui.Options {
t.Helper()
isolateCLIUserState(t)
var stdout, stderr bytes.Buffer
cwd := t.TempDir()
var launched tui.Options
exitCode := runWithDeps([]string{}, &stdout, &stderr, appDeps{
getwd: func() (string, error) { return cwd, nil },
resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) {
return config.ResolvedConfig{MaxTurns: 12, Sandbox: sandboxConfig}, nil
},
newProvider: func(config.ProviderProfile) (zeroruntime.Provider, error) {
t.Fatal("newProvider should not be called without a resolved provider")
return nil, nil
},
registerMCPTools: func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
return noopMCPRuntime{}, nil
},
selectSandboxBackend: unavailableTestSandbox,
runTUI: func(_ context.Context, options tui.Options) int {
launched = options
return 0
},
})
if exitCode != 0 {
t.Fatalf("exit code = %d, stderr %q", exitCode, stderr.String())
}
return launched
}
// The TUI shows the same notice when the session opens, taken from the engine
// its commands run through.
func TestTUILaunchCarriesTheDegradedNotice(t *testing.T) {
clearSandboxNestingMarkers(t)
for _, testCase := range []struct {
name string
sandbox config.SandboxConfig
want bool
}{
{name: "degraded", want: true},
{name: "turned off", sandbox: config.SandboxConfig{Enabled: func() *bool { off := false; return &off }()}},
} {
t.Run(testCase.name, func(t *testing.T) {
launched := launchTUIWithSandbox(t, testCase.sandbox)
var notices []string
for _, notice := range launched.StartupNotices {
if strings.TrimSpace(notice) != "" {
notices = append(notices, notice)
}
}
want := strings.TrimSuffix(strings.TrimPrefix(unavailableTestSandboxNotice, "[zero] "), "\n")
if testCase.want && (len(notices) != 1 || notices[0] != want) {
t.Fatalf("StartupNotices = %q, want exactly %q", notices, want)
}
if !testCase.want && len(notices) != 0 {
t.Fatalf("StartupNotices = %q, want none", notices)
}
})
}
}
@@ -62,10 +62,6 @@ func TestCLILaunchHelpersLeaveUserStateAlone(t *testing.T) {
{"execAdvertisesEscalateModel", func(t *testing.T) {
_ = execAdvertisesEscalateModel(t, []string{"--allow-escalation", "exec", "say hi"})
}},
{"runExecWithSandbox", func(t *testing.T) {
_, _, _ = runExecWithSandbox(t, []string{"exec", "hello"}, unavailableTestSandbox, config.SandboxConfig{})
}},
{"launchTUIWithSandbox", func(t *testing.T) { _ = launchTUIWithSandbox(t, config.SandboxConfig{}) }},
} {
t.Run(testCase.name, func(t *testing.T) {
seedRoot, configPath, seeded, before := seedEmulatedUserConfig(t)
-80
View File
@@ -244,86 +244,6 @@ func petJSONObject(encodedPet []byte) []byte {
return append(result, '}')
}
// setNotifyJSONObject replaces only the notify member's value, keeping the
// original bytes of every unrelated member so a notification save cannot
// reorder, reformat, or drop the user's other settings — including values whose
// EXPLICIT presence matters (tools.deferThreshold: 0, mcp servers'
// disabled: false) and keys the typed FileConfig does not model at all. The
// typed serializer cannot round-trip explicit zeros (omitempty drops them), so
// notification writes go through this byte-preserving editor instead
// (maintainer review, PR #1001). An empty notify block removes the member
// entirely so a reset leaves no `"notify": {}` husk behind.
//
// A hand-edited file may contain DUPLICATE notify members. Go's decoder
// MERGES the fields of duplicate object members into the same struct, so an
// earlier duplicate's field survives a replacement that omits it — e.g.
// {"notify":{"mode":"off"},"notify":{"focusMode":"focused"}} + a
// replacement {"focusMode":"always"} would leave mode "off" effective.
// The nonempty path therefore removes EVERY duplicate and inserts the
// replacement fresh (full-replace semantics: nothing inherited), and the
// reset path removes every member the same way. UpdateNotify's partial
// merge is unaffected: it seeds omitted fields from the stored value BEFORE
// calling this editor, so the value it passes is already complete.
// Mirrors setPetPreferenceJSON's duplicate handling.
func setNotifyJSONObject(data []byte, notify NotifyConfig) ([]byte, error) {
encoded, err := json.Marshal(notify)
if err != nil {
return nil, fmt.Errorf("encode notify preference: %w", err)
}
rootStart := skipJSONSpace(data, 0)
root, err := parseJSONObject(data, rootStart)
if err != nil {
return nil, err
}
notifyIndex := lastJSONMember(root.members, "notify")
// Count duplicates: the ordinary file has exactly one notify member, and
// for it a positional replace keeps the key where the user put it and
// preserves their formatting (byte-preserving contract). The
// remove-all-and-reinsert path below is ONLY for files that actually
// contain duplicates, where Go's field-merging decode would otherwise let
// an earlier block's fields leak into the replacement.
duplicates := 0
for _, member := range root.members {
if member.key == "notify" {
duplicates++
}
}
if string(encoded) == "{}" {
for notifyIndex >= 0 {
data = removeJSONMember(data, root, notifyIndex)
rootStart = skipJSONSpace(data, 0)
root, err = parseJSONObject(data, rootStart)
if err != nil {
return nil, err
}
notifyIndex = lastJSONMember(root.members, "notify")
}
return data, nil
}
if duplicates <= 1 {
if notifyIndex < 0 {
return insertJSONMember(data, root, "notify", encoded), nil
}
// Single member: replace its value in place — no key movement, no
// reformatting of the surrounding bytes.
return replaceJSONRange(data, root.members[notifyIndex].valueStart, root.members[notifyIndex].valueEnd, encoded), nil
}
// Multiple duplicates: Go merges their fields at decode time, so the
// replacement must not inherit anything — remove every duplicate and
// insert the complete value fresh.
for notifyIndex >= 0 {
data = removeJSONMember(data, root, notifyIndex)
rootStart = skipJSONSpace(data, 0)
root, err = parseJSONObject(data, rootStart)
if err != nil {
return nil, err
}
notifyIndex = lastJSONMember(root.members, "notify")
}
return insertJSONMember(data, root, "notify", encoded), nil
}
func replaceJSONRange(data []byte, start, end int, replacement []byte) []byte {
result := make([]byte, 0, len(data)-(end-start)+len(replacement))
result = append(result, data[:start]...)
+2 -20
View File
@@ -157,26 +157,8 @@ func Resolve(options ResolveOptions) (ResolvedConfig, error) {
// On ErrNoActiveProvider, providers may still hold the successfully
// normalized (but active-less) profile list — keep it so a caller can fall
// back to an already-configured usable provider instead of treating this
// like a config with nothing set up at all. The non-provider fields parsed
// before the failure (notify, sandbox, tools, preferences, ...) are kept
// too: the TUI's provider-recovery path forwards this partial config, and
// dropping them would let the TUI's unconfigured notify default overwrite
// a stored explicit opt-out on startup (maintainer review, PR #1001).
partial := ResolvedConfig{
Providers: providers,
MaxTurns: cfg.MaxTurns,
MCP: cfg.MCP,
Sandbox: cfg.Sandbox,
Notify: cfg.Notify,
Tools: cfg.Tools,
Swarm: cfg.Swarm,
Preferences: cfg.Preferences,
KeyBindings: cfg.KeyBindings,
LocalControl: cfg.LocalControl,
STT: cfg.STT,
CrossSessionInbound: cfg.CrossSessionInbound,
}
return partial, err
// like a config with nothing set up at all.
return ResolvedConfig{Providers: providers}, err
}
return ResolvedConfig{
+7 -60
View File
@@ -1151,13 +1151,8 @@ func TestResolveRejectsActiveProviderWithoutConfiguredProfiles(t *testing.T) {
if HasProviderProfile(resolved.Provider) {
t.Fatalf("Provider = %#v, want zero value", resolved.Provider)
}
// The error path now carries the parsed non-provider fields through
// (maintainer review, PR #1001): the TUI's provider-recovery startup
// forwards this partial config, and a zero MaxTurns there would silently
// discard the user's turn budget. Matches the no-providers SUCCESS path,
// which has always defaulted MaxTurns.
if resolved.MaxTurns != defaultMaxTurns {
t.Fatalf("MaxTurns = %d, want default %d carried through the partial resolve", resolved.MaxTurns, defaultMaxTurns)
if resolved.MaxTurns != 0 {
t.Fatalf("MaxTurns = %d, want zero on failed resolve", resolved.MaxTurns)
}
}
@@ -1869,43 +1864,14 @@ func TestResolveNotifyInvalidFocusMode(t *testing.T) {
}
}
// An unconfigured notify block must resolve EMPTY. Defaults live in the TUI
// (effectiveTUINotifyMode) and in `zero config notify`'s display, not in the
// resolver: a filled-in default here leaks into headless `zero exec` (BEL +
// OSC-9 bytes on stderr under -o json) and into the CLI/TUI preserve paths,
// which must treat "user never chose" differently from "user chose both"
// (maintainer review, PR #1001).
func TestResolveNotifyUnconfiguredStaysEmpty(t *testing.T) {
// No config file at all.
resolved, err := Resolve(ResolveOptions{Env: map[string]string{}})
func TestResolveNotifyDefaultEmpty(t *testing.T) {
path := writeConfig(t, `{}`)
resolved, err := Resolve(ResolveOptions{UserConfigPath: path, Env: map[string]string{}})
if err != nil {
t.Fatalf("Resolve: %v", err)
}
if resolved.Notify.Mode != "" || resolved.Notify.FocusMode != "" {
t.Fatalf("no config file: notify = %+v, want empty (resolver must not default)", resolved.Notify)
}
for name, body := range map[string]string{
"empty config": `{}`,
"empty block": `{"notify":{}}`,
"mode only": `{"notify":{"mode":"off"}}`,
} {
t.Run(name, func(t *testing.T) {
resolved, err := Resolve(ResolveOptions{UserConfigPath: writeConfig(t, body), Env: map[string]string{}})
if err != nil {
t.Fatalf("Resolve: %v", err)
}
if name == "mode only" {
if resolved.Notify.Mode != "off" {
t.Errorf("notify.mode = %q, want preserved off", resolved.Notify.Mode)
}
} else if resolved.Notify.Mode != "" {
t.Errorf("notify.mode = %q, want empty (resolver must not default)", resolved.Notify.Mode)
}
if resolved.Notify.FocusMode != "" {
t.Errorf("notify.focusMode = %q, want empty (resolver must not default)", resolved.Notify.FocusMode)
}
})
t.Fatalf("unset notify should be empty, got %+v", resolved.Notify)
}
}
@@ -2288,19 +2254,6 @@ func TestResolveRejectsInvalidCrossSessionInbound(t *testing.T) {
}
func TestResolveReportsExplicitMaxTurns(t *testing.T) {
// A nil ResolveOptions.Env means "read the real process environment" (see
// envValue), so a case that leaves it nil inherits whatever the developer
// happens to export. ANTHROPIC_BASE_URL alone is enough: applyProviderEnv
// synthesizes a second, keyless "anthropic" profile from a base URL on its
// own, the single-provider shortcut in normalizeProviders stops applying,
// and Resolve fails with `active provider "" not found` before this test
// reaches its assertion. Every variable the resolver reads does it, so the
// cases below state an EMPTY environment rather than no environment, and
// these three are set to a hostile value so the test proves it.
t.Setenv("ANTHROPIC_BASE_URL", "https://proxy.example/v1")
t.Setenv("OPENAI_API_KEY", "from-the-developer-shell")
t.Setenv(MaxTurnsEnv, "7")
const providerJSON = `{
"activeProvider": "p",
"providers": [{"name": "p", "provider_kind": "openai-compatible", "base_url": "https://x.example/v1", "apiKey": "k", "model": "m"}]
@@ -2313,15 +2266,10 @@ func TestResolveReportsExplicitMaxTurns(t *testing.T) {
set bool
}{
{"built-in default only", ResolveOptions{
Env: map[string]string{},
Overrides: Overrides{Provider: ProviderProfile{Name: "p", ProviderKind: "openai-compatible", BaseURL: "https://x.example/v1", APIKey: "k", Model: "m"}},
}, false},
{"user config", ResolveOptions{
Env: map[string]string{},
UserConfigPath: writeConfig(t, providerWithTurns),
}, true},
{"user config", ResolveOptions{UserConfigPath: writeConfig(t, providerWithTurns)}, true},
{"project config", ResolveOptions{
Env: map[string]string{},
UserConfigPath: writeConfig(t, providerJSON),
ProjectConfigPath: writeConfig(t, `{"maxTurns": 40}`),
}, true},
@@ -2330,7 +2278,6 @@ func TestResolveReportsExplicitMaxTurns(t *testing.T) {
Overrides: Overrides{Provider: ProviderProfile{Name: "p", ProviderKind: "openai-compatible", BaseURL: "https://x.example/v1", APIKey: "k", Model: "m"}},
}, true},
{"cli override", ResolveOptions{
Env: map[string]string{},
Overrides: Overrides{
MaxTurns: 42,
Provider: ProviderProfile{Name: "p", ProviderKind: "openai-compatible", BaseURL: "https://x.example/v1", APIKey: "k", Model: "m"},
-141
View File
@@ -9,7 +9,6 @@ import (
"sort"
"strings"
"github.com/Gitlawb/zero/internal/notify"
"github.com/Gitlawb/zero/internal/providercatalog"
)
@@ -757,146 +756,6 @@ func SetTheme(path string, theme string) (result FileConfig, err error) {
return cfg, nil
}
// UserNotify returns the notify block stored in the user's own config file at
// path, trimmed. A missing file or missing/empty block returns the zero value:
// callers that need to preserve "whatever the user already chose" on a partial
// update must seed from THIS value, not from the resolved view — the resolver
// merges project config and the TUI applies its own defaults, so seeding from
// resolved copies choices the user never made into their global file (a
// project's mode: off, or a pinned default) and across every other project.
func UserNotify(path string) (NotifyConfig, error) {
path = strings.TrimSpace(path)
if path == "" {
return NotifyConfig{}, nil
}
data, err := os.ReadFile(path)
if err != nil {
if os.IsNotExist(err) {
return NotifyConfig{}, nil
}
return NotifyConfig{}, fmt.Errorf("read config %s: %w", path, err)
}
var cfg FileConfig
if err := json.Unmarshal(data, &cfg); err != nil {
return NotifyConfig{}, fmt.Errorf("invalid config JSON %s: %w", path, err)
}
cfg.Notify.Mode = strings.TrimSpace(cfg.Notify.Mode)
cfg.Notify.FocusMode = strings.TrimSpace(cfg.Notify.FocusMode)
return cfg.Notify, nil
}
// validateNotify trims both fields and checks them against the accepted vocab
// (mode in {off,bell,notify,both}; focusMode in {unfocused,always,focused}) so
// a bad caller cannot write a value the resolver would later reject at
// startup. An empty Mode or FocusMode is stored as-is — blank means "use the
// built-in defaults" (the TUI's effectiveTUINotifyMode maps an empty mode to
// both; the notify package treats an empty focusMode as unfocused), not "off".
func validateNotify(value NotifyConfig) (NotifyConfig, error) {
value.Mode = strings.TrimSpace(value.Mode)
value.FocusMode = strings.TrimSpace(value.FocusMode)
if mode := value.Mode; mode != "" {
switch notify.Mode(mode) {
case notify.ModeOff, notify.ModeBell, notify.ModeNotify, notify.ModeBoth:
default:
return NotifyConfig{}, fmt.Errorf("invalid notify.mode %q: expected off, bell, notify, or both", mode)
}
}
if focus := value.FocusMode; focus != "" {
switch notify.FocusMode(focus) {
case notify.FocusUnfocused, notify.FocusAlways, notify.FocusFocused:
default:
return NotifyConfig{}, fmt.Errorf("invalid notify.focusMode %q: expected unfocused, always, or focused", focus)
}
}
return value, nil
}
// UpdateNotify runs one SERIALIZED read-modify-write transaction on the user's
// notify block. merge receives the currently stored block and returns the
// block to store; it must preserve any field the caller is not explicitly
// changing. The lock covers reading the stored block, applying merge, and
// replacing the file, so two concurrent partial updates (for example
// `--mode off` and `--focus always` from two terminals) cannot interleave
// their reads and writes and silently undo each other's explicit change
// (maintainer review, PR #1001 reproduced a 50/50 lost-update without it).
//
// The write replaces ONLY the notify member's bytes (setNotifyJSONObject),
// preserving every unrelated value and its explicit presence — including
// explicit zeros/falses like tools.deferThreshold: 0 or an MCP server's
// disabled: false, which the typed serializer's omitempty cannot round-trip —
// through the existing temp-file-and-rename atomic publish. Locking only the
// final write would leave the stale-field merge outside the transaction, so
// the whole sequence runs under one lock — the SAME lockConfigFileFn authority
// every other config mutator uses (provider/theme/credential/MCP writes), not
// a private lock: two authorities over the same file would not contend
// (maintainer review, PR #1001). The release error is joined into the
// result like every other writer: a failed unlock would leave the advisory
// lock held and later updates blocked while the caller is told success.
func UpdateNotify(path string, merge func(current NotifyConfig) NotifyConfig) (result NotifyConfig, err error) {
path = strings.TrimSpace(path)
if path == "" {
return NotifyConfig{}, fmt.Errorf("config path is required")
}
if merge == nil {
return NotifyConfig{}, fmt.Errorf("merge function is required")
}
release, err := lockConfigFileFn(path)
if err != nil {
return NotifyConfig{}, err
}
// Joined, not chosen between: a release failure annotates the result
// instead of masking the mutation error that actually explains what went
// wrong. Reporting success after a failed release would claim a state the
// next update cannot reproduce.
defer func() { err = errors.Join(err, release()) }()
current, err := UserNotify(path)
if err != nil {
return NotifyConfig{}, err
}
next, err := validateNotify(merge(current))
if err != nil {
return NotifyConfig{}, err
}
// Read the raw file and edit only the notify member. A missing file starts
// from an empty object so a first-time write does not need a pre-seeded
// config.
data, err := os.ReadFile(path)
if err != nil {
if !os.IsNotExist(err) {
return NotifyConfig{}, fmt.Errorf("read config %s: %w", path, err)
}
data = []byte("{}")
}
updated, err := setNotifyJSONObject(data, next)
if err != nil {
return NotifyConfig{}, fmt.Errorf("invalid config JSON %s: %w", path, err)
}
if err := writeConfigData(path, updated); err != nil {
return NotifyConfig{}, err
}
return next, nil
}
// SetNotify replaces the stored notification preference with value. It is a
// thin wrapper over the transactional UpdateNotify so existing callers keep
// their replace semantics under the same lock and byte-preserving write.
func SetNotify(path string, value NotifyConfig) (FileConfig, error) {
next, err := UpdateNotify(path, func(NotifyConfig) NotifyConfig { return value })
if err != nil {
return FileConfig{}, err
}
// Re-read the full config so callers that inspect unrelated fields (and
// existing tests) keep their FileConfig-shaped result.
cfg, err := loadConfigFile(path)
if err != nil {
return FileConfig{}, err
}
cfg.Notify = next
return cfg, nil
}
// SetPet persists only the terminal-pet preference while preserving every
// unrelated user setting through the config writer's atomic replace path.
func SetPet(path string, pet string) (result FileConfig, err error) {
-371
View File
@@ -3,7 +3,6 @@ package config
import (
"encoding/json"
"errors"
"fmt"
"io/fs"
"os"
"os/exec"
@@ -11,7 +10,6 @@ import (
"reflect"
"runtime"
"strings"
"sync"
"testing"
)
@@ -428,269 +426,6 @@ func TestSetThemePreservesUnknownTopLevelFields(t *testing.T) {
}
}
func TestSetNotifyPersistsValidValues(t *testing.T) {
path := filepath.Join(t.TempDir(), "zero.json")
writeConfigFixture(t, path, FileConfig{
ActiveProvider: "openai",
Providers: []ProviderProfile{
{Name: "openai", ProviderKind: ProviderKindOpenAI, Model: "gpt-4.1"},
},
}, 0o600)
cfg, err := SetNotify(path, NotifyConfig{Mode: " both ", FocusMode: " unfocused "})
if err != nil {
t.Fatalf("SetNotify() error = %v", err)
}
if cfg.Notify.Mode != "both" || cfg.Notify.FocusMode != "unfocused" {
t.Fatalf("Notify = %+v, want mode=both focusMode=unfocused (trimmed)", cfg.Notify)
}
persisted := readConfigFixture(t, path)
if persisted.Notify.Mode != "both" || persisted.Notify.FocusMode != "unfocused" {
t.Fatalf("persisted Notify = %+v, want mode=both focusMode=unfocused", persisted.Notify)
}
if persisted.ActiveProvider != "openai" || len(persisted.Providers) != 1 {
t.Fatalf("provider config was not preserved by SetNotify: %#v", persisted)
}
}
func TestSetNotifyRejectsInvalidMode(t *testing.T) {
path := filepath.Join(t.TempDir(), "zero.json")
writeConfigFixture(t, path, FileConfig{ActiveProvider: "openai"}, 0o600)
if _, err := SetNotify(path, NotifyConfig{Mode: "loud", FocusMode: "unfocused"}); err == nil {
t.Fatal("expected error for invalid notify.mode")
}
}
func TestSetNotifyRejectsInvalidFocusMode(t *testing.T) {
path := filepath.Join(t.TempDir(), "zero.json")
writeConfigFixture(t, path, FileConfig{ActiveProvider: "openai"}, 0o600)
if _, err := SetNotify(path, NotifyConfig{Mode: "off", FocusMode: "sideways"}); err == nil {
t.Fatal("expected error for invalid notify.focusMode")
}
}
func TestSetNotifyRejectsEmptyConfigPath(t *testing.T) {
if _, err := SetNotify("", NotifyConfig{Mode: "off"}); err == nil {
t.Fatal("expected error for empty config path")
}
}
func TestSetNotifyBlankValuesPreservedAsDefaults(t *testing.T) {
// An empty mode/focusMode stored on disk is a valid "use the built-in
// defaults" signal — SetNotify must not reject blanks, and they must round
// trip unchanged.
path := filepath.Join(t.TempDir(), "zero.json")
writeConfigFixture(t, path, FileConfig{ActiveProvider: "openai"}, 0o600)
if _, err := SetNotify(path, NotifyConfig{}); err != nil {
t.Fatalf("SetNotify({}) should accept blank values, got error: %v", err)
}
persisted := readConfigFixture(t, path)
if persisted.Notify.Mode != "" || persisted.Notify.FocusMode != "" {
t.Fatalf("blank notify values should round-trip, got %+v", persisted.Notify)
}
}
// Maintainer regression (PR #1001): a notification save must preserve unrelated
// values whose EXPLICIT presence matters. The typed serializer's omitempty
// cannot round-trip tools.deferThreshold: 0 (explicit "never defer" vs unset
// "use default 3") or an MCP server's disabled: false (explicitly enabled), so
// SetNotify/UpdateNotify edit only the notify member's bytes instead. The
// assertions run against the RAW file, because reading through FileConfig
// would hide exactly the presence loss being tested.
func TestSetNotifyPreservesExplicitUnrelatedValues(t *testing.T) {
dir := t.TempDir()
for name, body := range map[string]string{
"explicit zero deferThreshold": `{
"tools": {"deferThreshold": 0},
"mcp": {"servers": {"firecrawl": {"command": "npx", "disabled": false}}}
}`,
"unknown keys stay untouched": `{
"activeProvider": "openai",
"customTopLevel": {"nested": [1, 2, 3]},
"notify": {"mode": "both"}
}`,
} {
t.Run(name, func(t *testing.T) {
path := filepath.Join(dir, name+".json")
if err := os.WriteFile(path, []byte(body), 0o600); err != nil {
t.Fatal(err)
}
if _, err := SetNotify(path, NotifyConfig{Mode: "off", FocusMode: "always"}); err != nil {
t.Fatalf("SetNotify: %v", err)
}
raw, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
var obj map[string]json.RawMessage
if err := json.Unmarshal(raw, &obj); err != nil {
t.Fatalf("decode: %v", err)
}
if strings.Contains(name, "explicit zero") {
tools, ok := obj["tools"]
if !ok || !strings.Contains(string(tools), `"deferThreshold": 0`) {
t.Errorf("tools.deferThreshold: 0 lost through a notify write; tools = %s", string(tools))
}
mcp, ok := obj["mcp"]
if !ok || !strings.Contains(string(mcp), `"disabled": false`) {
t.Errorf("mcp disabled: false lost through a notify write; mcp = %s", string(mcp))
}
}
if strings.Contains(name, "unknown keys") {
if _, ok := obj["customTopLevel"]; !ok {
t.Error("unknown top-level key lost through a notify write")
}
}
})
}
}
// Maintainer regression (PR #1001, CodeRabbit follow-up): a hand-edited config
// may contain DUPLICATE notify members (JSON decoders tolerate them and the
// last occurrence wins). A reset must remove EVERY notify member — removing
// only the last leaves the earlier block as the new effective preference, so
// `zero config notify --reset` would report success while the old value still
// applies.
func TestSetNotifyResetRemovesEveryDuplicateNotifyMember(t *testing.T) {
path := filepath.Join(t.TempDir(), "zero.json")
duplicated := `{"notify":{"mode":"off"},"activeProvider":"openai","notify":{"mode":"bell"}}`
if err := os.WriteFile(path, []byte(duplicated), 0o600); err != nil {
t.Fatal(err)
}
if _, err := SetNotify(path, NotifyConfig{}); err != nil {
t.Fatalf("SetNotify({}) reset: %v", err)
}
raw, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
var obj map[string]json.RawMessage
if err := json.Unmarshal(raw, &obj); err != nil {
t.Fatalf("decode: %v", err)
}
if _, still := obj["notify"]; still {
t.Fatalf("reset left a notify member behind: %s", string(raw))
}
if obj["activeProvider"] == nil || string(obj["activeProvider"]) != `"openai"` {
t.Errorf("unrelated member lost through the reset: %s", string(raw))
}
// A reset on a file with NO notify member is a no-op, and a partial update
// against duplicates replaces the LAST member (the effective one under
// last-occurrence-wins decoding).
if err := os.WriteFile(path, []byte(duplicated), 0o600); err != nil {
t.Fatal(err)
}
if _, err := SetNotify(path, NotifyConfig{Mode: "off", FocusMode: "always"}); err != nil {
t.Fatalf("SetNotify replace: %v", err)
}
raw, err = os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
stored, err := UserNotify(path)
if err != nil {
t.Fatalf("UserNotify: %v", err)
}
if stored.Mode != "off" || stored.FocusMode != "always" {
t.Fatalf("after replace, effective notify = %+v, want off/always (last member replaced)", stored)
}
if !strings.Contains(string(raw), `"activeProvider"`) {
t.Errorf("unrelated member lost through the replace: %s", string(raw))
}
}
// Maintainer regression (PR #1001): concurrent partial updates must not lose
// each other's explicit change. Two `zero config notify` calls (--mode off and
// --focus always) racing from both/unfocused previously interleaved their
// read-merge-write and one write silently undid the other (reproduced 50/50
// without the lock). UpdateNotify serializes the whole transaction.
func TestUpdateNotifyConcurrentPartialUpdatesLoseNoField(t *testing.T) {
path := filepath.Join(t.TempDir(), "zero.json")
if err := os.WriteFile(path, []byte(`{"notify":{"mode":"both","focusMode":"unfocused"}}`), 0o600); err != nil {
t.Fatal(err)
}
var wg sync.WaitGroup
errs := make(chan error, 2)
wg.Add(2)
go func() {
defer wg.Done()
if _, err := UpdateNotify(path, func(current NotifyConfig) NotifyConfig {
current.Mode = "off"
return current
}); err != nil {
errs <- fmt.Errorf("mode update: %w", err)
}
}()
go func() {
defer wg.Done()
if _, err := UpdateNotify(path, func(current NotifyConfig) NotifyConfig {
current.FocusMode = "always"
return current
}); err != nil {
errs <- fmt.Errorf("focus update: %w", err)
}
}()
wg.Wait()
close(errs)
for err := range errs {
t.Fatalf("concurrent update failed: %v", err)
}
stored, err := UserNotify(path)
if err != nil {
t.Fatalf("read stored: %v", err)
}
// BOTH explicit changes must survive: mode from writer 1, focus from
// writer 2. The old race let the second writer's stale focus
// ("unfocused") or stale mode ("both") win.
if stored.Mode != "off" {
t.Errorf("mode = %q, want off (concurrent focus update must not clobber it)", stored.Mode)
}
if stored.FocusMode != "always" {
t.Errorf("focusMode = %q, want always (concurrent mode update must not clobber it)", stored.FocusMode)
}
}
// UserNotify reads the notify block from the user's own file. Partial updates
// seed from this value so they preserve what the USER chose (blank included)
// instead of copying a project config's setting or a pinned default into the
// global file (maintainer review, PR #1001).
func TestUserNotify(t *testing.T) {
dir := t.TempDir()
// Missing file: zero value, no error.
got, err := UserNotify(filepath.Join(dir, "missing.json"))
if err != nil {
t.Fatalf("missing file should not error: %v", err)
}
if got.Mode != "" || got.FocusMode != "" {
t.Fatalf("missing file = %+v, want zero value", got)
}
// Present block: trimmed values returned.
path := filepath.Join(dir, "zero.json")
writeConfigFixture(t, path, FileConfig{Notify: NotifyConfig{Mode: " bell ", FocusMode: " always "}}, 0o600)
got, err = UserNotify(path)
if err != nil {
t.Fatalf("UserNotify: %v", err)
}
if got.Mode != "bell" || got.FocusMode != "always" {
t.Fatalf("UserNotify = %+v, want bell/always (trimmed)", got)
}
// Blank path: zero value, no error.
if got, err = UserNotify(""); err != nil || got.Mode != "" || got.FocusMode != "" {
t.Fatalf("blank path = %+v err=%v, want zero value", got, err)
}
}
func TestRecapsPreferenceRoundTrips(t *testing.T) {
// Default (unset) is ON.
if !(PreferencesConfig{}).RecapsEnabled() {
@@ -1343,109 +1078,3 @@ func TestEditProviderRejectsCollisionAndUnknown(t *testing.T) {
t.Fatalf("config was rewritten by a rejected edit")
}
}
// Maintainer regression (PR #1001): notify writes must contend with EXISTING
// config mutators, not only each other. UpdateNotify originally took a private
// lock (acquireConfigLock); on Windows its LockFileEx range did not overlap
// lockutil's, so `zero config notify --mode off` could race a theme (or
// provider/credential/MCP) write and silently discard it — the later rename
// wins and the other acknowledged change is gone. UpdateNotify now uses the
// same lockConfigFileFn authority as every other mutator; this race holds it
// against SetTheme and requires BOTH acknowledged edits to survive.
func TestUpdateNotifyRacesExistingConfigMutatorWithoutLoss(t *testing.T) {
path := filepath.Join(t.TempDir(), "zero.json")
if err := os.WriteFile(path, []byte(`{"notify":{"mode":"both","focusMode":"unfocused"}}`), 0o600); err != nil {
t.Fatal(err)
}
var wg sync.WaitGroup
errs := make(chan error, 2)
wg.Add(2)
go func() {
defer wg.Done()
if _, err := UpdateNotify(path, func(NotifyConfig) NotifyConfig {
return NotifyConfig{Mode: "off", FocusMode: "always"}
}); err != nil {
errs <- fmt.Errorf("notify write: %w", err)
}
}()
go func() {
defer wg.Done()
if _, err := SetTheme(path, "dracula"); err != nil {
errs <- fmt.Errorf("theme write: %w", err)
}
}()
wg.Wait()
close(errs)
for err := range errs {
t.Fatal(err)
}
cfg := readConfigFixture(t, path)
if cfg.Notify.Mode != "off" || cfg.Notify.FocusMode != "always" {
t.Errorf("notify lost the race: %+v, want off/always — the two writers did not contend on one lock", cfg.Notify)
}
if cfg.Preferences.Theme != "dracula" {
t.Errorf("theme lost the race: %q, want dracula — the two writers did not contend on one lock", cfg.Preferences.Theme)
}
}
// Maintainer regression (PR #1001, jatmn review): Go's decoder MERGES the
// fields of duplicate JSON members into one struct, so a full replacement that
// omits a field could inherit it from an earlier duplicate —
// {"notify":{"mode":"off"},"notify":{"focusMode":"focused"}} +
// SetNotify({FocusMode:"always"}) must yield blank/always, not off/always.
// The editor removes every duplicate and inserts the replacement fresh.
func TestSetNotifyReplacementInheritsNoFieldFromEarlierDuplicate(t *testing.T) {
path := filepath.Join(t.TempDir(), "zero.json")
duplicated := `{"notify":{"mode":"off"},"notify":{"focusMode":"focused"},"activeProvider":"openai"}`
if err := os.WriteFile(path, []byte(duplicated), 0o600); err != nil {
t.Fatal(err)
}
if _, err := SetNotify(path, NotifyConfig{FocusMode: "always"}); err != nil {
t.Fatalf("SetNotify: %v", err)
}
stored, err := UserNotify(path)
if err != nil {
t.Fatal(err)
}
if stored.Mode != "" || stored.FocusMode != "always" {
t.Fatalf("replacement inherited a field from an earlier duplicate: %+v, want blank/always", stored)
}
// Unrelated bytes still survive the rewrite.
raw, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(string(raw), `"activeProvider":"openai"`) {
t.Fatalf("unrelated member lost through the replace: %s", string(raw))
}
}
// Maintainer regression (PR #1001, Vasanthdev2004 review): the duplicate-key
// fix must not reach the ORDINARY single-member file. Removing and re-inserting
// the notify member there moved the key to the end of the object and collapsed
// its formatting, so every /notify toggle reordered the user's config — JSON
// member order carries no meaning, but a byte-preserving write should not
// reformat anything it does not have to. A single notify member is replaced
// in place: key position and surrounding formatting unchanged.
func TestSetNotifySingleMemberReplacedInPlace(t *testing.T) {
path := filepath.Join(t.TempDir(), "zero.json")
original := "{\n \"activeProvider\": \"openai\",\n \"notify\": {\"mode\": \"both\", \"focusMode\": \"unfocused\"},\n \"preferences\": {\"theme\": \"dracula\"}\n}\n"
if err := os.WriteFile(path, []byte(original), 0o600); err != nil {
t.Fatal(err)
}
if _, err := SetNotify(path, NotifyConfig{Mode: "off", FocusMode: "always"}); err != nil {
t.Fatalf("SetNotify: %v", err)
}
raw, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
want := "{\n \"activeProvider\": \"openai\",\n \"notify\": {\"mode\":\"off\",\"focusMode\":\"always\"},\n \"preferences\": {\"theme\": \"dracula\"}\n}\n"
if string(raw) != want {
t.Fatalf("single-member replace moved or reformatted unrelated bytes:\n got: %q\nwant: %q", string(raw), want)
}
}
+1 -305
View File
@@ -1,61 +1,10 @@
package cron
import (
"bytes"
"encoding/json"
"errors"
"os"
"path/filepath"
"runtime"
"strings"
"sync"
"syscall"
"testing"
"github.com/Gitlawb/zero/internal/fsutil"
)
func TestAppendRunRetainsNewestThousand(t *testing.T) {
store := newTestStore(t)
job, err := store.Add(Job{Expr: "* * * * *", Prompt: "x"})
if err != nil {
t.Fatal(err)
}
// Seed an oversized legacy log, then cross the boundary again after compaction.
var history bytes.Buffer
for i := 0; i < 1207; i++ {
if err := json.NewEncoder(&history).Encode(RunRecord{ExitCode: i}); err != nil {
t.Fatal(err)
}
}
path := filepath.Join(store.jobDir(job.ID), "runs.jsonl")
if err := os.WriteFile(path, history.Bytes(), 0o600); err != nil {
t.Fatal(err)
}
for _, next := range []int{1207, 1208} {
if err := store.AppendRun(job.ID, RunRecord{ExitCode: next}); err != nil {
t.Fatal(err)
}
data, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
lines := bytes.Split(bytes.TrimSpace(data), []byte("\n"))
if len(lines) != 1000 {
t.Fatalf("retained %d records, want 1000", len(lines))
}
for i, line := range lines {
var rec RunRecord
if err := json.Unmarshal(line, &rec); err != nil {
t.Fatal(err)
}
if rec.ExitCode != next-999+i {
t.Fatalf("record %d = %d, want %d", i, rec.ExitCode, next-999+i)
}
}
}
}
// AppendRun on a job removed mid-run must NOT resurrect its directory (which the
// old unconditional MkdirAll did, leaving an orphaned runs.jsonl with no
// metadata.json).
@@ -75,7 +24,7 @@ func TestAppendRunDoesNotResurrectRemovedJob(t *testing.T) {
if _, err := os.Stat(store.jobDir(job.ID)); !os.IsNotExist(err) {
t.Fatalf("AppendRun resurrected the removed job directory (stat err = %v)", err)
}
runs, err := store.Runs(job.ID, 0)
runs, err := store.Runs(job.ID)
if err != nil {
t.Fatalf("Runs: %v", err)
}
@@ -83,256 +32,3 @@ func TestAppendRunDoesNotResurrectRemovedJob(t *testing.T) {
t.Fatalf("expected no runs for a removed job, got %d", len(runs))
}
}
func TestRunsTail(t *testing.T) {
store := newTestStore(t)
job, err := store.Add(Job{Prompt: "x"})
if err != nil {
t.Fatal(err)
}
var history bytes.Buffer
// A forward scan would fail on this old oversized line. A recent-window
// reader must stop before reaching it, rather than merely cap its output.
history.WriteString(strings.Repeat("x", 1024*1024+1) + "\n")
for i := 0; i < 1103; i++ {
if err := json.NewEncoder(&history).Encode(RunRecord{ExitCode: i, SessionTitle: strings.Repeat("界", 1500)}); err != nil {
t.Fatal(err)
}
}
history.WriteString("bad json\n\n{\"exitCode\":1103}") // no final newline
path := filepath.Join(store.jobDir(job.ID), "runs.jsonl")
if err := os.WriteFile(path, history.Bytes(), 0o600); err != nil {
t.Fatal(err)
}
for _, limit := range []int{1, 7, 999, 1000, 1001, 0, -1} {
runs, err := store.Runs(job.ID, limit)
want := limit
if want <= 0 || want > 1000 {
want = 1000
}
if err != nil || len(runs) != want {
t.Fatalf("limit %d: got %d records, err %v", limit, len(runs), err)
}
for i, run := range runs {
if run.ExitCode != 1104-want+i {
t.Fatalf("limit %d, record %d = %d", limit, i, run.ExitCode)
}
if run.ExitCode != 1103 && run.SessionTitle != strings.Repeat("界", 1500) {
t.Fatal("record spanning read blocks was corrupted")
}
}
}
}
func TestAppendRunFailurePreservesHistory(t *testing.T) {
store := newTestStore(t)
job, err := store.Add(Job{Prompt: "x"})
if err != nil {
t.Fatal(err)
}
path := filepath.Join(store.jobDir(job.ID), "runs.jsonl")
for _, data := range []string{"{\"exitCode\":7}\n", strings.Repeat("x", 1024*1024)} {
if err := os.WriteFile(path, []byte(data), 0o600); err != nil {
t.Fatal(err)
}
rec := RunRecord{ExitCode: 8, Error: strings.Repeat("x", 1024*1024)}
if err := store.AppendRun(job.ID, rec); err == nil {
t.Fatal("expected oversized new record error")
}
got, err := os.ReadFile(path)
if err != nil || string(got) != data {
t.Fatalf("failed append changed history: %v", err)
}
}
}
func TestAppendRunWithOversizedHistory(t *testing.T) {
store := newTestStore(t)
job, err := store.Add(Job{Prompt: "x"})
if err != nil {
t.Fatal(err)
}
path := filepath.Join(store.jobDir(job.ID), "runs.jsonl")
// An oversized unterminated legacy record must not block future outcomes.
if err := os.WriteFile(path, []byte(strings.Repeat("x", maxRunBytes)), 0o600); err != nil {
t.Fatal(err)
}
if err := store.AppendRun(job.ID, RunRecord{ExitCode: 42}); err != nil {
t.Fatal(err)
}
runs, err := store.Runs(job.ID, 1)
if err != nil || len(runs) != 1 || runs[0].ExitCode != 42 {
t.Fatalf("new outcome lost: %+v, %v", runs, err)
}
}
func TestRunsSkipOversizedLines(t *testing.T) {
store := newTestStore(t)
job, err := store.Add(Job{Prompt: "x"})
if err != nil {
t.Fatal(err)
}
path := filepath.Join(store.jobDir(job.ID), "runs.jsonl")
for _, size := range []int{maxRunBytes - 1, maxRunBytes, maxRunBytes + 4096} {
// Valid JSON plus trailing whitespace catches accidental decoding of a
// prefix after dropping only the oversized suffix. The exact boundary
// also distinguishes an oversized line from the largest allowed line.
line := `{"exitCode":99}`
line += strings.Repeat(" ", size-len(line))
data := line + "\n{\"exitCode\":7}\n" + line + "\n{\"exitCode\":42}\n" + line
if err := os.WriteFile(path, []byte(data), 0o600); err != nil {
t.Fatal(err)
}
runs, err := store.Runs(job.ID, 10)
want := []int{7, 42}
if size < maxRunBytes {
want = []int{99, 7, 99, 42, 99}
}
if err != nil || len(runs) != len(want) {
t.Fatalf("size %d: got %+v, err %v; want %v", size, runs, err, want)
}
for i, rec := range runs {
if rec.ExitCode != want[i] {
t.Fatalf("size %d: record %d = %d, want %d", size, i, rec.ExitCode, want[i])
}
}
}
}
// An outside reader holding runs.jsonl open makes the Windows replace fail;
// the outcome must still be recorded rather than dropped.
func TestAppendRunWhileHistoryHeldOpen(t *testing.T) {
store := newTestStore(t)
job, err := store.Add(Job{Prompt: "x"})
if err != nil {
t.Fatal(err)
}
path := filepath.Join(store.jobDir(job.ID), "runs.jsonl")
if err := os.WriteFile(path, []byte("{\"exitCode\":7}\n"), 0o600); err != nil {
t.Fatal(err)
}
held, err := os.Open(path)
if err != nil {
t.Fatal(err)
}
defer held.Close()
if err := store.AppendRun(job.ID, RunRecord{ExitCode: 42}); err != nil {
t.Fatalf("AppendRun with history held open: %v", err)
}
runs, err := store.Runs(job.ID, 10)
if err != nil || len(runs) != 2 || runs[0].ExitCode != 7 || runs[1].ExitCode != 42 {
t.Fatalf("outcome lost while history held open: %+v, %v", runs, err)
}
}
// Only a sharing or lock violation falls back to appending. Any other replace
// failure, including a persistent access denial, and a committed replacement
// whose backup cleanup failed must leave the log untouched, so history can
// neither grow past retention nor record an outcome twice.
func TestAppendRunFallbackOnlyForSharingViolation(t *testing.T) {
const seed = "{\"exitCode\":7}\n"
for _, tc := range []struct {
name string
err error
append bool
}{
{"sharing violation", syscall.Errno(32), runtime.GOOS == "windows"},
{"access denied", os.ErrPermission, false},
{"other", errors.New("boom"), false},
{"committed", &fsutil.CommittedReplacementCleanupError{Cause: syscall.Errno(32)}, false},
} {
t.Run(tc.name, func(t *testing.T) {
store := newTestStore(t)
job, err := store.Add(Job{Prompt: "x"})
if err != nil {
t.Fatal(err)
}
path := filepath.Join(store.jobDir(job.ID), "runs.jsonl")
if err := os.WriteFile(path, []byte(seed), 0o600); err != nil {
t.Fatal(err)
}
store.replace = func(string, string) error { return tc.err }
err = store.AppendRun(job.ID, RunRecord{ExitCode: 42})
data, readErr := os.ReadFile(path)
if readErr != nil {
t.Fatal(readErr)
}
if tc.append {
if err != nil || !strings.HasPrefix(string(data), seed) || !strings.Contains(string(data), "\"exitCode\":42") {
t.Fatalf("sharing violation dropped the outcome: err %v, log %q", err, data)
}
return
}
if !errors.Is(err, tc.err) {
t.Fatalf("AppendRun error = %v, want %v", err, tc.err)
}
if string(data) != seed {
t.Fatalf("non-transient replace failure changed the log: %q", data)
}
})
}
}
// The fallback append must start a fresh line even when the log ends in an
// unterminated record, so neither record is corrupted.
func TestAppendRunLineAfterUnterminatedTail(t *testing.T) {
store := newTestStore(t)
job, err := store.Add(Job{Prompt: "x"})
if err != nil {
t.Fatal(err)
}
path := filepath.Join(store.jobDir(job.ID), "runs.jsonl")
if err := os.WriteFile(path, []byte("{\"exitCode\":7}"), 0o600); err != nil {
t.Fatal(err)
}
if err := appendRunLine(path, []byte("{\"exitCode\":42}")); err != nil {
t.Fatal(err)
}
runs, err := store.Runs(job.ID, 10)
if err != nil || len(runs) != 2 || runs[0].ExitCode != 7 || runs[1].ExitCode != 42 {
t.Fatalf("fallback append corrupted history: %+v, %v", runs, err)
}
}
func TestRunHistoryConcurrentStores(t *testing.T) {
store := newTestStore(t)
job, err := store.Add(Job{Prompt: "x"})
if err != nil {
t.Fatal(err)
}
var history bytes.Buffer
for i := 0; i < 1000; i++ {
if err := json.NewEncoder(&history).Encode(RunRecord{ExitCode: -1}); err != nil {
t.Fatal(err)
}
}
if err := os.WriteFile(filepath.Join(store.jobDir(job.ID), "runs.jsonl"), history.Bytes(), 0o600); err != nil {
t.Fatal(err)
}
var wg sync.WaitGroup
for i := 0; i < 12; i++ {
wg.Go(func() {
other := NewStore(StoreOptions{RootDir: store.root})
if err := other.AppendRun(job.ID, RunRecord{ExitCode: i}); err != nil {
t.Error(err)
return
}
runs, err := other.Runs(job.ID, 1000)
if err != nil || len(runs) != 1000 {
t.Errorf("inconsistent history: %d records, err %v", len(runs), err)
}
})
}
wg.Wait()
runs, err := store.Runs(job.ID, 12)
if err != nil || len(runs) != 12 {
t.Fatalf("final tail: %d records, err %v", len(runs), err)
}
seen := make(map[int]bool)
for _, run := range runs {
if run.ExitCode < 0 || run.ExitCode >= 12 || seen[run.ExitCode] {
t.Fatalf("lost or duplicate concurrent append: %+v", runs)
}
seen[run.ExitCode] = true
}
}
+20 -127
View File
@@ -7,7 +7,6 @@ import (
"fmt"
"os"
"path/filepath"
"slices"
"strings"
"time"
@@ -17,10 +16,6 @@ import (
const (
StatusActive = "active"
StatusPaused = "paused"
// MaxRunHistory is the number of recent outcomes retained per job.
MaxRunHistory = 1000
maxRunBytes = 1024 * 1024
)
// ErrJobNotFound is returned (wrapped) by Get when a job's metadata file is
@@ -61,9 +56,6 @@ type StoreOptions struct {
type Store struct {
root string
now func() time.Time
// replace overrides the run-history replacement primitive in tests; nil
// uses the platform default.
replace func(src, dst string) error
}
func NewStore(opts StoreOptions) *Store {
@@ -236,7 +228,8 @@ func (s *Store) Update(job Job) error {
// Remove). It re-reads the CURRENT on-disk job and passes it to mutate along with
// any read error; mutate returns the job to persist. A removed job aborts with
// ErrJobNotFound (no recreate). A transient read error (not removal) is surfaced
// via readErr; the callback can return it to abort without overwriting state.
// via readErr so the caller can still persist a best-effort state — the fire path
// advances the schedule regardless, to avoid a re-fire.
func (s *Store) Mutate(id string, mutate func(current Job, readErr error) (Job, error)) (Job, error) {
if !validID(id) {
return Job{}, fmt.Errorf("invalid cron job id %q", id)
@@ -316,8 +309,10 @@ func (s *Store) AppendRun(id string, rec RunRecord) error {
return err
}
defer unlock()
// Do not resurrect a job removed mid-run. The lock serializes this check
// and publication against Remove and other history writers/readers.
// Bail if the job was removed (e.g. mid-run) — otherwise the MkdirAll below
// would resurrect a deleted job's directory with an orphaned runs.jsonl and no
// metadata.json (which Runs would then still return). Under the per-job lock
// this stat-then-write is race-free against a concurrent Remove.
if _, err := os.Stat(filepath.Join(s.jobDir(id), "metadata.json")); err != nil {
if errors.Is(err, os.ErrNotExist) {
return nil
@@ -325,98 +320,30 @@ func (s *Store) AppendRun(id string, rec RunRecord) error {
return err
}
dir := s.jobDir(id)
if err := os.MkdirAll(dir, 0o700); err != nil {
return err
}
f, err := os.OpenFile(filepath.Join(dir, "runs.jsonl"), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600)
if err != nil {
return err
}
line, err := json.Marshal(rec)
if err != nil {
return err
}
if len(line) >= maxRunBytes {
return fmt.Errorf("cron run record exceeds %d bytes", maxRunBytes-1)
}
runs, err := s.readRuns(id, MaxRunHistory-1)
if err != nil {
return err
}
// Publish a complete snapshot even below the retention boundary so readers
// never see a partial append. Legacy oversized logs compact on their next run.
f, err := os.CreateTemp(dir, "runs-*.tmp")
if err != nil {
return err
}
defer os.Remove(f.Name())
w := bufio.NewWriter(f)
for _, run := range runs {
if err := json.NewEncoder(w).Encode(run); err != nil {
_ = f.Close()
return err
}
}
if _, err := w.Write(append(line, '\n')); err != nil {
_ = f.Close()
return err
}
if err := w.Flush(); err != nil {
if _, err := f.Write(append(line, '\n')); err != nil {
_ = f.Close()
return err
}
// Surface a buffered-write failure that only materializes on Close.
if err := f.Close(); err != nil {
return err
}
path := filepath.Join(dir, "runs.jsonl")
// ReplaceWithRetry keeps a DACL applied to runs.jsonl itself; a rename would
// publish the temporary file's inherited directory DACL instead.
err = fsutil.ReplaceWithRetry(f.Name(), path, s.replace)
var committed *fsutil.CommittedReplacementCleanupError
if err == nil || errors.As(err, &committed) || !fsutil.IsSharingOrLockViolation(err) {
return err
}
// Windows refuses to replace a log another process holds open. Append the
// record instead so the outcome survives; compaction resumes on a later run.
// Other failures are returned so a log that can never be replaced cannot
// grow past the retention limit. The leading newline ends any unterminated
// tail, and readRuns skips the blank line it may leave. Zero's own readers
// hold the job lock.
if appendErr := appendRunLine(path, line); appendErr != nil {
return errors.Join(err, appendErr)
}
return nil
}
func appendRunLine(path string, line []byte) error {
f, err := os.OpenFile(path, os.O_WRONLY|os.O_APPEND, 0)
if err != nil {
return err
}
record := append(append([]byte{'\n'}, line...), '\n')
if _, err := f.Write(record); err != nil {
_ = f.Close()
return err
}
return f.Close()
}
// Runs returns the newest limit valid records in append order (oldest first).
// Non-positive limits and limits above MaxRunHistory use MaxRunHistory. It reads
// backwards from the tail, including for legacy logs not yet compacted.
func (s *Store) Runs(id string, limit int) ([]RunRecord, error) {
func (s *Store) Runs(id string) ([]RunRecord, error) {
if !validID(id) {
return nil, fmt.Errorf("invalid cron job id %q", id)
}
if limit <= 0 || limit > MaxRunHistory {
limit = MaxRunHistory
}
unlock, err := s.lockJob(id)
if err != nil {
return nil, err
}
defer unlock()
return s.readRuns(id, limit)
}
// readRuns requires the job lock. Memory is bounded by limit records plus one
// line. Malformed and oversized legacy lines are skipped so they cannot prevent
// recording new outcomes during compaction.
func (s *Store) readRuns(id string, limit int) ([]RunRecord, error) {
f, err := os.Open(filepath.Join(s.jobDir(id), "runs.jsonl"))
if errors.Is(err, os.ErrNotExist) {
return nil, nil
@@ -425,48 +352,14 @@ func (s *Store) readRuns(id string, limit int) ([]RunRecord, error) {
return nil, err
}
defer f.Close()
info, err := f.Stat()
if err != nil {
return nil, err
}
var runs []RunRecord
var line []byte
oversized := false
decode := func() {
slices.Reverse(line)
scanner := bufio.NewScanner(f)
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
for scanner.Scan() {
var rec RunRecord
if json.Unmarshal(line, &rec) == nil {
if json.Unmarshal(scanner.Bytes(), &rec) == nil {
runs = append(runs, rec)
}
line = line[:0]
}
var block [4096]byte
for end := info.Size(); end > 0 && len(runs) < limit; {
start := max(int64(0), end-int64(len(block)))
n, err := f.ReadAt(block[:end-start], start)
if err != nil {
return nil, err
}
for i := n - 1; i >= 0 && len(runs) < limit; i-- {
if block[i] == '\n' {
if !oversized {
decode()
}
oversized = false
} else if !oversized {
if len(line) >= maxRunBytes-1 {
line = line[:0]
oversized = true
continue
}
line = append(line, block[i])
}
}
end = start
}
if len(runs) < limit && len(line) > 0 {
decode()
}
slices.Reverse(runs)
return runs, nil
return runs, scanner.Err()
}
+2 -2
View File
@@ -238,7 +238,7 @@ func TestStoreAppendRun(t *testing.T) {
t.Fatalf("AppendRun: %v", err)
}
}
runs, err := s.Runs(job.ID, 0)
runs, err := s.Runs(job.ID)
if err != nil || len(runs) != 3 || runs[2].ExitCode != 2 {
t.Fatalf("Runs=%v err=%v", runs, err)
}
@@ -280,7 +280,7 @@ func TestStoreRejectsUnsafeID(t *testing.T) {
if _, err := s.Get(id); err == nil {
t.Fatalf("Get(%q) must be rejected", id)
}
if _, err := s.Runs(id, 0); err == nil {
if _, err := s.Runs(id); err == nil {
t.Fatalf("Runs(%q) must be rejected", id)
}
}
-73
View File
@@ -1,73 +0,0 @@
//go:build windows
package cron
import (
"os"
"path/filepath"
"strings"
"testing"
"golang.org/x/sys/windows"
)
// Compacting runs.jsonl publishes a temporary file carrying the job directory's
// inherited DACL. A rename would replace a restrictive DACL applied to the log
// itself, so AppendRun must keep the log's own descriptor.
func TestAppendRunPreservesHistoryDACL(t *testing.T) {
store := newTestStore(t)
job, err := store.Add(Job{Prompt: "x"})
if err != nil {
t.Fatal(err)
}
path := filepath.Join(store.jobDir(job.ID), "runs.jsonl")
if err := os.WriteFile(path, []byte("{\"exitCode\":7}\n"), 0o600); err != nil {
t.Fatal(err)
}
// A protected DACL granting only the owner: distinct from whatever the
// temporary file inherits from the directory.
restricted, err := windows.SecurityDescriptorFromString("D:P(A;;FA;;;OW)")
if err != nil {
t.Skipf("cannot build a test security descriptor: %v", err)
}
dacl, _, err := restricted.DACL()
if err != nil {
t.Skipf("cannot read the test DACL: %v", err)
}
if err := windows.SetNamedSecurityInfo(
path,
windows.SE_FILE_OBJECT,
windows.DACL_SECURITY_INFORMATION|windows.PROTECTED_DACL_SECURITY_INFORMATION,
nil, nil, dacl, nil,
); err != nil {
t.Skipf("cannot apply a restrictive DACL on this filesystem: %v", err)
}
want := describeDACL(t, path)
if inherited := describeDACL(t, store.jobDir(job.ID)); strings.Contains(inherited, want) {
t.Skip("the job directory already carries the same DACL; this filesystem cannot show the difference")
}
if err := store.AppendRun(job.ID, RunRecord{ExitCode: 42}); err != nil {
t.Fatal(err)
}
if got := describeDACL(t, path); got != want {
t.Fatalf("DACL after AppendRun = %q, want the log's own %q", got, want)
}
runs, err := store.Runs(job.ID, 10)
if err != nil || len(runs) != 2 || runs[1].ExitCode != 42 {
t.Fatalf("outcome not published by replacement: %+v, %v", runs, err)
}
}
func describeDACL(t *testing.T, path string) string {
t.Helper()
sd, err := windows.GetNamedSecurityInfo(path, windows.SE_FILE_OBJECT, windows.DACL_SECURITY_INFORMATION)
if err != nil {
t.Skipf("cannot read the security descriptor of %s: %v", path, err)
}
text := sd.String()
if index := strings.Index(text, "D:"); index >= 0 {
return text[index:]
}
return text
}
+9 -8
View File
@@ -72,10 +72,9 @@ func (w *execWorker) Kill() error {
if w.cmd.Process == nil {
return nil
}
// The group the worker was launched into, not one rediscovered from its PID:
// the pool may call this after the worker has exited and been waited, and its
// children can still be running. See background.TerminateCommandGroup.
return background.TerminateCommandGroup(w.cmd)
// background.TerminateProcess is the cross-platform terminate (kills the
// process group on POSIX, taskkill /T on Windows).
return background.TerminateProcess(w.cmd.Process.Pid)
}
// readerLines adapts a bufio.Reader to the Lines interface. Unlike a capped
@@ -155,11 +154,13 @@ func NewExecLauncher(cfg ExecLauncherConfig) (Launcher, error) {
background.ConfigureChildProcessGroup(cmd)
// CommandContext's default cancel sends os.Process.Kill to the LEADER only,
// orphaning the process group we just configured (a stuck worker's children
// would survive ctx cancellation). Terminate the whole group instead, the
// same launch-time group terminate Kill() uses (D11). Cancel runs while Wait
// is in progress, so it must not reap.
// would survive ctx cancellation). Terminate the whole group instead — the
// same cross-platform group terminate Kill() uses (D11).
cmd.Cancel = func() error {
return background.TerminateCommandGroup(cmd)
if cmd.Process == nil {
return nil
}
return background.TerminateProcess(cmd.Process.Pid)
}
stdout, err := cmd.StdoutPipe()
@@ -1,90 +0,0 @@
//go:build !windows
package daemon
import (
"context"
"errors"
"os/exec"
"strconv"
"strings"
"syscall"
"testing"
"time"
)
// A WORKER'S CHILDREN ARE STOPPED THROUGH THE GROUP IT WAS LAUNCHED INTO.
//
// Kill and the Cancel hook used to hand the worker's PID to
// background.TerminateProcess, which looks the group up again with Getpgid when it
// signals. Once the worker has exited and been waited, that lookup fails, only the
// dead PID is signalled, and whatever the worker left running survives; on Darwin
// the lookup already fails while the exited worker is still a zombie. The worker
// here is a shell that backgrounds a sleep, prints its PID and exits, so after the
// wait the sleep is the only process left in the worker's group. Reported in #861.
func TestExecWorkerStopsWhatItLeftRunningAfterItExits(t *testing.T) {
for _, stop := range []struct {
name string
run func(*execWorker) error
}{
{"Kill", func(worker *execWorker) error { return worker.Kill() }},
{"Cancel hook", func(worker *execWorker) error { return worker.cmd.Cancel() }},
} {
t.Run(stop.name, func(t *testing.T) {
launcher, err := NewExecLauncher(ExecLauncherConfig{
Executable: "/bin/sh",
BaseArgs: []string{"-c", "sleep 300 >/dev/null 2>&1 & echo $!"},
})
if err != nil {
t.Fatalf("NewExecLauncher: %v", err)
}
handle, err := launcher(context.Background(), WorkerSpec{})
if err != nil {
t.Fatalf("launch: %v", err)
}
worker, ok := handle.(*execWorker)
if !ok {
t.Fatalf("SETUP INVALID: the launcher returned %T, not the exec worker under test", handle)
}
line, ok, err := worker.Stdout().Next()
if err != nil || !ok {
t.Fatalf("read the child's pid: ok=%v err=%v", ok, err)
}
childPID, err := strconv.Atoi(strings.TrimSpace(line))
if err != nil {
t.Fatalf("parse the child's pid %q: %v", line, err)
}
t.Cleanup(func() { _ = syscall.Kill(childPID, syscall.SIGKILL) })
if _, err := worker.Wait(); err != nil {
t.Fatalf("wait for the worker: %v", err)
}
if childStopped(childPID) {
t.Fatalf("SETUP INVALID: the worker's child %d is not running, so there is nothing left to stop", childPID)
}
// The pool ignores what Kill returns; the outcome is what matters.
_ = stop.run(worker)
deadline := time.Now().Add(5 * time.Second)
for !childStopped(childPID) {
if time.Now().After(deadline) {
t.Fatalf("the worker's child %d survived %s after the worker was reaped", childPID, stop.name)
}
time.Sleep(20 * time.Millisecond)
}
})
}
}
// childStopped reports whether pid is gone or only waiting to be reaped.
func childStopped(pid int) bool {
if errors.Is(syscall.Kill(pid, syscall.Signal(0)), syscall.ESRCH) {
return true
}
state, err := exec.Command("ps", "-o", "stat=", "-p", strconv.Itoa(pid)).Output()
if err != nil {
return errors.Is(syscall.Kill(pid, syscall.Signal(0)), syscall.ESRCH)
}
return strings.HasPrefix(strings.TrimSpace(string(state)), "Z")
}
+1 -27
View File
@@ -17,28 +17,6 @@ const (
ExitPermanent = 76
)
// Exit codes `zero exec` (internal/cli/exec.go) uses to report a run that
// finished, or was stopped on purpose, rather than crashed. They are duplicated
// here because the daemon cannot import the CLI package; keep them in sync.
const (
exitUsage = 2 // bad flags or arguments: a retry fails the same way
exitProvider = 3 // provider failure the run already surfaced
exitIncomplete = 4 // the agent ran (edits, commands) and stopped unfinished
exitInterrupted = 130 // SIGINT: someone stopped it on purpose
)
// isFinalExit reports whether a worker exit code means the run is over and must
// not be retried. Retrying would repeat the agent's side effects and re-bill the
// provider, and append a second run to the same session stream. Only crash-type
// exits (1, signals) and launch/read errors are worth another attempt.
func isFinalExit(code int) bool {
switch code {
case exitUsage, exitProvider, exitIncomplete, exitInterrupted:
return true
}
return false
}
// defaultPoolSize, defaultMaxAttempts and defaultKillTimeout are used when a
// PoolOptions field is left zero.
const (
@@ -198,8 +176,7 @@ type Sink interface {
// Run leases a worker slot and dispatches spec to a worker, streaming its
// stream-json lines to sink. It is at-least-once with bounded retries: a worker
// that crashes (non-zero, non-permanent) is retried on a fresh worker after a
// backoff; ExitPermanent, a final `zero exec` exit (usage, provider, incomplete,
// interrupted) or exhausting MaxAttempts returns ErrPermanent. Run
// backoff; ExitPermanent or exhausting MaxAttempts returns ErrPermanent. Run
// queues when all slots are busy. The returned int is the final worker exit code.
func (p *Pool) Run(ctx context.Context, spec WorkerSpec, sink Sink) (int, error) {
if err := p.acquire(ctx); err != nil {
@@ -232,9 +209,6 @@ func (p *Pool) Run(ctx context.Context, spec WorkerSpec, sink Sink) (int, error)
case code == ExitPermanent:
p.logf("worker %d exited permanently (code=%d) — not retrying", stat.id, code)
return code, ErrPermanent
case isFinalExit(code):
p.logf("worker %d finished with exit code %d — not retrying", stat.id, code)
return code, fmt.Errorf("%w: worker %d exited code=%d", ErrPermanent, stat.id, code)
case code == ExitTempfail:
lastErr = fmt.Errorf("worker %d tempfail (code=%d)", stat.id, code)
p.logf("worker %d tempfail — retry after %s", stat.id, p.opts.TempfailDelay)
-30
View File
@@ -3,7 +3,6 @@ package daemon
import (
"context"
"errors"
"fmt"
"sync"
"sync/atomic"
"testing"
@@ -149,35 +148,6 @@ func TestPoolRunPermanentStopsImmediately(t *testing.T) {
}
}
// A `zero exec` run that ended with usage (2), provider (3), incomplete (4) or
// interrupted (130) is over: retrying would repeat its side effects and merge a
// second run into the same session stream.
func TestPoolRunFinalExitCodesDoNotRetry(t *testing.T) {
for _, code := range []int{2, 3, 4, 130} {
t.Run(fmt.Sprintf("exit_%d", code), func(t *testing.T) {
launcher, calls := seqLauncher(
&fakeWorker{pid: 1, exitCode: code, out: []string{"first run"}},
&fakeWorker{pid: 2, exitCode: 0, out: []string{"second run"}},
)
pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher, MaxAttempts: 5, Backoff: func(int) time.Duration { return 0 }})
sink := &collectSink{}
got, err := pool.Run(context.Background(), WorkerSpec{Session: "s"}, sink)
if !errors.Is(err, ErrPermanent) {
t.Fatalf("Run err = %v, want ErrPermanent", err)
}
if got != code {
t.Fatalf("Run code = %d, want the worker's exit code %d", got, code)
}
if *calls != 1 {
t.Fatalf("launcher calls = %d, want 1 (no retry)", *calls)
}
if len(sink.lines) != 1 || sink.lines[0] != "first run" {
t.Fatalf("sink lines = %v, want only the first run's output", sink.lines)
}
})
}
}
func TestPoolRunCapExhausted(t *testing.T) {
launcher, calls := seqLauncher(
&fakeWorker{pid: 1, exitCode: 1},
+1 -31
View File
@@ -32,9 +32,6 @@ const (
// ErrSessionExists is returned when starting a session ID already in use.
var ErrSessionExists = errors.New("daemon: session already exists")
// ErrSessionOverloaded means the manager's unfinished-session limit is reached.
var ErrSessionOverloaded = errors.New("daemon: pending session limit reached")
// ErrSessionNotFound is returned by Get/Attach for an unknown session ID.
var ErrSessionNotFound = errors.New("daemon: session not found")
@@ -222,9 +219,6 @@ func (s *Session) isFinished() bool {
// daemon does not accumulate finished sessions without limit.
const defaultMaxSessions = 256
// defaultMaxPending bounds queued and running sessions independently of history.
const defaultMaxPending = 256
// SessionManager owns the live session registry and routes each session to a
// pool worker. Mirrors session-manager.js + roster.js. Finished sessions are
// retained (so a late `attach` sees history) up to MaxSessions, past which the
@@ -233,10 +227,8 @@ type SessionManager struct {
pool *Pool
maxBuffer int
maxSessions int
maxPending int
mu sync.Mutex
pending int // queued + running; protected by mu
sessions map[string]*Session
order []string // creation order, for FIFO eviction of finished sessions
}
@@ -248,9 +240,6 @@ type SessionManagerOptions struct {
// MaxSessions caps the retained session registry; 0 => default. The oldest
// FINISHED sessions are evicted once the cap is exceeded.
MaxSessions int
// MaxPending caps unfinished sessions (queued plus running), independently
// of retained history. Non-positive values use the default of 256.
MaxPending int
}
// NewSessionManager builds a manager over pool.
@@ -262,15 +251,10 @@ func NewSessionManager(opts SessionManagerOptions) (*SessionManager, error) {
if maxSessions <= 0 {
maxSessions = defaultMaxSessions
}
maxPending := opts.MaxPending
if maxPending <= 0 {
maxPending = defaultMaxPending
}
return &SessionManager{
pool: opts.Pool,
maxBuffer: opts.MaxBuffer,
maxSessions: maxSessions,
maxPending: maxPending,
sessions: map[string]*Session{},
}, nil
}
@@ -278,23 +262,13 @@ func NewSessionManager(opts SessionManagerOptions) (*SessionManager, error) {
// Start creates a session for spec and dispatches it to the pool. It returns the
// Session immediately (non-blocking); the run proceeds in the background and the
// session transitions queued -> running -> done/failed. A duplicate ID is
// rejected with ErrSessionExists; a full admission limit returns
// ErrSessionOverloaded without creating a session or background goroutine.
// rejected with ErrSessionExists.
func (m *SessionManager) Start(ctx context.Context, spec WorkerSpec) (*Session, error) {
m.mu.Lock()
if _, exists := m.sessions[spec.Session]; exists {
m.mu.Unlock()
return nil, ErrSessionExists
}
if err := ctx.Err(); err != nil {
m.mu.Unlock()
return nil, err
}
if m.pending >= m.maxPending {
m.mu.Unlock()
return nil, ErrSessionOverloaded
}
m.pending++
sess := newSession(spec.Session, spec.Cwd, m.maxBuffer)
m.sessions[spec.Session] = sess
m.order = append(m.order, spec.Session)
@@ -303,11 +277,7 @@ func (m *SessionManager) Start(ctx context.Context, spec WorkerSpec) (*Session,
go func() {
code, err := m.pool.Run(ctx, spec, sess)
m.mu.Lock()
sess.finish(code, err)
m.pending--
m.pruneLocked()
m.mu.Unlock()
}()
return sess, nil
}
-126
View File
@@ -3,8 +3,6 @@ package daemon
import (
"context"
"errors"
"fmt"
"sync"
"sync/atomic"
"testing"
"time"
@@ -210,130 +208,6 @@ func TestSessionManagerKeepsRunningOverCap(t *testing.T) {
}
}
func TestSessionManagerAdmissionBound(t *testing.T) {
block := make(chan struct{})
launcher, _ := seqLauncher(&fakeWorker{pid: 1, waitCh: block})
pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher})
mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool, MaxSessions: 2})
ctx, cancel := context.WithCancel(context.Background())
defer func() {
cancel()
close(block)
testutil.WaitFor(t, "all admitted sessions finished", func() bool {
mgr.mu.Lock()
defer mgr.mu.Unlock()
return mgr.pending == 0
})
}()
var wg sync.WaitGroup
var accepted atomic.Int32
for i := 0; i < 300; i++ {
wg.Add(1)
go func() {
defer wg.Done()
if _, err := mgr.Start(ctx, WorkerSpec{Session: fmt.Sprintf("s%d", i)}); err == nil {
accepted.Add(1)
}
}()
}
wg.Wait()
if got := accepted.Load(); got != 256 {
t.Fatalf("accepted %d unfinished sessions, want default bound 256", got)
}
if got := len(mgr.Statuses()); got != 256 {
t.Fatalf("registered %d sessions, want 256", got)
}
}
func TestSessionManagerCancellationReleasesAdmission(t *testing.T) {
block := make(chan struct{})
var release sync.Once
defer release.Do(func() { close(block) })
launcher, _ := seqLauncher(&fakeWorker{pid: 1, waitCh: block}, &fakeWorker{pid: 2, exitCode: ExitPermanent}, &fakeWorker{pid: 3})
pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher})
mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool, MaxPending: 2, MaxSessions: 10})
first, err := mgr.Start(context.Background(), WorkerSpec{Session: "running"})
if err != nil {
t.Fatal(err)
}
testutil.WaitFor(t, "first running", func() bool { return first.State() == SessionRunning })
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
queued, err := mgr.Start(ctx, WorkerSpec{Session: "queued"})
if err != nil {
t.Fatal(err)
}
if s, err := mgr.Start(context.Background(), WorkerSpec{Session: "retry"}); s != nil || !errors.Is(err, ErrSessionOverloaded) {
t.Fatalf("overloaded Start = (%v, %v)", s, err)
}
if _, ok := mgr.Get("retry"); ok {
t.Fatal("rejected session was registered")
}
if _, err := mgr.Start(context.Background(), WorkerSpec{Session: "running"}); !errors.Is(err, ErrSessionExists) {
t.Fatalf("duplicate at capacity: %v", err)
}
cancel()
testutil.WaitFor(t, "queued cancellation", func() bool { return queued.State() == SessionFailed })
if !errors.Is(queued.Err(), context.Canceled) {
t.Fatalf("queued error = %v", queued.Err())
}
if s, err := mgr.Start(ctx, WorkerSpec{Session: "canceled"}); s != nil || !errors.Is(err, context.Canceled) {
t.Fatalf("canceled Start = (%v, %v)", s, err)
}
retry, err := mgr.Start(context.Background(), WorkerSpec{Session: "retry"})
if err != nil {
t.Fatalf("capacity not released after cancellation: %v", err)
}
release.Do(func() { close(block) })
testutil.WaitFor(t, "first completed", func() bool { return first.State() == SessionDone })
testutil.WaitFor(t, "retry failed", func() bool { return retry.State() == SessionFailed })
if !errors.Is(retry.Err(), ErrPermanent) {
t.Fatalf("retry error = %v", retry.Err())
}
last, err := mgr.Start(context.Background(), WorkerSpec{Session: "last"})
if err != nil {
t.Fatal(err)
}
testutil.WaitFor(t, "last completed", func() bool { return last.State() == SessionDone })
mgr.mu.Lock()
defer mgr.mu.Unlock()
if mgr.pending != 0 || len(mgr.sessions) != 4 {
t.Fatalf("pending=%d retained=%d, want 0 and 4", mgr.pending, len(mgr.sessions))
}
}
func TestSessionManagerPrunesOnCompletion(t *testing.T) {
block := make(chan struct{})
var release sync.Once
defer release.Do(func() { close(block) })
launcher, _ := seqLauncher(
&fakeWorker{pid: 1, waitCh: block},
&fakeWorker{pid: 2},
&fakeWorker{pid: 3, exitCode: ExitPermanent},
)
pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher})
mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool, MaxSessions: 2})
sessions := make([]*Session, 3)
for i := range sessions {
var err error
sessions[i], err = mgr.Start(context.Background(), WorkerSpec{Session: fmt.Sprintf("s%d", i)})
if err != nil {
t.Fatal(err)
}
}
release.Do(func() { close(block) })
for _, s := range sessions {
select {
case <-s.Done():
case <-time.After(2 * time.Second):
t.Fatal("session did not finish")
}
}
if got := len(mgr.Statuses()); got != 2 {
t.Fatalf("retained %d completed sessions without another Start, want 2", got)
}
}
func TestSessionStatuses(t *testing.T) {
launcher, _ := seqLauncher(&fakeWorker{pid: 1, out: []string{"l"}, exitCode: 0})
pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher})
-8
View File
@@ -75,14 +75,6 @@ func RenameWithRetry(src, dst string, rename func(src, dst string) error) error
return err
}
// IsSharingOrLockViolation reports whether err is a Windows sharing or lock
// violation: another open handle on the file denies the requested access. It
// is always false on other platforms, where the same errno values mean
// unrelated errors.
func IsSharingOrLockViolation(err error) bool {
return runtime.GOOS == "windows" && isWindowsSharingOrLockViolation(err)
}
func isWindowsSharingOrLockViolation(err error) bool {
var errno syscall.Errno
if errors.As(err, &errno) {
-20
View File
@@ -2,7 +2,6 @@ package fsutil
import (
"errors"
"fmt"
"runtime"
"syscall"
"testing"
@@ -46,22 +45,3 @@ func TestRenameWithRetryNonRetryableError(t *testing.T) {
t.Errorf("expected only 1 attempt for non-retryable error, got %d", attempts)
}
}
// Errno 32 and 33 are EPIPE and EDOM on Unix, so only Windows may classify them.
func TestIsSharingOrLockViolation(t *testing.T) {
onWindows := runtime.GOOS == "windows"
for _, tc := range []struct {
err error
want bool
}{
{syscall.Errno(32), onWindows},
{fmt.Errorf("wrapped: %w", syscall.Errno(33)), onWindows},
{syscall.Errno(5), false}, // ERROR_ACCESS_DENIED
{errors.New("other"), false},
{nil, false},
} {
if got := IsSharingOrLockViolation(tc.err); got != tc.want {
t.Errorf("IsSharingOrLockViolation(%v) = %v, want %v", tc.err, got, tc.want)
}
}
}
+23 -153
View File
@@ -3,13 +3,11 @@ package mcp
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"io"
"math"
"net/http"
"os"
"os/exec"
"strconv"
@@ -18,8 +16,6 @@ import (
"time"
"github.com/Gitlawb/zero/internal/execution"
"github.com/Gitlawb/zero/internal/imageinput"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
type RemoteTool struct {
@@ -31,16 +27,10 @@ type RemoteTool struct {
type Content struct {
Type string `json:"type"`
Text string `json:"text,omitempty"`
// MimeType names what a non-text block holds. Additive and omitempty, so
// servers that never send it decode exactly as before. Image blocks with
// valid data are forwarded on Result.Images; MimeType is what lets a
// remaining dropped block (audio, resource, failed decode) be described
// instead of vanishing (#823).
// MimeType names what a non-text block holds. Decoded but not yet forwarded:
// it is what lets a dropped block be described to the model instead of
// vanishing (#823). Servers that omit it still decode fine.
MimeType string `json:"mimeType,omitempty"`
// Data is the MCP image block's base64 payload. Additive and omitempty so a
// result that never sent it still unmarshals. Decoded only for type
// "image"; other types leave it unread.
Data string `json:"data,omitempty"`
}
type CallToolResult struct {
@@ -391,6 +381,10 @@ func (client *Client) request(ctx context.Context, method string, params any, ta
}
}
func (client *Client) ensureWriter() {
_ = client.startWriter()
}
func (client *Client) startWriter() error {
client.writeMu.Lock()
defer client.writeMu.Unlock()
@@ -752,95 +746,24 @@ func TextContent(content []Content) string {
return strings.TrimSpace(strings.Join(parts, "\n"))
}
// itemDisp is the per-item forwarding/drop disposition produced by the
// single-pass conversion. DroppedContentSummary is built from this so a
// valid image is never base64-decoded a second time just to name what was
// kept versus dropped.
type itemDisp uint8
const (
dispText itemDisp = iota
dispForwarded
dispDropped
dispBudgetExceeded
dispUninspected
)
// decodeImageBase64 is the MCP image payload decoder. Tests replace it to
// count decode attempts; production uses standard base64.
var decodeImageBase64 = base64.StdEncoding.DecodeString
// Bound provider content-block overhead independently of decoded image size.
const maxForwardedImages = 16
// Bound decode work even when invalid or oversized-for-the-residue candidates
// do not consume the forwarding budgets.
const maxInspectedImages = 32
type imageLimitReason int
const (
imageLimitNone imageLimitReason = iota
imageLimitBudget
imageLimitForwardCount
imageLimitInspectionCount
)
func forwardImages(content []Content) ([]zeroruntime.ImageBlock, []itemDisp, imageLimitReason) {
disp := make([]itemDisp, len(content))
var images []zeroruntime.ImageBlock
remaining := imageinput.MaxImageBytes
inspected := 0
var limitReason imageLimitReason
for i, item := range content {
if item.Type == "text" {
disp[i] = dispText
continue
}
if item.Type == "image" {
if remaining == 0 {
disp[i] = dispUninspected
if limitReason == imageLimitNone {
limitReason = imageLimitBudget
}
continue
}
if len(images) >= maxForwardedImages {
disp[i] = dispUninspected
if limitReason == imageLimitNone {
limitReason = imageLimitForwardCount
}
continue
}
if inspected >= maxInspectedImages {
disp[i] = dispUninspected
if limitReason == imageLimitNone {
limitReason = imageLimitInspectionCount
}
continue
}
inspected++
if image, ok := imageBlockFromContent(item); ok {
if len(image.Data) <= remaining {
images = append(images, image)
remaining -= len(image.Data)
disp[i] = dispForwarded
continue
}
disp[i] = dispBudgetExceeded
continue
}
}
disp[i] = dispDropped
}
return images, disp, limitReason
}
func droppedContentNote(content []Content, disp []itemDisp, kinds ...itemDisp) string {
// DroppedContentSummary describes the blocks TextContent discards, e.g.
// "1 image/png block" or "2 resource blocks, 1 audio/wav block". It returns ""
// when a result is entirely text, so a caller adds nothing to the ordinary case.
//
// This exists because dropping silently is the worst available behaviour. A
// screenshot server returns a valid image, TextContent keeps nothing, and the
// call is reported as "(empty MCP tool result)" — so the model concludes the
// tool produced nothing and usually retries, burning another call on the same
// empty answer. Naming what came back costs nothing and ends that loop even
// though the payload still cannot be forwarded.
//
// Counts are grouped by mime type and ordered by first appearance, so the same
// result always produces the same sentence.
func DroppedContentSummary(content []Content) string {
labels := make([]string, 0, len(content))
counts := make(map[string]int, len(content))
for i, item := range content {
if i >= len(disp) || !dispKind(disp[i], kinds) {
for _, item := range content {
if item.Type == "text" {
continue
}
// Prefer the mime type: "image/png" tells the reader more than "image".
@@ -871,56 +794,3 @@ func droppedContentNote(content []Content, disp []itemDisp, kinds ...itemDisp) s
}
return strings.Join(parts, ", ")
}
func dispKind(got itemDisp, kinds []itemDisp) bool {
for _, kind := range kinds {
if got == kind {
return true
}
}
return false
}
func dispCount(disp []itemDisp, kinds ...itemDisp) int {
count := 0
for _, d := range disp {
if dispKind(d, kinds) {
count++
}
}
return count
}
func imageBlockFromContent(item Content) (zeroruntime.ImageBlock, bool) {
if item.Type != "image" {
return zeroruntime.ImageBlock{}, false
}
raw := strings.TrimSpace(item.Data)
if raw == "" {
return zeroruntime.ImageBlock{}, false
}
// EncodedLen(MaxImageBytes) is the encoded size of an image that decodes
// to exactly the inclusive cap, including "==" padding. DecodedLen is an
// upper bound and reports cap+2 for that input, so using it here would
// reject a valid at-limit PNG. The post-decode len(data) check is the
// exact backstop.
if len(raw) > base64.StdEncoding.EncodedLen(imageinput.MaxImageBytes) {
return zeroruntime.ImageBlock{}, false
}
data, err := decodeImageBase64(raw)
if err != nil {
return zeroruntime.ImageBlock{}, false
}
if len(data) == 0 || len(data) > imageinput.MaxImageBytes {
return zeroruntime.ImageBlock{}, false
}
sniffLen := len(data)
if sniffLen > 512 {
sniffLen = 512
}
mediaType := zeroruntime.NormalizeImageMediaType(http.DetectContentType(data[:sniffLen]))
if mediaType == "" {
return zeroruntime.ImageBlock{}, false
}
return zeroruntime.ImageBlock{MediaType: mediaType, Data: data}, true
}
+1 -1
View File
@@ -1420,7 +1420,7 @@ func TestWriterLoopExitsOnRepeatedClose(t *testing.T) {
reader: newMessageReader(inReader),
writer: newMessageWriter(outWriter),
}
_ = client.startWriter()
client.ensureWriter()
if err := client.writeMessage(context.Background(), rpcMessage{Method: "notifications/ping"}); err != nil {
t.Fatalf("writeMessage: %v", err)
}
-19
View File
@@ -1,19 +0,0 @@
package mcp
import "github.com/Gitlawb/zero/internal/zeroruntime"
// DroppedContentSummary describes the blocks that were not forwarded, e.g.
// "1 audio/wav block" or "2 resource blocks, 1 image/png block". It returns ""
// when every block is text or an image that ImageBlocks successfully
// forwarded, so a caller adds nothing to the ordinary case.
func DroppedContentSummary(content []Content) string {
_, disp, _ := forwardImages(content)
return droppedContentNote(content, disp, dispDropped, dispBudgetExceeded, dispUninspected)
}
// ImageBlocks converts MCP image content into the same ImageBlock channel
// capture tools already use.
func ImageBlocks(content []Content) []zeroruntime.ImageBlock {
images, _, _ := forwardImages(content)
return images
}
+6 -9
View File
@@ -657,10 +657,11 @@ func decodeSSERPCMessage(reader io.Reader) (rpcMessage, error) {
return rpcMessage{}, fmt.Errorf("missing MCP SSE response data")
}
// maxSSEEventBytes bounds a single SSE line/event. It accommodates multi-image
// responses (such as 8 MiB + 4 MiB) in base64 (~16.8 MiB) plus JSON-RPC envelope,
// metadata, and SSE framing overhead.
const maxSSEEventBytes = 32 * 1024 * 1024
// maxSSEEventBytes bounds a single SSE line/event. The previous 1 MiB cap made a
// large but legitimate MCP message (e.g. a big tool result) hit bufio.ErrTooLong,
// which failed the request permanently with no recovery. Raise it to a generous
// bound that still protects against an unbounded remote server.
const maxSSEEventBytes = 8 * 1024 * 1024
func scanSSEEvents(reader io.Reader, handle func(sseEvent) bool) error {
scanner := bufio.NewScanner(reader)
@@ -870,11 +871,7 @@ func (source *storeTokenSource) Refresh(ctx context.Context) (string, error) {
if !ok {
return "", fmt.Errorf("no stored OAuth token for MCP server %s", source.server.Name)
}
cfg, client, err := refreshSettings(ctx, source.httpClient, source.server.URL, source.config(), token)
if err != nil {
return "", err
}
refreshed, err := refreshAccessToken(ctx, client, cfg, token, source.now)
refreshed, err := refreshAccessToken(ctx, source.httpClient, source.config(), token, source.now)
if err != nil {
return "", err
}
-78
View File
@@ -2,7 +2,6 @@ package mcp
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"net/url"
@@ -441,80 +440,3 @@ func TestDeliverEventMessageSkipsMethodPresence(t *testing.T) {
t.Fatal("true response must complete pending")
}
}
func TestScanSSEEventsLargeImagePayload(t *testing.T) {
// Construct a 10 MiB image payload (~13.98 MiB in base64).
tenMiB := 10 * 1024 * 1024
imgB64 := paddedPNGBase64(tenMiB)
resultJSON, err := json.Marshal(map[string]any{
"content": []map[string]any{
{"type": "text", "text": "analysis screenshot"},
{"type": "image", "mimeType": "image/png", "data": imgB64},
},
})
if err != nil {
t.Fatal(err)
}
rpcJSON, err := json.Marshal(map[string]any{
"jsonrpc": "2.0",
"id": 1,
"result": json.RawMessage(resultJSON),
})
if err != nil {
t.Fatal(err)
}
t.Run("Single data line with 10 MiB image", func(t *testing.T) {
stream := "event: message\ndata: " + string(rpcJSON) + "\n\n"
msg, err := decodeSSERPCMessage(strings.NewReader(stream))
if err != nil {
t.Fatalf("decodeSSERPCMessage failed on single-line 10 MiB image: %v", err)
}
if !rpcIDMatches(msg.ID, 1) {
t.Fatalf("expected id 1, got %#v", msg.ID)
}
var decoded CallToolResult
if err := json.Unmarshal(msg.Result, &decoded); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if len(decoded.Content) != 2 || len(decoded.Content[1].Data) != len(imgB64) {
t.Fatalf("image payload did not survive the SSE round trip: got %d bytes, want %d",
len(decoded.Content[1].Data), len(imgB64))
}
})
t.Run("Multi data lines with 10 MiB image", func(t *testing.T) {
stream := "event: message\ndata: {\"jsonrpc\":\"2.0\",\"id\":1,\ndata: \"result\":" + string(resultJSON) + "}\n\n"
msg, err := decodeSSERPCMessage(strings.NewReader(stream))
if err != nil {
t.Fatalf("decodeSSERPCMessage failed on multi-line 10 MiB image: %v", err)
}
if !rpcIDMatches(msg.ID, 1) {
t.Fatalf("expected id 1, got %#v", msg.ID)
}
var decoded CallToolResult
if err := json.Unmarshal(msg.Result, &decoded); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if len(decoded.Content) != 2 || len(decoded.Content[1].Data) != len(imgB64) {
t.Fatalf("image payload did not survive the SSE round trip: got %d bytes, want %d",
len(decoded.Content[1].Data), len(imgB64))
}
})
t.Run("Oversized event exceeding 32 MiB rejected cleanly", func(t *testing.T) {
oversizedRPC, marshalErr := json.Marshal(map[string]any{
"jsonrpc": "2.0",
"id": 1,
"result": map[string]any{"padding": strings.Repeat("A", 33*1024*1024)},
})
if marshalErr != nil {
t.Fatal(marshalErr)
}
stream := "event: message\ndata: " + string(oversizedRPC) + "\n\n"
_, err := decodeSSERPCMessage(strings.NewReader(stream))
if err == nil {
t.Fatal("expected error for 33 MiB event, got nil")
}
})
}
+7 -581
View File
@@ -2,35 +2,20 @@ package mcp
import (
"context"
"encoding/base64"
"encoding/json"
"strings"
"testing"
"github.com/Gitlawb/zero/internal/imageinput"
"github.com/Gitlawb/zero/internal/tools"
)
// tinyPNGBase64 is a 1x1 PNG. http.DetectContentType sniffs it as image/png.
const tinyPNGBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
// paddedPNGBase64 is a decodable image/png of decoded length n. The 8-byte
// PNG magic is enough for http.DetectContentType; the rest is padding so
// tests can hit size caps without committing a multi-mebibyte fixture.
func paddedPNGBase64(n int) string {
raw := make([]byte, n)
copy(raw, "\x89PNG\r\n\x1a\n")
return base64.StdEncoding.EncodeToString(raw)
}
// A server that returns an image WITHOUT payload still cannot be forwarded, so
// the delivered result must name the block rather than report "(empty MCP tool
// result)". The model then usually retries, which is the worst outcome (#823).
// A server that returns only an image currently reports "(empty MCP tool
// result)": TextContent keeps text blocks and drops the rest, so a successful
// call looks like it produced nothing. The model then usually retries, which is
// the worst outcome, and the user is never told an image existed (#823).
//
// Image blocks that DO carry data are forwarded on Result.Images (see the
// payload tests below). This case is the remaining drop path, and it has to be
// true of the DELIVERED result, so this drives registryTool.Run rather than
// the helper alone.
// Carrying the payload is a separate change. Naming what was dropped is what
// stops the retry loop, and it has to be true of the DELIVERED result, so this
// drives registryTool.Run rather than the helper alone.
func TestAnImageOnlyResultSaysWhatItReturned(t *testing.T) {
tool := registryTool{
client: &nonTextClient{content: []Content{
@@ -149,24 +134,6 @@ func TestDroppedContentSummaryNamesTheBlocks(t *testing.T) {
},
want: "2 resource blocks, 1 audio/wav block",
},
{
name: "successfully forwarded image is not named as dropped",
content: []Content{{Type: "image", MimeType: "image/png", Data: tinyPNGBase64}},
want: "",
},
{
name: "malformed image data is still named",
content: []Content{{Type: "image", MimeType: "image/png", Data: "%%%not-base64%%%"}},
want: "1 image/png block",
},
{
name: "forwarded image plus audio names only the audio",
content: []Content{
{Type: "image", MimeType: "image/png", Data: tinyPNGBase64},
{Type: "audio", MimeType: "audio/wav"},
},
want: "1 audio/wav block",
},
}
for _, test := range tests {
@@ -189,544 +156,3 @@ func (client *nonTextClient) CallTool(context.Context, string, map[string]any) (
}
func (client *nonTextClient) Close() error { return nil }
func TestAnImageWithPayloadIsForwarded(t *testing.T) {
tool := registryTool{
client: &nonTextClient{content: []Content{
{Type: "image", MimeType: "image/png", Data: tinyPNGBase64},
}},
server: Server{Name: "shots"},
remote: RemoteTool{Name: "screenshot"},
}
result := tool.Run(context.Background(), map[string]any{})
if result.Output != "[image returned by tool]" {
t.Fatalf("image-only Output = %q, want [image returned by tool] so the tool_result is not an empty body", result.Output)
}
if strings.Contains(result.Output, "cannot forward") {
t.Fatalf("a forwarded image is still described as unforwardable:\n%s", result.Output)
}
if len(result.Images) != 1 {
t.Fatalf("Images len = %d, want 1", len(result.Images))
}
if result.Images[0].MediaType != "image/png" {
t.Errorf("MediaType = %q, want image/png", result.Images[0].MediaType)
}
if len(result.Images[0].Data) == 0 {
t.Fatal("forwarded image has empty Data")
}
if DroppedContentSummary([]Content{{Type: "image", MimeType: "image/png", Data: tinyPNGBase64}}) != "" {
t.Fatal("drop summary named a successfully forwarded image")
}
if result.Status != tools.StatusOK {
t.Errorf("status = %v, want OK", result.Status)
}
}
func TestAudioIsStillDroppedAndNamed(t *testing.T) {
tool := registryTool{
client: &nonTextClient{content: []Content{
{Type: "audio", MimeType: "audio/wav"},
}},
server: Server{Name: "shots"},
remote: RemoteTool{Name: "clip"},
}
result := tool.Run(context.Background(), map[string]any{})
if len(result.Images) != 0 {
t.Fatalf("audio was forwarded as an image: %#v", result.Images)
}
if !strings.Contains(result.Output, "audio/wav") {
t.Errorf("the output does not name the dropped audio:\n%s", result.Output)
}
if !strings.Contains(result.Output, "cannot forward yet") {
t.Errorf("the output does not say the audio cannot be forwarded:\n%s", result.Output)
}
}
func TestTextAndImageKeepsTextAndForwardsImage(t *testing.T) {
tool := registryTool{
client: &nonTextClient{content: []Content{
{Type: "text", Text: "captured the page"},
{Type: "image", MimeType: "image/png", Data: tinyPNGBase64},
}},
server: Server{Name: "shots"},
remote: RemoteTool{Name: "screenshot"},
}
result := tool.Run(context.Background(), map[string]any{})
if !strings.Contains(result.Output, "captured the page") {
t.Errorf("the text block was lost:\n%s", result.Output)
}
if strings.Contains(result.Output, "cannot forward") {
t.Errorf("a forwarded image is still described as unforwardable:\n%s", result.Output)
}
if strings.Contains(result.Output, "[image forwarded]") {
t.Errorf("text+image result substituted a placeholder over the text:\n%s", result.Output)
}
if len(result.Images) != 1 {
t.Fatalf("Images len = %d, want 1", len(result.Images))
}
if result.Images[0].MediaType != "image/png" {
t.Errorf("MediaType = %q, want image/png", result.Images[0].MediaType)
}
}
func TestMalformedImageDataDoesNotPanic(t *testing.T) {
tool := registryTool{
client: &nonTextClient{content: []Content{
{Type: "image", MimeType: "image/png", Data: "%%%not-base64%%%"},
}},
server: Server{Name: "shots"},
remote: RemoteTool{Name: "screenshot"},
}
result := tool.Run(context.Background(), map[string]any{})
if len(result.Images) != 0 {
t.Fatalf("malformed image was forwarded: %#v", result.Images)
}
if !strings.Contains(result.Output, "image/png") {
t.Errorf("malformed image was not named in the drop summary:\n%s", result.Output)
}
if strings.Contains(result.Output, "(empty MCP tool result)") {
t.Errorf("malformed image still reported as empty:\n%s", result.Output)
}
}
func TestImageContentJSONDecodesDataAndStaysCompatibleWithoutIt(t *testing.T) {
var withData CallToolResult
raw := []byte(`{"content":[{"type":"image","mimeType":"image/png","data":"` + tinyPNGBase64 + `"}]}`)
if err := json.Unmarshal(raw, &withData); err != nil {
t.Fatalf("unmarshal image content: %v", err)
}
if len(withData.Content) != 1 {
t.Fatalf("content len = %d, want 1", len(withData.Content))
}
if withData.Content[0].Type != "image" || withData.Content[0].MimeType != "image/png" {
t.Fatalf("decoded fields = %+v", withData.Content[0])
}
if withData.Content[0].Data != tinyPNGBase64 {
t.Fatalf("data = %q, want tiny PNG base64", withData.Content[0].Data)
}
var withoutData CallToolResult
if err := json.Unmarshal([]byte(`{"content":[{"type":"image","mimeType":"image/png"}]}`), &withoutData); err != nil {
t.Fatalf("unmarshal image content without data: %v", err)
}
if withoutData.Content[0].Data != "" {
t.Fatalf("absent data decoded as %q, want empty", withoutData.Content[0].Data)
}
}
func TestAnExactlyMaxImageBytesPaddedPNGIsForwarded(t *testing.T) {
payload := paddedPNGBase64(imageinput.MaxImageBytes)
if got := base64.StdEncoding.DecodedLen(len(payload)); got <= imageinput.MaxImageBytes {
t.Fatalf("fixture DecodedLen = %d, want > %d so the old bound would reject it", got, imageinput.MaxImageBytes)
}
tool := registryTool{
client: &nonTextClient{content: []Content{
{Type: "image", MimeType: "image/png", Data: payload},
}},
server: Server{Name: "shots"},
remote: RemoteTool{Name: "screenshot"},
}
result := tool.Run(context.Background(), map[string]any{})
if len(result.Images) != 1 {
t.Fatalf("at-limit padded PNG was dropped: images=%d output=%q", len(result.Images), result.Output)
}
if got := len(result.Images[0].Data); got != imageinput.MaxImageBytes {
t.Fatalf("forwarded size = %d, want %d", got, imageinput.MaxImageBytes)
}
}
func TestAnOversizedImageIsDroppedAndNamed(t *testing.T) {
tool := registryTool{
client: &nonTextClient{content: []Content{
{Type: "image", MimeType: "image/png", Data: paddedPNGBase64(imageinput.MaxImageBytes + 1)},
}},
server: Server{Name: "shots"},
remote: RemoteTool{Name: "screenshot"},
}
result := tool.Run(context.Background(), map[string]any{})
if len(result.Images) != 0 {
t.Fatalf("oversized image was forwarded: %#v", result.Images)
}
if !strings.Contains(result.Output, "image/png") {
t.Errorf("oversized image was not named in the drop summary:\n%s", result.Output)
}
if !strings.Contains(result.Output, "cannot forward yet") {
t.Errorf("individually oversized image should stay unforwardable, not a budget skip:\n%s", result.Output)
}
if strings.Contains(result.Output, "image budget") {
t.Errorf("individually oversized image was described as a budget skip:\n%s", result.Output)
}
}
func TestAggregateImageBudgetForwardsTheFirstAndNamesTheRest(t *testing.T) {
// Each payload is under the per-image cap; together they exceed the
// aggregate MaxImageBytes budget for one result. Identical bytes are
// deliberate: DroppedContentSummary must name the second even though
// imageBlockFromContent would accept it in isolation.
payload := paddedPNGBase64(imageinput.MaxImageBytes/2 + 1)
tool := registryTool{
client: &nonTextClient{content: []Content{
{Type: "image", MimeType: "image/png", Data: payload},
{Type: "image", MimeType: "image/png", Data: payload},
}},
server: Server{Name: "shots"},
remote: RemoteTool{Name: "screenshot"},
}
result := tool.Run(context.Background(), map[string]any{})
if len(result.Images) != 1 {
t.Fatalf("Images len = %d, want 1 (first fits the aggregate budget)", len(result.Images))
}
if got := len(result.Images[0].Data); got != imageinput.MaxImageBytes/2+1 {
t.Errorf("forwarded image size = %d, want %d", got, imageinput.MaxImageBytes/2+1)
}
if !strings.Contains(result.Output, "[image returned by tool]") {
t.Errorf("the forwarded first image has no placeholder:\n%s", result.Output)
}
if !strings.Contains(result.Output, "image/png") {
t.Errorf("the dropped second image was not named:\n%s", result.Output)
}
if !strings.Contains(result.Output, "which exceeded this result's remaining image budget") {
t.Errorf("the dropped second image is not described as a budget exceeded:\n%s", result.Output)
}
if !strings.Contains(result.Output, "Retrying with fewer images can recover this payload.") {
t.Errorf("the output does not tell the model a retry can recover the payload:\n%s", result.Output)
}
if strings.Contains(result.Output, "cannot forward yet") {
t.Errorf("a budget-exceeded image is described as unforwardable:\n%s", result.Output)
}
if strings.Contains(result.Output, "Retrying cannot recover this payload.") {
t.Errorf("a budget-exceeded image is described as unrecoverable:\n%s", result.Output)
}
}
func TestImagePayloadsAreDecodedOnceAndNotPastTheBudget(t *testing.T) {
orig := decodeImageBase64
t.Cleanup(func() { decodeImageBase64 = orig })
var n int
decodeImageBase64 = func(s string) ([]byte, error) {
n++
return orig(s)
}
// Two images fill the 10 MiB aggregate exactly, so remaining hits 0 and
// the later candidates must not be decoded at all.
half := imageinput.MaxImageBytes / 2
payload := paddedPNGBase64(half)
content := []Content{
{Type: "image", MimeType: "image/png", Data: payload},
{Type: "image", MimeType: "image/png", Data: payload},
{Type: "image", MimeType: "image/png", Data: payload},
{Type: "image", MimeType: "image/png", Data: "not-even-valid-base64!"},
{Type: "image", MimeType: "image/png", Data: ""},
{Type: "audio", MimeType: "audio/wav"},
}
n = 0
images := ImageBlocks(content)
if n != 2 {
t.Fatalf("ImageBlocks decoded %d payloads, want 2 (budget fills after two %d-byte images)", n, half)
}
if len(images) != 2 {
t.Fatalf("ImageBlocks len = %d, want 2", len(images))
}
n = 0
if got := DroppedContentSummary(content); got != "3 image/png blocks, 1 audio/wav block" {
t.Fatalf("DroppedContentSummary() = %q, want the three skipped images and the audio", got)
}
if n != 2 {
t.Fatalf("DroppedContentSummary decoded %d payloads, want 2 (same single pass as ImageBlocks)", n)
}
n = 0
result := registryTool{
client: &nonTextClient{content: content},
server: Server{Name: "shots"},
remote: RemoteTool{Name: "screenshot"},
}.Run(context.Background(), map[string]any{})
if n != 2 {
t.Fatalf("Run decoded %d payloads, want 2 (one pass; drop note must not decode again)", n)
}
if len(result.Images) != 2 {
t.Fatalf("Images len = %d, want 2", len(result.Images))
}
if !strings.Contains(result.Output, "[image returned by tool]") {
t.Errorf("forwarded images have no placeholder:\n%s", result.Output)
}
if !strings.Contains(result.Output, "image/png") {
t.Errorf("skipped images were not named:\n%s", result.Output)
}
if !strings.Contains(result.Output, "which were not inspected because the aggregate image budget was reached") {
t.Errorf("uninspected images are not described correctly:\n%s", result.Output)
}
if strings.Contains(result.Output, "Retrying with fewer images can recover this payload.") {
t.Errorf("uninspected images must not make unsupported recovery claims:\n%s", result.Output)
}
if !strings.Contains(result.Output, "audio/wav") {
t.Errorf("audio was not named:\n%s", result.Output)
}
if !strings.Contains(result.Output, "cannot forward yet") {
t.Errorf("audio is not described as unforwardable:\n%s", result.Output)
}
if strings.Contains(result.Output, "(empty MCP tool result)") {
t.Errorf("forwarded images still reported as empty:\n%s", result.Output)
}
n = 0
one := []Content{{Type: "image", MimeType: "image/png", Data: payload}}
oneResult := registryTool{
client: &nonTextClient{content: one},
server: Server{Name: "shots"},
remote: RemoteTool{Name: "screenshot"},
}.Run(context.Background(), map[string]any{})
if n != 1 {
t.Fatalf("one image decoded %d times, want 1", n)
}
if len(oneResult.Images) != 1 {
t.Fatalf("one-image Images len = %d, want 1", len(oneResult.Images))
}
if oneResult.Output != "[image returned by tool]" {
t.Fatalf("one-image Output = %q, want [image returned by tool]", oneResult.Output)
}
if strings.Contains(oneResult.Output, "cannot forward") {
t.Fatalf("a forwarded image is still described as unforwardable:\n%s", oneResult.Output)
}
}
func TestImageBudgetNonZeroResidueAllowsSmallerLaterImage(t *testing.T) {
// First image: 8 MiB (fits, 2 MiB left)
// Second image: 3 MiB (exceeds remaining 2 MiB, budgetExceeded)
// Third image: 1 MiB (fits in remaining 2 MiB, forwarded, 1 MiB left)
img8 := paddedPNGBase64(8 * 1024 * 1024)
img3 := paddedPNGBase64(3 * 1024 * 1024)
img1 := paddedPNGBase64(1 * 1024 * 1024)
content := []Content{
{Type: "image", MimeType: "image/png", Data: img8},
{Type: "image", MimeType: "image/png", Data: img3},
{Type: "image", MimeType: "image/png", Data: img1},
}
images, disp, _ := forwardImages(content)
if len(images) != 2 {
t.Fatalf("forwardImages len = %d, want 2 (8 MiB + 1 MiB)", len(images))
}
if disp[0] != dispForwarded || disp[1] != dispBudgetExceeded || disp[2] != dispForwarded {
t.Fatalf("dispositions = %v, want [forwarded, budgetExceeded, forwarded]", disp)
}
result := registryTool{
client: &nonTextClient{content: content},
server: Server{Name: "shots"},
remote: RemoteTool{Name: "screenshot"},
}.Run(context.Background(), map[string]any{})
if len(result.Images) != 2 {
t.Fatalf("result Images len = %d, want 2", len(result.Images))
}
if !strings.Contains(result.Output, "exceeded this result's remaining image budget") {
t.Fatalf("expected remaining budget notice in output:\n%s", result.Output)
}
if !strings.Contains(result.Output, "Retrying with fewer images can recover this payload.") {
t.Fatalf("expected retry recovery guidance for validated exceeded image:\n%s", result.Output)
}
}
func TestImageCountBudgetPreservesTextAndSkipsFurtherDecoding(t *testing.T) {
content := []Content{{Type: "image", Data: "invalid"}}
for range 17 {
content = append(content, Content{Type: "image", MimeType: "image/png", Data: tinyPNGBase64})
}
content = append(content, Content{Type: "text", Text: "trailing text"})
previous := decodeImageBase64
decodes := 0
decodeImageBase64 = func(s string) ([]byte, error) {
decodes++
return previous(s)
}
t.Cleanup(func() { decodeImageBase64 = previous })
result := registryTool{
client: &nonTextClient{content: content},
server: Server{Name: "shots"}, remote: RemoteTool{Name: "screenshot"},
}.Run(context.Background(), map[string]any{})
if len(result.Images) != 16 || decodes != 17 {
t.Fatalf("forwarded %d images with %d decodes, want 16 images and 17 attempts including invalid input", len(result.Images), decodes)
}
if !strings.Contains(result.Output, "trailing text") || !strings.Contains(result.Output, "1 image/png block, which was not inspected") {
t.Fatalf("missing text or count-budget notice: %s", result.Output)
}
}
func TestImageInspectionBudgetBoundsRejectedCandidates(t *testing.T) {
// Use the decoder seam to exercise accounting without allocating a large
// payload for every rejected candidate.
previous := decodeImageBase64
decodes := 0
large, _ := previous(paddedPNGBase64(8 * 1024 * 1024))
rejected, _ := previous(paddedPNGBase64(3 * 1024 * 1024))
decodeImageBase64 = func(string) ([]byte, error) {
decodes++
if decodes == 1 {
return large, nil
}
return rejected, nil
}
t.Cleanup(func() { decodeImageBase64 = previous })
content := make([]Content, 40)
for i := range content {
content[i] = Content{Type: "image", MimeType: "image/png", Data: tinyPNGBase64}
}
content = append(content, Content{Type: "text", Text: "trailing text"})
images, disp, _ := forwardImages(content)
if decodes != 32 || len(images) != 1 {
t.Fatalf("got %d decodes and %d forwarded images, want 32 and 1", decodes, len(images))
}
for i := 1; i < 40; i++ {
want := dispBudgetExceeded
if i >= 32 {
want = dispUninspected
}
if disp[i] != want {
t.Errorf("disposition[%d] = %v, want %v", i, disp[i], want)
}
}
if disp[40] != dispText {
t.Fatal("inspection limit discarded trailing text")
}
}
func BenchmarkForwardImagesFourHalfBudget(b *testing.B) {
payload := paddedPNGBase64(imageinput.MaxImageBytes / 2)
content := []Content{
{Type: "image", MimeType: "image/png", Data: payload},
{Type: "image", MimeType: "image/png", Data: payload},
{Type: "image", MimeType: "image/png", Data: payload},
{Type: "image", MimeType: "image/png", Data: payload},
}
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
_, _, _ = forwardImages(content)
}
}
func TestMCPResultUninspectedImageLimitingReasonNotes(t *testing.T) {
t.Run("forwarded count limit reached", func(t *testing.T) {
content := make([]Content, 20)
for i := range content {
content[i] = Content{Type: "image", MimeType: "image/png", Data: tinyPNGBase64}
}
result := registryTool{
client: &nonTextClient{content: content},
server: Server{Name: "shots"},
remote: RemoteTool{Name: "screenshot"},
}.Run(context.Background(), map[string]any{})
if len(result.Images) != 16 {
t.Fatalf("expected 16 forwarded images, got %d", len(result.Images))
}
if !strings.Contains(result.Output, "which were not inspected because the maximum forwarded image count was reached.") {
t.Fatalf("expected maximum forwarded image count message, got:\n%s", result.Output)
}
})
t.Run("inspection limit reached", func(t *testing.T) {
previous := decodeImageBase64
large := make([]byte, 1024)
rejected := make([]byte, imageinput.MaxImageBytes+1)
decodes := 0
decodeImageBase64 = func(s string) ([]byte, error) {
decodes++
if decodes == 1 {
return large, nil
}
return rejected, nil
}
t.Cleanup(func() { decodeImageBase64 = previous })
content := make([]Content, 35)
for i := range content {
content[i] = Content{Type: "image", MimeType: "image/png", Data: tinyPNGBase64}
}
result := registryTool{
client: &nonTextClient{content: content},
server: Server{Name: "shots"},
remote: RemoteTool{Name: "screenshot"},
}.Run(context.Background(), map[string]any{})
if !strings.Contains(result.Output, "which were not inspected because the maximum image inspection limit was reached.") {
t.Fatalf("expected maximum image inspection limit message, got:\n%s", result.Output)
}
})
}
// The pre-decode encoded-length cap is the decode-DoS guard: it keeps a
// hostile server from making Zero decode a payload of any size (up to the
// inspection cap times per result) before the post-decode backstop drops it.
// Removing the cap changes nothing that is forwarded, so no existing test
// notices; this pins the guard on the decode hook instead. A payload one
// byte over the encoded cap must be rejected without the decoder ever being
// called.
func TestOverCapEncodedPayloadIsRejectedBeforeDecoding(t *testing.T) {
previous := decodeImageBase64
decodes := 0
decodeImageBase64 = func(s string) ([]byte, error) {
decodes++
return previous(s)
}
t.Cleanup(func() { decodeImageBase64 = previous })
// An at-limit PNG would decode to exactly MaxImageBytes and forward; a
// single extra base64 character pushes it one byte over the encoded cap,
// so only the pre-decode guard can be responsible for rejecting it.
payload := paddedPNGBase64(imageinput.MaxImageBytes) + "A"
content := []Content{{Type: "image", MimeType: "image/png", Data: payload}}
if images := ImageBlocks(content); len(images) != 0 {
t.Fatalf("ImageBlocks forwarded %d images from an over-cap payload, want 0", len(images))
}
if decodes != 0 {
t.Fatalf("decoder ran %d times, want 0: the pre-decode cap must reject before any decode work", decodes)
}
if got := DroppedContentSummary(content); got != "1 image/png block" {
t.Fatalf("DroppedContentSummary() = %q, want the over-cap payload named as dropped", got)
}
}
// The media type is sniffed from the decoded bytes, never taken from the
// block's declared mimeType: a hostile server must not be able to relabel
// HTML or a script as image/png and have it forwarded to the provider as an
// image. Every other fixture pairs real PNG bytes with image/png, so
// sniffing and trusting the declared type agree there and no test can tell
// them apart; this fixture makes them disagree.
func TestDeclaredMimeTypeDisagreeingWithSniffedBytesIsDropped(t *testing.T) {
html := []byte("<!doctype html><html><body><script>alert(1)</script></body></html>")
tool := registryTool{
client: &nonTextClient{content: []Content{
{Type: "image", MimeType: "image/png", Data: base64.StdEncoding.EncodeToString(html)},
}},
server: Server{Name: "shots"},
remote: RemoteTool{Name: "screenshot"},
}
result := tool.Run(context.Background(), map[string]any{})
if len(result.Images) != 0 {
t.Fatalf("forwarded %d images from an HTML payload declared as image/png, want 0", len(result.Images))
}
if strings.Contains(result.Output, "[image returned by tool]") {
t.Errorf("a non-image payload declared as image/png is still announced as a returned image:\n%s", result.Output)
}
if !strings.Contains(result.Output, "image/png") {
t.Errorf("the dropped block is not named in the drop summary:\n%s", result.Output)
}
}
+26 -85
View File
@@ -59,11 +59,6 @@ func tokenToStored(t oauth.Token) StoredToken {
TokenType: t.TokenType,
Scopes: t.Scopes,
ExpiresAt: t.ExpiresAt,
TokenEndpoint: t.TokenEndpoint,
ProtectedTokenEndpoint: t.ProtectedTokenEndpoint,
ClientID: t.ClientID,
ClientSecret: t.ClientSecret,
}
}
@@ -346,53 +341,7 @@ func refreshAccessToken(ctx context.Context, client *http.Client, cfg OAuthConfi
if err != nil {
return StoredToken{}, err
}
refreshed := tokenToStored(token)
// The token response never carries what Login learned about the client and
// endpoint, so carry it over or the next refresh would lose it.
refreshed.TokenEndpoint = current.TokenEndpoint
refreshed.ProtectedTokenEndpoint = current.ProtectedTokenEndpoint
refreshed.ClientID = current.ClientID
refreshed.ClientSecret = current.ClientSecret
return refreshed, nil
}
// refreshSettings fills the gaps in the configured OAuth settings from what the
// login stored: the discovered token endpoint and the dynamically registered
// client. Explicit config always wins. A stored endpoint is re-validated before
// credentials are posted to it, and one that came from server-advertised
// discovery is also held to the public-network policy, with the returned client
// pinned to it, exactly as Login did.
func refreshSettings(ctx context.Context, client *http.Client, resourceURL string, cfg OAuthConfig, stored StoredToken) (OAuthConfig, *http.Client, error) {
if strings.TrimSpace(cfg.ClientID) == "" && strings.TrimSpace(stored.ClientID) != "" {
cfg.ClientID = stored.ClientID
if strings.TrimSpace(cfg.ClientSecret) == "" {
cfg.ClientSecret = stored.ClientSecret
}
}
if strings.TrimSpace(cfg.TokenEndpoint) != "" || strings.TrimSpace(stored.TokenEndpoint) == "" {
return cfg, client, nil
}
endpoint := strings.TrimSpace(stored.TokenEndpoint)
if err := oauth.ValidateEndpointURL(endpoint); err != nil {
return cfg, client, fmt.Errorf("mcp oauth: stored token endpoint: %w", err)
}
if stored.ProtectedTokenEndpoint {
metadata := authServerMetadata{
TokenEndpoint: endpoint,
protectedResourceURL: strings.TrimSpace(resourceURL),
protectTokenEndpoint: true,
}
if err := validateProtectedDiscoveredEndpoints(ctx, metadata); err != nil {
return cfg, client, err
}
protected, err := newAdvertisedOAuthDiscoveryClient(client, metadata.protectedResourceURL)
if err != nil {
return cfg, client, err
}
client = protected
}
cfg.TokenEndpoint = endpoint
return cfg, client, nil
return tokenToStored(token), nil
}
// registerClient performs dynamic client registration against the registration
@@ -497,7 +446,6 @@ func Login(ctx context.Context, options LoginOptions) (StoredToken, error) {
tokenClient = protectedEndpointClient
}
registeredClientID, registeredClientSecret := "", ""
if strings.TrimSpace(cfg.ClientID) == "" {
if registration := strings.TrimSpace(metadata.RegistrationEndpoint); registration != "" {
registrationClient := httpClient
@@ -509,10 +457,8 @@ func Login(ctx context.Context, options LoginOptions) (StoredToken, error) {
return StoredToken{}, regErr
}
cfg.ClientID = clientID
registeredClientID = clientID
if clientSecret != "" {
cfg.ClientSecret = clientSecret
registeredClientSecret = clientSecret
}
}
}
@@ -548,28 +494,31 @@ func Login(ctx context.Context, options LoginOptions) (StoredToken, error) {
err error
}
resultChan := make(chan callbackResult, 1)
// The shared callback server bounds header reads, force-closes a client
// still connected when shutdown runs out of time, and waits for Serve to
// return. The bare http.Server this used to build did none of that, so a
// client stalled mid-header outlived Login (#1028).
server := oauth.StartCallbackServer(listener, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/callback" {
http.NotFound(w, r)
return
}
code, parseErr := flow.parseCallback(r.URL.Query())
if parseErr != nil {
w.WriteHeader(http.StatusBadRequest)
_, _ = io.WriteString(w, "Authorization failed. You may close this window.")
} else {
_, _ = io.WriteString(w, "Authorization complete. You may close this window.")
}
select {
case resultChan <- callbackResult{code: code, err: parseErr}:
default:
}
}))
defer server.Close()
server := &http.Server{
Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/callback" {
http.NotFound(w, r)
return
}
code, parseErr := flow.parseCallback(r.URL.Query())
if parseErr != nil {
w.WriteHeader(http.StatusBadRequest)
_, _ = io.WriteString(w, "Authorization failed. You may close this window.")
} else {
_, _ = io.WriteString(w, "Authorization complete. You may close this window.")
}
select {
case resultChan <- callbackResult{code: code, err: parseErr}:
default:
}
}),
}
go func() { _ = server.Serve(listener) }()
defer func() {
shutdownCtx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
_ = server.Shutdown(shutdownCtx)
}()
open := options.OpenBrowser
if open == nil {
@@ -588,14 +537,6 @@ func Login(ctx context.Context, options LoginOptions) (StoredToken, error) {
if err != nil {
return StoredToken{}, err
}
// Persist only what config does not already carry, so refresh can reuse
// it without copying configured secrets into the token store.
if strings.TrimSpace(options.Config.TokenEndpoint) == "" {
token.TokenEndpoint = metadata.TokenEndpoint
token.ProtectedTokenEndpoint = metadata.protectTokenEndpoint
}
token.ClientID = registeredClientID
token.ClientSecret = registeredClientSecret
return token, nil
case <-loginCtx.Done():
return StoredToken{}, fmt.Errorf("timed out waiting for OAuth authorization callback: %w", loginCtx.Err())
-315
View File
@@ -1,315 +0,0 @@
package mcp
import (
"context"
"encoding/json"
"errors"
"net/http"
"net/http/httptest"
"net/url"
"path/filepath"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/Gitlawb/zero/internal/oauth"
)
const testServerIdentity = "0123456789abcdef0123456789abcdef"
// discoveryOAuthServer serves a standard `auth: oauth` setup: protected-resource
// metadata, authorization-server metadata with a registration endpoint, dynamic
// client registration, and a token endpoint that records refresh requests.
type discoveryOAuthServer struct {
*httptest.Server
refreshForms chan url.Values
}
func newDiscoveryOAuthServer(t *testing.T) *discoveryOAuthServer {
t.Helper()
fixture := &discoveryOAuthServer{refreshForms: make(chan url.Values, 4)}
fixture.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
base := "http://" + r.Host
switch r.URL.Path {
case "/mcp":
w.Header().Set("WWW-Authenticate", `Bearer resource_metadata="`+base+`/resource-metadata"`)
w.WriteHeader(http.StatusUnauthorized)
case "/resource-metadata":
_ = json.NewEncoder(w).Encode(map[string]any{
"resource": base + "/mcp",
"authorization_servers": []string{base},
})
case "/.well-known/oauth-authorization-server":
_ = json.NewEncoder(w).Encode(map[string]any{
"issuer": base,
"authorization_endpoint": base + "/authorize",
"token_endpoint": base + "/token",
"registration_endpoint": base + "/register",
})
case "/register":
_ = json.NewEncoder(w).Encode(map[string]any{
"client_id": "registered-client",
"client_secret": "registered-secret",
})
case "/token":
if err := r.ParseForm(); err != nil {
t.Errorf("parse token form: %v", err)
}
if r.Form.Get("grant_type") == "refresh_token" {
fixture.refreshForms <- r.Form
_ = json.NewEncoder(w).Encode(map[string]any{
"access_token": "access-refreshed",
"refresh_token": "refresh-rotated",
"token_type": "Bearer",
"expires_in": 3600,
})
return
}
_ = json.NewEncoder(w).Encode(map[string]any{
"access_token": "access-initial",
"refresh_token": "refresh-initial",
"token_type": "Bearer",
"expires_in": 3600,
})
default:
http.NotFound(w, r)
}
}))
t.Cleanup(fixture.Close)
return fixture
}
func driveLoopbackCallback(authURL string) error {
parsed, err := url.Parse(authURL)
if err != nil {
return err
}
callbackURL := parsed.Query().Get("redirect_uri") + "?code=auth-code&state=" + url.QueryEscape(parsed.Query().Get("state"))
go func() {
for attempt := 0; attempt < 20; attempt++ {
response, requestErr := http.Get(callbackURL)
if requestErr == nil {
response.Body.Close()
return
}
time.Sleep(25 * time.Millisecond)
}
}()
return nil
}
func TestLoginRecordsDiscoveredEndpointAndRegisteredClient(t *testing.T) {
fixture := newDiscoveryOAuthServer(t)
token, err := Login(context.Background(), LoginOptions{
ServerName: "notion-style",
ServerURL: fixture.URL + "/mcp",
Config: OAuthConfig{},
HTTPClient: fixture.Client(),
OpenBrowser: driveLoopbackCallback,
Timeout: 5 * time.Second,
Now: time.Now,
})
if err != nil {
t.Fatalf("Login() error = %v", err)
}
if token.TokenEndpoint != fixture.URL+"/token" || !token.ProtectedTokenEndpoint {
t.Fatalf("stored token endpoint = %q protected=%v, want the discovered protected endpoint", token.TokenEndpoint, token.ProtectedTokenEndpoint)
}
if token.ClientID != "registered-client" || token.ClientSecret != "registered-secret" {
t.Fatalf("stored client = %q/%q, want the dynamically registered client", token.ClientID, token.ClientSecret)
}
}
func TestLoginDoesNotCopyConfiguredClientOrEndpointIntoStore(t *testing.T) {
fixture := newDiscoveryOAuthServer(t)
token, err := Login(context.Background(), LoginOptions{
ServerName: "configured",
ServerURL: fixture.URL + "/mcp",
Config: OAuthConfig{
ClientID: "configured-client",
ClientSecret: "configured-secret",
TokenEndpoint: fixture.URL + "/token",
},
HTTPClient: fixture.Client(),
OpenBrowser: driveLoopbackCallback,
Timeout: 5 * time.Second,
Now: time.Now,
})
if err != nil {
t.Fatalf("Login() error = %v", err)
}
if token.TokenEndpoint != "" || token.ProtectedTokenEndpoint || token.ClientID != "" || token.ClientSecret != "" {
t.Fatalf("configured values were copied into the token store: %#v", token)
}
}
// The standard `auth: oauth` setup configures no token endpoint or client, so
// refresh must work from what the discovery-based login stored.
func TestStoreTokenSourceRefreshUsesStoredDiscoveryAndClient(t *testing.T) {
fixture := newDiscoveryOAuthServer(t)
server := Server{Name: "notion-style", URL: fixture.URL + "/mcp", Identity: testServerIdentity, Auth: ServerAuthOAuth}
loggedIn, err := Login(context.Background(), LoginOptions{
ServerName: server.Name,
ServerURL: server.URL,
HTTPClient: fixture.Client(),
OpenBrowser: driveLoopbackCallback,
Timeout: 5 * time.Second,
Now: time.Now,
})
if err != nil {
t.Fatalf("Login() error = %v", err)
}
store, err := NewTokenStore(TokenStoreOptions{FilePath: filepath.Join(t.TempDir(), "oauth-tokens.json")})
if err != nil {
t.Fatalf("NewTokenStore() error = %v", err)
}
if err := store.SaveForServer(server, loggedIn); err != nil {
t.Fatalf("SaveForServer() error = %v", err)
}
source := &storeTokenSource{server: server, store: store, httpClient: http.DefaultClient, now: time.Now}
access, err := source.Refresh(context.Background())
if err != nil {
t.Fatalf("Refresh() error = %v", err)
}
if access != "access-refreshed" {
t.Fatalf("Refresh() access token = %q", access)
}
form := <-fixture.refreshForms
if form.Get("client_id") != "registered-client" || form.Get("client_secret") != "registered-secret" || form.Get("refresh_token") != "refresh-initial" {
t.Fatalf("refresh form = %v, want the registered client and stored refresh token", form)
}
// The refresh must not drop what Login stored, or the next one would fail.
saved, ok, err := store.LoadForServer(server)
if err != nil || !ok {
t.Fatalf("LoadForServer() = %v, %v", ok, err)
}
if saved.RefreshToken != "refresh-rotated" || saved.TokenEndpoint != fixture.URL+"/token" || !saved.ProtectedTokenEndpoint ||
saved.ClientID != "registered-client" || saved.ClientSecret != "registered-secret" {
t.Fatalf("refreshed token lost the stored endpoint or client: %#v", saved)
}
if _, err := source.Refresh(context.Background()); err != nil {
t.Fatalf("second Refresh() error = %v", err)
}
}
func TestRefreshSettingsConfigOverridesStored(t *testing.T) {
cfg, client, err := refreshSettings(context.Background(), http.DefaultClient, "https://mcp.example.com/mcp",
OAuthConfig{ClientID: "configured-client", ClientSecret: "configured-secret", TokenEndpoint: "https://configured.example.com/token"},
StoredToken{TokenEndpoint: "https://stored.example.com/token", ProtectedTokenEndpoint: true, ClientID: "stored-client", ClientSecret: "stored-secret"})
if err != nil {
t.Fatalf("refreshSettings() error = %v", err)
}
if cfg.TokenEndpoint != "https://configured.example.com/token" || cfg.ClientID != "configured-client" || cfg.ClientSecret != "configured-secret" {
t.Fatalf("config did not win over stored values: %#v", cfg)
}
if client != http.DefaultClient {
t.Fatal("an explicitly configured endpoint must not switch to the protected-discovery client")
}
}
func TestRefreshSettingsKeepsConfiguredSecretWithStoredClientID(t *testing.T) {
cfg, _, err := refreshSettings(context.Background(), http.DefaultClient, "https://mcp.example.com/mcp",
OAuthConfig{ClientSecret: "configured-secret"},
StoredToken{ClientID: "stored-client", ClientSecret: "stored-secret"})
if err != nil {
t.Fatalf("refreshSettings() error = %v", err)
}
if cfg.ClientID != "stored-client" || cfg.ClientSecret != "configured-secret" {
t.Fatalf("client = %q/%q, want stored id with the configured secret", cfg.ClientID, cfg.ClientSecret)
}
}
func TestRefreshSettingsRejectsUnsafeStoredEndpoint(t *testing.T) {
for _, tc := range []struct {
name string
stored StoredToken
want error
}{
{
name: "cleartext endpoint",
stored: StoredToken{TokenEndpoint: "http://auth.example.com/token", ClientID: "c"},
want: oauth.ErrInsecureTokenEndpoint,
},
{
name: "advertised loopback endpoint",
stored: StoredToken{TokenEndpoint: "https://127.0.0.1/token", ProtectedTokenEndpoint: true, ClientID: "c"},
want: errUnsafeOAuthDiscoveryTarget,
},
{
name: "advertised private endpoint",
stored: StoredToken{TokenEndpoint: "https://10.0.0.5/token", ProtectedTokenEndpoint: true, ClientID: "c"},
want: errUnsafeOAuthDiscoveryTarget,
},
} {
t.Run(tc.name, func(t *testing.T) {
_, _, err := refreshSettings(context.Background(), http.DefaultClient, "https://mcp.example.com/mcp", OAuthConfig{}, tc.stored)
if !errors.Is(err, tc.want) {
t.Fatalf("refreshSettings() error = %v, want %v", err, tc.want)
}
})
}
}
func TestStoreTokenSourceRefreshDoesNotPostToUnsafeStoredEndpoint(t *testing.T) {
var hits atomic.Int64
loopback := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { hits.Add(1) }))
defer loopback.Close()
// A public MCP resource must not be able to point refresh at loopback.
server := Server{Name: "public", URL: "https://mcp.example.com/mcp", Identity: testServerIdentity, Auth: ServerAuthOAuth}
store, err := NewTokenStore(TokenStoreOptions{FilePath: filepath.Join(t.TempDir(), "oauth-tokens.json")})
if err != nil {
t.Fatalf("NewTokenStore() error = %v", err)
}
if err := store.SaveForServer(server, StoredToken{
AccessToken: "a", RefreshToken: "r", ClientID: "c",
TokenEndpoint: loopback.URL + "/token", ProtectedTokenEndpoint: true,
}); err != nil {
t.Fatalf("SaveForServer() error = %v", err)
}
source := &storeTokenSource{server: server, store: store, httpClient: http.DefaultClient, now: time.Now}
if _, err := source.Refresh(context.Background()); !errors.Is(err, errUnsafeOAuthDiscoveryTarget) {
t.Fatalf("Refresh() error = %v, want the discovery policy refusal", err)
}
if hits.Load() != 0 {
t.Fatalf("loopback endpoint received %d request(s)", hits.Load())
}
}
func TestStoredTokenPersistsDiscoveryFields(t *testing.T) {
store, err := NewTokenStore(TokenStoreOptions{FilePath: filepath.Join(t.TempDir(), "oauth-tokens.json")})
if err != nil {
t.Fatalf("NewTokenStore() error = %v", err)
}
want := StoredToken{
AccessToken: "a", RefreshToken: "r",
TokenEndpoint: "https://auth.example.com/token", ProtectedTokenEndpoint: true,
ClientID: "client", ClientSecret: "secret",
}
if err := store.Save("demo", want); err != nil {
t.Fatalf("Save() error = %v", err)
}
got, ok, err := store.Load("demo")
if err != nil || !ok {
t.Fatalf("Load() = %v, %v", ok, err)
}
if got.TokenEndpoint != want.TokenEndpoint || !got.ProtectedTokenEndpoint || got.ClientID != want.ClientID || got.ClientSecret != want.ClientSecret {
t.Fatalf("round trip lost discovery fields: %#v", got)
}
statuses, err := store.Status()
if err != nil {
t.Fatalf("Status() error = %v", err)
}
encoded, _ := json.Marshal(statuses)
for _, secret := range []string{"secret", "https://auth.example.com/token"} {
if strings.Contains(string(encoded), secret) {
t.Fatalf("status output leaked %q: %s", secret, encoded)
}
}
}
-14
View File
@@ -29,15 +29,6 @@ type StoredToken struct {
TokenType string `json:"token_type,omitempty"`
Scopes []string `json:"scopes,omitempty"`
ExpiresAt time.Time `json:"expires_at,omitempty"`
// TokenEndpoint, ClientID and ClientSecret hold what Login discovered or
// dynamically registered, so Refresh works for the standard `auth: oauth`
// setup whose config carries none of them. ClientSecret is sensitive.
// ProtectedTokenEndpoint is set when the endpoint came from server-advertised
// discovery and must pass the public-network policy again before use.
TokenEndpoint string `json:"token_endpoint,omitempty"`
ProtectedTokenEndpoint bool `json:"protected_token_endpoint,omitempty"`
ClientID string `json:"client_id,omitempty"`
ClientSecret string `json:"client_secret,omitempty"`
}
// TokenStatus is a redaction-safe summary of a stored token. It deliberately
@@ -402,11 +393,6 @@ func storedToOAuth(s StoredToken) oauth.Token {
TokenType: s.TokenType,
Scopes: s.Scopes,
ExpiresAt: s.ExpiresAt,
TokenEndpoint: s.TokenEndpoint,
ProtectedTokenEndpoint: s.ProtectedTokenEndpoint,
ClientID: s.ClientID,
ClientSecret: s.ClientSecret,
}
}
-85
View File
@@ -6,8 +6,6 @@ import (
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"net"
"net/http"
"net/http/httptest"
"net/url"
@@ -670,86 +668,3 @@ func TestFullFlowUsesProtectedResourceDiscovery(t *testing.T) {
t.Fatalf("flow hits = probe:%d resource:%d authorization:%d token:%d", resourceProbeHits.Load(), resourceMetadataHits.Load(), authorizationMetadataHits.Load(), tokenHits.Load())
}
}
// A CLIENT STALLED MID-HEADER MUST NOT OUTLIVE LOGIN (#1028).
//
// Login built a bare http.Server and ended it with Shutdown under a one-second
// deadline, and Shutdown only closes idle connections. A client that has sent
// part of a request header is not idle, so it survived the deadline and kept
// its goroutine, and the Serve goroutine, alive after Login returned.
//
// No sleep is needed to get the stalled connection accepted before Login ends:
// it is dialled before the real callback's connection, Serve accepts in order,
// and Login cannot finish until the real callback has been served.
func TestLoginDoesNotLeaveAStalledCallbackClientOpen(t *testing.T) {
tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_ = json.NewEncoder(w).Encode(map[string]any{
"access_token": "access-final", "token_type": "Bearer", "expires_in": 3600,
})
}))
defer tokenServer.Close()
var stalled net.Conn
opener := func(authURL string) error {
parsed, err := url.Parse(authURL)
if err != nil {
return err
}
redirect, err := url.Parse(parsed.Query().Get("redirect_uri"))
if err != nil {
return err
}
stalled, err = net.Dial("tcp", redirect.Host)
if err != nil {
return err
}
if _, err := fmt.Fprintf(stalled, "GET /callback HTTP/1.1\r\nHost: %s\r\nX-Stall:", redirect.Host); err != nil {
return err
}
cbURL := redirect.String() + "?code=auth-code&state=" + url.QueryEscape(parsed.Query().Get("state"))
go func() {
if resp, err := http.Get(cbURL); err == nil {
resp.Body.Close()
}
}()
return nil
}
_, err := Login(context.Background(), LoginOptions{
ServerName: "demo",
ServerURL: "https://issuer.invalid",
Config: OAuthConfig{
ClientID: "client-123",
AuthorizationEndpoint: "https://issuer.invalid/authorize",
TokenEndpoint: tokenServer.URL,
},
HTTPClient: http.DefaultClient,
OpenBrowser: opener,
Timeout: 5 * time.Second,
Now: time.Now,
})
if err != nil {
t.Fatalf("Login() error = %v", err)
}
if stalled == nil {
t.Fatal("SETUP INVALID: the stalled client never connected, so this proves nothing")
}
defer stalled.Close()
requireCallbackClientClosed(t, stalled)
}
// requireCallbackClientClosed fails unless the server end of conn is already
// closed: a read has to fail with something other than a timeout.
func requireCallbackClientClosed(t *testing.T, conn net.Conn) {
t.Helper()
_ = conn.SetReadDeadline(time.Now().Add(2 * time.Second))
var one [1]byte
_, err := conn.Read(one[:])
var netErr net.Error
if errors.As(err, &netErr) && netErr.Timeout() {
t.Fatal("the stalled callback client was still connected after Login returned")
}
if err == nil {
t.Fatal("the stalled callback client received data instead of being closed")
}
}
+15 -59
View File
@@ -331,56 +331,23 @@ func (tool registryTool) Run(ctx context.Context, args map[string]any) tools.Res
status = tools.StatusError
}
output := TextContent(result.Content)
images, disp, limitReason := forwardImages(result.Content)
// Image blocks with valid data ride Result.Images, the same channel capture
// tools already use. Everything else non-text is still named rather than
// silently dropped, because Zero still has nowhere to put audio, embedded
// resources, or a block whose payload could not be decoded (#823).
// Say what was thrown away. Without this an image-only result reads as
// "(empty MCP tool result)", the model concludes the call produced nothing
// and retries, and the user never learns an image came back (#823). The note
// is appended only when something was actually dropped, so a text-only
// result is byte-for-byte what it was before.
//
// Conversion happens once: the drop note is built from the same pass's
// per-item disposition, so a valid image is not decoded again just to
// decide whether it was forwarded.
//
// Image-only success has no text block, so a one-line placeholder keeps
// the tool_result self-describing. ModelOutput and finalizeToolOutcome
// copy Output through verbatim; an empty string would hand the model an
// empty body next to the image.
//
// Notes are appended only when something was actually dropped, so a
// text-only result is byte-for-byte what it was before, and a successfully
// forwarded image is not described as unforwardable.
//
// Unforwardable blocks (audio, resource, failed decode) say retrying
// cannot RECOVER the payload: Zero still has nowhere to put them.
// Budget-skipped images are different — Zero can forward them, and a
// retry with fewer images can recover the payload — so they get their
// own sentence.
if output == "" && len(images) > 0 {
output = "[image returned by tool]"
}
if exceeded := droppedContentNote(result.Content, disp, dispBudgetExceeded); exceeded != "" {
output = appendServerNote(output, exceeded,
", which exceeded this result's remaining image budget. Retrying with fewer images can recover this payload.")
}
if uninspected := droppedContentNote(result.Content, disp, dispUninspected); uninspected != "" {
verb := ", which were not inspected"
if dispCount(disp, dispUninspected) == 1 {
verb = ", which was not inspected"
// It says retrying cannot RECOVER the payload rather than that a retry
// returns the same thing. Each retry is a fresh call, so the server may well
// answer differently; what cannot change is that Zero still has nowhere to
// put a non-text block. Claiming the response would be identical would be a
// promise this code is in no position to make.
if dropped := DroppedContentSummary(result.Content); dropped != "" {
note := "[zero] this server also returned " + dropped + ", which Zero cannot forward yet. Retrying cannot recover this payload."
if output == "" {
note = "[zero] this server returned " + dropped + ", which Zero cannot forward yet. Retrying cannot recover this payload."
}
var cause string
switch limitReason {
case imageLimitForwardCount:
cause = " because the maximum forwarded image count was reached."
case imageLimitInspectionCount:
cause = " because the maximum image inspection limit was reached."
default:
cause = " because the aggregate image budget was reached."
}
output = appendServerNote(output, uninspected, verb+cause)
}
if dropped := droppedContentNote(result.Content, disp, dispDropped); dropped != "" {
output = appendServerNote(output, dropped,
", which Zero cannot forward yet. Retrying cannot recover this payload.")
output = strings.TrimSpace(output + "\n\n" + note)
}
if output == "" {
output = "(empty MCP tool result)"
@@ -388,7 +355,6 @@ func (tool registryTool) Run(ctx context.Context, args map[string]any) tools.Res
return tools.Result{
Status: status,
Output: output,
Images: images,
Meta: tool.meta(),
}
}
@@ -433,13 +399,3 @@ func isPersistentlyApproved(store *PermissionStore, server Server, toolName stri
})
return err == nil && approved
}
// appendServerNote joins a [zero] note to output, choosing "returned" or
// "also returned" based on whether anything precedes it.
func appendServerNote(output, subject, tail string) string {
verb := "this server also returned "
if output == "" {
verb = "this server returned "
}
return strings.TrimSpace(output + "\n\n[zero] " + verb + subject + tail)
}
-18
View File
@@ -197,24 +197,6 @@ func cachedModelsDevProviders() map[string]map[string]modelsDevModel {
return modelsDevCached
}
// startModelsDevRefresh starts a background RefreshModelsDevCache. Tests
// replace it to see whether a refresh was started.
var startModelsDevRefresh = func() {
go func() { _ = RefreshModelsDevCache(context.Background()) }()
}
// StartModelsDevRefresh refreshes the models.dev cache in the background when
// the fetch is on. ZERO_DISABLE_MODELS_FETCH is read here, before anything is
// started: a refresh goroutine reads the environment when it runs, so one
// started by a caller that had turned the fetch off (a test, through t.Setenv)
// could run after that setting was restored and fetch into the real cache.
func StartModelsDevRefresh() {
if strings.TrimSpace(os.Getenv("ZERO_DISABLE_MODELS_FETCH")) != "" {
return
}
startModelsDevRefresh()
}
// RefreshModelsDevCache fetches models.dev/api.json into the on-disk cache
// when the cache is missing or older than modelsDevRefreshAfter. It is safe to
// call fire-and-forget from startup (use a goroutine); it never affects the
-22
View File
@@ -256,25 +256,3 @@ func TestDefaultModelEntriesAppliesFreshCache(t *testing.T) {
}
t.Fatal("claude-sonnet-4.5 not found")
}
// With the fetch off, no refresh is started at all, so none is left running to
// read the environment after a test has restored it and fetch into the real
// cache.
func TestStartModelsDevRefreshStartsNothingWhenTheFetchIsOff(t *testing.T) {
started := 0
previous := startModelsDevRefresh
startModelsDevRefresh = func() { started++ }
t.Cleanup(func() { startModelsDevRefresh = previous })
t.Setenv("ZERO_DISABLE_MODELS_FETCH", "1")
StartModelsDevRefresh()
if started != 0 {
t.Fatalf("started %d refreshes with the fetch off, want none", started)
}
t.Setenv("ZERO_DISABLE_MODELS_FETCH", "")
StartModelsDevRefresh()
if started != 1 {
t.Fatalf("started %d refreshes with the fetch on, want 1", started)
}
}
+1 -10
View File
@@ -1,13 +1,6 @@
package modelregistry
import (
"regexp"
"strings"
)
// Match explicit multimodal local families, including common Ollama spellings.
// Gemma 3 requires a vision-capable size: 270M and 1B are text-only.
var localVisionFamily = regexp.MustCompile(`^(gemma-?3[-:](4b|12b|27b)|llama-?4|mistral-small-?3\.1|phi-?4-multimodal)($|[-:_.])`)
import "strings"
// SupportsVision reports whether the model identified by modelID accepts image
// input. A model in the curated catalog is authoritative (its declared
@@ -43,8 +36,6 @@ func VisionCapableByName(modelID string) bool {
id = id[slash+1:] // drop a "provider/" prefix
}
switch {
case localVisionFamily.MatchString(id):
return true
case strings.Contains(id, "gemini"):
return true // every Gemini model is multimodal
case strings.Contains(id, "claude-3"), strings.Contains(id, "claude-4"),
@@ -10,15 +10,10 @@ func TestVisionCapableByName(t *testing.T) {
"claude-sonnet-4.5", "claude-3-haiku",
"MiniMax-M3",
"llava:13b", "qwen2.5-vl-7b", "llama3.2-vision", "pixtral-12b", "moondream",
"google/gemma-3-27b-it", "gemma3:27b", "gemma3:4b", "gemma-3-12b-it",
"llama-4-scout", "llama4:17b", "mistral-small-3.1", "mistral-small3.1:24b",
"phi-4-multimodal", "phi4-multimodal-instruct",
}
textOnly := []string{
"gpt-oss:120b", "kimi-for-coding", "deepseek-coder", "qwen2.5-coder",
"codestral", "llama3.1-8b", "grok-text-only", "mistral-large",
"gemma3:1b", "gemma-3-1b-it", "gemma3:270m", "gemma-3-270m-it",
"gemma2:27b", "mistral-small3:24b", "phi4:14b",
// Negated "vision" names must NOT match the bare-"vision" fallback.
"my-custom-vision-less-model", "no-vision-model", "grok-vision-less",
}
+6 -19
View File
@@ -57,7 +57,7 @@ type Sink interface {
// attached Sinks. Safe for concurrent use.
type Notifier struct {
w io.Writer
cfg Config // swap at runtime via Configure; reads outside the lock are safe
cfg Config // immutable after New; reads outside the lock are safe
mu sync.Mutex
focused bool
@@ -71,16 +71,6 @@ func New(w io.Writer, cfg Config) *Notifier {
return &Notifier{w: w, cfg: cfg}
}
// Configure swaps the mode/focus policy at runtime. Sinks, focus state, and the
// writer are preserved, so an in-session preference change (e.g. the TUI's
// /notify command) applies from the next Notify call. Safe to call concurrently
// with Notify.
func (n *Notifier) Configure(cfg Config) {
n.mu.Lock()
n.cfg = cfg
n.mu.Unlock()
}
// AddSink registers an additional destination that receives every eligible
// event (subject to the same mode/focus policy as the terminal). Sinks fire
// even when the Notifier has no terminal writer, so a headless CI run can still
@@ -112,21 +102,18 @@ func (n *Notifier) SetFocused(focused bool) {
// Sinks are invoked outside the lock so a slow/blocking sink cannot stall a
// concurrent Notify or SetFocused.
func (n *Notifier) Notify(event Event, message string) {
n.mu.Lock()
// cfg is mutable at runtime (Configure), so every read happens under the
// lock; shouldEmit/sequence work on the local copy.
cfg := n.cfg
if cfg.Mode == ModeOff || cfg.Mode == "" {
n.mu.Unlock()
if n.cfg.Mode == ModeOff || n.cfg.Mode == "" {
return
}
eligible := shouldEmit(cfg, event, n.focused)
n.mu.Lock()
eligible := shouldEmit(n.cfg, event, n.focused)
var sinks []Sink
if eligible && len(n.sinks) > 0 {
sinks = append(sinks, n.sinks...)
}
if eligible && n.w != nil {
if seq := sequence(cfg.Mode, message); seq != "" {
if seq := sequence(n.cfg.Mode, message); seq != "" {
_, _ = io.WriteString(n.w, seq)
}
}
-53
View File
@@ -50,44 +50,6 @@ func TestSequence(t *testing.T) {
}
}
// Configure swaps the policy at runtime: after turning a silent notifier on,
// the next Notify emits; after turning a noisy one off, it stays silent. Sinks
// survive the swap (the TUI relies on this when /notify reconfigures mid-run).
func TestConfigureAppliesImmediatelyAndKeepsSinks(t *testing.T) {
var buf bytes.Buffer
n := New(&buf, Config{Mode: ModeOff})
n.SetFocused(true)
n.Notify(Completion, "x")
if buf.Len() != 0 {
t.Fatalf("off should be silent, got %q", buf.String())
}
sink := &recordingSink{}
n.AddSink(sink)
n.Configure(Config{Mode: ModeBell, FocusMode: FocusAlways})
n.Notify(Completion, "x")
if buf.String() != "\x07" {
t.Fatalf("after Configure(bell) should bell, got %q", buf.String())
}
sink.mu.Lock()
got := len(sink.events)
sink.mu.Unlock()
if got != 1 {
t.Fatalf("sink should still receive events after Configure, got %d", got)
}
n.Configure(Config{Mode: ModeOff})
buf.Reset()
n.Notify(Completion, "x")
sink.mu.Lock()
got = len(sink.events)
sink.mu.Unlock()
if buf.Len() != 0 || got != 1 {
t.Fatalf("after Configure(off) should be fully silent, buf=%q events=%d", buf.String(), got)
}
}
func TestSanitizeMessage(t *testing.T) {
if got := sanitizeMessage("ok\x1b]0;evil\x07more\nx"); got != "ok]0;evilmorex" {
t.Fatalf("sanitize=%q", got)
@@ -147,21 +109,6 @@ func TestNotifyRaceSafe(t *testing.T) {
wg.Wait()
}
// Configure mutates cfg under the lock while Notify reads it; this pair must
// be race-clean (run under -race). Regression for the unsynchronized cfg read
// Notify used to perform before acquiring the lock.
func TestConfigureConcurrentWithNotify(t *testing.T) {
n := New(&bytes.Buffer{}, Config{Mode: ModeBell, FocusMode: FocusAlways})
var wg sync.WaitGroup
for i := 0; i < 50; i++ {
wg.Add(3)
go func() { defer wg.Done(); n.Configure(Config{Mode: ModeBoth, FocusMode: FocusAlways}) }()
go func() { defer wg.Done(); n.Configure(Config{Mode: ModeOff}) }()
go func() { defer wg.Done(); n.Notify(AwaitingInput, "x") }()
}
wg.Wait()
}
func TestDefaultMessage(t *testing.T) {
if DefaultMessage(Completion) != "Zero: ready" {
t.Fatal("completion message")
-70
View File
@@ -1,70 +0,0 @@
package oauth
import (
"context"
"net"
"net/http"
"time"
)
// callbackReadHeaderTimeout bounds how long an accepted connection may take to
// send its request header. A single-use loopback callback server has no reason
// to wait longer, and without a bound a client that stalls mid-header holds the
// connection, and the server goroutine, open indefinitely. A var so the
// package tests can shorten it.
var callbackReadHeaderTimeout = 5 * time.Second
// callbackShutdownBudget is how long Close waits for in-flight requests before
// it force-closes whatever is still open.
const callbackShutdownBudget = time.Second
// CallbackServer is the lifecycle every loopback OAuth callback server needs,
// in one place.
//
// It exists because the same server was written out three times (the shared
// LoopbackListener, MCP's Login and OpenRouter's login) and a fix to one of
// them did not reach the others. http.Server.Shutdown only closes idle
// connections, so a client that has sent part of a request header is active,
// survives the shutdown deadline, and keeps its goroutine and the Serve
// goroutine alive past the end of the flow (#1028). #1051 closed that for
// LoopbackListener; the other two kept the original construction.
type CallbackServer struct {
server *http.Server
done chan struct{}
}
// StartCallbackServer serves handler on listener until Close.
func StartCallbackServer(listener net.Listener, handler http.Handler) *CallbackServer {
return startCallbackServer(listener, handler, nil)
}
// startCallbackServer is StartCallbackServer with a ConnState hook, which the
// package's own tests use to see a connection reach a given state.
func startCallbackServer(listener net.Listener, handler http.Handler, connState func(net.Conn, http.ConnState)) *CallbackServer {
s := &CallbackServer{
server: &http.Server{
Handler: handler,
ReadHeaderTimeout: callbackReadHeaderTimeout,
ConnState: connState,
},
done: make(chan struct{}),
}
go func() {
_ = s.server.Serve(listener)
close(s.done)
}()
return s
}
// Close shuts the server down within callbackShutdownBudget. A connection still
// open when the budget runs out (a stalled request header, say) is force-closed
// rather than leaked, and Close returns only after the Serve goroutine has
// finished, so nothing this server started outlives the call. Idempotent.
func (s *CallbackServer) Close() {
ctx, cancel := context.WithTimeout(context.Background(), callbackShutdownBudget)
defer cancel()
if err := s.server.Shutdown(ctx); err != nil {
_ = s.server.Close()
}
<-s.done
}
-100
View File
@@ -1,100 +0,0 @@
package oauth
import (
"errors"
"fmt"
"net"
"net/http"
"testing"
"time"
)
// startStalledCallbackClient starts a CallbackServer and connects a client that
// sends part of a request header and then stops, which is the shape #1028 is
// about. It returns once the server has ACCEPTED the connection, using the
// ConnState hook rather than a sleep. StateNew, not StateActive: Go reports
// StateActive only after a request read returns, and a stalled header read
// returns only when the header timeout fires, so waiting for StateActive would
// make the setup depend on the very bound the test is about.
func startStalledCallbackClient(t *testing.T) (*CallbackServer, net.Conn) {
t.Helper()
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
accepted := make(chan struct{}, 1)
server := startCallbackServer(listener, http.NotFoundHandler(), func(_ net.Conn, state http.ConnState) {
if state == http.StateNew {
select {
case accepted <- struct{}{}:
default:
}
}
})
conn, err := net.Dial("tcp", listener.Addr().String())
if err != nil {
server.Close()
t.Fatal(err)
}
if _, err := fmt.Fprintf(conn, "GET /callback HTTP/1.1\r\nHost: %s\r\nX-Stall:", listener.Addr()); err != nil {
t.Fatal(err)
}
select {
case <-accepted:
case <-time.After(5 * time.Second):
t.Fatal("SETUP INVALID: the server never accepted the stalled client")
}
return server, conn
}
// connClosedWithin reports whether the server end of conn closes before the
// deadline: a read must fail with something other than a timeout.
func connClosedWithin(conn net.Conn, deadline time.Duration) bool {
_ = conn.SetReadDeadline(time.Now().Add(deadline))
var one [1]byte
_, err := conn.Read(one[:])
var netErr net.Error
if errors.As(err, &netErr) && netErr.Timeout() {
return false
}
return err != nil
}
// THE HEADER READ IS BOUNDED EVEN BEFORE CLOSE. The force-close in Close ends a
// stalled client when the flow ends; the header timeout ends it while the flow
// is still waiting for its callback, which can be minutes. Close is not called
// until after the assertion, so this can only pass on ReadHeaderTimeout.
func TestCallbackServerBoundsAStalledHeaderWithoutClose(t *testing.T) {
previous := callbackReadHeaderTimeout
callbackReadHeaderTimeout = 200 * time.Millisecond
defer func() { callbackReadHeaderTimeout = previous }()
server, conn := startStalledCallbackClient(t)
defer server.Close()
defer conn.Close()
if !connClosedWithin(conn, 3*time.Second) {
t.Fatal("a client stalled mid-header was still connected with Close never called; the header read is unbounded")
}
}
// CLOSE RETURNS ONLY AFTER SERVE HAS. Its doc comment promises that nothing the
// server started outlives the call. With no connections open, Shutdown closes
// the listener and returns at once, so this is the case where Close could beat
// the Serve goroutine back; the goroutine closing done is the last thing it
// does, so done has to be closed by the time Close returns.
func TestCallbackServerCloseWaitsForServe(t *testing.T) {
for attempt := 0; attempt < 50; attempt++ {
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
server := StartCallbackServer(listener, http.NotFoundHandler())
server.Close()
select {
case <-server.done:
default:
t.Fatalf("attempt %d: Close returned while the Serve goroutine was still running", attempt)
}
}
}
+23 -5
View File
@@ -9,6 +9,7 @@ import (
"net/http"
"net/url"
"strings"
"time"
)
// LoopbackListener is a single-use loopback HTTP server that captures an OAuth
@@ -18,7 +19,8 @@ type LoopbackListener struct {
listener net.Listener
state string
result chan callbackResult
server *CallbackServer
server *http.Server
done chan struct{}
}
type callbackResult struct {
@@ -56,8 +58,20 @@ func newLoopbackListener(state string, port int, connState func(net.Conn, http.C
listener: ln,
state: state,
result: make(chan callbackResult, 1),
done: make(chan struct{}),
}
l.server = startCallbackServer(ln, http.HandlerFunc(l.handle), connState)
l.server = &http.Server{
Handler: http.HandlerFunc(l.handle),
// This is a single-use, loopback-only callback server: a client that
// never finishes sending its request header must not be able to hold
// the accepted connection (and Close, below) open indefinitely.
ReadHeaderTimeout: 5 * time.Second,
ConnState: connState,
}
go func() {
_ = l.server.Serve(ln)
close(l.done)
}()
return l, nil
}
@@ -108,10 +122,14 @@ func (l *LoopbackListener) Wait(ctx context.Context) (string, error) {
// Close shuts the listener down (bounded), idempotent. If a client leaves a
// connection open past the shutdown deadline (e.g. a stalled request header),
// Close forcibly closes it rather than leaking it past the end of the flow,
// and waits for the Serve goroutine to finish before returning. See
// CallbackServer.Close.
// and waits for the Serve goroutine to finish before returning.
func (l *LoopbackListener) Close() {
l.server.Close()
shutdownCtx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := l.server.Shutdown(shutdownCtx); err != nil {
_ = l.server.Close()
}
<-l.done
}
// parseCallback validates the redirect query and returns the authorization code,
-10
View File
@@ -45,16 +45,6 @@ type Token struct {
// `chatgpt_account_id`) that the request path needs as headers. Treated as
// sensitive like the access token: never logged, persisted 0600.
IDToken string `json:"id_token,omitempty"`
// TokenEndpoint, ClientID and ClientSecret record what a login learned that
// config does not carry (a discovered token endpoint, a dynamically
// registered client) so a later refresh can reuse it. ClientSecret is
// sensitive like the tokens. ProtectedTokenEndpoint marks an endpoint
// advertised by the server, which must be re-checked against the public
// network policy before a refresh posts credentials to it.
TokenEndpoint string `json:"token_endpoint,omitempty"`
ProtectedTokenEndpoint bool `json:"protected_token_endpoint,omitempty"`
ClientID string `json:"client_id,omitempty"`
ClientSecret string `json:"client_secret,omitempty"`
}
// Expired reports whether the token has an expiry that is at or before now.
@@ -375,14 +375,6 @@ type modelsResponseItem struct {
AdditionalSpeedTiers []string `json:"additional_speed_tiers"`
DefaultServiceTier string `json:"default_service_tier"`
SupportedParameters []string `json:"supported_parameters"`
Modalities struct {
Input []string `json:"input"`
Output []string `json:"output"`
} `json:"modalities"`
Architecture struct {
InputModalities []string `json:"input_modalities"`
OutputModalities []string `json:"output_modalities"`
} `json:"architecture"`
}
type modelsResponse struct {
@@ -436,14 +428,6 @@ func parseModelsResponse(body []byte) ([]Model, error) {
discoveryContainsFold(item.SupportedParameters, "include_reasoning")
efforts := reasoningEfforts(item)
serviceTiers := modelServiceTiers(item)
inputModalities := cleanDiscoveryStrings(item.Modalities.Input)
if len(inputModalities) == 0 {
inputModalities = cleanDiscoveryStrings(item.Architecture.InputModalities)
}
outputModalities := cleanDiscoveryStrings(item.Modalities.Output)
if len(outputModalities) == 0 {
outputModalities = cleanDiscoveryStrings(item.Architecture.OutputModalities)
}
models = append(models, Model{
ID: id,
Description: description,
@@ -454,8 +438,6 @@ func parseModelsResponse(body []byte) ([]Model, error) {
DefaultReasoningEffort: strings.TrimSpace(item.DefaultReasoning),
ServiceTiers: serviceTiers,
DefaultServiceTier: normalizeServiceTier(item.DefaultServiceTier),
InputModalities: inputModalities,
OutputModalities: outputModalities,
Source: "live",
})
}
@@ -592,12 +574,6 @@ func mergeLiveModels(provider providercatalog.Descriptor, liveModels []Model, ca
if live.DefaultServiceTier != "" {
catalog.DefaultServiceTier = live.DefaultServiceTier
}
if len(catalog.InputModalities) == 0 && len(live.InputModalities) > 0 {
catalog.InputModalities = append([]string{}, live.InputModalities...)
}
if len(catalog.OutputModalities) == 0 && len(live.OutputModalities) > 0 {
catalog.OutputModalities = append([]string{}, live.OutputModalities...)
}
catalog.Source = firstDiscoverySource(catalog.Source, "live")
result = append(result, catalog)
continue
@@ -624,20 +600,6 @@ func mergeLiveModels(provider providercatalog.Descriptor, liveModels []Model, ca
return result
}
func cleanDiscoveryStrings(values []string) []string {
result := make([]string, 0, len(values))
seen := map[string]bool{}
for _, value := range values {
value = strings.TrimSpace(value)
if value == "" || seen[value] {
continue
}
seen[value] = true
result = append(result, value)
}
return result
}
func discoveryContainsFold(values []string, want string) bool {
for _, value := range values {
if strings.EqualFold(strings.TrimSpace(value), want) {
@@ -109,49 +109,6 @@ func TestParseModelsResponseSupportsChatGPTCatalog(t *testing.T) {
}
}
func TestParseModelsResponseCapturesModalities(t *testing.T) {
models, err := parseModelsResponse([]byte(`{
"data": [
{
"id": "openrouter/custom-multimodal",
"architecture": {
"input_modalities": ["text", "image"],
"output_modalities": ["text"]
}
},
{
"id": "opengateway/flat-multimodal",
"modalities": {
"input": ["text", "image"],
"output": ["text", "image"]
}
},
{
"id": "text-only",
"architecture": {
"input_modalities": ["text"]
}
}
]
}`))
if err != nil {
t.Fatalf("parseModelsResponse: %v", err)
}
byID := map[string]Model{}
for _, m := range models {
byID[m.ID] = m
}
if got := strings.Join(byID["openrouter/custom-multimodal"].InputModalities, ","); got != "text,image" {
t.Fatalf("openrouter input modalities = %q, want text,image", got)
}
if got := strings.Join(byID["opengateway/flat-multimodal"].InputModalities, ","); got != "text,image" {
t.Fatalf("opengateway input modalities = %q, want text,image", got)
}
if got := strings.Join(byID["text-only"].InputModalities, ","); got != "text" {
t.Fatalf("text-only input modalities = %q, want text", got)
}
}
func TestParseModelsResponseNormalizesLegacyFastTier(t *testing.T) {
models, err := parseModelsResponse([]byte(`{"data":[{"id":"gpt-test","additional_speed_tiers":["fast","priority"]}]}`))
if err != nil {
+8 -7
View File
@@ -79,11 +79,7 @@ func OpenRouterLogin(ctx context.Context, opts OpenRouterOptions) (string, error
codeCh := make(chan string, 1)
errCh := make(chan error, 1)
// The shared callback server bounds header reads, force-closes a client
// still connected when shutdown runs out of time, and waits for Serve to
// return. The bare http.Server this used to build did none of that, so a
// client stalled mid-header outlived the login (#1028).
server := oauth.StartCallbackServer(listener, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
server := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/callback" {
http.NotFound(w, r)
return
@@ -102,8 +98,13 @@ func OpenRouterLogin(ctx context.Context, opts OpenRouterOptions) (string, error
case errCh <- errors.New("provideroauth: callback missing authorization code"):
default:
}
}))
defer server.Close()
})}
go func() { _ = server.Serve(listener) }()
defer func() {
shutdownCtx, cancelShutdown := context.WithTimeout(context.Background(), time.Second)
defer cancelShutdown()
_ = server.Shutdown(shutdownCtx)
}()
authURL := base + openRouterAuthPath + "?" + url.Values{
"callback_url": {callbackURL},
-62
View File
@@ -4,15 +4,12 @@ import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"net/url"
"strings"
"testing"
"time"
)
// simulateBrowser returns an OpenBrowser func that authorizes by GETting the
@@ -132,62 +129,3 @@ func TestOpenRouterLoginBrowserErrorPropagates(t *testing.T) {
t.Fatal("a browser-open failure should propagate")
}
}
// A CLIENT STALLED MID-HEADER MUST NOT OUTLIVE THE LOGIN (#1028). Same defect
// as MCP's Login: a bare http.Server ended by Shutdown alone, which leaves a
// connection that has sent part of a request header open after the flow ends.
// The stalled client is dialled before the real callback, so Serve has
// accepted it by the time the login can complete; no sleep is involved.
func TestOpenRouterLoginDoesNotLeaveAStalledCallbackClientOpen(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]string{"key": "sk-or-mock-123"})
}))
defer srv.Close()
var stalled net.Conn
open := func(authURL string) error {
u, err := url.Parse(authURL)
if err != nil {
return err
}
cb, err := url.Parse(u.Query().Get("callback_url"))
if err != nil {
return err
}
stalled, err = net.Dial("tcp", cb.Host)
if err != nil {
return err
}
if _, err := fmt.Fprintf(stalled, "GET /callback HTTP/1.1\r\nHost: %s\r\nX-Stall:", cb.Host); err != nil {
return err
}
resp, err := http.Get(cb.String() + "?code=TESTCODE") //nolint:noctx // test loopback
if err != nil {
return err
}
_ = resp.Body.Close()
return nil
}
if _, err := OpenRouterLogin(context.Background(), OpenRouterOptions{
BaseURL: srv.URL, HTTPClient: srv.Client(), OpenBrowser: open,
}); err != nil {
t.Fatalf("OpenRouterLogin: %v", err)
}
if stalled == nil {
t.Fatal("SETUP INVALID: the stalled client never connected, so this proves nothing")
}
defer stalled.Close()
_ = stalled.SetReadDeadline(time.Now().Add(2 * time.Second))
var one [1]byte
_, err := stalled.Read(one[:])
var netErr net.Error
if errors.As(err, &netErr) && netErr.Timeout() {
t.Fatal("the stalled callback client was still connected after OpenRouterLogin returned")
}
if err == nil {
t.Fatal("the stalled callback client received data instead of being closed")
}
}
-547
View File
@@ -1,547 +0,0 @@
package redaction
import (
"fmt"
"regexp"
"regexp/syntax"
"sort"
"strings"
"unicode"
"unicode/utf8"
)
// A CREDENTIAL IS STILL A CREDENTIAL WITH AN INVISIBLE BYTE IN THE MIDDLE.
//
// The shape matchers (openaiKeyPattern, textSecretPatterns) all describe a
// CONTIGUOUS run of body characters, and none of their character classes admits
// a control or format character. So a single NUL, ESC or zero-width space
// dropped into the body of a key ends the match, and RedactString emits the two
// fragments untouched. Whoever reads that output next, a terminal that eats the
// escape, a log viewer, a JSON consumer that strips control bytes, or a person
// copying the text, sees the original key back, because nothing about those
// characters was ever displayed. Reported in #969 for NUL and ESC; every other
// Cc and Cf character behaves the same way, which is why this is written against
// the classes rather than against the two bytes named there.
//
// HOW: each shape is rewritten so that a run of separators may sit between any
// two characters it consumes, and matched against the ORIGINAL text. Nothing is
// removed before matching, so there are no offsets to map back, and the leading
// \b of every shape is judged by the regexp engine against the real neighbour:
// a key written after a NUL starts at a boundary, exactly as the unsplit key
// does, and a key glued to the end of a word does not.
//
// An earlier version removed the separators first and matched the compacted
// copy. That lost the difference between a separator inside a credential and one
// next to it, and the fixes for that (check the boundary against the source
// afterwards, map offsets back) left a hole: in "xAKIA<NUL>AKIAIOSF<SOH>ODNN7..."
// the leftmost compacted match began at the decoy AKIA, was then rejected for
// the "x" in front of it, and had already consumed the start of the real key, so
// the real key was never considered. Resuming the scan one byte after each
// rejection would close that, but it is quadratic for the greedy shapes on text
// like "ask-ask-ask-..." with a single NUL in it. Keeping the boundary inside
// the pattern lets RE2 do it in one linear pass.
//
// This pass runs AFTER the contiguous matchers, over their output, and never
// across a replacement. A credential that is whole is claimed by its own strict
// match, so two whole credentials with a separator between them stay two
// replacements with the separator still between them. Two split credentials do
// too: see spansFor.
//
// A match never starts or ends on a separator, so the separators on either side
// of a credential stay where they were: "<NUL>AKIA..." becomes
// "<NUL>[REDACTED]", as it does for the unsplit key.
//
// ONE KNOWN LIMIT, from running the contiguous matchers first. When the part of
// a split key BEFORE its first separator is already a complete credential on
// its own, the contiguous matcher claims that part and stops at the separator,
// and what follows is judged by itself. That is the rest of the key, which is no
// credential alone, so it stays visible: a fragment of the key, never the whole
// of it, because the part in front is replaced. For the same reason a credential
// glued straight onto that rest is left, where the unsplit key would have run on
// and swallowed it. Carrying the contiguous match across the separator instead
// would bring back what running first exists to prevent, since nothing tells the
// rest of a key from prose written after a whole one: that prose would be eaten,
// and a following JWT would lose its header.
//
// WHAT THIS DELIBERATELY DOES NOT TREAT AS A SEPARATOR: tab, newline and
// carriage return. Those three are real text structure rather than invisible
// filler, so a credential broken across a line boundary is still two visible
// fragments here and is not covered. Treating them as gaps would also let a
// match span a line break and collapse unrelated lines into one replacement. A
// credential split that way is visible to whoever reads it; the ones handled
// here are not. Invalid UTF-8 is out for the same reason: the engine reads an
// invalid byte as U+FFFD, which is drawn.
// splitSecretSeparator reports a character that can sit inside a credential
// without being seen. Cc covers the C0 and C1 controls (NUL, ESC, DEL, NEL, CSI)
// and Cf the format characters (zero-width space and joiner, word joiner, BOM,
// soft hyphen). Tab, newline and carriage return are excluded: see above.
func splitSecretSeparator(r rune) bool {
switch r {
case '\t', '\n', '\r':
return false
}
return unicode.Is(unicode.Cc, r) || unicode.Is(unicode.Cf, r)
}
// containsSplitSeparator is the cheap gate in front of the gap-tolerant pass, so
// text with nothing invisible in it, which is nearly all text, costs one scan and
// no allocation. ASCII is settled from the byte alone; only a byte that could
// begin a multi-byte rune needs decoding, and the C1 and Cf blocks all start
// with a lead byte at or above 0xC2.
func containsSplitSeparator(value string) bool {
for i := 0; i < len(value); i++ {
b := value[i]
if b < utf8.RuneSelf {
if b < 0x20 || b == 0x7f {
if b == '\t' || b == '\n' || b == '\r' {
continue
}
return true
}
continue
}
r, width := utf8.DecodeRuneInString(value[i:])
if splitSecretSeparator(r) {
return true
}
i += width - 1
}
return false
}
// stripSplitSeparators returns value with every splitSecretSeparator removed.
// Used only on a matched credential, to judge it by its visible characters.
func stripSplitSeparators(value string) string {
var builder strings.Builder
builder.Grow(len(value))
for offset := 0; offset < len(value); {
r, width := utf8.DecodeRuneInString(value[offset:])
if !splitSecretSeparator(r) {
builder.WriteString(value[offset : offset+width])
}
offset += width
}
return builder.String()
}
// splitSeparatorClass is splitSecretSeparator as a regexp character class,
// generated from the same Unicode tables and the same three exclusions so the
// class and the predicate cannot disagree.
var splitSeparatorClass = func() string {
var class strings.Builder
class.WriteByte('[')
emit := func(lo, hi rune) {
fmt.Fprintf(&class, `\x{%x}-\x{%x}`, lo, hi)
}
for _, table := range []*unicode.RangeTable{unicode.Cc, unicode.Cf} {
visit := func(lo, hi, stride uint32) {
for r := lo; r <= hi; r += stride {
start := rune(r)
if !splitSecretSeparator(start) {
continue
}
end := start
if stride == 1 {
for next := rune(r) + 1; uint32(next) <= hi && splitSecretSeparator(next); next++ {
end = next
}
r = uint32(end)
}
emit(start, end)
}
}
for _, r := range table.R16 {
visit(uint32(r.Lo), uint32(r.Hi), uint32(r.Stride))
}
for _, r := range table.R32 {
visit(r.Lo, r.Hi, r.Stride)
}
}
class.WriteByte(']')
return class.String()
}()
// splitOpenAIPattern and splitTextPatterns are the shape matchers with a run of
// separators allowed between every two characters they consume, derived from
// the originals rather than written out a second time so the two lists cannot
// drift apart.
var (
splitOpenAIPattern = gapTolerant(openaiKeyPattern)
splitTextPatterns = gapTolerantAll(textSecretPatterns)
)
// splitShape is one credential shape as the split pass uses it.
type splitShape struct {
// free finds the shape anywhere, leftmost first. anchored matches it only
// where the text begins, and whole only when it spans the entire text.
free, anchored, whole *regexp.Regexp
// prefix is the literal every match of the shape begins with.
prefix string
// accept is the test a match has to pass to be a credential, or nil when
// every match is one.
accept func(match string) bool
}
var splitShapes = func() []splitShape {
shapes := []splitShape{newSplitShape(openaiKeyPattern, splitOpenAIPattern, func(match string) bool {
return openAIShapeIsSecret(stripSplitSeparators(match))
})}
for i, pattern := range textSecretPatterns {
shapes = append(shapes, newSplitShape(pattern, splitTextPatterns[i], nil))
}
return shapes
}()
func newSplitShape(original, free *regexp.Regexp, accept func(string) bool) splitShape {
return splitShape{
free: free,
anchored: regexp.MustCompile(`\A(?:` + free.String() + `)`),
whole: regexp.MustCompile(`\A(?:` + free.String() + `)\z`),
prefix: leadingLiteral(original),
accept: accept,
}
}
// leadingLiteral is the literal every match of pattern begins with, read past a
// leading \b, or "" when the pattern does not open with one.
func leadingLiteral(pattern *regexp.Regexp) string {
tree, err := syntax.Parse(pattern.String(), syntax.Perl)
if err != nil {
panic(fmt.Sprintf("redaction: parse %q: %v", pattern.String(), err))
}
nodes := []*syntax.Regexp{tree}
if tree.Op == syntax.OpConcat {
nodes = tree.Sub
}
for _, node := range nodes {
if node.Op == syntax.OpWordBoundary {
continue
}
if node.Op == syntax.OpLiteral {
return string(node.Rune)
}
return ""
}
return ""
}
// maxRunOnCandidates bounds the anchored matches tried inside one match, so a
// match full of lookalike starts still costs a fixed number of them.
const maxRunOnCandidates = 8
// spansFor turns one match of the shape into the stretches to replace.
//
// A MATCH CAN RUN ON INTO THE NEXT CREDENTIAL OF ITS OWN SHAPE. Gaps are allowed
// inside a match, and a separator between two credentials looks exactly like one
// inside a credential, so a greedy body runs through it into whatever follows. A
// credential of another shape is still found, because every shape is matched on
// its own and the union is replaced. One of the SAME shape is not: matches of one
// pattern never overlap, so the second credential cannot start inside the first
// one's match. Most bodies swallow it whole, so nothing leaks, but a JWT stops at
// the next token's first dot, and two split JWTs side by side left the second
// one's payload and signature in the clear. And a match the filter rejects hides
// whatever it ran into, so kebab-case prose in front of a split key took the key
// down with it. Reported by @jatmn.
//
// So when a match runs through a run of separators into a position where the
// same shape matches again and reaches at least as far, that credential is taken
// on its own, and the first is cut at the separators if it is a complete match by
// itself. The separators between the two then stay outside both replacements, as
// they do between two whole credentials, and each piece is filtered on its own.
// When the first is not complete alone, it keeps the whole of its match and the
// union joins the two into one replacement.
//
// ONE STEP PER MATCH, NOT THE WHOLE CHAIN. The second credential's own match can
// run on into a third, and following that from here would walk the rest of a
// chain once for every match inside it, which is quadratic: a megabyte of split
// JWTs ran for minutes. The scan's own next match starts inside the second
// credential and takes the step after it, so the chain is still covered, and in
// a chain of three or more the markers after the first pair can join into one.
func (shape splitShape) spansFor(text string, start, end int) [][]int {
next, gap, nextEnd := shape.runOn(text, start, end)
if next < 0 {
return shape.keep(nil, text, start, end)
}
var spans [][]int
if shape.whole.MatchString(text[start:gap]) {
spans = shape.keep(spans, text, start, gap)
} else {
spans = shape.keep(spans, text, start, end)
}
return shape.keep(spans, text, next, nextEnd)
}
func (shape splitShape) keep(spans [][]int, text string, start, end int) [][]int {
if shape.accept != nil && !shape.accept(text[start:end]) {
return spans
}
return append(spans, []int{start, end})
}
// runOn looks inside the match [start, end) for where it ran through a run of
// separators into another credential of the same shape. It returns where that
// credential begins, where the separator run in front of it begins, and where the
// credential's own match ends, or next = -1 when there is none.
//
// Latest first, because a match reaches into the next credential only as far as
// its own pattern allows. Only where the text after the separators opens with the
// shape's literal prefix, which is where a match of the shape can begin; the
// separator in front supplies the word boundary, so matching from there judges it
// the same as the whole text would.
func (shape splitShape) runOn(text string, start, end int) (next, gap, nextEnd int) {
tried := 0
for at := end; at > start && tried < maxRunOnCandidates; {
r, size := utf8.DecodeLastRuneInString(text[start:at])
if !splitSecretSeparator(r) {
at -= size
continue
}
candidate := at
for at > start {
r, size = utf8.DecodeLastRuneInString(text[start:at])
if !splitSecretSeparator(r) {
break
}
at -= size
}
if !hasGappedPrefix(text[candidate:], shape.prefix) {
continue
}
tried++
if match := shape.anchored.FindStringIndex(text[candidate:]); match != nil && candidate+match[1] >= end {
return candidate, at, candidate + match[1]
}
}
return -1, 0, 0
}
// hasGappedPrefix reports whether text opens with prefix, allowing a run of
// separators between its characters the way the shapes themselves do. A plain
// prefix test missed a token split inside its first few characters, and the run
// that swallowed it then went unnoticed.
func hasGappedPrefix(text, prefix string) bool {
at := 0
for index, want := range prefix {
if index > 0 {
for at < len(text) {
r, size := utf8.DecodeRuneInString(text[at:])
if !splitSecretSeparator(r) {
break
}
at += size
}
}
if at >= len(text) {
return false
}
r, size := utf8.DecodeRuneInString(text[at:])
if r != want {
return false
}
at += size
}
return true
}
func gapTolerantAll(patterns []*regexp.Regexp) []*regexp.Regexp {
tolerant := make([]*regexp.Regexp, 0, len(patterns))
for _, pattern := range patterns {
tolerant = append(tolerant, gapTolerant(pattern))
}
return tolerant
}
// gapTolerant rewrites pattern so that a run of separators may appear before
// every character it consumes except the first. Zero-width assertions, the
// leading \b in particular, are left exactly where they were.
func gapTolerant(pattern *regexp.Regexp) *regexp.Regexp {
tree, err := syntax.Parse(pattern.String(), syntax.Perl)
if err != nil {
panic(fmt.Sprintf("redaction: parse %q: %v", pattern.String(), err))
}
separator, err := syntax.Parse(splitSeparatorClass, syntax.Perl)
if err != nil {
panic(fmt.Sprintf("redaction: parse separator class: %v", err))
}
consumed := false
return regexp.MustCompile(withGaps(tree, separator, &consumed).String())
}
// withGaps inserts the separator run before each consuming atom. consumed
// tracks whether anything has been consumed yet on this path, because the first
// character of a match must be a real one: a leading gap would let a match start
// on a separator and swallow it.
func withGaps(node, separator *syntax.Regexp, consumed *bool) *syntax.Regexp {
switch node.Op {
case syntax.OpLiteral:
parts := make([]*syntax.Regexp, 0, len(node.Rune))
for _, r := range node.Rune {
atom := &syntax.Regexp{Op: syntax.OpLiteral, Rune: []rune{r}, Flags: node.Flags}
parts = append(parts, gapBefore(atom, separator, consumed))
}
if len(parts) == 1 {
return parts[0]
}
return &syntax.Regexp{Op: syntax.OpConcat, Sub: parts}
case syntax.OpCharClass, syntax.OpAnyChar, syntax.OpAnyCharNotNL:
return gapBefore(node, separator, consumed)
case syntax.OpConcat, syntax.OpCapture:
for i := range node.Sub {
node.Sub[i] = withGaps(node.Sub[i], separator, consumed)
}
return node
case syntax.OpAlternate:
start, after := *consumed, *consumed
for i := range node.Sub {
branch := start
node.Sub[i] = withGaps(node.Sub[i], separator, &branch)
after = after || branch
}
*consumed = after
return node
case syntax.OpStar, syntax.OpPlus, syntax.OpQuest, syntax.OpRepeat:
// Every shape opens with a literal prefix, so a repetition is never the
// first thing a match consumes. If one ever is, its first iteration would
// need no gap and the rest would; fail loudly rather than guess.
if !*consumed {
panic("redaction: gap-tolerant rewrite reached a repetition before any literal")
}
node.Sub[0] = withGaps(node.Sub[0], separator, consumed)
return node
case syntax.OpWordBoundary, syntax.OpNoWordBoundary, syntax.OpBeginLine, syntax.OpEndLine,
syntax.OpBeginText, syntax.OpEndText, syntax.OpEmptyMatch:
return node
default:
panic(fmt.Sprintf("redaction: gap-tolerant rewrite does not handle %v", node.Op))
}
}
func gapBefore(atom, separator *syntax.Regexp, consumed *bool) *syntax.Regexp {
if !*consumed {
*consumed = true
return atom
}
gap := &syntax.Regexp{Op: syntax.OpStar, Sub: []*syntax.Regexp{separator}}
return &syntax.Regexp{Op: syntax.OpConcat, Sub: []*syntax.Regexp{gap, atom}}
}
// openAIShapeIsSecret is the digit/kebab-case filter RedactString applies to an
// openaiKeyPattern match, named so both callers state the same rule once.
func openAIShapeIsSecret(match string) bool {
if knownOpenAIKeyPrefix(match) || secretMatchHasDigit(match) {
return true
}
return !strings.Contains(strings.TrimPrefix(match, "sk-"), "-")
}
// redactControlSplitSecrets replaces credentials whose body is interrupted by
// invisible characters. It is a no-op for text containing none of them, which is
// nearly all text, so the contiguous path keeps its behavior exactly.
//
// It runs over the contiguous matchers' output and stops at every replacement
// they left behind: those mark text already accounted for, and a match reaching
// across one would be reaching out of one credential and into the next.
//
// It runs at most twice. The second pass exists for a credential glued straight
// onto the end of another split one: it has no word boundary of its own until
// the first is replaced, and the replacement then supplies one, the same way the
// contiguous shapes cascade into each other. Two passes and no more keeps the
// cost linear; a longer chain of glued split credentials is left as the
// contiguous matchers leave a chain of glued whole ones.
func redactControlSplitSecrets(value, replacement string) string {
if !containsSplitSeparator(value) {
return value
}
once := redactSplitRegions(value, replacement, true)
if once == value {
return value
}
// Ungated: the credential the second pass exists for is often contiguous
// itself (a whole JWT glued onto a split AWS key), so its region can hold no
// separator at all once the first pass has cut the text around it.
return redactSplitRegions(once, replacement, false)
}
func redactSplitRegions(value, replacement string, gated bool) string {
if replacement == "" {
return redactSplitRegion(value, replacement, 0, gated)
}
regions := strings.Split(value, replacement)
if len(regions) == 1 {
return redactSplitRegion(value, replacement, 0, gated)
}
previous := byte(0)
for i, region := range regions {
regions[i] = redactSplitRegion(region, replacement, previous, gated)
previous = replacement[len(replacement)-1]
}
return strings.Join(regions, replacement)
}
// redactSplitRegion runs every gap-tolerant shape over one stretch of text with
// no replacement inside it and replaces the UNION of what they found.
//
// Independently, not one after another. An unbounded body with gaps allowed will
// run on across a separator into whatever follows, and when that is the start of
// another credential it takes that credential's prefix with it: in
// "sk-ant-…<ESC>build<ESC>eyJ….eyJ….sig" the Anthropic shape stops only at the
// JWT's first dot. Applied in sequence, the JWT shape then finds its header gone
// and the payload and signature are left in the clear. Matched against the same
// text, the JWT shape still finds the whole token, and the union covers both.
//
// previous is the byte immediately before the region in the original, or zero at
// the start of the string. It is lent to the engine in front of the region so a
// leading \b is judged against the real neighbour, and a match is never allowed
// to begin on it.
func redactSplitRegion(value, replacement string, previous byte, gated bool) string {
if gated && !containsSplitSeparator(value) {
return value
}
text, lead := value, 0
if previous != 0 {
text, lead = string([]byte{previous})+value, 1
}
var spans [][]int
for _, shape := range splitShapes {
for _, match := range shape.free.FindAllStringIndex(text, -1) {
if match[0] >= lead {
spans = append(spans, shape.spansFor(text, match[0], match[1])...)
}
}
}
return spliceSpans(text, spans, replacement)[lead:]
}
// spliceSpans replaces every named range of value with replacement, merging
// ranges that overlap so one stretch of text is never replaced twice.
func spliceSpans(value string, spans [][]int, replacement string) string {
if len(spans) == 0 {
return value
}
sort.Slice(spans, func(i, j int) bool {
if spans[i][0] != spans[j][0] {
return spans[i][0] < spans[j][0]
}
return spans[i][1] > spans[j][1]
})
var builder strings.Builder
builder.Grow(len(value))
written := 0
for _, span := range spans {
if span[0] < written {
// Overlaps the stretch just replaced: extend it rather than emit a
// second marker, which would claim two credentials where there was one
// stretch of text.
if span[1] > written {
written = span[1]
}
continue
}
builder.WriteString(value[written:span[0]])
builder.WriteString(replacement)
written = span[1]
}
builder.WriteString(value[written:])
return builder.String()
}
-622
View File
@@ -1,622 +0,0 @@
package redaction
import (
"fmt"
"regexp"
"runtime"
"strings"
"testing"
"unicode"
)
const (
awsKey = "AKIAIOSFODNN7EXAMPLE"
anthropicKey = "sk-ant-api03-aaaaaaaaaaaaaaaaaaaaaaaa0123456789ABCD"
githubKey = "ghp_cccccccccccccccccccccccccccccccccccc"
openaiKey = "sk-proj-bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"
)
// rejoin drops every character a reader would never see, which is what a
// terminal, a log viewer or a copy-paste does to this output. If the secret is
// back after that, the redaction did not happen.
func rejoin(value string) string {
return strings.Map(func(r rune) rune {
if splitSecretSeparator(r) {
return -1
}
return r
}, value)
}
// AN INVISIBLE BYTE INSIDE A KEY MUST NOT END THE MATCH.
//
// Every shape matcher describes a contiguous body, so one NUL or ESC in the
// middle of a key ended the match and both fragments were emitted. Dropping the
// byte — which is what the next reader does, because it was never displayed —
// put the original key back. Reported in #969 for NUL and ESC; this covers the
// classes those two belong to, because every member behaves the same way.
func TestRedactStringRedactsSecretsSplitByInvisibleCharacters(t *testing.T) {
separators := map[string]string{
"NUL": "\x00",
"ESC": "\x1b",
"BEL": "\x07",
"backspace": "\x08",
"vertical tab": "\x0b",
"form feed": "\x0c",
"DEL": "\x7f",
"C1 NEL": "\u0085",
"C1 CSI": "\u009b",
"zero width space": "\u200b",
"zero width joiner": "\u200d",
"word joiner": "\u2060",
"byte order mark": "\ufeff",
"soft hyphen": "\u00ad",
}
// Every entry has to BE a separator, or its leg passes vacuously: inserting an
// ordinary character into a key also stops the secret matching, and the rejoin
// below would then find nothing. An escape mangled on the way into this file is
// exactly how that happens.
for name, separator := range separators {
if runes := []rune(separator); len(runes) != 1 || !splitSecretSeparator(runes[0]) {
t.Fatalf("SETUP INVALID: %s is %q, which is not a single invisible separator", name, separator)
}
}
secrets := map[string]string{"aws": awsKey, "anthropic": anthropicKey, "github": githubKey, "openai": openaiKey}
for secretName, secret := range secrets {
// The premise: the unsplit form already redacts, so a failure below is
// about the separator and not about the shape being unknown.
if out := RedactString(secret, Options{}); strings.Contains(out, secret) {
t.Fatalf("SETUP INVALID: the unsplit %s key is not redacted at all: %q", secretName, out)
}
for sepName, separator := range separators {
for _, at := range []int{1, len(secret) / 2, len(secret) - 1} {
split := secret[:at] + separator + secret[at:]
out := RedactString(split, Options{})
if got := rejoin(out); strings.Contains(got, secret) {
t.Errorf("%s key split by %s at byte %d survives: %q rejoins to %q", secretName, sepName, at, out, got)
}
}
}
}
}
// The replacement has to cover the ORIGINAL bytes, separators included. Leaving
// the separator behind would emit `[REDACTED]\x00[REDACTED]` and tell a reader
// there were two secrets, and leaving either fragment behind is the leak itself.
func TestRedactStringCoversTheSeparatorsInsideASplitSecret(t *testing.T) {
for _, testCase := range []struct {
name string
value string
want string
}{
{"one separator", "key=" + awsKey[:8] + "\x00" + awsKey[8:], "key=[REDACTED]"},
{"several separators", "key=" + strings.Join(strings.Split(awsKey, ""), "\x1b"), "key=[REDACTED]"},
{"adjacent separators", "key=" + awsKey[:8] + "\x00\x1b\u200b" + awsKey[8:], "key=[REDACTED]"},
{"two split secrets", awsKey[:8] + "\x00" + awsKey[8:] + " and " + githubKey[:8] + "\x00" + githubKey[8:], "[REDACTED] and [REDACTED]"},
{"separator before the secret", "\x00" + awsKey, "\x00[REDACTED]"},
{"separator after the secret", awsKey + "\x00", "[REDACTED]\x00"},
{"text around it survives", "before " + awsKey[:4] + "\x00" + awsKey[4:] + " after", "before [REDACTED] after"},
} {
t.Run(testCase.name, func(t *testing.T) {
if got := RedactString(testCase.value, Options{}); got != testCase.want {
t.Errorf("RedactString(...) = %q, want %q", got, testCase.want)
}
})
}
}
// NOTHING ELSE MOVES. The compaction exists to find a match; it must never
// reach text that has no secret in it, and text with no invisible character at
// all must take exactly the path it took before.
func TestRedactStringLeavesOrdinaryTextWithControlCharactersAlone(t *testing.T) {
for _, value := range []string{
"a plain sentence",
"a sentence\x00with a NUL in it",
"\x1b[31mcolored output\x1b[0m",
"tab\tseparated\tcolumns",
"line one\nline two\r\nline three",
"zero\u200bwidth\u200bspaces",
"", "\x00", "\x00\x00\x00",
"not-a-secret-just-a-long-kebab-case-identifier-here",
"sk-some-kebab-case-value-without-any-digits-at-all",
} {
if got := RedactString(value, Options{}); got != value {
t.Errorf("RedactString(%q) = %q, want it unchanged", value, got)
}
}
}
// The openai filter that keeps kebab-case identifiers out of the redactor has
// to apply on the split path too, or working around the false positive would
// become a matter of inserting a control byte.
func TestSplitPathKeepsTheOpenAIKebabCaseFilter(t *testing.T) {
kebab := "sk-some-kebab-case-value-without-digits"
split := kebab[:6] + "\x00" + kebab[6:]
if got := RedactString(split, Options{}); got != split {
t.Errorf("a kebab-case identifier split by a NUL was redacted: %q", got)
}
// ... while a real key of the same family, split the same way, is not spared.
realKey := openaiKey[:6] + "\x00" + openaiKey[6:]
if got := RedactString(realKey, Options{}); strings.Contains(rejoin(got), openaiKey) {
t.Errorf("a real openai key split by a NUL survived: %q", got)
}
}
// TAB, NEWLINE AND CARRIAGE RETURN ARE OUT OF SCOPE ON PURPOSE, and this says
// so in a place that fails if someone changes it without meaning to. They are
// real text structure: stripping them would let one match span a line break and
// replace unrelated lines, and a credential broken across a line is visible to
// whoever reads it rather than hidden from them.
func TestLineStructureIsNotTreatedAsAnInvisibleSeparator(t *testing.T) {
for name, separator := range map[string]string{"tab": "\t", "newline": "\n", "carriage return": "\r"} {
if splitSecretSeparator([]rune(separator)[0]) {
t.Errorf("%s is treated as an invisible separator; if that is now wanted, this test and the note in control_split.go both need to change", name)
}
split := awsKey[:8] + separator + awsKey[8:]
if got := RedactString(split, Options{}); got != split {
t.Errorf("%s split was redacted: %q. That may be an improvement, but it is a scope change this test is here to make deliberate", name, got)
}
}
}
// Malformed UTF-8 must not move the span mapping, which works in the source's
// own byte offsets. A decoder that measured an invalid byte as the three bytes
// of U+FFFD would shift every later span.
// The exact output is asserted, not just the absence of the key. An invalid
// byte measured as three would still cover the secret while moving the
// replacement's edges onto the bytes around it, which reads as a pass if the
// test only asks whether the secret is gone.
func TestSplitRedactionHandlesMalformedUTF8(t *testing.T) {
for _, testCase := range []struct {
name string
value string
want string
}{
{"invalid byte before the secret", "\xff" + awsKey[:8] + "\x00" + awsKey[8:], "\xff[REDACTED]"},
// An invalid byte is NOT one of the separators: it decodes to U+FFFD and
// is drawn, so a key broken by one is visible to whoever reads it, the
// same argument that keeps tab and newline out. It is kept in the
// compacted text and so it still ends the match.
{"invalid byte inside the split", awsKey[:8] + "\x00\xff" + awsKey[8:], awsKey[:8] + "\x00\xff" + awsKey[8:]},
{"invalid byte after the secret", awsKey[:8] + "\x00" + awsKey[8:] + "\xfe", "[REDACTED]\xfe"},
{"invalid bytes on both sides", "\xff\xfe" + awsKey[:8] + "\x00" + awsKey[8:] + "\xfd\xfc", "\xff\xfe[REDACTED]\xfd\xfc"},
{"invalid byte then text then the secret", "\xffpre " + awsKey[:8] + "\x00" + awsKey[8:] + " post", "\xffpre [REDACTED] post"},
{"only invalid bytes", "\xff\xfe\x00\xff", "\xff\xfe\x00\xff"},
} {
t.Run(testCase.name, func(t *testing.T) {
got := RedactString(testCase.value, Options{})
if got != testCase.want {
t.Errorf("RedactString(%q) = %q, want %q", testCase.value, got, testCase.want)
}
if strings.Contains(rejoin(got), awsKey) {
t.Errorf("the key survived: %q", got)
}
})
}
}
// The compaction is only a lens for finding the match. A string with no
// separator at all must come back byte-identical, including its invalid bytes.
func TestCompactionDoesNotRewriteTextItDoesNotRedact(t *testing.T) {
for _, value := range []string{"plain", "with\xffinvalid", "tab\there", "\xef\xbb\xbfbom at the front"} {
if got := redactControlSplitSecrets(value, RedactedSecret); got != value {
t.Errorf("redactControlSplitSecrets(%q) = %q, want it unchanged", value, got)
}
}
}
// The separators the tests below need by name, built from their code points so
// no escape in this file can be mangled into the character it stands for.
var (
nulSeparator = string(rune(0x00))
escSeparator = string(rune(0x1b))
zwspSeparator = string(rune(0x200b))
)
const (
slackKey = "xoxb-EXAMPLE-NOT-A-REAL-TOKEN-AAAAAAAAAA"
googleKey = "AIzaSyD-0123456789abcdefghijklmnopqrstuv"
jwtToken = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.dBjftJeZ4CVPmB92K27uhbUJU1p1r0W1gFWFOEjXkPY"
otherJWT = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiI5ODc2NTQzMjEwIn0.QWxpY2VCb2JDYXJvbERhdmVFdmVGcmFua0dyYWNlSA"
)
// splitInTheMiddle puts one separator inside the body of a secret.
func splitInTheMiddle(secret, separator string) string {
half := len(secret) / 2
return secret[:half] + separator + secret[half:]
}
// A SEPARATOR IN FRONT OF A KEY IS A DELIMITER, ONE INSIDE IT IS FILLER.
//
// #969 asks for one thing: the split form is redacted the same as the unsplit
// form. That is the whole assertion here, run over every kind of byte a key can
// follow. It catches the two ways of getting it wrong from opposite sides. Drop
// the separators and then ask the pattern for a word boundary, and a key written
// after a NUL stops redacting the moment a second separator lands inside it,
// even though the unsplit key after that same NUL redacts today. Ignore the
// boundary instead, and a key glued to the end of a word starts redacting when
// it is split, although the unsplit form never did.
func TestSplitRedactionMatchesTheUnsplitVerdictAtEveryLeadingBoundary(t *testing.T) {
leaders := map[string]string{
"start of string": "",
"space": "before ",
"quote": `{"k":"`,
"equals": "key=",
"NUL": "prefix" + nulSeparator,
"ESC": "prefix" + escSeparator,
"zero width": "prefix" + zwspSeparator,
"letter": "prefix",
"digit": "7",
"underscore": "field_",
}
secrets := map[string]string{"aws": awsKey, "github": githubKey, "openai": openaiKey, "slack": slackKey, "google": googleKey}
for secretName, secret := range secrets {
for leaderName, leader := range leaders {
unsplit := leader + secret
split := leader + splitInTheMiddle(secret, escSeparator)
wantRedacted := !strings.Contains(RedactString(unsplit, Options{}), secret)
out := RedactString(split, Options{})
gotRedacted := !strings.Contains(rejoin(out), secret)
if gotRedacted == wantRedacted {
continue
}
if wantRedacted {
t.Errorf("%s key after a %s leader: the unsplit form redacts and the split form does not: %q rejoins to %q",
secretName, leaderName, out, rejoin(out))
continue
}
t.Errorf("%s key after a %s leader: the unsplit form is left alone and the split form is redacted: %q",
secretName, leaderName, out)
}
}
}
// A MATCH MUST NOT REACH OUT OF ONE CREDENTIAL AND INTO THE NEXT.
//
// The JWT shapes end on a run of body characters with no trailing boundary, so
// a matcher reading a copy with the separators removed can start in the key in
// front and run through the header of the JWT behind it, replacing both with one
// marker and eating the delimiter between them. The contiguous pass claims whole
// credentials before this one runs, which is what keeps them apart.
func TestSplitRedactionKeepsNeighbouringCredentialsApart(t *testing.T) {
separators := map[string]string{"NUL": nulSeparator, "ESC": escSeparator, "zero width": zwspSeparator}
leaders := map[string]string{
"jwt": jwtToken,
"github": githubKey,
"anthropic": anthropicKey,
"openai": openaiKey,
"slack": slackKey,
"google": googleKey,
"aws": awsKey,
}
for leaderName, leader := range leaders {
for sepName, separator := range separators {
value := leader + separator + otherJWT
want := RedactedSecret + separator + RedactedSecret
if got := RedactString(value, Options{}); got != want {
t.Errorf("a %s key %s a JWT: RedactString(...) = %q, want %q", leaderName, sepName, got, want)
}
}
}
}
// splitThroughout puts a separator a quarter of the way into every dot-separated
// segment of a token. Early, so the part in front is too short to be a
// credential by itself and the strict pass cannot claim it; a split whose front
// part is complete alone is the documented limit, not what these tests are about.
func splitThroughout(secret, separator string) string {
parts := strings.Split(secret, ".")
for i, part := range parts {
quarter := len(part) / 4
parts[i] = part[:quarter] + separator + part[quarter:]
}
return strings.Join(parts, ".")
}
// TWO SPLIT CREDENTIALS OF ONE SHAPE ARE TWO CREDENTIALS.
//
// Matches of one pattern never overlap, so when the first one's body runs on
// through the separator between them, the second cannot start a match of its
// own. Whole neighbours are safe because the strict pass claims them first;
// these are both split. A JWT body stops at the next token's first dot and a
// classic GitHub body at the next token's underscore, which left the rest of the
// second credential in the clear, and the bodies that swallow the next key whole
// joined both keys and the separator between them into one replacement.
func TestSplitRedactionKeepsTwoSplitCredentialsOfOneShapeApart(t *testing.T) {
separators := map[string]string{"NUL": nulSeparator, "ESC": escSeparator, "zero width": zwspSeparator}
pairs := map[string][2]string{
"jwt": {jwtToken, otherJWT},
"github": {githubKey, "ghp_ffffffffffffffffffffffffffffffffffff"},
"openai": {openaiKey, "sk-proj-dddddddddddddddddddddddddddddddd"},
"anthropic": {anthropicKey, "sk-ant-api03-eeeeeeeeeeeeeeeeeeeeeeee9876543210WXYZ"},
"slack": {slackKey, "xoxb-ANOTHER-NOT-REAL-TOKEN-BBBBBBBBBBBB"},
}
for name, pair := range pairs {
for sepName, separator := range separators {
value := splitThroughout(pair[0], escSeparator) + separator + splitThroughout(pair[1], escSeparator)
want := RedactedSecret + separator + RedactedSecret
if got := RedactString(value, Options{}); got != want {
t.Errorf("two split %s credentials with %s between them: RedactString(...) = %q, want %q", name, sepName, got, want)
}
}
}
}
// The search for where the next credential begins is bounded, so it has to be
// spent on positions where one can begin. Here the second token's header is split
// at every character inside the part the first match swallowed, which puts dozens
// of separator runs between the end of that match and the real boundary.
func TestSplitRedactionFindsTheNextCredentialPastItsOwnSplits(t *testing.T) {
parts := strings.Split(otherJWT, ".")
header := strings.Join(strings.Split(parts[0], ""), escSeparator)
second := header + "." + strings.Join(parts[1:], ".")
value := splitThroughout(jwtToken, escSeparator) + nulSeparator + second
want := RedactedSecret + nulSeparator + RedactedSecret
if got := RedactString(value, Options{}); got != want {
t.Errorf("RedactString(...) = %q, want %q", got, want)
}
}
// And a chain of them leaks nothing, however the markers fall.
func TestSplitRedactionLeavesNothingOfAChainOfSplitJWTs(t *testing.T) {
chain := splitThroughout(jwtToken, escSeparator) + nulSeparator +
splitThroughout(otherJWT, escSeparator) + nulSeparator +
splitThroughout(jwtToken, zwspSeparator)
out := rejoin(RedactString(chain, Options{}))
for _, segment := range append(strings.Split(jwtToken, "."), strings.Split(otherJWT, ".")...) {
if strings.Contains(out, segment) {
t.Errorf("a JWT segment survived a chain of three split tokens: %q in %q", segment, out)
}
}
}
// A CHAIN OF SPLIT CREDENTIALS COSTS WHAT ITS LENGTH COSTS.
//
// The run-on search once followed a chain to its end from every match inside it,
// which is quadratic in the chain: a megabyte of split JWTs back to back
// allocated seventeen gigabytes and ran for ten minutes. Measured as allocation,
// which that walk multiplies with the chain and a linear pass does not.
func TestSplitRedactionStaysLinearOnAChainOfSplitJWTs(t *testing.T) {
unit := splitThroughout(jwtToken, escSeparator) + nulSeparator
chain := strings.Repeat(unit, (64<<10)/len(unit))
signature := strings.Split(jwtToken, ".")[2]
if strings.Contains(rejoin(RedactString(chain, Options{})), signature) {
t.Fatal("SETUP INVALID: the chain is not redacted, so the run-on search never ran")
}
limit := int64(200 * len(chain))
if used := redactAllocBytes(chain); used > limit {
t.Errorf("redacting a %d byte chain of split JWTs allocated %d bytes, over the %d byte limit", len(chain), used, limit)
}
}
// A MATCH THE FILTER REJECTS MUST NOT HIDE THE KEY IT RAN INTO.
//
// The OpenAI shape is filtered, because kebab-case with no digit in it is prose
// and not a key. A rejected match still spans the text it matched, and its body
// runs on through a separator into whatever follows, so a split key after a
// kebab-case name was rejected along with the name and printed in full.
func TestSplitRedactionDoesNotLetARejectedNeighbourHideASplitKey(t *testing.T) {
const prose = "sk-some-kebab-case-name-here-ok"
const key = "sk-ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefgh"
if !strings.Contains(RedactString(prose, Options{}), prose) {
t.Fatal("SETUP INVALID: the prose is redacted on its own, so it is not a rejected match")
}
if strings.Contains(RedactString(key, Options{}), key) {
t.Fatal("SETUP INVALID: the key is not redacted even whole, so it is not a credential")
}
value := prose + nulSeparator + splitInTheMiddle(key, escSeparator)
want := prose + nulSeparator + RedactedSecret
if got := RedactString(value, Options{}); got != want {
t.Errorf("RedactString(...) = %q, want %q", got, want)
}
}
// ORDINARY TEXT AFTER A KEY IS NOT PART OF THE KEY. Removing the separator
// between them joins the word behind it onto the end of the key, and an
// unbounded shape then carries the replacement over text that was never secret.
func TestSplitRedactionLeavesTextBehindAKeyAlone(t *testing.T) {
for _, testCase := range []struct {
name string
value string
want string
}{
{"word", openaiKey + nulSeparator + "ordinary prose here", RedactedSecret + nulSeparator + "ordinary prose here"},
{"escape then word", openaiKey + escSeparator + "ordinary", RedactedSecret + escSeparator + "ordinary"},
{"zero width then word", githubKey + zwspSeparator + "ordinary", RedactedSecret + zwspSeparator + "ordinary"},
{"digits", awsKey + nulSeparator + "0123456789", RedactedSecret + nulSeparator + "0123456789"},
} {
t.Run(testCase.name, func(t *testing.T) {
if got := RedactString(testCase.value, Options{}); got != testCase.want {
t.Errorf("RedactString(...) = %q, want %q", got, testCase.want)
}
})
}
}
// THE OFFSET TABLE IS SIZED BY THE MATCHES, NOT BY THE TEXT IT SEARCHED.
//
// RedactString is handed whole command output, verification stdout and
// notification bodies, none of which is bounded before it runs, and an index
// carrying a start and an end for every source byte costs sixteen bytes per byte
// of input on a 64-bit build. Measured against the same text with no separator
// in it, so the regex work and the output copy cancel out and what is left is
// what the split path added.
func TestSplitRedactionDoesNotIndexEverySourceByte(t *testing.T) {
const line = "plain log output line with no secret in it "
filler := strings.Repeat(line, (1<<20)/len(line))
split := filler + " " + splitInTheMiddle(awsKey, escSeparator)
plain := filler + " " + awsKey
if got := RedactString(split, Options{}); !strings.HasSuffix(got, RedactedSecret) {
t.Fatalf("SETUP INVALID: the split key in the large input is not redacted, so nothing indexes it: %q", got[len(got)-64:])
}
// Four times the input still fails an index of two ints per source byte,
// which needs sixteen, and clears the copy this path actually makes.
limit := int64(4 * len(split))
overhead := redactAllocBytes(split) - redactAllocBytes(plain)
if overhead > limit {
t.Errorf("redacting %d bytes with a split key allocated %d bytes more than the same text without one, over the %d byte limit",
len(split), overhead, limit)
}
}
// redactAllocBytes is the smallest number of bytes any of three RedactString
// calls allocated. Smallest rather than mean because a garbage collection or an
// unrelated goroutine can only ever add to the figure.
func redactAllocBytes(value string) int64 {
smallest := int64(-1)
for i := 0; i < 3; i++ {
var before, after runtime.MemStats
runtime.GC()
runtime.ReadMemStats(&before)
out := RedactString(value, Options{})
runtime.ReadMemStats(&after)
if len(out) == 0 {
panic("RedactString returned nothing")
}
used := int64(after.TotalAlloc - before.TotalAlloc)
if smallest < 0 || used < smallest {
smallest = used
}
}
return smallest
}
// A REPLACEMENT ALREADY ON THE PAGE IS A BARRIER. The default marker is spelled
// with brackets, which no credential body admits, so it stops a match by
// itself. A caller is free to supply one made of ordinary word characters, and
// then the text left by the contiguous pass would read as the middle of a key
// and carry the replacement out over both delimiters and the text behind them.
func TestSplitRedactionStopsAtACallerSuppliedReplacement(t *testing.T) {
options := Options{Replacement: "REDACTED"}
value := "sk-aaaaa" + nulSeparator + awsKey + nulSeparator + "bbbbbbbbbbbbbbb"
want := "sk-aaaaa" + nulSeparator + "REDACTED" + nulSeparator + "bbbbbbbbbbbbbbb"
if got := RedactString(value, options); got != want {
t.Errorf("RedactString(...) = %q, want %q", got, want)
}
}
// A REJECTED MATCH MUST NOT SHADOW THE REAL KEY BEHIND IT. Raised by CodeRabbit
// on #1067. With the separators removed before matching, the leftmost AWS match
// in "xAKIA<NUL>AKIAIOSF<SOH>ODNN7EXAMPLE" began at the decoy AKIA, was rejected
// for the "x" in front of it, and had already consumed the start of the real key,
// so the real key was never tried. "ask-" in front of a split OpenAI key does the
// same through "sk-", which is the ordinary-prose version of it. Each input is
// checked against its unsplit counterpart, which main already redacts.
func TestSplitRedactionDoesNotLetARejectedMatchShadowTheRealKey(t *testing.T) {
soh := string(rune(0x01))
for _, testCase := range []struct {
name string
split string
unsplit string
}{
{"decoy AKIA before a split key", "xAKIA" + nulSeparator + "AKIAIOSFODNN7E" + soh + "XAMPLE", "xAKIA" + nulSeparator + awsKey},
{"ask- before a split OpenAI key", "ask-" + nulSeparator + splitInTheMiddle(openaiKey, soh), "ask-" + nulSeparator + openaiKey},
{"task- before a split OpenAI key", "task-" + escSeparator + splitInTheMiddle(openaiKey, zwspSeparator), "task-" + escSeparator + openaiKey},
} {
t.Run(testCase.name, func(t *testing.T) {
want := RedactString(testCase.unsplit, Options{})
if strings.Contains(want, awsKey) || strings.Contains(want, openaiKey) {
t.Fatalf("SETUP INVALID: main does not redact the unsplit form either: %q", want)
}
if got := RedactString(testCase.split, Options{}); got != want {
t.Errorf("RedactString(split) = %q, want the unsplit form's %q", got, want)
}
})
}
}
// AN UNBOUNDED BODY MUST NOT TAKE A NEIGHBOUR'S PREFIX WITH IT. With gaps allowed
// between body characters, a split key's body runs on across a separator into
// whatever follows. When that is a JWT, it stops only at the JWT's first dot, and
// if the shapes were applied one after another the JWT shape would then find its
// header gone and leave the payload and signature in the clear. The shapes are
// matched against the same text and their union replaced, so it cannot.
func TestSplitRedactionKeepsANeighbourWhosePrefixASplitBodyWouldSwallow(t *testing.T) {
cut := len(jwtToken) / 2
splitJWT := jwtToken[:cut] + nulSeparator + jwtToken[cut:]
for _, testCase := range []struct {
name string
value string
}{
{"split Anthropic key, a word, then a split JWT", splitInTheMiddle(anthropicKey, zwspSeparator) + escSeparator + "build" + escSeparator + splitJWT},
{"split GitHub key, then a split JWT", splitInTheMiddle(githubKey, nulSeparator) + escSeparator + splitJWT},
} {
t.Run(testCase.name, func(t *testing.T) {
got := rejoin(RedactString(testCase.value, Options{}))
for _, piece := range strings.Split(jwtToken, ".") {
if strings.Contains(got, piece) {
t.Errorf("part of the JWT survives: %q in %q", piece, got)
}
}
})
}
}
// A CREDENTIAL GLUED ONTO A SPLIT ONE GETS ITS BOUNDARY FROM THE REPLACEMENT.
// Contiguous shapes cascade: once an AWS key is replaced, a JWT written straight
// after it has the "]" in front of it and matches. A split key is replaced by
// the second pass, so the JWT behind it needs a pass of its own after that one,
// and the JWT is often whole, which means its region may hold no separator at
// all by then. The second case adds a whole key elsewhere so that the strict
// pass has already cut the text into regions.
func TestSplitRedactionCascadesIntoACredentialGluedOntoASplitOne(t *testing.T) {
csi := string(rune(0x9b))
for _, testCase := range []struct {
name string
value string
}{
{"whole JWT glued onto a split AWS key", splitInTheMiddle(awsKey, escSeparator) + jwtToken},
{"same, with a whole key elsewhere in the text", "user?q=" + "ASIAIOSFODNN7EXA" + csi + "MPLE" + jwtToken + "'" + awsKey + string(rune(0x7f)) + "ok"},
} {
t.Run(testCase.name, func(t *testing.T) {
got := rejoin(RedactString(testCase.value, Options{}))
for _, piece := range strings.Split(jwtToken, ".") {
if strings.Contains(got, piece) {
t.Errorf("part of the JWT survives: %q in %q", piece, got)
}
}
})
}
}
// THE CLASS AND THE PREDICATE ARE ONE FACT. splitSeparatorClass is generated
// from the same tables as splitSecretSeparator; this walks every code point and
// fails on the first one they disagree about.
func TestSplitSeparatorClassAgreesWithThePredicate(t *testing.T) {
class := regexp.MustCompile("^" + splitSeparatorClass + "$")
for r := rune(0); r <= unicode.MaxRune; r++ {
if r >= 0xd800 && r <= 0xdfff {
continue // surrogates are not valid in a Go string
}
if got, want := class.MatchString(string(r)), splitSecretSeparator(r); got != want {
t.Fatalf("U+%04X: class matches = %v, predicate says %v", r, got, want)
}
}
}
// ON TEXT WITH NO SEPARATOR, EVERY GAP-TOLERANT SHAPE IS ITS ORIGINAL. The
// rewrite only adds optional gaps, so a contiguous key has to produce exactly the
// same match; anything else means the rewrite changed what a shape is.
func TestGapTolerantShapesMatchContiguousKeysLikeTheOriginals(t *testing.T) {
inputs := []string{
"key=" + awsKey + " and " + githubKey,
"x" + awsKey,
anthropicKey + "." + openaiKey,
jwtToken + " then " + slackKey + ", " + googleKey,
"sk-this-is-kebab-case-prose-not-a-key",
"github_pat_11ABCDEFG0123456789_abcdefghijklmnopqrstuvwxyzABCDEFGH/glpat-abcdefghij0123456789",
}
originals := append([]*regexp.Regexp{openaiKeyPattern}, textSecretPatterns...)
tolerant := append([]*regexp.Regexp{splitOpenAIPattern}, splitTextPatterns...)
if len(originals) != len(tolerant) {
t.Fatalf("SETUP INVALID: %d original shapes and %d gap-tolerant ones", len(originals), len(tolerant))
}
for i := range originals {
for _, input := range inputs {
want := fmt.Sprint(originals[i].FindAllStringIndex(input, -1))
if got := fmt.Sprint(tolerant[i].FindAllStringIndex(input, -1)); got != want {
t.Errorf("shape %s on %q: gap-tolerant matched %s, original %s", originals[i], input, got, want)
}
}
}
}
+4 -71
View File
@@ -96,69 +96,6 @@ var textSecretPatterns = []*regexp.Regexp{
regexp.MustCompile(`\beyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}`),
}
// anchoredSecretPatterns mirrors textSecretPatterns with the leading \b
// replaced by \A, so each shape can be tested at an exact byte offset. A
// credential that starts exactly where a redacted span ends (two AWS keys
// glued end to end, e.g. AKIA...AKIA...) has no word boundary on its left,
// so the plain pattern misses it; the anchored variant catches it.
var anchoredSecretPatterns = func() []*regexp.Regexp {
out := make([]*regexp.Regexp, len(textSecretPatterns))
for i, p := range textSecretPatterns {
src := p.String()
if strings.HasPrefix(src, `\b`) {
src = `\A(?:` + src[len(`\b`):] + `)`
}
out[i] = regexp.MustCompile(src)
}
return out
}()
// redactAdjacent redacts every match of pattern, plus any match of anchored
// that begins exactly where a redacted span ends. This catches credentials
// glued end to end (AKIA...AKIA...) without weakening the leading-boundary
// rule elsewhere: a mid-word occurrence that does not abut a redacted span
// still does not match. Chained runs (three or more glued credentials) are
// followed to the end.
func redactAdjacent(s string, pattern, anchored *regexp.Regexp, replacement string) string {
spans := pattern.FindAllStringIndex(s, -1)
// Bound the outer loop to the original match count. The inner loop
// appends chained spans to the same slice; re-reading len(spans) would
// re-chain every appended span, growing the span count (and the
// FindStringIndex calls) exponentially on long glued runs.
n := len(spans)
for i := 0; i < n; i++ {
end := spans[i][1]
for end < len(s) {
loc := anchored.FindStringIndex(s[end:])
if loc == nil || loc[0] != 0 || loc[1] <= 0 {
break
}
newEnd := end + loc[1]
spans = append(spans, []int{end, newEnd})
end = newEnd
}
}
if len(spans) == 0 {
return s
}
sort.Slice(spans, func(i, j int) bool { return spans[i][0] < spans[j][0] })
var out strings.Builder
pos := 0
for _, sp := range spans {
if sp[0] < pos {
continue // overlapping span, already covered
}
out.WriteString(s[pos:sp[0]])
// s[sp[0]:sp[1]] is a match (anchored spans match at \b when taken
// alone), so ReplaceAllString expands $ references exactly as the
// plain loop below would.
out.WriteString(pattern.ReplaceAllString(s[sp[0]:sp[1]], replacement))
pos = sp[1]
}
out.WriteString(s[pos:])
return out.String()
}
var (
privateKeyPattern = regexp.MustCompile(`(?s)-----BEGIN [A-Z ]*PRIVATE KEY-----.*?-----END [A-Z ]*PRIVATE KEY-----`)
jsonStringPattern = regexp.MustCompile(`("([^"\\]*(?:\\.[^"\\]*)*)"\s*:\s*)"([^"\\]*(?:\\.[^"\\]*)*)"`)
@@ -290,19 +227,15 @@ func RedactString(value string, options Options) string {
// openai keys first so the filter can drop kebab-case false positives
// before any other pattern rewrites nearby text.
redacted = openaiKeyPattern.ReplaceAllStringFunc(redacted, func(match string) string {
if !openAIShapeIsSecret(match) {
if !knownOpenAIKeyPrefix(match) && !secretMatchHasDigit(match) &&
strings.Contains(strings.TrimPrefix(match, "sk-"), "-") {
return match
}
return replacement
})
for i, pattern := range textSecretPatterns {
redacted = redactAdjacent(redacted, pattern, anchoredSecretPatterns[i], replacement)
for _, pattern := range textSecretPatterns {
redacted = pattern.ReplaceAllString(redacted, replacement)
}
// Credentials broken up by invisible characters last, over what the
// contiguous matchers left: a whole credential is already claimed by its own
// strict match, so this pass cannot reach out of one and into the next. It is
// a no-op for text with no such character, which is nearly all text.
redacted = redactControlSplitSecrets(redacted, replacement)
return redacted
}
-44
View File
@@ -4,32 +4,8 @@ import (
"errors"
"strings"
"testing"
"time"
)
func TestRedactStringRedactsAdjacentAWSKeys(t *testing.T) {
// Two AWS keys glued end to end have no word boundary between them, so a
// plain \b-anchored pattern only matches the first. Both must be redacted.
// gitleaks:allow -- synthetic redaction fixtures below
key1 := "AKIAIOSFODNN7EXAMPLE"
key2 := "ASIAIOSFODNN7EXAMPLE"
input := key1 + key2
if got := RedactString(input, Options{}); got != RedactedSecret+RedactedSecret {
t.Fatalf("adjacent AWS keys not both redacted: got %q", got)
}
// Three glued keys chain as well.
input3 := key1 + key2 + key1
if got := RedactString(input3, Options{}); got != RedactedSecret+RedactedSecret+RedactedSecret {
t.Fatalf("three adjacent AWS keys not all redacted: got %q", got)
}
// A mid-word AKIA that does not abut a redacted span still must not match.
if got := RedactString("prefixAKIAIOSFODNN7EXAMPLE", Options{}); got != "prefixAKIAIOSFODNN7EXAMPLE" {
t.Fatalf("mid-word AKIA should not be redacted: got %q", got)
}
}
func TestRedactStringCoversCommonSecretShapes(t *testing.T) {
input := strings.Join([]string{
`{"apiKey":"sk-proj-abcdefghijklmnopqrstuvwxyz"}`,
@@ -165,23 +141,3 @@ func containsCircular(v any) bool {
}
return false
}
func TestRedactStringLongAdjacentKeyRunCompletesPromptly(t *testing.T) {
// A long run of glued AWS keys must not blow up the span-chaining loop in
// redactAdjacent: the outer loop is bounded to the original match count,
// so appended spans are never re-chained. On the old code this input
// grows the span list exponentially and never finishes.
// The key is assembled at runtime so the literal never appears in source.
// gitleaks:allow -- synthetic redaction fixture below
key := "AKIA" + strings.Repeat("A", 16)
const count = 64
input := strings.Repeat(key, count)
start := time.Now()
got := RedactString(input, Options{})
if elapsed := time.Since(start); elapsed > 5*time.Second {
t.Fatalf("RedactString took %v on %d glued keys", elapsed, count)
}
if want := strings.Repeat(RedactedSecret, count); got != want {
t.Fatalf("expected %d redaction markers, got %d", count, strings.Count(got, RedactedSecret))
}
}
-15
View File
@@ -207,21 +207,6 @@ func (backend Backend) EnforcementLevel(policy Policy) EnforcementLevel {
return EnforcementDegraded
}
// DegradedNotice is one line saying the plan's enforcement is degraded and why,
// or "" at any other level. Only degraded: a disabled sandbox is the user's own
// choice, and the unelevated Windows tier still enforces the write jail and
// says what it lacks through `zero sandbox policy`.
func (plan BackendPlan) DegradedNotice() string {
if plan.EnforcementLevel != EnforcementDegraded {
return ""
}
reason := strings.TrimRight(strings.TrimSpace(plan.DowngradeReason), ".")
if reason == "" {
reason = "no native sandbox is available"
}
return "Sandbox enforcement is degraded: " + reason + ". See `zero sandbox policy --effective`."
}
func (backend Backend) DowngradeReason(policy Policy) string {
if policy.Mode == "" {
policy = DefaultPolicy()
-187
View File
@@ -1,187 +0,0 @@
package sandbox
import (
"fmt"
"path/filepath"
"strings"
"testing"
)
// Only DEGRADED says anything. A disabled sandbox is the user's own choice, and
// the unelevated Windows tier still enforces the write jail and reports what it
// lacks through `zero sandbox policy`; a notice on every session for either
// would stop meaning anything.
func TestDegradedNoticeSpeaksOnlyForADegradedPlan(t *testing.T) {
for _, level := range []EnforcementLevel{EnforcementNative, EnforcementUnelevated, EnforcementDisabled, ""} {
plan := BackendPlan{EnforcementLevel: level, DowngradeReason: "some reason"}
if notice := plan.DegradedNotice(); notice != "" {
t.Errorf("level %q produced a notice: %q", level, notice)
}
}
plan := BackendPlan{EnforcementLevel: EnforcementDegraded, DowngradeReason: "Linux sandbox helper is not available"}
want := "Sandbox enforcement is degraded: Linux sandbox helper is not available. See `zero sandbox policy --effective`."
if got := plan.DegradedNotice(); got != want {
t.Errorf("DegradedNotice() = %q, want %q", got, want)
}
}
// A reason that already ends a sentence keeps one period, and a missing reason
// still says what is wrong.
func TestDegradedNoticeWordsItsReasonOnce(t *testing.T) {
withPeriod := BackendPlan{EnforcementLevel: EnforcementDegraded, DowngradeReason: "Run `zero sandbox setup` first."}
if got := withPeriod.DegradedNotice(); strings.Contains(got, "..") {
t.Errorf("the reason's own period was doubled: %q", got)
}
empty := BackendPlan{EnforcementLevel: EnforcementDegraded}
if got := empty.DegradedNotice(); !strings.Contains(got, "no native sandbox is available") {
t.Errorf("a degraded plan with no reason gave no reason: %q", got)
}
}
// THE SESSION'S NOTICE AND `zero sandbox policy` DESCRIBE ONE SANDBOX. The engine
// builds its notice from its own backend and policy through BuildPlan, the
// function the policy command prints from, rather than deciding "degraded" a
// second way.
func TestEngineDegradedNoticeComesFromThePolicyCommandsPlan(t *testing.T) {
clearNestingMarkers(t)
root := t.TempDir()
unavailable := Backend{Name: BackendUnavailable, Platform: "linux", Fallback: true, Message: "Linux sandbox helper is not available"}
engine := NewEngine(EngineOptions{WorkspaceRoot: root, Policy: DefaultPolicy(), Backend: unavailable})
want := unavailable.BuildPlan(root, DefaultPolicy()).DegradedNotice()
if want == "" {
t.Fatal("SETUP INVALID: an unavailable backend under the default policy is not degraded")
}
if got := engine.DegradedNotice(); got != want {
t.Errorf("engine notice = %q, want the policy command's %q", got, want)
}
disabled := DefaultPolicy()
disabled.Mode = ModeDisabled
off := NewEngine(EngineOptions{WorkspaceRoot: root, Policy: disabled, Backend: unavailable})
if got := off.DegradedNotice(); got != "" {
t.Errorf("a sandbox the user turned off still produced a notice: %q", got)
}
var none *Engine
if got := none.DegradedNotice(); got != "" {
t.Errorf("a nil engine produced a notice: %q", got)
}
}
// clearNestingMarkers keeps a test from inheriting the markers when the suite
// itself runs inside a Zero sandbox.
func clearNestingMarkers(t *testing.T) {
t.Helper()
t.Setenv(EnvSandboxed, "")
t.Setenv(EnvSandboxBackend, "")
}
// THE NESTING MARKERS SKIP WRAPPING, SO THEY HAVE TO BE SAID.
//
// With ZERO_SANDBOXED and ZERO_SANDBOX_BACKEND set, the runner passes every
// command through unwrapped on the word of an outer sandbox nothing verifies
// (#727). A native backend's plan knows nothing of that and reports full
// enforcement, so the session used to run every command unwrapped in silence.
func TestEngineDegradedNoticeSpeaksWhenTheNestingMarkersSkipWrapping(t *testing.T) {
root := t.TempDir()
native := Backend{Name: BackendLinuxBwrap, Platform: "linux", Available: true, NativeIsolation: true, CommandWrapping: true, Executable: "/usr/bin/bwrap"}
clearNestingMarkers(t)
if got := NewEngine(EngineOptions{WorkspaceRoot: root, Policy: DefaultPolicy(), Backend: native}).DegradedNotice(); got != "" {
t.Fatalf("SETUP INVALID: a native backend without the markers already gives a notice: %q", got)
}
t.Setenv(EnvSandboxed, "1")
t.Setenv(EnvSandboxBackend, string(BackendLinuxBwrap))
got := NewEngine(EngineOptions{WorkspaceRoot: root, Policy: DefaultPolicy(), Backend: native}).DegradedNotice()
if !strings.Contains(got, EnvSandboxed) || !strings.Contains(got, "cannot verify") {
t.Errorf("with the nesting markers set the notice = %q, want one naming the markers and the unverified outer sandbox", got)
}
disabled := DefaultPolicy()
disabled.Mode = ModeDisabled
if got := NewEngine(EngineOptions{WorkspaceRoot: root, Policy: disabled, Backend: native}).DegradedNotice(); got != "" {
t.Errorf("a sandbox the user turned off still produced a notice under the markers: %q", got)
}
}
// NO GRANT REACHES A DEGRADED PLAN THE SESSION'S NOTICE MISSED.
//
// The notice is decided once, from the policy the session starts with, and
// grants made during the session change the policy each command runs under.
// They cannot open a gap: a grant only opens the network or adds allow and
// deny paths, never the mode, and every mode but disabled builds a restricted
// filesystem profile that needs the platform sandbox from the start. So a
// degraded command plan means the notice already spoke, and a disabled sandbox
// stays disabled whatever is granted. A deny path the backend cannot enforce is
// refused at the command rather than run degraded. Raised by CodeRabbit.
func TestNoGrantReachesADegradedPlanTheSessionNoticeMissed(t *testing.T) {
clearNestingMarkers(t)
cacheRoot, tempRoot := t.TempDir(), t.TempDir()
original := sandboxUserCacheDir
sandboxUserCacheDir = func() (string, error) { return cacheRoot, nil }
t.Cleanup(func() { sandboxUserCacheDir = original })
t.Setenv("TMPDIR", tempRoot)
t.Setenv("TMP", tempRoot)
t.Setenv("TEMP", tempRoot)
root := t.TempDir()
extra := t.TempDir()
secret := filepath.Join(root, "secret.txt")
enabled := true
grants := []struct {
name string
profile RequestPermissionProfile
}{
{"network", RequestPermissionProfile{Network: &NetworkPermissions{Enabled: &enabled}}},
{"write root", RequestPermissionProfile{FileSystem: &FileSystemPermissions{Write: []string{extra}}}},
{"deny read", RequestPermissionProfile{FileSystem: &FileSystemPermissions{DenyRead: []string{secret}}}},
}
unavailable := Backend{Name: BackendUnavailable, Platform: "linux", Fallback: true, Message: "Linux sandbox helper is not available"}
degradedAndSaid, refused := 0, 0
for _, mode := range []PolicyMode{ModeEnforce, ModeDisabled} {
for _, network := range []NetworkMode{NetworkAllow, NetworkDeny} {
for _, grant := range grants {
for _, scope := range []PermissionGrantScope{PermissionGrantScopeSession, PermissionGrantScopeTurn} {
name := fmt.Sprintf("%s/%s/%s/%s", mode, network, grant.name, scope)
policy := DefaultPolicy()
policy.Mode, policy.Network = mode, network
engine := NewEngine(EngineOptions{WorkspaceRoot: root, Policy: policy, Backend: unavailable})
notice := engine.DegradedNotice()
undo, err := engine.GrantRequestPermissions(grant.profile, scope)
if err != nil {
t.Fatalf("%s: SETUP INVALID: grant: %v", name, err)
}
plan, err := engine.BuildCommandPlan(CommandSpec{Name: "/bin/sh", Args: []string{"-c", "true"}, Dir: root})
undo()
// The plan holds the runtime lease until it is cleaned up, and a lease
// file still open fails the temp directory removal on Windows.
plan.Cleanup()
if err != nil {
if grant.name != "deny read" || mode == ModeDisabled {
t.Errorf("%s: refused a command it should have planned: %v", name, err)
}
refused++
continue
}
if plan.EnforcementLevel == EnforcementDegraded {
if notice == "" {
t.Errorf("%s: the command plan is degraded but the session started with no notice", name)
}
degradedAndSaid++
}
if mode == ModeDisabled && plan.EnforcementLevel != EnforcementDisabled {
t.Errorf("%s: a grant took a disabled sandbox to %q", name, plan.EnforcementLevel)
}
}
}
}
}
// Not vacuous: the enforced modes did reach degraded plans, and the deny
// grants were refused rather than degraded.
if degradedAndSaid == 0 || refused == 0 {
t.Fatalf("SETUP INVALID: %d degraded plans and %d refusals, so the invariant was never exercised", degradedAndSaid, refused)
}
}
-27
View File
@@ -80,33 +80,6 @@ func (engine *Engine) CanPersistGrants() bool {
return engine != nil && engine.store != nil
}
// DegradedNotice is the line a session shows when this engine's sandbox cannot
// enforce what its policy asks for, or "" when it can or the sandbox is off.
//
// Built from the engine's own backend and policy through the same plan `zero
// sandbox policy` prints, so the notice and the command that explains it cannot
// describe two different sandboxes.
//
// EXCEPT WHERE THE RUNNER TAKES A DIFFERENT PATH. With the nesting markers set,
// BuildCommandPlan passes every command through unwrapped, trusting an outer
// sandbox to contain it, and the markers are ambient environment that nothing
// authenticates (#727). The plan above does not model that, so a session started
// with them ran every command unwrapped and said nothing. Reported by CodeRabbit.
func (engine *Engine) DegradedNotice() string {
if engine == nil || engine.policy.Mode == ModeDisabled {
return ""
}
if IsAlreadySandboxed() {
return nestedSandboxNotice
}
return engine.backend.BuildPlan(engine.workspaceRoot, engine.policy).DegradedNotice()
}
// nestedSandboxNotice says why a session with the nesting markers set does not
// wrap its commands, and that the outer sandbox it relies on is unverified.
const nestedSandboxNotice = "Sandbox enforcement is degraded: " + EnvSandboxed + " and " + EnvSandboxBackend +
" say this session already runs inside a sandbox, so it does not wrap commands again, and Zero cannot verify that outer sandbox."
func (engine *Engine) ConsumeGrantMigrationNotice() (string, error) {
if engine == nil || engine.store == nil {
return "", nil
-42
View File
@@ -246,62 +246,20 @@ func RedactResult(result Result) Result {
return redacted
}
// redactMetadata redacts EVERY text field of the metadata, not the subset a
// reader happens to think is user-facing. `zero search --json` prints this
// straight to stdout, and the fields carrying the most free text are the ones
// least likely to be on a hand-kept list: SpecUserComment and SpecRejectReason
// are review text the user typed, SpecDraftReasoning is the model's own prose,
// and Goal.Objective and Goal.StatusReason are whatever the goal was phrased
// as. A field added to sessions.Metadata later is covered by
// TestRedactMetadataCoversEveryTextField, which walks the struct rather than
// repeating this list, so the two cannot drift apart silently.
func redactMetadata(session sessions.Metadata, options redaction.Options) sessions.Metadata {
session.SessionID = redaction.RedactString(session.SessionID, options)
session.SessionKind = sessions.SessionKind(redaction.RedactString(string(session.SessionKind), options))
session.Title = redaction.RedactString(session.Title, options)
session.Cwd = redaction.RedactString(session.Cwd, options)
session.ModelID = redaction.RedactString(session.ModelID, options)
session.Provider = redaction.RedactString(session.Provider, options)
session.Tag = redaction.RedactString(session.Tag, options)
session.ParentSessionID = redaction.RedactString(session.ParentSessionID, options)
session.RootSessionID = redaction.RedactString(session.RootSessionID, options)
session.AgentName = redaction.RedactString(session.AgentName, options)
session.TaskID = redaction.RedactString(session.TaskID, options)
session.ForkedFromEventID = redaction.RedactString(session.ForkedFromEventID, options)
session.SpawnedFromEventID = redaction.RedactString(session.SpawnedFromEventID, options)
session.SpecID = redaction.RedactString(session.SpecID, options)
session.SpecFilePath = redaction.RedactString(session.SpecFilePath, options)
session.SpecStatus = sessions.SpecStatus(redaction.RedactString(string(session.SpecStatus), options))
session.SpecDraftModelID = redaction.RedactString(session.SpecDraftModelID, options)
session.SpecDraftReasoning = redaction.RedactString(session.SpecDraftReasoning, options)
session.SpecUserComment = redaction.RedactString(session.SpecUserComment, options)
session.SpecRejectReason = redaction.RedactString(session.SpecRejectReason, options)
session.SpecSourceSessionID = redaction.RedactString(session.SpecSourceSessionID, options)
session.SpecImplSessionID = redaction.RedactString(session.SpecImplSessionID, options)
session.CreatedAt = redaction.RedactString(session.CreatedAt, options)
session.UpdatedAt = redaction.RedactString(session.UpdatedAt, options)
session.LastEventType = sessions.EventType(redaction.RedactString(string(session.LastEventType), options))
session.Goal = redactGoal(session.Goal, options)
return session
}
// redactGoal returns a redacted COPY. Metadata is passed by value, but Goal is
// a pointer inside it, so redacting through it would rewrite the goal the
// caller still holds: a search would quietly edit the live session record it
// was only supposed to report on.
func redactGoal(goal *sessions.Goal, options redaction.Options) *sessions.Goal {
if goal == nil {
return nil
}
redacted := *goal
redacted.Objective = redaction.RedactString(redacted.Objective, options)
redacted.Status = sessions.GoalStatus(redaction.RedactString(string(redacted.Status), options))
redacted.StatusReason = redaction.RedactString(redacted.StatusReason, options)
redacted.CreatedAt = redaction.RedactString(redacted.CreatedAt, options)
redacted.UpdatedAt = redaction.RedactString(redacted.UpdatedAt, options)
return &redacted
}
func redactEventSummary(event EventSummary, options redaction.Options) EventSummary {
event.ID = redaction.RedactString(event.ID, options)
event.Type = sessions.EventType(redaction.RedactString(string(event.Type), options))
-156
View File
@@ -1,15 +1,12 @@
package search
import (
"fmt"
"os"
"path/filepath"
"reflect"
"strings"
"testing"
"time"
"github.com/Gitlawb/zero/internal/redaction"
"github.com/Gitlawb/zero/internal/sessions"
)
@@ -110,156 +107,3 @@ func fixedSearchClock(value string) func() time.Time {
}
return func() time.Time { return parsed }
}
// EVERY TEXT FIELD, NOT A HAND-KEPT SUBSET.
//
// `zero search --json` prints Hit.Session straight to stdout, so a field
// redactMetadata forgets is a field that reaches the terminal with whatever was
// in it. The list in redactMetadata is written out by hand; this walks the
// struct instead, so a string field added to sessions.Metadata (or to the Goal
// hanging off it) fails here on the day it lands rather than the day someone
// notices it in output. It named twenty uncovered fields when it was written,
// including the review text the user types and the model's own draft reasoning.
func TestRedactMetadataCoversEveryTextField(t *testing.T) {
const secret = "sk-ant-api03-aaaaaaaaaaaaaaaaaaaaaaaa0123456789ABCD"
// The premise: the secret is one RedactString already recognizes, so a field
// that still carries it was not passed through the redactor at all.
if redaction.RedactString(secret, redaction.Options{}) == secret {
t.Fatalf("SETUP INVALID: %q is not a shape RedactString redacts, so every leg below passes vacuously", secret)
}
var session sessions.Metadata
filled := fillMetadataText(t, reflect.ValueOf(&session).Elem(), secret, 0)
var leaked []string
var checked int
walkMetadataText(t, reflect.ValueOf(redactMetadata(session, redaction.Options{})), "", 0, func(name, value string) {
checked++
if strings.Contains(value, secret) {
leaked = append(leaked, name)
}
})
// Filling and walking are separate traversals of the same shape, so a field
// one of them understands and the other does not would otherwise show up as
// a quiet pass rather than a failure.
if checked != filled {
t.Fatalf("SETUP INVALID: filled %d text fields and walked %d, so the two traversals disagree about the shape of Metadata", filled, checked)
}
if filled == 0 {
t.Fatal("SETUP INVALID: no text fields were populated, so this proves nothing")
}
if len(leaked) > 0 {
t.Errorf("redactMetadata left the secret in %d of %d text fields: %s",
len(leaked), checked, strings.Join(leaked, ", "))
}
}
// fillMetadataText writes text into every text-bearing field reachable from
// value, allocating nil pointers and one-element slices on the way, and returns
// how many strings it set.
//
// Recursive on purpose. The first version of this walked only the direct string
// fields of one struct and hand-allocated the Goal next to it, which left the
// guard with the same blind spot it exists to close: a nested struct added to
// sessions.Metadata later would arrive empty, the walk would find nothing in it
// and the test would pass while the new field went unredacted. Raised by
// CodeRabbit on the PR.
//
// It fails rather than skips on a kind it does not understand, because a field
// this cannot populate is a field the guard is not guarding.
func fillMetadataText(t *testing.T, value reflect.Value, text string, depth int) int {
t.Helper()
if depth > 6 {
t.Fatalf("fillMetadataText recursed past depth %d, which suggests a cycle in the type", depth)
}
if !value.CanSet() {
return 0
}
switch value.Kind() {
case reflect.String:
value.SetString(text)
return 1
case reflect.Pointer:
if value.IsNil() {
value.Set(reflect.New(value.Type().Elem()))
}
return fillMetadataText(t, value.Elem(), text, depth+1)
case reflect.Struct:
filled := 0
for index := 0; index < value.NumField(); index++ {
if !value.Type().Field(index).IsExported() {
continue
}
filled += fillMetadataText(t, value.Field(index), text, depth+1)
}
return filled
case reflect.Slice:
value.Set(reflect.MakeSlice(value.Type(), 1, 1))
return fillMetadataText(t, value.Index(0), text, depth+1)
case reflect.Bool, reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64,
reflect.Float32, reflect.Float64:
return 0
default:
t.Fatalf("fillMetadataText does not know how to populate %s (kind %s); teach it and walkMetadataText together, or this guard stops covering that field",
value.Type(), value.Kind())
return 0
}
}
// walkMetadataText visits every text-bearing field reachable from value, in the
// same shape fillMetadataText populates. The two are deliberately symmetric and
// the caller compares their counts.
func walkMetadataText(t *testing.T, value reflect.Value, prefix string, depth int, visit func(name, value string)) {
t.Helper()
if depth > 6 {
t.Fatalf("walkMetadataText recursed past depth %d, which suggests a cycle in the type", depth)
}
switch value.Kind() {
case reflect.String:
visit(prefix, value.String())
case reflect.Pointer:
if !value.IsNil() {
walkMetadataText(t, value.Elem(), prefix, depth+1, visit)
}
case reflect.Struct:
for index := 0; index < value.NumField(); index++ {
field := value.Type().Field(index)
if !field.IsExported() {
continue
}
name := field.Name
if prefix != "" {
name = prefix + "." + name
}
walkMetadataText(t, value.Field(index), name, depth+1, visit)
}
case reflect.Slice:
for index := 0; index < value.Len(); index++ {
walkMetadataText(t, value.Index(index), fmt.Sprintf("%s[%d]", prefix, index), depth+1, visit)
}
}
}
// Metadata is passed by value, but Goal is a pointer inside it. Redacting
// through that pointer would edit the session the caller still holds, so a
// search would quietly rewrite the live record it was asked to report on.
func TestRedactMetadataDoesNotRewriteTheCallersGoal(t *testing.T) {
const secret = "sk-ant-api03-aaaaaaaaaaaaaaaaaaaaaaaa0123456789ABCD"
goal := &sessions.Goal{Objective: "ship " + secret, StatusReason: "blocked on " + secret}
session := sessions.Metadata{SessionID: "s", Goal: goal}
redacted := redactMetadata(session, redaction.Options{})
if goal.Objective != "ship "+secret || goal.StatusReason != "blocked on "+secret {
t.Errorf("redactMetadata rewrote the caller's goal: %+v", *goal)
}
if redacted.Goal == goal {
t.Error("redacted metadata still points at the caller's goal")
}
if strings.Contains(redacted.Goal.Objective, secret) || strings.Contains(redacted.Goal.StatusReason, secret) {
t.Errorf("goal text was not redacted: %+v", *redacted.Goal)
}
}
-7
View File
@@ -218,13 +218,6 @@ func btwCommandUnavailable(command parsedCommand) bool {
return arg != "" && arg != "status"
case commandTheme:
return arg != "list"
case commandNotify:
// Mutations reconfigure the shared notifier pointer and write the
// global preference; inside a BTW side conversation that would mutate
// the hidden parent's live policy and the user's global file while
// only the side surface's display fields change. Read-only list stays
// available (maintainer review, PR #1001).
return arg != "list"
case commandConfig:
return arg != ""
default:
-62
View File
@@ -1,7 +1,6 @@
package tui
import (
"bytes"
"context"
"os"
"path/filepath"
@@ -10,8 +9,6 @@ import (
tea "charm.land/bubbletea/v2"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/notify"
"github.com/Gitlawb/zero/internal/sessions"
)
@@ -244,12 +241,6 @@ func TestBTWBlocksPersistentConfigurationCommands(t *testing.T) {
"/mcp",
"/rewind",
"/compact",
// Notify mutations reconfigure the parent's notifier (shared pointer)
// and write the global preference; they must stay blocked in a BTW
// side conversation (maintainer review, PR #1001).
"/notify",
"/notify off",
"/notify bell always",
} {
t.Run(input, func(t *testing.T) {
m := newBTWTestModel(t)
@@ -274,7 +265,6 @@ func TestBTWAllowsReadOnlyConfigurationCommands(t *testing.T) {
"/profile status",
"/theme list",
"/config",
"/notify list",
} {
t.Run(input, func(t *testing.T) {
m := newBTWTestModel(t)
@@ -288,58 +278,6 @@ func TestBTWAllowsReadOnlyConfigurationCommands(t *testing.T) {
}
}
// Maintainer regression (PR #1001): running notify mutations inside a BTW
// side conversation must not touch the PARENT's live notifier or the stored
// preference. The side surface shallow-copies the parent model and shares its
// notifier pointer, so an unblocked mutation would reconfigure the parent and
// write the global file while only the side's display fields change; after
// returning, the parent could report one policy in /notify list while emitting
// another. Checks both the parent's actual notification output and the stored
// preference, per the review.
func TestBTWNotifyMutationsLeaveParentAndStoredPreferenceIntact(t *testing.T) {
cfgPath := filepath.Join(t.TempDir(), "config.json")
if err := os.WriteFile(cfgPath, []byte(`{"notify":{"mode":"bell","focusMode":"always"}}`), 0o600); err != nil {
t.Fatal(err)
}
parent := newBTWTestModel(t)
parent.userConfigPath = cfgPath
parent.notifyMode = "bell"
parent.notifyFocusMode = "always"
var buf bytes.Buffer
parent.notifier = notify.New(&buf, notify.Config{Mode: notify.ModeBell, FocusMode: notify.FocusAlways})
parent.notifier.SetFocused(true)
// Enter the side conversation and attempt a mutation.
side, _ := parent.handleBTWCommand("")
updated, _ := side.dispatchCommand(parseCommand("/notify off"))
got := updated.(model)
if got.picker != nil {
t.Fatal("/notify inside BTW must not open the picker")
}
if !transcriptContains(got.transcript, "unavailable in a BTW conversation") {
t.Fatal("expected the BTW blocked-command guidance")
}
// Return to the parent: the live policy still bells (actual output), and
// the stored preference is unchanged.
returned, _ := got.handleBTWCommand("")
parent = returned
parent.notifier.Notify(notify.Completion, "x")
if buf.String() != "\x07" {
t.Fatalf("parent notifier was reconfigured by the side conversation: got %q, want the original bell", buf.String())
}
stored, err := config.UserNotify(cfgPath)
if err != nil {
t.Fatalf("read stored notify: %v", err)
}
if stored.Mode != "bell" || stored.FocusMode != "always" {
t.Fatalf("stored preference changed from inside BTW: %+v", stored)
}
if parent.notifyMode != "bell" || parent.notifyFocusMode != "always" {
t.Fatalf("parent live policy changed from inside BTW: %q/%q", parent.notifyMode, parent.notifyFocusMode)
}
}
func TestBTWExitBlockedWhileParentRunActive(t *testing.T) {
m := newBTWTestModel(t)
m.pending = true
-8
View File
@@ -37,7 +37,6 @@ const (
commandFast
commandStyle
commandTheme
commandNotify
commandTranscript
commandBash
commandImage
@@ -409,13 +408,6 @@ var commandDefinitions = []commandDefinition{
description: "Pick a color theme (no arg opens the picker; auto detects the terminal background).",
kind: commandTheme,
},
{
name: "/notify",
usage: "/notify [list|off|bell|notify|both [unfocused|always|focused]]",
group: commandGroupSession,
description: "Pick when Zero alerts you (completion and needs-input). No arg opens the picker.",
kind: commandNotify,
},
{
name: "/exit",
aliases: []string{"/quit"},
+17
View File
@@ -332,6 +332,23 @@ func TestDiffViewerFallsBackForOversizedSourceLine(t *testing.T) {
}
}
func TestDiffViewerSyntaxHighlighterKeepsSpans(t *testing.T) {
requireDiffViewerTrueColor(t)
styled, ok := highlightCodeForPathWithSpans(
[]string{"func greet(name string) string { return \"new\" }"},
"example.go",
200,
zeroTheme.addLine.GetBackground(),
[]highlightSpan{{line: 0, start: 39, end: 44, background: zeroTheme.addLineWord.GetBackground()}},
)
if !ok || len(styled) != 1 {
t.Fatalf("highlightCodeForPathWithSpans = %#v, %v", styled, ok)
}
if !strings.Contains(styled[0], "46;101;77") {
t.Fatalf("highlighted span lost its word-diff background: %q", styled[0])
}
}
func TestDiffViewerHighlightsEachFileWithItsOwnPath(t *testing.T) {
requireDiffViewerTrueColor(t)
rawLines := []string{
+25 -29
View File
@@ -14,7 +14,6 @@ import (
"github.com/Gitlawb/zero/internal/imageinput"
"github.com/Gitlawb/zero/internal/modelregistry"
"github.com/Gitlawb/zero/internal/providermodeldiscovery"
"github.com/Gitlawb/zero/internal/terminalpet"
"github.com/Gitlawb/zero/internal/zeroruntime"
_ "golang.org/x/image/webp"
@@ -103,44 +102,41 @@ func stripMatchingQuotes(s string) (string, bool) {
// fetched it) — this carries InputModalities from models.dev, which
// includes "image" for vision-capable models
// 3. Falls back to the name heuristic for unknown models
func (m model) modelSupportsVisionFor(modelID string) bool {
trimmed := strings.TrimSpace(modelID)
func (m model) modelSupportsVisionTUI() bool {
trimmed := strings.TrimSpace(m.modelName)
if trimmed == "" {
return false
}
// Check the discovered model list for the ACTIVE provider route first.
if models, ok := m.discoveredModelsForActiveRoute(); ok {
if supported, ok := discoveredVisionSupport(models, trimmed); ok {
return supported
}
}
// The curated catalog is authoritative when active-provider discovery is absent or inconclusive.
// The curated catalog is authoritative only when it knows the model.
if entry, known := m.modelCatalog.Resolve(trimmed); known {
return entry.Supports(modelregistry.ModelCapabilityVision)
}
// Fall back to curated catalog or the name heuristic.
return modelregistry.SupportsVision(m.modelCatalog, trimmed)
}
func discoveredVisionSupport(models []providermodeldiscovery.Model, modelID string) (bool, bool) {
for _, dm := range models {
if strings.EqualFold(strings.TrimSpace(dm.ID), modelID) {
if len(dm.InputModalities) == 0 {
return false, false
}
for _, modality := range dm.InputModalities {
if strings.EqualFold(strings.TrimSpace(modality), "image") {
return true, true
// Check the discovered model list (from models.dev) for InputModalities
// containing "image". This covers custom/ollama/cloud models not in the
// curated catalog — models.dev knows their capabilities.
for _, models := range m.modelPickerLiveByProvider {
for _, dm := range models {
if strings.EqualFold(strings.TrimSpace(dm.ID), trimmed) {
// A provider's authenticated listing may only establish which models
// are available, without repeating modalities. Treat an empty list as
// unknown rather than as an explicit image-input denial, so the
// curated registry/name capability fallback remains available while
// models.dev metadata is temporarily unavailable.
if len(dm.InputModalities) == 0 {
continue
}
for _, modality := range dm.InputModalities {
if strings.EqualFold(strings.TrimSpace(modality), "image") {
return true
}
}
return false // found the model in discovered list, no image modality
}
return false, true // found model with explicit modalities, no image modality
}
}
return false, false
}
func (m model) modelSupportsVisionTUI() bool {
return m.modelSupportsVisionFor(m.modelName)
// Fall back to the name heuristic for models not in the catalog or
// discovered list.
return modelregistry.VisionCapableByName(trimmed)
}
// attachClipboardImage attaches an image read from the OS clipboard (a
-357
View File
@@ -10,10 +10,6 @@ import (
"strings"
"testing"
tea "charm.land/bubbletea/v2"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/providercatalog"
"github.com/Gitlawb/zero/internal/providermodeldiscovery"
"github.com/Gitlawb/zero/internal/tools"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
@@ -576,356 +572,3 @@ func TestRetryResendsAttachments(t *testing.T) {
t.Fatalf("retried prompt should include the remembered user text, got:\n%s", last.Content)
}
}
func TestModelSupportsVision_ActiveProviderPrecedence(t *testing.T) {
m := newModel(t.Context(), Options{ModelName: "gpt-4.1", ProviderName: "openai"})
// gpt-4.1 is in curated catalog as supporting vision.
// Override active provider to explicitly report it as text-only.
m.modelPickerLiveByProvider = map[string][]providermodeldiscovery.Model{
"openai": {
{
ID: "gpt-4.1",
InputModalities: []string{"text"},
},
},
"other-provider": {
{
ID: "custom-text-model",
InputModalities: []string{"image", "text"},
},
},
}
// Active provider explicit modalities must win over catalog.
if m.modelSupportsVisionFor("gpt-4.1") {
t.Fatal("expected active provider text-only explicit modality to override catalog vision support")
}
// Foreign provider discovery record must NOT decide capabilities for active provider.
if m.modelSupportsVisionFor("custom-text-model") {
t.Fatal("expected foreign provider discovery record to be ignored for active provider")
}
}
func TestRunAgentWithOptions_DiscoverySnapshotRace(t *testing.T) {
ctx, cancel := context.WithCancel(t.Context())
defer cancel()
m := newModel(ctx, Options{ModelName: "gpt-4.1", ProviderName: "openai"})
m.modelPickerLiveByProvider = map[string][]providermodeldiscovery.Model{
"openai": {{ID: "gpt-4.1", InputModalities: []string{"text"}}},
}
// Schedule the command on the main thread (synchronously snapshotting discovered models)
_ = m.runAgentWithOptions(1, ctx, "hello", nil, tuiAgentRunOptions{})
// Concurrently simulate discovery resolution via applyModelPickerModelsDiscovered
done := make(chan struct{})
go func() {
defer close(done)
localM := m
for i := 0; i < 200; i++ {
localM = localM.applyModelPickerModelsDiscovered(modelPickerModelsDiscoveredMsg{
providerID: fmt.Sprintf("provider-%d", i%10),
models: []providermodeldiscovery.Model{
{ID: "discovered-model", InputModalities: []string{"image"}},
},
})
}
}()
for i := 0; i < 50; i++ {
cmd := m.runAgentWithOptions(i+2, ctx, "hello", nil, tuiAgentRunOptions{})
if cmd == nil {
t.Fatal("expected non-nil cmd")
}
}
<-done
}
func TestModelSupportsVision_CustomEndpointScoping(t *testing.T) {
desc, ok := providercatalog.Get("custom-openai-compatible")
if !ok {
t.Fatal("custom-openai-compatible descriptor missing from catalog")
}
profileA := config.ProviderProfile{
Name: "custom-endpoint-a",
CatalogID: desc.ID,
BaseURL: "https://a.example.com/v1",
}
profileB := config.ProviderProfile{
Name: "custom-endpoint-b",
CatalogID: desc.ID,
BaseURL: "https://b.example.com/v1",
}
// Discovery returns gpt-4.1 as text-only on endpoint A, but multimodal on endpoint B.
endpointAModels := []providermodeldiscovery.Model{
{ID: "gpt-4.1", InputModalities: []string{"text"}},
{ID: "endpoint-a-model", InputModalities: []string{"text"}},
}
endpointBModels := []providermodeldiscovery.Model{
{ID: "gpt-4.1", InputModalities: []string{"text", "image"}},
{ID: "endpoint-b-model", InputModalities: []string{"text", "image"}},
}
provider := &fakeProvider{events: []zeroruntime.StreamEvent{{Type: zeroruntime.StreamEventDone}}}
// Case 1: Endpoint B is active.
m := newModel(t.Context(), Options{
ProviderName: profileB.Name,
ModelName: "gpt-4.1",
ProviderProfile: profileB,
Provider: provider,
SavedProviders: []config.ProviderProfile{profileA, profileB},
})
// Simulate discovery delivery for both endpoints.
m = m.applyModelPickerModelsDiscovered(modelPickerModelsDiscoveredMsg{
providerID: desc.ID,
endpointKey: providerEndpointKey(profileA, desc),
models: endpointAModels,
})
m = m.applyModelPickerModelsDiscovered(modelPickerModelsDiscoveredMsg{
providerID: desc.ID,
endpointKey: providerEndpointKey(profileB, desc),
models: endpointBModels,
})
// Synchronous gate on active endpoint B must see B's multimodal discovery,
// ignoring endpoint A's text-only result.
if !m.modelSupportsVisionFor("gpt-4.1") {
t.Fatal("expected gpt-4.1 to support vision on endpoint B")
}
if !m.modelSupportsVisionFor("endpoint-b-model") {
t.Fatal("expected endpoint-b-model to support vision on endpoint B")
}
if m.modelSupportsVisionFor("endpoint-a-model") {
t.Fatal("endpoint A model must not be recognized on active endpoint B")
}
// Assert options.SupportsVision wired in runAgentWithOptions also reports true for endpoint B.
var capturedVision func(string) bool
m.captureRunSupportsVision = func(fn func(string) bool) {
capturedVision = fn
}
_ = execCmd(m.runAgentWithOptions(1, t.Context(), "test", nil, tuiAgentRunOptions{}))
if capturedVision == nil {
t.Fatal("expected captureRunSupportsVision to be called")
}
if !capturedVision("gpt-4.1") {
t.Fatal("expected run agent options.SupportsVision to report true for gpt-4.1 on endpoint B")
}
if !capturedVision("endpoint-b-model") {
t.Fatal("expected run agent options.SupportsVision to report true for endpoint-b-model on endpoint B")
}
if capturedVision("endpoint-a-model") {
t.Fatal("run agent options.SupportsVision must not inherit endpoint A model capabilities")
}
// Case 2: Switch to Endpoint A.
m.providerProfile = profileA
m.providerName = profileA.Name
// Synchronous gate on active endpoint A must restrict gpt-4.1 to text-only.
if m.modelSupportsVisionFor("gpt-4.1") {
t.Fatal("expected gpt-4.1 to be text-only on endpoint A")
}
if m.modelSupportsVisionFor("endpoint-b-model") {
t.Fatal("endpoint B model must not be recognized on active endpoint A")
}
capturedVision = nil
_ = execCmd(m.runAgentWithOptions(2, t.Context(), "test", nil, tuiAgentRunOptions{}))
if capturedVision == nil {
t.Fatal("expected captureRunSupportsVision to be called")
}
if capturedVision("gpt-4.1") {
t.Fatal("expected run agent options.SupportsVision to report false for gpt-4.1 on endpoint A")
}
}
func TestModelSupportsVision_UnnamedEndpointScoping(t *testing.T) {
desc, ok := providercatalog.Get("custom-openai-compatible")
if !ok {
t.Fatal("custom-openai-compatible descriptor missing from catalog")
}
profile1 := config.ProviderProfile{
CatalogID: desc.ID,
BaseURL: "https://host-one.example.com/v1",
}
profile2 := config.ProviderProfile{
CatalogID: desc.ID,
BaseURL: "https://host-two.example.com/v1",
}
key1 := providerEndpointKey(profile1, desc)
key2 := providerEndpointKey(profile2, desc)
if key1 == key2 || key1 == "" || key2 == "" {
t.Fatalf("unnamed custom profiles must have distinct non-empty endpoint keys, got %q vs %q", key1, key2)
}
m := newModel(t.Context(), Options{
ModelName: "gpt-4.1",
ProviderProfile: profile1,
})
m.modelPickerLiveByProvider = map[string][]providermodeldiscovery.Model{
key1: {{ID: "unnamed-model-1", InputModalities: []string{"image"}}},
key2: {{ID: "unnamed-model-2", InputModalities: []string{"image"}}},
}
if !m.modelSupportsVisionFor("unnamed-model-1") {
t.Fatal("active endpoint 1 should support unnamed-model-1")
}
if m.modelSupportsVisionFor("unnamed-model-2") {
t.Fatal("active endpoint 1 should NOT support unnamed-model-2 from endpoint 2")
}
}
func TestModelPickerDiscoveryCmds_OrderIndependence(t *testing.T) {
desc, ok := providercatalog.Get("custom-openai-compatible")
if !ok {
t.Fatal("custom-openai-compatible descriptor missing from catalog")
}
profileA := config.ProviderProfile{
Name: "custom-a",
CatalogID: desc.ID,
BaseURL: "https://a.example.com/v1",
}
profileB := config.ProviderProfile{
Name: "custom-b",
CatalogID: desc.ID,
BaseURL: "https://b.example.com/v1",
}
for _, order := range [][]config.ProviderProfile{
{profileA, profileB},
{profileB, profileA},
} {
m := newModel(t.Context(), Options{
SavedProviders: order,
})
cmd := m.modelPickerDiscoveryCmds()
if cmd == nil {
t.Fatalf("expected discovery commands for order %v", order)
}
msg := cmd()
batch, ok := msg.(tea.BatchMsg)
if !ok {
t.Fatalf("expected tea.BatchMsg, got %T", msg)
}
if len(batch) != 2 {
t.Fatalf("expected 2 discovery commands (one per custom endpoint), got %d", len(batch))
}
}
}
func TestModelSupportsVision_DelayedDiscoveryAfterProfileSwitch(t *testing.T) {
desc, ok := providercatalog.Get("custom-openai-compatible")
if !ok {
t.Fatal("custom-openai-compatible descriptor missing from catalog")
}
profileA := config.ProviderProfile{
Name: "endpoint-a",
CatalogID: desc.ID,
BaseURL: "https://a.example.com/v1",
}
profileB := config.ProviderProfile{
Name: "endpoint-b",
CatalogID: desc.ID,
BaseURL: "https://b.example.com/v1",
}
// Start with endpoint A active.
m := newModel(t.Context(), Options{
ProviderName: profileA.Name,
ModelName: "gpt-4.1",
ProviderProfile: profileA,
SavedProviders: []config.ProviderProfile{profileA, profileB},
})
// User switches active provider to endpoint B before discovery completes.
m.providerProfile = profileB
m.providerName = profileB.Name
// Now delayed discovery for endpoint A arrives with text-only gpt-4.1.
m = m.applyModelPickerModelsDiscovered(modelPickerModelsDiscoveredMsg{
providerID: desc.ID,
endpointKey: providerEndpointKey(profileA, desc),
models: []providermodeldiscovery.Model{
{ID: "gpt-4.1", InputModalities: []string{"text"}},
},
})
// Active endpoint B has not discovered yet, so it should fall back to catalog authority
// (where gpt-4.1 supports vision), NOT get contaminated by endpoint A's text-only result.
if !m.modelSupportsVisionFor("gpt-4.1") {
t.Fatal("expected active endpoint B to fall back to catalog vision support, not contaminated by endpoint A")
}
// Generic descriptor key must NOT have been populated with endpoint A's discovery.
if _, genericFound := m.modelPickerLiveByProvider[desc.ID]; genericFound {
t.Fatalf("generic provider ID %q must not be stored in discovery map when endpointKey is set", desc.ID)
}
}
func TestRunAgentWithOptions_SnapshotIsolation(t *testing.T) {
desc, ok := providercatalog.Get("custom-openai-compatible")
if !ok {
t.Fatal("custom-openai-compatible descriptor missing from catalog")
}
profileB := config.ProviderProfile{
Name: "endpoint-b",
CatalogID: desc.ID,
BaseURL: "https://b.example.com/v1",
}
provider := &fakeProvider{events: []zeroruntime.StreamEvent{{Type: zeroruntime.StreamEventDone}}}
m := newModel(t.Context(), Options{
ProviderName: profileB.Name,
ModelName: "gpt-4.1",
ProviderProfile: profileB,
Provider: provider,
})
// Launch run 1 when endpoint B has NO discovery yet.
var capturedVisionRun1 func(string) bool
m.captureRunSupportsVision = func(fn func(string) bool) {
capturedVisionRun1 = fn
}
_ = execCmd(m.runAgentWithOptions(1, t.Context(), "run1", nil, tuiAgentRunOptions{}))
if capturedVisionRun1 == nil {
t.Fatal("expected captureRunSupportsVision for run 1")
}
// Before discovery: catalog authority confirms gpt-4.1 vision support.
if !capturedVisionRun1("gpt-4.1") {
t.Fatal("run 1 should support vision via catalog fallback")
}
// Discovery arrives for endpoint B reporting gpt-4.1 as text-only.
m = m.applyModelPickerModelsDiscovered(modelPickerModelsDiscoveredMsg{
providerID: desc.ID,
endpointKey: providerEndpointKey(profileB, desc),
models: []providermodeldiscovery.Model{
{ID: "gpt-4.1", InputModalities: []string{"text"}},
},
})
// Newly launched run 2 must see text-only discovery.
var capturedVisionRun2 func(string) bool
m.captureRunSupportsVision = func(fn func(string) bool) {
capturedVisionRun2 = fn
}
_ = execCmd(m.runAgentWithOptions(2, t.Context(), "run2", nil, tuiAgentRunOptions{}))
if capturedVisionRun2 == nil {
t.Fatal("expected captureRunSupportsVision for run 2")
}
if capturedVisionRun2("gpt-4.1") {
t.Fatal("run 2 must report false for gpt-4.1 due to live text-only discovery")
}
// CRITICAL: Run 1's captured predicate MUST STILL report true (snapshot isolation!).
if !capturedVisionRun1("gpt-4.1") {
t.Fatal("run 1 must retain its snapshot and still report true for gpt-4.1")
}
}
+17 -122
View File
@@ -73,7 +73,6 @@ type model struct {
ctx context.Context
cwd string
appVersion string
startupNotices []string // Options.StartupNotices, said again when a session is resumed
userCommands []usercommands.Command // file-sourced /commands (.zero/commands)
loadSkills func() []skills.Skill // lazy installed-skills loader for /skills + /<skill-name>
userConfigPath string
@@ -216,26 +215,16 @@ type model struct {
keyBindings keyBindings
themeMode themeMode // palette preference: system (default) or named palette
hasDarkBg bool // last terminal background-detection result, if one is delivered
// notifyMode and notifyFocusMode track the user's in-session notify
// preference (from options.Notify, updated by /notify). The notifier built
// at startup is independent, so changes apply on the NEXT permission prompt.
notifyMode string
// notifyConfiguredMode holds the RAW configured mode ("" when the user
// never set one) so state output can label the effective default as
// "(default)" — mirroring the CLI's `mode: (default)` for the same state
// (maintainer review, PR #1001). Cleared once the user picks a mode.
notifyConfiguredMode string
notifyFocusMode string
userAgent string
compactRequests int
compactInFlight bool
compactFrame int
lastCompactResult *CompactResult
lastCompactError string
unpricedRequests int
unpricedTokens int
lastUsage usage.Normalized
lastUsageSeen bool
userAgent string
compactRequests int
compactInFlight bool
compactFrame int
lastCompactResult *CompactResult
lastCompactError string
unpricedRequests int
unpricedTokens int
lastUsage usage.Normalized
lastUsageSeen bool
// turnLatencySum / turnLatencyCount accumulate completed-run wall time so
// /context can show a rolling average turn latency (the "is it slow?" signal).
// Reset by /new.
@@ -600,11 +589,6 @@ type model struct {
// with. Nil in production; used by tests to assert image threading without a
// real provider round-trip.
captureRunImages func([]zeroruntime.ImageBlock)
// captureRunSupportsVision, when set, is invoked with the vision-support
// predicate wired into the run's agent options. Nil in production; used by
// tests to assert snapshot isolation and discovery scoping.
captureRunSupportsVision func(func(string) bool)
}
type agentTextMsg struct {
@@ -914,22 +898,6 @@ type tuiAgentRunOptions struct {
specDraft bool
}
// effectiveTUINotifyMode returns the notification mode the TUI should use. An
// empty/unconfigured mode falls back to the TUI's own effective default
// ("both": terminal bell + OSC-9 desktop notification) so the needs-input
// alert works for new users without requiring them to hand-edit config.json.
// The default lives HERE, deliberately not in config.Resolve: the resolver is
// shared with headless `zero exec`, which must stay byte-silent when
// unconfigured (maintainer review, PR #1001). The /notify command persists
// explicit choices; blank stays blank on disk.
func effectiveTUINotifyMode(mode string) notify.Mode {
m := notify.Mode(strings.TrimSpace(mode))
if m == "" {
return notify.ModeBoth
}
return m
}
func newModel(ctx context.Context, options Options) model {
if ctx == nil {
ctx = context.Background()
@@ -1001,7 +969,7 @@ func newModel(ctx context.Context, options Options) model {
runSpinner.Spinner.FPS = activeAnimationFrameInterval
notifier := notify.New(os.Stderr, notify.Config{
Mode: effectiveTUINotifyMode(options.Notify.Mode),
Mode: notify.Mode(strings.TrimSpace(options.Notify.Mode)),
FocusMode: notify.FocusMode(strings.TrimSpace(options.Notify.FocusMode)),
})
// Opt-in webhook fan-out (ZERO_NOTIFY_WEBHOOK_URL). Delivery failures stay
@@ -1016,7 +984,6 @@ func newModel(ctx context.Context, options Options) model {
ctx: ctx,
cwd: cwd,
appVersion: strings.TrimSpace(options.Version),
startupNotices: options.StartupNotices,
swarmDoneAt: map[string]time.Time{},
userCommands: loadedUserCommands,
loadSkills: options.LoadSkills,
@@ -1065,9 +1032,6 @@ func newModel(ctx context.Context, options Options) model {
keyBindings: resolvedKeyBindings,
themeMode: resolveThemeMode(options.Theme, os.Getenv("ZERO_THEME"), options.SavedTheme),
hasDarkBg: true,
notifyMode: string(effectiveTUINotifyMode(options.Notify.Mode)),
notifyConfiguredMode: strings.TrimSpace(options.Notify.Mode),
notifyFocusMode: strings.TrimSpace(options.Notify.FocusMode),
userAgent: options.UserAgent,
usageTracker: usageTracker,
transcript: initialTranscript(),
@@ -1131,28 +1095,9 @@ func newModel(ctx context.Context, options Options) model {
for _, warning := range keyBindingWarnings {
m = m.appendSystemNotice(warning)
}
m.transcript = m.withStartupNotices(m.transcript)
return m
}
// withStartupNotices appends the session's startup notices to rows as system
// rows, in order, skipping blank ones so a caller can pass one unconditionally.
//
// EVERY SESSION THIS PROCESS OPENS IS TOLD. The notices describe the sandbox
// every session here runs under, so they open the first session, and they come
// back when /new or /resume rebuilds the transcript for another one; both
// started from an empty transcript and dropped them. Rewind, compaction and
// /clear stay in the same session and keep the notice in the scrollback above
// their divider. Reported by CodeRabbit.
func (m model) withStartupNotices(rows []transcriptRow) []transcriptRow {
for _, notice := range m.startupNotices {
if strings.TrimSpace(notice) != "" {
rows = appendRow(rows, rowSystem, notice)
}
}
return rows
}
func (m model) doctorOptions(connectivity bool) doctor.Options {
var health *providerhealth.Result
if connectivity && m.probeProviderHealth != nil && config.HasProviderProfile(m.providerProfile) {
@@ -4586,14 +4531,15 @@ func (m model) choosePicker() (tea.Model, tea.Cmd) {
previousProvider, previousModel := m.providerName, m.modelName
text := ""
var switchPersistErr error
// The row's owner label is derived from the same decision, so it always
// names the provider used here.
if owner := m.modelPickerSwitchOwner(item); owner != "" {
owner := strings.TrimSpace(item.OwnerProvider)
_, ownerIsSavedProvider := m.savedProviderByName(owner)
if owner != "" && !strings.EqualFold(owner, strings.TrimSpace(m.providerName)) && ownerIsSavedProvider {
// A model from another saved provider: switch provider + model together.
m, text, _, cmd, switchPersistErr = m.switchProviderModel(owner, item.Value)
} else {
// Apply against the active provider instead of attempting an
// unresolvable provider switch.
// OwnerProvider is blank, matches the active provider, or (registry-fallback
// / stale-history rows) doesn't resolve to any saved provider: apply against
// the active provider instead of attempting an unresolvable provider switch.
m, text, switchPersistErr = m.handleModelCommand(item.Value)
}
if m.providerName != previousProvider || m.modelName != previousModel {
@@ -4646,14 +4592,6 @@ func (m model) choosePicker() (tea.Model, tea.Cmd) {
return m.showTransientNotice(m.themeAppliedNotice(), transientNoticeSuccess)
}
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: text})
case pickerNotify:
// The picker item's Value is "<mode> <focus>"; reusing the text handler
// keeps validation, persistence, and the user-facing message in one
// place. There is no live preview for notify, so no follow-up command
// is needed here.
text := ""
m, text = m.handleNotifyCommand(item.Value)
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: text})
}
return m, cmd
}
@@ -5068,18 +5006,6 @@ func (m model) dispatchCommand(command parsedCommand) (tea.Model, tea.Cmd) {
}
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: text})
return m, nil
case commandNotify:
// Bare `/notify` opens the popup picker so the user can pick mode + focus
// with arrow keys, matching /model and /theme. An explicit
// `/notify off|bell|notify|both [unfocused|always]` runs the text handler.
if strings.TrimSpace(command.text) == "" {
m.picker = m.newNotifyPicker()
return m, nil
}
text := ""
m, text = m.handleNotifyCommand(command.text)
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: text})
return m, nil
case commandImage:
m = m.handleImageCommand(command.text)
return m, nil
@@ -5555,19 +5481,6 @@ func selfCorrectAutonomyForMode(mode agent.PermissionMode) string {
}
func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt string, images []zeroruntime.ImageBlock, runOptions tuiAgentRunOptions) tea.Cmd {
var activeDiscovered []providermodeldiscovery.Model
if models, ok := m.discoveredModelsForActiveRoute(); ok {
activeDiscovered = make([]providermodeldiscovery.Model, len(models))
for i, dm := range models {
copied := dm
if len(dm.InputModalities) > 0 {
copied.InputModalities = append([]string{}, dm.InputModalities...)
}
activeDiscovered[i] = copied
}
}
catalog := m.modelCatalog
return func() tea.Msg {
started := m.now()
if m.turnTimer != nil {
@@ -5683,24 +5596,6 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
func(modelID string) { usageModelID = modelID },
)
}
options.SupportsVision = func(modelID string) bool {
trimmed := strings.TrimSpace(modelID)
if trimmed == "" {
return false
}
if len(activeDiscovered) > 0 {
if supported, ok := discoveredVisionSupport(activeDiscovered, trimmed); ok {
return supported
}
}
if entry, known := catalog.Resolve(trimmed); known {
return entry.Supports(modelregistry.ModelCapabilityVision)
}
return modelregistry.SupportsVision(catalog, trimmed)
}
if m.captureRunSupportsVision != nil {
m.captureRunSupportsVision(options.SupportsVision)
}
// Post-edit self-correction is on by default in the TUI but kept FAST: it
// runs LSP diagnostics over the changed files only — cheap, change-scoped,
-52
View File
@@ -5,7 +5,6 @@ import (
"context"
"encoding/json"
"errors"
"os"
"path/filepath"
"reflect"
"strings"
@@ -3201,31 +3200,6 @@ func TestScrimViewportLine(t *testing.T) {
}
}
func TestEffectiveTUINotifyMode(t *testing.T) {
cases := []struct {
in string
want notify.Mode
}{
// Empty input falls through to the TUI's own effective default
// ("both": bell + OSC-9 desktop notification) so the needs-input
// alert works for users who never configured notify. The default
// lives here, NOT in config.Resolve — headless runs stay silent
// when unconfigured (maintainer review, PR #1001).
{"", notify.ModeBoth},
{" ", notify.ModeBoth},
{"off", notify.ModeOff},
{"bell", notify.ModeBell},
{"notify", notify.ModeNotify},
{"both", notify.ModeBoth},
{" bell ", notify.ModeBell},
}
for _, c := range cases {
if got := effectiveTUINotifyMode(c.in); got != c.want {
t.Errorf("effectiveTUINotifyMode(%q) = %q, want %q", c.in, got, c.want)
}
}
}
func TestOverlayViewportLinesCompositesAndPreservesBackdropText(t *testing.T) {
width := 40
lines := make([]string, 9)
@@ -3395,29 +3369,3 @@ func TestMultilineBurstSubmits(t *testing.T) {
t.Fatal("multiline burst got nil cmd")
}
}
// Maintainer regression (PR #1001, rebase review): main's pickerTheme case ends
// by appending the handler's text to the transcript, so a theme choice whose
// preference cannot be saved still reports "could not save theme preference
// (...)". A rebase accidentally dropped that line, leaving a failed save
// completely silent when the theme was picked from the picker (the /theme text
// path kept its line, so no test caught it). This pins the picker's failure
// path: the config path sits under a regular file so the write must fail.
func TestThemePickerAppendNoteOnFailedSave(t *testing.T) {
defer applyTheme(themeDark, true)
cfg := filepath.Join(t.TempDir(), "blocked", "zero.json")
if err := os.WriteFile(filepath.Dir(cfg), []byte("not a directory"), 0o600); err != nil {
t.Fatal(err)
}
m := newModel(context.Background(), Options{UserConfigPath: cfg})
m.picker = &commandPicker{kind: pickerTheme, items: []pickerItem{{Label: "dracula", Value: "dracula"}}}
updated, _ := m.Update(testKey(tea.KeyEnter)) // choosePicker
next := updated.(model)
row, ok := findTranscriptRow(next.transcript, rowSystem)
if !ok {
t.Fatalf("picker theme save failure should append a transcript note, got rows %#v", next.transcript)
}
if !strings.Contains(row.text, "could not save theme preference") {
t.Fatalf("transcript note = %q, want the could-not-save note", row.text)
}
}
+12 -20
View File
@@ -474,31 +474,23 @@ func (wizard *mcpAddWizardState) mcpAddWizardResultActionStartRow() int {
func (m *model) selectModelPickerAtMouse(msg tea.MouseMsg) (mouseSelectionTarget, bool) {
width := chatWidth(m.width)
overlay, start, count := m.modelPickerOverlayLayout(width)
hit, ok := m.overlayMouseHit(msg, overlay, width)
hit, ok := m.overlayMouseHit(msg, m.modelPickerOverlay(width), width)
if !ok {
return mouseSelectionTarget{}, false
}
line := 3
if strings.TrimSpace(m.modelPickerLoadError) != "" {
line++
maxVisible := minInt(pickerOverlayMaxVisible, len(m.picker.items))
start := selectableListStart(len(m.picker.items), maxVisible, clampInt(m.picker.selected, 0, len(m.picker.items)-1))
rowStart := 3
if m.modelPickerLoadError != "" {
rowStart++
}
innerWidth := maxInt(1, modelPickerOverlayWidth(width, m.picker)-4)
lastGroup := ""
for index := start; index < start+count; index++ {
item := m.picker.items[index]
if item.Group != "" && item.Group != lastGroup {
line++
lastGroup = item.Group
}
rowLines := len(viewLines(renderModelPickerRow(innerWidth, index == m.picker.selected, item)))
if hit.y >= line && hit.y < line+rowLines {
m.picker.selected = index
return mouseSelectionTarget{Scope: "picker", Kind: int(m.picker.kind), Value: item.Value, Index: index}, true
}
line += rowLines
row := hit.y - rowStart
if row < 0 || row >= maxVisible {
return mouseSelectionTarget{}, false
}
return mouseSelectionTarget{}, false
index := start + row
m.picker.selected = index
return mouseSelectionTarget{Scope: "picker", Kind: int(m.picker.kind), Value: m.picker.items[index].Value, Index: index}, true
}
func (m *model) selectGenericPickerAtMouse(msg tea.MouseMsg) (mouseSelectionTarget, bool) {
-50
View File
@@ -1014,56 +1014,6 @@ func composerMousePoint(t *testing.T, m model, column int) (int, int) {
return x, y
}
func TestModelPickerMouseUsesRenderedRowBoundaries(t *testing.T) {
for _, line := range []int{6, 7} {
m := mouseTestModel()
m.width, m.height = 30, 40
m.picker = &commandPicker{kind: pickerModel, items: []pickerItem{
{Group: "Recent", Label: "model-subscription-one", Value: "one", OwnerLabel: "subscription-profile"},
{Group: "Recent", Label: "model-subscription-two", Value: "two", OwnerLabel: "subscription-profile"},
}}
overlay := m.modelPickerOverlay(30)
rect := m.overlayMouseRect(len(viewLines(overlay)), 30)
// Border/search/rule/header occupy lines 0–3; each model has two lines.
target, ok := m.selectModelPickerAtMouse(testMouseClick(tea.MouseLeft, 15, rect.y+line))
if !ok || target.Value != "two" {
t.Fatalf("line %d selected %#v, ok=%v; want second model", line, target, ok)
}
if _, ok := m.selectModelPickerAtMouse(testMouseClick(tea.MouseLeft, 15, rect.y+3)); ok {
t.Fatal("group header must not select a model")
}
}
}
func TestModelPickerFitsMultilineRowsToViewport(t *testing.T) {
for _, selected := range []int{0, 5, 9} {
m := mouseTestModel()
m.width, m.height = 30, 24
m.picker = &commandPicker{kind: pickerModel, selected: selected}
for i := 0; i < 10; i++ {
m.picker.items = append(m.picker.items, pickerItem{Group: "Recent", Label: "model-subscription-" + string(rune('a'+i)), OwnerLabel: "work-subscription-provider-east", OwnerNumber: 1})
}
overlay := m.modelPickerOverlay(30)
frame := m.scrollableTranscriptFrame(m.pinnedTitleBar(30), m.footerView(30))
if got := len(viewLines(overlay)); got > frame.bodyHeight {
t.Fatalf("selected=%d: overlay height %d exceeds body %d", selected, got, frame.bodyHeight)
}
plain := plainRender(t, overlay)
assertContains(t, plain, "❯ "+m.picker.items[selected].Label)
assertContains(t, plain, "Enter select")
rect := m.overlayMouseRect(len(viewLines(overlay)), 30)
for line, text := range strings.Split(plain, "\n") {
if !strings.Contains(text, "❯ ") {
continue
}
target, ok := m.selectModelPickerAtMouse(testMouseClick(tea.MouseLeft, 15, rect.y+line+1))
if !ok || target.Index != selected {
t.Fatalf("scrolled owner-line click selected %#v, ok=%v; want %d", target, ok, selected)
}
}
}
}
func mouseTestModel() model {
m := newModel(context.Background(), Options{})
m.width = 100
-221
View File
@@ -1,221 +0,0 @@
package tui
import (
"strings"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/notify"
)
// notifyChoice is one row in the /notify picker: a (mode, focusMode) pair and
// the label the user reads.
type notifyChoice struct {
label string
mode string
focusMode string
}
// notifyChoiceSubtitle renders the human explanation shown as the picker row's
// Meta text: "<what the mode does> — <when the focus fires>".
func (c notifyChoice) subtitle() string {
modeDescriptions := map[string]string{
string(notify.ModeOff): "silent",
string(notify.ModeBell): "terminal bell only",
string(notify.ModeNotify): "desktop notification only",
string(notify.ModeBoth): "terminal bell + desktop notification",
}
focusDescriptions := map[string]string{
string(notify.FocusUnfocused): "only when the terminal is in the background",
string(notify.FocusAlways): "every time",
string(notify.FocusFocused): "only while the terminal is focused",
}
return modeDescriptions[c.mode] + " — " + focusDescriptions[c.focusMode]
}
// notifyPickerChoices enumerates the FULL mode x focus space (4 modes x 3
// focus modes = 12 rows), the way newThemePicker enumerates every theme. The
// earlier 4-row curated list could not represent the other 8 valid pairs, so
// opening the picker on one of them fell through to row 0 and Enter silently
// committed a different setting than the user's current one (maintainer
// review, PR #1001). Every valid pair must have a row so Enter always keeps
// (or explicitly changes) the user's actual setting.
func notifyPickerChoices() []notifyChoice {
modes := []string{string(notify.ModeBoth), string(notify.ModeBell), string(notify.ModeNotify), string(notify.ModeOff)}
foci := []string{string(notify.FocusUnfocused), string(notify.FocusAlways), string(notify.FocusFocused)}
choices := make([]notifyChoice, 0, len(modes)*len(foci))
for _, mode := range modes {
for _, focus := range foci {
choices = append(choices, notifyChoice{
label: mode + " · " + focus,
mode: mode,
focusMode: focus,
})
}
}
return choices
}
// handleNotifyCommand implements /notify [list|off|bell|notify|both [focus]].
// Bare `/notify` opens the picker at the dispatch layer.
//
// Two contracts share this handler and must stay separate (maintainer review,
// PR #1001):
//
// - LIVE policy (m.notifyMode/m.notifyFocusMode, m.notifier): what the
// current session uses, initialized by newModel from the RESOLVED pair
// (user config after project precedence). A mode-only change preserves the
// live focus so a project's focus rule keeps applying to the session, and
// the change reaches the notifier immediately.
// - PERSISTED policy (the user's own config file): a partial write seeds
// omitted fields from config.UserNotify — the user's global choice —
// never from the resolved view, so a project's value cannot leak into the
// global file and a blank stays blank.
//
// Mirrors handleThemeCommand.
func (m model) handleNotifyCommand(args string) (model, string) {
tokens := strings.Fields(strings.TrimSpace(args))
if len(tokens) == 0 {
return m, m.notifyStateText()
}
if tokens[0] == "list" {
if len(tokens) != 1 {
return m, "Notify\nToo many arguments: " + args + " (usage: /notify list)"
}
return m, m.notifyStateText()
}
if len(tokens) > 2 {
return m, "Notify\nToo many arguments: " + args + " (usage: /notify <off|bell|notify|both> [unfocused|always|focused])"
}
mode := strings.ToLower(strings.TrimSpace(tokens[0]))
if !isValidNotifyMode(mode) {
return m, "Notify\nUnknown mode: " + tokens[0] + " (expected off, bell, notify, or both; run /notify with no argument to pick from the list)"
}
// LIVE focus: an explicit token wins; otherwise preserve the in-session
// value (which started as the resolved pair, project precedence included).
liveFocus := m.notifyFocusMode
focusExplicit := false
if len(tokens) > 1 {
focus := strings.ToLower(strings.TrimSpace(tokens[1]))
if !isValidNotifyFocusMode(focus) {
return m, "Notify\nUnknown focus mode: " + tokens[1] + " (expected unfocused, always, or focused)"
}
liveFocus = focus
focusExplicit = true
}
m.notifyMode = mode
m.notifyConfiguredMode = mode
m.notifyFocusMode = liveFocus
// Apply to the live notifier so the change takes effect on the next
// notification in this session, not just after a restart. A blank
// focusMode is fine here: the notifier treats blank as "unfocused".
if m.notifier != nil {
m.notifier.Configure(notify.Config{
Mode: notify.Mode(mode),
FocusMode: notify.FocusMode(liveFocus),
})
}
lines := []string{
"Notify",
"active mode: " + effectiveModeLabel(mode, m.notifyConfiguredMode) + ", focus: " + effectiveFocusLabel(liveFocus),
}
if note := m.persistNotifyPreference(mode, liveFocus, focusExplicit); note != "" {
lines = append(lines, note)
}
return m, strings.Join(lines, "\n")
}
// effectiveFocusLabel renders a focus value for the state line: blank means the
// built-in "unfocused" default, so say so instead of showing an empty string.
func effectiveFocusLabel(focus string) string {
if strings.TrimSpace(focus) == "" {
return "unfocused (default)"
}
return focus
}
// effectiveModeLabel renders the notify mode for state lines: when the mode in
// effect came from the built-in default rather than an explicit choice (stored
// or in-session), label it "(default)" so the TUI agrees with the CLI's
// `mode: (default)` for the same configuration (maintainer review, PR #1001).
// A mode the user actually chose — including a previously stored "both" —
// renders bare.
func effectiveModeLabel(mode, configuredMode string) string {
if strings.TrimSpace(mode) == strings.TrimSpace(configuredMode) {
return mode
}
return mode + " (default)"
}
// persistNotifyPreference writes the choice to user config so it survives a
// restart, as ONE serialized read-merge-write (config.UpdateNotify): the lock
// covers reading the stored block and replacing the file, so a concurrent
// partial update cannot interleave (maintainer review, PR #1001).
//
// Persist contract: an explicit pair replaces both fields. A mode-only change
// preserves the focus stored in the USER'S OWN file (blank stays blank) — the
// merge leaves current.FocusMode untouched rather than pre-reading it, which
// would reopen the stale-preservation window the lock is there to close. Best-
// effort: returns a short note to surface on failure, or "" on success / when
// there is no config path (e.g. tests).
func (m model) persistNotifyPreference(mode string, focus string, focusExplicit bool) string {
if strings.TrimSpace(m.userConfigPath) == "" {
return ""
}
_, err := config.UpdateNotify(m.userConfigPath, func(current config.NotifyConfig) config.NotifyConfig {
current.Mode = mode
if focusExplicit {
current.FocusMode = focus
}
return current
})
if err != nil {
return "note: could not save notify preference (" + err.Error() + ")"
}
return ""
}
// notifyStateText renders the /notify state view: the LIVE session policy
// (model fields, initialized from the resolved pair and updated by /notify) —
// not the stored file, which can legitimately disagree when a project config
// overrides notify for this session (maintainer review, PR #1001).
func (m model) notifyStateText() string {
sections := []commandSection{{
Title: "State",
Lines: []string{
"active mode: " + effectiveModeLabel(m.notifyMode, m.notifyConfiguredMode),
"active focus: " + effectiveFocusLabel(m.notifyFocusMode),
},
}}
rows := make([]string, 0, 12)
for _, c := range notifyPickerChoices() {
rows = append(rows, c.label)
}
sections = append(sections, commandSection{
Title: "Available",
Lines: rows,
})
return renderCommandOutput(commandOutput{
Title: "Notify",
Status: commandStatusOK,
Sections: sections,
Hints: []string{"run /notify with no argument to open the picker, or /notify <mode> [focus] to change directly"},
})
}
// isValidNotifyMode reports whether s names one of the four notification modes.
func isValidNotifyMode(s string) bool {
switch s {
case string(notify.ModeOff), string(notify.ModeBell), string(notify.ModeNotify), string(notify.ModeBoth):
return true
}
return false
}
// isValidNotifyFocusMode reports whether s names one of the three focus modes.
func isValidNotifyFocusMode(s string) bool {
switch s {
case string(notify.FocusUnfocused), string(notify.FocusAlways), string(notify.FocusFocused):
return true
}
return false
}
-472
View File
@@ -1,472 +0,0 @@
package tui
import (
"bytes"
"context"
"encoding/json"
"os"
"path/filepath"
"strings"
"testing"
tea "charm.land/bubbletea/v2"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/notify"
)
// A committed /notify choice is written to user config and reloaded at startup,
// so a /notify choice survives restart, just like /theme.
func TestNotifyChoicePersistsAcrossRestart(t *testing.T) {
cfgPath := filepath.Join(t.TempDir(), "config.json")
// First session: pick a notify pair via the text handler (the same commit
// path the picker uses via choosePicker).
m := newModel(context.Background(), Options{UserConfigPath: cfgPath})
m, out := m.handleNotifyCommand("off always")
if m.notifyMode != "off" {
t.Fatalf("notifyMode = %q, want off", m.notifyMode)
}
if !strings.Contains(out, "Notify") {
t.Fatalf("output should announce the change, got: %s", out)
}
data, err := os.ReadFile(cfgPath)
if err != nil {
t.Fatalf("notify commit should have written config: %v", err)
}
var cfg struct {
Notify config.NotifyConfig `json:"notify"`
}
if err := json.Unmarshal(data, &cfg); err != nil {
t.Fatalf("config is not valid JSON: %v", err)
}
if cfg.Notify.Mode != "off" || cfg.Notify.FocusMode != "always" {
t.Fatalf("notify = %+v, want mode=off focusMode=always", cfg.Notify)
}
// Second session: the persisted notify block seeds startup (options.Notify
// is populated by the resolver from the same file).
restarted := newModel(context.Background(), Options{UserConfigPath: cfgPath, Notify: config.NotifyConfig{Mode: "off", FocusMode: "always"}})
if restarted.notifyMode != "off" || restarted.notifyFocusMode != "always" {
t.Fatalf("restarted = mode %q focus %q, want off/always (from saved config)", restarted.notifyMode, restarted.notifyFocusMode)
}
}
// Mode-only `/notify <mode>` keeps the LIVE and PERSISTED policies separate
// (maintainer review, PR #1001): the live session preserves its in-session
// focus (initialized from the resolved pair, so a project's focus rule keeps
// applying), while the global write seeds from the USER'S OWN file — blank
// stays blank, and a project-derived value is never copied into the global
// config. The three states here deliberately disagree:
//
// user file: focusMode=focused, no mode
// resolved: both/unfocused (as a project override would produce)
// live start: both/unfocused
func TestNotifyModeOnlySplitsLiveAndPersistedFocus(t *testing.T) {
cfgPath := filepath.Join(t.TempDir(), "config.json")
if err := os.WriteFile(cfgPath, []byte(`{"notify":{"focusMode":"focused"}}`), 0o600); err != nil {
t.Fatal(err)
}
m := newModel(context.Background(), Options{
UserConfigPath: cfgPath,
Notify: config.NotifyConfig{Mode: "both", FocusMode: "unfocused"},
})
m, _ = m.handleNotifyCommand("bell")
// LIVE: mode changes, focus keeps the in-session (resolved) value — the
// user-file focus must NOT bleed into the session.
if m.notifyMode != "bell" {
t.Errorf("live mode = %q, want bell", m.notifyMode)
}
if m.notifyFocusMode != "unfocused" {
t.Errorf("live focus = %q, want the in-session unfocused (not the user-file focused)", m.notifyFocusMode)
}
// PERSISTED: the mode is written; the focus seeded from the user's file
// (their explicit focused), never the resolved/live value.
persisted := readNotifyBlock(t, cfgPath)
if persisted.Mode != "bell" {
t.Errorf("persisted mode = %q, want bell", persisted.Mode)
}
if persisted.FocusMode != "focused" {
t.Errorf("persisted focusMode = %q, want the user-file value focused", persisted.FocusMode)
}
// Blank stays blank: with nothing stored, a mode-only change must not pin
// the default focus as an explicit choice, and live still keeps its value.
blankPath := filepath.Join(t.TempDir(), "config.json")
m = newModel(context.Background(), Options{
UserConfigPath: blankPath,
Notify: config.NotifyConfig{Mode: "off", FocusMode: "focused"},
})
m, _ = m.handleNotifyCommand("off")
if m.notifyFocusMode != "focused" {
t.Errorf("live focus = %q, want the in-session focused preserved", m.notifyFocusMode)
}
persisted = readNotifyBlock(t, blankPath)
if persisted.Mode != "off" {
t.Errorf("persisted mode = %q, want off", persisted.Mode)
}
if persisted.FocusMode != "" {
t.Errorf("persisted focusMode = %q, want blank (unspecified stays unspecified)", persisted.FocusMode)
}
}
// The full contract matrix when the three notify states disagree (blank/different
// user file + project-resolved pair), per the maintainer review: display and
// preselection report the LIVE policy; partial updates change live and persist
// independently; restart precedence never copies project values into the
// global file.
func TestNotifyLivePolicyContractMatrix(t *testing.T) {
cfgPath := filepath.Join(t.TempDir(), "config.json")
// User file: nothing stored at all.
if err := os.WriteFile(cfgPath, []byte(`{}`), 0o600); err != nil {
t.Fatal(err)
}
// Project-resolved pair: off/focused (a project override shape).
m := newModel(context.Background(), Options{
UserConfigPath: cfgPath,
Notify: config.NotifyConfig{Mode: "off", FocusMode: "focused"},
})
// `/notify list` reports the live resolved pair, not the (blank) file.
state := m.notifyStateText()
if !strings.Contains(state, "active mode: off") || !strings.Contains(state, "active focus: focused") {
t.Errorf("list should report the live off/focused pair, got:\n%s", state)
}
// The picker preselects the live pair; Enter without navigation does not
// change it.
m.input.SetValue("/notify")
updated, cmd := m.Update(testKey(tea.KeyEnter))
if cmd != nil {
t.Fatalf("opening the picker should not emit a cmd, got %T", cmd)
}
m = updated.(model)
if sel := m.picker.items[m.picker.selected]; sel.Value != "off focused" {
t.Fatalf("preselected = %q, want the live pair %q", sel.Value, "off focused")
}
updated, _ = m.Update(testKey(tea.KeyEnter))
m = updated.(model)
if m.notifyMode != "off" || m.notifyFocusMode != "focused" {
t.Errorf("Enter on preselected live pair changed live state: %q/%q", m.notifyMode, m.notifyFocusMode)
}
persisted := readNotifyBlock(t, cfgPath)
if persisted.Mode != "off" || persisted.FocusMode != "focused" {
t.Errorf("Enter on the live pair persisted %q/%q, want off/focused (explicit row commit)", persisted.Mode, persisted.FocusMode)
}
// Reset the file to blank and re-seed the session for the partial-update
// checks (fresh states: blank file, live off/focused).
if err := os.WriteFile(cfgPath, []byte(`{}`), 0o600); err != nil {
t.Fatal(err)
}
m = newModel(context.Background(), Options{
UserConfigPath: cfgPath,
Notify: config.NotifyConfig{Mode: "off", FocusMode: "focused"},
})
// `/notify bell` (mode-only): live becomes bell/focused; the file gets
// mode bell and keeps its focus BLANK — the project's focused is not
// copied into the global file.
m, _ = m.handleNotifyCommand("bell")
if m.notifyMode != "bell" || m.notifyFocusMode != "focused" {
t.Errorf("after mode-only: live = %q/%q, want bell/focused", m.notifyMode, m.notifyFocusMode)
}
persisted = readNotifyBlock(t, cfgPath)
if persisted.Mode != "bell" {
t.Errorf("after mode-only: persisted mode = %q, want bell", persisted.Mode)
}
if persisted.FocusMode != "" {
t.Errorf("after mode-only: persisted focusMode = %q, want blank (project value must not leak)", persisted.FocusMode)
}
// `/notify bell always` (explicit pair): both live fields change and both
// explicit values persist.
m, _ = m.handleNotifyCommand("bell always")
if m.notifyMode != "bell" || m.notifyFocusMode != "always" {
t.Errorf("after explicit pair: live = %q/%q, want bell/always", m.notifyMode, m.notifyFocusMode)
}
persisted = readNotifyBlock(t, cfgPath)
if persisted.Mode != "bell" || persisted.FocusMode != "always" {
t.Errorf("after explicit pair: persisted = %q/%q, want bell/always", persisted.Mode, persisted.FocusMode)
}
// Restart precedence: a fresh resolve of the user file (now bell/always,
// no project) yields exactly the stored pair — nothing else leaked in.
resolved, err := config.Resolve(config.ResolveOptions{UserConfigPath: cfgPath, Env: map[string]string{}})
if err != nil {
t.Fatalf("resolve: %v", err)
}
if resolved.Notify.Mode != "bell" || resolved.Notify.FocusMode != "always" {
t.Errorf("restart resolve = %q/%q, want bell/always", resolved.Notify.Mode, resolved.Notify.FocusMode)
}
}
// `/notify bell unfocused` updates both fields in one call.
func TestNotifyCommandSetsModeAndFocus(t *testing.T) {
m := newModel(context.Background(), Options{})
m, _ = m.handleNotifyCommand("bell unfocused")
if m.notifyMode != "bell" {
t.Errorf("notifyMode = %q, want bell", m.notifyMode)
}
if m.notifyFocusMode != "unfocused" {
t.Errorf("notifyFocusMode = %q, want unfocused", m.notifyFocusMode)
}
}
// The choice reaches the LIVE notifier immediately, so the change applies on
// the next permission prompt in this session (not only after a restart).
func TestNotifyCommandAppliesToLiveNotifier(t *testing.T) {
var buf bytes.Buffer
m := newModel(context.Background(), Options{Notify: config.NotifyConfig{Mode: "off", FocusMode: "always"}})
m.notifier = notify.New(&buf, notify.Config{Mode: notify.ModeOff, FocusMode: notify.FocusAlways})
m.notifier.SetFocused(true)
m, _ = m.handleNotifyCommand("bell always")
m.notifier.Notify(notify.Completion, "x")
if buf.String() != "\x07" {
t.Fatalf("live notifier should bell after /notify bell, got %q", buf.String())
}
m, _ = m.handleNotifyCommand("off always")
m.notifier.Notify(notify.Completion, "x")
if buf.String() != "\x07" {
t.Fatalf("live notifier should go silent after /notify off, got %q", buf.String())
}
}
// `/notify off always typo` (more than two tokens) is rejected, not silently
// accepted as a successful change.
func TestNotifyCommandRejectsTrailingArguments(t *testing.T) {
m := newModel(context.Background(), Options{})
m.notifyMode = "both"
m, out := m.handleNotifyCommand("off always typo")
if m.notifyMode != "both" {
t.Errorf("trailing args should not mutate state, got %q", m.notifyMode)
}
if !strings.Contains(out, "Too many arguments") {
t.Errorf("output should explain the rejection, got: %s", out)
}
}
// `/notify list junk` is a usage error, not a successful list — matching
// /theme, /effort, and /style, which recognize only an exact `list`
// (maintainer review, PR #1001). Exact `/notify list` still shows state.
func TestNotifyCommandListRequiresExactlyOneToken(t *testing.T) {
m := newModel(context.Background(), Options{})
m.notifyMode = "off"
m.notifyFocusMode = "focused"
m, out := m.handleNotifyCommand("list junk")
if !strings.Contains(out, "Too many arguments") {
t.Errorf("`list junk` should be a usage error, got: %s", out)
}
if m.notifyMode != "off" || m.notifyFocusMode != "focused" {
t.Errorf("`list junk` should not mutate state, got %q/%q", m.notifyMode, m.notifyFocusMode)
}
_, out = m.handleNotifyCommand("list")
if !strings.Contains(out, "active mode: off") || !strings.Contains(out, "active focus: focused") {
t.Errorf("exact `/notify list` should still show the live state, got: %s", out)
}
}
// `/notify loud` (invalid) returns an error message; the model's notifyMode
// is NOT mutated, so a typo cannot accidentally turn the alert off.
func TestNotifyCommandRejectsInvalidMode(t *testing.T) {
m := newModel(context.Background(), Options{})
m.notifyMode = "both"
m, out := m.handleNotifyCommand("loud")
if m.notifyMode != "both" {
t.Errorf("invalid mode should not mutate state, got %q", m.notifyMode)
}
if !strings.Contains(out, "Unknown mode") {
t.Errorf("output should explain the error, got: %s", out)
}
}
// `/notify bell sideways` fails validation before persisting, so neither field
// changes.
func TestNotifyCommandRejectsInvalidFocus(t *testing.T) {
m := newModel(context.Background(), Options{})
m.notifyMode = "bell"
m.notifyFocusMode = "always"
m, out := m.handleNotifyCommand("bell sideways")
if m.notifyMode != "bell" || m.notifyFocusMode != "always" {
t.Errorf("invalid focus should not mutate state, got mode=%q focus=%q", m.notifyMode, m.notifyFocusMode)
}
if !strings.Contains(out, "Unknown focus mode") {
t.Errorf("output should explain the error, got: %s", out)
}
}
// The picker enumerates the FULL mode x focus space (12 rows), so every valid
// pair is representable and Enter can never silently commit a different pair
// than the user's current one. Maintainer review, PR #1001: with a 4-row
// curated list, opening the picker on (off, always) preselected row 0
// (both, unfocused) and Enter changed the setting.
func TestNotifyPickerEnumeratesFullSpace(t *testing.T) {
m := newModel(context.Background(), Options{})
picker := m.newNotifyPicker()
if len(picker.items) != 12 {
t.Fatalf("picker has %d items, want 12 (4 modes x 3 focus modes)", len(picker.items))
}
seen := map[string]bool{}
for _, item := range picker.items {
if seen[item.Value] {
t.Errorf("duplicate picker row %q", item.Value)
}
seen[item.Value] = true
}
for _, mode := range []string{"off", "bell", "notify", "both"} {
for _, focus := range []string{"unfocused", "always", "focused"} {
if !seen[mode+" "+focus] {
t.Errorf("picker is missing row for valid pair %q %q", mode, focus)
}
}
}
}
// Every pair the picker preselects must also be a row: Enter on an open
// picker must keep (or explicitly change) the user's actual setting. This is
// the maintainer's suggested regression: send Enter to an open picker from a
// pair that is NOT in a 4-row curated list — with full enumeration the
// preselected row IS the current pair and the commit is a no-op change.
func TestNotifyPickerEnterOnUnlistedPairKeepsSetting(t *testing.T) {
cfgPath := filepath.Join(t.TempDir(), "config.json")
if err := os.WriteFile(cfgPath, []byte(`{"notify":{"mode":"off","focusMode":"always"}}`), 0o600); err != nil {
t.Fatal(err)
}
m := newModel(context.Background(), Options{
UserConfigPath: cfgPath,
Notify: config.NotifyConfig{Mode: "off", FocusMode: "always"},
})
m.input.SetValue("/notify")
updated, cmd := m.Update(testKey(tea.KeyEnter))
m = updated.(model)
if cmd != nil {
t.Fatalf("opening the notify picker should not emit a cmd, got %T", cmd)
}
if m.picker == nil || m.picker.kind != pickerNotify {
t.Fatalf("expected the notify picker to open, got %#v", m.picker)
}
// (off, always) must be preselected — it is a valid pair even though the
// old curated list could not represent it.
if sel := m.picker.items[m.picker.selected]; sel.Value != "off always" {
t.Fatalf("preselected = %q, want the current pair %q", sel.Value, "off always")
}
// Enter commits the preselected row: the setting is unchanged.
updated, _ = m.Update(testKey(tea.KeyEnter))
m = updated.(model)
if m.picker != nil {
t.Fatal("picker should close on Enter")
}
persisted := readNotifyBlock(t, cfgPath)
if persisted.Mode != "off" || persisted.FocusMode != "always" {
t.Fatalf("Enter changed the setting: got %+v, want off/always unchanged", persisted)
}
}
// The picker's Value strings are the same "<mode> <focus>" form the text
// handler accepts, so the commit path can be shared. This is the contract that
// lets choosePicker dispatch to handleNotifyCommand without translation.
func TestNotifyPickerValuesAreValidCommandArgs(t *testing.T) {
m := newModel(context.Background(), Options{})
picker := m.newNotifyPicker()
if len(picker.items) == 0 {
t.Fatal("picker has no items")
}
for _, item := range picker.items {
tokens := strings.Fields(item.Value)
if len(tokens) != 2 {
t.Errorf("item %q has %d tokens, want 2 (mode focus)", item.Value, len(tokens))
continue
}
if !isValidNotifyMode(tokens[0]) {
t.Errorf("item %q: mode %q is not a valid notify mode", item.Value, tokens[0])
}
if !isValidNotifyFocusMode(tokens[1]) {
t.Errorf("item %q: focus %q is not a valid focus mode", item.Value, tokens[1])
}
}
}
// The /notify state view reports the LIVE session policy (initialized from the
// resolved pair), not the stored file — the two can legitimately disagree when
// a project config overrides notify (maintainer review, PR #1001). A blank
// live focus renders as the unfocused default.
func TestNotifyStateTextShowsLivePair(t *testing.T) {
m := newModel(context.Background(), Options{Notify: config.NotifyConfig{Mode: "bell", FocusMode: "always"}})
state := m.notifyStateText()
if !strings.Contains(state, "active mode: bell") {
t.Errorf("state should show the live mode, got: %s", state)
}
if !strings.Contains(state, "active focus: always") {
t.Errorf("state should show the live focus, got: %s", state)
}
// Blank live focus renders as the default, never an empty string.
m = newModel(context.Background(), Options{Notify: config.NotifyConfig{Mode: "bell"}})
state = m.notifyStateText()
if !strings.Contains(state, "active focus: unfocused (default)") {
t.Errorf("state should label a blank live focus as the default, got: %s", state)
}
}
func readNotifyBlock(t *testing.T, path string) config.NotifyConfig {
t.Helper()
data, err := os.ReadFile(path)
if err != nil {
t.Fatalf("read config %s: %v", path, err)
}
var cfg struct {
Notify config.NotifyConfig `json:"notify"`
}
if err := json.Unmarshal(data, &cfg); err != nil {
t.Fatalf("decode config %s: %v", path, err)
}
return cfg.Notify
}
// Maintainer regression (PR #1001): a blank stored mode must read as a
// DEFAULT on every first-party surface — the TUI says "(default)" like the
// CLI's `mode: (default)`, not a bare effective value the user cannot
// distinguish from a stored "both". The mode line mirrors focus, which
// already labels its built-in default ("unfocused (default)").
func TestNotifyStateLabelsEffectiveDefaultMode(t *testing.T) {
// Nothing stored, nothing resolved: the TUI applies its effective
// default ("both") and must say so.
m := newModel(context.Background(), Options{})
state := m.notifyStateText()
if !strings.Contains(state, "active mode: both (default)") {
t.Errorf("unconfigured state should label the effective mode as (default), got:\n%s", state)
}
// A stored/resolved explicit "both" renders bare — the user chose it.
// (The focus line may still say "unfocused (default)"; only the MODE
// line is under assertion.)
m = newModel(context.Background(), Options{Notify: config.NotifyConfig{Mode: "both"}})
state = m.notifyStateText()
if strings.Contains(state, "active mode: both (default)") {
t.Errorf("an explicit both must not carry the (default) label, got:\n%s", state)
}
if !strings.Contains(state, "active mode: both") {
t.Errorf("explicit both should still show the mode, got:\n%s", state)
}
// A mode-only change from unconfigured commits an explicit choice: the
// change line and later state output render it bare, never as a default.
m = newModel(context.Background(), Options{})
m, out := m.handleNotifyCommand("bell")
if strings.Contains(out, "active mode: bell (default)") {
t.Errorf("explicit /notify bell change line should render bare, got:\n%s", out)
}
if !strings.Contains(out, "active mode: bell") {
t.Errorf("change line should show the new mode, got:\n%s", out)
}
if state := m.notifyStateText(); strings.Contains(state, "active mode: bell (default)") {
t.Errorf("state after an explicit choice should render bare, got:\n%s", state)
}
}
-5
View File
@@ -99,11 +99,6 @@ type Options struct {
// KeyBindingsConfig means "use built-in defaults" for each action.
KeyBindings config.KeyBindingsConfig
// StartupNotices are shown as system notices when the session opens: facts
// about the session the user should see before the first prompt, decided by
// the caller that built it. Blank entries are skipped.
StartupNotices []string
// STT configures speech-to-text dictation (§ docs/dictation.md).
STT config.STTConfig
// BuildDictationTranscriber constructs the transcriber for the current STT
+16 -215
View File
@@ -11,7 +11,6 @@ import (
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/modelregistry"
"github.com/Gitlawb/zero/internal/notify"
"github.com/Gitlawb/zero/internal/providercatalog"
"github.com/Gitlawb/zero/internal/providermodelcatalog"
"github.com/Gitlawb/zero/internal/providermodeldiscovery"
@@ -30,7 +29,6 @@ const (
pickerSTTModel
pickerSTTDownload
pickerPet
pickerNotify
)
// pickerItem is one selectable row: Label is shown, Value is passed to the
@@ -50,16 +48,9 @@ type pickerItem struct {
// /model picker can switch providers when a model from a non-active provider is
// chosen. Empty for non-model items.
OwnerProvider string
// OwnerLabel names the provider Enter will actually use for a row in a
// mixed-provider group (Recent, Favorites), whose section header cannot
// convey it. Empty for provider-grouped rows.
OwnerLabel string
// OwnerNumber distinguishes abbreviated names, even when their visible
// suffixes coincide. It is the saved-profile position (1-based).
OwnerNumber int
Remote bool
Local bool
Favorite bool
Remote bool
Local bool
Favorite bool
}
// commandPicker is a generic single-select overlay reused by /model and /effort
@@ -298,19 +289,10 @@ func (m model) savedProviderModelPickerItems(profile config.ProviderProfile, act
descriptor, hasDescriptor := m.descriptorForProfile(profile)
group := modelPickerProviderGroup(profile, descriptor, hasDescriptor)
endpointKey := providerEndpointKey(profile, descriptor)
var discovered []providermodeldiscovery.Model
if endpointKey != "" {
discovered = m.modelPickerLiveByProvider[endpointKey]
}
if len(discovered) == 0 && hasDescriptor && !genericProviderCatalogID(descriptor.ID) {
discovered = m.modelPickerLiveByProvider[descriptor.ID]
}
var raw []pickerItem
switch {
case len(discovered) > 0:
for _, model := range discovered {
case hasDescriptor && len(m.modelPickerLiveByProvider[descriptor.ID]) > 0:
for _, model := range m.modelPickerLiveByProvider[descriptor.ID] {
if strings.TrimSpace(model.ID) == "" {
continue
}
@@ -393,24 +375,17 @@ func (m model) openModelPicker() (model, tea.Cmd) {
}
// modelPickerDiscoveryCmds dispatches a live model-discovery command for each
// usable provider (deduped by endpoint key), so /model shows the same real
// usable provider (deduped by catalog descriptor), so /model shows the same real
// models the provider-setup wizard discovers.
func (m model) modelPickerDiscoveryCmds() tea.Cmd {
cmds := []tea.Cmd{}
seen := map[string]bool{}
for _, profile := range m.modelPickerProviders() {
descriptor, ok := m.descriptorForProfile(profile)
if !ok {
if !ok || seen[descriptor.ID] {
continue
}
endpointKey := providerEndpointKey(profile, descriptor)
if endpointKey == "" {
endpointKey = descriptor.ID
}
if seen[endpointKey] {
continue
}
seen[endpointKey] = true
seen[descriptor.ID] = true
if cmd := m.modelPickerProviderDiscoveryCmd(descriptor, profile); cmd != nil {
cmds = append(cmds, cmd)
}
@@ -442,7 +417,6 @@ func (m model) modelPickerProviderDiscoveryCmd(descriptor providercatalog.Descri
}
}
providerID := descriptor.ID
endpointKey := providerEndpointKey(profile, descriptor)
return func() tea.Msg {
ctx, cancel := context.WithTimeout(m.ctx, 8*time.Second)
defer cancel()
@@ -453,12 +427,7 @@ func (m model) modelPickerProviderDiscoveryCmd(descriptor providercatalog.Descri
}
}
models, err := discover(ctx, providerWizardDiscoveryProfile(descriptor, k, authed.BaseURL))
return modelPickerModelsDiscoveredMsg{
providerID: providerID,
endpointKey: endpointKey,
models: models,
err: err,
}
return modelPickerModelsDiscoveredMsg{providerID: providerID, models: models, err: err}
}
}
@@ -493,30 +462,6 @@ func pickerItemDedupKey(item pickerItem) string {
return item.Value
}
// modelPickerSwitchOwner returns the saved provider that choosing item switches
// to, or "" when the model is applied to the active provider: the owner is
// blank, is the active provider, or (registry-fallback / stale-history rows
// after a rename or removal) no longer resolves to a saved provider.
func (m model) modelPickerSwitchOwner(item pickerItem) string {
owner := strings.TrimSpace(item.OwnerProvider)
if owner == "" || strings.EqualFold(owner, strings.TrimSpace(m.providerName)) {
return ""
}
if _, ok := m.savedProviderByName(owner); !ok {
return ""
}
return owner
}
// modelPickerEffectiveOwner names the provider choosing item will use, so a
// row's owner label never promises a provider that selection falls back from.
func (m model) modelPickerEffectiveOwner(item pickerItem) string {
if owner := m.modelPickerSwitchOwner(item); owner != "" {
return owner
}
return strings.TrimSpace(m.providerName)
}
func (m model) assembleModelPickerItems(recent []pickerItem, catalog []pickerItem) []pickerItem {
result := []pickerItem{}
// Favorites keep the pre-provider-aware semantics: one row per favorited
@@ -531,7 +476,6 @@ func (m model) assembleModelPickerItems(recent []pickerItem, catalog []pickerIte
}
item.Group = "Favorites"
item.Favorite = true
item.OwnerLabel = m.modelPickerEffectiveOwner(item)
result = append(result, item)
favoriteSeen[item.Value] = true
}
@@ -554,7 +498,6 @@ func (m model) assembleModelPickerItems(recent []pickerItem, catalog []pickerIte
}
item.Group = "Recent"
item.Favorite = m.favoriteModels[item.Value]
item.OwnerLabel = m.modelPickerEffectiveOwner(item)
result = append(result, item)
seen[key] = true
}
@@ -570,17 +513,6 @@ func (m model) assembleModelPickerItems(recent []pickerItem, catalog []pickerIte
result = append(result, item)
seen[key] = true
}
for i := range result {
if result[i].OwnerLabel == "" {
continue
}
for index, profile := range m.modelPickerProviders() {
if strings.EqualFold(strings.TrimSpace(profile.Name), result[i].OwnerLabel) {
result[i].OwnerNumber = index + 1
break
}
}
}
return result
}
@@ -779,15 +711,7 @@ func modelPickerTitleWord(word string) string {
}
func (m model) activeProviderDescriptor() (providercatalog.Descriptor, bool) {
if descriptor, ok := m.descriptorForProfile(m.providerProfile); ok {
return descriptor, true
}
if name := strings.TrimSpace(m.providerName); name != "" {
if descriptor, ok := providercatalog.Get(name); ok {
return descriptor, true
}
}
return providercatalog.Descriptor{}, false
return m.descriptorForProfile(m.providerProfile)
}
func customProviderDescriptorForProfile(profile config.ProviderProfile) (providercatalog.Descriptor, bool) {
@@ -858,87 +782,10 @@ func genericProviderCatalogID(id string) bool {
return strings.HasPrefix(strings.TrimSpace(id), "custom-")
}
// providerEndpointKey derives an endpoint route identity for profile discovery.
// Profiles with an explicit Name use that name so distinct endpoints sharing a
// generic descriptor (e.g. custom-openai-compatible) never collide.
// Unnamed profiles qualify by base URL or descriptor ID.
func providerEndpointKey(profile config.ProviderProfile, descriptor providercatalog.Descriptor) string {
if name := strings.TrimSpace(profile.Name); name != "" {
return name
}
if baseURL := normalizeProviderBaseURL(profile.BaseURL); baseURL != "" {
if descriptor.ID != "" {
return descriptor.ID + "@" + baseURL
}
return baseURL
}
if descriptor.ID != "" {
return descriptor.ID
}
if catID := strings.TrimSpace(profile.CatalogID); catID != "" {
return catID
}
return ""
}
// activeDiscoveryKeys returns candidate lookup keys for live-discovered models
// belonging to the active provider endpoint route, in order of specificity.
// Generic descriptors (custom-*) are excluded as fallback keys so one custom
// endpoint's discovery cannot determine another endpoint's capabilities.
func (m model) activeDiscoveryKeys() []string {
var keys []string
seen := map[string]bool{}
add := func(k string) {
k = strings.TrimSpace(k)
if k != "" && !seen[k] {
seen[k] = true
keys = append(keys, k)
}
}
descriptor, hasDescriptor := m.activeProviderDescriptor()
if key := providerEndpointKey(m.providerProfile, descriptor); key != "" {
add(key)
}
if m.providerProfile.Name != "" {
add(m.providerProfile.Name)
}
if m.providerName != "" {
add(m.providerName)
}
if m.providerProfile.CatalogID != "" && !genericProviderCatalogID(m.providerProfile.CatalogID) {
add(m.providerProfile.CatalogID)
}
if hasDescriptor && !genericProviderCatalogID(descriptor.ID) {
add(descriptor.ID)
}
if len(keys) == 0 && len(m.modelPickerLiveByProvider) == 1 {
// When no active provider or profile is configured (e.g. dummy models in unit tests),
// accept the single available provider entry.
for id := range m.modelPickerLiveByProvider {
add(id)
break
}
}
return keys
}
// discoveredModelsForActiveRoute returns the live-discovered models for the active
// endpoint route, if available in modelPickerLiveByProvider.
func (m model) discoveredModelsForActiveRoute() ([]providermodeldiscovery.Model, bool) {
for _, key := range m.activeDiscoveryKeys() {
if models, ok := m.modelPickerLiveByProvider[key]; ok && len(models) > 0 {
return models, true
}
}
return nil, false
}
type modelPickerModelsDiscoveredMsg struct {
providerID string
endpointKey string
models []providermodeldiscovery.Model
err error
providerID string
models []providermodeldiscovery.Model
err error
}
// ollamaContextWindowDiscoveredMsg carries the result of an async /api/show
@@ -1029,18 +876,10 @@ func (m model) applyModelPickerModelsDiscovered(msg modelPickerModelsDiscoveredM
if msg.err != nil || len(msg.models) == 0 {
return m
}
newMap := make(map[string][]providermodeldiscovery.Model, len(m.modelPickerLiveByProvider)+2)
for k, v := range m.modelPickerLiveByProvider {
newMap[k] = v
if m.modelPickerLiveByProvider == nil {
m.modelPickerLiveByProvider = map[string][]providermodeldiscovery.Model{}
}
modelsCopy := append([]providermodeldiscovery.Model{}, msg.models...)
if msg.endpointKey != "" {
newMap[msg.endpointKey] = modelsCopy
}
if msg.providerID != "" && (!genericProviderCatalogID(msg.providerID) || msg.endpointKey == "") {
newMap[msg.providerID] = modelsCopy
}
m.modelPickerLiveByProvider = newMap
m.modelPickerLiveByProvider[msg.providerID] = append([]providermodeldiscovery.Model{}, msg.models...)
// Rebuild the open picker so this provider's section shows its live models,
// preserving the current query + selection.
if m.picker != nil && m.picker.kind == pickerModel {
@@ -1242,44 +1081,6 @@ func (m model) newThemePicker() *commandPicker {
return &commandPicker{kind: pickerTheme, title: "Choose a theme", items: items, allItems: append([]pickerItem{}, items...), selected: selected}
}
// newNotifyPicker lists the FULL (mode, focus) space from notifyPickerChoices
// (every valid pair has a row, so a setting like (off, always) is always
// preselectable and Enter can never silently commit a different pair). Each
// row's Value is the same synthetic string the text /notify handler accepts
// ("<mode> <focus>"), so /notify with no arg and the picker share one commit
// path through handleNotifyCommand. Preselection uses the LIVE session pair
// (m.notifyMode/m.notifyFocusMode, initialized from the resolved config) — not
// the stored file, which can disagree when a project config overrides notify
// for this session. Blank live fields resolve to their effective defaults for
// preselection only; committing a row is an explicit choice and persists that
// pair. There is no live preview — notify affects the next notification, not
// the current view.
func (m model) newNotifyPicker() *commandPicker {
choices := notifyPickerChoices()
items := make([]pickerItem, 0, len(choices))
selected := 0
activeMode := m.notifyMode
if strings.TrimSpace(activeMode) == "" {
activeMode = string(notify.ModeBoth)
}
activeFocus := m.notifyFocusMode
if strings.TrimSpace(activeFocus) == "" {
activeFocus = string(notify.FocusUnfocused)
}
for _, c := range choices {
items = append(items, pickerItem{
Group: "When Zero needs your input",
Label: c.label,
Value: c.mode + " " + c.focusMode,
Meta: c.subtitle(),
})
if c.mode == activeMode && c.focusMode == activeFocus {
selected = len(items) - 1
}
}
return &commandPicker{kind: pickerNotify, title: "select notify mode", items: items, allItems: append([]pickerItem{}, items...), selected: selected}
}
// pickerMoved advances the open picker's cursor by delta. Theme candidates render
// only in the picker preview; their active palette is applied only after Enter.
// Safe to call with no picker open. Callers mutate through m.picker (a pointer),
-121
View File
@@ -1351,127 +1351,6 @@ func TestModelPickerListsAllSavedProviders(t *testing.T) {
}
}
// A Recent or Favorites row's owner label must name the provider Enter actually
// uses. A history owner that was renamed or removed no longer resolves, so
// selection falls back to the active provider; the label must follow it rather
// than show the stale name.
func TestModelPickerOwnerLabelMatchesSelectedProvider(t *testing.T) {
const staleModel = "anthropic/claude-sonnet-4.5"
for _, group := range []string{"Recent", "Favorites"} {
t.Run(group, func(t *testing.T) {
active := config.ProviderProfile{
Name: "openrouter",
CatalogID: "openrouter",
ProviderKind: config.ProviderKindOpenAICompatible,
Model: "google/gemini-2.5-pro",
APIKeyEnv: "OPENROUTER_API_KEY",
BaseURL: "https://openrouter.ai/api/v1",
APIFormat: "chat-completions",
}
options := Options{
ProviderName: "openrouter",
ModelName: "google/gemini-2.5-pro",
Provider: &fakeProvider{},
ProviderProfile: active,
SavedProviders: []config.ProviderProfile{active, {Name: "xai", CatalogID: "xai", Model: "grok-4"}},
// "chatgpt" was removed after this model was used from it.
RecentModels: []config.RecentModelEntry{
{Provider: "chatgpt", Model: staleModel},
{Provider: "xai", Model: "grok-4"},
},
NewProvider: func(config.ProviderProfile) (zeroruntime.Provider, error) {
return &fakeProvider{}, nil
},
}
if group == "Favorites" {
options.FavoriteModels = []string{staleModel}
}
m := newModel(context.Background(), options)
m.picker = m.newModelPicker()
if m.picker == nil {
t.Fatal("expected model picker")
}
target := -1
for index, item := range m.picker.items {
if item.Group == group && item.Value == staleModel {
target = index
}
if item.Group == "Recent" && item.Value == "grok-4" && item.OwnerLabel != "xai" {
t.Fatalf("resolvable owner label = %q, want xai", item.OwnerLabel)
}
}
if target < 0 {
t.Fatalf("expected %s row for %s, got %#v", group, staleModel, m.picker.items)
}
if got := m.picker.items[target].OwnerLabel; got != "openrouter" {
t.Fatalf("stale owner label = %q, want the active provider", got)
}
row := plainRender(t, renderModelPickerRow(60, false, m.picker.items[target]))
assertContains(t, row, "openrouter · ")
assertNotContains(t, row, "chatgpt")
m.picker.selected = target
updated, _ := m.Update(testKey(tea.KeyEnter))
m = updated.(model)
if m.providerName != "openrouter" || m.modelName != staleModel {
t.Fatalf("selection used %s/%s, want the labeled provider openrouter/%s", m.providerName, m.modelName, staleModel)
}
})
}
}
func TestModelPickerAbbreviatedOwnersMatchSelection(t *testing.T) {
t.Setenv(config.ActiveProviderEnv, "")
home := t.TempDir()
for _, variable := range []string{"HOME", "XDG_CONFIG_HOME", "XDG_CACHE_HOME", "APPDATA", "LOCALAPPDATA"} {
t.Setenv(variable, home)
}
for _, names := range [][]string{
{"work-subscription-provider-east", "work-subscription-provider-west"},
{"east-subscription-provider-shared", "west-subscription-provider-shared"},
{"東-subscription-provider-shared", "西-subscription-provider-shared"},
} {
for _, favorite := range []bool{false, true} {
for _, terminalWidth := range []int{76, 30} {
rows := map[string]bool{}
for selectedOwner, name := range names {
profiles := []config.ProviderProfile{}
for _, owner := range names {
profiles = append(profiles, config.ProviderProfile{Name: owner, CatalogID: "openrouter", ProviderKind: config.ProviderKindOpenAICompatible, Model: "gpt-5.6", APIKey: "test-key", BaseURL: "https://openrouter.ai/api/v1", APIFormat: "chat-completions"})
}
m := newModel(context.Background(), Options{ProviderName: names[0], ProviderProfile: profiles[0], ModelName: "gpt-5.6", Provider: &fakeProvider{}, SavedProviders: profiles,
NewProvider: func(config.ProviderProfile) (zeroruntime.Provider, error) { return &fakeProvider{}, nil },
})
m.favoriteModels = map[string]bool{"gpt-5.6": favorite}
// Put each owner first in turn: Favorites retains only one row per model.
recent := []pickerItem{{Label: "GPT-5.6", Value: "gpt-5.6", OwnerProvider: name}, {Label: "GPT-5.6", Value: "gpt-5.6", OwnerProvider: names[1-selectedOwner]}}
items := m.assembleModelPickerItems(recent, nil)
m.picker = &commandPicker{kind: pickerModel, title: "Choose a model", items: items, allItems: items}
width := modelPickerOverlayWidth(terminalWidth, m.picker)
for _, selected := range []bool{false, true} {
row := plainRender(t, renderModelPickerRow(width-4, selected, items[0]))
assertContains(t, row, "GPT-5.6")
if selected {
if rows[row] {
t.Fatalf("different destinations have identical rows: %q", row)
}
rows[row] = true
}
}
overlay := plainRender(t, m.modelPickerOverlay(terminalWidth))
joined := strings.NewReplacer(" ", "", "\n", "", "│", "").Replace(overlay)
assertContains(t, joined, fmt.Sprintf("[%d]%s", selectedOwner+1, name))
updated, _ := m.Update(testKey(tea.KeyEnter))
got := updated.(model)
if got.providerName != name || got.modelName != "gpt-5.6" {
t.Fatalf("selection = %s/%s, want %s/gpt-5.6", got.providerName, got.modelName, name)
}
}
}
}
}
}
func pickerGroups(items []pickerItem) []string {
groups := []string{}
seen := map[string]bool{}
+1 -91
View File
@@ -1894,100 +1894,10 @@ func TestModelPickerRowsCarryCapabilityMeta(t *testing.T) {
}
}
func TestModelPickerRowShowsOwnerInMixedGroups(t *testing.T) {
for _, group := range []string{"Recent", "Favorites"} {
for _, owner := range []string{"chatgpt", "openai", "work-openai"} {
for _, selected := range []bool{false, true} {
item := pickerItem{Group: group, Label: "GPT-5.6", Value: "gpt-5.6", Provider: "openai", OwnerProvider: owner, OwnerLabel: owner, Favorite: group == "Favorites"}
want := owner + " · GPT-5.6"
if item.Favorite {
want = "* " + want
}
got := plainRender(t, renderModelPickerRow(60, selected, item))
if !strings.Contains(got, want) {
t.Errorf("group=%s selected=%v: row = %q, want %q", group, selected, got, want)
}
}
}
}
item := pickerItem{Group: "Recent", Value: "custom-model"}
if got := strings.TrimSpace(plainRender(t, renderModelPickerRow(60, false, item))); got != "custom-model" {
t.Fatalf("ownerless fallback row = %q", got)
}
}
func TestModelPickerWidthIncludesOwner(t *testing.T) {
item := pickerItem{Group: "Recent", Label: strings.Repeat("m", 40), OwnerProvider: "subscription-profile", OwnerLabel: "subscription-profile", Favorite: true}
picker := &commandPicker{items: []pickerItem{item}}
width := modelPickerOverlayWidth(120, picker)
got := plainRender(t, renderModelPickerRow(width-4, false, item))
if want := "* subscription-profile · " + item.Label; !strings.Contains(got, want) {
t.Fatalf("row clipped at overlay width %d: %q, want %q", width, got, want)
}
}
// Profile names are unbounded; a long owner must not consume the model name
// once the overlay reaches its maximum width.
func TestModelPickerLongOwnerKeepsModelVisible(t *testing.T) {
owner := strings.Repeat("o", 70)
for _, group := range []string{"Recent", "Favorites"} {
for _, label := range []string{"GPT-5.6", strings.Repeat("m", 40)} {
item := pickerItem{Group: group, Label: label, Value: label, OwnerProvider: owner, OwnerLabel: owner, Favorite: group == "Favorites"}
wide := pickerItem{Group: "Recent", Label: strings.Repeat("w", 100)}
width := modelPickerOverlayWidth(200, &commandPicker{items: []pickerItem{item, wide}})
if width != modelPickerOverlayMaxWidth {
t.Fatalf("overlay width = %d, want the %d cap", width, modelPickerOverlayMaxWidth)
}
got := plainRender(t, renderModelPickerRow(width-4, false, item))
want := "…" + strings.Repeat("o", 19) + " · " + label
if !strings.Contains(got, want) || !strings.Contains(want, "…") {
t.Fatalf("group=%s: row = %q, want %q with the owner capped", group, got, want)
}
}
}
}
func TestModelPickerNarrowOwnerKeepsModelVisible(t *testing.T) {
for _, favorite := range []bool{false, true} {
for _, selected := range []bool{false, true} {
item := pickerItem{Label: "GPT-5.6", OwnerLabel: "subscription-profile", Favorite: favorite}
width := modelPickerOverlayWidth(30, &commandPicker{items: []pickerItem{item}})
got := plainRender(t, renderModelPickerRow(width-4, selected, item))
assertContains(t, got, "GPT-5.6")
assertContains(t, got, "profile · ")
}
}
}
func TestModelPickerNarrowLongModelKeepsSuffix(t *testing.T) {
for _, favorite := range []bool{false, true} {
for _, selected := range []bool{false, true} {
for _, label := range []string{"model-subscription-one", "model-subscription-two"} {
item := pickerItem{Label: label, OwnerLabel: "subscription-profile", OwnerNumber: 1, Favorite: favorite}
got := plainRender(t, renderModelPickerRow(26, selected, item))
assertContains(t, got, label)
assertContains(t, got, "subscription-profile")
if len(strings.Split(got, "\n")) != 2 {
t.Fatalf("expected separate model and owner lines: %q", got)
}
}
}
}
}
func TestModelPickerClippedModelCannotHideOwnerKey(t *testing.T) {
m := limeTestModel()
item := pickerItem{Group: "Recent", Label: strings.Repeat("model-", 10) + "work-subscription-profile", OwnerLabel: "work-subscription-profile", OwnerNumber: 1}
m.picker = &commandPicker{kind: pickerModel, items: []pickerItem{item}}
got := plainRender(t, m.modelPickerOverlay(30))
joined := strings.NewReplacer(" ", "", "\n", "", "│", "").Replace(got)
assertContains(t, joined, "[1]work-subscription-profile")
}
func TestModelPickerRowOmitsProviderTag(t *testing.T) {
// The provider is shown as a section header above each group, so a row renders
// just the model label — no repeated right-aligned provider tag.
item := pickerItem{Group: "anthropic", Label: "Claude Sonnet 4.6", Value: "claude-sonnet-4-6", Provider: "anthropic", OwnerProvider: "anthropic", Remote: true}
item := pickerItem{Label: "Claude Sonnet 4.6", Value: "claude-sonnet-4-6", Provider: "anthropic", Remote: true}
got := plainRender(t, renderModelPickerRow(60, false, item))
if !strings.Contains(got, "Claude Sonnet 4.6") {
t.Fatalf("row = %q, missing model label", got)
+1 -17
View File
@@ -2,7 +2,6 @@ package tui
import (
"context"
"path/filepath"
"strings"
"testing"
@@ -113,22 +112,7 @@ func TestMouseWheelOnClippedFooterStatusDoesNotMoveComposerCursor(t *testing.T)
}
func TestAltScreenTranscriptScrollKeepsFooterFixed(t *testing.T) {
// Cwd is pinned because an empty Options.Cwd makes newModel take the real
// working directory, which at test time is this package inside whoever's
// checkout, and the branch assertion below is not a monotonic function of
// that path's length. The title bar picks the widest candidate that fits:
// workspace carries branch AND path, cwdOnly carries only the path, and
// compactLeft carries only the branch. At width 90 with this branch, a cwd
// of up to 52 bytes keeps branch and path, 53 to 74 falls to path-only and
// loses the branch, and 75 or more falls again to branch-only and gets it
// back. So this failed from a checkout 74 bytes deep and passed from one 80
// bytes deep, while CI sits at 40 and never sees it.
m := newModel(context.Background(), Options{
AltScreen: true,
ProviderName: "openai",
ModelName: "gpt-4.1",
Cwd: filepath.Join(string(filepath.Separator), "workspace", "zero"),
})
m := newModel(context.Background(), Options{AltScreen: true, ProviderName: "openai", ModelName: "gpt-4.1"})
m.width = 90
m.height = 10
m.gitBranch = "feat/pinned-header"
Loaded 100 of 112 files, more files were not shown because too many files have changed in this diff. Show more