Compare commits

..
105 changed files with 1358 additions and 6176 deletions

No files matched your search

+3 -7
View File
@@ -371,13 +371,9 @@ go run ./cmd/zero-perf-bench
Experimental: `ZERO_OPENAI_TURN_SESSION=1` enables the optimized OpenAI turn
session (background connection prewarm + request-prefix telemetry) for headless
`zero exec` runs against official OpenAI profiles. Off by default; `0`, `false`,
or `off` disable it. A/B-benchmark it by running the same `zero-perf-bench` suite
with the variable unset and set.
Native ChatGPT Responses sessions are enabled by default. Set
`ZERO_CHATGPT_TURN_SESSION=0`, `false`, or `off` to restore stateless HTTP/SSE
transport.
`zero exec` runs against official OpenAI profiles. Off by default; `0`/`false`
disable. A/B-benchmark it by running the same `zero-perf-bench` suite with the
variable unset and set.
### Code Quality and Security Checks
+1 -1
View File
@@ -4,7 +4,7 @@ go 1.26.6
require (
charm.land/bubbles/v2 v2.1.1
charm.land/bubbletea/v2 v2.0.9
charm.land/bubbletea/v2 v2.0.8
charm.land/lipgloss/v2 v2.0.5
github.com/Microsoft/go-winio v0.6.2
github.com/alecthomas/chroma/v2 v2.27.0
+4 -2
View File
@@ -1,7 +1,7 @@
charm.land/bubbles/v2 v2.1.1 h1:7r55WzBxpo/R3z98hGmY7KKPd3ET6vsf0Fb9sDHOV60=
charm.land/bubbles/v2 v2.1.1/go.mod h1:GE6M31gaWZVXzGw73OeuTTgy4lX+OtkH0E5ymnNsHxo=
charm.land/bubbletea/v2 v2.0.9 h1:DpJCMWKgzQK8SJv4zbKKFHAI10ymWy/evClPFk0k0f8=
charm.land/bubbletea/v2 v2.0.9/go.mod h1:2SkdgoTXluXJHOUwAoRlRXF/28vklb1rFl6GcgV1/ss=
charm.land/bubbletea/v2 v2.0.8 h1:SxTJMhCAI3lbPmy4SgX5LWZ24AdINr4I6UEqzZvYJuY=
charm.land/bubbletea/v2 v2.0.8/go.mod h1:2SkdgoTXluXJHOUwAoRlRXF/28vklb1rFl6GcgV1/ss=
charm.land/lipgloss/v2 v2.0.5 h1:kbNxgeeUOYv5J0YdpxFjfvf3dFvqH8Aci4zB6xqFtrY=
charm.land/lipgloss/v2 v2.0.5/go.mod h1:9oqhxt4yxIMe6q5A4kHr44DremZk7J9UNh74GlWa5nc=
github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY=
@@ -64,6 +64,8 @@ github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavM
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
golang.org/x/exp v0.0.0-20260611194520-c48552f49976 h1:X8Hz2ImujgbmetVuW+w2YkyZChE3cBpZi2P158rTG9M=
golang.org/x/exp v0.0.0-20260611194520-c48552f49976/go.mod h1:vnf4pv9iKZXY58sQE1L86zmNWJ4159e1RkcWiLCkeEY=
golang.org/x/image v0.44.0 h1:+tDekMZED9+LrtB3G5xzRggpVh9CARjZqROla3R3R+I=
golang.org/x/image v0.44.0/go.mod h1:V8K3KE9KKKE+pLpQDOeN18w9oacNSvy1tDOirTu4xtY=
golang.org/x/image v0.45.0 h1:FMb1nTbH5H9vF55SriQHgFw5GnNL9Jg6L25BwXKzhB0=
golang.org/x/image v0.45.0/go.mod h1:n62x/7RqlwXDvGsSU4u6IUTUf6KghUZ9Bt7cG/T9Fx4=
golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek=
-6
View File
@@ -144,17 +144,11 @@ func estimateToolDefTokens(tools []zeroruntime.ToolDefinition) int {
for _, tool := range tools {
total += ApproxTextTokens(tool.Name)
total += ApproxTextTokens(tool.Description)
total += ApproxTextTokens(string(tool.Type))
if len(tool.Parameters) > 0 {
if encoded, err := json.Marshal(tool.Parameters); err == nil {
total += ApproxTextTokens(string(encoded))
}
}
if tool.Format != nil {
total += ApproxTextTokens(tool.Format.Type)
total += ApproxTextTokens(tool.Format.Syntax)
total += ApproxTextTokens(tool.Format.Definition)
}
total += 4 // per-tool overhead
}
return total
+1 -4
View File
@@ -126,11 +126,8 @@ func capCompactionBrief(brief string) (string, bool) {
func compactionToolCallLine(call zeroruntime.ToolCall) string {
detail := ""
if call.Freeform && call.Name == "apply_patch" {
detail = strings.Join(patchPaths(call.Arguments), ", ")
}
var arguments map[string]any
if detail == "" && json.Unmarshal([]byte(strings.TrimSpace(call.Arguments)), &arguments) == nil {
if json.Unmarshal([]byte(strings.TrimSpace(call.Arguments)), &arguments) == nil {
for _, key := range []string{"file_path", "path", "url", "query", "pattern", "prompt", "description", "question"} {
if value, ok := arguments[key].(string); ok && strings.TrimSpace(value) != "" {
detail = value
+1 -2
View File
@@ -73,7 +73,6 @@ func TestProjectCompactionInputRetainsAskUserExchangeAndSemanticArguments(t *tes
{Role: zeroruntime.MessageRoleTool, ToolCallID: "ask", Content: "Which database should remain?: Postgres only; do not touch MySQL."},
{Role: zeroruntime.MessageRoleAssistant, ToolCalls: []zeroruntime.ToolCall{
{ID: "patch", Name: "apply_patch", Arguments: `{"patch":"*** Begin Patch\n*** Update File: internal/db.go\n*** End Patch"}`},
{ID: "raw-patch", Name: "apply_patch", Freeform: true, Arguments: "*** Begin Patch\n*** Update File: internal/api.go\n*** Add File: internal/config.go\n*** End Patch\n"},
{ID: "fetch", Name: "web_fetch", Arguments: `{"url":"https://example.com/spec"}`},
{ID: "task", Name: "task", Arguments: `{"prompt":"Inspect the Windows implementation"}`},
}},
@@ -81,7 +80,7 @@ func TestProjectCompactionInputRetainsAskUserExchangeAndSemanticArguments(t *tes
}
content := projectCompactionInput(messages).messages[0].Content
for _, want := range []string{"Which database should remain?", "Postgres only; do not touch MySQL", "internal/db.go", "internal/api.go", "internal/config.go", "https://example.com/spec", "Inspect the Windows implementation"} {
for _, want := range []string{"Which database should remain?", "Postgres only; do not touch MySQL", "internal/db.go", "https://example.com/spec", "Inspect the Windows implementation"} {
if !strings.Contains(content, want) {
t.Fatalf("projection missing %q: %q", want, content)
}
-21
View File
@@ -275,27 +275,6 @@ func TestEstimateToolDefTokensCountsDefinitions(t *testing.T) {
if estimateToolDefTokens(one) <= 0 {
t.Fatal("a tool definition (name + description + schema) should estimate > 0 tokens")
}
typed := append([]zeroruntime.ToolDefinition(nil), one...)
typed[0].Type = zeroruntime.ToolDefinitionFreeform
if estimateToolDefTokens(typed) <= estimateToolDefTokens(one) {
t.Fatal("a non-empty tool type must increase the estimated request cost")
}
for _, test := range []struct {
name string
format zeroruntime.ToolDefinitionFormat
}{
{name: "type", format: zeroruntime.ToolDefinitionFormat{Type: "grammar"}},
{name: "syntax", format: zeroruntime.ToolDefinitionFormat{Syntax: "lark"}},
{name: "definition", format: zeroruntime.ToolDefinitionFormat{Definition: "start: WORD"}},
} {
t.Run("format "+test.name, func(t *testing.T) {
formatted := append([]zeroruntime.ToolDefinition(nil), one...)
formatted[0].Format = &test.format
if estimateToolDefTokens(formatted) <= estimateToolDefTokens(one) {
t.Fatalf("a non-empty format %s must increase the estimated request cost", test.name)
}
})
}
two := append(append([]zeroruntime.ToolDefinition{}, one...), zeroruntime.ToolDefinition{
Name: "write_file",
Description: "Write contents to a file in the workspace.",
-4
View File
@@ -102,10 +102,6 @@ func copyToolDefinitions(toolDefs []zeroruntime.ToolDefinition) []zeroruntime.To
for index, definition := range toolDefs {
copied[index] = definition
copied[index].Parameters = copySchemaMap(definition.Parameters)
if definition.Format != nil {
format := *definition.Format
copied[index].Format = &format
}
}
return copied
}
-3
View File
@@ -15,7 +15,6 @@ func TestContextPlannerPreservesProviderRequest(t *testing.T) {
}
toolDefs := []zeroruntime.ToolDefinition{{
Name: "read_file", Description: "Read a file",
Format: &zeroruntime.ToolDefinitionFormat{Type: "grammar", Syntax: "lark", Definition: "start: PATH"},
Parameters: map[string]any{
"type": "object",
"required": []string{"path"},
@@ -54,12 +53,10 @@ func TestContextPlannerPreservesProviderRequest(t *testing.T) {
toolDefs[0].Parameters["required"].([]string)[0] = "other"
toolDefs[0].Parameters["properties"].(map[string]any)["path"].(map[string]any)["type"] = "integer"
toolDefs[0].Parameters["anyOf"].([]any)[0].(map[string]any)["type"] = "array"
toolDefs[0].Format.Definition = "changed"
if plan.Request.Messages[1].Content != "inspect this" || plan.Request.Messages[1].Images[0].Data[0] != 1 {
t.Fatalf("planned messages alias caller state: %#v", plan.Request.Messages[1])
}
if plan.Request.Tools[0].Name != "read_file" ||
plan.Request.Tools[0].Format == nil || plan.Request.Tools[0].Format.Definition != "start: PATH" ||
plan.Request.Tools[0].Parameters["required"].([]string)[0] != "path" ||
plan.Request.Tools[0].Parameters["properties"].(map[string]any)["path"].(map[string]any)["type"] != "string" ||
plan.Request.Tools[0].Parameters["anyOf"].([]any)[0].(map[string]any)["type"] != "object" {
-378
View File
@@ -1,378 +0,0 @@
package agent
import (
"context"
"os"
"path/filepath"
"runtime"
"strings"
"testing"
"github.com/Gitlawb/zero/internal/sandbox"
"github.com/Gitlawb/zero/internal/tools"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
type customApplyPatchTool struct {
calls int
}
type fixedPathScope struct {
roots []string
}
func (scope fixedPathScope) Roots() []string {
return append([]string(nil), scope.roots...)
}
func newFixedPathScope(t *testing.T, roots ...string) fixedPathScope {
t.Helper()
resolved := make([]string, 0, len(roots))
for _, root := range roots {
absolute, err := filepath.Abs(root)
if err != nil {
t.Fatalf("Abs(%q): %v", root, err)
}
if physical, err := filepath.EvalSymlinks(absolute); err == nil {
absolute = physical
}
resolved = append(resolved, absolute)
}
return fixedPathScope{roots: resolved}
}
func (*customApplyPatchTool) Name() string { return "apply_patch" }
func (*customApplyPatchTool) Description() string { return "custom JSON tool" }
func (*customApplyPatchTool) Parameters() tools.Schema {
return tools.Schema{Type: "object", Properties: map[string]tools.PropertySchema{"value": {Type: "string"}}}
}
func (*customApplyPatchTool) Safety() tools.Safety {
return tools.Safety{Permission: tools.PermissionAllow, SideEffect: tools.SideEffectRead}
}
func (tool *customApplyPatchTool) Run(context.Context, map[string]any) tools.Result {
tool.calls++
return tools.Result{Status: tools.StatusOK}
}
func TestFreeformApplyPatchUsesExistingToolExecutionPath(t *testing.T) {
root := t.TempDir()
registry := tools.NewRegistry()
registry.Register(tools.NewScopedApplyPatchTool(root, nil))
patch := "*** Begin Patch\n*** Add File: hello.txt\n+hello\n*** End Patch\n"
result, err := executeToolCall(context.Background(), registry, zeroruntime.ToolCall{
ID: "call-1",
Name: "apply_patch",
Arguments: patch,
Freeform: true,
}, PermissionModeUnsafe, Options{Cwd: root, FileTracker: tools.NewFileTracker()})
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusOK {
t.Fatalf("freeform apply_patch status = %s: %s", result.Status, result.Output)
}
content, err := os.ReadFile(filepath.Join(root, "hello.txt"))
if err != nil {
t.Fatalf("read patched file: %v", err)
}
if string(content) != "hello\n" {
t.Fatalf("patched content = %q, want hello newline", content)
}
}
func TestFreeformApplyPatchSupportsGrantedExtraRoot(t *testing.T) {
root := t.TempDir()
extra := t.TempDir()
scope, err := sandbox.NewScope(root, []string{extra})
if err != nil {
t.Fatalf("NewScope: %v", err)
}
registry := tools.NewRegistry()
registry.Register(tools.NewScopedApplyPatchTool(root, scope))
target := filepath.Join(extra, "hello.txt")
patch := "*** Begin Patch\n*** Add File: " + target + "\n+hello\n*** End Patch\n"
result, err := executeToolCall(context.Background(), registry, zeroruntime.ToolCall{
ID: "call-1", Name: "apply_patch", Arguments: patch, Freeform: true,
}, PermissionModeUnsafe, Options{Cwd: root, FileTracker: tools.NewFileTracker()})
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusOK {
t.Fatalf("freeform extra-root apply_patch status = %s: %s", result.Status, result.Output)
}
content, err := os.ReadFile(target)
if err != nil {
t.Fatalf("read extra-root patched file: %v", err)
}
if string(content) != "hello\n" {
t.Fatalf("patched content = %q, want hello newline", content)
}
}
func TestFreeformApplyPatchRejectsMixedWorkspaceAndExtraRoots(t *testing.T) {
for _, extraHasMatchingPath := range []bool{false, true} {
name := "extra target missing"
if extraHasMatchingPath {
name = "extra target exists"
}
t.Run(name, func(t *testing.T) {
workspace := t.TempDir()
extra := t.TempDir()
if err := os.MkdirAll(filepath.Join(workspace, "src"), 0o755); err != nil {
t.Fatal(err)
}
workspaceTarget := filepath.Join(workspace, "src", "a.go")
if err := os.WriteFile(workspaceTarget, []byte("old\n"), 0o644); err != nil {
t.Fatal(err)
}
extraTarget := filepath.Join(extra, "src", "a.go")
if extraHasMatchingPath {
if err := os.MkdirAll(filepath.Dir(extraTarget), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(extraTarget, []byte("old\n"), 0o644); err != nil {
t.Fatal(err)
}
}
scope, err := sandbox.NewScope(workspace, []string{extra})
if err != nil {
t.Fatal(err)
}
registry := tools.NewRegistry()
registry.Register(tools.NewScopedApplyPatchTool(workspace, scope))
extraCreated := filepath.Join(extra, "b.go")
patch := strings.Join([]string{
"*** Begin Patch",
"*** Update File: src/a.go",
"@@",
"-old",
"+new",
"*** Add File: " + extraCreated,
"+extra",
"*** End Patch",
}, "\n")
result, err := executeToolCall(context.Background(), registry, zeroruntime.ToolCall{
ID: "call-1", Name: "apply_patch", Arguments: patch, Freeform: true,
}, PermissionModeUnsafe, Options{Cwd: workspace, FileTracker: tools.NewFileTracker()})
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusError || !strings.Contains(result.Output, "mixes workspace-relative paths with an extra write root") {
t.Fatalf("mixed-root patch result = %s: %s", result.Status, result.Output)
}
if content, err := os.ReadFile(workspaceTarget); err != nil || string(content) != "old\n" {
t.Fatalf("workspace target changed: content=%q err=%v", content, err)
}
if extraHasMatchingPath {
if content, err := os.ReadFile(extraTarget); err != nil || string(content) != "old\n" {
t.Fatalf("extra-root target changed: content=%q err=%v", content, err)
}
}
if _, err := os.Stat(extraCreated); !os.IsNotExist(err) {
t.Fatalf("mixed-root patch created extra file: %v", err)
}
})
}
}
func TestFreeformApplyPatchAcceptsOneSemanticRoot(t *testing.T) {
for _, test := range []struct {
name string
useExtraRoot bool
}{
{name: "workspace"},
{name: "extra root", useExtraRoot: true},
} {
t.Run(test.name, func(t *testing.T) {
workspace := t.TempDir()
extra := t.TempDir()
scope, err := sandbox.NewScope(workspace, []string{extra})
if err != nil {
t.Fatal(err)
}
root := workspace
if test.useExtraRoot {
root = extra
}
if err := os.MkdirAll(filepath.Join(root, "src"), 0o755); err != nil {
t.Fatal(err)
}
target := filepath.Join(root, "src", "a.go")
if err := os.WriteFile(target, []byte("old\n"), 0o644); err != nil {
t.Fatal(err)
}
updatePath := target
if !test.useExtraRoot {
updatePath = filepath.Join("src", "a.go")
}
created := filepath.Join(root, "b.go")
patch := strings.Join([]string{
"*** Begin Patch",
"*** Update File: " + updatePath,
"@@",
"-old",
"+new",
"*** Add File: " + created,
"+created",
"*** End Patch",
}, "\n")
registry := tools.NewRegistry()
registry.Register(tools.NewScopedApplyPatchTool(workspace, scope))
result, err := executeToolCall(context.Background(), registry, zeroruntime.ToolCall{
ID: "call-1", Name: "apply_patch", Arguments: patch, Freeform: true,
}, PermissionModeUnsafe, Options{Cwd: workspace, FileTracker: tools.NewFileTracker()})
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusOK {
t.Fatalf("single-root patch status = %s: %s", result.Status, result.Output)
}
if content, err := os.ReadFile(target); err != nil || string(content) != "new\n" {
t.Fatalf("updated content=%q err=%v", content, err)
}
if content, err := os.ReadFile(created); err != nil || string(content) != "created\n" {
t.Fatalf("created content=%q err=%v", content, err)
}
})
}
}
func TestFreeformApplyPatchRejectsAbsolutePathOutsideGrantedRoots(t *testing.T) {
root := t.TempDir()
granted := t.TempDir()
outside := t.TempDir()
scope := newFixedPathScope(t, root, granted)
registry := tools.NewRegistry()
registry.Register(tools.NewScopedApplyPatchTool(root, scope))
target := filepath.Join(outside, "blocked.txt")
patch := "*** Begin Patch\n*** Add File: " + target + "\n+blocked\n*** End Patch\n"
result, err := executeToolCall(context.Background(), registry, zeroruntime.ToolCall{
ID: "call-1", Name: "apply_patch", Arguments: patch, Freeform: true,
}, PermissionModeUnsafe, Options{Cwd: root})
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusError {
t.Fatalf("ungranted absolute patch status = %s, want error: %s", result.Status, result.Output)
}
if _, err := os.Stat(target); !os.IsNotExist(err) {
t.Fatalf("ungranted absolute patch created target: %v", err)
}
}
func TestPreparedFreeformPatchRejectsIntermediateSymlinkSwap(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("creating directory symlinks requires privileges on Windows")
}
root := t.TempDir()
extra := t.TempDir()
outside := t.TempDir()
intermediate := filepath.Join(extra, "nested")
if err := os.Mkdir(intermediate, 0o755); err != nil {
t.Fatalf("Mkdir: %v", err)
}
scope := newFixedPathScope(t, root, extra)
tool := tools.NewScopedApplyPatchTool(root, scope)
target := filepath.Join(intermediate, "blocked.txt")
patch := "*** Begin Patch\n*** Add File: " + target + "\n+blocked\n*** End Patch\n"
args, err := tools.PrepareFreeformApplyPatchArguments(tool, patch)
if err != nil {
t.Fatalf("PrepareFreeformApplyPatchArguments: %v", err)
}
if err := os.Remove(intermediate); err != nil {
t.Fatalf("Remove intermediate: %v", err)
}
if err := os.Symlink(outside, intermediate); err != nil {
t.Fatalf("Symlink: %v", err)
}
result := tool.Run(context.Background(), args)
if result.Status != tools.StatusError {
t.Fatalf("symlink-swapped patch status = %s, want error: %s", result.Status, result.Output)
}
if _, err := os.Stat(filepath.Join(outside, "blocked.txt")); !os.IsNotExist(err) {
t.Fatalf("symlink-swapped patch wrote outside root: %v", err)
}
}
func TestUnknownFreeformToolFailsClosed(t *testing.T) {
result, err := executeToolCall(context.Background(), tools.NewRegistry(), zeroruntime.ToolCall{
ID: "call-1", Name: "unknown", Arguments: "raw", Freeform: true,
}, PermissionModeUnsafe, Options{})
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusError {
t.Fatalf("unknown freeform status = %s, want error", result.Status)
}
}
func TestCustomApplyPatchCannotUseFreeformContract(t *testing.T) {
custom := &customApplyPatchTool{}
registry := tools.NewRegistry()
registry.Register(custom)
result, err := executeToolCall(context.Background(), registry, zeroruntime.ToolCall{
ID: "call-1", Name: "apply_patch", Arguments: "raw patch", Freeform: true,
}, PermissionModeUnsafe, Options{})
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusError || !strings.Contains(result.Output, "Unsupported freeform tool call") {
t.Fatalf("custom freeform apply_patch result = %#v, want unsupported error", result)
}
if custom.calls != 0 {
t.Fatalf("custom apply_patch executed %d times", custom.calls)
}
}
func TestDeniedFreeformApplyPatchDoesNotWrite(t *testing.T) {
root := t.TempDir()
registry := tools.NewRegistry()
registry.Register(tools.NewScopedApplyPatchTool(root, nil))
patch := "*** Begin Patch\n*** Add File: denied.txt\n+must not exist\n*** End Patch\n"
result, err := executeToolCall(context.Background(), registry, zeroruntime.ToolCall{
ID: "call-1",
Name: "apply_patch",
Arguments: patch,
Freeform: true,
}, PermissionModeAsk, Options{
Cwd: root,
OnPermissionRequest: func(context.Context, PermissionRequest) (PermissionDecision, error) {
return PermissionDecision{Action: PermissionDecisionDeny, Reason: "test denial"}, nil
},
})
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusError {
t.Fatalf("denied freeform status = %s, want error", result.Status)
}
if _, err := os.Stat(filepath.Join(root, "denied.txt")); !os.IsNotExist(err) {
t.Fatalf("denied freeform patch changed the workspace: %v", err)
}
}
func TestHistorySafeToolCallsPreservesRawFreeformArguments(t *testing.T) {
raw := "*** Begin Patch\n*** Update File: internal/main.go\n@@\n-old\n+new\n*** End Patch\n"
safe := historySafeToolCalls([]ToolCall{{
ID: "call-1", ProviderCallID: "provider-1", Name: "apply_patch", Arguments: raw, Freeform: true,
}})
if len(safe) != 1 || safe[0].Arguments != raw || safe[0].ProviderCallID != "provider-1" || !safe[0].Freeform {
t.Fatalf("freeform history changed raw call: %#v", safe)
}
}
func TestAbortedToolResultPreservesProviderCallID(t *testing.T) {
messages := appendAbortedToolResults(nil, []ToolCall{{
ID: "call-1", ProviderCallID: "provider-1", Name: "read_file",
}})
if len(messages) != 1 || messages[0].ToolCallID != "call-1" || messages[0].ToolCallProviderID != "provider-1" || !messages[0].IsError {
t.Fatalf("aborted tool result lost provider identity: %#v", messages)
}
}
+9 -8
View File
@@ -33,11 +33,11 @@ const (
// knows before spending more tool turns.
toolOnlyProgressReminderAt = 6
// planReminderTurn and planReminderToolThreshold keep bounded implementation
// work plan-free while still nudging a genuinely sustained run that has not
// established any structured progress. Both thresholds must be met.
planReminderTurn = 7
planReminderToolThreshold = 7
// planReminderTurn is the turn (1-based) by the end of which a multi-step
// task should have called update_plan; if it hasn't, a one-time reminder is
// injected. Set to 3 (not 2) so short, legitimate two-step tasks finish
// without a spurious planning nag.
planReminderTurn = 3
// planToolName is the planning tool the loop watches for by name.
planToolName = "update_plan"
@@ -591,12 +591,13 @@ func (state *guardState) planReminder(turn int) string {
return planStaleReminder(state.toolCallsSincePlanUpdate)
}
// NOT-CALLED reminder: once both the turn and tool-call thresholds are met,
// sustained work should have called update_plan. One-shot for the whole run.
// NOT-CALLED reminder: by the end of planReminderTurn the model should have
// called update_plan if it's doing a multi-step task (>=1 other tool call).
// One-shot for the whole run.
if !state.notCalledReminderSent &&
!state.planEverCalled &&
turn >= planReminderTurn &&
state.totalToolCalls >= planReminderToolThreshold {
state.totalToolCalls >= 1 {
state.notCalledReminderSent = true
return planNotCalledReminder()
}
+8 -43
View File
@@ -257,16 +257,19 @@ func TestRunDoesNotCountDroppedToolCallTurnsAsEmpty(t *testing.T) {
}
func TestRunInjectsPlanNotCalledReminderForMultiStepTask(t *testing.T) {
const expectedThreshold = 7
if planReminderToolThreshold != expectedThreshold {
t.Fatalf("plan reminder tool threshold = %d, want %d", planReminderToolThreshold, expectedThreshold)
}
root := t.TempDir()
writeAgentTestFile(t, root+"/notes.txt", "alpha")
registry := tools.NewRegistry()
registry.Register(tools.NewScopedReadFileTool(root, nil))
provider := &mockProvider{turns: append(repeatedReadTurns(expectedThreshold), textTurn("done"))}
provider := &mockProvider{
turns: [][]zeroruntime.StreamEvent{
toolTurn("call-1", "read_file", `{"path":"notes.txt"}`), // turn 1: other tool call
toolTurn("call-2", "read_file", `{"path":"notes.txt"}`), // turn 2: still no update_plan
toolTurn("call-3", "read_file", `{"path":"notes.txt"}`), // turn 3: reminder fires here
textTurn("done"),
},
}
result, err := Run(context.Background(), "go", provider, Options{
Registry: registry,
@@ -284,44 +287,6 @@ func TestRunInjectsPlanNotCalledReminderForMultiStepTask(t *testing.T) {
}
}
func TestRunDoesNotInjectPlanReminderForBoundedToolSequence(t *testing.T) {
const expectedThreshold = 7
root := t.TempDir()
writeAgentTestFile(t, root+"/notes.txt", "alpha")
registry := tools.NewRegistry()
registry.Register(tools.NewScopedReadFileTool(root, nil))
provider := &mockProvider{turns: append(repeatedReadTurns(expectedThreshold-1), textTurn("done"))}
result, err := Run(context.Background(), "go", provider, Options{Registry: registry, MaxTurns: 12})
if err != nil {
t.Fatal(err)
}
if result.FinalAnswer != "done" {
t.Fatalf("expected final answer, got %q", result.FinalAnswer)
}
if count := countUserMessagesContaining(result.Messages, planNotCalledReminderMarker); count != 0 {
t.Fatalf("bounded tool sequence must not draw a plan reminder, got %d", count)
}
}
func TestPlanReminderWaitsForTurnThresholdAfterEnoughToolCalls(t *testing.T) {
state := guardState{totalToolCalls: 7}
if got := state.planReminder(6); got != "" {
t.Fatalf("planReminder(6) = %q, want no reminder before turn threshold", got)
}
if state.notCalledReminderSent {
t.Fatal("early turn marked the not-called reminder as sent")
}
}
func repeatedReadTurns(count int) [][]zeroruntime.StreamEvent {
turns := make([][]zeroruntime.StreamEvent, 0, count)
for range count {
turns = append(turns, toolTurn("call", "read_file", `{"path":"notes.txt"}`))
}
return turns
}
func TestRunDoesNotInjectPlanReminderForTrivialTask(t *testing.T) {
root := t.TempDir()
writeAgentTestFile(t, root+"/notes.txt", "alpha")
+12 -41
View File
@@ -141,7 +141,6 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
}
options.Trace.Counter(trace.CounterInputTokens, int64(usage.InputTokens))
options.Trace.Counter(trace.CounterCachedInputTokens, int64(usage.CachedInputTokens))
options.Trace.Counter(trace.CounterCacheWriteTokens, int64(usage.CacheWriteTokens))
options.Trace.Counter(trace.CounterOutputTokens, int64(usage.OutputTokens))
}
}
@@ -698,12 +697,11 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
turnRequestedModel = toolResult.RequestedModel
}
messages = append(messages, zeroruntime.Message{
Role: zeroruntime.MessageRoleTool,
Content: toolResult.ModelOutput(),
ToolCallID: toolResult.ToolCallID,
ToolCallProviderID: call.ProviderCallID,
IsError: toolResult.Status == tools.StatusError,
ChangedFiles: append([]string(nil), toolResult.ChangedFiles...),
Role: zeroruntime.MessageRoleTool,
Content: toolResult.ModelOutput(),
ToolCallID: toolResult.ToolCallID,
IsError: toolResult.Status == tools.StatusError,
ChangedFiles: append([]string(nil), toolResult.ChangedFiles...),
})
// Images ride a following USER message rather than the tool result
// above. Every provider drops images on a tool-role message —
@@ -1037,9 +1035,6 @@ func historySafeToolCalls(calls []ToolCall) []ToolCall {
safe := make([]ToolCall, len(calls))
for i, call := range calls {
safe[i] = call
if call.Freeform {
continue
}
args := strings.TrimSpace(call.Arguments)
switch {
case args == "":
@@ -1123,28 +1118,8 @@ func decodeToolArguments(arguments string, v any) error {
}
func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCall, permissionMode PermissionMode, options Options) (ToolResult, error) {
tool, toolFound := registry.Get(call.Name)
args := map[string]any{}
if call.Freeform {
if !toolFound || !tools.IsBuiltInApplyPatch(tool) {
return ToolResult{
ToolCallID: call.ID,
Name: call.Name,
Status: tools.StatusError,
Output: "Error: Unsupported freeform tool call for " + call.Name + ".",
}, nil
}
prepared, err := tools.PrepareFreeformApplyPatchArguments(tool, call.Arguments)
if err != nil {
return ToolResult{
ToolCallID: call.ID,
Name: call.Name,
Status: tools.StatusError,
Output: "Error: Invalid freeform apply_patch input: " + err.Error(),
}, nil
}
args = prepared
} else if call.Arguments != "" {
if call.Arguments != "" {
if err := decodeToolArguments(call.Arguments, &args); err != nil {
return ToolResult{
ToolCallID: call.ID,
@@ -1172,6 +1147,7 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
DenialReason: DenialFiltered,
}, nil
}
tool, toolFound := registry.Get(call.Name)
if (permissionMode == PermissionModeSpecDraft || permissionMode == PermissionModePlan) && toolFound && !ToolAdvertised(tool, permissionMode) {
modeName := string(permissionMode)
return ToolResult{
@@ -3234,15 +3210,11 @@ func cachedRuntimeToolDefinition(defCache map[string]zeroruntime.ToolDefinition,
// runtimeToolDefinition renders a tool's advertised definition (name, description,
// JSON-schema parameters) as sent to the provider.
func runtimeToolDefinition(tool tools.Tool) zeroruntime.ToolDefinition {
definition := zeroruntime.ToolDefinition{
return zeroruntime.ToolDefinition{
Name: tool.Name(),
Description: tool.Description(),
Parameters: schemaToRuntimeMap(tool.Parameters()),
}
if tools.IsBuiltInApplyPatch(tool) {
definition.Type = zeroruntime.ToolDefinitionFreeform
}
return definition
}
func ToolVisible(tool tools.Tool, permissionMode PermissionMode, enabledTools []string, disabledTools []string) bool {
@@ -3397,11 +3369,10 @@ func loadedToolsFromResult(meta map[string]string) []string {
func appendAbortedToolResults(messages []Message, remaining []ToolCall) []Message {
for _, call := range remaining {
messages = append(messages, zeroruntime.Message{
Role: zeroruntime.MessageRoleTool,
Content: abortedToolResultNotice,
ToolCallID: call.ID,
ToolCallProviderID: call.ProviderCallID,
IsError: true,
Role: zeroruntime.MessageRoleTool,
Content: abortedToolResultNotice,
ToolCallID: call.ID,
IsError: true,
})
}
return messages
+1 -4
View File
@@ -4096,7 +4096,7 @@ func TestRunDoesNotFlagCleanToolOutput(t *testing.T) {
// caller's callback AND stamps token counters, and the run completes.
func TestRunTracingWrapperStampsUsage(t *testing.T) {
provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{{
{Type: zeroruntime.StreamEventUsage, Usage: zeroruntime.Usage{PromptTokens: 100, CachedInputTokens: 20, CacheWriteTokens: 10, CompletionTokens: 40}},
{Type: zeroruntime.StreamEventUsage, Usage: zeroruntime.Usage{InputTokens: 100, CachedInputTokens: 20, OutputTokens: 40}},
{Type: zeroruntime.StreamEventText, Content: "done"},
{Type: zeroruntime.StreamEventDone},
}}}
@@ -4135,9 +4135,6 @@ func TestRunTracingWrapperStampsUsage(t *testing.T) {
if got := tr.Counter(trace.CounterCachedInputTokens); got != 20 {
t.Fatalf("cached input token counter = %d, want 20", got)
}
if got := tr.Counter(trace.CounterCacheWriteTokens); got != 10 {
t.Fatalf("cache-write token counter = %d, want 10", got)
}
if got := tr.Counter(trace.CounterOutputTokens); got != 40 {
t.Fatalf("output token counter = %d, want 40", got)
}
-13
View File
@@ -31,19 +31,6 @@ func (t countingSchemaTool) Run(_ context.Context, _ map[string]any) tools.Resul
return tools.Result{Status: tools.StatusOK}
}
func TestRuntimeToolDefinitionMarksOnlyBuiltInApplyPatchFreeform(t *testing.T) {
builtIn := runtimeToolDefinition(tools.NewScopedApplyPatchTool(t.TempDir(), nil))
if builtIn.Type != zeroruntime.ToolDefinitionFreeform {
t.Fatalf("built-in apply_patch type = %q, want freeform", builtIn.Type)
}
calls := 0
custom := runtimeToolDefinition(countingSchemaTool{name: "apply_patch", calls: &calls})
if custom.Type != "" {
t.Fatalf("custom apply_patch type = %q, want legacy function shape", custom.Type)
}
}
// The cache renders a tool's schema once and reuses it across calls, and its
// output is identical to the uncached path.
func TestPartitionToolsCacheRendersOnceAndMatchesUncached(t *testing.T) {
+4 -4
View File
@@ -68,7 +68,7 @@ func TestEagerToolSchemaTokenBudget(t *testing.T) {
}
}
func TestAgentAdvertisesBothEditTools(t *testing.T) {
func TestAgentAdvertisesPatchInsteadOfAmbiguousStringReplacement(t *testing.T) {
registry := tools.NewRegistry()
for _, tool := range tools.CoreToolsScoped(t.TempDir(), nil) {
registry.Register(tool)
@@ -79,9 +79,9 @@ func TestAgentAdvertisesBothEditTools(t *testing.T) {
names[definition.Name] = true
}
if !names["apply_patch"] {
t.Fatal("agent must retain apply_patch for multi-hunk changes")
t.Fatal("agent must retain apply_patch for existing-file changes")
}
if !names["edit_file"] {
t.Fatal("agent must receive edit_file for targeted exact replacements")
if names["edit_file"] {
t.Fatal("agent must not receive the ambiguous string-replacement tool")
}
}
-6
View File
@@ -58,18 +58,12 @@ func toolSubstrings(exposed []zeroruntime.ToolDefinition) (toolsSubstr, schemaSu
for _, def := range exposed {
writeLengthPrefixed(&toolsSB, def.Name)
writeLengthPrefixed(&toolsSB, def.Description)
writeLengthPrefixed(&toolsSB, string(def.Type))
// Parameters is rendered to a canonical JSON string per tool. The
// schema-render cache in internal/agent (used by partitionToolsCached)
// guarantees the same render is byte-identical across turns for the
// same tool, so this substring is a stable hash input.
writeLengthPrefixed(&schemaSB, def.Name)
writeLengthPrefixed(&schemaSB, canonicalSchemaString(def.Parameters))
if def.Format != nil {
writeLengthPrefixed(&schemaSB, def.Format.Type)
writeLengthPrefixed(&schemaSB, def.Format.Syntax)
writeLengthPrefixed(&schemaSB, def.Format.Definition)
}
}
return toolsSB.String(), schemaSB.String()
}
-32
View File
@@ -73,38 +73,6 @@ func TestComputePrefixFingerprintSchemaChangeChangesSchemaHash(t *testing.T) {
}
}
func TestComputePrefixFingerprintIncludesToolTypeAndFormat(t *testing.T) {
opts := Options{Cwd: t.TempDir(), SystemPrompt: "core"}
base := zeroruntime.ToolDefinition{Name: "apply_patch", Description: "Patch files"}
baseFingerprint := ComputePrefixFingerprint(opts, []zeroruntime.ToolDefinition{base})
typed := base
typed.Type = zeroruntime.ToolDefinitionFreeform
if got := ComputePrefixFingerprint(opts, []zeroruntime.ToolDefinition{typed}); got.ToolsHash == baseFingerprint.ToolsHash {
t.Fatal("ToolsHash must change when tool Type changes")
}
formatCases := []struct {
name string
format zeroruntime.ToolDefinitionFormat
}{
{name: "type", format: zeroruntime.ToolDefinitionFormat{Type: "grammar"}},
{name: "syntax", format: zeroruntime.ToolDefinitionFormat{Syntax: "lark"}},
{name: "definition", format: zeroruntime.ToolDefinitionFormat{Definition: "start: PATCH"}},
{name: "empty", format: zeroruntime.ToolDefinitionFormat{}},
}
for _, test := range formatCases {
t.Run(test.name, func(t *testing.T) {
withFormat := base
withFormat.Format = &test.format
got := ComputePrefixFingerprint(opts, []zeroruntime.ToolDefinition{withFormat})
if got.SchemaHash == baseFingerprint.SchemaHash {
t.Fatalf("SchemaHash must change for %s format input", test.name)
}
})
}
}
// TestComputePrefixFingerprintSystemPromptChangeChangesBaseHash asserts the
// most common drift path: a system-prompt edit (e.g. an agent.md update, a
// model addendum change) must move the base-instructions hash, which moves
+24 -32
View File
@@ -30,10 +30,13 @@ work.
read-before-edit discipline: inspect the target file and nearby callers,
tests, or config before you modify behavior. Never edit a file you have not
read.
2. **Plan.** Use update_plan for multi-component, sequencing-uncertain, or long
work. Skip it for bounded changes in one component. Update only when the
overall phase changes, keep at most one item in_progress, and never create a
plan after the work is complete.
2. **Plan.** Use update_plan for implementation or investigation that has at
least three meaningful dependent steps, then keep it current at milestones.
Mark completed work and the next in-progress step together when practical;
do not spend separate turns updating it after every file or command. Keep at
most one item in_progress. Skip plans for simple lookups, explanations, code
navigation, and other short tasks. Never create a plan after the work is
already complete merely to describe what you did.
3. **Implement.** Make focused changes that match the surrounding code's style,
naming, and conventions. Prefer the smallest change that fully solves the
problem. Avoid broad refactors, unrelated rewrites, dependency churn, and
@@ -49,48 +52,38 @@ work.
- Choose the narrowest tool that safely accomplishes the step. Prefer native
file tools - read_file, read_minified_file, list_directory, glob, grep,
edit_file, apply_patch, write_file - over shelling out to
write_file, apply_patch - over shelling out to
cat/sed/awk/python for file operations.
They are safer, reviewable, and produce clean diffs.
- Use read_minified_file only to explore large or unfamiliar source. Use
read_file for small or likely edit targets and exact text, comments, or line
numbers. Do not read both unless a new need for exact content appears.
- Prefer read_minified_file when initially exploring source code; it preserves
code while removing comments and redundant whitespace. Use read_file when
exact text, comments, or line numbers are needed.
- Keep edits focused and reviewable. A single patch may update several related
files when they form one coherent change; do not hide unrelated edits in a
bulk shell or script rewrite.
- For edits to existing files, use edit_file for a targeted change (old_string
must match the file exactly and be unique, so include a few surrounding
lines) and apply_patch when one coherent change spans several hunks or
files. Use write_file only to create a file or when most of it changes; do
not rewrite a whole file to change a few lines. Match the existing
indentation, imports, and idioms. Match the file's comment density: do not
add explanatory comments unless the user asks or the code is already
- For edits to existing files, prefer apply_patch with minimal, targeted hunks
and enough unchanged context to identify the intended location. Match the
existing indentation, imports, and idioms. Match the file's comment density:
do not add explanatory comments unless the user asks or the code is already
comment-dense.
- A successful edit result already confirms the change; do not re-read a file
just to verify an edit that succeeded. If an edit fails, read the error, fix
the old_string or hunk, and retry the same tool rather than switching to a
full rewrite.
- Solve the problem as posed, not a more general version of it. Add no
speculative abstraction, configurability, or handling for cases that cannot
occur, and nothing the user did not ask for. A small diff can still be
over-built; if a 200-line solution could be 50, rewrite it. Once the required
behavior passes validation, stop instead of adding optional cleanup.
over-built; if a 200-line solution could be 50, rewrite it.
- Preserve behavior you were not asked to change. Do not delete or rewrite code
you did not author unless the task requires it; if you must, say so.
## Testing gate (mandatory)
- After code changes, run the project's documented appropriate validators.
Scope iteration checks, combine compatible commands when diagnostics stay
useful, and run the final set once after the last edit. Rerun only after a
change or failure.
- If unsure which entry points apply, inspect the Makefile, manifests, or CI;
do not invent a parallel flow.
- After any change to code, verify after edits by running the project's
validators before you summarize or commit: tests, type-checks, linters, and/or
the build, as appropriate. Scope them to the change while iterating; reserve
full-suite runs for milestones.
- If you are unsure which validators apply, search the repo (Makefile, package
manifests, CI config) to find them.
- Never claim a task is done, and never commit, while validators are failing. If
they fail, fix the cause and rerun; do not paper over it. If you could not run
a validator, say so explicitly rather than implying success.
- A successful native patch confirms the edit. Reread only when the next change
needs exact content or the result was ambiguous.
## Tool use
@@ -112,7 +105,7 @@ work.
you need to clean up a running foreground command yourself, use write_stdin.
- write_stdin's session_id is only ever an id returned by a still-running
exec_command; never guess or probe ids. If you have no such session, start one
with exec_command, or use edit_file/apply_patch/write_file for file changes.
with exec_command, or use write_file/apply_patch for file changes.
- write_stdin with empty input polls an existing exec_command session, and
`\u0003` interrupts it. Sending other stdin bytes may require approval because
it can drive the running process beyond the original command. Non-tty sessions
@@ -148,7 +141,6 @@ work.
- Default to concise, skimmable output. Lead with the answer or the result.
- Use GitHub-flavored Markdown: headings to structure longer replies, fenced
code blocks for code, and `inline code` for file paths, commands, symbols, and
short snippets. Reference code with file paths and include line numbers when
already known; do not look them up solely for the final summary.
short snippets. Reference code as `file:line` so it is clickable.
- Report outcomes faithfully: if tests failed, show it; if a step was skipped,
say so; when something is done and verified, state it plainly without hedging.
+3 -6
View File
@@ -60,17 +60,14 @@ const openAIPromptAddendum = `<model_guidance>
longer answers, fenced code blocks for code, and ` + "`inline code`" + ` for paths,
commands, and symbols.
- Strongly prefer the native file tools (read_file, list_directory, grep, glob,
edit_file, apply_patch, write_file) over shelling out to cat/sed/awk/python for
file work. Use them directly; never run apply_patch through exec_command.
read_file, edit_file and write_file take one file per call, while one coherent
apply_patch may span several files. Keep independent native calls in the same
turn, and do not batch file writes into a script.
write_file, apply_patch) over shelling out to cat/sed/awk/python for
file work. Make one tool call per file; do not batch file writes into a script.
- Persist until the task is fully handled this turn: gather context, implement,
run the validators, and report — do not stop at a partial result.
</model_guidance>`
const geminiPromptAddendum = `<model_guidance>
- Prefer the dedicated tools (read_file, grep, glob, edit_file, apply_patch) over
- Prefer the dedicated tools (read_file, grep, glob, apply_patch) over
equivalent shell commands; they are safer and produce cleaner diffs.
- Be concise and concrete. When you run a shell command with side effects, state
in one short clause why it is needed.
@@ -29,9 +29,6 @@ func TestBuildSystemPromptAppendsModelAddendum(t *testing.T) {
if got := buildSystemPrompt(Options{Model: "gpt-5"}); !strings.Contains(got, openAIPromptAddendum) {
t.Fatalf("expected the OpenAI addendum in the gpt-5 prompt")
}
if got := strings.Join(strings.Fields(strings.ToLower(buildSystemPrompt(Options{Model: "gpt-5"}))), " "); !strings.Contains(got, "never run apply_patch through exec_command") {
t.Fatalf("OpenAI prompt must route edits through the native patch tool:\n%s", got)
}
// Claude is aligned with the core prompt and gets no family addendum now that
// comment discipline is universal; it must not pick up another family's block.
claude := buildSystemPrompt(Options{Model: "claude-opus-4-6"})
+2 -12
View File
@@ -29,17 +29,14 @@ func TestMain(m *testing.M) {
}
func TestCoreSystemPromptIncludesCodingQualityRules(t *testing.T) {
prompt := strings.Join(strings.Fields(strings.ToLower(buildSystemPrompt(Options{}))), " ")
if strings.Contains(prompt, "clickable file paths") {
t.Fatalf("system prompt must not claim inline-code file paths are clickable: %s", prompt)
}
prompt := strings.ToLower(buildSystemPrompt(Options{}))
for _, want := range []string{
"read-before-edit",
"inspect the target file",
"plan then act",
"choose the narrowest tool",
"for edits to existing files, use edit_file",
"for edits to existing files, prefer apply_patch",
"verify after edits",
"honor the active permission mode",
"avoid broad refactors",
@@ -47,13 +44,6 @@ func TestCoreSystemPromptIncludesCodingQualityRules(t *testing.T) {
"do not recognize",
"scaled to the work",
"comment density",
"skip it for bounded changes in one component",
"do not read both unless a new need for exact content appears",
"combine compatible commands when diagnostics stay useful",
"do not invent a parallel flow",
"a successful native patch confirms the edit",
"once the required behavior passes validation, stop",
"do not look them up solely for the final summary",
} {
if !strings.Contains(prompt, want) {
t.Fatalf("expected core system prompt to include %q, got:\n%s", want, buildSystemPrompt(Options{}))
+12 -48
View File
@@ -839,20 +839,18 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
mcpTokenStore = nil
err = nil
}
criticalMCPConfig, optionalMCPConfig := splitMCPStartupConfig(mcpConfig)
mcpRuntime := mcpToolRuntime(noopMCPRuntime{})
if len(criticalMCPConfig.Servers) > 0 {
runtime, registerErr := deps.registerMCPTools(context.Background(), registry, criticalMCPConfig, mcp.RegisterOptions{
if len(mcpConfig.Servers) > 0 {
mcpRuntime, err = deps.registerMCPTools(context.Background(), registry, mcpConfig, mcp.RegisterOptions{
PermissionStore: mcpPermissionStore,
Autonomy: mcp.AutonomyLow,
Execution: executionRunner,
WorkspaceRoot: workspaceRoot,
})
if registerErr != nil {
closeMCPRuntime(stderr, runtime)
return writeAppError(stderr, redaction.ErrorMessage(registerErr, redaction.Options{}), 1)
}
mcpRuntime = runtime
}
if err != nil {
closeMCPRuntime(stderr, mcpRuntime)
return writeAppError(stderr, err.Error(), 1)
}
defer closeMCPRuntime(stderr, mcpRuntime)
// A server that could not be reached or validated is skipped, not fatal (one
@@ -911,35 +909,11 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
}
// Activate deferred MCP-tool loading for the interactive run only when the
// VISIBLE deferred-eligible count meets the resolved threshold, matching exec.
// This first pass covers every prompt-critical tool. Optional built-in MCP
// defaults re-run the same gate after publishing their atomic batch, so a
// later catalog generation cannot miss its loader. The interactive surface
// applies no operator tool filters, so enabled/disabled are nil — matching the
// AgentOptions below.
// The registry is complete (core + specialist + MCP + plugins) here, so the
// count is accurate; below threshold this is a no-op and the surface is
// unchanged. The interactive surface applies no operator tool filters, so
// enabled/disabled are nil — matching the AgentOptions below.
registerToolSearchIfEligible(registry, resolved.Tools.DeferThreshold, permissionMode, nil, nil)
var optionalMCPRuntime *optionalMCPStartup
var awaitToolReadiness func(context.Context)
if len(optionalMCPConfig.Servers) > 0 {
optionalMCPRuntime = startOptionalMCP(
context.Background(),
registry,
optionalMCPConfig,
mcp.RegisterOptions{
PermissionStore: mcpPermissionStore,
Autonomy: mcp.AutonomyLow,
Execution: executionRunner,
WorkspaceRoot: workspaceRoot,
},
deps.registerMCPTools,
func() {
registerToolSearchIfEligible(registry, resolved.Tools.DeferThreshold, permissionMode, nil, nil)
},
)
defer closeMCPRuntime(stderr, optionalMCPRuntime)
awaitToolReadiness = func(ctx context.Context) {
optionalMCPRuntime.Await(ctx, optionalMCPPromptGrace)
}
}
lastKnownMCPConfig := mcpConfig
fileTracker := tools.NewFileTracker()
var scratchBaseline scratchFileBaseline
@@ -976,17 +950,8 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
RecapsEnabled: resolved.Preferences.RecapsEnabled(),
Provider: provider,
NewProvider: deps.newProvider,
NewTurnSessionProvider: func(profile config.ProviderProfile, provider zeroruntime.Provider) zeroruntime.TurnSessionProvider {
if provider == nil {
return nil
}
if optimized, ok := providers.OptimizedTurnSessions(profile, provider, providers.Options{}); ok {
return optimized
}
return providers.DefaultTurnSessions(profile, provider, providers.Options{})
},
ProbeProviderHealth: deps.probeProviderHealth,
UserAgent: userAgent(),
ProbeProviderHealth: deps.probeProviderHealth,
UserAgent: userAgent(),
PrepareRunCompletionWarning: func() {
scratchBaseline = scratchFileSnapshot(workspaceRoot)
},
@@ -994,7 +959,6 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
return scratchFileWarning(workspaceRoot, scratchBaseline)
},
Registry: registry,
AwaitToolReadiness: awaitToolReadiness,
SessionStore: sessionStore,
PeerService: peerService,
SandboxStore: sandboxStore,
+1 -77
View File
@@ -17,7 +17,6 @@ import (
"github.com/Gitlawb/zero/internal/agent"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/mcp"
"github.com/Gitlawb/zero/internal/redaction"
"github.com/Gitlawb/zero/internal/tools"
"github.com/Gitlawb/zero/internal/tui"
"github.com/Gitlawb/zero/internal/update"
@@ -487,74 +486,6 @@ func TestRunNoArgsLaunchesTUIWithMCPState(t *testing.T) {
}
}
func TestRunNoArgsPaintsBeforeOptionalDefaultMCPIsReady(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer
cwd := t.TempDir()
userConfigPath := filepath.Join(t.TempDir(), "zero", "config.json")
permissionStore, err := mcp.NewPermissionStore(mcp.StoreOptions{FilePath: filepath.Join(t.TempDir(), "mcp-permissions.json")})
if err != nil {
t.Fatalf("NewPermissionStore() error = %v", err)
}
tokenStore, err := mcp.NewTokenStore(mcp.TokenStoreOptions{FilePath: filepath.Join(t.TempDir(), "mcp-oauth.json")})
if err != nil {
t.Fatalf("NewTokenStore() error = %v", err)
}
started := make(chan struct{})
release := make(chan struct{})
configResult := make(chan error, 1)
exitCode := runWithDeps([]string{}, &stdout, &stderr, appDeps{
getwd: func() (string, error) { return cwd, nil },
userConfigPath: func() (string, error) { return userConfigPath, nil },
resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) {
return config.ResolvedConfig{MaxTurns: 8}, nil
},
resolveMCPConfig: func(string, bool) (config.MCPConfig, error) {
return config.MCPConfig{Servers: config.DefaultMCPServers()}, nil
},
newMCPStore: func() (*mcp.PermissionStore, error) { return permissionStore, nil },
newMCPTokenStore: func() (*mcp.TokenStore, error) { return tokenStore, nil },
registerMCPTools: func(_ context.Context, registry *tools.Registry, cfg config.MCPConfig, _ mcp.RegisterOptions) (mcpToolRuntime, error) {
var configErr error
if _, ok := cfg.Servers["exa"]; !ok || len(cfg.Servers) != 1 {
configErr = fmt.Errorf("optional MCP config = %#v, want only exa", cfg.Servers)
}
configResult <- configErr
close(started)
<-release
registry.Register(cliFakeMCPRegistryTool{})
return noopMCPRuntime{}, nil
},
runTUI: func(ctx context.Context, options tui.Options) int {
select {
case <-started:
case <-time.After(time.Second):
t.Fatal("optional MCP initialization did not start")
}
if options.AwaitToolReadiness == nil {
t.Fatal("TUI did not receive optional tool readiness barrier")
}
if _, ok := options.Registry.Get("mcp_docs_lookup"); ok {
t.Fatal("optional MCP tool was published before registration completed")
}
close(release)
options.AwaitToolReadiness(ctx)
if _, ok := options.Registry.Get("mcp_docs_lookup"); !ok {
t.Fatal("optional MCP tool was not visible after readiness completed")
}
return 0
},
})
if exitCode != 0 {
t.Fatalf("exitCode = %d stderr=%s", exitCode, stderr.String())
}
if configErr := <-configResult; configErr != nil {
t.Fatal(configErr)
}
}
func TestTUIMCPCommandUsesLastGoodConfigOnRefreshError(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer
@@ -634,7 +565,6 @@ func TestRunNoArgsClosesPartialMCPRuntimeWhenRegistrationFails(t *testing.T) {
t.Fatalf("NewTokenStore() error = %v", err)
}
runtimeClosed := false
secret := "sk-proj-" + strings.Repeat("a", 23) + "0"
exitCode := runWithDeps([]string{}, &stdout, &stderr, appDeps{
getwd: func() (string, error) {
@@ -661,7 +591,7 @@ func TestRunNoArgsClosesPartialMCPRuntimeWhenRegistrationFails(t *testing.T) {
return closeFunc(func() error {
runtimeClosed = true
return nil
}), fmt.Errorf("register mcp tools failed: %s", secret)
}), errors.New("register mcp tools failed")
},
runTUI: func(ctx context.Context, options tui.Options) int {
t.Fatal("TUI should not launch after MCP registration fails")
@@ -675,12 +605,6 @@ func TestRunNoArgsClosesPartialMCPRuntimeWhenRegistrationFails(t *testing.T) {
if !strings.Contains(stderr.String(), "register mcp tools failed") {
t.Fatalf("stderr missing registration error: %s", stderr.String())
}
if strings.Contains(stderr.String(), secret) {
t.Fatalf("stderr leaked MCP registration secret: %s", stderr.String())
}
if !strings.Contains(stderr.String(), redaction.RedactedSecret) {
t.Fatalf("stderr missing redaction marker: %s", stderr.String())
}
if !runtimeClosed {
t.Fatal("partial MCP runtime was not closed after registration error")
}
-3
View File
@@ -939,9 +939,6 @@ func registerToolSearchIfEligible(registry *tools.Registry, deferThreshold int,
if deferThreshold <= 0 {
return
}
if _, registered := registry.Get(tools.ToolSearchToolName); registered {
return
}
if deferredEligibleCount(registry, permissionMode, enabledTools, disabledTools) < deferThreshold {
return
}
+9 -15
View File
@@ -267,29 +267,23 @@ func streamJSONPermissionEventType(event agent.PermissionEvent) streamjson.Event
func (writer *execEventWriter) usage(usage agent.Usage) {
if writer.format == execOutputJSON {
writer.writeJSON(map[string]any{
"type": "usage",
"prompt_tokens": usage.PromptTokens,
"completion_tokens": usage.CompletionTokens,
"cached_input_tokens": usage.CachedInputTokens,
"cache_write_tokens": usage.CacheWriteTokens,
"total_tokens": usage.TotalTokens(),
"type": "usage",
"prompt_tokens": usage.PromptTokens,
"completion_tokens": usage.CompletionTokens,
"total_tokens": usage.TotalTokens(),
})
return
}
if writer.format == execOutputStreamJSON {
promptTokens := usage.EffectiveInputTokens()
completionTokens := usage.EffectiveOutputTokens()
cachedInputTokens := usage.CachedInputTokens
cacheWriteTokens := usage.CacheWriteTokens
totalTokens := usage.TotalTokens()
writer.writeStreamJSON(streamjson.Event{
Type: streamjson.EventUsage,
RunID: writer.runID,
PromptTokens: &promptTokens,
CompletionTokens: &completionTokens,
CachedInputTokens: &cachedInputTokens,
CacheWriteTokens: &cacheWriteTokens,
TotalTokens: &totalTokens,
Type: streamjson.EventUsage,
RunID: writer.runID,
PromptTokens: &promptTokens,
CompletionTokens: &completionTokens,
TotalTokens: &totalTokens,
})
}
}
-36
View File
@@ -52,39 +52,3 @@ func TestExecWriterPropagatesToolResultTruncation(t *testing.T) {
})
}
}
func TestExecWriterPropagatesCacheUsage(t *testing.T) {
for _, format := range []execOutputFormat{execOutputJSON, execOutputStreamJSON} {
t.Run(string(format), func(t *testing.T) {
var stdout bytes.Buffer
writer := execEventWriter{
stdout: &stdout,
format: format,
runID: "run_usage",
streamedText: &strings.Builder{},
}
writer.usage(agent.Usage{
InputTokens: 100,
OutputTokens: 20,
PromptTokens: 100,
CompletionTokens: 20,
CachedInputTokens: 60,
CacheWriteTokens: 10,
})
if writer.err != nil {
t.Fatalf("usage: %v", writer.err)
}
var payload map[string]any
if err := json.Unmarshal(bytes.TrimSpace(stdout.Bytes()), &payload); err != nil {
t.Fatalf("decode output %q: %v", stdout.String(), err)
}
cachedKey, writeKey := "cached_input_tokens", "cache_write_tokens"
if format == execOutputStreamJSON {
cachedKey, writeKey = "cachedInputTokens", "cacheWriteTokens"
}
if payload[cachedKey] != float64(60) || payload[writeKey] != float64(10) {
t.Fatalf("cache usage missing from payload: %#v", payload)
}
})
}
}
-147
View File
@@ -1,147 +0,0 @@
package cli
import (
"context"
"sync"
"time"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/mcp"
"github.com/Gitlawb/zero/internal/tools"
)
const (
optionalMCPPromptGrace = time.Second
optionalMCPCloseGrace = 250 * time.Millisecond
)
// splitMCPStartupConfig keeps explicitly configured servers on the
// prompt-critical path. Only unchanged built-in defaults may initialize after
// first paint; their failure was already intentionally silent.
func splitMCPStartupConfig(cfg config.MCPConfig) (critical config.MCPConfig, optional config.MCPConfig) {
critical.Servers = map[string]config.MCPServerConfig{}
optional.Servers = map[string]config.MCPServerConfig{}
for name, server := range cfg.Servers {
if server.Disabled {
continue
}
if config.IsUnconfiguredDefault(name, server) {
optional.Servers[name] = server
continue
}
critical.Servers[name] = server
}
return critical, optional
}
type optionalMCPStartup struct {
cancel context.CancelFunc
done chan struct{}
mu sync.Mutex
runtime mcpToolRuntime
closeOnce sync.Once
runtimeCloseOnce sync.Once
closeErr error
}
func startOptionalMCP(
parent context.Context,
registry *tools.Registry,
cfg config.MCPConfig,
options mcp.RegisterOptions,
register func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error),
onReady func(),
) *optionalMCPStartup {
ctx, cancel := context.WithCancel(parent)
startup := &optionalMCPStartup{cancel: cancel, done: make(chan struct{})}
go func() {
runtime, err := register(ctx, registry, cfg, options)
if err == nil && onReady != nil {
onReady()
}
startup.mu.Lock()
startup.runtime = runtime
startup.mu.Unlock()
close(startup.done)
}()
return startup
}
// Await gives optional tools one bounded opportunity to join the next turn's
// immutable registry snapshot. A timeout leaves this turn on the existing core
// tool set; a later turn observes the completed generation.
func (startup *optionalMCPStartup) Await(ctx context.Context, timeout time.Duration) bool {
if startup == nil {
return true
}
select {
case <-startup.done:
return true
default:
}
if timeout <= 0 {
select {
case <-startup.done:
return true
case <-ctx.Done():
return false
}
}
timer := time.NewTimer(timeout)
defer timer.Stop()
select {
case <-startup.done:
return true
case <-timer.C:
return false
case <-ctx.Done():
return false
}
}
func (startup *optionalMCPStartup) Close() error {
if startup == nil {
return nil
}
startup.closeOnce.Do(func() {
startup.cancel()
reaped := make(chan struct{})
go func() {
defer close(reaped)
<-startup.done
startup.closeRuntime()
}()
timer := time.NewTimer(optionalMCPCloseGrace)
defer timer.Stop()
select {
case <-reaped:
case <-timer.C:
}
})
startup.mu.Lock()
defer startup.mu.Unlock()
return startup.closeErr
}
func (startup *optionalMCPStartup) closeRuntime() {
startup.runtimeCloseOnce.Do(func() {
startup.mu.Lock()
runtime := startup.runtime
startup.mu.Unlock()
if runtime == nil {
return
}
err := runtime.Close()
startup.mu.Lock()
if err != nil && startup.closeErr == nil {
startup.closeErr = err
}
startup.mu.Unlock()
})
}
func (startup *optionalMCPStartup) Skipped() []mcp.SkippedServer {
return nil
}
-187
View File
@@ -1,187 +0,0 @@
package cli
import (
"context"
"errors"
"testing"
"time"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/mcp"
"github.com/Gitlawb/zero/internal/tools"
)
type closeTrackingMCPRuntime struct {
closed chan struct{}
err error
}
func (runtime *closeTrackingMCPRuntime) Close() error {
close(runtime.closed)
return runtime.err
}
func (*closeTrackingMCPRuntime) Skipped() []mcp.SkippedServer { return nil }
func TestSplitMCPStartupConfigDefersOnlyUnconfiguredDefaults(t *testing.T) {
defaults := config.DefaultMCPServers()
cfg := config.MCPConfig{Servers: map[string]config.MCPServerConfig{
"exa": defaults["exa"],
"docs": {Type: "http", URL: "https://docs.example.test/mcp"},
"disabled": {Type: "http", URL: "https://disabled.example.test/mcp", Disabled: true},
}}
critical, optional := splitMCPStartupConfig(cfg)
if _, ok := critical.Servers["docs"]; !ok || len(critical.Servers) != 1 {
t.Fatalf("critical servers = %#v, want only docs", critical.Servers)
}
if _, ok := optional.Servers["exa"]; !ok || len(optional.Servers) != 1 {
t.Fatalf("optional servers = %#v, want only exa", optional.Servers)
}
}
func TestOptionalMCPAwaitTimesOutThenObservesPublishedTools(t *testing.T) {
registry := tools.NewRegistry()
release := make(chan struct{})
started := make(chan struct{})
startup := startOptionalMCP(
context.Background(),
registry,
config.MCPConfig{Servers: config.DefaultMCPServers()},
mcp.RegisterOptions{},
func(ctx context.Context, registry *tools.Registry, _ config.MCPConfig, _ mcp.RegisterOptions) (mcpToolRuntime, error) {
close(started)
select {
case <-release:
registry.RegisterBatch([]tools.Tool{tools.NewScopedReadFileTool(t.TempDir(), nil)})
return noopMCPRuntime{}, nil
case <-ctx.Done():
return noopMCPRuntime{}, ctx.Err()
}
},
nil,
)
defer func() { _ = startup.Close() }()
<-started
if startup.Await(t.Context(), time.Millisecond) {
t.Fatal("optional startup unexpectedly completed before release")
}
if got := len(registry.All()); got != 0 {
t.Fatalf("registry published tools before readiness: %d", got)
}
close(release)
if !startup.Await(t.Context(), time.Second) {
t.Fatal("optional startup did not complete after release")
}
if _, ok := registry.Get("read_file"); !ok {
t.Fatal("completed optional startup did not publish its tool batch")
}
}
func TestOptionalMCPFailedRegistrationSkipsReadinessCallback(t *testing.T) {
readyCalls := 0
startup := startOptionalMCP(
context.Background(),
tools.NewRegistry(),
config.MCPConfig{Servers: config.DefaultMCPServers()},
mcp.RegisterOptions{},
func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
return noopMCPRuntime{}, errors.New("connect failed")
},
func() { readyCalls++ },
)
defer func() { _ = startup.Close() }()
if !startup.Await(t.Context(), time.Second) {
t.Fatal("failed optional startup did not settle")
}
if readyCalls != 0 {
t.Fatalf("readiness callback ran after failed registration: %d", readyCalls)
}
}
func TestOptionalMCPCloseCancelsBlockedStartup(t *testing.T) {
started := make(chan struct{})
startup := startOptionalMCP(
context.Background(),
tools.NewRegistry(),
config.MCPConfig{Servers: config.DefaultMCPServers()},
mcp.RegisterOptions{},
func(ctx context.Context, _ *tools.Registry, _ config.MCPConfig, _ mcp.RegisterOptions) (mcpToolRuntime, error) {
close(started)
<-ctx.Done()
return noopMCPRuntime{}, nil
},
nil,
)
<-started
if err := startup.Close(); err != nil {
t.Fatalf("Close() error = %v", err)
}
}
func TestOptionalMCPCloseReturnsRuntimeErrorWithinGrace(t *testing.T) {
closeErr := errors.New("close failed")
runtime := &closeTrackingMCPRuntime{closed: make(chan struct{}), err: closeErr}
startup := startOptionalMCP(
context.Background(),
tools.NewRegistry(),
config.MCPConfig{Servers: config.DefaultMCPServers()},
mcp.RegisterOptions{},
func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
return runtime, nil
},
nil,
)
if !startup.Await(t.Context(), time.Second) {
t.Fatal("optional startup did not settle")
}
if err := startup.Close(); !errors.Is(err, closeErr) {
t.Fatalf("Close() error = %v, want %v", err, closeErr)
}
}
func TestOptionalMCPCloseBoundsRegistrationThatIgnoresCancellation(t *testing.T) {
started := make(chan struct{})
release := make(chan struct{})
released := false
defer func() {
if !released {
close(release)
}
}()
runtime := &closeTrackingMCPRuntime{closed: make(chan struct{})}
startup := startOptionalMCP(
context.Background(),
tools.NewRegistry(),
config.MCPConfig{Servers: config.DefaultMCPServers()},
mcp.RegisterOptions{},
func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
close(started)
<-release
return runtime, nil
},
nil,
)
<-started
closeResult := make(chan error, 1)
go func() { closeResult <- startup.Close() }()
select {
case err := <-closeResult:
if err != nil {
t.Fatalf("Close() error = %v", err)
}
case <-time.After(time.Second):
t.Fatal("Close blocked on registration after cancellation")
}
close(release)
released = true
select {
case <-runtime.closed:
case <-time.After(time.Second):
t.Fatal("late optional MCP runtime was not closed")
}
}
+9 -66
View File
@@ -5,31 +5,23 @@ import (
"fmt"
"os"
"path/filepath"
"sync/atomic"
"time"
"github.com/Gitlawb/zero/internal/lockutil"
)
// Cross-process lock tuning. The lock is held only for a single metadata
// read-modify-write (milliseconds), never across a job's exec, so the timeout is
// generous and the stale threshold sits far above any real hold.
// read-modify-write (milliseconds), never across a job's exec, so the timeout
// is generous relative to any real hold.
const (
cronLockTimeout = 10 * time.Second
cronLockStaleAfter = 60 * time.Second
cronLockRetryDelay = 20 * time.Millisecond
)
var cronLockSeq atomic.Uint64
// lockJob takes a cross-process exclusive lock for one job by O_EXCL-creating a
// sibling "<id>.lock" file next to the job directory (so Remove's RemoveAll of
// the job dir never deletes a live lock). It serializes the read-modify-write of
// a job's metadata across concurrent schedulers and commands. The lock file is
// removed on release; a stale lock from a crashed holder (older than
// cronLockStaleAfter) is reclaimed. Wall-clock time.Now is used deliberately
// (not the injectable Store.now) so lock timing never depends on a frozen test
// clock and the stale check compares against real file mtimes.
// lockJob takes a cross-process advisory lock for one job on a stable sibling
// "<id>.lock" file next to the job directory. It serializes the read-modify-write
// of a job's metadata across concurrent schedulers and commands. The kernel
// releases the lock when a holder exits; the file is never renamed or removed.
func (s *Store) lockJob(id string) (func(), error) {
if !validID(id) {
return nil, fmt.Errorf("invalid cron job id %q", id)
@@ -38,64 +30,15 @@ func (s *Store) lockJob(id string) (func(), error) {
if err := os.MkdirAll(filepath.Dir(lockPath), 0o700); err != nil {
return nil, err
}
token := fmt.Sprintf("%d-%d-%d", os.Getpid(), time.Now().UnixNano(), cronLockSeq.Add(1))
deadline := time.Now().Add(cronLockTimeout)
for {
f, err := os.OpenFile(lockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
lock, err := lockutil.TryAcquireFileLock(lockPath)
if err == nil {
// A partial write would leave a lock file without our token, so the
// releaser could never delete it — stranding the lock. Fail closed.
if _, werr := f.WriteString(token); werr != nil {
_ = f.Close()
_ = lockutil.RemoveLockFile(lockPath)
return nil, fmt.Errorf("cron: write job lock: %w", werr)
}
if cerr := f.Close(); cerr != nil {
_ = lockutil.RemoveLockFile(lockPath)
return nil, fmt.Errorf("cron: close job lock: %w", cerr)
}
var released bool
return func() {
if released {
return
}
released = true
if data, rerr := os.ReadFile(lockPath); rerr == nil && string(data) == token {
_ = lockutil.RemoveLockFile(lockPath)
}
}, nil
return func() { _ = lock.Release() }, nil
}
// On Windows a concurrent holder's os.Remove leaves the lock file in a
// "delete pending" state, so an O_EXCL create races it with
// ERROR_ACCESS_DENIED (os.ErrPermission) rather than ErrExist. Treat that
// as contention and retry, exactly like ErrExist — otherwise the lock
// spuriously fails under concurrency on Windows.
if !errors.Is(err, os.ErrExist) && !errors.Is(err, os.ErrPermission) {
if !errors.Is(err, lockutil.ErrLockHeld) {
return nil, fmt.Errorf("cron: acquire job lock: %w", err)
}
// Reclaim a stale lock left by a crashed holder — atomically (H3). A blind
// Remove lets two racers both "reclaim" and recreate, so both hold the lock;
// reclaimStaleLock renames the file aside (only one rename of a given file
// wins) and restores it if it turns out fresh, so a live lock is never deleted
// out from under its holder.
if info, statErr := os.Stat(lockPath); statErr == nil && time.Since(info.ModTime()) > cronLockStaleAfter {
cleared, rerr := lockutil.ReclaimStaleLock(lockPath, token, func(reclaimedPath string) bool {
info, err := os.Stat(reclaimedPath)
return err == nil && time.Since(info.ModTime()) <= cronLockStaleAfter
})
if rerr != nil {
// Reclaim hit a hard failure: the rename aside failed outright, or a
// live holder's lock could not be put back (the lock path may be
// missing, so re-acquiring would break mutual exclusion). Fail closed
// instead of spinning to the deadline.
return nil, fmt.Errorf("cron: reclaim stale job lock: %w", rerr)
}
if cleared {
continue // cleared a genuinely stale lock; retry the O_EXCL create now
}
// Lost the reclaim race (or it was actually fresh) — fall through to the
// bounded wait instead of hot-spinning on a reclaim that never wins (L13).
}
if !time.Now().Before(deadline) {
return nil, fmt.Errorf("cron: timed out acquiring job lock for %q", id)
}
+1 -1
View File
@@ -64,7 +64,7 @@ func TestMutateRemovedJobReturnsNotFound(t *testing.T) {
if !errors.Is(err, ErrJobNotFound) {
t.Fatalf("Mutate on removed job = %v, want ErrJobNotFound", err)
}
// And the lock file left no litter / did not resurrect the job dir.
// The stable advisory-lock file does not resurrect the removed job dir.
if _, gerr := store.Get(job.ID); !errors.Is(gerr, ErrJobNotFound) {
t.Fatalf("job should remain removed, Get = %v", gerr)
}
+35 -67
View File
@@ -3,107 +3,60 @@ package daemon
import (
"errors"
"fmt"
"io/fs"
"os"
"strconv"
"strings"
"sync/atomic"
"github.com/Gitlawb/zero/internal/lockutil"
)
// Single-instance lock. Mirrors reference-daemon-code-agent-js/supervisor.js's
// lock file: a PID file created with O_EXCL. A second start fails; a STALE lock
// left by a dead daemon (the recorded PID is no longer alive) is reclaimed so the
// daemon recovers from an unclean shutdown without manual cleanup.
// Single-instance lock. The stable PID file carries a kernel advisory lock for
// the daemon's full lifetime. Process exit releases the lock automatically, so
// stale metadata needs no rename/remove recovery protocol.
// ErrAlreadyRunning is returned when a live daemon already holds the lock.
var ErrAlreadyRunning = errors.New("daemon: another instance is already running")
// fileLock is an acquired single-instance lock.
type fileLock struct {
path string
lock *lockutil.FileLock
}
// processAlive reports whether pid is a live process. Implemented per-platform
// (lock_posix.go / lock_windows.go). It is a package var so tests can stub it.
var processAlive = osProcessAlive
// acquireLock takes the single-instance lock at path, reclaiming a stale lock
// whose recorded PID is dead. isAlive overrides the liveness check (tests pass a
// stub); nil uses the real processAlive.
// acquireLock takes the single-instance advisory lock at path. isAlive is used
// only to enrich a contention error with the current holder's PID; the kernel
// lock, not PID metadata, is authoritative.
func acquireLock(path string, isAlive func(pid int) bool) (*fileLock, error) {
if isAlive == nil {
isAlive = processAlive
}
// At most two passes: create, or detect-stale-then-retry once.
for attempt := 0; attempt < 2; attempt++ {
f, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
if err == nil {
// A failed PID write would leave a malformed lock file that another
// process reads as stale (unparsable PID) and wrongly reclaims, breaking
// the single-instance guarantee — so on write failure, remove it and fail.
if _, werr := fmt.Fprintf(f, "%d\n", os.Getpid()); werr != nil {
_ = f.Close()
_ = lockutil.RemoveLockFile(path)
return nil, werr
}
if cerr := f.Close(); cerr != nil {
_ = lockutil.RemoveLockFile(path)
return nil, cerr
}
return &fileLock{path: path}, nil
}
if !errors.Is(err, fs.ErrExist) && !errors.Is(err, os.ErrPermission) {
lock, err := lockutil.TryAcquireFileLock(path)
if err != nil {
if !errors.Is(err, lockutil.ErrLockHeld) {
return nil, err
}
pid, perr := readPidFile(path)
if perr == nil && pid > 0 && isAlive(pid) {
return nil, fmt.Errorf("%w (pid %d)", ErrAlreadyRunning, pid)
}
// Stale lock (dead PID or unreadable) — reclaim it atomically, then retry the
// O_EXCL create. A blind Remove here races: two daemons starting at once could
// both read the stale PID, both Remove, and then one Removes the OTHER's
// freshly-created lock — leaving both "holding" the single-instance lock.
// reclaimStaleLock renames the file aside so only one racer wins the rename,
// and restores it if a live holder reacquired in the gap (D6).
if _, rerr := reclaimStaleLock(path, isAlive); rerr != nil {
// Reclaim hit a hard failure: the rename aside failed outright, or a
// live holder's lock could not be put back (the lock path may be
// missing, so re-acquiring would break the single-instance guarantee).
// Fail closed instead of spinning to the deadline.
return nil, fmt.Errorf("daemon: reclaim stale lock: %w", rerr)
}
return nil, ErrAlreadyRunning
}
return nil, ErrAlreadyRunning
if err := lock.WriteMetadata([]byte(fmt.Sprintf("%d\n", os.Getpid()))); err != nil {
return nil, errors.Join(err, lock.Release())
}
return &fileLock{lock: lock}, nil
}
// daemonLockSeq makes each reclaim attempt's sidelined filename unique per process.
var daemonLockSeq atomic.Uint64
// reclaimStaleLock atomically reclaims a single-instance lock whose recorded
// PID is dead, via lockutil.ReclaimStaleLock: the lock file is renamed aside
// (only one racer can win the rename) and stolen only if the moved file's PID
// is still dead; a holder that reacquired the lock in the gap between the
// stale check and the rename carries a LIVE pid and is restored rather than
// stolen. Returns true only when a genuinely stale lock was removed. A
// non-nil error means the caller must fail closed instead of re-acquiring
// (see lockutil.ReclaimStaleLock).
func reclaimStaleLock(path string, isAlive func(pid int) bool) (bool, error) {
suffix := fmt.Sprintf("%d-%d", os.Getpid(), daemonLockSeq.Add(1))
return lockutil.ReclaimStaleLock(path, suffix, func(reclaimedPath string) bool {
pid, err := readPidFile(reclaimedPath)
return err == nil && pid > 0 && isAlive(pid)
})
}
// release removes the lock file. Safe to call once. An already-missing lock
// file is not an error (lockutil.RemoveLockFile swallows it on every platform).
// release drops the kernel lock and closes its handle. The stable PID file is
// deliberately retained so contenders always address the same inode.
func (l *fileLock) release() error {
if l == nil || l.path == "" {
if l == nil || l.lock == nil {
return nil
}
return lockutil.RemoveLockFile(l.path)
return l.lock.Release()
}
// readPidFile reads and parses the PID recorded in a lock file.
@@ -112,5 +65,20 @@ func readPidFile(path string) (int, error) {
if err != nil {
return 0, err
}
return strconv.Atoi(strings.TrimSpace(string(data)))
metadata := strings.TrimSpace(string(data))
if pid, err := strconv.Atoi(metadata); err == nil {
return pid, nil
}
pidText, sequenceText, ok := strings.Cut(metadata, "-")
if !ok || pidText == "" || sequenceText == "" || strings.Contains(sequenceText, "-") {
return 0, fmt.Errorf("daemon: parse lock PID metadata %q", metadata)
}
if _, err := strconv.ParseUint(sequenceText, 10, 64); err != nil {
return 0, fmt.Errorf("daemon: parse lock holder sequence: %w", err)
}
pid, err := strconv.Atoi(pidText)
if err != nil {
return 0, fmt.Errorf("daemon: parse lock PID: %w", err)
}
return pid, nil
}
+54 -42
View File
@@ -32,27 +32,48 @@ func TestLockSingleInstance(t *testing.T) {
_ = l2.release()
}
func TestLockStaleRecovery(t *testing.T) {
// A PID liveness probe is diagnostic only. Even if it reports the recorded PID
// dead, a contended kernel lock proves that the holder is active and must win.
// The old rename-aside protocol trusted this stale observation and admitted a
// second daemon while the first one was still inside its critical section.
func TestLockDoesNotOverrideKernelHolderWithStalePIDProbe(t *testing.T) {
path := filepath.Join(t.TempDir(), "daemon.lock")
// Simulate a stale lock from a crashed daemon: a PID file whose process is
// dead.
first, err := acquireLock(path, func(int) bool { return true })
if err != nil {
t.Fatalf("first acquireLock: %v", err)
}
defer first.release()
second, err := acquireLock(path, func(int) bool { return false })
if second != nil {
_ = second.release()
}
if !errors.Is(err, ErrAlreadyRunning) {
t.Fatalf("contended acquire with stale PID probe = %v, want ErrAlreadyRunning", err)
}
}
func TestLockIgnoresStaleMetadataWhenKernelLockIsFree(t *testing.T) {
path := filepath.Join(t.TempDir(), "daemon.lock")
// A crashed daemon may leave PID metadata behind, but its kernel lock is
// released automatically, so the next daemon can acquire without moving it.
if err := os.WriteFile(path, []byte("4242\n"), 0o600); err != nil {
t.Fatalf("write stale lock: %v", err)
}
dead := func(int) bool { return false }
l, err := acquireLock(path, dead)
if err != nil {
t.Fatalf("stale-lock recovery failed: %v", err)
t.Fatalf("acquire over stale metadata: %v", err)
}
// The lock now records OUR pid, not the stale one.
data, _ := os.ReadFile(path)
if strings.TrimSpace(string(data)) != strconv.Itoa(os.Getpid()) {
t.Fatalf("reclaimed lock pid = %q, want %d", strings.TrimSpace(string(data)), os.Getpid())
t.Fatalf("lock pid = %q, want %d", strings.TrimSpace(string(data)), os.Getpid())
}
_ = l.release()
}
func TestLockReleaseRemovesFile(t *testing.T) {
func TestLockReleaseKeepsStableFile(t *testing.T) {
path := filepath.Join(t.TempDir(), "daemon.lock")
l, err := acquireLock(path, func(int) bool { return true })
if err != nil {
@@ -61,46 +82,37 @@ func TestLockReleaseRemovesFile(t *testing.T) {
if err := l.release(); err != nil {
t.Fatalf("release: %v", err)
}
if _, err := os.Stat(path); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("lock file still present after release: %v", err)
if _, err := os.Stat(path); err != nil {
t.Fatalf("release removed stable lock file: %v", err)
}
}
func TestReclaimStaleLockRemovesDeadHolder(t *testing.T) {
func TestReadPidFileMetadataFormats(t *testing.T) {
path := filepath.Join(t.TempDir(), "daemon.lock")
if err := os.WriteFile(path, []byte("4242\n"), 0o600); err != nil {
t.Fatalf("write stale lock: %v", err)
}
if ok, err := reclaimStaleLock(path, func(int) bool { return false }); err != nil || !ok {
t.Fatalf("reclaimStaleLock must report a genuinely stale (dead-PID) lock reclaimed (ok=%v err=%v)", ok, err)
}
if _, err := os.Stat(path); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("a reclaimed stale lock must be removed, stat = %v", err)
}
}
func TestReclaimStaleLockRestoresLiveHolder(t *testing.T) {
// If a holder reacquires the lock in the gap between the stale check and the
// rename, the moved file carries a LIVE pid; reclaim must restore it untouched
// rather than steal it — otherwise two daemons both "hold" the lock (D6).
path := filepath.Join(t.TempDir(), "daemon.lock")
if err := os.WriteFile(path, []byte("4242\n"), 0o600); err != nil {
t.Fatalf("write lock: %v", err)
}
if ok, err := reclaimStaleLock(path, func(int) bool { return true }); err != nil || ok {
t.Fatalf("reclaimStaleLock must NOT report a live-PID lock reclaimed (ok=%v err=%v)", ok, err)
}
data, err := os.ReadFile(path)
if err != nil {
t.Fatalf("live lock must be restored in place, read = %v", err)
}
if strings.TrimSpace(string(data)) != "4242" {
t.Fatalf("restored lock content = %q, want %q (unchanged)", strings.TrimSpace(string(data)), "4242")
}
// No sidelined ".stale" leftovers.
matches, _ := filepath.Glob(path + ".stale.*")
if len(matches) != 0 {
t.Fatalf("reclaim left sidelined files: %v", matches)
for _, test := range []struct {
name string
metadata string
wantPID int
wantErr bool
}{
{name: "legacy PID", metadata: "4242\n", wantPID: 4242},
{name: "PID and holder sequence", metadata: "4242-7\n", wantPID: 4242},
{name: "missing sequence", metadata: "4242-", wantErr: true},
{name: "non-numeric sequence", metadata: "4242-next", wantErr: true},
{name: "extra component", metadata: "4242-7-8", wantErr: true},
} {
t.Run(test.name, func(t *testing.T) {
if err := os.WriteFile(path, []byte(test.metadata), 0o600); err != nil {
t.Fatal(err)
}
got, err := readPidFile(path)
if (err != nil) != test.wantErr {
t.Fatalf("readPidFile() error = %v, wantErr %v", err, test.wantErr)
}
if !test.wantErr && got != test.wantPID {
t.Fatalf("readPidFile() = %d, want %d", got, test.wantPID)
}
})
}
}
+6 -2
View File
@@ -124,12 +124,16 @@ func TestServerEndToEnd(t *testing.T) {
case <-time.After(3 * time.Second):
t.Fatal("Serve did not return after shutdown")
}
// Socket, status, and lock files are removed on exit.
for _, p := range []string{paths.Socket, paths.Status, paths.Lock} {
// Runtime endpoints are removed on exit. The advisory lock file remains as
// a stable inode, but the kernel lock on it has been released.
for _, p := range []string{paths.Socket, paths.Status} {
if _, err := os.Stat(p); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("file %s not cleaned up after shutdown: %v", p, err)
}
}
if _, err := os.Stat(paths.Lock); err != nil {
t.Fatalf("stable lock file missing after shutdown: %v", err)
}
}
func TestServerSecondInstanceFails(t *testing.T) {
+9 -63
View File
@@ -5,93 +5,39 @@ import (
"fmt"
"os"
"path/filepath"
"sync/atomic"
"time"
"github.com/Gitlawb/zero/internal/lockutil"
)
// Cross-process lock tuning for the audit log. The lock is held only across a
// single read-then-append (milliseconds), so the timeout is generous and the
// stale threshold sits far above any real hold.
// single read-then-append (milliseconds), so the timeout is generous relative
// to any real hold.
const (
auditLockTimeout = 10 * time.Second
auditLockStaleAfter = 60 * time.Second
auditLockRetryDelay = 20 * time.Millisecond
)
var auditLockSeq atomic.Uint64
// lockAudit takes a cross-process exclusive lock on the audit log by
// O_EXCL-creating a sibling "<auditPath>.lock" file (removed on release). It makes
// lockAudit takes a cross-process advisory lock on the audit log through a
// stable sibling "<auditPath>.lock" file. It makes
// the lastSequence()+append in append() atomic across processes; the audit log is
// a shared global file, so without it two processes can read the same last
// sequence and write duplicate numbers. A stale lock from a crashed holder (older
// than auditLockStaleAfter) is reclaimed. Wall-clock time is used deliberately so
// lock timing never depends on the store's injectable clock and the stale check
// compares against real file mtimes. This mirrors internal/cron/lock.go and
// internal/oauth/lock.go. Uses internal/lockutil for lock restore and reclaim.
// sequence and write duplicate numbers. Kernel lock lifetime provides crash
// recovery without moving or deleting the lock path.
func (store *AuditStore) lockAudit() (func(), error) {
lockPath := store.auditPath + ".lock"
if err := os.MkdirAll(filepath.Dir(lockPath), 0o700); err != nil {
return nil, err
}
token := fmt.Sprintf("%d-%d-%d", os.Getpid(), time.Now().UnixNano(), auditLockSeq.Add(1))
deadline := time.Now().Add(auditLockTimeout)
for {
f, err := os.OpenFile(lockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
lock, err := lockutil.TryAcquireFileLock(lockPath)
if err == nil {
// A partial write would leave a lock file without our token, so the
// releaser could never delete it — stranding the lock. Fail closed.
if _, werr := f.WriteString(token); werr != nil {
_ = f.Close()
_ = lockutil.RemoveLockFile(lockPath)
return nil, fmt.Errorf("hooks: write audit lock: %w", werr)
}
if cerr := f.Close(); cerr != nil {
_ = lockutil.RemoveLockFile(lockPath)
return nil, fmt.Errorf("hooks: close audit lock: %w", cerr)
}
var released bool
return func() {
if released {
return
}
released = true
if data, rerr := os.ReadFile(lockPath); rerr == nil && string(data) == token {
_ = lockutil.RemoveLockFile(lockPath)
}
}, nil
return func() { _ = lock.Release() }, nil
}
// On Windows a concurrent holder's os.Remove leaves the lock file in a
// "delete pending" state, so an O_EXCL create races it with
// ERROR_ACCESS_DENIED (os.ErrPermission) rather than ErrExist. Treat both
// as contention and retry.
if !errors.Is(err, os.ErrExist) && !errors.Is(err, os.ErrPermission) {
if !errors.Is(err, lockutil.ErrLockHeld) {
return nil, fmt.Errorf("hooks: acquire audit lock: %w", err)
}
// Reclaim a stale lock left by a crashed holder — atomically (H3). A blind
// Remove lets two racers both reclaim + recreate and so both hold the lock;
// reclaimStaleLock renames the file aside (only one rename wins) and restores
// it if it turns out fresh, so a live lock is never deleted out from under it.
if info, statErr := os.Stat(lockPath); statErr == nil && time.Since(info.ModTime()) > auditLockStaleAfter {
cleared, rerr := lockutil.ReclaimStaleLock(lockPath, token, func(reclaimedPath string) bool {
info, err := os.Stat(reclaimedPath)
return err == nil && time.Since(info.ModTime()) <= auditLockStaleAfter
})
if rerr != nil {
// Reclaim hit a hard failure: the rename aside failed outright, or a
// live holder's lock could not be put back (the lock path may be
// missing, so re-acquiring would break mutual exclusion). Fail closed
// instead of spinning to the deadline.
return nil, fmt.Errorf("hooks: reclaim stale audit lock: %w", rerr)
}
if cleared {
continue
}
// Lost the reclaim race (or it was actually fresh) — fall through to the
// bounded wait rather than hot-spinning on a reclaim that never wins.
}
if !time.Now().Before(deadline) {
return nil, fmt.Errorf("hooks: timed out acquiring audit lock")
}
+111 -48
View File
@@ -1,64 +1,127 @@
// Package lockutil provides the platform-specific file primitives behind the
// O_EXCL lock files in cron, daemon, hooks, oauth, and swarm: a no-overwrite
// restore for locks that were sidelined during a stale-reclaim attempt, and a
// lock file remover with one cross-platform contract (missing files are a
// no-op; Windows retries transient sharing violations).
// Package lockutil provides cross-process advisory file locks used by cron,
// daemon, hooks, oauth, and swarm. Locks are held by the kernel for the lifetime
// of an open file handle, so process exit releases them without a stale-file
// reclaim protocol. The lock file itself is stable: holders never rename or
// remove it while another process may still have its inode locked.
package lockutil
import (
"io"
"errors"
"fmt"
"os"
"path/filepath"
"strings"
"sync"
"sync/atomic"
)
// restoreByCopy restores reclaimed to path without overwriting an existing
// path, as a fallback for when the platform's primary no-replace primitive
// (hard link on POSIX, MoveFileEx on Windows) fails for a reason other than
// the destination existing. It stages a full copy under a private name next
// to reclaimed and publishes it to path with publish (the same no-replace
// primitive the caller's platform uses for the primary restore), so path
// never appears in a partially-copied state: a crash between staging and
// publish leaves path exactly as it was (missing or, if a competing holder
// created it, untouched), never a truncated file that a PID/content-based
// liveness check could mistake for dead. Leaving path missing here would let
// the next O_EXCL create succeed while the sidelined holder is still live,
// breaking mutual exclusion. publish keeps the no-overwrite guarantee: a new
// holder that appeared in the meantime wins and this returns os.ErrExist. The
// copy resets the lock's mtime to now, which only makes the restored lock
// look fresher; that is safe, since it is being handed back to a live holder.
func restoreByCopy(reclaimed, path string, publish func(from, to string) error) error {
staged := reclaimed + ".copy"
src, err := os.Open(reclaimed)
var (
// ErrLockHeld reports that another process or goroutine currently holds the
// advisory lock. Callers may wait and retry until their own deadline.
ErrLockHeld = errors.New("lockutil: lock is held")
// ErrLockReleased reports an attempt to update metadata after Release.
ErrLockReleased = errors.New("lockutil: lock is released")
)
var holderSequence atomic.Uint64
// FileLock is an exclusive advisory lock held through an open file handle.
// Release is idempotent. The lock path remains on disk after release; deleting
// or replacing it would let different processes lock different inodes.
type FileLock struct {
mu sync.Mutex
file *os.File
state platformLockState
released bool
releaseErr error
}
// TryAcquireFileLock attempts to take an exclusive advisory lock without
// waiting. A contended lock returns ErrLockHeld. The kernel releases a lock if
// its process exits, so an old unlocked file never needs stale reclamation.
func TryAcquireFileLock(path string) (*FileLock, error) {
return TryAcquireFileLockAt(filepath.Dir(path), path)
}
// TryAcquireFileLockAt attempts to lock path while confining every component
// below root to a handle-relative traversal. root is the caller's trusted
// boundary and may itself be reached through a legitimate link; links and
// reparse points below it are rejected before metadata can be written.
func TryAcquireFileLockAt(root, path string) (*FileLock, error) {
if root == "" {
return nil, errors.New("lockutil: lock root is empty")
}
relative, err := filepath.Rel(root, path)
if err != nil {
return err
return nil, fmt.Errorf("lockutil: resolve lock path relative to root: %w", err)
}
dst, err := os.OpenFile(staged, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
relative = filepath.Clean(relative)
if relative == "." || relative == ".." || filepath.IsAbs(relative) || strings.HasPrefix(relative, ".."+string(filepath.Separator)) {
return nil, fmt.Errorf("lockutil: lock path %q escapes root %q", path, root)
}
file, err := openLockFileAt(root, relative, path)
if err != nil {
_ = src.Close()
return err
return nil, fmt.Errorf("lockutil: open lock file: %w", err)
}
_, err = io.Copy(dst, src)
// Close the source before removing anything below: Go opens files without
// FILE_SHARE_DELETE on Windows, so deleting reclaimed or staged while src
// is open would fail with a sharing violation.
_ = src.Close()
state, contended, err := tryLockFile(file)
if err != nil {
_ = dst.Close()
_ = os.Remove(staged)
return err
_ = file.Close()
return nil, fmt.Errorf("lockutil: acquire advisory lock: %w", err)
}
if err := dst.Close(); err != nil {
_ = os.Remove(staged)
return err
if contended {
_ = file.Close()
return nil, ErrLockHeld
}
if err := publish(staged, path); err != nil {
_ = os.Remove(staged)
return err
lock := &FileLock{file: file, state: state}
metadata := []byte(fmt.Sprintf("%d-%d\n", os.Getpid(), holderSequence.Add(1)))
if err := lock.WriteMetadata(metadata); err != nil {
return nil, errors.Join(err, lock.Release())
}
return lock, nil
}
// WriteMetadata replaces the diagnostic contents of the lock file while the
// caller holds it. Metadata does not establish ownership; the kernel lock does.
func (lock *FileLock) WriteMetadata(data []byte) error {
if lock == nil {
return ErrLockReleased
}
lock.mu.Lock()
defer lock.mu.Unlock()
if lock.released || lock.file == nil {
return ErrLockReleased
}
if err := lock.file.Truncate(0); err != nil {
return fmt.Errorf("lockutil: truncate lock metadata: %w", err)
}
if _, err := lock.file.Seek(0, 0); err != nil {
return fmt.Errorf("lockutil: seek lock metadata: %w", err)
}
if _, err := lock.file.Write(data); err != nil {
return fmt.Errorf("lockutil: write lock metadata: %w", err)
}
// publish may leave staged behind (a hard-link publish keeps both names),
// so remove it explicitly; a move-based publish already made this a no-op.
_ = os.Remove(staged)
// The lock is back at path, so the restore has succeeded; failing to clean
// up the sidelined name must not be reported as a failed restore.
_ = RemoveLockFile(reclaimed)
return nil
}
// Release unlocks and closes the held file handle. It deliberately leaves the
// stable lock file in place. Repeated calls return the first release result.
func (lock *FileLock) Release() error {
if lock == nil {
return nil
}
lock.mu.Lock()
defer lock.mu.Unlock()
if lock.released {
return lock.releaseErr
}
lock.released = true
unlockErr := unlockFile(lock.file, &lock.state)
closeErr := lock.file.Close()
lock.file = nil
if err := errors.Join(unlockErr, closeErr); err != nil {
lock.releaseErr = fmt.Errorf("lockutil: release advisory lock: %w", err)
}
return lock.releaseErr
}
+82 -33
View File
@@ -4,43 +4,92 @@ package lockutil
import (
"errors"
"io/fs"
"fmt"
"os"
"path/filepath"
"strings"
"golang.org/x/sys/unix"
)
// RestoreLockFile restores a sidelined lock file on non-Windows platforms. It
// uses os.Link so a competing lock created at path in the meantime wins: the
// link fails with os.ErrExist instead of overwriting it. If the link fails for
// any other reason (hardlink-incapable filesystems such as FAT or some
// FUSE/network mounts, ENOSPC, EPERM), it falls back to an O_EXCL copy rather
// than leaving path missing and the live holder's lock stranded in reclaimed.
// Once the lock is back at path the restore has succeeded, so a failed cleanup
// of the sidelined name is not reported as an error; the leftover file is
// invisible to the lock protocol.
func RestoreLockFile(reclaimed, path string) error {
err := os.Link(reclaimed, path)
type platformLockState struct{}
func openLockFileAt(root, relative, displayPath string) (_ *os.File, resultErr error) {
parentFD, err := unix.Open(root, unix.O_RDONLY|unix.O_DIRECTORY|unix.O_CLOEXEC, 0)
if err != nil {
return nil, fmt.Errorf("open lock root: %w", err)
}
defer func() {
if parentFD >= 0 {
resultErr = errors.Join(resultErr, unix.Close(parentFD))
}
}()
components := strings.Split(filepath.Clean(relative), string(filepath.Separator))
for _, component := range components[:len(components)-1] {
nextFD, openErr := unix.Openat(parentFD, component, unix.O_RDONLY|unix.O_DIRECTORY|unix.O_CLOEXEC|unix.O_NOFOLLOW, 0)
if openErr != nil {
return nil, fmt.Errorf("open lock path component %q: %w", component, openErr)
}
if closeErr := unix.Close(parentFD); closeErr != nil {
_ = unix.Close(nextFD)
parentFD = -1
return nil, fmt.Errorf("close lock path component: %w", closeErr)
}
parentFD = nextFD
}
name := components[len(components)-1]
fd := -1
// Darwin can transiently return ENOENT when many callers race to create the
// same O_NOFOLLOW entry. The bound parent handle cannot be redirected, so a
// bounded retry preserves first-use contention without weakening no-follow.
for attempt := 0; attempt < 10; attempt++ {
fd, err = unix.Openat(parentFD, name, unix.O_CREAT|unix.O_RDWR|unix.O_CLOEXEC|unix.O_NOFOLLOW, 0o600)
if !errors.Is(err, unix.ENOENT) {
break
}
}
if err != nil {
return nil, err
}
file := os.NewFile(uintptr(fd), displayPath)
if file == nil {
_ = unix.Close(fd)
return nil, errors.New("convert lock handle to file")
}
var stat unix.Stat_t
if err := unix.Fstat(fd, &stat); err != nil {
return nil, errors.Join(fmt.Errorf("inspect lock file: %w", err), file.Close())
}
if stat.Mode&unix.S_IFMT != unix.S_IFREG {
return nil, errors.Join(errors.New("refusing non-regular lock file"), file.Close())
}
if stat.Nlink != 1 {
return nil, errors.Join(errors.New("refusing multiply-linked lock file"), file.Close())
}
if err := unix.Close(parentFD); err != nil {
parentFD = -1
return nil, errors.Join(fmt.Errorf("close lock parent: %w", err), file.Close())
}
parentFD = -1
return file, nil
}
func tryLockFile(file *os.File) (platformLockState, bool, error) {
err := unix.Flock(int(file.Fd()), unix.LOCK_EX|unix.LOCK_NB)
if err == nil {
_ = RemoveLockFile(reclaimed)
return platformLockState{}, false, nil
}
if errors.Is(err, unix.EWOULDBLOCK) || errors.Is(err, unix.EAGAIN) {
return platformLockState{}, true, nil
}
return platformLockState{}, false, err
}
func unlockFile(file *os.File, _ *platformLockState) error {
if file == nil {
return nil
}
if errors.Is(err, os.ErrExist) {
return err
}
return restoreByCopy(reclaimed, path, os.Link)
}
// isReclaimContended reports whether a failed rename-aside of a suspected
// stale lock means it was lost to a racer rather than a hard failure. POSIX
// rename has no contention errno (a lost race surfaces only as ENOENT, which
// ReclaimStaleLock already treats as benign), so nothing extra is benign here.
func isReclaimContended(error) bool { return false }
// RemoveLockFile removes a lock file on non-Windows platforms. Removing an
// already-missing file reports nil, matching the Windows implementation, so
// callers see one cross-platform contract.
func RemoveLockFile(path string) error {
if err := os.Remove(path); err != nil && !errors.Is(err, fs.ErrNotExist) {
return err
}
return nil
return unix.Flock(int(file.Fd()), unix.LOCK_UN)
}
+205 -101
View File
@@ -1,116 +1,220 @@
package lockutil
import (
"bufio"
"errors"
"fmt"
"io"
"os"
"os/exec"
"path/filepath"
"runtime"
"strings"
"sync"
"testing"
)
func TestRestoreLockFile(t *testing.T) {
tempDir := t.TempDir()
path := filepath.Join(tempDir, "lock")
reclaimed := filepath.Join(tempDir, "lock.stale.token")
// 1. Successful restore when target does not exist
if err := os.WriteFile(reclaimed, []byte("token"), 0o600); err != nil {
t.Fatal(err)
}
if err := RestoreLockFile(reclaimed, path); err != nil {
t.Fatalf("RestoreLockFile failed: %v", err)
}
if data, err := os.ReadFile(path); err != nil || string(data) != "token" {
t.Fatalf("restored lock content = %q, err %v; want %q intact", data, err, "token")
}
if _, err := os.Stat(reclaimed); !os.IsNotExist(err) {
t.Fatalf("expected sidelined lock to be cleaned up/removed: %v", err)
}
// 2. Fail when target already exists (prevent overwrite)
if err := os.WriteFile(reclaimed, []byte("token"), 0o600); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte("competing"), 0o600); err != nil {
t.Fatal(err)
}
err := RestoreLockFile(reclaimed, path)
if err == nil {
t.Fatal("expected RestoreLockFile to fail when target exists")
}
if !errors.Is(err, os.ErrExist) {
t.Fatalf("expected os.ErrExist, got: %v", err)
}
// The sidelined lock must still exist on failure, and the competing lock
// must be untouched.
if _, err := os.Stat(reclaimed); err != nil {
t.Fatalf("expected sidelined lock to still exist: %v", err)
}
if data, _ := os.ReadFile(path); string(data) != "competing" {
t.Fatalf("competing lock was clobbered: %q", data)
}
}
func TestRestoreByCopy(t *testing.T) {
tempDir := t.TempDir()
path := filepath.Join(tempDir, "lock")
reclaimed := filepath.Join(tempDir, "lock.stale.token")
// 1. Restores content and cleans up the sidelined file when the target is
// free (the fallback taken when the platform's no-replace primitive fails
// for a reason other than the destination existing).
if err := os.WriteFile(reclaimed, []byte("token"), 0o600); err != nil {
t.Fatal(err)
}
if err := restoreByCopy(reclaimed, path, os.Link); err != nil {
t.Fatalf("restoreByCopy failed: %v", err)
}
if data, err := os.ReadFile(path); err != nil || string(data) != "token" {
t.Fatalf("restored lock content = %q, err %v; want %q intact", data, err, "token")
}
if _, err := os.Stat(reclaimed); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("expected sidelined lock to be cleaned up after copy: %v", err)
}
if _, err := os.Stat(reclaimed + ".copy"); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("expected staged copy to be cleaned up: %v", err)
}
// 2. Keeps the no-overwrite guarantee: a competing lock at the target wins
// and the sidelined lock survives for the caller's ErrExist cleanup.
if err := os.WriteFile(reclaimed, []byte("token"), 0o600); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte("competing"), 0o600); err != nil {
t.Fatal(err)
}
if err := restoreByCopy(reclaimed, path, os.Link); !errors.Is(err, os.ErrExist) {
t.Fatalf("expected os.ErrExist when target exists, got: %v", err)
}
if data, _ := os.ReadFile(path); string(data) != "competing" {
t.Fatalf("competing lock was clobbered: %q", data)
}
if _, err := os.Stat(reclaimed); err != nil {
t.Fatalf("expected sidelined lock to still exist: %v", err)
}
if _, err := os.Stat(reclaimed + ".copy"); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("expected staged copy to be cleaned up after a failed publish: %v", err)
}
}
func TestRemoveLockFile(t *testing.T) {
func TestFileLockSerializesAndKeepsStablePath(t *testing.T) {
path := filepath.Join(t.TempDir(), "lock")
if err := os.WriteFile(path, []byte("token"), 0o600); err != nil {
t.Fatal(err)
first, err := TryAcquireFileLock(path)
if err != nil {
t.Fatalf("first acquire: %v", err)
}
if err := RemoveLockFile(path); err != nil {
t.Fatalf("RemoveLockFile failed: %v", err)
firstInfo, err := os.Stat(path)
if err != nil {
t.Fatalf("stat first lock path: %v", err)
}
if _, err := os.Stat(path); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("expected lock file to be removed: %v", err)
if _, err := TryAcquireFileLock(path); !errors.Is(err, ErrLockHeld) {
t.Fatalf("contended acquire = %v, want ErrLockHeld", err)
}
// Removing an already-missing lock file is a no-op on every platform.
if err := RemoveLockFile(path); err != nil {
t.Fatalf("RemoveLockFile on a missing file = %v, want nil", err)
if err := first.Release(); err != nil {
t.Fatalf("release first: %v", err)
}
releasedInfo, err := os.Stat(path)
if err != nil {
t.Fatalf("release removed stable lock path: %v", err)
}
if !os.SameFile(firstInfo, releasedInfo) {
t.Fatal("release replaced the stable lock file")
}
second, err := TryAcquireFileLock(path)
if err != nil {
t.Fatalf("acquire after release: %v", err)
}
secondInfo, err := os.Stat(path)
if err != nil {
t.Fatalf("stat second lock path: %v", err)
}
if !os.SameFile(firstInfo, secondInfo) {
t.Fatal("reacquire used a different lock file")
}
if err := second.Release(); err != nil {
t.Fatalf("release second: %v", err)
}
if err := second.Release(); err != nil {
t.Fatalf("idempotent release: %v", err)
}
}
func TestFileLockMetadata(t *testing.T) {
path := filepath.Join(t.TempDir(), "lock")
lock, err := TryAcquireFileLock(path)
if err != nil {
t.Fatal(err)
}
if err := lock.WriteMetadata([]byte("holder\n")); err != nil {
t.Fatalf("write metadata: %v", err)
}
data, err := os.ReadFile(path)
if err != nil || string(data) != "holder\n" {
t.Fatalf("metadata = %q, err %v", data, err)
}
if err := lock.Release(); err != nil {
t.Fatal(err)
}
if err := lock.WriteMetadata(nil); !errors.Is(err, ErrLockReleased) {
t.Fatalf("write after release = %v, want ErrLockReleased", err)
}
}
func TestFileLockConcurrentFirstAcquisition(t *testing.T) {
path := filepath.Join(t.TempDir(), "lock")
const contenders = 200
start := make(chan struct{})
results := make(chan error, contenders)
locks := make(chan *FileLock, contenders)
var ready sync.WaitGroup
ready.Add(contenders)
for range contenders {
go func() {
ready.Done()
<-start
lock, err := TryAcquireFileLock(path)
if lock != nil {
locks <- lock
}
results <- err
}()
}
ready.Wait()
close(start)
acquired := 0
for range contenders {
err := <-results
switch {
case err == nil:
acquired++
case !errors.Is(err, ErrLockHeld):
t.Errorf("concurrent first acquire: %v", err)
}
}
close(locks)
for lock := range locks {
if err := lock.Release(); err != nil {
t.Error(err)
}
}
if acquired != 1 {
t.Fatalf("concurrent first acquire: got %d holders, want 1", acquired)
}
}
func TestFileLockRefusesMultiplyLinkedEntry(t *testing.T) {
dir := t.TempDir()
target := filepath.Join(dir, "target")
lockPath := filepath.Join(dir, "lock")
const sentinel = "do not overwrite"
if err := os.WriteFile(target, []byte(sentinel), 0o600); err != nil {
t.Fatal(err)
}
if err := os.Link(target, lockPath); err != nil {
t.Fatalf("create redirected hard link: %v", err)
}
if lock, err := TryAcquireFileLock(lockPath); err == nil {
_ = lock.Release()
t.Fatal("TryAcquireFileLock accepted a multiply-linked entry")
}
data, err := os.ReadFile(target)
if err != nil {
t.Fatal(err)
}
if string(data) != sentinel {
t.Fatalf("redirected target was modified: %q", data)
}
}
func TestFileLockAtRejectsEscapingPath(t *testing.T) {
base := t.TempDir()
outside := filepath.Join(filepath.Dir(base), "outside.lock")
if lock, err := TryAcquireFileLockAt(base, outside); err == nil {
_ = lock.Release()
t.Fatal("TryAcquireFileLockAt accepted a path outside its root")
}
if _, err := os.Stat(outside); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("escaping path was created: %v", err)
}
}
// TestFileLockReleasedByProcessExit proves crash recovery is provided by the
// kernel lock lifetime. The child exits without calling Release; the next
// process can then lock the same persistent file without moving or deleting it.
func TestFileLockReleasedByProcessExit(t *testing.T) {
path := filepath.Join(t.TempDir(), "lock")
cmd := exec.Command(os.Args[0], "-test.run=^TestFileLockHelperProcess$")
cmd.Env = append(os.Environ(), "ZERO_LOCKUTIL_HELPER=1", "ZERO_LOCKUTIL_PATH="+path)
stdout, err := cmd.StdoutPipe()
if err != nil {
t.Fatal(err)
}
stdin, err := cmd.StdinPipe()
if err != nil {
t.Fatal(err)
}
cmd.Stderr = os.Stderr
if err := cmd.Start(); err != nil {
t.Fatal(err)
}
ready, err := bufio.NewReader(stdout).ReadString('\n')
if err != nil || strings.TrimSpace(ready) != "ready" {
_ = cmd.Process.Kill()
_ = cmd.Wait()
t.Fatalf("helper readiness = %q, err %v", ready, err)
}
if _, err := TryAcquireFileLock(path); !errors.Is(err, ErrLockHeld) {
_ = cmd.Process.Kill()
_ = cmd.Wait()
t.Fatalf("acquire while helper holds lock = %v, want ErrLockHeld", err)
}
if err := stdin.Close(); err != nil {
t.Fatal(err)
}
if err := cmd.Wait(); err != nil {
t.Fatalf("helper exit: %v", err)
}
lock, err := TryAcquireFileLock(path)
if err != nil {
t.Fatalf("acquire after helper exit: %v", err)
}
if err := lock.Release(); err != nil {
t.Fatal(err)
}
}
func TestFileLockHelperProcess(t *testing.T) {
if os.Getenv("ZERO_LOCKUTIL_HELPER") != "1" {
return
}
lock, err := TryAcquireFileLock(os.Getenv("ZERO_LOCKUTIL_PATH"))
if err != nil {
fmt.Fprintln(os.Stderr, err)
os.Exit(2)
}
if _, err := fmt.Fprintln(os.Stdout, "ready"); err != nil {
os.Exit(3)
}
_, _ = io.Copy(io.Discard, os.Stdin)
runtime.KeepAlive(lock) // deliberately exit without Release
}
+183 -58
View File
@@ -4,77 +4,202 @@ package lockutil
import (
"errors"
"fmt"
"os"
"syscall"
"time"
"path/filepath"
"strings"
"unsafe"
"golang.org/x/sys/windows"
)
// RestoreLockFile restores a sidelined lock file on Windows. It moves the file
// with a no-replace MoveFileEx so a competing lock created at path in the
// meantime wins: the move fails with os.ErrExist instead of overwriting it. If
// the move fails for any other reason (ERROR_SHARING_VIOLATION or
// ERROR_ACCESS_DENIED under concurrent access to the file), it falls back to
// an O_EXCL copy rather than leaving path missing and the live holder's lock
// stranded in reclaimed.
func RestoreLockFile(reclaimed, path string) error {
err := moveFileNoReplace(reclaimed, path)
if err == nil || errors.Is(err, os.ErrExist) {
return err
}
return restoreByCopy(reclaimed, path, moveFileNoReplace)
type platformLockState struct {
overlapped windows.Overlapped
}
// moveFileNoReplace renames from to to, failing if to already exists. It calls
// MoveFileExW directly via golang.org/x/sys/windows (the standard library's
// syscall package does not export MoveFileEx on this platform) with no flags,
// so an existing destination fails with ERROR_ALREADY_EXISTS, which satisfies
// errors.Is(err, os.ErrExist) (pinned by
// TestMoveFileNoReplaceMapsAlreadyExistsToErrExist), instead of overwriting it
// the way os.Rename does on Windows (it passes MOVEFILE_REPLACE_EXISTING to
// match POSIX rename semantics cross-platform).
func moveFileNoReplace(from, to string) error {
fromPtr, err := windows.UTF16PtrFromString(from)
const lockByteOffsetHigh = 1
func openLockFileAt(root, relative, displayPath string) (_ *os.File, resultErr error) {
rootPtr, err := windows.UTF16PtrFromString(root)
if err != nil {
return err
return nil, err
}
toPtr, err := windows.UTF16PtrFromString(to)
parent, err := windows.CreateFile(
rootPtr,
windows.FILE_READ_ATTRIBUTES|windows.FILE_TRAVERSE|windows.SYNCHRONIZE,
windows.FILE_SHARE_READ|windows.FILE_SHARE_WRITE|windows.FILE_SHARE_DELETE,
nil,
windows.OPEN_EXISTING,
windows.FILE_FLAG_BACKUP_SEMANTICS,
0,
)
if err != nil {
return err
return nil, fmt.Errorf("open lock root: %w", err)
}
return windows.MoveFileEx(fromPtr, toPtr, 0)
}
// isReclaimContended reports whether a failed rename-aside of a suspected
// stale lock means the file was concurrently open or pending deletion (which
// surfaces as ERROR_SHARING_VIOLATION or ERROR_ACCESS_DENIED on Windows)
// rather than a hard failure, so ReclaimStaleLock reports a lost race instead
// of an error. Neither errno satisfies errors.Is(err, os.ErrExist), and
// ERROR_SHARING_VIOLATION does not satisfy errors.Is(err, os.ErrPermission)
// either, so the classification compares raw errnos like RemoveLockFile below.
func isReclaimContended(err error) bool {
var errno syscall.Errno
return errors.As(err, &errno) && (errno == windows.ERROR_SHARING_VIOLATION || errno == windows.ERROR_ACCESS_DENIED)
}
// RemoveLockFile removes a lock file on Windows, retrying on the sharing
// violation or access denied errors that are common under heavy concurrent
// contention. Removing an already-missing file reports nil, matching the
// non-Windows implementation, so callers see one cross-platform contract.
func RemoveLockFile(path string) error {
var err error
for i := 0; i < 15; i++ {
err = os.Remove(path)
if err == nil || os.IsNotExist(err) {
return nil
defer func() {
if parent != 0 {
resultErr = errors.Join(resultErr, windows.CloseHandle(parent))
}
var errno syscall.Errno
if errors.As(err, &errno) && (errno == windows.ERROR_SHARING_VIOLATION || errno == windows.ERROR_ACCESS_DENIED) {
time.Sleep(5 * time.Millisecond)
}()
components := strings.Split(filepath.Clean(relative), string(filepath.Separator))
for _, component := range components[:len(components)-1] {
next, openErr := openLockDirectoryAt(parent, component)
if openErr != nil {
return nil, fmt.Errorf("open lock path component %q: %w", component, openErr)
}
if closeErr := windows.CloseHandle(parent); closeErr != nil {
_ = windows.CloseHandle(next)
parent = 0
return nil, fmt.Errorf("close lock path component: %w", closeErr)
}
parent = next
}
handle, err := openLockFileHandleAt(parent, components[len(components)-1])
if err != nil {
return nil, err
}
file := os.NewFile(uintptr(handle), displayPath)
if file == nil {
_ = windows.CloseHandle(handle)
return nil, errors.New("convert lock handle to file")
}
if err := windows.CloseHandle(parent); err != nil {
parent = 0
return nil, errors.Join(fmt.Errorf("close lock parent: %w", err), file.Close())
}
parent = 0
return file, nil
}
func openLockDirectoryAt(parent windows.Handle, name string) (windows.Handle, error) {
handle, err := openLockPathAt(
parent,
name,
windows.FILE_READ_ATTRIBUTES|windows.FILE_TRAVERSE|windows.SYNCHRONIZE,
windows.FILE_OPEN,
windows.FILE_ATTRIBUTE_DIRECTORY,
windows.FILE_DIRECTORY_FILE|windows.FILE_SYNCHRONOUS_IO_NONALERT|windows.FILE_OPEN_REPARSE_POINT,
)
if err != nil {
return 0, err
}
if err := validateWindowsLockHandle(handle, true); err != nil {
return 0, errors.Join(err, windows.CloseHandle(handle))
}
return handle, nil
}
func openLockFileHandleAt(parent windows.Handle, name string) (windows.Handle, error) {
for attempt := 0; attempt < 10; attempt++ {
handle, err := openLockPathAt(
parent,
name,
windows.GENERIC_READ|windows.GENERIC_WRITE|windows.SYNCHRONIZE,
windows.FILE_OPEN,
windows.FILE_ATTRIBUTE_NORMAL,
windows.FILE_NON_DIRECTORY_FILE|windows.FILE_SYNCHRONOUS_IO_NONALERT|windows.FILE_OPEN_REPARSE_POINT,
)
if err == nil {
if err := validateWindowsLockHandle(handle, false); err != nil {
return 0, errors.Join(err, windows.CloseHandle(handle))
}
return handle, nil
}
if !windowsLockPathMissing(err) {
return 0, err
}
handle, err = openLockPathAt(
parent,
name,
windows.GENERIC_READ|windows.GENERIC_WRITE|windows.SYNCHRONIZE,
windows.FILE_CREATE,
windows.FILE_ATTRIBUTE_NORMAL,
windows.FILE_NON_DIRECTORY_FILE|windows.FILE_SYNCHRONOUS_IO_NONALERT|windows.FILE_OPEN_REPARSE_POINT,
)
if err == windows.STATUS_OBJECT_NAME_COLLISION || errors.Is(err, windows.ERROR_ALREADY_EXISTS) {
continue
}
return err
if err != nil {
return 0, err
}
if err := validateWindowsLockHandle(handle, false); err != nil {
return 0, errors.Join(err, windows.CloseHandle(handle))
}
return handle, nil
}
return err
return 0, errors.New("lock file creation did not stabilize")
}
func windowsLockPathMissing(err error) bool {
return err == windows.STATUS_NO_SUCH_FILE ||
err == windows.STATUS_OBJECT_NAME_NOT_FOUND ||
err == windows.STATUS_OBJECT_PATH_NOT_FOUND ||
errors.Is(err, windows.ERROR_FILE_NOT_FOUND) ||
errors.Is(err, windows.ERROR_PATH_NOT_FOUND)
}
func openLockPathAt(parent windows.Handle, name string, access, disposition, attributes, options uint32) (windows.Handle, error) {
objectName, err := windows.NewNTUnicodeString(name)
if err != nil {
return 0, err
}
objectAttributes := &windows.OBJECT_ATTRIBUTES{
Length: uint32(unsafe.Sizeof(windows.OBJECT_ATTRIBUTES{})),
RootDirectory: parent,
ObjectName: objectName,
Attributes: windows.OBJ_CASE_INSENSITIVE | windows.OBJ_DONT_REPARSE,
}
var handle windows.Handle
err = windows.NtCreateFile(
&handle,
access,
objectAttributes,
&windows.IO_STATUS_BLOCK{},
nil,
attributes,
windows.FILE_SHARE_READ|windows.FILE_SHARE_WRITE|windows.FILE_SHARE_DELETE,
disposition,
options,
0,
0,
)
return handle, err
}
func validateWindowsLockHandle(handle windows.Handle, wantDirectory bool) error {
var info windows.ByHandleFileInformation
if err := windows.GetFileInformationByHandle(handle, &info); err != nil {
return fmt.Errorf("inspect lock path handle: %w", err)
}
isDirectory := info.FileAttributes&windows.FILE_ATTRIBUTE_DIRECTORY != 0
if info.FileAttributes&windows.FILE_ATTRIBUTE_REPARSE_POINT != 0 || isDirectory != wantDirectory {
return errors.New("refusing non-file or reparse-point lock path")
}
if !wantDirectory && info.NumberOfLinks != 1 {
return errors.New("refusing multiply-linked lock file")
}
return nil
}
func tryLockFile(file *os.File) (platformLockState, bool, error) {
state := platformLockState{overlapped: windows.Overlapped{OffsetHigh: lockByteOffsetHigh}}
flags := uint32(windows.LOCKFILE_EXCLUSIVE_LOCK | windows.LOCKFILE_FAIL_IMMEDIATELY)
err := windows.LockFileEx(windows.Handle(file.Fd()), flags, 0, 1, 0, &state.overlapped)
if err == nil {
return state, false, nil
}
if errors.Is(err, windows.ERROR_LOCK_VIOLATION) || errors.Is(err, windows.ERROR_SHARING_VIOLATION) {
return platformLockState{}, true, nil
}
return platformLockState{}, false, err
}
func unlockFile(file *os.File, state *platformLockState) error {
if file == nil {
return nil
}
return windows.UnlockFileEx(windows.Handle(file.Fd()), 0, 1, 0, &state.overlapped)
}
@@ -1,83 +0,0 @@
//go:build windows
package lockutil
import (
"errors"
"os"
"path/filepath"
"testing"
"time"
"golang.org/x/sys/windows"
)
// TestMoveFileNoReplaceMapsAlreadyExistsToErrExist pins the errno coupling
// every caller keys its cleanup on: with no flags, MoveFileEx fails with
// ERROR_ALREADY_EXISTS when the destination exists, and that errno must
// satisfy errors.Is(err, os.ErrExist).
func TestMoveFileNoReplaceMapsAlreadyExistsToErrExist(t *testing.T) {
dir := t.TempDir()
from := filepath.Join(dir, "from")
to := filepath.Join(dir, "to")
if err := os.WriteFile(from, []byte("sidelined"), 0o600); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(to, []byte("competing"), 0o600); err != nil {
t.Fatal(err)
}
err := moveFileNoReplace(from, to)
if !errors.Is(err, windows.ERROR_ALREADY_EXISTS) {
t.Fatalf("MoveFileEx with an existing destination = %v, want ERROR_ALREADY_EXISTS", err)
}
if !errors.Is(err, os.ErrExist) {
t.Fatalf("ERROR_ALREADY_EXISTS must satisfy errors.Is(err, os.ErrExist), got: %v", err)
}
if data, _ := os.ReadFile(to); string(data) != "competing" {
t.Fatalf("destination was overwritten: %q", data)
}
}
// TestRemoveLockFileRetriesSharingViolation holds the file open briefly (Go
// opens files without FILE_SHARE_DELETE on Windows, so a concurrent DeleteFile
// fails with ERROR_SHARING_VIOLATION while the handle is open); RemoveLockFile
// must ride its retry loop (15 x 5ms) past the window and succeed once the
// handle closes.
func TestRemoveLockFileRetriesSharingViolation(t *testing.T) {
path := filepath.Join(t.TempDir(), "lock")
if err := os.WriteFile(path, []byte("token"), 0o600); err != nil {
t.Fatal(err)
}
f, err := os.Open(path)
if err != nil {
t.Fatal(err)
}
go func() {
time.Sleep(10 * time.Millisecond)
_ = f.Close()
}()
if err := RemoveLockFile(path); err != nil {
t.Fatalf("RemoveLockFile should succeed once the contending handle closes: %v", err)
}
if _, err := os.Stat(path); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("expected lock file to be removed: %v", err)
}
}
// TestRemoveLockFileSharingViolationExhausts verifies the retry loop is
// bounded: with the contending handle held for the whole call, RemoveLockFile
// surfaces the sharing violation instead of retrying forever.
func TestRemoveLockFileSharingViolationExhausts(t *testing.T) {
path := filepath.Join(t.TempDir(), "lock")
if err := os.WriteFile(path, []byte("token"), 0o600); err != nil {
t.Fatal(err)
}
f, err := os.Open(path)
if err != nil {
t.Fatal(err)
}
defer f.Close()
if err := RemoveLockFile(path); !errors.Is(err, windows.ERROR_SHARING_VIOLATION) {
t.Fatalf("RemoveLockFile with a held handle = %v, want ERROR_SHARING_VIOLATION", err)
}
}
-82
View File
@@ -1,82 +0,0 @@
package lockutil
import (
"errors"
"os"
)
// restoreLockFile is swappable so tests can force the fail-closed path of
// ReclaimStaleLock, which requires both the fast restore and its no-replace
// fallback (with its own copy fallback) to fail; that cannot be provoked
// portably on a healthy filesystem.
var restoreLockFile = restoreLiveLock
// restoreLiveLock puts a lock that turned out to be live back at path after
// ReclaimStaleLock moved it aside to inspect it. It first tries a fast,
// replacing rename straight from reclaimed to path: a single syscall, which
// keeps the window during which path does not exist (and so is open to a
// fourth process's unrelated O_EXCL create landing on it) as short as
// possible. RestoreLockFile's no-replace restore (with its slower copy
// fallback on some failures) is a correctness-preserving fallback for when
// the fast path itself fails (e.g. a cross-device sidelined name): it is a
// much longer version of the identical race, but still detects rather than
// silently clobbers a competing lock, which the fast path's replacing rename
// cannot do. Neither path makes the race impossible, only unlikely; see
// ReclaimStaleLock's doc comment.
func restoreLiveLock(reclaimed, path string) error {
if err := os.Rename(reclaimed, path); err == nil {
return nil
}
return RestoreLockFile(reclaimed, path)
}
// ReclaimStaleLock atomically reclaims a suspected-stale lock file. It renames
// lockPath aside to "<lockPath>.stale.<suffix>" (only one racer can win the
// rename of a given file, so two racers can never both reclaim the same lock),
// then consults isLive on the moved file; if the lock turns out to be live (a
// holder reacquired it in the gap between the caller's stale check and the
// rename) it is restored rather than stolen. The suffix must be unique per
// acquirer attempt. Returns true only when a genuinely stale lock was removed,
// so the caller knows it is safe to retry its exclusive create immediately; on
// a lost race it returns false. A non-nil error means either the rename aside
// failed for a reason that is not contention (so retrying cannot help and the
// caller should fail fast instead of spinning to its deadline), or a live
// holder's lock could not be restored (both restoreLiveLock's fast path and
// its no-replace fallback failed), so lockPath may be missing; callers must
// fail closed instead of re-acquiring. The sidelined file is removed on every
// restore failure: once the restore has failed it has no protocol function
// (release only consults the lock path), so keeping it would only leak files.
//
// The live-restore path has an inherent, unclosed race: between the rename
// aside above and restoreLiveLock putting the lock back, lockPath does not
// exist, so an unrelated caller's O_EXCL create can legitimately succeed
// there. restoreLiveLock's fast path then silently overwrites that new
// claimant's lock file, which does not corrupt release (it is
// ownership-aware, so the new claimant's later release safely no-ops against
// content it no longer owns) but does mean the new claimant can still run
// its critical section concurrently with the original live holder. Making
// this race actually impossible would need an OS-level advisory lock (flock
// / LockFileEx) held for a holder's whole critical section, checked
// non-destructively instead of by moving the file; restoreLiveLock only
// shrinks the window to roughly one syscall, it does not close it.
func ReclaimStaleLock(lockPath, suffix string, isLive func(reclaimedPath string) bool) (bool, error) {
reclaimed := lockPath + ".stale." + suffix
if err := os.Rename(lockPath, reclaimed); err != nil {
if errors.Is(err, os.ErrNotExist) || isReclaimContended(err) {
return false, nil // another racer already moved/removed it, or it vanished
}
return false, err
}
if isLive(reclaimed) {
// Put the live lock back instead of stealing it, and let the caller wait.
if rerr := restoreLockFile(reclaimed, lockPath); rerr != nil {
_ = RemoveLockFile(reclaimed)
if !errors.Is(rerr, os.ErrExist) {
return false, rerr
}
}
return false, nil
}
_ = RemoveLockFile(reclaimed)
return true, nil
}
-31
View File
@@ -1,31 +0,0 @@
//go:build !windows
package lockutil
import (
"os"
"path/filepath"
"testing"
)
// TestReclaimStaleLockPropagatesRenameFailure pins the fail-fast contract of
// the rename aside: a failure that is not a lost race must reach the caller.
// Mapping it to the benign lost-race result would leave acquirers spinning
// until their deadline and reporting a timeout instead of the real cause. A
// directory squatting on the sidelined name makes os.Rename fail with EISDIR,
// which matches neither os.ErrNotExist nor a contention errno.
func TestReclaimStaleLockPropagatesRenameFailure(t *testing.T) {
lockPath := filepath.Join(t.TempDir(), "lock")
if err := os.WriteFile(lockPath, []byte("stale-holder"), 0o600); err != nil {
t.Fatal(err)
}
if err := os.Mkdir(lockPath+".stale.tok", 0o755); err != nil {
t.Fatal(err)
}
if ok, err := ReclaimStaleLock(lockPath, "tok", func(string) bool { return false }); err == nil || ok {
t.Fatalf("a hard rename failure must surface an error (ok=%v err=%v)", ok, err)
}
if _, err := os.Stat(lockPath); err != nil {
t.Fatalf("the lock file must be left in place on a failed rename: %v", err)
}
}
-112
View File
@@ -1,112 +0,0 @@
package lockutil
import (
"errors"
"os"
"path/filepath"
"testing"
)
func TestReclaimStaleLock(t *testing.T) {
lockPath := filepath.Join(t.TempDir(), "lock")
dead := func(string) bool { return false }
live := func(string) bool { return true }
// A lock the predicate reports dead is reclaimed and removed.
if err := os.WriteFile(lockPath, []byte("crashed-holder"), 0o600); err != nil {
t.Fatal(err)
}
if ok, err := ReclaimStaleLock(lockPath, "tok-a", dead); err != nil || !ok {
t.Fatalf("a dead lock should be reclaimed (ok=%v err=%v)", ok, err)
}
if _, err := os.Stat(lockPath); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("reclaimed dead lock should be gone, stat err=%v", err)
}
// A LIVE lock (a holder reacquired in the gap) must be RESTORED intact.
if err := os.WriteFile(lockPath, []byte("live-holder"), 0o600); err != nil {
t.Fatal(err)
}
if ok, err := ReclaimStaleLock(lockPath, "tok-b", live); err != nil || ok {
t.Fatalf("a live lock must not be reclaimed (ok=%v err=%v)", ok, err)
}
if data, err := os.ReadFile(lockPath); err != nil || string(data) != "live-holder" {
t.Fatalf("live lock must be left intact, got %q err %v", data, err)
}
// A missing lock reports no reclaim (nothing to steal).
_ = os.Remove(lockPath)
if ok, err := ReclaimStaleLock(lockPath, "tok-c", live); err != nil || ok {
t.Fatalf("a missing lock should not report a reclaim (ok=%v err=%v)", ok, err)
}
}
// restoreLiveLock's fast path is a single replacing rename, chosen to
// minimize (not eliminate) the window during which lockPath is absent and
// open to an unrelated O_EXCL create; see ReclaimStaleLock's doc comment for
// why that residual race exists. This test pins the resulting behavior: the
// fast path overwrites a competing file at path rather than detecting it,
// unlike RestoreLockFile's own no-replace contract.
func TestRestoreLiveLockFastPathOverwritesCompetingFile(t *testing.T) {
dir := t.TempDir()
reclaimed := filepath.Join(dir, "lock.stale.tok")
path := filepath.Join(dir, "lock")
if err := os.WriteFile(reclaimed, []byte("original-holder"), 0o600); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte("new-claimant"), 0o600); err != nil {
t.Fatal(err)
}
if err := restoreLiveLock(reclaimed, path); err != nil {
t.Fatalf("restoreLiveLock failed: %v", err)
}
data, err := os.ReadFile(path)
if err != nil || string(data) != "original-holder" {
t.Fatalf("path = %q, err %v; want the original holder's content restored", data, err)
}
if _, err := os.Stat(reclaimed); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("expected reclaimed to be consumed by the rename: %v", err)
}
}
func TestReclaimStaleLockFailsClosedOnRestoreError(t *testing.T) {
// When both the no-replace restore and its copy fallback fail (only
// provokable via the seam; a healthy filesystem cannot produce it), the
// caller must receive an error so it fails closed instead of re-acquiring a
// missing lock path, and the sidelined file must not leak.
restoreLockFile = func(reclaimed, path string) error { return errors.New("restore failed") }
defer func() { restoreLockFile = restoreLiveLock }()
lockPath := filepath.Join(t.TempDir(), "lock")
if err := os.WriteFile(lockPath, []byte("live-holder"), 0o600); err != nil {
t.Fatal(err)
}
ok, err := ReclaimStaleLock(lockPath, "tok", func(string) bool { return true })
if err == nil || ok {
t.Fatalf("a failed restore must surface an error (ok=%v err=%v)", ok, err)
}
if matches, _ := filepath.Glob(lockPath + ".stale.*"); len(matches) != 0 {
t.Fatalf("a failed restore must not leak sidelined files: %v", matches)
}
}
func TestReclaimStaleLockDropsSidelinedWhenNewHolderWins(t *testing.T) {
// An os.ErrExist restore failure means a new holder recreated the lock
// path; that is not an error for the caller, and the sidelined file is
// dropped rather than leaked.
restoreLockFile = func(reclaimed, path string) error { return os.ErrExist }
defer func() { restoreLockFile = restoreLiveLock }()
lockPath := filepath.Join(t.TempDir(), "lock")
if err := os.WriteFile(lockPath, []byte("live-holder"), 0o600); err != nil {
t.Fatal(err)
}
ok, err := ReclaimStaleLock(lockPath, "tok", func(string) bool { return true })
if err != nil || ok {
t.Fatalf("losing to a new holder is not an error (ok=%v err=%v)", ok, err)
}
if matches, _ := filepath.Glob(lockPath + ".stale.*"); len(matches) != 0 {
t.Fatalf("the sidelined file must be dropped when a new holder wins: %v", matches)
}
}
-50
View File
@@ -1,50 +0,0 @@
//go:build windows
package lockutil
import (
"os"
"path/filepath"
"testing"
)
// TestReclaimStaleLockTreatsSharingViolationAsLostRace pins the errno
// classification the Windows acquire loops depend on: a rename aside that
// fails because the lock file is concurrently open (Go opens files without
// FILE_SHARE_DELETE, so the rename fails with ERROR_SHARING_VIOLATION or
// ERROR_ACCESS_DENIED) is contention, not a hard failure, and must report a
// benign lost race so the caller keeps waiting.
func TestReclaimStaleLockTreatsSharingViolationAsLostRace(t *testing.T) {
lockPath := filepath.Join(t.TempDir(), "lock")
if err := os.WriteFile(lockPath, []byte("held"), 0o600); err != nil {
t.Fatal(err)
}
f, err := os.Open(lockPath)
if err != nil {
t.Fatal(err)
}
defer f.Close()
if ok, rerr := ReclaimStaleLock(lockPath, "tok", func(string) bool { return true }); rerr != nil || ok {
t.Fatalf("a sharing-violation rename must report a benign lost race (ok=%v err=%v)", ok, rerr)
}
if _, err := os.Stat(lockPath); err != nil {
t.Fatalf("the lock file must be left in place: %v", err)
}
}
// TestReclaimStaleLockPropagatesRenameFailure pins the fail-fast contract of
// the rename aside on Windows: a failure that is not contention must reach
// the caller instead of masquerading as a lost race. A quote is an invalid
// filename character, so the sidelined name fails with ERROR_INVALID_NAME.
func TestReclaimStaleLockPropagatesRenameFailure(t *testing.T) {
lockPath := filepath.Join(t.TempDir(), "lock")
if err := os.WriteFile(lockPath, []byte("stale-holder"), 0o600); err != nil {
t.Fatal(err)
}
if ok, err := ReclaimStaleLock(lockPath, `bad"name`, func(string) bool { return false }); err == nil || ok {
t.Fatalf("a hard rename failure must surface an error (ok=%v err=%v)", ok, err)
}
if _, err := os.Stat(lockPath); err != nil {
t.Fatalf("the lock file must be left in place on a failed rename: %v", err)
}
}
+2 -4
View File
@@ -176,11 +176,9 @@ func RegisterTools(ctx context.Context, registry *tools.Registry, cfg config.MCP
staged = append(staged, tool)
}
}
registered := make([]tools.Tool, 0, len(staged))
for index := range staged {
registered = append(registered, staged[index])
for _, tool := range staged {
registry.Register(tool)
}
registry.RegisterBatch(registered)
return runtime, nil
}
-56
View File
@@ -70,62 +70,6 @@ func TestRegisterToolsAddsPromptGatedMCPTools(t *testing.T) {
}
}
func TestRegisterToolsPublishesServerToolsInOneGeneration(t *testing.T) {
registry := tools.NewRegistry()
before := registry.Snapshot().Generation
client := &fakeToolClient{listed: []RemoteTool{
{Name: "lookup", Description: "Lookup documentation"},
{Name: "search", Description: "Search documentation"},
}}
runtime, err := RegisterTools(context.Background(), registry, config.MCPConfig{Servers: map[string]config.MCPServerConfig{
"docs": {Type: "stdio", Command: "docs-mcp"},
}}, RegisterOptions{ClientFactory: func(context.Context, Server) (ToolClient, error) {
return client, nil
}})
if err != nil {
t.Fatalf("RegisterTools() error = %v", err)
}
defer runtime.Close()
after := registry.Snapshot()
if after.Generation != before+1 {
t.Fatalf("registry generation = %d, want %d", after.Generation, before+1)
}
for _, name := range []string{"mcp_docs_lookup", "mcp_docs_search"} {
if _, ok := registry.Get(name); !ok {
t.Fatalf("expected %s to be published in the batch", name)
}
}
}
func TestRegisterToolsPublishesNoneWhenServerToolValidationFails(t *testing.T) {
registry := tools.NewRegistry()
before := registry.Snapshot().Generation
client := &fakeToolClient{listed: []RemoteTool{
{Name: "lookup", Description: "Lookup documentation"},
{Name: "", Description: "Invalid nameless tool"},
}}
runtime, err := RegisterTools(context.Background(), registry, config.MCPConfig{Servers: map[string]config.MCPServerConfig{
"docs": {Type: "stdio", Command: "docs-mcp"},
}}, RegisterOptions{ClientFactory: func(context.Context, Server) (ToolClient, error) {
return client, nil
}})
if err != nil {
t.Fatalf("RegisterTools() error = %v", err)
}
defer runtime.Close()
if got := registry.Snapshot().Generation; got != before {
t.Fatalf("registry generation = %d, want unchanged %d", got, before)
}
if _, ok := registry.Get("mcp_docs_lookup"); ok {
t.Fatal("valid prefix of rejected server batch was published")
}
if skipped := runtime.Skipped(); len(skipped) != 1 || skipped[0].Name != "docs" || skipped[0].Err == nil {
t.Fatalf("Skipped() = %#v, want one validation failure for docs", skipped)
}
}
func TestRegisterToolsMarksPersistentlyApprovedToolsAllow(t *testing.T) {
store, err := NewPermissionStore(StoreOptions{
FilePath: filepath.Join(t.TempDir(), "permissions.json"),
+7 -64
View File
@@ -5,24 +5,16 @@ import (
"fmt"
"os"
"path/filepath"
"sync/atomic"
"time"
"github.com/Gitlawb/zero/internal/lockutil"
)
const (
fileLockTimeout = 5 * time.Second
fileLockStaleAfter = 30 * time.Second
)
const fileLockTimeout = 5 * time.Second
var lockSeq atomic.Uint64
// acquireFileLock takes a cross-process exclusive lock by creating lockPath with
// O_EXCL. It retries with a short backoff until a timeout, breaking a lock whose
// file is older than fileLockStaleAfter (so a crashed holder cannot deadlock the
// store). Release is ownership-aware: it removes the lock only if it still holds
// our token, so a stale-broken holder cannot delete a newer holder's lock.
// acquireFileLock takes a cross-process advisory lock through a stable lock
// file. It retries with a short backoff until a timeout. Kernel lock lifetime
// provides crash recovery, so the path is never renamed or removed.
func acquireFileLock(lockPath string, now func() time.Time) (func(), error) {
if now == nil {
now = time.Now
@@ -30,64 +22,15 @@ func acquireFileLock(lockPath string, now func() time.Time) (func(), error) {
if err := os.MkdirAll(filepath.Dir(lockPath), 0o700); err != nil {
return nil, err
}
token := fmt.Sprintf("%d-%d-%d", os.Getpid(), now().UnixNano(), lockSeq.Add(1))
deadline := now().Add(fileLockTimeout)
for {
f, err := os.OpenFile(lockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
lock, err := lockutil.TryAcquireFileLock(lockPath)
if err == nil {
// A partial/failed write would leave a lock file without our token, so
// the releaser below could never delete it — stranding the lock for
// other processes. Fail closed: remove the file and surface the error.
if _, werr := f.WriteString(token); werr != nil {
_ = f.Close()
_ = lockutil.RemoveLockFile(lockPath)
return nil, fmt.Errorf("oauth: write token lock: %w", werr)
}
if cerr := f.Close(); cerr != nil {
_ = lockutil.RemoveLockFile(lockPath)
return nil, fmt.Errorf("oauth: close token lock: %w", cerr)
}
var released bool
return func() {
if released {
return
}
released = true
if data, rerr := os.ReadFile(lockPath); rerr == nil && string(data) == token {
_ = lockutil.RemoveLockFile(lockPath)
}
}, nil
return func() { _ = lock.Release() }, nil
}
// On Windows a concurrent holder's os.Remove leaves the lock file in a
// "delete pending" state, so an O_EXCL create races it with
// ERROR_ACCESS_DENIED (os.ErrPermission) rather than ErrExist. Treat that
// as contention and retry, exactly like ErrExist — otherwise the lock
// spuriously fails under concurrency on Windows.
if !errors.Is(err, os.ErrExist) && !errors.Is(err, os.ErrPermission) {
if !errors.Is(err, lockutil.ErrLockHeld) {
return nil, fmt.Errorf("oauth: acquire token lock: %w", err)
}
// Reclaim a stale lock left by a crashed holder — atomically (H3). A blind
// Remove lets two racers both reclaim + recreate and so both hold the lock;
// reclaimStaleLock renames the file aside (only one rename wins) and restores
// it if it turns out fresh, so a live lock is never deleted out from under it.
if info, statErr := os.Stat(lockPath); statErr == nil && time.Since(info.ModTime()) > fileLockStaleAfter {
cleared, rerr := lockutil.ReclaimStaleLock(lockPath, token, func(reclaimedPath string) bool {
info, err := os.Stat(reclaimedPath)
return err == nil && time.Since(info.ModTime()) <= fileLockStaleAfter
})
if rerr != nil {
// Reclaim hit a hard failure: the rename aside failed outright, or a
// live holder's lock could not be put back (the lock path may be
// missing, so re-acquiring would break mutual exclusion). Fail closed
// instead of spinning to the deadline.
return nil, fmt.Errorf("oauth: reclaim stale token lock: %w", rerr)
}
if cleared {
continue
}
// Lost the reclaim race (or it was actually fresh) — fall through to the
// bounded wait rather than hot-spinning on a reclaim that never wins.
}
if now().After(deadline) {
return nil, fmt.Errorf("oauth: timed out acquiring token lock %s", filepath.Base(lockPath))
}
+2 -22
View File
@@ -191,7 +191,6 @@ type TurnBenchResult struct {
type TurnBenchTotals struct {
InputTokens int64 `json:"inputTokens"`
CachedInputTokens int64 `json:"cachedInputTokens"`
CacheWriteTokens int64 `json:"cacheWriteTokens"`
OutputTokens int64 `json:"outputTokens"`
ModelRequests int64 `json:"modelRequests"`
ToolCalls int64 `json:"toolCalls"`
@@ -480,7 +479,6 @@ func topLatencySources(perSpan map[string]SpanStats, top int) []LatencySource {
func aggregateTotals(totals *TurnBenchTotals, tr *trace.TurnTrace) {
totals.InputTokens += tr.Counter(trace.CounterInputTokens)
totals.CachedInputTokens += tr.Counter(trace.CounterCachedInputTokens)
totals.CacheWriteTokens += tr.Counter(trace.CounterCacheWriteTokens)
totals.OutputTokens += tr.Counter(trace.CounterOutputTokens)
totals.ModelRequests += tr.Counter(trace.CounterModelRequests)
totals.ToolCalls += tr.Counter(trace.CounterToolCalls)
@@ -533,9 +531,8 @@ func FormatTurnBenchSummary(result TurnBenchResult) string {
lines = append(lines, fmt.Sprintf(" %-18s %10s %5.1f%%", src.Span, FormatMetric(src.TotalMs, "ms"), src.Share*100))
}
}
lines = append(lines, fmt.Sprintf("totals: in=%d (cache-read %d, cache-write %d, uncached %d) out=%d | requests=%d tools=%d retries=%d reconnects=%d compactions=%d",
result.Totals.InputTokens, result.Totals.CachedInputTokens, result.Totals.CacheWriteTokens,
uncachedInputTokens(result.Totals), result.Totals.OutputTokens,
lines = append(lines, fmt.Sprintf("totals: in=%d (cached %d) out=%d | requests=%d tools=%d retries=%d reconnects=%d compactions=%d",
result.Totals.InputTokens, result.Totals.CachedInputTokens, result.Totals.OutputTokens,
result.Totals.ModelRequests, result.Totals.ToolCalls, result.Totals.Retries,
result.Totals.Reconnects, result.Totals.Compactions))
for _, class := range sortedClasses(result.PerClass) {
@@ -548,23 +545,6 @@ func FormatTurnBenchSummary(result TurnBenchResult) string {
return strings.Join(lines, "\n")
}
func uncachedInputTokens(totals TurnBenchTotals) int64 {
remaining := totals.InputTokens
if remaining <= 0 {
return 0
}
for _, cached := range []int64{totals.CachedInputTokens, totals.CacheWriteTokens} {
if cached <= 0 {
continue
}
if cached >= remaining {
return 0
}
remaining -= cached
}
return remaining
}
// classTier returns the oracle tier label a class was classified into, so the
// per-class roll-up states explicitly which oracle (if any) its "passed" count
// is measured against.
-33
View File
@@ -74,7 +74,6 @@ func TestRunTurnBenchAggregation(t *testing.T) {
"t3": cannedTrace(150, 20, 1500),
"t4": cannedTrace(300, 10, 1000),
}
canned["t1"].Counters = append(canned["t1"].Counters, trace.Counter{Name: trace.CounterCacheWriteTokens, Value: 125})
cfg := TurnBenchConfig{
Model: "fake-model",
Iterations: 1,
@@ -156,9 +155,6 @@ func TestRunTurnBenchAggregation(t *testing.T) {
if result.Totals.InputTokens != 5500 {
t.Fatalf("inputTokens = %d, want 5500", result.Totals.InputTokens)
}
if result.Totals.CacheWriteTokens != 125 {
t.Fatalf("cacheWriteTokens = %d, want 125", result.Totals.CacheWriteTokens)
}
if result.Totals.OutputTokens != 2750 {
t.Fatalf("outputTokens = %d, want 2750", result.Totals.OutputTokens)
}
@@ -338,35 +334,6 @@ func TestFormatTurnBenchSummaryNamesTopSources(t *testing.T) {
}
}
func TestUncachedInputTokensClampsInconsistentTotals(t *testing.T) {
tests := []struct {
name string
totals TurnBenchTotals
want int64
}{
{name: "normal", totals: TurnBenchTotals{InputTokens: 100, CachedInputTokens: 30, CacheWriteTokens: 20}, want: 50},
{name: "cached exceeds input", totals: TurnBenchTotals{InputTokens: 100, CachedInputTokens: 120}, want: 0},
{name: "write exceeds remainder", totals: TurnBenchTotals{InputTokens: 100, CachedInputTokens: 20, CacheWriteTokens: 90}, want: 0},
{name: "overflowing inconsistent counters", totals: TurnBenchTotals{CachedInputTokens: 1<<63 - 1, CacheWriteTokens: 1<<63 - 1}, want: 0},
{name: "overflowing counters with input", totals: TurnBenchTotals{InputTokens: 100, CachedInputTokens: 1<<63 - 1, CacheWriteTokens: 1<<63 - 1}, want: 0},
{name: "negative cache counters ignored", totals: TurnBenchTotals{InputTokens: 100, CachedInputTokens: -20, CacheWriteTokens: -10}, want: 100},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
if got := uncachedInputTokens(test.totals); got != test.want {
t.Fatalf("uncachedInputTokens(%+v) = %d, want %d", test.totals, got, test.want)
}
})
}
summary := FormatTurnBenchSummary(TurnBenchResult{Totals: TurnBenchTotals{
CachedInputTokens: 1<<63 - 1,
CacheWriteTokens: 1<<63 - 1,
}})
if !strings.Contains(summary, "uncached 0") {
t.Fatalf("summary did not clamp inconsistent cache totals:\n%s", summary)
}
}
func TestLoadBaselineManifest(t *testing.T) {
path := filepath.Join("manifests", "baseline.json")
set, err := LoadTaskSet(path)
+60 -238
View File
@@ -62,7 +62,6 @@ const (
responsesEventReasoningDelta = "response.reasoning_summary_text.delta"
responsesEventOutputTextDone = "response.output_text.done"
responsesEventFunctionArgsDelta = "response.function_call_arguments.delta"
responsesEventCustomInputDelta = "response.custom_tool_call_input.delta"
responsesEventContentPartDone = "response.content_part.done"
responsesEventOutputItemDone = "response.output_item.done"
responsesEventCompleted = "response.completed"
@@ -71,37 +70,17 @@ const (
responsesEventIncomplete = "response.incomplete"
)
const applyPatchLarkGrammar = `start: begin_patch hunk+ end_patch
begin_patch: "*** Begin Patch" LF
end_patch: "*** End Patch" LF?
hunk: add_hunk | delete_hunk | update_hunk
add_hunk: "*** Add File: " filename LF add_line+
delete_hunk: "*** Delete File: " filename LF
update_hunk: "*** Update File: " filename LF change_move? change?
filename: /(.+)/
add_line: "+" /(.*)/ LF -> line
change_move: "*** Move to: " filename LF
change: (change_context | change_line)+ eof_line?
change_context: ("@@" | "@@ " /(.+)/) LF
change_line: ("+" | "-" | " ") /(.*)/ LF
eof_line: "*** End of File" LF
%import common.LF`
// responsesRequest is the wire shape POSTed to {baseURL}/responses.
type responsesRequest struct {
Type string `json:"type,omitempty"`
Model string `json:"model"`
Instructions string `json:"instructions"`
PreviousResponseID string `json:"previous_response_id,omitempty"`
Input []inputItem `json:"input"`
Stream bool `json:"stream"`
Store bool `json:"store"`
ParallelToolCalls bool `json:"parallel_tool_calls"`
MaxOutputTokens int `json:"max_output_tokens,omitempty"`
Tools []responsesTool `json:"tools,omitempty"`
Reasoning *responsesReasoning `json:"reasoning,omitempty"`
ServiceTier string `json:"service_tier,omitempty"`
PromptCacheKey string `json:"prompt_cache_key,omitempty"`
Model string `json:"model"`
Instructions string `json:"instructions"`
Input []inputItem `json:"input"`
Stream bool `json:"stream"`
Store bool `json:"store"`
MaxOutputTokens int `json:"max_output_tokens,omitempty"`
Tools []responsesTool `json:"tools,omitempty"`
Reasoning *responsesReasoning `json:"reasoning,omitempty"`
ServiceTier string `json:"service_tier,omitempty"`
}
// responsesReasoning carries the reasoning controls for the Responses API. The
@@ -126,7 +105,6 @@ type inputItem struct {
CallID string `json:"call_id,omitempty"`
Name string `json:"name,omitempty"`
Arguments string `json:"arguments,omitempty"`
Input string `json:"input,omitempty"`
// function_call_output fields
Output any `json:"output,omitempty"`
}
@@ -140,11 +118,10 @@ type contentItem struct {
}
type responsesTool struct {
Type string `json:"type"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
Parameters map[string]any `json:"parameters,omitempty"`
Format *zeroruntime.ToolDefinitionFormat `json:"format,omitempty"`
Type string `json:"type"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
Parameters map[string]any `json:"parameters,omitempty"`
}
// responsesEvent is the decoded form of one Responses SSE data payload.
@@ -155,7 +132,6 @@ type responsesEvent struct {
Delta string `json:"delta,omitempty"`
// item payloads (response.output_item.added / done)
ItemID string `json:"item_id,omitempty"`
CallID string `json:"call_id,omitempty"`
// OutputIndex is a *int so a real 0 (the first output) is distinguishable from
// "absent" — a plain int defaulting to 0 dropped a no-item_id call's args (M1).
OutputIndex *int `json:"output_index,omitempty"`
@@ -174,7 +150,6 @@ type itemPayload struct {
CallID string `json:"call_id,omitempty"`
Name string `json:"name,omitempty"`
Arguments string `json:"arguments,omitempty"`
Input string `json:"input,omitempty"`
Status string `json:"status,omitempty"`
}
@@ -198,8 +173,7 @@ type outputDetails struct {
}
type inputDetails struct {
CachedTokens int `json:"cached_tokens"`
CacheWriteTokens int `json:"cache_write_tokens"`
CachedTokens int `json:"cached_tokens"`
}
type errorPayload struct {
@@ -213,10 +187,8 @@ type errorPayload struct {
// response.output_item.done.
type toolCallBuilder struct {
ID string
CallID string
Name string
Arguments strings.Builder
Freeform bool
started bool
ended bool
}
@@ -225,57 +197,15 @@ type toolCallBuilder struct {
// event has been emitted. It is the Responses-API equivalent of the
// chat-completions toolState in provider.go.
type responsesState struct {
toolCalls map[string]*toolCallBuilder
toolOrder []string
text strings.Builder
responseID string
usage *usagePayload
emitted bool
done bool
toolCalls map[string]*toolCallBuilder
usage *usagePayload
done bool
}
func newResponsesState() *responsesState {
return &responsesState{toolCalls: map[string]*toolCallBuilder{}}
}
func (state *responsesState) outputInputItems() []inputItem {
items := make([]inputItem, 0, len(state.toolOrder)+1)
if state.text.Len() > 0 {
items = append(items, inputItem{
Type: "message",
Role: "assistant",
Content: []contentItem{{Type: "output_text", Text: state.text.String()}},
})
}
for _, key := range state.toolOrder {
builder := state.toolCalls[key]
if builder == nil || builder.ID == "" || builder.Name == "" {
continue
}
callID := builder.CallID
if callID == "" {
callID = builder.ID
}
itemType := "function_call"
if builder.Freeform {
itemType = "custom_tool_call"
}
item := inputItem{
Type: itemType,
ID: builder.ID,
CallID: callID,
Name: builder.Name,
}
if builder.Freeform {
item.Input = builder.Arguments.String()
} else {
item.Arguments = builder.Arguments.String()
}
items = append(items, item)
}
return items
}
// buildResponsesRequest converts a runtime CompletionRequest into the
// Responses API wire format. System messages fold into a top-level
// `instructions` field; everything else (user/assistant turns, tool
@@ -285,14 +215,11 @@ func (p *CodexProvider) buildResponsesRequest(request zeroruntime.CompletionRequ
return nil, errors.New("codex provider: model is required")
}
req := &responsesRequest{
Model: p.inner.model,
Stream: true,
ParallelToolCalls: true,
MaxOutputTokens: p.inner.maxTokens,
PromptCacheKey: request.PromptCacheKey,
Model: p.inner.model,
Stream: true,
MaxOutputTokens: p.inner.maxTokens,
}
instructions := []string{}
customCallIDs := map[string]string{}
for _, msg := range request.Messages {
switch msg.Role {
case zeroruntime.MessageRoleSystem:
@@ -307,31 +234,13 @@ func (p *CodexProvider) buildResponsesRequest(request zeroruntime.CompletionRequ
return nil, err
}
req.Input = append(req.Input, items...)
for _, call := range msg.ToolCalls {
if call.Freeform && call.ID != "" {
callID := call.ProviderCallID
if callID == "" {
callID = call.ID
}
customCallIDs[call.ID] = callID
}
}
case zeroruntime.MessageRoleTool:
if msg.ToolCallID == "" {
return nil, fmt.Errorf("codex provider: tool message missing tool call id")
}
outputType := "function_call_output"
callID := msg.ToolCallProviderID
if callID == "" {
callID = msg.ToolCallID
}
if customID, custom := customCallIDs[msg.ToolCallID]; custom {
outputType = "custom_tool_call_output"
callID = customID
}
req.Input = append(req.Input, inputItem{
Type: outputType,
CallID: callID,
Type: "function_call_output",
CallID: msg.ToolCallID,
Output: msg.Content,
})
default:
@@ -340,25 +249,12 @@ func (p *CodexProvider) buildResponsesRequest(request zeroruntime.CompletionRequ
}
req.Instructions = strings.Join(instructions, "\n\n")
for _, tool := range request.Tools {
definition := tool
if definition.Name == "apply_patch" && definition.Type == zeroruntime.ToolDefinitionFreeform {
definition.Description = "The apply_patch tool edits files from a raw structured patch. Use absolute file paths in patch headers to target a granted extra write root. Do not wrap the patch in JSON."
definition.Format = &zeroruntime.ToolDefinitionFormat{Type: "grammar", Syntax: "lark", Definition: applyPatchLarkGrammar}
}
wireTool := responsesTool{
Type: string(definition.Type),
Name: definition.Name,
Description: definition.Description,
Parameters: definition.Parameters,
Format: definition.Format,
}
if wireTool.Type == "" {
wireTool.Type = string(zeroruntime.ToolDefinitionFunction)
}
if wireTool.Type == string(zeroruntime.ToolDefinitionFreeform) {
wireTool.Parameters = nil
}
req.Tools = append(req.Tools, wireTool)
req.Tools = append(req.Tools, responsesTool{
Type: "function",
Name: tool.Name,
Description: tool.Description,
Parameters: tool.Parameters,
})
}
// Codex / o-series reasoning models take the effort tier nested under
// `reasoning` on the Responses API (the chat-completions `reasoning_effort`
@@ -416,26 +312,13 @@ func (p *CodexProvider) assistantInputItems(msg zeroruntime.Message) ([]inputIte
if tc.Name == "" {
return nil, errors.New("codex provider: assistant tool call missing name")
}
callID := tc.ProviderCallID
if callID == "" {
callID = tc.ID
}
itemType := "function_call"
if tc.Freeform {
itemType = "custom_tool_call"
}
item := inputItem{
Type: itemType,
ID: tc.ID,
CallID: callID,
Name: tc.Name,
}
if tc.Freeform {
item.Input = tc.Arguments
} else {
item.Arguments = tc.Arguments
}
items = append(items, item)
items = append(items, inputItem{
Type: "function_call",
ID: tc.ID,
CallID: tc.ID,
Name: tc.Name,
Arguments: tc.Arguments,
})
}
return items, nil
}
@@ -596,18 +479,6 @@ func (p *CodexProvider) emitResponsesEvent(
state.done = true
return false
}
return p.emitParsedResponsesEvent(ctx, &event, state, events)
}
// emitParsedResponsesEvent converts one already-decoded Responses event into
// zero or more runtime events. WebSocket sessions use this path to avoid
// decoding each frame twice; the SSE path retains emitResponsesEvent above.
func (p *CodexProvider) emitParsedResponsesEvent(
ctx context.Context,
event *responsesEvent,
state *responsesState,
events chan<- zeroruntime.StreamEvent,
) bool {
if event.Type == "" {
// A payload that parses as JSON but carries no `type` field is a
// malformed Responses event — the type is the discriminator that
@@ -629,17 +500,12 @@ func (p *CodexProvider) emitParsedResponsesEvent(
// the runtime. response.incomplete is folded into a length finish at
// scan-end (state.done stays false so the wrapper emits the
// StreamEventDone with FinishReasonLength).
if event.Response != nil && event.Response.ID != "" {
state.responseID = event.Response.ID
}
return true
case responsesEventOutputItemAdded:
p.handleOutputItemAdded(ctx, event, state, events)
p.handleOutputItemAdded(ctx, &event, state, events)
return true
case responsesEventOutputTextDelta:
if event.Delta != "" {
state.emitted = true
state.text.WriteString(event.Delta)
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventText,
Content: event.Delta,
@@ -651,7 +517,6 @@ func (p *CodexProvider) emitParsedResponsesEvent(
// phase shows progress (and keeps the activity clock fresh) instead of
// looking like a hang. Requested via reasoning.summary="auto".
if event.Delta != "" {
state.emitted = true
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventReasoning,
Content: event.Delta,
@@ -659,16 +524,13 @@ func (p *CodexProvider) emitParsedResponsesEvent(
}
return true
case responsesEventFunctionArgsDelta:
p.handleToolInputDelta(ctx, event, state, events, false)
return true
case responsesEventCustomInputDelta:
p.handleToolInputDelta(ctx, event, state, events, true)
p.handleFunctionArgsDelta(ctx, &event, state, events)
return true
case responsesEventOutputItemDone:
p.handleOutputItemDone(ctx, event, state, events)
p.handleOutputItemDone(ctx, &event, state, events)
return true
case responsesEventCompleted, responsesEventFailed:
return p.handleTerminalResponse(ctx, event, state, events)
return p.handleTerminalResponse(ctx, &event, state, events)
case responsesEventError:
message := event.Message
if event.Code != "" {
@@ -708,38 +570,22 @@ func (p *CodexProvider) handleOutputItemAdded(
return
}
switch event.Item.Type {
case "function_call", "custom_tool_call":
case "function_call":
key := p.toolCallKey(event)
if key == "" {
return
}
builder, exists := state.toolCalls[key]
if !exists {
builder = &toolCallBuilder{ID: key, Freeform: event.Item.Type == "custom_tool_call"}
state.toolCalls[key] = builder
state.toolOrder = append(state.toolOrder, key)
}
if event.Item.CallID != "" {
builder.CallID = event.Item.CallID
}
if event.Item.Name != "" {
builder.Name = event.Item.Name
}
builder.Freeform = builder.Freeform || event.Item.Type == "custom_tool_call"
if builder.Freeform && event.Item.Input != "" && builder.Arguments.Len() == 0 {
builder.Arguments.WriteString(event.Item.Input)
}
state.emitted = true
if !builder.started {
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventToolCallStart,
ToolCallID: key,
ToolCallProviderID: builder.CallID,
ToolName: builder.Name,
ToolCallFreeform: builder.Freeform,
})
builder.started = true
builder := &toolCallBuilder{
ID: key,
Name: event.Item.Name,
}
state.toolCalls[key] = builder
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventToolCallStart,
ToolCallID: key,
ToolName: event.Item.Name,
})
builder.started = true
}
}
@@ -747,12 +593,11 @@ func (p *CodexProvider) handleOutputItemAdded(
// to the in-flight builder and emits StreamEventToolCallDelta. The
// Codex backend uses either `item_id` or `output_index` to attribute
// the delta; we honor both.
func (p *CodexProvider) handleToolInputDelta(
func (p *CodexProvider) handleFunctionArgsDelta(
ctx context.Context,
event *responsesEvent,
state *responsesState,
events chan<- zeroruntime.StreamEvent,
freeform bool,
) {
key := p.toolCallKey(event)
if key == "" {
@@ -765,16 +610,10 @@ func (p *CodexProvider) handleToolInputDelta(
// an intermediate proxy). Treat the delta as the start of a new
// tool call with no known name yet — the eventual
// response.output_item.done carries the final name.
builder = &toolCallBuilder{ID: key, Freeform: freeform}
builder = &toolCallBuilder{ID: key}
state.toolCalls[key] = builder
state.toolOrder = append(state.toolOrder, key)
}
builder.Freeform = builder.Freeform || freeform
if event.CallID != "" && builder.CallID == "" {
builder.CallID = event.CallID
}
builder.Arguments.WriteString(event.Delta)
state.emitted = true
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventToolCallDelta,
ToolCallID: key,
@@ -793,7 +632,7 @@ func (p *CodexProvider) handleOutputItemDone(
state *responsesState,
events chan<- zeroruntime.StreamEvent,
) {
if event.Item == nil || (event.Item.Type != "function_call" && event.Item.Type != "custom_tool_call") {
if event.Item == nil || event.Item.Type != "function_call" {
return
}
key := p.toolCallKey(event)
@@ -802,31 +641,20 @@ func (p *CodexProvider) handleOutputItemDone(
}
builder, ok := state.toolCalls[key]
if !ok {
builder = &toolCallBuilder{ID: key, Freeform: event.Item.Type == "custom_tool_call"}
builder = &toolCallBuilder{ID: key}
state.toolCalls[key] = builder
state.toolOrder = append(state.toolOrder, key)
}
if event.Item.Name != "" {
builder.Name = event.Item.Name
}
if event.Item.CallID != "" {
builder.CallID = event.Item.CallID
}
builder.Freeform = builder.Freeform || event.Item.Type == "custom_tool_call"
finalInput := event.Item.Arguments
if builder.Freeform {
finalInput = event.Item.Input
}
if finalInput != "" && builder.Arguments.Len() == 0 {
builder.Arguments.WriteString(finalInput)
if event.Item.Arguments != "" && builder.Arguments.Len() == 0 {
builder.Arguments.WriteString(event.Item.Arguments)
}
if !builder.started {
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventToolCallStart,
ToolCallID: key,
ToolCallProviderID: builder.CallID,
ToolName: builder.Name,
ToolCallFreeform: builder.Freeform,
Type: zeroruntime.StreamEventToolCallStart,
ToolCallID: key,
ToolName: builder.Name,
})
builder.started = true
}
@@ -836,7 +664,6 @@ func (p *CodexProvider) handleOutputItemDone(
ToolName: builder.Name,
})
builder.ended = true
state.emitted = true
}
// handleTerminalResponse emits the final usage + done event for a
@@ -866,7 +693,6 @@ func (p *CodexProvider) handleTerminalResponse(
return false
}
state.usage = event.Response.Usage
state.responseID = event.Response.ID
if event.Response.Usage != nil {
usage := zeroruntime.Usage{
InputTokens: event.Response.Usage.InputTokens,
@@ -874,7 +700,6 @@ func (p *CodexProvider) handleTerminalResponse(
}
if event.Response.Usage.InputTokensDetails != nil {
usage.CachedInputTokens = event.Response.Usage.InputTokensDetails.CachedTokens
usage.CacheWriteTokens = event.Response.Usage.InputTokensDetails.CacheWriteTokens
}
if event.Response.Usage.OutputTokensDetails != nil {
usage.ReasoningTokens = event.Response.Usage.OutputTokensDetails.ReasoningTokens
@@ -910,7 +735,7 @@ func (p *CodexProvider) handleTerminalResponse(
state.done = true
return false
}
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone, ResponseID: event.Response.ID})
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone})
state.done = true
return false
}
@@ -929,9 +754,6 @@ func (p *CodexProvider) toolCallKey(event *responsesEvent) string {
if event.Item != nil && event.Item.CallID != "" {
return event.Item.CallID
}
if event.CallID != "" {
return event.CallID
}
if event.OutputIndex != nil {
// output_index is present (a *int distinguishes a real 0 — the first output
// in the response — from "absent"), so key on it even when 0 instead of
-421
View File
@@ -1,421 +0,0 @@
package openai
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"net/url"
"reflect"
"strings"
"sync"
"time"
"github.com/coder/websocket"
"github.com/Gitlawb/zero/internal/providers/providerio"
"github.com/Gitlawb/zero/internal/trace"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
const (
responsesWebSocketBetaHeader = "OpenAI-Beta"
responsesWebSocketBetaValue = "responses_websockets=2026-02-06"
responsesWebSocketRequest = "response.create"
previousResponseNotFoundCode = "previous_response_not_found"
webSocketConnectionLimitCode = "websocket_connection_limit_reached"
responsesWebSocketReadLimit = 16 << 20
)
var errCodexTurnSessionBusy = errors.New("codex turn session already has an active stream")
// NewCodexTurnSessionProvider enables one Responses WebSocket connection for
// each agent run. The session falls back to the existing HTTP/SSE provider when
// the endpoint cannot upgrade or loses its response chain.
func NewCodexTurnSessionProvider(provider *CodexProvider, caps zeroruntime.ProviderCapabilities) zeroruntime.TurnSessionProvider {
return &codexTurnSessionProvider{provider: provider, caps: caps}
}
type codexTurnSessionProvider struct {
provider *CodexProvider
caps zeroruntime.ProviderCapabilities
}
func (provider *codexTurnSessionProvider) OpenTurnSession(ctx context.Context) (zeroruntime.TurnSession, error) {
if provider == nil || provider.provider == nil || provider.provider.inner == nil {
return nil, errors.New("codex turn session requires a provider")
}
sessionCtx, cancel := context.WithCancel(ctx)
return &codexTurnSession{
ctx: sessionCtx,
cancel: cancel,
provider: provider.provider,
prewarmDone: make(chan struct{}),
}, nil
}
func (provider *codexTurnSessionProvider) Capabilities() zeroruntime.ProviderCapabilities {
return provider.caps
}
type codexTurnSession struct {
ctx context.Context
cancel context.CancelFunc
provider *CodexProvider
mu sync.Mutex
connection *websocket.Conn
webSocketOff bool
lastRequest *responsesRequest
lastResponseID string
lastOutput []inputItem
streamActive bool
prewarmOnce sync.Once
prewarmDone chan struct{}
closeOnce sync.Once
}
// Prewarm starts the WebSocket handshake without blocking prompt assembly. The
// first Stream waits for this bounded attempt; failure simply selects HTTP/SSE.
func (session *codexTurnSession) Prewarm(ctx context.Context) error {
session.prewarmOnce.Do(func() {
go func() {
defer close(session.prewarmDone)
prewarmCtx, cancel := context.WithTimeout(session.ctx, prewarmTimeout)
defer cancel()
recorder := trace.FromContext(ctx)
span := recorder.Span(trace.SpanProviderPrewarm)
defer span.End()
recorder.Counter(trace.CounterPrewarmAttempts, 1)
connection, err := session.provider.dialResponsesWebSocket(prewarmCtx)
session.mu.Lock()
defer session.mu.Unlock()
if err != nil {
session.webSocketOff = true
return
}
session.connection = connection
}()
})
return nil
}
func (session *codexTurnSession) Stream(ctx context.Context, request zeroruntime.CompletionRequest) (<-chan zeroruntime.StreamEvent, error) {
_ = session.Prewarm(ctx)
select {
case <-session.prewarmDone:
case <-ctx.Done():
return nil, ctx.Err()
}
fullRequest, err := session.provider.buildResponsesRequest(request)
if err != nil {
return nil, err
}
session.mu.Lock()
if session.streamActive {
session.mu.Unlock()
return nil, errCodexTurnSessionBusy
}
if session.webSocketOff || session.connection == nil {
session.mu.Unlock()
trace.FromContext(ctx).Counter(trace.CounterResponsesHTTPFallback, 1)
return session.provider.StreamCompletion(ctx, request)
}
session.streamActive = true
connection := session.connection
wireRequest := *fullRequest
wireRequest.Type = responsesWebSocketRequest
if delta, ok := responseInputDelta(session.lastRequest, session.lastOutput, fullRequest); ok && session.lastResponseID != "" {
wireRequest.PreviousResponseID = session.lastResponseID
wireRequest.Input = delta
trace.FromContext(ctx).Counter(trace.CounterResponseChainReused, 1)
} else if session.lastRequest != nil {
trace.FromContext(ctx).Counter(trace.CounterResponseChainReset, 1)
session.lastRequest = nil
session.lastResponseID = ""
session.lastOutput = nil
}
session.mu.Unlock()
body, err := json.Marshal(wireRequest)
if err != nil {
session.releaseStream()
return nil, fmt.Errorf("encode codex websocket request: %w", err)
}
events := make(chan zeroruntime.StreamEvent, 16)
go func() {
defer close(events)
defer session.releaseStream()
session.streamWebSocket(ctx, connection, body, fullRequest, request, events)
}()
return events, nil
}
func (session *codexTurnSession) releaseStream() {
session.mu.Lock()
session.streamActive = false
session.mu.Unlock()
}
func (session *codexTurnSession) Compact(context.Context, zeroruntime.CompletionRequest) ([]zeroruntime.Message, error) {
return nil, zeroruntime.ErrCompactionUnsupported
}
func (session *codexTurnSession) Close() error {
session.closeOnce.Do(func() {
session.cancel()
session.prewarmOnce.Do(func() { close(session.prewarmDone) })
select {
case <-session.prewarmDone:
case <-time.After(prewarmTimeout):
}
session.mu.Lock()
connection := session.connection
session.connection = nil
session.webSocketOff = true
session.mu.Unlock()
if connection != nil {
_ = connection.CloseNow()
}
})
return nil
}
func responseInputDelta(previous *responsesRequest, output []inputItem, current *responsesRequest) ([]inputItem, bool) {
if previous == nil || current == nil || !responsesRequestPropertiesEqual(previous, current) {
return nil, false
}
prefix := make([]inputItem, 0, len(previous.Input)+len(output))
prefix = append(prefix, previous.Input...)
prefix = append(prefix, output...)
if len(current.Input) <= len(prefix) || !reflect.DeepEqual(current.Input[:len(prefix)], prefix) {
return nil, false
}
return append([]inputItem(nil), current.Input[len(prefix):]...), true
}
func responsesRequestPropertiesEqual(left, right *responsesRequest) bool {
return left.Model == right.Model &&
left.Instructions == right.Instructions &&
left.Stream == right.Stream &&
left.Store == right.Store &&
left.ParallelToolCalls == right.ParallelToolCalls &&
left.MaxOutputTokens == right.MaxOutputTokens &&
reflect.DeepEqual(left.Tools, right.Tools) &&
reflect.DeepEqual(left.Reasoning, right.Reasoning) &&
left.ServiceTier == right.ServiceTier &&
left.PromptCacheKey == right.PromptCacheKey
}
func (session *codexTurnSession) streamWebSocket(
ctx context.Context,
connection *websocket.Conn,
body []byte,
fullRequest *responsesRequest,
runtimeRequest zeroruntime.CompletionRequest,
events chan<- zeroruntime.StreamEvent,
) {
if err := connection.Write(ctx, websocket.MessageText, body); err != nil {
session.disableWebSocket(connection)
session.forwardHTTP(ctx, runtimeRequest, events)
return
}
state := newResponsesState()
for {
readCtx := ctx
cancel := func() {}
idleTimeout := session.provider.inner.streamIdleTimeout
if idleTimeout > 0 {
readCtx, cancel = context.WithTimeout(ctx, idleTimeout)
}
messageType, data, err := connection.Read(readCtx)
readTimedOut := errors.Is(readCtx.Err(), context.DeadlineExceeded) && ctx.Err() == nil
cancel()
if err != nil {
session.disableWebSocket(connection)
if ctx.Err() != nil {
return
}
if !state.emitted {
session.forwardHTTP(ctx, runtimeRequest, events)
return
}
errorMessage := err.Error()
if readTimedOut {
errorMessage = providerio.StreamTimeoutMessage(providerio.ErrStreamIdle, idleTimeout)
}
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventError,
Error: session.provider.redact("provider stream error: " + errorMessage),
})
return
}
if messageType != websocket.MessageText {
continue
}
var responseEvent responsesEvent
decodeErr := json.Unmarshal(data, &responseEvent)
if decodeErr == nil {
if (responseEvent.Code == previousResponseNotFoundCode || responseEvent.Code == webSocketConnectionLimitCode) &&
!state.emitted {
session.disableWebSocket(connection)
session.forwardHTTP(ctx, runtimeRequest, events)
return
}
if responseEvent.Type == responsesEventIncomplete {
session.resetResponseChain(connection)
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventDone,
FinishReason: zeroruntime.FinishReasonLength,
})
return
}
}
keepReading := false
if decodeErr == nil {
keepReading = session.provider.emitParsedResponsesEvent(ctx, &responseEvent, state, events)
} else {
keepReading = session.provider.emitResponsesEvent(ctx, string(data), state, events)
}
if keepReading {
continue
}
if responseEvent.Type == responsesEventCompleted && responseEvent.Response != nil &&
responseEvent.Response.Error == nil && responseEvent.Response.Status != "failed" && state.responseID != "" {
session.mu.Lock()
if session.connection == connection && !session.webSocketOff {
requestCopy := *fullRequest
requestCopy.Input = append([]inputItem(nil), fullRequest.Input...)
session.lastRequest = &requestCopy
session.lastResponseID = state.responseID
session.lastOutput = state.outputInputItems()
}
session.mu.Unlock()
}
return
}
}
func (session *codexTurnSession) resetResponseChain(connection *websocket.Conn) {
session.mu.Lock()
defer session.mu.Unlock()
if session.connection != connection {
return
}
session.lastRequest = nil
session.lastResponseID = ""
session.lastOutput = nil
}
func (session *codexTurnSession) forwardHTTP(ctx context.Context, request zeroruntime.CompletionRequest, events chan<- zeroruntime.StreamEvent) {
trace.FromContext(ctx).Counter(trace.CounterResponsesHTTPFallback, 1)
stream, err := session.provider.StreamCompletion(ctx, request)
if err != nil {
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventError,
Error: session.provider.redact(err.Error()),
})
return
}
for event := range stream {
providerio.SendEvent(ctx, events, event)
}
}
func (session *codexTurnSession) disableWebSocket(connection *websocket.Conn) {
session.mu.Lock()
if session.connection == connection {
session.connection = nil
session.webSocketOff = true
session.lastRequest = nil
session.lastResponseID = ""
session.lastOutput = nil
}
session.mu.Unlock()
_ = connection.CloseNow()
}
func (provider *CodexProvider) dialResponsesWebSocket(ctx context.Context) (*websocket.Conn, error) {
endpoint, err := responsesWebSocketURL(provider.inner.endpoint)
if err != nil {
return nil, err
}
for authTry := 0; authTry < 2; authTry++ {
headers, err := provider.responsesWebSocketHeaders(ctx, authTry == 1)
if err != nil {
return nil, err
}
connection, response, err := websocket.Dial(ctx, endpoint, &websocket.DialOptions{
HTTPClient: provider.inner.httpClient,
HTTPHeader: headers,
CompressionMode: websocket.CompressionNoContextTakeover,
})
if err == nil {
connection.SetReadLimit(responsesWebSocketReadLimit)
return connection, nil
}
if response == nil || response.StatusCode != http.StatusUnauthorized || provider.inner.oauthResolver == nil || authTry == 1 {
return nil, err
}
}
return nil, errors.New("codex websocket authentication failed")
}
func (provider *CodexProvider) responsesWebSocketHeaders(ctx context.Context, forceRefresh bool) (http.Header, error) {
request, err := http.NewRequestWithContext(ctx, http.MethodGet, provider.inner.endpoint, nil)
if err != nil {
return nil, err
}
base := providerio.AuthHeaders{
APIKey: provider.inner.apiKey,
DefaultAuthHeader: "Authorization",
DefaultAuthScheme: "Bearer",
AuthHeader: provider.inner.authHeader,
AuthScheme: provider.inner.authScheme,
AuthHeaderValue: provider.inner.authHeaderValue,
CustomHeaders: provider.inner.customHeaders,
}
if resolver := provider.inner.oauthResolver; resolver != nil {
header, value, ok, err := resolver(ctx, forceRefresh)
if err != nil {
return nil, err
}
if ok {
providerio.ApplyAuthHeaders(request, providerio.AuthHeaders{CustomHeaders: provider.inner.customHeaders})
if strings.TrimSpace(header) == "" {
header = "Authorization"
}
request.Header.Set(header, value)
} else {
providerio.ApplyAuthHeaders(request, base)
}
} else {
providerio.ApplyAuthHeaders(request, base)
}
provider.injectCodexHeaders(request)
request.Header.Set(responsesWebSocketBetaHeader, responsesWebSocketBetaValue)
return request.Header.Clone(), nil
}
func responsesWebSocketURL(endpoint string) (string, error) {
parsed, err := url.Parse(strings.TrimSpace(endpoint))
if err != nil {
return "", err
}
switch parsed.Scheme {
case "http":
parsed.Scheme = "ws"
case "https":
parsed.Scheme = "wss"
default:
return "", fmt.Errorf("unsupported Responses websocket URL scheme %q", parsed.Scheme)
}
return parsed.String(), nil
}
@@ -1,590 +0,0 @@
package openai
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"time"
"github.com/coder/websocket"
"github.com/Gitlawb/zero/internal/trace"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
func TestCodexTurnSessionChainsOnlyAppendOnlyRequests(t *testing.T) {
var mu sync.Mutex
requests := []map[string]any{}
betaHeader := ""
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
if !strings.EqualFold(request.Header.Get("Upgrade"), "websocket") {
http.Error(writer, "websocket required", http.StatusUpgradeRequired)
return
}
mu.Lock()
betaHeader = request.Header.Get(responsesWebSocketBetaHeader)
mu.Unlock()
connection, err := websocket.Accept(writer, request, nil)
if err != nil {
return
}
defer connection.CloseNow()
for index := 0; index < 3; index++ {
_, body, err := connection.Read(request.Context())
if err != nil {
return
}
var payload map[string]any
if err := json.Unmarshal(body, &payload); err != nil {
return
}
mu.Lock()
requests = append(requests, payload)
mu.Unlock()
switch index {
case 0:
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.created","response":{"id":"resp-1","status":"in_progress"}}`,
`{"type":"response.output_item.added","item_id":"call-1","item":{"type":"function_call","id":"call-1","call_id":"call-1","name":"read_file"}}`,
`{"type":"response.function_call_arguments.delta","item_id":"call-1","delta":"{}"}`,
`{"type":"response.output_item.done","item_id":"call-1","item":{"type":"function_call","id":"call-1","call_id":"call-1","name":"read_file","arguments":"{}"}}`,
`{"type":"response.completed","response":{"id":"resp-1","status":"completed"}}`,
)
case 1:
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.output_text.delta","delta":"done"}`,
`{"type":"response.completed","response":{"id":"resp-2","status":"completed"}}`,
)
case 2:
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.completed","response":{"id":"resp-3","status":"completed"}}`,
)
}
}
}))
defer server.Close()
provider := newCodexSessionTestProvider(t, server)
// A disabled stream watchdog must use the parent context rather than creating
// an immediately-expired read deadline.
provider.inner.streamIdleTimeout = 0
session := openCodexSession(t, provider)
defer session.Close()
recorder := trace.NewRecorder("session-1", "run-1", "test")
recorder.Start()
streamContext := trace.WithContext(t.Context(), recorder)
base := zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{
{Role: zeroruntime.MessageRoleSystem, Content: "Use tools carefully."},
{Role: zeroruntime.MessageRoleUser, Content: "Read the file."},
},
Tools: []zeroruntime.ToolDefinition{{Name: "read_file", Parameters: map[string]any{"type": "object"}}},
PromptCacheKey: "session-1",
}
first := collectCodexSessionEvents(t, session, base)
if !hasToolCallEnd(first, "call-1") {
t.Fatalf("first stream did not produce the tool call: %#v", first)
}
second := base
second.Messages = append(append([]zeroruntime.Message(nil), base.Messages...),
zeroruntime.Message{Role: zeroruntime.MessageRoleAssistant, ToolCalls: []zeroruntime.ToolCall{{ID: "call-1", Name: "read_file", Arguments: "{}"}}},
zeroruntime.Message{Role: zeroruntime.MessageRoleTool, ToolCallID: "call-1", Content: "contents"},
)
secondEvents := collectCodexSessionEventsContext(t, streamContext, session, second)
if got := joinedText(secondEvents); got != "done" {
t.Fatalf("second stream text = %q, want done", got)
}
third := second
third.ReasoningEffort = "high"
third.Messages = append(append([]zeroruntime.Message(nil), second.Messages...),
zeroruntime.Message{Role: zeroruntime.MessageRoleAssistant, Content: "done"},
zeroruntime.Message{Role: zeroruntime.MessageRoleUser, Content: "Continue."},
)
collectCodexSessionEventsContext(t, streamContext, session, third)
finishedTrace := recorder.Finish()
if got := finishedTrace.Counter(trace.CounterResponseChainReused); got != 1 {
t.Fatalf("response-chain reuse count = %d, want 1", got)
}
if got := finishedTrace.Counter(trace.CounterResponseChainReset); got != 1 {
t.Fatalf("response-chain reset count = %d, want 1", got)
}
mu.Lock()
defer mu.Unlock()
if len(requests) != 3 {
t.Fatalf("websocket requests = %d, want 3", len(requests))
}
if betaHeader != responsesWebSocketBetaValue {
t.Fatalf("OpenAI-Beta = %q, want %q", betaHeader, responsesWebSocketBetaValue)
}
if requests[0]["type"] != responsesWebSocketRequest || requests[0]["previous_response_id"] != nil {
t.Fatalf("first request chaining fields = %#v", requests[0])
}
if requests[0]["parallel_tool_calls"] != true || requests[0]["prompt_cache_key"] != "session-1" {
t.Fatalf("first request stable fields = %#v", requests[0])
}
if requests[1]["previous_response_id"] != "resp-1" {
t.Fatalf("second previous_response_id = %#v", requests[1]["previous_response_id"])
}
secondInput, _ := requests[1]["input"].([]any)
if len(secondInput) != 1 || secondInput[0].(map[string]any)["type"] != "function_call_output" {
t.Fatalf("second incremental input = %#v", requests[1]["input"])
}
if requests[2]["previous_response_id"] != nil {
t.Fatalf("incompatible request reused response chain: %#v", requests[2])
}
thirdInput, _ := requests[2]["input"].([]any)
if len(thirdInput) <= 1 {
t.Fatalf("incompatible request did not send full input: %#v", thirdInput)
}
}
func TestCodexTurnSessionFallsBackWhenPreviousResponseIsMissing(t *testing.T) {
var mu sync.Mutex
webSocketRequests := 0
httpRequests := 0
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
if strings.EqualFold(request.Header.Get("Upgrade"), "websocket") {
connection, err := websocket.Accept(writer, request, nil)
if err != nil {
return
}
defer connection.CloseNow()
for index := 0; index < 2; index++ {
if _, _, err := connection.Read(request.Context()); err != nil {
return
}
mu.Lock()
webSocketRequests++
mu.Unlock()
if index == 0 {
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.output_item.added","item_id":"call-1","item":{"type":"function_call","id":"call-1","call_id":"call-1","name":"read_file"}}`,
`{"type":"response.output_item.done","item_id":"call-1","item":{"type":"function_call","id":"call-1","call_id":"call-1","name":"read_file","arguments":"{}"}}`,
`{"type":"response.completed","response":{"id":"resp-1","status":"completed"}}`,
)
continue
}
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.error","code":"previous_response_not_found","message":"missing"}`,
)
}
return
}
mu.Lock()
httpRequests++
mu.Unlock()
writer.Header().Set("Content-Type", "text/event-stream")
_, _ = fmt.Fprint(writer, "data: {\"type\":\"response.output_text.delta\",\"delta\":\"fallback\"}\n\n")
_, _ = fmt.Fprint(writer, "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-http\",\"status\":\"completed\"}}\n\n")
}))
defer server.Close()
provider := newCodexSessionTestProvider(t, server)
session := openCodexSession(t, provider)
defer session.Close()
base := zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "Read it."}},
Tools: []zeroruntime.ToolDefinition{{Name: "read_file", Parameters: map[string]any{"type": "object"}}},
}
collectCodexSessionEvents(t, session, base)
second := base
second.Messages = append(append([]zeroruntime.Message(nil), base.Messages...),
zeroruntime.Message{Role: zeroruntime.MessageRoleAssistant, ToolCalls: []zeroruntime.ToolCall{{ID: "call-1", Name: "read_file", Arguments: "{}"}}},
zeroruntime.Message{Role: zeroruntime.MessageRoleTool, ToolCallID: "call-1", Content: "contents"},
)
events := collectCodexSessionEvents(t, session, second)
if got := joinedText(events); got != "fallback" {
t.Fatalf("fallback stream text = %q, want fallback; events=%#v", got, events)
}
for _, event := range events {
if event.Type == zeroruntime.StreamEventError {
t.Fatalf("response-chain recovery leaked an error: %#v", event)
}
}
mu.Lock()
defer mu.Unlock()
if webSocketRequests != 2 || httpRequests != 1 {
t.Fatalf("requests = websocket:%d http:%d, want 2/1", webSocketRequests, httpRequests)
}
}
func TestCodexTurnSessionFallsBackWhenSocketClosesBeforeOutput(t *testing.T) {
var mu sync.Mutex
httpRequests := 0
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
if strings.EqualFold(request.Header.Get("Upgrade"), "websocket") {
connection, err := websocket.Accept(writer, request, nil)
if err != nil {
return
}
if _, _, err := connection.Read(request.Context()); err == nil {
_ = connection.CloseNow()
}
return
}
mu.Lock()
httpRequests++
mu.Unlock()
writer.Header().Set("Content-Type", "text/event-stream")
_, _ = fmt.Fprint(writer, "data: {\"type\":\"response.output_text.delta\",\"delta\":\"recovered\"}\n\n")
_, _ = fmt.Fprint(writer, "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-http\",\"status\":\"completed\"}}\n\n")
}))
defer server.Close()
provider := newCodexSessionTestProvider(t, server)
session := openCodexSession(t, provider)
defer session.Close()
recorder := trace.NewRecorder("session-fallback", "run-fallback", "test")
recorder.Start()
events := collectCodexSessionEventsContext(t, trace.WithContext(t.Context(), recorder), session, zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "Hello"}},
})
if got := joinedText(events); got != "recovered" {
t.Fatalf("fallback text = %q, want recovered", got)
}
mu.Lock()
defer mu.Unlock()
if httpRequests != 1 {
t.Fatalf("HTTP fallback requests = %d, want 1", httpRequests)
}
if got := recorder.Finish().Counter(trace.CounterResponsesHTTPFallback); got != 1 {
t.Fatalf("responses HTTP fallback counter = %d, want 1", got)
}
}
func TestCodexTurnSessionKeepsWebSocketAfterIncompleteResponse(t *testing.T) {
var mu sync.Mutex
webSocketRequests := 0
httpRequests := 0
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
if strings.EqualFold(request.Header.Get("Upgrade"), "websocket") {
connection, err := websocket.Accept(writer, request, nil)
if err != nil {
return
}
defer connection.CloseNow()
for index := 0; index < 2; index++ {
if _, _, err := connection.Read(request.Context()); err != nil {
return
}
mu.Lock()
webSocketRequests++
mu.Unlock()
if index == 0 {
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.incomplete","response":{"id":"resp-1","status":"incomplete"}}`,
)
continue
}
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.output_text.delta","delta":"websocket"}`,
`{"type":"response.completed","response":{"id":"resp-2","status":"completed"}}`,
)
}
return
}
mu.Lock()
httpRequests++
mu.Unlock()
writer.Header().Set("Content-Type", "text/event-stream")
_, _ = fmt.Fprint(writer, "data: {\"type\":\"response.output_text.delta\",\"delta\":\"fallback\"}\n\n")
_, _ = fmt.Fprint(writer, "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-http\",\"status\":\"completed\"}}\n\n")
}))
defer server.Close()
provider := newCodexSessionTestProvider(t, server)
session := openCodexSession(t, provider)
defer session.Close()
request := zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "Continue."}},
}
first := collectCodexSessionEvents(t, session, request)
if len(first) != 1 || first[0].Type != zeroruntime.StreamEventDone || first[0].FinishReason != zeroruntime.FinishReasonLength {
t.Fatalf("incomplete response events = %#v, want length completion", first)
}
second := collectCodexSessionEvents(t, session, request)
if got := joinedText(second); got != "websocket" {
t.Fatalf("second turn text = %q, want websocket; events=%#v", got, second)
}
mu.Lock()
defer mu.Unlock()
if webSocketRequests != 2 || httpRequests != 0 {
t.Fatalf("requests = websocket:%d http:%d, want 2/0", webSocketRequests, httpRequests)
}
}
func TestCodexTurnSessionDoesNotReplayAfterVisibleOutput(t *testing.T) {
var mu sync.Mutex
httpRequests := 0
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
if strings.EqualFold(request.Header.Get("Upgrade"), "websocket") {
connection, err := websocket.Accept(writer, request, nil)
if err != nil {
return
}
if _, _, err := connection.Read(request.Context()); err == nil {
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.output_text.delta","delta":"visible"}`,
)
}
_ = connection.Close(websocket.StatusInternalError, "interrupted")
return
}
mu.Lock()
httpRequests++
mu.Unlock()
}))
defer server.Close()
provider := newCodexSessionTestProvider(t, server)
session := openCodexSession(t, provider)
defer session.Close()
stream, err := session.Stream(t.Context(), zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "Hello"}},
})
if err != nil {
t.Fatalf("Stream: %v", err)
}
sawText, sawError := false, false
for event := range stream {
sawText = sawText || event.Type == zeroruntime.StreamEventText
sawError = sawError || event.Type == zeroruntime.StreamEventError
}
if !sawText || !sawError {
t.Fatalf("visible interrupted stream = text:%v error:%v", sawText, sawError)
}
mu.Lock()
defer mu.Unlock()
if httpRequests != 0 {
t.Fatalf("visible interrupted stream was replayed over HTTP %d time(s)", httpRequests)
}
}
func TestCodexTurnSessionReportsWebSocketIdleTimeout(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
connection, err := websocket.Accept(writer, request, nil)
if err != nil {
return
}
defer connection.CloseNow()
if _, _, err := connection.Read(request.Context()); err != nil {
return
}
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.output_text.delta","delta":"visible"}`,
)
<-request.Context().Done()
}))
defer server.Close()
provider := newCodexSessionTestProvider(t, server)
provider.inner.streamIdleTimeout = 200 * time.Millisecond
session := openCodexSession(t, provider)
defer session.Close()
stream, err := session.Stream(t.Context(), zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "Hello"}},
})
if err != nil {
t.Fatalf("Stream: %v", err)
}
var streamError string
for event := range stream {
if event.Type == zeroruntime.StreamEventError {
streamError = event.Error
break
}
}
if !strings.Contains(streamError, "idle timeout after 200ms") {
t.Fatalf("stream error = %q, want idle-timeout detail", streamError)
}
}
func TestCodexTurnSessionRejectsOverlappingStreams(t *testing.T) {
started := make(chan struct{})
release := make(chan struct{})
var releaseOnce sync.Once
releaseStream := func() { releaseOnce.Do(func() { close(release) }) }
t.Cleanup(releaseStream)
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
connection, err := websocket.Accept(writer, request, nil)
if err != nil {
return
}
defer connection.CloseNow()
if _, _, err := connection.Read(request.Context()); err != nil {
return
}
close(started)
<-release
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.completed","response":{"id":"resp-1","status":"completed"}}`,
)
}))
defer server.Close()
provider := newCodexSessionTestProvider(t, server)
session := openCodexSession(t, provider)
defer session.Close()
request := zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "Hello"}},
}
first, err := session.Stream(t.Context(), request)
if err != nil {
t.Fatalf("first Stream: %v", err)
}
select {
case <-started:
case <-time.After(time.Second):
t.Fatal("first WebSocket stream did not start")
}
if _, err := session.Stream(t.Context(), request); !errors.Is(err, errCodexTurnSessionBusy) {
t.Fatalf("overlapping Stream error = %v, want %v", err, errCodexTurnSessionBusy)
}
releaseStream()
for range first {
}
}
func TestCodexTurnSessionHTTPFallbackRedactsSetupError(t *testing.T) {
const secret = "sk-secret-fallback"
provider, err := NewCodexProvider(CodexOptions{Options: Options{
APIKey: secret,
BaseURL: "https://chatgpt.example/backend-api/codex",
Model: "gpt-test",
}})
if err != nil {
t.Fatalf("NewCodexProvider: %v", err)
}
session := &codexTurnSession{provider: provider}
events := make(chan zeroruntime.StreamEvent, 1)
session.forwardHTTP(t.Context(), zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRole(secret), Content: "invalid"}},
}, events)
select {
case event := <-events:
if event.Type != zeroruntime.StreamEventError {
t.Fatalf("fallback event type = %s, want error", event.Type)
}
if strings.Contains(event.Error, secret) {
t.Fatalf("fallback error leaked credential: %q", event.Error)
}
case <-time.After(time.Second):
t.Fatal("fallback did not emit setup error")
}
}
func TestCodexTurnSessionCloseDoesNotStartPrewarm(t *testing.T) {
provider, err := NewCodexProvider(CodexOptions{Options: Options{
APIKey: "test-token",
BaseURL: "https://chatgpt.example/backend-api/codex",
Model: "gpt-test",
}})
if err != nil {
t.Fatalf("NewCodexProvider: %v", err)
}
session, err := NewCodexTurnSessionProvider(provider, zeroruntime.ProviderCapabilities{}).OpenTurnSession(t.Context())
if err != nil {
t.Fatalf("OpenTurnSession: %v", err)
}
done := make(chan error, 1)
go func() { done <- session.Close() }()
select {
case err := <-done:
if err != nil {
t.Fatalf("Close() error = %v", err)
}
case <-time.After(200 * time.Millisecond):
t.Fatal("Close blocked without a prewarm attempt")
}
}
func newCodexSessionTestProvider(t *testing.T, server *httptest.Server) *CodexProvider {
t.Helper()
provider, err := NewCodexProvider(CodexOptions{
Options: Options{
APIKey: "test-token",
BaseURL: server.URL,
Model: "gpt-test",
HTTPClient: server.Client(),
},
AccountID: "account-test",
})
if err != nil {
t.Fatalf("NewCodexProvider: %v", err)
}
return provider
}
func openCodexSession(t *testing.T, provider *CodexProvider) zeroruntime.TurnSession {
t.Helper()
session, err := NewCodexTurnSessionProvider(provider, zeroruntime.ProviderCapabilities{}).OpenTurnSession(t.Context())
if err != nil {
t.Fatalf("OpenTurnSession: %v", err)
}
if err := session.Prewarm(t.Context()); err != nil {
t.Fatalf("Prewarm: %v", err)
}
return session
}
func collectCodexSessionEvents(t *testing.T, session zeroruntime.TurnSession, request zeroruntime.CompletionRequest) []zeroruntime.StreamEvent {
t.Helper()
return collectCodexSessionEventsContext(t, t.Context(), session, request)
}
func collectCodexSessionEventsContext(t *testing.T, ctx context.Context, session zeroruntime.TurnSession, request zeroruntime.CompletionRequest) []zeroruntime.StreamEvent {
t.Helper()
stream, err := session.Stream(ctx, request)
if err != nil {
t.Fatalf("Stream: %v", err)
}
events := []zeroruntime.StreamEvent{}
for event := range stream {
events = append(events, event)
if event.Type == zeroruntime.StreamEventError {
t.Fatalf("stream error: %s", event.Error)
}
}
return events
}
func writeWebSocketEvents(ctx context.Context, connection *websocket.Conn, events ...string) {
for _, event := range events {
if err := connection.Write(ctx, websocket.MessageText, []byte(event)); err != nil {
return
}
}
}
func hasToolCallEnd(events []zeroruntime.StreamEvent, id string) bool {
for _, event := range events {
if event.Type == zeroruntime.StreamEventToolCallEnd && event.ToolCallID == id {
return true
}
}
return false
}
func joinedText(events []zeroruntime.StreamEvent) string {
var builder strings.Builder
for _, event := range events {
if event.Type == zeroruntime.StreamEventText {
builder.WriteString(event.Content)
}
}
return builder.String()
}
@@ -2,8 +2,6 @@ package openai
import (
"context"
"reflect"
"strings"
"testing"
"github.com/Gitlawb/zero/internal/zeroruntime"
@@ -27,126 +25,6 @@ func TestToolCallKeyOutputIndexZero(t *testing.T) {
}
}
func TestResponsesStatePreservesDistinctCallID(t *testing.T) {
provider := &CodexProvider{}
state := newResponsesState()
events := make(chan zeroruntime.StreamEvent, 4)
provider.handleOutputItemAdded(context.Background(), &responsesEvent{
ItemID: "item-1",
Item: &itemPayload{
Type: "function_call",
ID: "item-1",
CallID: "call-1",
Name: "read_file",
},
}, state, events)
items := state.outputInputItems()
if len(items) != 1 || items[0].ID != "item-1" || items[0].CallID != "call-1" {
t.Fatalf("captured function call = %#v, want item id item-1 and call id call-1", items)
}
}
func TestCodexCustomApplyPatchStreamProducesFreeformCall(t *testing.T) {
provider := &CodexProvider{}
state := newResponsesState()
events := make(chan zeroruntime.StreamEvent, 16)
patch := "*** Begin Patch\n*** Add File: hello.txt\n+hello\n*** End Patch\n"
for _, payload := range []string{
`{"type":"response.output_item.added","item_id":"item-1","item":{"type":"custom_tool_call","call_id":"call-1","name":"apply_patch"}}`,
`{"type":"response.custom_tool_call_input.delta","item_id":"item-1","call_id":"call-1","delta":"*** Begin Patch\n*** Add File: hello.txt\n"}`,
`{"type":"response.custom_tool_call_input.delta","item_id":"item-1","call_id":"call-1","delta":"+hello\n*** End Patch\n"}`,
`{"type":"response.output_item.done","item_id":"item-1","item":{"type":"custom_tool_call","call_id":"call-1","name":"apply_patch","input":"*** Begin Patch\n*** Add File: hello.txt\n+hello\n*** End Patch\n"}}`,
`{"type":"response.completed","response":{"id":"resp-1","status":"completed"}}`,
} {
if keepReading := provider.emitResponsesEvent(context.Background(), payload, state, events); !keepReading && !state.done {
t.Fatalf("stream stopped before terminal event for %s", payload)
}
}
close(events)
collected := zeroruntime.CollectStream(context.Background(), events)
if len(collected.ToolCalls) != 1 {
t.Fatalf("tool calls = %#v, want one", collected.ToolCalls)
}
call := collected.ToolCalls[0]
if call.ID != "item-1" || call.ProviderCallID != "call-1" || call.Name != "apply_patch" || !call.Freeform || call.Arguments != patch {
t.Fatalf("custom apply_patch call = %#v", call)
}
}
func TestCodexApplyPatchUsesFreeformWireContract(t *testing.T) {
provider, err := NewCodexProvider(CodexOptions{Options: Options{
APIKey: "test-token",
BaseURL: "https://chatgpt.example/backend-api/codex",
Model: "gpt-test",
}})
if err != nil {
t.Fatalf("NewCodexProvider: %v", err)
}
patch := "*** Begin Patch\n*** Add File: hello.txt\n+hello\n*** End Patch\n"
request, err := provider.buildResponsesRequest(zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{
{Role: zeroruntime.MessageRoleUser, Content: "Add the file."},
{Role: zeroruntime.MessageRoleAssistant, ToolCalls: []zeroruntime.ToolCall{{
ID: "item-1", ProviderCallID: "call-1", Name: "apply_patch", Arguments: patch, Freeform: true,
}}},
{Role: zeroruntime.MessageRoleTool, ToolCallID: "item-1", ToolCallProviderID: "call-1", Content: "Done!"},
},
Tools: []zeroruntime.ToolDefinition{{
Name: "apply_patch", Description: "Apply a patch.", Type: zeroruntime.ToolDefinitionFreeform,
}},
})
if err != nil {
t.Fatalf("buildResponsesRequest: %v", err)
}
if len(request.Tools) != 1 || request.Tools[0].Type != string(zeroruntime.ToolDefinitionFreeform) ||
request.Tools[0].Format == nil || request.Tools[0].Format.Syntax != "lark" || request.Tools[0].Parameters != nil {
t.Fatalf("apply_patch tool wire shape = %#v", request.Tools)
}
if !strings.Contains(request.Tools[0].Format.Definition, "*** Begin Patch") {
t.Fatalf("apply_patch grammar = %q", request.Tools[0].Format.Definition)
}
if !strings.Contains(request.Tools[0].Description, "absolute file paths") || !strings.Contains(request.Tools[0].Description, "granted extra write root") {
t.Fatalf("apply_patch description does not explain extra-root paths: %q", request.Tools[0].Description)
}
if len(request.Input) != 3 || request.Input[1].Type != "custom_tool_call" || request.Input[1].ID != "item-1" ||
request.Input[1].CallID != "call-1" || request.Input[1].Input != patch ||
request.Input[2].Type != "custom_tool_call_output" || request.Input[2].CallID != "call-1" {
t.Fatalf("custom apply_patch replay input = %#v", request.Input)
}
}
func TestCodexPreservesCustomApplyPatchFunction(t *testing.T) {
provider, err := NewCodexProvider(CodexOptions{Options: Options{
APIKey: "test-token",
BaseURL: "https://chatgpt.example/backend-api/codex",
Model: "gpt-test",
}})
if err != nil {
t.Fatalf("NewCodexProvider: %v", err)
}
parameters := map[string]any{
"type": "object",
"properties": map[string]any{"value": map[string]any{"type": "string"}},
}
request, err := provider.buildResponsesRequest(zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "Call the custom tool."}},
Tools: []zeroruntime.ToolDefinition{{
Name: "apply_patch", Description: "Custom integration tool.", Parameters: parameters,
}},
})
if err != nil {
t.Fatalf("buildResponsesRequest: %v", err)
}
if len(request.Tools) != 1 {
t.Fatalf("tools = %#v, want one", request.Tools)
}
tool := request.Tools[0]
if tool.Type != string(zeroruntime.ToolDefinitionFunction) || tool.Format != nil ||
tool.Description != "Custom integration tool." || !reflect.DeepEqual(tool.Parameters, parameters) {
t.Fatalf("custom apply_patch function was rewritten: %#v", tool)
}
}
func TestHandleTerminalResponseNilPayload(t *testing.T) {
p := &CodexProvider{}
@@ -212,14 +90,11 @@ func TestHandleTerminalResponseFailedPayloadWithoutError(t *testing.T) {
// A response.completed with a payload and no error remains a clean done.
ok := make(chan zeroruntime.StreamEvent, 8)
p.handleTerminalResponse(context.Background(),
&responsesEvent{Type: responsesEventCompleted, Response: &responsePayload{ID: "resp-1", Status: "completed"}}, &responsesState{}, ok)
&responsesEvent{Type: responsesEventCompleted, Response: &responsePayload{Status: "completed"}}, &responsesState{}, ok)
close(ok)
for ev := range ok {
if ev.Type == zeroruntime.StreamEventError {
t.Error("response.completed with a non-error payload must NOT emit an error")
}
if ev.Type == zeroruntime.StreamEventDone && ev.ResponseID != "resp-1" {
t.Errorf("done response id = %q, want resp-1", ev.ResponseID)
}
}
}
+3 -3
View File
@@ -771,7 +771,7 @@ func TestCodexProviderParsesResponsesTextDeltas(t *testing.T) {
`{"type":"response.output_text.delta","delta":"hello "}`,
`{"type":"response.output_text.delta","delta":"world"}`,
`{"type":"response.output_text.delta","delta":"!"}`,
`{"type":"response.completed","response":{"id":"resp-1","status":"completed","usage":{"input_tokens":7,"output_tokens":3,"output_tokens_details":{"reasoning_tokens":1},"input_tokens_details":{"cached_tokens":2,"cache_write_tokens":1}}}}`,
`{"type":"response.completed","response":{"id":"resp-1","status":"completed","usage":{"input_tokens":7,"output_tokens":3,"output_tokens_details":{"reasoning_tokens":1},"input_tokens_details":{"cached_tokens":2}}}}`,
)
defer srv.Close()
@@ -809,8 +809,8 @@ func TestCodexProviderParsesResponsesTextDeltas(t *testing.T) {
if usage == nil {
t.Fatal("expected a usage event, got none")
}
if usage.InputTokens != 7 || usage.OutputTokens != 3 || usage.CachedInputTokens != 2 || usage.CacheWriteTokens != 1 || usage.ReasoningTokens != 1 {
t.Fatalf("usage = %+v, want input=7 output=3 cached=2 cache-write=1 reasoning=1", *usage)
if usage.InputTokens != 7 || usage.OutputTokens != 3 || usage.CachedInputTokens != 2 || usage.ReasoningTokens != 1 {
t.Fatalf("usage = %+v, want input=7 output=3 cached=2 reasoning=1", *usage)
}
if !gotDone {
t.Fatal("expected a done event, got none")
-1
View File
@@ -357,7 +357,6 @@ func (provider *Provider) emitChunk(
PromptTokens: chunk.Usage.PromptTokens,
CompletionTokens: chunk.Usage.CompletionTokens,
CachedInputTokens: chunk.Usage.PromptTokensDetails.CachedTokens,
CacheWriteTokens: chunk.Usage.PromptTokensDetails.CacheWriteTokens,
ReasoningTokens: chunk.Usage.CompletionTokensDetails.ReasoningTokens,
},
})
+2 -2
View File
@@ -260,14 +260,14 @@ func TestStreamCompletionAppliesCustomAuthAndHeaders(t *testing.T) {
func TestStreamCompletionEmitsTextUsageAndDone(t *testing.T) {
provider := newTestProvider(t, func(w http.ResponseWriter, r *http.Request) {
writeSSE(w, `{"choices":[{"delta":{"content":"hello "}}]}`)
writeSSE(w, `{"choices":[{"delta":{"content":"zero"}}],"usage":{"prompt_tokens":12,"completion_tokens":5,"prompt_tokens_details":{"cached_tokens":3,"cache_write_tokens":2}}}`)
writeSSE(w, `{"choices":[{"delta":{"content":"zero"}}],"usage":{"prompt_tokens":12,"completion_tokens":5,"prompt_tokens_details":{"cached_tokens":3}}}`)
writeSSE(w, `[DONE]`)
})
events := collectProviderEvents(t, provider)
assertEvent(t, events[0], zeroruntime.StreamEventText, "hello ")
assertEvent(t, events[1], zeroruntime.StreamEventText, "zero")
if events[2].Type != zeroruntime.StreamEventUsage || events[2].Usage.PromptTokens != 12 || events[2].Usage.CompletionTokens != 5 || events[2].Usage.CachedInputTokens != 3 || events[2].Usage.CacheWriteTokens != 2 {
if events[2].Type != zeroruntime.StreamEventUsage || events[2].Usage.PromptTokens != 12 || events[2].Usage.CompletionTokens != 5 || events[2].Usage.CachedInputTokens != 3 {
t.Fatalf("unexpected usage event: %#v", events[2])
}
if events[3].Type != zeroruntime.StreamEventDone {
+1 -6
View File
@@ -195,12 +195,7 @@ func (s *turnSession) computeFingerprint(request zeroruntime.CompletionRequest)
func wireToolsDigest(tools []zeroruntime.ToolDefinition) string {
rendered := make([]string, 0, len(tools))
for _, tool := range tools {
wireShape := struct {
Type zeroruntime.ToolDefinitionType `json:"type,omitempty"`
Parameters map[string]any `json:"parameters,omitempty"`
Format *zeroruntime.ToolDefinitionFormat `json:"format,omitempty"`
}{Type: tool.Type, Parameters: tool.Parameters, Format: tool.Format}
schema, err := json.Marshal(wireShape)
schema, err := json.Marshal(tool.Parameters)
if err != nil {
rendered = append(rendered, tool.Name+"\n"+tool.Description+"\n__non_json:"+tool.Name)
continue
+1 -2
View File
@@ -109,8 +109,7 @@ type usage struct {
}
type promptTokenDetails struct {
CachedTokens int `json:"cached_tokens"`
CacheWriteTokens int `json:"cache_write_tokens"`
CachedTokens int `json:"cached_tokens"`
}
type completionTokenDetails struct {
+13 -44
View File
@@ -15,27 +15,9 @@ import (
// value (e.g. "1") to enable — the same boolean idiom as ZERO_FORMAT_ON_WRITE.
const openaiTurnSessionEnv = "ZERO_OPENAI_TURN_SESSION"
// ChatGPT Responses sessions are enabled by default for their native provider.
// Setting this to 0/false/off restores the stateless HTTP/SSE transport.
const chatGPTTurnSessionEnv = "ZERO_CHATGPT_TURN_SESSION"
func openaiTurnSessionEnabled() bool {
value := strings.TrimSpace(os.Getenv(openaiTurnSessionEnv))
return value != "" && !explicitEnvFalse(value)
}
func chatGPTTurnSessionEnabled() bool {
value := strings.TrimSpace(os.Getenv(chatGPTTurnSessionEnv))
return value == "" || !explicitEnvFalse(value)
}
func explicitEnvFalse(value string) bool {
switch strings.ToLower(strings.TrimSpace(value)) {
case "0", "false", "off":
return true
default:
return false
}
return value != "" && value != "0" && !strings.EqualFold(value, "false")
}
// OptimizedTurnSessions returns the gated optimized TurnSessionProvider for an
@@ -43,35 +25,22 @@ func explicitEnvFalse(value string) bool {
// is ineligible. Callers leave agent Options.TurnSessionProvider nil on false,
// so the loop's default wrap keeps today's behavior byte-identical.
//
// Eligibility is deliberately narrow: native ChatGPT Responses uses its
// concrete Codex provider and defaults on with a kill switch; official OpenAI
// chat-completions keeps its existing opt-in prewarm session. Compatible
// gateways and foreign provider values always retain the default adapter.
// Eligibility is deliberately narrow: the resolved provider kind must be
// official OpenAI (NOT openai-compatible gateways — the constructor merges the
// two, so this branches on the resolved kind explicitly), not the ChatGPT
// Codex catalog, and the provider value must be the concrete *openai.Provider
// (a fake, a Codex provider, or nil falls back to the default path safely).
// No base-URL check on top of the kind: prewarm and fingerprint telemetry are
// harmless against any host, and kind==openai mirrors the existing
// official-OpenAI precedent used for prompt_cache_key.
func OptimizedTurnSessions(profile config.ProviderProfile, provider zeroruntime.Provider, options Options) (zeroruntime.TurnSessionProvider, bool) {
resolved, err := resolveProfile(profile, options)
if err != nil {
return nil, false
}
if isCodexCatalog(profile, resolved) {
if !chatGPTTurnSessionEnabled() {
return nil, false
}
concrete, ok := provider.(*openai.CodexProvider)
if !ok {
return nil, false
}
caps, err := resolveCapabilities(profile, options)
if err != nil {
return nil, false
}
return openai.NewCodexTurnSessionProvider(concrete, caps), true
}
if resolved.providerKind != config.ProviderKindOpenAI {
return nil, false
}
if !openaiTurnSessionEnabled() {
return nil, false
}
resolved, err := resolveProfile(profile, options)
if err != nil || resolved.providerKind != config.ProviderKindOpenAI || isCodexCatalog(profile, resolved) {
return nil, false
}
concrete, ok := provider.(*openai.Provider)
if !ok {
return nil, false
+5 -72
View File
@@ -5,7 +5,6 @@ import (
"testing"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/providers/openai"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
@@ -46,7 +45,7 @@ func TestOptimizedTurnSessionsDefaultOff(t *testing.T) {
func TestOptimizedTurnSessionsFalseyValues(t *testing.T) {
profile := openaiEligibleProfile()
provider := buildGateProvider(t, profile)
for _, value := range []string{"0", "false", "FALSE", "False", "off", "OFF", " ", ""} {
for _, value := range []string{"0", "false", "FALSE", "False", " ", ""} {
t.Setenv(openaiTurnSessionEnv, value)
if _, ok := OptimizedTurnSessions(profile, provider, Options{}); ok {
t.Fatalf("gate enabled for falsey value %q", value)
@@ -60,21 +59,6 @@ func TestOptimizedTurnSessionsFalseyValues(t *testing.T) {
}
}
func TestChatGPTTurnSessionGateValues(t *testing.T) {
for _, value := range []string{"0", "false", "FALSE", "off", "OFF"} {
t.Setenv(chatGPTTurnSessionEnv, value)
if chatGPTTurnSessionEnabled() {
t.Fatalf("ChatGPT gate enabled for falsey value %q", value)
}
}
for _, value := range []string{"", " ", "1", "true", "on"} {
t.Setenv(chatGPTTurnSessionEnv, value)
if !chatGPTTurnSessionEnabled() {
t.Fatalf("ChatGPT gate disabled for default/truthy value %q", value)
}
}
}
func TestOptimizedTurnSessionsOpenAIEligible(t *testing.T) {
t.Setenv(openaiTurnSessionEnv, "1")
profile := openaiEligibleProfile()
@@ -104,64 +88,13 @@ func TestOptimizedTurnSessionsRejectsCompatible(t *testing.T) {
}
func TestOptimizedTurnSessionsRejectsCodexCatalog(t *testing.T) {
t.Setenv(chatGPTTurnSessionEnv, "0")
profile := openaiEligibleProfile()
profile.Name = "chatgpt"
profile.CatalogID = "chatgpt"
if _, ok := OptimizedTurnSessions(profile, buildChatGPTGateProvider(t, profile), Options{}); ok {
t.Fatal("the ChatGPT Codex catalog must honor the turn-session kill switch")
}
}
func TestOptimizedTurnSessionsEnablesChatGPTResponsesByDefault(t *testing.T) {
t.Setenv(chatGPTTurnSessionEnv, "")
profile := openaiEligibleProfile()
profile.Name = "chatgpt"
profile.CatalogID = "chatgpt"
provider := buildChatGPTGateProvider(t, profile)
turnSessions, ok := OptimizedTurnSessions(profile, provider, Options{})
if !ok || turnSessions == nil {
t.Fatal("ChatGPT Responses provider did not receive its default turn session")
}
if got := turnSessions.Capabilities().Model; got != profile.Model {
t.Fatalf("Capabilities().Model = %q, want %q", got, profile.Model)
}
}
func TestOptimizedTurnSessionsRecognizesCompatibleLabeledChatGPTProfile(t *testing.T) {
t.Setenv(chatGPTTurnSessionEnv, "")
profile := openaiEligibleProfile()
profile.Name = "chatgpt"
profile.CatalogID = "chatgpt"
profile.ProviderKind = config.ProviderKindOpenAICompatible
if turnSessions, ok := OptimizedTurnSessions(profile, buildChatGPTGateProvider(t, profile), Options{}); !ok || turnSessions == nil {
t.Fatal("ChatGPT catalog profile must use its native session regardless of the stored compatibility label")
}
}
func buildChatGPTGateProvider(t *testing.T, profile config.ProviderProfile) *openai.CodexProvider {
t.Helper()
provider, err := openai.NewCodexProvider(openai.CodexOptions{
Options: openai.Options{
APIKey: "test-token",
BaseURL: "https://chatgpt.example/backend-api/codex",
Model: profile.Model,
},
AccountID: "account-test",
})
if err != nil {
t.Fatalf("NewCodexProvider: %v", err)
}
return provider
}
func TestOptimizedTurnSessionsRejectsNonCodexProviderForChatGPT(t *testing.T) {
t.Setenv(chatGPTTurnSessionEnv, "1")
t.Setenv(openaiTurnSessionEnv, "1")
profile := openaiEligibleProfile()
profile.CatalogID = "chatgpt"
// Build the provider from the base profile: the point under test is the
// catalog check, not Codex construction (which needs OAuth wiring).
if _, ok := OptimizedTurnSessions(profile, buildGateProvider(t, openaiEligibleProfile()), Options{}); ok {
t.Fatal("a ChatGPT catalog profile with a non-Codex provider must use the default session")
t.Fatal("the ChatGPT Codex catalog must not get the optimized session")
}
}
@@ -1,96 +0,0 @@
package sandbox
import (
"path/filepath"
"strings"
"testing"
)
func TestApplyPatchPathBlockOnlyRejectsRelativeTraversal(t *testing.T) {
root := t.TempDir()
inside := filepath.Join(root, "main.js")
structured := func(path string) string {
return strings.Join([]string{"*** Begin Patch", "*** Update File: " + path, "@@", "-a", "+b", "*** End Patch"}, "\n")
}
for name, patch := range map[string]string{
"absolute inside workspace": structured(inside),
"relative": structured("main.js"),
"decorated markers": "*** Begin Patch ***\n*** Update File: main.js\n@@\n-a\n+b\n*** End Patch ***",
"no-space marker": "***Begin Patch\n*** Update File: main.js\n@@\n-a\n+b\n***End Patch",
} {
request := Request{ToolName: "apply_patch", WorkspaceRoot: root, SideEffect: SideEffectWrite, Args: map[string]any{"patch": patch}}
if block := applyPatchPathBlock(request); block != nil {
t.Fatalf("%s: unexpected block %+v", name, block)
}
}
for _, path := range []string{"../escape.js", ".."} {
for name, patch := range map[string]string{"canonical": structured(path), "no-space": "***Begin Patch\n*** Update File: " + path + "\n@@\n-a\n+b\n***End Patch"} {
request := Request{ToolName: "apply_patch", WorkspaceRoot: root, SideEffect: SideEffectWrite, Args: map[string]any{"patch": patch}}
block := applyPatchPathBlock(request)
if block == nil || block.Code != BlockOutsideWorkspace {
t.Fatalf("%s %q must be blocked as traversal, got %+v", name, path, block)
}
}
}
}
// Every marker spelling the tool applies must be classified as structured at
// the sandbox boundary too; otherwise the boundary scans the patch as a unified
// diff, extracts no targets, and validates nothing (fail-open).
func TestStructuredPatchClassifierMatchesToolSpellings(t *testing.T) {
for _, header := range []string{"*** Begin Patch", "*** Begin Patch ***", "***Begin Patch", " *** Begin Patch ", "\ufeff*** Begin Patch"} {
patch := header + "\n*** Update File: main.js\n@@\n-a\n+b\n*** End Patch"
if !IsStructuredPatch(patch) {
t.Fatalf("%q must classify as a structured patch", header)
}
if paths := applyPatchPaths(patch); len(paths) != 1 || paths[0] != "main.js" {
t.Fatalf("%q: sandbox must extract the structured target, got %v", header, paths)
}
}
for _, header := range []string{"--- a/x", "Begin Patch", "*** Begin Patchwork", "*** Update File: x"} {
if IsStructuredPatch(header + "\n-a\n+b") {
t.Fatalf("%q must not classify as a structured patch", header)
}
}
if StructuredPatchMarker("*** End Patch ***") != "end" || StructuredPatchMarker("***End Patch") != "end" {
t.Fatal("decorated end markers must classify as end")
}
}
func TestApplyPatchRequestPathsCarryAbsolutePathsToScopeValidation(t *testing.T) {
root := t.TempDir()
// NewScope also grants the system temp dir, so a sibling t.TempDir() is
// legitimately in scope; pick a path under the filesystem root instead.
outside, err := filepath.Abs(filepath.Join(string(filepath.Separator), "zero-outside-workspace-test", "escape.js"))
if err != nil {
t.Fatal(err)
}
scope, err := NewScope(root, nil)
if err != nil {
t.Fatal(err)
}
inside := filepath.Join(root, "main.js")
structured := func(header, footer, path string) string {
return strings.Join([]string{header, "*** Update File: " + path, "@@", "-a", "+b", footer}, "\n")
}
for name, spelling := range map[string][2]string{"canonical": {"*** Begin Patch", "*** End Patch"}, "no-space": {"***Begin Patch", "***End Patch"}} {
// Failure path: the exact path the boundary parsed must be denied by the
// scope. structuredPatchHeaderPaths normalises separators to "/", so the
// parsed form is compared in slash form and then validated as-is.
paths := applyPatchRequestPaths(map[string]any{"patch": structured(spelling[0], spelling[1], outside)})
if len(paths) != 1 || paths[0] != filepath.ToSlash(outside) {
t.Fatalf("%s: absolute patch path must reach scope validation unchanged, got %v", name, paths)
}
if block := scope.validate(paths[0]); block == nil || block.Code != BlockOutsideWorkspace {
t.Fatalf("%s: scope must deny the parsed outside path %q, got %+v", name, paths[0], block)
}
// Success path: the parsed inside path must be accepted by the scope.
paths = applyPatchRequestPaths(map[string]any{"patch": structured(spelling[0], spelling[1], inside)})
if len(paths) != 1 || paths[0] != filepath.ToSlash(inside) {
t.Fatalf("%s: inside patch path must reach scope validation unchanged, got %v", name, paths)
}
if block := scope.validate(paths[0]); block != nil {
t.Fatalf("%s: scope must accept the parsed inside path %q, got %+v", name, paths[0], block)
}
}
}
+2 -33
View File
@@ -281,15 +281,11 @@ func applyPatchPathBlock(request Request) *pathBlock {
if patch == "" {
return nil
}
// Only relative traversal is rejected up front. Absolute paths flow through
// the regular workspace-scope validation below (requestPaths), which accepts
// one inside the workspace and denies one outside — a model that echoes the
// absolute path read_file showed it must not be blocked for that alone.
for _, path := range applyPatchPaths(patch) {
if path == "" || path == "/dev/null" {
continue
}
if path == ".." || strings.HasPrefix(path, "../") {
if filepath.IsAbs(path) || path == ".." || strings.HasPrefix(path, "../") {
return &pathBlock{
Code: BlockOutsideWorkspace,
Path: path,
@@ -300,35 +296,8 @@ func applyPatchPathBlock(request Request) *pathBlock {
return nil
}
// structuredPatchMarkerPattern is the single classifier for structured-patch
// markers, shared by the sandbox boundary and the apply_patch tool (which
// imports this package). It accepts the canonical "*** Begin Patch" /
// "*** End Patch" and the decorated spellings models emit ("*** Begin Patch ***",
// "***Begin Patch", trailing whitespace). Both sides must agree: a spelling the
// tool would apply but the sandbox did not recognise would make the sandbox
// scan the patch as a unified diff, extract no targets, and validate nothing.
var structuredPatchMarkerPattern = regexp.MustCompile(`^\*{3}\s*(Begin|End) Patch\s*\**\s*$`)
// StructuredPatchMarker classifies a line as the "begin" or "end" marker of a
// structured patch, or "" when it is neither.
func StructuredPatchMarker(line string) string {
match := structuredPatchMarkerPattern.FindStringSubmatch(strings.TrimSpace(line))
if match == nil {
return ""
}
return strings.ToLower(match[1])
}
// IsStructuredPatch reports whether patch opens with a structured begin marker.
// The tool applies exactly the patches this returns true for, so the sandbox
// extracts structured header paths for exactly the same set.
func IsStructuredPatch(patch string) bool {
first, _, _ := strings.Cut(strings.TrimSpace(strings.TrimPrefix(patch, "\ufeff")), "\n")
return StructuredPatchMarker(first) == "begin"
}
func applyPatchPaths(patch string) []string {
if IsStructuredPatch(patch) {
if strings.HasPrefix(strings.TrimSpace(strings.TrimPrefix(patch, "\ufeff")), "*** Begin Patch") {
return structuredPatchHeaderPaths(patch)
}
return patchHeaderPaths(patch)
-2
View File
@@ -94,8 +94,6 @@ type Event struct {
Meta map[string]string `json:"meta,omitempty"`
PromptTokens *int `json:"promptTokens,omitempty"`
CompletionTokens *int `json:"completionTokens,omitempty"`
CachedInputTokens *int `json:"cachedInputTokens,omitempty"`
CacheWriteTokens *int `json:"cacheWriteTokens,omitempty"`
TotalTokens *int `json:"totalTokens,omitempty"`
CostUSD *float64 `json:"costUsd,omitempty"`
Text string `json:"text,omitempty"`
-15
View File
@@ -1,15 +0,0 @@
//go:build !windows
package swarm
import (
"errors"
"os"
)
// isLockContended reports whether an O_EXCL lock-create error means the lock is
// currently held (so the caller should retry/wait) rather than a hard failure.
// On POSIX an existing lock surfaces only as os.ErrExist.
func isLockContended(err error) bool {
return errors.Is(err, os.ErrExist)
}
@@ -0,0 +1,18 @@
//go:build !windows
package swarm
import (
"os"
"path/filepath"
"testing"
)
func makeRedirectedLockPath(t *testing.T, root, target string) string {
t.Helper()
lockPath := filepath.Join(root, "redirected.lock")
if err := os.Symlink(target, lockPath); err != nil {
t.Fatalf("create redirected lock path: %v", err)
}
return lockPath
}
@@ -0,0 +1,20 @@
//go:build windows
package swarm
import (
"os/exec"
"path/filepath"
"strings"
"testing"
)
func makeRedirectedLockPath(t *testing.T, root, target string) string {
t.Helper()
linkDir := filepath.Join(root, "redirected")
out, err := exec.Command("cmd", "/c", "mklink", "/J", linkDir, filepath.Dir(target)).CombinedOutput()
if err != nil {
t.Fatalf("create redirected lock junction: %v %s", err, strings.TrimSpace(string(out)))
}
return filepath.Join(linkDir, filepath.Base(target))
}
-30
View File
@@ -1,30 +0,0 @@
//go:build windows
package swarm
import (
"errors"
"os"
"syscall"
)
// Windows error codes that an O_EXCL create (or a racing delete) can surface
// when another writer holds the lock file open, instead of ERROR_FILE_EXISTS.
// Treating them as "contended" makes the lock retry/wait rather than fail the
// whole send/read spuriously under contention.
const (
errorAccessDenied = syscall.Errno(5) // ERROR_ACCESS_DENIED
errorSharingViolation = syscall.Errno(32) // ERROR_SHARING_VIOLATION
)
// isLockContended reports whether an O_EXCL lock-create error means the lock is
// currently held (so the caller should retry/wait) rather than a hard failure.
// On Windows a held lock can appear as ERROR_FILE_EXISTS (mapped to os.ErrExist)
// but also as ERROR_ACCESS_DENIED / ERROR_SHARING_VIOLATION when the file is
// concurrently open or pending deletion.
func isLockContended(err error) bool {
if errors.Is(err, os.ErrExist) {
return true
}
return errors.Is(err, errorAccessDenied) || errors.Is(err, errorSharingViolation)
}
+12 -58
View File
@@ -8,7 +8,6 @@ import (
"path/filepath"
"regexp"
"strings"
"sync/atomic"
"time"
"github.com/Gitlawb/zero/internal/fsutil"
@@ -17,10 +16,10 @@ import (
// Mailbox is a per-agent, per-team message inbox persisted as a JSON array on
// disk at <baseDir>/<team>/inboxes/<agent>.json, written atomically under a
// lock file. It is hardened for untrusted input:
// advisory lock. It is hardened for untrusted input:
//
// - inbox files and their parent dirs are owner-only (0600/0700);
// - every write takes an exclusive lock file (bounded retry, stale-break);
// - every write takes an exclusive advisory lock (bounded retry);
// - message bodies and whole inboxes are size-capped (fail closed on oversize);
// - team/agent names are sanitized so a name can never escape the base dir;
// - malformed inbox JSON is rejected rather than silently reset (fail closed).
@@ -44,7 +43,6 @@ const (
defaultMaxMessageBytes = 64 * 1024 // 64 KiB per message
defaultMaxMessages = 1000 // per inbox
defaultLockTimeout = 5 * time.Second
lockStaleAfter = 30 * time.Second
inboxFileMode = 0o600
inboxDirMode = 0o700
)
@@ -206,7 +204,7 @@ func (m *Mailbox) Send(team, recipient string, msg Message) error {
if err := m.ensureInboxDir(path); err != nil {
return err
}
release, err := acquireLock(path+".lock", m.lockTimeout())
release, err := acquireLock(m.BaseDir, path+".lock", m.lockTimeout())
if err != nil {
return err
}
@@ -233,12 +231,12 @@ func (m *Mailbox) ReadAndConsume(team, recipient string) ([]Message, error) {
if err != nil {
return nil, err
}
// The lock file lives beside the inbox; create the dir so it can be taken
// The stable lock file lives beside the inbox; create the dir so it can be taken
// even on first read of a not-yet-existing inbox.
if err := m.ensureInboxDir(path); err != nil {
return nil, err
}
release, err := acquireLock(path+".lock", m.lockTimeout())
release, err := acquireLock(m.BaseDir, path+".lock", m.lockTimeout())
if err != nil {
return nil, err
}
@@ -334,63 +332,19 @@ func (m *Mailbox) atomicWriteJSON(path string, data any) error {
return nil
}
// lockSeq makes each lock token unique within the process even when two
// goroutines acquire in the same nanosecond.
var lockSeq atomic.Uint64
// acquireLock takes an exclusive lock by creating lockPath with O_EXCL. It
// retries with a short backoff until timeout, breaking a lock whose file is
// older than lockStaleAfter (so a crashed holder cannot deadlock the swarm).
//
// Each holder writes a unique token into the lock file and release is
// OWNERSHIP-AWARE: it removes the lock only if it still holds our token. After a
// stale-break another writer may legitimately own the lock; unconditionally
// removing it would let a third writer in and create split-brain writers that
// can corrupt or drop mailbox updates.
func acquireLock(lockPath string, timeout time.Duration) (func(), error) {
token := fmt.Sprintf("%d-%d-%d", os.Getpid(), time.Now().UnixNano(), lockSeq.Add(1))
// acquireLock takes an exclusive advisory lock through a stable lock file. It
// retries with a short backoff until timeout. The kernel releases the lock when
// a holder exits, so crash recovery never moves or deletes the canonical path.
func acquireLock(root, lockPath string, timeout time.Duration) (func(), error) {
deadline := time.Now().Add(timeout)
for {
f, err := os.OpenFile(lockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, inboxFileMode)
lock, err := lockutil.TryAcquireFileLockAt(root, lockPath)
if err == nil {
// Record our ownership token; the exclusive create already holds the
// lock, so a write error only weakens diagnostics, not safety.
_, _ = f.WriteString(token)
f.Close()
var released bool
return func() {
if released {
return
}
released = true
if data, rerr := os.ReadFile(lockPath); rerr == nil && string(data) == token {
_ = lockutil.RemoveLockFile(lockPath)
}
}, nil
return func() { _ = lock.Release() }, nil
}
if !isLockContended(err) {
if !errors.Is(err, lockutil.ErrLockHeld) {
return nil, fmt.Errorf("swarm: acquire lock: %w", err)
}
// Lock held: break it if stale via the atomic rename-with-verify in
// lockutil.ReclaimStaleLock (AUDIT-M13), otherwise wait.
if info, statErr := os.Stat(lockPath); statErr == nil && time.Since(info.ModTime()) > lockStaleAfter {
cleared, rerr := lockutil.ReclaimStaleLock(lockPath, token, func(reclaimedPath string) bool {
info, err := os.Stat(reclaimedPath)
return err == nil && time.Since(info.ModTime()) <= lockStaleAfter
})
if rerr != nil {
// Reclaim hit a hard failure: the rename aside failed outright, or a
// live holder's lock could not be put back (the lock path may be
// missing, so re-acquiring would break mutual exclusion). Fail closed
// instead of spinning to the deadline.
return nil, fmt.Errorf("swarm: reclaim stale lock: %w", rerr)
}
if cleared {
continue // cleared a genuinely stale lock; retry the O_EXCL create now
}
// Lost the reclaim race (or it was actually fresh) — fall through to the
// bounded wait instead of hot-spinning on a reclaim that never wins.
}
if time.Now().After(deadline) {
return nil, fmt.Errorf("swarm: timed out acquiring lock %s", filepath.Base(lockPath))
}
+38 -27
View File
@@ -219,64 +219,75 @@ func TestMailboxRejectsSymlinkedInboxDir(t *testing.T) {
}
}
func TestMailboxLockReleaseIsOwnershipAware(t *testing.T) {
func TestMailboxLockReleaseKeepsStablePath(t *testing.T) {
mb := newTestMailbox(t)
path, _ := mb.inboxPath("team", "bob")
if err := mb.ensureInboxDir(path); err != nil {
t.Fatalf("ensureInboxDir: %v", err)
}
lockPath := path + ".lock"
// Writer A acquires the lock.
releaseA, err := acquireLock(lockPath, time.Second)
releaseA, err := acquireLock(mb.BaseDir, lockPath, time.Second)
if err != nil {
t.Fatalf("acquire A: %v", err)
}
// Simulate a stale-break + takeover by writer B: overwrite the lock content
// with B's token (as a fresh acquire after a break would).
if err := os.WriteFile(lockPath, []byte("writer-B-token"), 0o600); err != nil {
t.Fatalf("simulate B takeover: %v", err)
}
// A's release must NOT delete B's lock (ownership-aware).
releaseA()
if _, err := os.Stat(lockPath); err != nil {
t.Fatalf("A's release deleted B's lock (split-brain): %v", err)
t.Fatalf("release removed stable advisory-lock path: %v", err)
}
// B's own release removes it.
os.Remove(lockPath)
}
func TestAcquireLockReclaimsStaleLock(t *testing.T) {
// A crashed holder's stale lock (old mtime) must be reclaimed via the atomic
// rename-with-verify path, leaving no sidelined .stale.* file (AUDIT-M13).
func TestAcquireLockIgnoresOldUnlockedFile(t *testing.T) {
// A crashed holder can leave an old file behind, but the kernel releases its
// advisory lock. Acquisition must succeed without renaming the stable path.
lockPath := filepath.Join(t.TempDir(), "x.lock")
if err := os.WriteFile(lockPath, []byte("dead-holder"), 0o600); err != nil {
t.Fatalf("seed stale lock: %v", err)
t.Fatalf("seed old lock metadata: %v", err)
}
old := time.Now().Add(-2 * lockStaleAfter)
old := time.Now().Add(-time.Hour)
if err := os.Chtimes(lockPath, old, old); err != nil {
t.Fatalf("chtimes: %v", err)
}
release, err := acquireLock(lockPath, time.Second)
release, err := acquireLock(filepath.Dir(lockPath), lockPath, time.Second)
if err != nil {
t.Fatalf("acquireLock should reclaim a genuinely stale lock, got %v", err)
t.Fatalf("acquire over old unlocked file: %v", err)
}
release()
if matches, _ := filepath.Glob(lockPath + ".stale.*"); len(matches) != 0 {
t.Fatalf("reclaim left sidelined files: %v", matches)
if _, err := os.Stat(lockPath); err != nil {
t.Fatalf("stable lock path missing after release: %v", err)
}
}
func TestAcquireLockDoesNotBreakFreshLock(t *testing.T) {
// A fresh, held lock (recent mtime) must never be broken — the stale-break must
// not steal a live lock (AUDIT-M13).
func TestAcquireLockDoesNotBreakHeldLock(t *testing.T) {
// File age is irrelevant while the kernel reports an active holder.
lockPath := filepath.Join(t.TempDir(), "x.lock")
release, err := acquireLock(lockPath, time.Second)
release, err := acquireLock(filepath.Dir(lockPath), lockPath, time.Second)
if err != nil {
t.Fatalf("first acquire: %v", err)
}
defer release()
if _, err := acquireLock(lockPath, 50*time.Millisecond); err == nil {
t.Fatal("acquireLock broke a fresh, held lock (split-brain risk)")
if _, err := acquireLock(filepath.Dir(lockPath), lockPath, 50*time.Millisecond); err == nil {
t.Fatal("acquireLock admitted a second holder")
}
}
func TestAcquireLockRefusesRedirectedPath(t *testing.T) {
root := t.TempDir()
target := filepath.Join(t.TempDir(), "target")
const sentinel = "do not overwrite"
if err := os.WriteFile(target, []byte(sentinel), 0o600); err != nil {
t.Fatal(err)
}
lockPath := makeRedirectedLockPath(t, root, target)
if release, err := acquireLock(root, lockPath, time.Second); err == nil {
release()
t.Fatal("acquireLock followed a redirected lock path")
}
data, err := os.ReadFile(target)
if err != nil {
t.Fatal(err)
}
if string(data) != sentinel {
t.Fatalf("redirected target was modified: %q", data)
}
}
+187 -127
View File
@@ -3,11 +3,11 @@ package tools
import (
"context"
"fmt"
"os"
"os/exec"
"path/filepath"
"strconv"
"strings"
"github.com/Gitlawb/zero/internal/sandbox"
)
type applyPatchTool struct {
@@ -16,110 +16,15 @@ type applyPatchTool struct {
scope PathScope
}
func (applyPatchTool) isBuiltInApplyPatch() {}
// PrepareFreeformApplyPatchArguments converts native structured-patch input
// into the ordinary apply_patch argument map used by permission checks and tool
// execution. Absolute patch-header paths select a granted write root and are
// rewritten relative to that root; relative-only patches keep workspace-root
// behavior.
func PrepareFreeformApplyPatchArguments(tool Tool, patch string) (map[string]any, error) {
preparer, ok := tool.(interface {
prepareFreeformArguments(string) (map[string]any, error)
})
if !ok {
return nil, fmt.Errorf("tool is not Zero's built-in apply_patch")
}
return preparer.prepareFreeformArguments(patch)
}
func (tool applyPatchTool) prepareFreeformArguments(patch string) (map[string]any, error) {
args := map[string]any{"patch": patch}
if !isStructuredPatch(patch) {
return args, nil
}
lines := strings.Split(strings.ReplaceAll(patch, "\r\n", "\n"), "\n")
selectedRoot := ""
hasRelativeHeader := false
for index, line := range lines {
prefix := ""
switch {
case strings.HasPrefix(line, structuredAddFile):
prefix = structuredAddFile
case strings.HasPrefix(line, structuredDeleteFile):
prefix = structuredDeleteFile
case strings.HasPrefix(line, structuredUpdateFile):
prefix = structuredUpdateFile
case strings.HasPrefix(line, structuredMoveTo):
prefix = structuredMoveTo
default:
continue
}
path := strings.TrimSpace(strings.TrimPrefix(line, prefix))
if !filepath.IsAbs(path) {
hasRelativeHeader = true
continue
}
root, relative, err := tool.freeformPatchRoot(path)
if err != nil {
return nil, err
}
if selectedRoot != "" && selectedRoot != root {
return nil, fmt.Errorf("freeform patch spans multiple write roots")
}
selectedRoot = root
lines[index] = prefix + filepath.ToSlash(relative)
}
if selectedRoot != "" {
roots, err := scopedRoots(tool.workspaceRoot, tool.scope)
if err != nil {
return nil, err
}
if hasRelativeHeader && selectedRoot != roots[0] {
return nil, fmt.Errorf("freeform patch mixes workspace-relative paths with an extra write root")
}
args["cwd"] = selectedRoot
args["patch"] = strings.Join(lines, "\n")
}
return args, nil
}
func (tool applyPatchTool) freeformPatchRoot(path string) (string, string, error) {
roots, err := scopedRoots(tool.workspaceRoot, tool.scope)
if err != nil {
return "", "", err
}
var firstErr error
for _, root := range roots {
candidate := sandbox.NormalizePrefixForRoot(path, root)
candidate, err = filepath.Abs(candidate)
if err != nil {
if firstErr == nil {
firstErr = err
}
continue
}
relative, err := filepath.Rel(root, candidate)
if err == nil && relative != ".." && !strings.HasPrefix(relative, ".."+string(filepath.Separator)) && !filepath.IsAbs(relative) {
return root, filepath.ToSlash(relative), nil
}
if firstErr == nil {
firstErr = outsideWorkspaceError(path)
}
}
return "", "", firstErr
}
func NewScopedApplyPatchTool(workspaceRoot string, scope PathScope) Tool {
return applyPatchTool{
baseTool: baseTool{
name: "apply_patch",
description: "Apply a multi-hunk or multi-file patch inside the workspace (or a granted extra write root). Structured format:\n*** Begin Patch\n*** Update File: src/app.js\n@@\n unchanged context line\n-removed line\n+added line\n*** End Patch\nUse \"*** Add File: path\" with \"+\" lines to create a file and \"*** Delete File: path\" to remove one; several sections may follow each other. A unified diff (---/+++ headers with a/ b/ prefixes, @@ hunks) is also accepted and applied in-process. Paths are workspace-relative; absolute paths inside the workspace are fine. For a single targeted change, edit_file is simpler.",
description: "Apply a patch inside the workspace or an explicitly granted extra write root.",
parameters: Schema{
Type: "object",
Properties: map[string]PropertySchema{
"patch": {Type: "string", Description: "The structured (*** Begin Patch) or unified-diff patch text."},
"patch": {Type: "string", Description: "A unified diff or a structured *** Begin Patch patch to apply."},
"cwd": {Type: "string", Description: "Directory where the patch should be applied. Relative paths stay in the workspace; use an absolute path to target a granted extra write root. Defaults to workspace root.", Default: "."},
},
Required: []string{"patch"},
@@ -154,14 +59,177 @@ func (tool applyPatchTool) RunWithOptions(ctx context.Context, args map[string]a
if isStructuredPatch(patch) {
return tool.runStructuredPatch(applyRoot, relativeRoot, patch, options)
}
// Unified diffs are translated into the same operations and applied by
// the same os.Root engine, so neither format opens a target by pathname
// after validation (no check-to-use window) and git is not needed.
operations, err := parseUnifiedPatch(patch)
if err := validatePatchPaths(applyRoot, patch); err != nil {
return errorResult("Error applying patch: " + err.Error())
}
tempFile, err := os.CreateTemp("", "zero-patch-*.patch")
if err != nil {
return errorResult("Error applying patch: " + err.Error())
}
return applyPatchOperations(applyRoot, relativeRoot, operations, options)
patchPath := tempFile.Name()
defer func() {
_ = os.Remove(patchPath)
}()
if _, err := tempFile.WriteString(patch); err != nil {
_ = tempFile.Close()
return errorResult("Error applying patch: " + err.Error())
}
if err := tempFile.Close(); err != nil {
return errorResult("Error applying patch: " + err.Error())
}
if err := recheckPatchWriteTargets(applyRoot, patch); err != nil {
return errorResult("Error applying patch: " + err.Error())
}
var createdTargets []string
var fullySuppliedTargets []string
wholeBefore := map[string]bool{}
if options.FileTracker != nil {
createdTargets = missingPatchTargets(applyRoot, patch)
fullySuppliedTargets = completeCreatedPatchTargets(applyRoot, patch)
for _, path := range patchHeaderPaths(patch) {
if path == "" || path == "/dev/null" {
continue
}
if absolute, _, rerr := resolveWorkspaceTargetPath(applyRoot, path); rerr == nil {
wholeBefore[absolute] = options.FileTracker.SeenWhole(absolute)
}
}
}
command := exec.CommandContext(ctx, "git", "apply", "--whitespace=nowarn", patchPath)
command.Dir = applyRoot
output, err := command.CombinedOutput()
if err != nil {
message := strings.TrimSpace(string(output))
if message == "" {
message = err.Error()
}
return errorResult("Error applying patch: " + message)
}
summary := "Patch applied successfully."
if relativeRoot != "." {
summary = "Patch applied successfully in " + relativeRoot + "."
}
result := okResult(summary)
result.ChangedFiles = changedFilesFromPatch(relativeRoot, patch)
result.Display = Display{Summary: summary, Kind: "diff", Preview: capPreviewDiff(patch)}
fullySupplied := make(map[string]bool, len(fullySuppliedTargets))
for _, absolute := range fullySuppliedTargets {
fullySupplied[absolute] = true
}
// Re-baseline files changed by this tool. When the model had already seen the
// whole input (or supplied a complete new file), the exact patch plus that
// input determines the whole output, so a follow-up edit needs no wasted read.
// Partial observations remain conservative and are cleared by Record.
for _, changed := range result.ChangedFiles {
if absolute, _, rerr := resolveScopedPath(tool.workspaceRoot, tool.scope, changed); rerr == nil {
content, readErr := os.ReadFile(absolute)
if readErr != nil {
options.FileTracker.Forget(absolute)
continue
}
info, _ := os.Stat(absolute)
wasWhole := wholeBefore[absolute] || fullySupplied[absolute]
options.FileTracker.Record(absolute, content, info)
if wasWhole {
lines := lineCount(string(content))
options.FileTracker.RecordSeenRange(absolute, 1, lines, lines)
}
}
}
recordCreatedPatchTargets(options.FileTracker, createdTargets)
return result
}
func missingPatchTargets(root string, patch string) []string {
seen := map[string]bool{}
var missing []string
for _, path := range patchHeaderPaths(patch) {
if path == "" || path == "/dev/null" {
continue
}
absolute, _, err := resolveWorkspaceTargetPath(root, path)
if err != nil || seen[absolute] {
continue
}
seen[absolute] = true
if _, err := os.Stat(absolute); os.IsNotExist(err) {
missing = append(missing, absolute)
}
}
return missing
}
// completeCreatedPatchTargets returns only files whose full contents are
// supplied by a /dev/null creation patch. A missing rename/copy destination is
// created by git too, but its bytes come from an unread source and must not gain
// whole-file observation credit.
func completeCreatedPatchTargets(root string, patch string) []string {
seen := map[string]bool{}
var created []string
oldRemaining, newRemaining := 0, 0
inHunk := false
fromDevNull := false
for _, line := range strings.Split(strings.ReplaceAll(patch, "\r\n", "\n"), "\n") {
if inHunk && (oldRemaining > 0 || newRemaining > 0) {
switch {
case strings.HasPrefix(line, "-"):
oldRemaining--
case strings.HasPrefix(line, "+"):
newRemaining--
case strings.HasPrefix(line, "\\"):
default:
oldRemaining--
newRemaining--
}
continue
}
inHunk = false
switch {
case strings.HasPrefix(line, "diff --git "):
fromDevNull = false
case strings.HasPrefix(line, "@@"):
oldRemaining, newRemaining = parseHunkCounts(line)
inHunk = oldRemaining > 0 || newRemaining > 0
case strings.HasPrefix(line, "--- "):
fromDevNull = patchFileHeaderPath(line) == "/dev/null"
case strings.HasPrefix(line, "+++ "):
path := patchFileHeaderPath(line)
if !fromDevNull || path == "" || path == "/dev/null" {
fromDevNull = false
continue
}
fromDevNull = false
absolute, _, err := resolveWorkspaceTargetPath(root, stripPatchPrefix(path))
if err != nil || seen[absolute] {
continue
}
if _, err := os.Stat(absolute); !os.IsNotExist(err) {
continue
}
seen[absolute] = true
created = append(created, absolute)
}
}
return created
}
func recordCreatedPatchTargets(tracker *FileTracker, missingBefore []string) {
if tracker == nil {
return
}
for _, absolute := range missingBefore {
if _, err := os.Stat(absolute); err != nil {
continue
}
if resolved, err := filepath.EvalSymlinks(absolute); err == nil {
absolute = resolved
}
tracker.RecordCreated(absolute)
}
}
// changedFilesFromPatch extracts the unique, WORKSPACE-relative paths a patch
@@ -192,35 +260,27 @@ func changedFilesFromPatch(relativeRoot string, patch string) []string {
return paths
}
// normalizePatchPathForRoot resolves platform-level symlinks (macOS /var ->
// /private/var) in the prefix of an absolute patch path that lies outside the
// apply root, so a path the model copied from read_file compares equal to the
// symlink-resolved root. Relative paths are returned unchanged.
func normalizePatchPathForRoot(root string, path string) string {
if !filepath.IsAbs(path) {
return path
}
resolvedRoot, err := filepath.Abs(root)
if err != nil {
return path
}
if evaluated, err := filepath.EvalSymlinks(resolvedRoot); err == nil {
resolvedRoot = evaluated
}
return sandbox.NormalizePrefixForRoot(path, resolvedRoot)
}
func validatePatchPaths(root string, patch string) error {
for _, path := range patchHeaderPaths(patch) {
if path == "" || path == "/dev/null" {
continue
}
// Relative traversal is rejected here; an absolute path is checked by
// resolveWorkspaceTargetPath, which only accepts one inside the root.
if path == ".." || strings.HasPrefix(path, "../") {
if filepath.IsAbs(path) || path == ".." || strings.HasPrefix(path, "../") {
return fmt.Errorf("patch path %q must stay inside the workspace", path)
}
if _, _, err := resolveWorkspaceTargetPath(root, normalizePatchPathForRoot(root, path)); err != nil {
if _, _, err := resolveWorkspaceTargetPath(root, path); err != nil {
return err
}
}
return nil
}
func recheckPatchWriteTargets(root string, patch string) error {
for _, path := range patchHeaderPaths(patch) {
if path == "" || path == "/dev/null" {
continue
}
if err := recheckWorkspaceWriteTarget(root, path); err != nil {
return err
}
}
@@ -291,7 +351,7 @@ func patchFileHeaderPath(line string) string {
// tokens. Scanning the whole line would let a crafted heading like
// "@@ -1,1 +1,1 @@ +1,999999" overwrite the real count, keep the parser stuck in
// hunk mode, and swallow later "--- "/"+++ " file headers so they escape
// validatePatchPaths — a workspace-confinement bypass.
// validatePatchPaths / recheckPatchWriteTargets — a workspace-confinement bypass.
func parseHunkCounts(line string) (int, int) {
_, rest, ok := strings.Cut(line, "@@")
if !ok {
@@ -1,712 +0,0 @@
package tools
import (
"context"
"os"
"path/filepath"
"strconv"
"strings"
"testing"
)
// Models routinely decorate the structured markers ("*** Begin Patch ***") and
// write unified-diff ranges after "@@"; both used to fail on the first line and
// pushed the model into whole-file rewrites.
func TestApplyPatchToleratesDecoratedMarkersAndRangeHeaders(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "hello.txt"), "hello\nold\nbye\n")
patch := strings.Join([]string{
"*** Begin Patch ***",
"*** Update File: hello.txt",
"@@ -1,3 +1,3 @@",
" hello",
"-old",
"+new",
"*** End Patch ***",
"",
}, "\n")
if !isStructuredPatch(patch) {
t.Fatal("decorated begin marker must still be recognised as a structured patch")
}
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": patch})
if result.Status != StatusOK {
t.Fatalf("decorated structured patch should apply, got %s: %s", result.Status, result.Output)
}
content, err := os.ReadFile(filepath.Join(root, "hello.txt"))
if err != nil {
t.Fatal(err)
}
if got := strings.ReplaceAll(string(content), "\r\n", "\n"); got != "hello\nnew\nbye\n" {
t.Fatalf("patched content = %q", got)
}
}
func TestStructuredHunkAnchorKeepsHeadingAndPlainContext(t *testing.T) {
cases := map[string]string{
"-12,4 +12,6": "",
"-12 +12": "",
"-12,4 +12,6 @@": "",
"-12,4 +12,6 @@ func main() {": "func main() {",
"func main() {": "func main() {",
"": "",
}
for input, want := range cases {
if got := structuredHunkAnchor(input); got != want {
t.Fatalf("structuredHunkAnchor(%q) = %q, want %q", input, got, want)
}
}
}
func TestStructuredPatchMarkerSpellings(t *testing.T) {
begin := []string{"*** Begin Patch", "*** Begin Patch ***", "***Begin Patch", " *** Begin Patch ", "*** Begin Patch **"}
for _, line := range begin {
if structuredPatchMarker(line) != "begin" {
t.Fatalf("%q must read as the begin marker", line)
}
}
if structuredPatchMarker("*** End Patch ***") != "end" {
t.Fatal("decorated end marker must read as the end marker")
}
for _, line := range []string{"*** Update File: x", "Begin Patch", "*** Begin Patchwork"} {
if structuredPatchMarker(line) != "" {
t.Fatalf("%q must not read as a marker", line)
}
}
}
func TestStructuredPatchHeaderErrorCarriesFormatHint(t *testing.T) {
_, err := parseStructuredPatch("--- a/x\n+++ b/x\n")
if err == nil || !strings.Contains(err.Error(), "*** Update File: path") {
t.Fatalf("header error should teach the format, got %v", err)
}
}
func TestApplyPatchAcceptsAbsoluteInWorkspacePaths(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "hello.txt"), "hello\nold\n")
absolute := filepath.Join(root, "hello.txt")
structured := strings.Join([]string{
"*** Begin Patch",
"*** Update File: " + absolute,
"@@",
" hello",
"-old",
"+structured",
"*** End Patch",
"",
}, "\n")
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": structured})
if result.Status != StatusOK {
t.Fatalf("absolute in-workspace structured path should apply, got %s: %s", result.Status, result.Output)
}
content, _ := os.ReadFile(absolute)
if got := strings.ReplaceAll(string(content), "\r\n", "\n"); got != "hello\nstructured\n" {
t.Fatalf("structured content = %q", got)
}
if len(result.ChangedFiles) != 1 || result.ChangedFiles[0] != "hello.txt" {
t.Fatalf("changed files should be workspace-relative, got %v", result.ChangedFiles)
}
unified := strings.Join([]string{
"--- " + absolute,
"+++ " + absolute,
"@@ -1,2 +1,2 @@",
" hello",
"-structured",
"+unified",
"",
}, "\n")
result = NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": unified})
if result.Status != StatusOK {
t.Fatalf("absolute in-workspace unified path should apply, got %s: %s", result.Status, result.Output)
}
content, _ = os.ReadFile(absolute)
if got := strings.ReplaceAll(string(content), "\r\n", "\n"); got != "hello\nunified\n" {
t.Fatalf("unified content = %q", got)
}
}
func TestApplyPatchStillRejectsAbsolutePathsOutsideWorkspace(t *testing.T) {
root := t.TempDir()
outside := filepath.Join(t.TempDir(), "escape.txt")
writeTestFile(t, outside, "old\n")
for name, patch := range map[string]string{
"structured": strings.Join([]string{"*** Begin Patch", "*** Update File: " + outside, "@@", "-old", "+new", "*** End Patch", ""}, "\n"),
"unified": strings.Join([]string{"--- " + outside, "+++ " + outside, "@@ -1 +1 @@", "-old", "+new", ""}, "\n"),
} {
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": patch})
if result.Status == StatusOK {
t.Fatalf("%s patch outside the workspace must be rejected", name)
}
if content, _ := os.ReadFile(outside); string(content) != "old\n" {
t.Fatalf("%s patch must not touch a file outside the workspace", name)
}
}
}
// Unified diffs are applied in-process through the same os.Root engine as
// structured patches: git C-quoted paths with whitespace resolve to one file,
// and the same header outside the workspace is refused before anything is
// written.
func TestUnifiedPatchQuotedWhitespacePaths(t *testing.T) {
root := t.TempDir()
dir := filepath.Join(root, "work tree")
if err := os.MkdirAll(dir, 0o755); err != nil {
t.Fatal(err)
}
target := filepath.Join(dir, "file.txt")
writeTestFile(t, target, "hello\nold\n")
quoted := func(path string) string { return strconv.Quote(path) }
inside := strings.Join([]string{
"diff --git " + quoted("a/"+filepath.ToSlash(target)) + " " + quoted("b/"+filepath.ToSlash(target)),
"--- " + quoted(target),
"+++ " + quoted(target),
"@@ -1,2 +1,2 @@",
" hello",
"-old",
"+new",
"",
}, "\n")
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": inside})
if result.Status != StatusOK {
t.Fatalf("quoted whitespace path inside the workspace should apply, got %s: %s", result.Status, result.Output)
}
content, _ := os.ReadFile(target)
if got := strings.ReplaceAll(string(content), "\r\n", "\n"); got != "hello\nnew\n" {
t.Fatalf("content = %q", got)
}
if len(result.ChangedFiles) != 1 || result.ChangedFiles[0] != "work tree/file.txt" {
t.Fatalf("changed files = %v", result.ChangedFiles)
}
outsideDir := filepath.Join(t.TempDir(), "other tree")
if err := os.MkdirAll(outsideDir, 0o755); err != nil {
t.Fatal(err)
}
outside := filepath.Join(outsideDir, "file.txt")
writeTestFile(t, outside, "hello\nold\n")
escape := strings.Join([]string{"--- " + quoted(outside), "+++ " + quoted(outside), "@@ -1,2 +1,2 @@", " hello", "-old", "+new", ""}, "\n")
result = NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": escape})
if result.Status == StatusOK {
t.Fatal("quoted whitespace path outside the workspace must be rejected")
}
if content, _ := os.ReadFile(outside); string(content) != "hello\nold\n" {
t.Fatal("outside file must be untouched")
}
}
// A workspace path that is swapped for a symlink escaping the root — the
// classic check-to-use attack — is refused by the os.Root engine for both
// formats, and the outside file is never modified.
func TestApplyPatchRefusesSymlinkEscapingWorkspace(t *testing.T) {
root := t.TempDir()
outside := filepath.Join(t.TempDir(), "secret.txt")
writeTestFile(t, outside, "hello\nold\n")
link := filepath.Join(root, "target.txt")
if err := os.Symlink(outside, link); err != nil {
t.Skipf("symlinks unavailable: %v", err)
}
for name, patch := range map[string]string{
"unified": strings.Join([]string{"--- a/target.txt", "+++ b/target.txt", "@@ -1,2 +1,2 @@", " hello", "-old", "+new", ""}, "\n"),
"structured": strings.Join([]string{"*** Begin Patch", "*** Update File: target.txt", "@@", " hello", "-old", "+new", "*** End Patch", ""}, "\n"),
} {
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": patch})
if result.Status == StatusOK {
t.Fatalf("%s patch through an escaping symlink must be refused", name)
}
if content, _ := os.ReadFile(outside); string(content) != "hello\nold\n" {
t.Fatalf("%s patch must not modify the file outside the workspace", name)
}
}
}
func TestUnifiedPatchCreatesDeletesAndHonoursNoNewlineMarker(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "gone.txt"), "bye\n")
writeTestFile(t, filepath.Join(root, "keep.txt"), "a\nb\n")
patch := strings.Join([]string{
"diff --git a/new.txt b/new.txt",
"new file mode 100644",
"--- /dev/null",
"+++ b/new.txt",
"@@ -0,0 +1,2 @@",
"+first",
"+second",
"\\ No newline at end of file",
"diff --git a/gone.txt b/gone.txt",
"deleted file mode 100644",
"--- a/gone.txt",
"+++ /dev/null",
"@@ -1 +0,0 @@",
"-bye",
"--- a/keep.txt",
"+++ b/keep.txt",
"@@ -1,2 +1,2 @@",
" a",
"-b",
"+c",
"\\ No newline at end of file",
"",
}, "\n")
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": patch})
if result.Status != StatusOK {
t.Fatalf("multi-file unified patch should apply, got %s: %s", result.Status, result.Output)
}
if content, _ := os.ReadFile(filepath.Join(root, "new.txt")); string(content) != "first\nsecond" {
t.Fatalf("created file = %q", string(content))
}
if _, err := os.Stat(filepath.Join(root, "gone.txt")); !os.IsNotExist(err) {
t.Fatal("deleted file must be gone")
}
if content, _ := os.ReadFile(filepath.Join(root, "keep.txt")); string(content) != "a\nc" {
t.Fatalf("updated file = %q", string(content))
}
if len(result.ChangedFiles) != 3 {
t.Fatalf("changed files = %v", result.ChangedFiles)
}
}
func TestUnifiedPatchUsesRangeHintThenFallsBackToContext(t *testing.T) {
root := t.TempDir()
// Two identical blocks: the hunk's range picks the second one, which
// content search alone would call ambiguous.
writeTestFile(t, filepath.Join(root, "dup.txt"), "x\ny\nx\ny\n")
patch := strings.Join([]string{"--- a/dup.txt", "+++ b/dup.txt", "@@ -3,2 +3,2 @@", " x", "-y", "+z", ""}, "\n")
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": patch})
if result.Status != StatusOK {
t.Fatalf("range-hinted hunk should apply, got %s: %s", result.Status, result.Output)
}
if content, _ := os.ReadFile(filepath.Join(root, "dup.txt")); string(content) != "x\ny\nx\nz\n" {
t.Fatalf("content = %q", string(content))
}
// A stale range (file grew above the hunk) still applies by unique context.
writeTestFile(t, filepath.Join(root, "moved.txt"), "header\nhello\nold\n")
patch = strings.Join([]string{"--- a/moved.txt", "+++ b/moved.txt", "@@ -1,2 +1,2 @@", " hello", "-old", "+new", ""}, "\n")
result = NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": patch})
if result.Status != StatusOK {
t.Fatalf("offset hunk should apply by context, got %s: %s", result.Status, result.Output)
}
if content, _ := os.ReadFile(filepath.Join(root, "moved.txt")); string(content) != "header\nhello\nnew\n" {
t.Fatalf("content = %q", string(content))
}
}
func TestUnifiedPatchRenameWithHunkAndRejectsBinary(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "x.txt"), "hello\nold\n")
rename := strings.Join([]string{
"diff --git a/x.txt b/y.txt",
"similarity index 90%",
"rename from x.txt",
"rename to y.txt",
"--- a/x.txt",
"+++ b/y.txt",
"@@ -1,2 +1,2 @@",
" hello",
"-old",
"+new",
"",
}, "\n")
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": rename})
if result.Status != StatusOK {
t.Fatalf("rename with hunk should apply, got %s: %s", result.Status, result.Output)
}
if _, err := os.Stat(filepath.Join(root, "x.txt")); !os.IsNotExist(err) {
t.Fatal("renamed source must be gone")
}
if content, _ := os.ReadFile(filepath.Join(root, "y.txt")); string(content) != "hello\nnew\n" {
t.Fatalf("renamed content = %q", string(content))
}
for name, patch := range map[string]string{
"binary": "diff --git a/x b/x\nBinary files a/x and b/x differ\n",
"empty": "diff --git a/x b/x\n",
} {
if result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": patch}); result.Status == StatusOK {
t.Fatalf("%s patch must be rejected", name)
}
}
}
func TestUnifiedPatchCopyOperation(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "src.txt"), "hello\nold\n")
copyPatch := func(from, to string, hunk bool) string {
lines := []string{"diff --git a/" + from + " b/" + to, "similarity index 90%", "copy from " + from, "copy to " + to}
if hunk {
lines = append(lines, "--- a/"+from, "+++ b/"+to, "@@ -1,2 +1,2 @@", " hello", "-old", "+new")
}
return strings.Join(append(lines, ""), "\n")
}
// Success: source kept, destination created with the hunk applied, and the
// destination inherits the source's tracker state (read whole -> whole).
tracker := NewFileTracker()
registry := NewRegistry()
registry.Register(NewScopedReadFileTool(root, nil))
registry.Register(NewScopedApplyPatchTool(root, nil))
opts := RunOptions{PermissionGranted: true, FileTracker: tracker}
if r := registry.RunWithOptions(context.Background(), "read_file", map[string]any{"path": "src.txt"}, opts); r.Status != StatusOK {
t.Fatalf("read: %s", r.Output)
}
result := registry.RunWithOptions(context.Background(), "apply_patch", map[string]any{"patch": copyPatch("src.txt", "dst.txt", true)}, opts)
if result.Status != StatusOK {
t.Fatalf("copy with hunk should apply, got %s: %s", result.Status, result.Output)
}
if content, _ := os.ReadFile(filepath.Join(root, "src.txt")); string(content) != "hello\nold\n" {
t.Fatalf("copy must keep the source, got %q", string(content))
}
if content, _ := os.ReadFile(filepath.Join(root, "dst.txt")); string(content) != "hello\nnew\n" {
t.Fatalf("copy destination = %q", string(content))
}
if len(result.ChangedFiles) != 2 || result.ChangedFiles[0] != "src.txt" || result.ChangedFiles[1] != "dst.txt" {
t.Fatalf("changed files = %v", result.ChangedFiles)
}
destination, err := filepath.EvalSymlinks(filepath.Join(root, "dst.txt"))
if err != nil {
t.Fatal(err)
}
if !tracker.SeenWhole(destination) {
t.Fatal("destination copied from a fully read source must be tracked as seen whole")
}
unreadSource, _ := filepath.EvalSymlinks(filepath.Join(root, "src.txt"))
if _, tracked := tracker.Version(unreadSource); !tracked {
t.Fatal("source must stay tracked after a copy")
}
// Failure: destination already exists — nothing is written.
writeTestFile(t, filepath.Join(root, "taken.txt"), "keep\n")
result = NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": copyPatch("src.txt", "taken.txt", false)})
if result.Status == StatusOK || !strings.Contains(result.Output, "already exists") {
t.Fatalf("copy onto an existing destination must be refused, got %s: %s", result.Status, result.Output)
}
if content, _ := os.ReadFile(filepath.Join(root, "taken.txt")); string(content) != "keep\n" {
t.Fatal("existing destination must be untouched")
}
// Failure: copy onto itself.
result = NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": copyPatch("src.txt", "src.txt", false)})
if result.Status == StatusOK || !strings.Contains(result.Output, "onto itself") {
t.Fatalf("copy onto itself must be refused, got %s: %s", result.Status, result.Output)
}
}
func TestUnifiedPatchRejectsOverDeclaredHunkCounts(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "a.txt"), "hello\nold\n")
writeTestFile(t, filepath.Join(root, "b.txt"), "x\n")
// "-1,3 +1,3" over-declares a two-line hunk, so without a guard the next
// file's ---/+++ headers would be swallowed as "-"/"+" content.
patch := strings.Join([]string{
"--- a/a.txt", "+++ b/a.txt", "@@ -1,3 +1,3 @@", " hello", "-old", "+new",
"--- a/b.txt", "+++ b/b.txt", "@@ -1 +1 @@", "-x", "+y", "",
}, "\n")
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": patch})
if result.Status == StatusOK || !strings.Contains(result.Output, "declared line counts") {
t.Fatalf("over-declared hunk must be reported as malformed, got %s: %s", result.Status, result.Output)
}
for name, want := range map[string]string{"a.txt": "hello\nold\n", "b.txt": "x\n"} {
if content, _ := os.ReadFile(filepath.Join(root, name)); string(content) != want {
t.Fatalf("%s must be untouched after a malformed patch, got %q", name, string(content))
}
}
// Truncated at end of input is reported the same way.
truncated := "--- a/a.txt\n+++ b/a.txt\n@@ -1,2 +1,2 @@\n hello\n-old\n"
if result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": truncated}); result.Status == StatusOK || !strings.Contains(result.Output, "declared line counts") {
t.Fatalf("truncated hunk must be reported as malformed, got %s: %s", result.Status, result.Output)
}
}
func TestUnifiedPatchNoNewlineMarkerOnlyAfterFinalHunk(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "two.txt"), "a\nb\nc\nd\ne\nf\n")
// A marker after the first of two hunks describes nothing real and must
// not leak into the second hunk's result: the patch is rejected untouched.
early := strings.Join([]string{
"--- a/two.txt", "+++ b/two.txt",
"@@ -1,2 +1,2 @@", " a", "-b", "+B", "\\ No newline at end of file",
"@@ -5,2 +5,2 @@", " e", "-f", "+F", "",
}, "\n")
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": early})
if result.Status == StatusOK || !strings.Contains(result.Output, "must follow the last hunk") {
t.Fatalf("marker before a later hunk must be rejected, got %s: %s", result.Status, result.Output)
}
if content, _ := os.ReadFile(filepath.Join(root, "two.txt")); string(content) != "a\nb\nc\nd\ne\nf\n" {
t.Fatal("rejected patch must not modify the file")
}
// The same marker after the final hunk applies and strips the newline.
final := strings.Join([]string{
"--- a/two.txt", "+++ b/two.txt",
"@@ -1,2 +1,2 @@", " a", "-b", "+B",
"@@ -5,2 +5,2 @@", " e", "-f", "+F", "\\ No newline at end of file", "",
}, "\n")
result = NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": final})
if result.Status != StatusOK {
t.Fatalf("two-hunk patch with a final marker should apply, got %s: %s", result.Status, result.Output)
}
if content, _ := os.ReadFile(filepath.Join(root, "two.txt")); string(content) != "a\nB\nc\nd\ne\nF" {
t.Fatalf("content = %q", string(content))
}
}
// A removed "-- …" line directly followed by an added "++ …" line looks like a
// ---/+++ header pair; the boundary check must not mistake it for one.
func TestUnifiedPatchKeepsAdjacentDashPlusContentLines(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "notes.sql"), "select 1;\n-- old comment\nselect 2;\n")
patch := strings.Join([]string{
"--- a/notes.sql", "+++ b/notes.sql",
"@@ -1,3 +1,3 @@", " select 1;", "--- old comment", "+++ new comment", " select 2;", "",
}, "\n")
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": patch})
if result.Status != StatusOK {
t.Fatalf("adjacent -- / ++ content lines must apply, got %s: %s", result.Status, result.Output)
}
if content, _ := os.ReadFile(filepath.Join(root, "notes.sql")); string(content) != "select 1;\n++ new comment\nselect 2;\n" {
t.Fatalf("content = %q", string(content))
}
}
// A unified deletion states the content it removes; it must be verified
// against the current file before anything is deleted.
func TestUnifiedPatchDeletionVerifiesExpectedContent(t *testing.T) {
deletion := func(path string, hunks ...string) string {
lines := []string{"diff --git a/" + path + " b/" + path, "deleted file mode 100644", "--- a/" + path, "+++ /dev/null"}
return strings.Join(append(append(lines, hunks...), ""), "\n")
}
t.Run("matching deletion succeeds", func(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "gone.txt"), "current\n")
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": deletion("gone.txt", "@@ -1 +0,0 @@", "-current")})
if result.Status != StatusOK {
t.Fatalf("matching deletion should apply, got %s: %s", result.Status, result.Output)
}
if _, err := os.Stat(filepath.Join(root, "gone.txt")); !os.IsNotExist(err) {
t.Fatal("file must be deleted")
}
})
t.Run("stale deletion is refused untouched", func(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "gone.txt"), "current\n")
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": deletion("gone.txt", "@@ -1 +0,0 @@", "-stale")})
if result.Status == StatusOK || !strings.Contains(result.Output, "does not match its current content") {
t.Fatalf("stale deletion must be refused, got %s: %s", result.Status, result.Output)
}
if content, _ := os.ReadFile(filepath.Join(root, "gone.txt")); string(content) != "current\n" {
t.Fatalf("file must be byte-for-byte unchanged, got %q", string(content))
}
})
t.Run("partial deletion is refused", func(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "gone.txt"), "one\ntwo\n")
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": deletion("gone.txt", "@@ -1 +0,0 @@", "-one")})
if result.Status == StatusOK || !strings.Contains(result.Output, "does not match its current content") {
t.Fatalf("deletion that does not cover the file must be refused, got %s: %s", result.Status, result.Output)
}
if content, _ := os.ReadFile(filepath.Join(root, "gone.txt")); string(content) != "one\ntwo\n" {
t.Fatal("file must be unchanged")
}
})
t.Run("header-only deletion is refused", func(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "gone.txt"), "current\n")
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": deletion("gone.txt")})
if result.Status == StatusOK || !strings.Contains(result.Output, "must include the hunk") {
t.Fatalf("header-only deletion must be refused, got %s: %s", result.Status, result.Output)
}
if _, err := os.Stat(filepath.Join(root, "gone.txt")); err != nil {
t.Fatal("file must still exist")
}
})
t.Run("multi-hunk deletion verifies everything before removing", func(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "gone.txt"), "a\nb\nc\nd\n")
good := deletion("gone.txt", "@@ -1,2 +0,0 @@", "-a", "-b", "@@ -3,2 +0,0 @@", "-c", "-d")
bad := deletion("gone.txt", "@@ -1,2 +0,0 @@", "-a", "-b", "@@ -3,2 +0,0 @@", "-c", "-stale")
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": bad})
if result.Status == StatusOK {
t.Fatalf("deletion with one stale hunk must be refused: %s", result.Output)
}
if content, _ := os.ReadFile(filepath.Join(root, "gone.txt")); string(content) != "a\nb\nc\nd\n" {
t.Fatal("no hunk may be applied when any hunk is stale")
}
result = NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": good})
if result.Status != StatusOK {
t.Fatalf("multi-hunk deletion should apply, got %s: %s", result.Status, result.Output)
}
if _, err := os.Stat(filepath.Join(root, "gone.txt")); !os.IsNotExist(err) {
t.Fatal("file must be deleted")
}
})
t.Run("structured delete stays unconditional", func(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "gone.txt"), "anything\n")
patch := "*** Begin Patch\n*** Delete File: gone.txt\n*** End Patch\n"
if result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": patch}); result.Status != StatusOK {
t.Fatalf("structured delete should apply, got %s: %s", result.Status, result.Output)
}
})
}
// The source is re-read through the root immediately before commit, so a
// change made after planning is refused instead of overwritten or removed.
func TestApplyStructuredPatchChangeRefusesSourceChangedAfterPlanning(t *testing.T) {
root := t.TempDir()
path := filepath.Join(root, "file.txt")
writeTestFile(t, path, "planned\n")
workspace, err := os.OpenRoot(root)
if err != nil {
t.Fatal(err)
}
defer workspace.Close()
target := structuredPatchTarget{requested: "file.txt", absolute: path, relative: "file.txt"}
destination := structuredPatchTarget{requested: "copy.txt", absolute: filepath.Join(root, "copy.txt"), relative: "copy.txt"}
for name, change := range map[string]structuredPatchChange{
"update": {kind: structuredPatchUpdate, from: target, to: target, before: "planned\n", after: "new\n", mode: 0o644},
"delete": {kind: structuredPatchDelete, from: target, to: target, before: "planned\n", mode: 0o644},
"copy": {kind: structuredPatchCopy, from: target, to: destination, before: "planned\n", after: "planned\n", mode: 0o644},
} {
writeTestFile(t, path, "changed after planning\n")
committed, err := applyStructuredPatchChange(workspace, change)
if err == nil || committed || !strings.Contains(err.Error(), "changed on disk between planning and commit") {
t.Fatalf("%s: expected a refusal, got committed=%v err=%v", name, committed, err)
}
if content, _ := os.ReadFile(path); string(content) != "changed after planning\n" {
t.Fatalf("%s: file must be untouched, got %q", name, string(content))
}
if _, statErr := os.Stat(filepath.Join(root, "copy.txt")); !os.IsNotExist(statErr) {
t.Fatalf("%s: refused change must not create copy.txt", name)
}
}
}
func TestUnifiedPatchDeletionIsByteExact(t *testing.T) {
deletion := func(hunks ...string) string {
return strings.Join(append(append([]string{"--- a/gone.txt", "+++ /dev/null"}, hunks...), ""), "\n")
}
for name, tc := range map[string]struct{ file, patch string }{
"trailing whitespace differs": {"current \n", deletion("@@ -1 +0,0 @@", "-current")},
"indentation differs": {" current\n", deletion("@@ -1 +0,0 @@", "-current")},
"blank line not covered": {"current\n\n", deletion("@@ -1 +0,0 @@", "-current")},
"space-only line not covered": {"current\n \n", deletion("@@ -1 +0,0 @@", "-current")},
"missing final newline": {"current", deletion("@@ -1 +0,0 @@", "-current")},
} {
t.Run(name, func(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "gone.txt"), tc.file)
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": tc.patch})
if result.Status == StatusOK {
t.Fatalf("deletion must be byte-exact; %q was deleted by a non-matching hunk", tc.file)
}
if content, _ := os.ReadFile(filepath.Join(root, "gone.txt")); string(content) != tc.file {
t.Fatalf("file must be unchanged, got %q", string(content))
}
})
}
t.Run("no-newline marker on the removed side matches exactly", func(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "gone.txt"), "current")
patch := deletion("@@ -1 +0,0 @@", "-current", "\\ No newline at end of file")
if result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": patch}); result.Status != StatusOK {
t.Fatalf("deletion with a matching no-newline marker should apply, got %s: %s", result.Status, result.Output)
}
if _, err := os.Stat(filepath.Join(root, "gone.txt")); !os.IsNotExist(err) {
t.Fatal("file must be deleted")
}
})
t.Run("CRLF file matches an LF patch of the same content", func(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "gone.txt"), "a\r\nb\r\n")
if result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": deletion("@@ -1,2 +0,0 @@", "-a", "-b")}); result.Status != StatusOK {
t.Fatalf("CRLF deletion should apply, got %s: %s", result.Status, result.Output)
}
})
}
// git describes an empty file's creation or deletion with header lines only.
func TestUnifiedPatchHeaderOnlyEmptyFileForms(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "empty.txt"), "")
writeTestFile(t, filepath.Join(root, "full.txt"), "content\n")
writeTestFile(t, filepath.Join(root, "keep.txt"), "a\n")
deleteHeader := func(path string) string {
return "diff --git a/" + path + " b/" + path + "\ndeleted file mode 100644\nindex e69de29..0000000\n"
}
// A non-empty file is not deleted by the header-only form.
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": deleteHeader("full.txt")})
if result.Status == StatusOK || !strings.Contains(result.Output, "expects an empty file") {
t.Fatalf("header-only deletion of a non-empty file must be refused, got %s: %s", result.Status, result.Output)
}
if _, err := os.Stat(filepath.Join(root, "full.txt")); err != nil {
t.Fatal("non-empty file must still exist")
}
// In a multi-file patch the empty-file deletion and creation are applied
// alongside a normal hunk.
patch := deleteHeader("empty.txt") +
"diff --git a/blank.txt b/blank.txt\nnew file mode 100644\nindex 0000000..e69de29\n" +
"diff --git a/keep.txt b/keep.txt\n--- a/keep.txt\n+++ b/keep.txt\n@@ -1 +1 @@\n-a\n+b\n"
result = NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": patch})
if result.Status != StatusOK {
t.Fatalf("multi-file patch with header-only forms should apply, got %s: %s", result.Status, result.Output)
}
if _, err := os.Stat(filepath.Join(root, "empty.txt")); !os.IsNotExist(err) {
t.Fatal("empty file must be deleted")
}
if content, err := os.ReadFile(filepath.Join(root, "blank.txt")); err != nil || len(content) != 0 {
t.Fatalf("empty file must be created, got err=%v content=%q", err, string(content))
}
if content, _ := os.ReadFile(filepath.Join(root, "keep.txt")); string(content) != "b\n" {
t.Fatalf("hunk file = %q", string(content))
}
if len(result.ChangedFiles) != 3 {
t.Fatalf("changed files = %v", result.ChangedFiles)
}
}
// When a later change fails, the error names exactly which files were already
// committed so the caller knows what changed and what did not.
func TestApplyPatchOperationsReportsCommittedPrefixOnFailure(t *testing.T) {
root := t.TempDir()
writeTestFile(t, filepath.Join(root, "first.txt"), "one\n")
writeTestFile(t, filepath.Join(root, "second.txt"), "two\n")
writeTestFile(t, filepath.Join(root, "third.txt"), "three\n")
// Make second.txt change after planning so its commit is refused.
structuredPatchBeforeCommit = func(change structuredPatchChange) {
if change.to.relative == "second.txt" {
writeTestFile(t, filepath.Join(root, "second.txt"), "changed\n")
}
}
defer func() { structuredPatchBeforeCommit = nil }()
patch := strings.Join([]string{
"*** Begin Patch",
"*** Update File: first.txt", "@@", "-one", "+ONE",
"*** Update File: second.txt", "@@", "-two", "+TWO",
"*** Update File: third.txt", "@@", "-three", "+THREE",
"*** End Patch", "",
}, "\n")
result := NewScopedApplyPatchTool(root, nil).Run(context.Background(), map[string]any{"patch": patch})
if result.Status == StatusOK {
t.Fatal("patch must fail when a later change is refused")
}
for _, want := range []string{"already committed: first.txt", "remaining files are unchanged"} {
if !strings.Contains(result.Output, want) {
t.Fatalf("error must report the committed prefix, got: %s", result.Output)
}
}
if strings.Contains(result.Output, "committed: first.txt, second.txt") || strings.Contains(result.Output, "third.txt") {
t.Fatalf("error must not list uncommitted files as committed: %s", result.Output)
}
if content, _ := os.ReadFile(filepath.Join(root, "first.txt")); string(content) != "ONE\n" {
t.Fatalf("first.txt must hold the committed change, got %q", string(content))
}
if content, _ := os.ReadFile(filepath.Join(root, "second.txt")); string(content) != "changed\n" {
t.Fatalf("second.txt must be untouched by the patch, got %q", string(content))
}
if content, _ := os.ReadFile(filepath.Join(root, "third.txt")); string(content) != "three\n" {
t.Fatalf("third.txt must be untouched, got %q", string(content))
}
}
+1 -1
View File
@@ -168,7 +168,7 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
options.FileTracker.Record(absolutePath, []byte(updated), newInfo)
if updated == modelKnownContent {
if previouslySeenWhole {
options.FileTracker.RecordSeenRange(absolutePath, 1, trackedLineTotal(updated), trackedLineTotal(updated))
options.FileTracker.RecordSeenRange(absolutePath, 1, lineCount(updated), lineCount(updated))
} else {
for _, span := range editedSpans {
options.FileTracker.RecordSeenBytes(absolutePath, span.start, span.end, len(updated))
+7 -7
View File
@@ -26,8 +26,8 @@ func TestReadFileToolReadsLineRanges(t *testing.T) {
}
for _, want := range []string{
"File: notes.txt (lines 2-3 of 4)",
"2→beta",
"3→gamma",
"2 | beta",
"3 | gamma",
} {
if !strings.Contains(result.Output, want) {
t.Fatalf("expected output to contain %q, got %q", want, result.Output)
@@ -54,7 +54,7 @@ func TestReadFileToolReadsCanonicalLineRange(t *testing.T) {
if result.Truncated {
t.Fatalf("canonical limit defines the requested range and must not mark it truncated: %#v", result.Meta)
}
if !strings.Contains(result.Output, "2→beta") || !strings.Contains(result.Output, "3→gamma") {
if !strings.Contains(result.Output, "2 | beta") || !strings.Contains(result.Output, "3 | gamma") {
t.Fatalf("canonical range returned the wrong lines: %q", result.Output)
}
if strings.Contains(result.Output, "alpha") || strings.Contains(result.Output, "delta") {
@@ -94,7 +94,7 @@ func TestReadFileToolMixedLegacyRangesPreferLines(t *testing.T) {
if result.Status != StatusOK {
t.Fatalf("mixed legacy range should recover, got %s: %s", result.Status, result.Output)
}
if !strings.Contains(result.Output, "2→beta") || !strings.Contains(result.Output, "3→gamma") {
if !strings.Contains(result.Output, "2 | beta") || !strings.Contains(result.Output, "3 | gamma") {
t.Fatalf("mixed legacy range did not prefer lines: %q", result.Output)
}
}
@@ -110,7 +110,7 @@ func TestReadFileToolCombinesLegacyStartWithCanonicalLimit(t *testing.T) {
if result.Status != StatusOK || result.Truncated {
t.Fatalf("mixed compatible range should be exact, got status=%s truncated=%v: %s", result.Status, result.Truncated, result.Output)
}
if !strings.Contains(result.Output, "2→beta") || !strings.Contains(result.Output, "3→gamma") {
if !strings.Contains(result.Output, "2 | beta") || !strings.Contains(result.Output, "3 | gamma") {
t.Fatalf("mixed compatible range returned the wrong lines: %q", result.Output)
}
if strings.Contains(result.Output, "alpha") || strings.Contains(result.Output, "delta") {
@@ -955,7 +955,7 @@ func TestReadFileToolRecoversBackwardsRange(t *testing.T) {
for _, want := range []string{
"end_line 2 was before start_line 3", // the note
"only line 3 was read",
"3→gamma",
"3 | gamma",
} {
if !strings.Contains(result.Output, want) {
t.Fatalf("expected %q in output, got %q", want, result.Output)
@@ -984,7 +984,7 @@ func TestReadFileToolValidRangeUnchangedByRecovery(t *testing.T) {
// Byte-for-byte: header, blank separator, the exact selected lines, and no
// recovery note or trailing content. An altered header/separator or a stray
// out-of-range line would fail this where a substring check would not.
const want = "File: notes.txt (lines 2-4 of 5)\n\n2→beta\n3→gamma\n4→delta"
const want = "File: notes.txt (lines 2-4 of 5)\n\n2 | beta\n3 | gamma\n4 | delta"
if result.Output != want {
t.Fatalf("valid-range output changed by the recovery path:\n got: %q\nwant: %q", result.Output, want)
}
+4 -6
View File
@@ -1,11 +1,9 @@
package tools
// ModelVisible reports whether a registered tool is advertised to an agent.
// Every registered tool is visible. edit_file (exact search/replace with a
// uniqueness guard, CRLF and fuzzy fallbacks) was hidden for a while in favour
// of apply_patch; head-to-head benchmarking showed that when a model misses the
// patch grammar it degrades to whole-file rewrites, so both edit tools are
// advertised and the prompt steers between them.
// edit_file remains implemented for compatibility with existing integrations,
// but its text-only replacement contract cannot safely disambiguate repeated
// source fragments. Agents use apply_patch for contextual existing-file edits.
func ModelVisible(tool Tool) bool {
return tool != nil
return tool != nil && tool.Name() != "edit_file"
}
+4 -15
View File
@@ -2,28 +2,17 @@ package tools
import "testing"
func TestToolSearchSurfacesBothEditTools(t *testing.T) {
func TestToolSearchDoesNotSurfaceLegacyStringReplacement(t *testing.T) {
registry := NewRegistry()
registry.Register(NewScopedEditFileTool(t.TempDir(), nil))
registry.Register(NewScopedApplyPatchTool(t.TempDir(), nil))
search := NewToolSearchTool(registry).(toolSearchTool)
eager := search.visibleEagerToolNames(nil, nil, "ask")
if !eager["edit_file"] {
t.Fatal("tool_search must describe edit_file as available to the model")
if eager["edit_file"] {
t.Fatal("tool_search must not describe edit_file as available to the model")
}
if !eager["apply_patch"] {
t.Fatal("tool_search must retain apply_patch as an available edit tool")
}
}
func TestModelVisibleAdvertisesEveryRegisteredTool(t *testing.T) {
if ModelVisible(nil) {
t.Fatal("a nil tool must not be visible")
}
for _, tool := range []Tool{NewScopedEditFileTool(t.TempDir(), nil), NewScopedApplyPatchTool(t.TempDir(), nil), NewScopedWriteFileTool(t.TempDir(), nil)} {
if !ModelVisible(tool) {
t.Fatalf("%s must be visible to the model", tool.Name())
}
t.Fatal("tool_search must retain apply_patch as the available edit tool")
}
}
-3
View File
@@ -9,9 +9,6 @@ import (
func TestUpdatePlanToolStoresAndFormatsPlan(t *testing.T) {
tool := NewUpdatePlanTool()
if description := strings.ToLower(tool.Description()); !strings.Contains(description, "skip it for bounded changes in one component") {
t.Fatalf("update_plan description must preserve the bounded-task fast path: %q", description)
}
result := tool.Run(context.Background(), map[string]any{
"plan": []any{
+9 -11
View File
@@ -9,14 +9,12 @@ import (
"io"
"os"
"strconv"
"strings"
"unicode/utf8"
)
const readFileByteChunkMax = 64 * 1024
// readFileLinePrefixSeparator follows the line number on every read_file line.
const readFileLinePrefixSeparator = "→"
type readFileTool struct {
baseTool
workspaceRoot string
@@ -33,7 +31,7 @@ func NewScopedReadFileTool(workspaceRoot string, scope PathScope) Tool {
return readFileTool{
baseTool: baseTool{
name: "read_file",
description: "Read exact file text, each line prefixed with its line number and →. Use directly for small files likely to be edited, or when comments, formatting, or line numbers matter. Prefer read_minified_file only for exploratory understanding of large or unfamiliar code. Use offset and limit for a line range.",
description: "Read exact file text with line numbers. Use for comments, formatting, or edits; prefer read_minified_file for initial code understanding. Use offset and limit for a line range.",
parameters: Schema{
Type: "object",
Properties: map[string]PropertySchema{
@@ -223,6 +221,7 @@ func renderReadFileRange(absolutePath string, relativePath string, total int, st
}
lastLine := startLine + selectedLines - 1
width := len(strconv.Itoa(lastLine))
header := fmt.Sprintf("File: %s (%d lines)", relativePath, total)
if startLine != 1 || endLine != total || maxLines > 0 {
header = fmt.Sprintf("File: %s (lines %d-%d of %d)", relativePath, startLine, lastLine, total)
@@ -245,7 +244,7 @@ func renderReadFileRange(absolutePath string, relativePath string, total int, st
budgetedOutput.WriteString("\n")
}
budgetedOutput.WriteString("\n")
if err := appendReadFileRange(budgetedOutput, absolutePath, startLine, selectedLines); err != nil {
if err := appendReadFileRange(budgetedOutput, absolutePath, startLine, selectedLines, width); err != nil {
return errorResult("Error reading file " + relativePath + ": " + err.Error())
}
if truncated {
@@ -358,7 +357,7 @@ func renderReadFileBytes(path, relativePath string, total, requestedStart, limit
return Result{Status: StatusOK, Output: output, Meta: map[string]string{"next_byte_offset": strconv.Itoa(end)}}, start, end
}
func appendReadFileRange(output *outputBudgetBuilder, path string, startLine int, selectedLines int) error {
func appendReadFileRange(output *outputBudgetBuilder, path string, startLine int, selectedLines int, width int) error {
file, err := os.Open(path)
if err != nil {
return err
@@ -381,11 +380,10 @@ func appendReadFileRange(output *outputBudgetBuilder, path string, startLine int
if emitted > 0 {
output.WriteString("\n")
}
// Compact "N→" prefix: the padded " N | " form cost ~19% of every
// read (about 1.5k tokens per 800-line file) and was re-sent on each
// later call; no consumer parses it, and models read "N→" natively.
output.WriteString(strconv.Itoa(lineNumber))
output.WriteString(readFileLinePrefixSeparator)
number := strconv.Itoa(lineNumber)
output.WriteString(strings.Repeat(" ", width-len(number)))
output.WriteString(number)
output.WriteString(" | ")
output.WriteString(string(trimLineBreak(raw, ended)))
emitted++
}
+9 -21
View File
@@ -28,7 +28,7 @@ func NewScopedReadMinifiedFileTool(workspaceRoot string, scope PathScope) Tool {
return readMinifiedFileTool{
baseTool: baseTool{
name: "read_minified_file",
description: "Read source code in a token-efficient, language-aware form. Use for exploratory understanding of large or unfamiliar source when no exact edit is planned; use read_file directly for likely edit targets. Do not immediately reread the same file exactly unless a new need for exact text or line numbers appears.",
description: "Read source code in a token-efficient, language-aware form. Prefer this for initial understanding; use read_file for exact text or line numbers.",
parameters: Schema{
Type: "object",
Properties: map[string]PropertySchema{
@@ -171,9 +171,14 @@ func selectSourceLines(content []byte, offset, limit int) sourceSelection {
}
func sourceLineCount(content []byte) int {
// One implementation for reader and writer: the tracker contract lives on
// trackedLineTotal, and this must never drift from it.
return trackedLineTotal(string(content))
if len(content) == 0 {
return 1
}
lines := strings.Count(string(content), "\n")
if content[len(content)-1] != '\n' {
lines++
}
return lines
}
// lineCount reports the number of newline-separated lines in s (an empty string
@@ -184,20 +189,3 @@ func lineCount(s string) int {
}
return strings.Count(s, "\n") + 1
}
// trackedLineTotal reports a file's line count the way read_file reports it
// to the FileTracker (newline-terminated lines, plus an unterminated last
// line; an empty file is one line). Writers must record this same number:
// RecordSeenRange resets every observation when the total changes, so a
// writer that recorded lineCount (one higher for a trailing newline) made the
// next partial read_file wipe whole-file knowledge and refuse the next edit.
func trackedLineTotal(s string) int {
if s == "" {
return 1
}
total := strings.Count(s, "\n")
if !strings.HasSuffix(s, "\n") {
total++
}
return total
}
-23
View File
@@ -8,29 +8,6 @@ import (
"testing"
)
func TestReadToolsDescribeExploratoryAndExactUseWithoutRedundantRereads(t *testing.T) {
root := t.TempDir()
minified := strings.ToLower(NewScopedReadMinifiedFileTool(root, nil).Description())
for _, want := range []string{
"exploratory understanding",
"use read_file directly for likely edit targets",
"unless a new need for exact text or line numbers appears",
} {
if !strings.Contains(minified, want) {
t.Fatalf("read_minified_file description missing %q: %s", want, minified)
}
}
exact := strings.ToLower(NewScopedReadFileTool(root, nil).Description())
for _, want := range []string{
"use directly for small files likely to be edited",
"prefer read_minified_file only for exploratory understanding",
} {
if !strings.Contains(exact, want) {
t.Fatalf("read_file description missing %q: %s", want, exact)
}
}
}
func TestReadMinifiedFileStripsCommentsAndLineNumbers(t *testing.T) {
dir := t.TempDir()
src := "package demo\n\nimport \"fmt\"\n\n// secret doc comment\nfunc F() { fmt.Println(\"x\") }\n"
+2 -59
View File
@@ -5,7 +5,6 @@ import (
"fmt"
"sort"
"strings"
"sync"
"github.com/Gitlawb/zero/internal/redaction"
"github.com/Gitlawb/zero/internal/sandbox"
@@ -13,17 +12,7 @@ import (
)
type Registry struct {
mu sync.RWMutex
tools map[string]Tool
generation uint64
}
// RegistrySnapshot is one immutable registry generation. Tools is a detached,
// name-sorted slice: later registrations cannot change the snapshot observed by
// an in-flight agent turn.
type RegistrySnapshot struct {
Tools []Tool
Generation uint64
tools map[string]Tool
}
type RunOptions struct {
@@ -118,69 +107,23 @@ func IsDeferralEligible(t Tool) bool {
}
func (registry *Registry) Register(tool Tool) {
registry.mu.Lock()
defer registry.mu.Unlock()
registry.tools[tool.Name()] = tool
registry.generation++
}
// RegisterBatch publishes a group of tools as one registry generation. Readers
// observe either the complete prior generation or the complete new one.
func (registry *Registry) RegisterBatch(batch []Tool) {
if len(batch) == 0 {
return
}
registry.mu.Lock()
defer registry.mu.Unlock()
for _, tool := range batch {
registry.tools[tool.Name()] = tool
}
registry.generation++
}
func (registry *Registry) Get(name string) (Tool, bool) {
registry.mu.RLock()
defer registry.mu.RUnlock()
tool, ok := registry.tools[name]
return tool, ok
}
func (registry *Registry) All() []Tool {
return registry.Snapshot().Tools
}
// Snapshot returns a stable, sorted view of the registry at one generation.
func (registry *Registry) Snapshot() RegistrySnapshot {
registry.mu.RLock()
tools := make([]Tool, 0, len(registry.tools))
for _, tool := range registry.tools {
tools = append(tools, tool)
}
generation := registry.generation
registry.mu.RUnlock()
sort.Slice(tools, func(left, right int) bool {
return tools[left].Name() < tools[right].Name()
})
return RegistrySnapshot{Tools: tools, Generation: generation}
}
// Clone returns an independent registry containing exactly one source snapshot.
func (registry *Registry) Clone() *Registry {
clone := NewRegistry()
if registry == nil {
return clone
}
snapshot := registry.Snapshot()
clonedTools := make([]Tool, 0, len(snapshot.Tools))
for _, tool := range snapshot.Tools {
if cloner, ok := tool.(interface{ cloneForRegistry(*Registry) Tool }); ok {
tool = cloner.cloneForRegistry(clone)
}
clonedTools = append(clonedTools, tool)
}
clone.RegisterBatch(clonedTools)
clone.generation = snapshot.Generation
return clone
return tools
}
func (registry *Registry) Run(ctx context.Context, name string, args map[string]any) Result {
-77
View File
@@ -2,12 +2,10 @@ package tools
import (
"context"
"fmt"
"os"
"path/filepath"
"slices"
"strings"
"sync"
"testing"
"github.com/Gitlawb/zero/internal/sandbox"
@@ -97,81 +95,6 @@ func TestRegistryAllReturnsToolsSortedByName(t *testing.T) {
}
}
func TestRegistrySnapshotAndCloneStayOnOneGeneration(t *testing.T) {
registry := NewRegistry()
registry.Register(fakePlainTool{baseTool: baseTool{name: "alpha", description: "alpha"}})
first := registry.Snapshot()
registry.RegisterBatch([]Tool{
fakePlainTool{baseTool: baseTool{name: "charlie", description: "charlie"}},
fakePlainTool{baseTool: baseTool{name: "bravo", description: "bravo"}},
})
second := registry.Snapshot()
clone := registry.Clone()
registry.Register(fakePlainTool{baseTool: baseTool{name: "delta", description: "delta"}})
if got := toolNames(first.Tools); !slices.Equal(got, []string{"alpha"}) {
t.Fatalf("first snapshot changed after registration: %v", got)
}
if second.Generation != first.Generation+1 {
t.Fatalf("batch generation = %d, want %d", second.Generation, first.Generation+1)
}
if got := toolNames(second.Tools); !slices.Equal(got, []string{"alpha", "bravo", "charlie"}) {
t.Fatalf("second snapshot tools = %v", got)
}
if got := toolNames(clone.All()); !slices.Equal(got, toolNames(second.Tools)) {
t.Fatalf("clone changed with source registry: %v", got)
}
if clone.Snapshot().Generation != second.Generation {
t.Fatalf("clone generation = %d, want %d", clone.Snapshot().Generation, second.Generation)
}
}
func TestRegistryBatchIsAtomicForConcurrentReaders(t *testing.T) {
registry := NewRegistry()
registry.Register(fakePlainTool{baseTool: baseTool{name: "base", description: "base"}})
batch := make([]Tool, 128)
for index := range batch {
name := fmt.Sprintf("tool_%03d", index)
batch[index] = fakePlainTool{baseTool: baseTool{name: name, description: name}}
}
start := make(chan struct{})
done := make(chan struct{})
var readers sync.WaitGroup
for range 4 {
readers.Add(1)
go func() {
defer readers.Done()
<-start
for {
count := len(registry.Snapshot().Tools)
if count != 1 && count != len(batch)+1 {
t.Errorf("reader observed partial batch with %d tools", count)
return
}
select {
case <-done:
return
default:
}
}
}()
}
close(start)
registry.RegisterBatch(batch)
close(done)
readers.Wait()
}
func toolNames(toolset []Tool) []string {
names := make([]string, 0, len(toolset))
for _, tool := range toolset {
names = append(names, tool.Name())
}
return names
}
func TestCoreNetworkToolsExposeSafetyMetadata(t *testing.T) {
// web_search is only registered when a backend is configured.
t.Setenv("ZERO_WEBSEARCH_BASE_URL", "https://search.example/api")
+23 -219
View File
@@ -6,11 +6,9 @@ import (
"io"
"os"
"path/filepath"
"regexp"
"strings"
"github.com/Gitlawb/zero/internal/pathjail"
"github.com/Gitlawb/zero/internal/sandbox"
)
const (
@@ -30,9 +28,6 @@ const (
structuredPatchAdd structuredPatchKind = iota
structuredPatchDelete
structuredPatchUpdate
// structuredPatchCopy creates movePath from path's content (plus any
// hunks) and keeps path; produced by a unified diff's "copy from/to".
structuredPatchCopy
)
type structuredPatchOperation struct {
@@ -42,29 +37,8 @@ type structuredPatchOperation struct {
contents string
chunks []structuredPatchChunk
line int
// eofNewline lets a unified diff's "\ No newline at end of file" marker
// force the trailing-newline state of the result; structured patches keep
// the file's existing state.
eofNewline eofNewlineMode
// verifyDelete marks a deletion that carries the expected old content
// (a unified diff's "+++ /dev/null" hunks, or git's header-only form for
// an empty file): the removed lines must equal the current file byte for
// byte, otherwise the deletion is stale and refused. A structured
// "*** Delete File" means "delete this path" and has no chunks.
verifyDelete bool
// oldNoNewline records a unified diff's "\ No newline at end of file"
// on the removed side, i.e. the old content did not end with a newline.
oldNoNewline bool
}
type eofNewlineMode uint8
const (
eofNewlineKeep eofNewlineMode = iota
eofNewlinePresent
eofNewlineAbsent
)
type structuredPatchChunk struct {
context string
hasContext bool
@@ -72,11 +46,6 @@ type structuredPatchChunk struct {
new []string
newSourceOffsets []int
endOfFile bool
// hint is the 0-based line the hunk is expected at (from a unified diff's
// "@@ -a,b" range). When the expected lines match there it is used as-is;
// otherwise the hunk is located by content like a structured hunk.
hint int
hasHint bool
}
type structuredPatchTarget struct {
@@ -94,36 +63,8 @@ type structuredPatchChange struct {
mode os.FileMode
}
// unifiedHunkRangePattern recognises a unified-diff range header ("-12,4 +12,6",
// optionally followed by "@@ heading") written inside a structured hunk marker.
// The line numbers carry no information for the context-anchored grammar, so
// the range is dropped and only an explicit heading is kept as the anchor.
var unifiedHunkRangePattern = regexp.MustCompile(`^-\d+(?:,\d+)?\s+\+\d+(?:,\d+)?\s*(?:@@\s*(.*))?$`)
const structuredPatchFormatHint = ` (format: "*** Begin Patch", then "*** Update File: path" / "*** Add File: path" / "*** Delete File: path" sections whose hunks start with "@@" and use " " context, "-" removed and "+" added lines, then "*** End Patch")`
// structuredPatchMarker classifies a line as the "begin" or "end" marker of a
// structured patch, or "" when it is neither. The classifier is owned by the
// sandbox package so the boundary check and the tool accept exactly the same
// spellings; a strict byte-equal match here once made every patch from some
// models fail on line 1 and pushed them into whole-file rewrites.
func structuredPatchMarker(line string) string {
return sandbox.StructuredPatchMarker(line)
}
// structuredHunkAnchor normalises the text after a hunk's "@@ " marker. A
// unified-diff range is dropped (its optional heading survives as the anchor);
// anything else is used verbatim.
func structuredHunkAnchor(context string) string {
context = strings.TrimSpace(context)
if match := unifiedHunkRangePattern.FindStringSubmatch(context); match != nil {
return strings.TrimSpace(match[1])
}
return context
}
func isStructuredPatch(patch string) bool {
return sandbox.IsStructuredPatch(patch)
return strings.HasPrefix(strings.TrimSpace(strings.TrimPrefix(patch, "\ufeff")), structuredPatchBegin)
}
func (tool applyPatchTool) runStructuredPatch(applyRoot, relativeRoot, patch string, options RunOptions) Result {
@@ -131,14 +72,6 @@ func (tool applyPatchTool) runStructuredPatch(applyRoot, relativeRoot, patch str
if err != nil {
return errorResult("Error applying patch: " + err.Error())
}
return applyPatchOperations(applyRoot, relativeRoot, operations, options)
}
// applyPatchOperations applies parsed operations (from either patch format)
// through an opened workspace root: every stat, read, create and write is
// descriptor-relative and refuses to follow a link out of the root, so there
// is no pathname check-to-use window between validation and write.
func applyPatchOperations(applyRoot, relativeRoot string, operations []structuredPatchOperation, options RunOptions) Result {
workspace, err := os.OpenRoot(applyRoot)
if err != nil {
return errorResult("Error applying patch: " + err.Error())
@@ -151,7 +84,7 @@ func applyPatchOperations(applyRoot, relativeRoot string, operations []structure
wholeBefore := make(map[string]bool, len(changes))
if options.FileTracker != nil {
for _, change := range changes {
if change.kind == structuredPatchUpdate || change.kind == structuredPatchCopy {
if change.kind == structuredPatchUpdate {
wholeBefore[change.from.absolute] = options.FileTracker.SeenWhole(change.from.absolute)
}
}
@@ -170,10 +103,6 @@ func applyPatchOperations(applyRoot, relativeRoot string, operations []structure
wasWhole := wholeBefore[change.from.absolute]
options.FileTracker.Forget(change.from.absolute)
recordStructuredPatchFile(options.FileTracker, change.to.absolute, false, wasWhole)
case structuredPatchCopy:
// The source is untouched; the destination inherits only what the
// model had actually seen of the source.
recordStructuredPatchFile(options.FileTracker, change.to.absolute, true, wholeBefore[change.from.absolute])
}
}
@@ -193,11 +122,11 @@ func parseStructuredPatch(patch string) ([]structuredPatchOperation, error) {
return nil, fmt.Errorf("structured patch is empty")
}
lines := strings.Split(normalized, "\n")
if structuredPatchMarker(lines[0]) != "begin" {
return nil, fmt.Errorf("the first line of a structured patch must be %q%s", structuredPatchBegin, structuredPatchFormatHint)
if strings.TrimSpace(lines[0]) != structuredPatchBegin {
return nil, fmt.Errorf("the first line of a structured patch must be %q", structuredPatchBegin)
}
if structuredPatchMarker(lines[len(lines)-1]) != "end" {
return nil, fmt.Errorf("the last line of a structured patch must be %q%s", structuredPatchEnd, structuredPatchFormatHint)
if strings.TrimSpace(lines[len(lines)-1]) != structuredPatchEnd {
return nil, fmt.Errorf("the last line of a structured patch must be %q", structuredPatchEnd)
}
var operations []structuredPatchOperation
@@ -262,12 +191,8 @@ func parseStructuredPatch(patch string) ([]structuredPatchOperation, error) {
}
if trimmed == "@@" || strings.HasPrefix(trimmed, structuredHunkContext) {
chunk := structuredPatchChunk{}
rest := ""
if trimmed != "@@" {
rest = strings.TrimPrefix(trimmed, structuredHunkContext)
}
if anchor := structuredHunkAnchor(rest); anchor != "" {
chunk.context = anchor
chunk.context = strings.TrimPrefix(trimmed, structuredHunkContext)
chunk.hasContext = true
}
op.chunks = append(op.chunks, chunk)
@@ -323,7 +248,7 @@ func structuredPatchPath(header, prefix string, line int) (string, error) {
}
func isStructuredPatchHeader(line string) bool {
return structuredPatchMarker(line) == "end" || strings.HasPrefix(line, structuredAddFile) || strings.HasPrefix(line, structuredDeleteFile) || strings.HasPrefix(line, structuredUpdateFile)
return line == structuredPatchEnd || strings.HasPrefix(line, structuredAddFile) || strings.HasPrefix(line, structuredDeleteFile) || strings.HasPrefix(line, structuredUpdateFile)
}
func structuredPatchLineError(line int, message string) error {
@@ -393,7 +318,7 @@ func planStructuredPatch(root *os.Root, operations []structuredPatchOperation, t
return nil, err
}
change.after = operation.contents
case structuredPatchDelete, structuredPatchUpdate, structuredPatchCopy:
case structuredPatchDelete, structuredPatchUpdate:
info, err := root.Stat(from.relative)
if err != nil {
return nil, fmt.Errorf("stating %s: %w", from.relative, err)
@@ -407,20 +332,12 @@ func planStructuredPatch(root *os.Root, operations []structuredPatchOperation, t
return nil, fmt.Errorf("%s", fileConflictMessage(from.relative))
}
change.before = string(content)
if operation.kind == structuredPatchDelete && operation.verifyDelete {
if err := verifyUnifiedDeletion(change.before, from.relative, operation); err != nil {
return nil, err
}
}
if operation.kind == structuredPatchUpdate || operation.kind == structuredPatchCopy {
if operation.kind == structuredPatchUpdate {
updated, err := applyStructuredPatchUpdate(change.before, from.relative, operation.chunks)
if err != nil {
return nil, err
}
change.after = applyEOFNewline(updated, operation.eofNewline)
if operation.kind == structuredPatchCopy && from.absolute == to.absolute {
return nil, fmt.Errorf("cannot copy %s onto itself", from.relative)
}
change.after = updated
if from.absolute != to.absolute {
if _, err := root.Lstat(to.relative); err == nil {
return nil, fmt.Errorf("cannot move %s to %s because the destination already exists", from.relative, to.relative)
@@ -436,14 +353,10 @@ func planStructuredPatch(root *os.Root, operations []structuredPatchOperation, t
}
func resolveStructuredPatchTarget(root, path string) (structuredPatchTarget, error) {
// Relative traversal is rejected outright. An absolute path is allowed when
// it resolves inside the apply root (models routinely echo the absolute path
// they were shown by read_file); resolveWorkspaceTargetPath rejects anything
// that lands outside.
if path == ".." || strings.HasPrefix(filepath.ToSlash(path), "../") {
if filepath.IsAbs(path) || path == ".." || strings.HasPrefix(filepath.ToSlash(path), "../") {
return structuredPatchTarget{}, fmt.Errorf("patch path %q must stay inside the workspace", path)
}
absolute, relative, err := resolveWorkspaceTargetPath(root, normalizePatchPathForRoot(root, path))
absolute, relative, err := resolveWorkspaceTargetPath(root, path)
if err != nil {
return structuredPatchTarget{}, err
}
@@ -476,15 +389,9 @@ func applyStructuredPatchUpdate(content, path string, chunks []structuredPatchCh
if len(chunk.old) == 0 {
// A context anchor identifies the line immediately before a pure
// insertion. Only an explicit end-of-file hunk (or an unanchored
// insertion) belongs at the file end. A unified diff names the
// insertion point in its range instead; a zero-context hunk (as
// from `git diff -U0`) has nothing to verify the position against,
// so it is position-trusting by design — the same semantics git
// apply uses without fuzz. Hunks that carry context are verified.
// insertion) belongs at the file end.
start := lineIndex
if chunk.hasHint {
start = max(chunk.hint, lineIndex)
} else if chunk.endOfFile || !chunk.hasContext {
if chunk.endOfFile || !chunk.hasContext {
start = len(lines)
}
if start > len(lines) {
@@ -498,13 +405,7 @@ func applyStructuredPatchUpdate(content, path string, chunks []structuredPatchCh
lineIndex = start
continue
}
var index int
ambiguous := false
if chunk.hasHint && chunk.hint >= lineIndex && sequenceMatchesAt(lines, chunk.old, chunk.hint) {
index = chunk.hint
} else {
index, ambiguous = findStructuredPatchSequence(lines, chunk.old, lineIndex, chunk.endOfFile)
}
index, ambiguous := findStructuredPatchSequence(lines, chunk.old, lineIndex, chunk.endOfFile)
if ambiguous {
return "", fmt.Errorf("expected lines are ambiguous in %s; provide more surrounding context:\n%s", path, strings.Join(chunk.old, "\n"))
}
@@ -539,64 +440,6 @@ func applyStructuredPatchUpdate(content, path string, chunks []structuredPatchCh
return updated, nil
}
// verifyUnifiedDeletion checks, byte for byte, that the lines a unified
// deletion removes are exactly the file's current content. A deletion is not
// recoverable, so none of the whitespace tolerance used to locate update
// hunks applies here; only the file's own line-ending style is normalised.
// With no hunks (git's header-only form) the file must be empty.
func verifyUnifiedDeletion(current, path string, operation structuredPatchOperation) error {
var removed []string
for _, chunk := range operation.chunks {
if len(chunk.new) > 0 {
return fmt.Errorf("deletion of %s must not add lines", path)
}
removed = append(removed, chunk.old...)
}
expected := strings.Join(removed, "\n")
if len(removed) > 0 && !operation.oldNoNewline {
expected += "\n"
}
actual := current
if structuredPatchLineEnding(current) == "\r\n" {
actual = strings.ReplaceAll(current, "\r\n", "\n")
}
if actual != expected {
if len(removed) == 0 {
return fmt.Errorf("deletion of %s expects an empty file but it has content; include the removed lines in the patch", path)
}
return fmt.Errorf("deletion of %s does not match its current content; the removed lines must equal the whole file", path)
}
return nil
}
// sequenceMatchesAt reports whether wanted appears verbatim at lines[start:].
func sequenceMatchesAt(lines, wanted []string, start int) bool {
if start < 0 || start+len(wanted) > len(lines) {
return false
}
for offset, line := range wanted {
if lines[start+offset] != line {
return false
}
}
return true
}
// applyEOFNewline forces or strips a single trailing line ending when a
// unified diff's end-of-file marker asked for it.
func applyEOFNewline(content string, mode eofNewlineMode) string {
switch mode {
case eofNewlinePresent:
if content != "" && !strings.HasSuffix(content, "\n") {
return content + structuredPatchLineEnding(content)
}
case eofNewlineAbsent:
content = strings.TrimSuffix(content, "\n")
content = strings.TrimSuffix(content, "\r")
}
return content
}
func structuredPatchReplacement(chunk structuredPatchChunk, source []string, start int) ([]string, error) {
if len(chunk.newSourceOffsets) != len(chunk.new) {
return nil, fmt.Errorf("context mapping has %d entries for %d output lines", len(chunk.newSourceOffsets), len(chunk.new))
@@ -673,19 +516,14 @@ func findStructuredPatchSequence(lines, wanted []string, start int, endOfFile bo
}
func applyStructuredPatchChanges(root *os.Root, changes []structuredPatchChange, tracker *FileTracker) error {
// committed lists, in order, the paths whose change reached disk before a
// later change failed, so the caller (and the model) knows exactly which
// files now hold the patched content and which were never touched.
var committed []string
mutated := false
for _, change := range changes {
done, err := applyStructuredPatchChange(root, change)
if done {
committed = append(committed, structuredPatchChangePaths(change)...)
}
committed, err := applyStructuredPatchChange(root, change)
mutated = mutated || committed
if err != nil {
forgetStructuredPatchFiles(tracker, changes)
if len(committed) > 0 {
return fmt.Errorf("%w; patch was partially applied — already committed: %s; the remaining files are unchanged; re-read the committed files before retrying", err, strings.Join(committed, ", "))
if mutated {
return fmt.Errorf("%w; structured patch was partially applied; re-read affected files before retrying", err)
}
return err
}
@@ -693,19 +531,6 @@ func applyStructuredPatchChanges(root *os.Root, changes []structuredPatchChange,
return nil
}
// structuredPatchChangePaths names the workspace-relative paths a committed
// change touched: the destination, plus the source of a move or copy.
func structuredPatchChangePaths(change structuredPatchChange) []string {
if change.kind == structuredPatchDelete {
return []string{change.from.relative}
}
paths := []string{change.to.relative}
if change.from.absolute != change.to.absolute && change.from.relative != "" {
paths = append([]string{change.from.relative}, paths...)
}
return paths
}
func forgetStructuredPatchFiles(tracker *FileTracker, changes []structuredPatchChange) {
if tracker == nil {
return
@@ -718,35 +543,14 @@ func forgetStructuredPatchFiles(tracker *FileTracker, changes []structuredPatchC
}
}
// structuredPatchBeforeCommit, when set, runs before each change's pre-commit
// recheck. Tests use it to alter the filesystem between planning and commit
// deterministically; it is nil in production.
var structuredPatchBeforeCommit func(change structuredPatchChange)
func applyStructuredPatchChange(root *os.Root, change structuredPatchChange) (bool, error) {
if structuredPatchBeforeCommit != nil {
structuredPatchBeforeCommit(change)
}
// The planner validated the source against the bytes it read; re-read
// through the same root immediately before committing so a change made
// by another process in between is refused rather than overwritten or
// removed.
if change.kind != structuredPatchAdd {
current, err := root.ReadFile(change.from.relative)
if err != nil {
return false, fmt.Errorf("re-reading %s before commit: %w", change.from.relative, err)
}
if string(current) != change.before {
return false, fmt.Errorf("%s changed on disk between planning and commit; re-read it and retry", change.from.relative)
}
}
switch change.kind {
case structuredPatchDelete:
if err := root.Remove(change.from.relative); err != nil {
return false, fmt.Errorf("deleting %s: %w", change.from.relative, err)
}
return true, nil
case structuredPatchAdd, structuredPatchCopy:
case structuredPatchAdd:
return writeStructuredPatchFile(root, change.to, change.after, change.mode, true)
case structuredPatchUpdate:
moving := change.from.absolute != change.to.absolute
@@ -902,7 +706,7 @@ func recordStructuredPatchFile(tracker *FileTracker, absolute string, created, s
info, _ := os.Stat(absolute)
tracker.Record(absolute, content, info)
if seenWhole {
lines := trackedLineTotal(string(content))
lines := lineCount(string(content))
tracker.RecordSeenRange(absolute, 1, lines, lines)
}
if created {
-5
View File
@@ -58,11 +58,6 @@ func NewToolSearchTool(registry *Registry) Tool {
}
}
func (tool toolSearchTool) cloneForRegistry(registry *Registry) Tool {
tool.registry = registry
return tool
}
// Run satisfies the Tool interface; actual dispatch goes through RunWithOptions.
func (tool toolSearchTool) Run(ctx context.Context, args map[string]any) Result {
return tool.RunWithOptions(ctx, args, RunOptions{})
-31
View File
@@ -73,37 +73,6 @@ func newDeferredFixtureRegistry() *Registry {
return reg
}
func TestRegistryCloneKeepsToolSearchOnFrozenSnapshot(t *testing.T) {
live := newDeferredFixtureRegistry()
live.Register(NewToolSearchTool(live))
runRegistry := live.Clone()
live.Register(searchFakeTool{
name: "late_lookup",
description: "Registered after the run snapshot.",
parameters: Schema{Type: "object", AdditionalProperties: false},
})
result := runRegistry.Run(context.Background(), ToolSearchToolName, map[string]any{
"query": "select:late_lookup",
})
if result.Status != StatusOK {
t.Fatalf("tool_search status = %s: %s", result.Status, result.Output)
}
if loaded := result.Meta["load_tools"]; loaded != "" {
t.Fatalf("snapshot tool_search loaded late tool %q", loaded)
}
if direct := runRegistry.Run(context.Background(), "late_lookup", map[string]any{}); direct.Status != StatusError {
t.Fatalf("late tool unexpectedly executable in run snapshot: %#v", direct)
}
known := runRegistry.Run(context.Background(), ToolSearchToolName, map[string]any{
"query": "select:weather_lookup",
})
if loaded := known.Meta["load_tools"]; loaded != "weather_lookup" {
t.Fatalf("snapshot tool_search lost existing tool: load_tools=%q output=%q", loaded, known.Output)
}
}
func TestToolSearchExposesExpectedSafetyAndSchema(t *testing.T) {
tool := NewToolSearchTool(NewRegistry())
-59
View File
@@ -1,59 +0,0 @@
package tools
import (
"context"
"fmt"
"path/filepath"
"testing"
)
func TestTrackedLineTotalMatchesReadFileStats(t *testing.T) {
root := t.TempDir()
for name, content := range map[string]string{"trailing": "a\nb\nc\n", "unterminated": "a\nb\nc", "single": "x\n", "empty": "", "crlf": "a\r\nb\r\n"} {
path := filepath.Join(root, name+".txt")
writeTestFile(t, path, content)
stats, err := scanReadFileStats(path)
if err != nil {
t.Fatal(err)
}
if got := trackedLineTotal(content); got != stats.lines {
t.Fatalf("%s: trackedLineTotal = %d, read_file reports %d", name, got, stats.lines)
}
}
}
// A whole-file read, an edit, then a partial re-read must keep the file "seen
// whole": the writer and reader now agree on the line total, so the partial
// read accumulates instead of resetting and a later edit far from the re-read
// window is not refused.
func TestEditThenPartialReadKeepsWholeFileObservation(t *testing.T) {
root := t.TempDir()
path := filepath.Join(root, "notes.txt")
content := ""
for i := 1; i <= 40; i++ {
content += fmt.Sprintf("line %02d\n", i)
}
writeTestFile(t, path, content)
tracker := NewFileTracker()
registry := NewRegistry()
registry.Register(NewScopedReadFileTool(root, nil))
registry.Register(NewScopedEditFileTool(root, nil))
opts := RunOptions{PermissionGranted: true, FileTracker: tracker}
if r := registry.RunWithOptions(context.Background(), "read_file", map[string]any{"path": "notes.txt"}, opts); r.Status != StatusOK {
t.Fatalf("read: %s", r.Output)
}
if r := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{"path": "notes.txt", "old_string": "line 01\nline 02\n", "new_string": "line 01\nline two\n"}, opts); r.Status != StatusOK {
t.Fatalf("first edit: %s", r.Output)
}
if r := registry.RunWithOptions(context.Background(), "read_file", map[string]any{"path": "notes.txt", "limit": 5}, opts); r.Status != StatusOK {
t.Fatalf("partial read: %s", r.Output)
}
absolute, _ := filepath.EvalSymlinks(path)
if !tracker.SeenWhole(absolute) {
t.Fatal("partial re-read of an unchanged file must not discard whole-file knowledge")
}
if r := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{"path": "notes.txt", "old_string": "line 39\nline 40\n", "new_string": "line 39\nline forty\n"}, opts); r.Status != StatusOK {
t.Fatalf("edit far from the re-read window must succeed: %s", r.Output)
}
}
-8
View File
@@ -211,14 +211,6 @@ type Tool interface {
Run(ctx context.Context, args map[string]any) Result
}
// IsBuiltInApplyPatch reports whether tool is Zero's scoped patch tool. The
// marker method is package-private so external and MCP tools cannot claim this
// identity merely by sharing the public name "apply_patch".
func IsBuiltInApplyPatch(tool Tool) bool {
_, ok := tool.(interface{ isBuiltInApplyPatch() })
return ok
}
// ArgsPermissioner is an optional interface a Tool can implement to refine its
// permission for a SPECIFIC call based on its arguments. When a tool implements
// it, the agent loop consults PermissionForArgs(args) instead of the static
-393
View File
@@ -1,393 +0,0 @@
package tools
import (
"fmt"
"strings"
)
// parseUnifiedPatch converts a unified diff into the same operations a
// structured patch produces, so both formats are applied by Zero itself
// through an opened workspace root (descriptor-relative, no-follow writes)
// instead of handing pathnames to git apply after validation.
//
// Supported: file modifications, creations (--- /dev/null) and deletions
// (+++ /dev/null, verified against the content the hunks expect to remove),
// multiple hunks per file, "\ No newline at end of file"
// markers, a/ and b/ prefixes, git C-quoted paths, CRLF patches, and git's
// "rename from/to" and "copy from/to" headers (with or without hunks). Headers
// git emits around a hunk (diff --git, index, mode lines) are skipped; binary
// patches are rejected.
func parseUnifiedPatch(patch string) ([]structuredPatchOperation, error) {
normalized := strings.TrimPrefix(strings.ReplaceAll(patch, "\r\n", "\n"), "\ufeff")
lines := strings.Split(normalized, "\n")
// The element after a final "\n" is the line terminator, not an empty
// context line; counting it would let a truncated hunk look complete.
if len(lines) > 0 && lines[len(lines)-1] == "" {
lines = lines[:len(lines)-1]
}
var operations []structuredPatchOperation
var current *structuredPatchOperation
var chunk *structuredPatchChunk
var added []string // collected "+" lines for a file creation
oldPath, newPath := "", ""
pendingFrom, pendingKind := "", structuredPatchUpdate
// git's header-only forms: "deleted file mode" / "new file mode" after a
// "diff --git" line with no ---/+++ pair describe an empty file.
diffPath, headerOnly := "", byte(0) // headerOnly is 'd' (deleted) or 'n' (new)
oldRemaining, newRemaining := 0, 0
inHunk := false
lastSide := byte(0) // '+' or '-' for the most recent content line
finish := func() error {
if current == nil {
return nil
}
switch current.kind {
case structuredPatchAdd:
if len(added) > 0 {
current.contents = strings.Join(added, "\n")
if current.eofNewline != eofNewlineAbsent {
current.contents += "\n"
}
}
case structuredPatchUpdate, structuredPatchCopy:
if !allStructuredPatchChunksHaveContent(current.chunks) {
return fmt.Errorf("unified diff for %s has an empty hunk", current.path)
}
if len(current.chunks) == 0 && current.movePath == "" {
return fmt.Errorf("unified diff for %s has no hunk lines", current.path)
}
case structuredPatchDelete:
if len(current.chunks) == 0 || !allStructuredPatchChunksHaveContent(current.chunks) {
return fmt.Errorf("unified deletion of %s must include the hunk with the content being removed", current.path)
}
}
operations = append(operations, *current)
current, chunk, added = nil, nil, nil
return nil
}
flushHeaderOnly := func(line int) error {
if headerOnly == 0 {
return nil
}
if diffPath == "" {
return fmt.Errorf("invalid unified diff at line %d: file mode header without a diff --git path", line)
}
if err := finish(); err != nil {
return err
}
switch headerOnly {
case 'd':
operations = append(operations, structuredPatchOperation{kind: structuredPatchDelete, path: diffPath, line: line, verifyDelete: true})
case 'n':
operations = append(operations, structuredPatchOperation{kind: structuredPatchAdd, path: diffPath, line: line})
}
headerOnly, diffPath = 0, ""
return nil
}
startFile := func(line int) error {
// A ---/+++ pair means the file has hunks; it is not header-only.
headerOnly, diffPath = 0, ""
if oldPath == "" || newPath == "" {
return fmt.Errorf("invalid unified diff at line %d: hunk before a ---/+++ header pair", line)
}
// A ---/+++ pair after a rename/copy header names the same files; keep
// accumulating that operation's hunks instead of starting a new one.
if current != nil && current.movePath != "" && current.path == oldPath && (current.movePath == newPath || oldPath == newPath) {
return nil
}
if err := finish(); err != nil {
return err
}
op := structuredPatchOperation{line: line}
switch {
case oldPath == "/dev/null" && newPath == "/dev/null":
return fmt.Errorf("invalid unified diff at line %d: both sides are /dev/null", line)
case oldPath == "/dev/null":
op.kind, op.path = structuredPatchAdd, newPath
case newPath == "/dev/null":
// A unified deletion states the content it expects to remove; keep
// its hunks so the planner verifies them before deleting.
op.kind, op.path, op.verifyDelete = structuredPatchDelete, oldPath, true
default:
op.kind, op.path = structuredPatchUpdate, oldPath
if newPath != oldPath {
op.movePath = newPath
}
}
current = &op
return nil
}
for index, raw := range lines {
lineNumber := index + 1
if inHunk && (oldRemaining > 0 || newRemaining > 0) {
// A "--- " line followed by "+++ " and then a "@@" hunk header is
// the next file's header pair, not a removed "-- x" line followed by
// an added "++ y" line: an over-declared count would otherwise
// swallow both and the hunk would later fail against the wrong
// file. The "@@" requirement keeps a genuine adjacent pair intact.
if strings.HasPrefix(raw, "--- ") && index+2 < len(lines) && strings.HasPrefix(lines[index+1], "+++ ") && strings.HasPrefix(lines[index+2], "@@") {
return nil, fmt.Errorf("invalid unified diff at line %d: hunk ended before its declared line counts", lineNumber)
}
switch {
case strings.HasPrefix(raw, "\\"):
// "\ No newline at end of file" qualifies the previous line's side.
if current != nil && lastSide == '+' {
current.eofNewline = eofNewlineAbsent
} else if current != nil && lastSide == '-' {
current.oldNoNewline = true
if current.eofNewline == eofNewlineKeep {
current.eofNewline = eofNewlinePresent
}
}
continue
case strings.HasPrefix(raw, "-"):
oldRemaining--
lastSide = '-'
if chunk != nil {
chunk.old = append(chunk.old, raw[1:])
}
case strings.HasPrefix(raw, "+"):
newRemaining--
lastSide = '+'
if current != nil && current.kind == structuredPatchAdd {
added = append(added, raw[1:])
} else if chunk != nil {
chunk.new = append(chunk.new, raw[1:])
chunk.newSourceOffsets = append(chunk.newSourceOffsets, -1)
}
default:
content := raw
if strings.HasPrefix(raw, " ") {
content = raw[1:]
} else if raw != "" {
return nil, fmt.Errorf("invalid unified diff at line %d: hunk lines must start with ' ', '+', '-' or '\\'", lineNumber)
}
oldRemaining--
newRemaining--
lastSide = ' '
if chunk != nil {
offset := len(chunk.old)
chunk.old = append(chunk.old, content)
chunk.new = append(chunk.new, content)
chunk.newSourceOffsets = append(chunk.newSourceOffsets, offset)
}
}
continue
}
inHunk = false
trimmed := strings.TrimSpace(raw)
if oldRemaining > 0 || newRemaining > 0 {
// The hunk ended before its declared counts were consumed. Report
// it here rather than absorbing the next file's ---/+++ headers as
// "-"/"+" content and failing later against the wrong file.
return nil, fmt.Errorf("invalid unified diff at line %d: hunk ended before its declared line counts", lineNumber)
}
switch {
case trimmed == "":
continue
case strings.HasPrefix(raw, "\\"):
// "\ No newline at end of file" directly after a hunk's last line.
if current != nil && lastSide == '+' {
current.eofNewline = eofNewlineAbsent
} else if current != nil && lastSide == '-' {
current.oldNoNewline = true
if current.eofNewline == eofNewlineKeep {
current.eofNewline = eofNewlinePresent
}
}
continue
case strings.HasPrefix(raw, "diff --git "):
if err := flushHeaderOnly(lineNumber); err != nil {
return nil, err
}
diffPath = diffGitNewPath(raw)
continue
case strings.HasPrefix(raw, "deleted file mode "):
headerOnly = 'd'
continue
case strings.HasPrefix(raw, "new file mode "):
headerOnly = 'n'
continue
case strings.HasPrefix(raw, "index "), strings.HasPrefix(raw, "old mode "), strings.HasPrefix(raw, "new mode "):
continue
case strings.HasPrefix(raw, "similarity index "), strings.HasPrefix(raw, "dissimilarity index "):
continue
case strings.HasPrefix(raw, "rename from "), strings.HasPrefix(raw, "copy from "):
pendingFrom = strings.TrimSpace(unquoteGitPath(strings.TrimPrefix(strings.TrimPrefix(raw, "rename from "), "copy from ")))
pendingKind = structuredPatchUpdate
if strings.HasPrefix(raw, "copy from ") {
pendingKind = structuredPatchCopy
}
if pendingFrom == "" {
return nil, fmt.Errorf("invalid unified diff at line %d: missing source path", lineNumber)
}
case strings.HasPrefix(raw, "rename to "), strings.HasPrefix(raw, "copy to "):
to := strings.TrimSpace(unquoteGitPath(strings.TrimPrefix(strings.TrimPrefix(raw, "rename to "), "copy to ")))
if pendingFrom == "" || to == "" {
return nil, fmt.Errorf("invalid unified diff at line %d: rename/copy destination without a source", lineNumber)
}
if err := finish(); err != nil {
return nil, err
}
current = &structuredPatchOperation{kind: pendingKind, path: pendingFrom, movePath: to, line: lineNumber}
oldPath, newPath, pendingFrom = pendingFrom, to, ""
case strings.HasPrefix(raw, "Binary files "), strings.HasPrefix(raw, "GIT binary patch"):
return nil, fmt.Errorf("invalid unified diff at line %d: binary patches are not supported", lineNumber)
case strings.HasPrefix(raw, "--- "):
oldPath = stripPatchPrefix(patchFileHeaderPath(raw))
newPath = ""
if oldPath == "" {
return nil, fmt.Errorf("invalid unified diff at line %d: missing path in --- header", lineNumber)
}
case strings.HasPrefix(raw, "+++ "):
newPath = stripPatchPrefix(patchFileHeaderPath(raw))
if newPath == "" {
return nil, fmt.Errorf("invalid unified diff at line %d: missing path in +++ header", lineNumber)
}
if err := startFile(lineNumber); err != nil {
return nil, err
}
case strings.HasPrefix(raw, "@@"):
if current == nil {
return nil, fmt.Errorf("invalid unified diff at line %d: hunk before a ---/+++ header pair", lineNumber)
}
oldStart, oldCount, newCount, ok := parseHunkRange(raw)
if !ok {
return nil, fmt.Errorf("invalid unified diff at line %d: malformed hunk header", lineNumber)
}
// "\ No newline at end of file" describes the end of the file, so it
// may only follow a file's final hunk; a marker before another hunk
// would otherwise leak into that later hunk's result.
if current.eofNewline != eofNewlineKeep {
return nil, fmt.Errorf("invalid unified diff at line %d: \"\\ No newline at end of file\" must follow the last hunk of a file", lineNumber)
}
oldRemaining, newRemaining = oldCount, newCount
inHunk = oldRemaining > 0 || newRemaining > 0
lastSide = 0
if current.kind != structuredPatchAdd {
next := structuredPatchChunk{hasHint: true, hint: oldStart - 1}
if oldCount == 0 {
// A pure insertion's range names the line after which the
// new lines go, so the insertion index is oldStart itself.
next.hint = oldStart
}
if next.hint < 0 {
next.hint = 0
}
current.chunks = append(current.chunks, next)
chunk = &current.chunks[len(current.chunks)-1]
}
default:
return nil, fmt.Errorf("invalid unified diff at line %d: unexpected %q", lineNumber, trimmed)
}
}
// Strict at end of input on purpose: a hunk cut short (typically a model
// response that stopped mid-patch) would otherwise apply as a partial
// change — removals without their replacements — which is worse than a
// clear error and a retry. Miscounted ranges with complete content are
// still tolerated at every other boundary.
if inHunk && (oldRemaining > 0 || newRemaining > 0) {
return nil, fmt.Errorf("invalid unified diff at line %d: hunk ended before its declared line counts", len(lines))
}
if err := flushHeaderOnly(len(lines)); err != nil {
return nil, err
}
if err := finish(); err != nil {
return nil, err
}
if len(operations) == 0 {
return nil, fmt.Errorf("unified diff contains no file changes")
}
return operations, nil
}
// diffGitNewPath returns the post-image path named by a "diff --git a/X b/Y"
// line, handling git's C-quoted form; "" when it cannot be determined.
func diffGitNewPath(line string) string {
rest := strings.TrimSpace(strings.TrimPrefix(line, "diff --git "))
if strings.HasPrefix(rest, "\"") {
// Two quoted tokens: skip the first, unquote the second.
end := strings.Index(rest[1:], "\"")
for end > 0 && rest[end] == '\\' {
next := strings.Index(rest[end+2:], "\"")
if next < 0 {
return ""
}
end += next + 2
}
if end < 0 || end+2 > len(rest) {
return ""
}
rest = strings.TrimSpace(rest[end+2:])
return stripPatchPrefix(unquoteGitPath(rest))
}
fields := strings.Fields(rest)
if len(fields) < 2 {
return ""
}
last := fields[len(fields)-1]
if strings.HasPrefix(last, "\"") {
return stripPatchPrefix(unquoteGitPath(last))
}
return stripPatchPrefix(last)
}
// parseHunkRange reads "@@ -a[,b] +c[,d] @@" and returns a, b and d; a missing
// count means 1 per unified-diff convention.
func parseHunkRange(line string) (oldStart, oldCount, newCount int, ok bool) {
_, rest, found := strings.Cut(line, "@@")
if !found {
return 0, 0, 0, false
}
rangeSection := rest
if before, _, found := strings.Cut(rest, "@@"); found {
rangeSection = before
}
fields := strings.Fields(rangeSection)
if len(fields) != 2 || !strings.HasPrefix(fields[0], "-") || !strings.HasPrefix(fields[1], "+") {
return 0, 0, 0, false
}
parse := func(spec string) (int, int, bool) {
startText, countText, hasCount := strings.Cut(spec, ",")
start, err := parseNonNegativeInt(startText)
if err != nil {
return 0, 0, false
}
count := 1
if hasCount {
if count, err = parseNonNegativeInt(countText); err != nil {
return 0, 0, false
}
}
return start, count, true
}
oldStart, oldCount, ok = parse(fields[0][1:])
if !ok {
return 0, 0, 0, false
}
_, newCount, ok = parse(fields[1][1:])
if !ok {
return 0, 0, 0, false
}
return oldStart, oldCount, newCount, true
}
func parseNonNegativeInt(text string) (int, error) {
if text == "" {
return 0, fmt.Errorf("empty number")
}
value := 0
for _, r := range text {
if r < '0' || r > '9' {
return 0, fmt.Errorf("not a number: %q", text)
}
value = value*10 + int(r-'0')
if value > 1<<30 {
return 0, fmt.Errorf("number too large: %q", text)
}
}
return value, nil
}
+1 -1
View File
@@ -26,7 +26,7 @@ func NewUpdatePlanTool() *updatePlanTool {
return &updatePlanTool{
baseTool: baseTool{
name: "update_plan",
description: "Create or update the in-memory plan for work spanning multiple components or workstreams, sequencing uncertainty, or sustained tracking across many tool calls. Skip it for bounded changes in one component, simple lookups, explanations, and code navigation. " +
description: "Create or update the in-memory plan for implementation or investigation with at least three meaningful dependent steps. Do not use it for simple lookups, explanations, or code navigation. " +
"Pass the full ordered list of steps each call; it replaces the previous plan. " +
"Each item needs a `content` string; `status` defaults to \"pending\" and `id` is " +
"auto-numbered, so you only need to supply `content` (and `status` as the task progresses). " +
+1 -1
View File
@@ -120,7 +120,7 @@ func (tool writeFileTool) RunWithOptions(ctx context.Context, args map[string]an
newInfo, _ := os.Stat(absolutePath)
options.FileTracker.Record(absolutePath, []byte(content), newInfo)
if content == modelKnownContent {
options.FileTracker.RecordSeenRange(absolutePath, 1, trackedLineTotal(content), trackedLineTotal(content))
options.FileTracker.RecordSeenRange(absolutePath, 1, lineCount(content), lineCount(content))
}
if !existed {
options.FileTracker.RecordCreated(absolutePath)
+5 -20
View File
@@ -866,9 +866,7 @@ func TestStructuredPatchFailedDeleteDoesNotRecreateMissingFile(t *testing.T) {
}
}
// A move whose source vanishes after planning is refused by the pre-commit
// recheck before the destination is published, so nothing is left half-done.
func TestStructuredPatchMoveWithMissingSourceIsRefusedBeforePublishing(t *testing.T) {
func TestStructuredPatchMoveFailureLeavesPublishedDestination(t *testing.T) {
root := t.TempDir()
workspace, err := os.OpenRoot(root)
if err != nil {
@@ -877,35 +875,22 @@ func TestStructuredPatchMoveWithMissingSourceIsRefusedBeforePublishing(t *testin
defer workspace.Close()
sourcePath := filepath.Join(root, "source.txt")
destinationPath := filepath.Join(root, "destination.txt")
writeTestFile(t, sourcePath, "removed source\n")
change := structuredPatchChange{
kind: structuredPatchUpdate,
from: structuredPatchTarget{absolute: sourcePath, relative: "source.txt"},
to: structuredPatchTarget{absolute: destinationPath, relative: "destination.txt"},
before: "removed source\n", after: "moved content\n", mode: 0o644,
}
// The source exists at planning time and disappears just before commit.
removed := false
structuredPatchBeforeCommit = func(structuredPatchChange) {
if err := os.Remove(sourcePath); err != nil {
t.Fatal(err)
}
removed = true
}
defer func() { structuredPatchBeforeCommit = nil }()
err = applyStructuredPatchChanges(workspace, []structuredPatchChange{change}, nil)
if !removed {
t.Fatal("pre-commit hook did not run")
}
if err == nil || !strings.Contains(err.Error(), "before commit") || strings.Contains(err.Error(), "partially applied") {
t.Fatalf("move with a missing source = %v, want a pre-commit refusal with nothing applied", err)
if err == nil || !strings.Contains(err.Error(), "partially applied") {
t.Fatalf("move with a missing source = %v, want partial-application error", err)
}
if _, statErr := os.Stat(sourcePath); !os.IsNotExist(statErr) {
t.Fatalf("failed move recreated a missing source: %v", statErr)
}
if _, statErr := os.Stat(destinationPath); !os.IsNotExist(statErr) {
t.Fatalf("refused move must not publish the destination: %v", statErr)
if got := mustReadTestFile(t, destinationPath); got != "moved content\n" {
t.Fatalf("published move destination = %q", got)
}
}
-11
View File
@@ -53,7 +53,6 @@ const (
CounterModelSwitches = "model_switches"
CounterInputTokens = "input_tokens"
CounterCachedInputTokens = "cached_input_tokens"
CounterCacheWriteTokens = "cache_write_tokens"
CounterOutputTokens = "output_tokens"
CounterPrewarmAttempts = "prewarm_attempts"
CounterPrefixStable = "prefix_stable"
@@ -65,12 +64,6 @@ const (
CounterPostureEscalations = "posture_escalations"
)
const (
CounterResponseChainReused = "response_chain_reused"
CounterResponseChainReset = "response_chain_reset"
CounterResponsesHTTPFallback = "responses_http_fallback"
)
// Span is one named wall interval attributed to part of a run. Each stamp is
// its own entry (spans are not merged by name): a run that streams two model
// completions has two generation entries, which lets the recorder derive
@@ -297,7 +290,6 @@ func OptionalEventKeys() []string {
"span:" + SpanVerification,
"counter:" + CounterToolCalls,
"counter:" + CounterCachedInputTokens,
"counter:" + CounterCacheWriteTokens,
"counter:" + CounterRetryCount,
"counter:" + CounterReconnectCount,
"counter:" + CounterCompactionCount,
@@ -309,9 +301,6 @@ func OptionalEventKeys() []string {
"counter:" + CounterPrewarmAttempts,
"counter:" + CounterPrefixStable,
"counter:" + CounterPrefixDrift,
"counter:" + CounterResponseChainReused,
"counter:" + CounterResponseChainReset,
"counter:" + CounterResponsesHTTPFallback,
"counter:" + CounterPostureEscalations,
"event:prefix_hash",
"task_state",
+4 -14
View File
@@ -614,21 +614,11 @@ func TestCounterPrecisionRoundTrip(t *testing.T) {
}
}
func TestOptionalEventKeysIncludeRuntimeCounters(t *testing.T) {
keys := make(map[string]bool)
func TestOptionalEventKeysIncludePostureEscalations(t *testing.T) {
for _, key := range OptionalEventKeys() {
keys[key] = true
}
for _, counter := range []string{
CounterCacheWriteTokens,
CounterResponseChainReused,
CounterResponseChainReset,
CounterResponsesHTTPFallback,
CounterPostureEscalations,
} {
key := "counter:" + counter
if !keys[key] {
t.Errorf("%s missing from OptionalEventKeys: %v", key, OptionalEventKeys())
if key == "counter:"+CounterPostureEscalations {
return
}
}
t.Fatalf("posture_escalations missing from OptionalEventKeys: %v", OptionalEventKeys())
}
+25 -85
View File
@@ -85,12 +85,10 @@ type model struct {
savedProviders []config.ProviderProfile
provider zeroruntime.Provider
newProvider func(config.ProviderProfile) (zeroruntime.Provider, error)
newTurnSessionProvider func(config.ProviderProfile, zeroruntime.Provider) zeroruntime.TurnSessionProvider
probeProviderHealth func(context.Context, providerhealth.Options) providerhealth.Result
discoverProviderModels func(context.Context, config.ProviderProfile) ([]providermodeldiscovery.Model, error)
discoverOllamaContextWindow func(ctx context.Context, baseURL string, model string) (int, error)
registry *tools.Registry
awaitToolReadiness func(context.Context)
// lspManager is created once per session and reused across prompts so gopls (and
// other language servers) stay warm — a fresh manager per run would cold-start
// the server on the first edit of every turn. Nil when cwd is unknown; runs then
@@ -235,20 +233,17 @@ type model struct {
leaderHelpOverlay bool
// leaderPending is true after Ctrl+X until a second key, Esc, or timeout
// resolves the chord (see leader.go). leaderSeq invalidates a stale tick.
leaderPending bool
leaderSeq int
transcriptBodyHeights *transcriptBodyHeightCache
input textinput.Model
composer composerState
composerActive bool
composerCursorVisible bool
composerBlinkSeq int
composerBlinkIdleTicks int
composerBlinkTicking bool
composerPastePreviews []composerPastePreview
composerSelection composerSelectionState
dictation dictationController
sttKeyPrompt *sttKeyPromptState
leaderPending bool
leaderSeq int
transcriptBodyHeights *transcriptBodyHeightCache
input textinput.Model
composer composerState
composerActive bool
composerCursorVisible bool
composerPastePreviews []composerPastePreview
composerSelection composerSelectionState
dictation dictationController
sttKeyPrompt *sttKeyPromptState
// plan holds the sticky plan panel state (steps, expansion, timings)
// synced from the update_plan tool. See plan_panel.go.
plan planPanelState
@@ -496,11 +491,10 @@ type model struct {
// transientNotice is a single, replaceable confirmation shown above the
// composer. Unlike copyStatus it is shared by lightweight slash commands
// and is never persisted into the transcript.
transientNotice transientNotice
transientNoticeSeq int
transientNoticeTimerSeq int
exitConfirmActive bool
exitConfirmSeq int
transientNotice transientNotice
transientNoticeSeq int
exitConfirmActive bool
exitConfirmSeq int
// cancelConfirmActive/cancelConfirmSeq mirror exitConfirmActive/exitConfirmSeq
// (same seq-gated tea.Tick pattern) but guard a DIFFERENT action: Esc
// cancelling a running turn. The two are deliberately separate state (not a
@@ -933,13 +927,6 @@ func newModel(ctx context.Context, options Options) model {
// terminal bracketed paste (Paste: true) is the single paste path.
input.KeyMap.Paste.SetEnabled(false)
input.Focus()
// The main composer paints its own cursor through composerCursorVisible.
// Keep the embedded textinput cursor static so the shared textinput.Blink
// message cannot start a second, invisible 530ms tick chain. Setup inputs
// still use the normal Bubbles cursor and continue blinking independently.
inputStyles := input.Styles()
inputStyles.Cursor.Blink = false
input.SetStyles(inputStyles)
runSpinner := spinner.New(spinner.WithSpinner(spinner.MiniDot))
runSpinner.Spinner.FPS = activeAnimationFrameInterval
@@ -964,8 +951,6 @@ func newModel(ctx context.Context, options Options) model {
userCommands: loadedUserCommands,
loadSkills: options.LoadSkills,
composerCursorVisible: true,
composerBlinkSeq: 1,
composerBlinkTicking: true,
terminalFocused: true,
userConfigPath: options.UserConfigPath,
doctorUserConfigPath: doctorUserConfigPath,
@@ -981,12 +966,10 @@ func newModel(ctx context.Context, options Options) model {
recapsEnabled: options.RecapsEnabled,
provider: options.Provider,
newProvider: options.NewProvider,
newTurnSessionProvider: options.NewTurnSessionProvider,
probeProviderHealth: options.ProbeProviderHealth,
discoverProviderModels: options.DiscoverProviderModels,
discoverOllamaContextWindow: options.DiscoverOllamaContextWindow,
registry: registry,
awaitToolReadiness: options.AwaitToolReadiness,
sessionStore: sessionStore,
peerService: options.PeerService,
sandboxStore: sandboxStore,
@@ -1107,11 +1090,6 @@ const (
// composerCursorBlinkInterval is the on/off period of the composer text cursor.
const composerCursorBlinkInterval = 530 * time.Millisecond
// composerBlinkIdleTicksBeforeSettle bounds the software cursor animation.
// Once the terminal is focused but otherwise idle, the cursor finishes a short
// blink sequence and settles visible instead of waking the TUI forever.
const composerBlinkIdleTicksBeforeSettle = 4
// composerTypingIdleThreshold is how long a typing pause must last before the
// cursor resumes blinking; comfortably above normal inter-keystroke gaps
// (~150-300ms) so it won't flicker mid-sentence.
@@ -1120,27 +1098,16 @@ const composerTypingIdleThreshold = 500 * time.Millisecond
// composerBlinkMsg toggles the composer cursor's visibility each tick. The custom
// composer render draws its own cursor (not textinput's), so it drives its own
// blink rather than relying on textinput.Blink.
type composerBlinkMsg struct{ seq int }
type composerBlinkMsg struct{}
func composerBlinkCmd(seq int) tea.Cmd {
func composerBlinkCmd() tea.Cmd {
return tea.Tick(composerCursorBlinkInterval, func(time.Time) tea.Msg {
return composerBlinkMsg{seq: seq}
return composerBlinkMsg{}
})
}
func (m model) armComposerBlink() (model, tea.Cmd) {
m.composerCursorVisible = true
m.composerBlinkIdleTicks = 0
if m.composerBlinkTicking {
return m, nil
}
m.composerBlinkTicking = true
m.composerBlinkSeq++
return m, composerBlinkCmd(m.composerBlinkSeq)
}
func (m model) Init() tea.Cmd {
cmds := []tea.Cmd{textinput.Blink, composerBlinkCmd(m.composerBlinkSeq)}
cmds := []tea.Cmd{textinput.Blink, composerBlinkCmd()}
if m.petAnimation != nil && !m.reducedMotion {
cmds = append(cmds, petTickCmd(m.petTickSeq, m.petFrameDelay()))
}
@@ -1323,15 +1290,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
nm = nm.syncChatScroll()
nm, mouseCmd := nm.syncMouseCapture()
nm, flushCmd := nm.settleTranscript()
var composerCmd tea.Cmd
switch msg.(type) {
case tea.KeyPressMsg, tea.PasteMsg, tea.FocusMsg:
if nm.terminalFocused {
nm, composerCmd = nm.armComposerBlink()
}
}
nm, noticeCmd := nm.ensureTransientNoticeTimer()
return nm, batchCommands(cmd, mouseCmd, flushCmd, recapActivityCmd, composerCmd, noticeCmd)
return nm, batchCommands(cmd, mouseCmd, flushCmd, recapActivityCmd)
}
func batchCommands(cmds ...tea.Cmd) tea.Cmd {
@@ -1394,27 +1353,16 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
}
return m, nil
case composerBlinkMsg:
if !m.composerBlinkTicking || msg.seq != m.composerBlinkSeq {
return m, nil
}
m = m.expireTransientNotice()
switch {
case !m.terminalFocused:
m.composerCursorVisible = false
m.composerBlinkTicking = false
return m, nil
m.composerCursorVisible = false // hidden while unfocused
case m.now().Sub(m.lastCharTime) < composerTypingIdleThreshold:
m.composerCursorVisible = true
m.composerBlinkIdleTicks = 0
m.composerCursorVisible = true // solid while actively typing
default:
m.composerBlinkIdleTicks++
if m.composerBlinkIdleTicks >= composerBlinkIdleTicksBeforeSettle {
m.composerCursorVisible = true
m.composerBlinkTicking = false
return m, nil
}
m.composerCursorVisible = !m.composerCursorVisible
m.composerCursorVisible = !m.composerCursorVisible // idle + focused: blink as before
}
return m, composerBlinkCmd(m.composerBlinkSeq)
return m, composerBlinkCmd()
case tea.BackgroundColorMsg:
// A background reading changes only the contrast direction used for named
// palettes. It never repaints the terminal canvas.
@@ -2241,8 +2189,6 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
}
m.terminalFocused = false
m.composerCursorVisible = false
m.composerBlinkTicking = false
m.composerBlinkSeq++
if m.notifier != nil {
m.notifier.SetFocused(false)
}
@@ -5426,9 +5372,6 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
sessionEvents := []pendingSessionEvent{}
usageModelID := m.modelName
var specReview *pendingSpecReviewPrompt
if m.awaitToolReadiness != nil {
m.awaitToolReadiness(runCtx)
}
options := m.agentOptions
options.Registry = cloneToolRegistry(m.registry)
goalAwareRun := !runOptions.specDraft && m.activeLoopID == "" &&
@@ -5471,9 +5414,6 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
options.ResponseStyle = m.responseStyle
options.Cwd = m.cwd
options.Images = images
if m.newTurnSessionProvider != nil && m.provider != nil {
options.TurnSessionProvider = m.newTurnSessionProvider(m.providerProfile, m.provider)
}
if m.captureRunImages != nil {
m.captureRunImages(images)
}
+12 -157
View File
@@ -109,81 +109,6 @@ func TestPromptSubmitInjectsLiveSessionModelContext(t *testing.T) {
}
}
func TestPromptWaitsForToolReadinessBeforeSnapshot(t *testing.T) {
t.Setenv("XDG_DATA_HOME", t.TempDir())
registry := tools.NewRegistry()
provider := &fakeProvider{events: []zeroruntime.StreamEvent{
{Type: zeroruntime.StreamEventText, Content: "done"},
{Type: zeroruntime.StreamEventDone},
}}
awaited := false
m := newModel(context.Background(), Options{
Cwd: t.TempDir(),
ProviderName: "test",
ModelName: "test-model",
Provider: provider,
Registry: registry,
AwaitToolReadiness: func(context.Context) {
awaited = true
registry.Register(tools.NewScopedReadFileTool(t.TempDir(), nil))
},
})
m.input.SetValue("inspect the workspace")
updated, cmd := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if cmd == nil {
t.Fatal("expected prompt submit to start an agent run")
}
next.Update(execCmd(cmd))
if !awaited {
t.Fatal("agent run did not await tool readiness")
}
if len(provider.requests) != 1 {
t.Fatalf("provider requests = %d, want 1", len(provider.requests))
}
for _, definition := range provider.requests[0].Tools {
if definition.Name == "read_file" {
return
}
}
t.Fatal("provider request omitted tool published by readiness callback")
}
func TestPromptBuildsTurnSessionForCurrentProvider(t *testing.T) {
t.Setenv("XDG_DATA_HOME", t.TempDir())
provider := &fakeProvider{events: []zeroruntime.StreamEvent{
{Type: zeroruntime.StreamEventText, Content: "done"},
{Type: zeroruntime.StreamEventDone},
}}
profile := config.ProviderProfile{Name: "chatgpt", Model: "gpt-test"}
called := false
m := newModel(context.Background(), Options{
Cwd: t.TempDir(),
ProviderName: profile.Name,
ModelName: profile.Model,
ProviderProfile: profile,
Provider: provider,
Registry: tools.NewRegistry(),
NewTurnSessionProvider: func(gotProfile config.ProviderProfile, gotProvider zeroruntime.Provider) zeroruntime.TurnSessionProvider {
called = true
if gotProfile.Name != profile.Name || gotProfile.Model != profile.Model || gotProvider != provider {
t.Fatalf("turn session factory inputs = %#v/%#v", gotProfile, gotProvider)
}
return zeroruntime.NewProviderTurnSessionProvider(gotProvider, zeroruntime.ProviderCapabilities{})
},
})
m.input.SetValue("inspect the workspace")
updated, cmd := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
next.Update(execCmd(cmd))
if !called {
t.Fatal("agent run did not build a turn session for the current provider")
}
}
func TestPromptSubmitStoresReasoningSeparatelyFromAnswer(t *testing.T) {
t.Setenv("XDG_DATA_HOME", t.TempDir())
provider := &fakeProvider{events: []zeroruntime.StreamEvent{
@@ -623,10 +548,10 @@ func TestModelCommandSwitchesSessionModel(t *testing.T) {
})
m.input.SetValue("/model gpt-4.1-mini")
updated, _ := m.Update(testKey(tea.KeyEnter))
updated, cmd := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if next.pending {
if cmd != nil {
t.Fatal("expected /model to be handled without starting an agent run")
}
if next.modelName != "gpt-4.1-mini" || next.provider != nextProvider {
@@ -666,10 +591,10 @@ func TestModelCommandAcceptsChatGPTCatalogModelID(t *testing.T) {
}
m.input.SetValue("/model openai/gpt-5.6-sol")
updated, _ := m.Update(testKey(tea.KeyEnter))
updated, cmd := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if next.pending {
if cmd != nil {
t.Fatal("expected /model to be handled without starting an agent run")
}
if next.modelName != "gpt-5.6-sol" || next.provider != nextProvider {
@@ -738,10 +663,10 @@ func TestModelCommandPersistsSelectedModelToUserConfig(t *testing.T) {
})
m.input.SetValue("/model gpt-4.1-mini")
updated, _ := m.Update(testKey(tea.KeyEnter))
updated, cmd := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if next.pending {
if cmd != nil {
t.Fatal("expected /model to be handled without starting an agent run")
}
if next.modelName != "gpt-4.1-mini" {
@@ -2514,7 +2439,7 @@ func TestAssistantNarrationBeforeToolCardGetsBlankSeparator(t *testing.T) {
m.transcript = append(m.transcript,
transcriptRow{kind: rowUser, text: "run it"},
transcriptRow{kind: rowAssistant, text: "I'll inspect the existing file, then run it."},
transcriptRow{kind: rowToolResult, id: "t1", tool: "read_file", status: tools.StatusOK, detail: "File: time_test.py\n\n1→print('x')"},
transcriptRow{kind: rowToolResult, id: "t1", tool: "read_file", status: tools.StatusOK, detail: "File: time_test.py\n\n 1 | print('x')"},
)
items := m.transcriptBodyItems(m.chatColumnWidth(), "", false)
toolIdx := -1
@@ -2943,8 +2868,6 @@ func TestComposerBlinkStaysSolidWhileTyping(t *testing.T) {
terminalFocused: true,
lastCharTime: base,
composerCursorVisible: true,
composerBlinkSeq: 1,
composerBlinkTicking: true,
}
// Each iteration simulates a keystroke (refreshing lastCharTime) followed by
@@ -2953,7 +2876,7 @@ func TestComposerBlinkStaysSolidWhileTyping(t *testing.T) {
for i := 0; i < 3; i++ {
now = now.Add(200 * time.Millisecond)
m.lastCharTime = now
updated, _ := m.Update(composerBlinkMsg{seq: m.composerBlinkSeq})
updated, _ := m.Update(composerBlinkMsg{})
m = updated.(model)
if !m.composerCursorVisible {
t.Fatalf("tick %d: expected cursor to stay visible while typing, got hidden", i)
@@ -2961,13 +2884,6 @@ func TestComposerBlinkStaysSolidWhileTyping(t *testing.T) {
}
}
func TestMainComposerDisablesEmbeddedCursorBlink(t *testing.T) {
m := newModel(t.Context(), Options{})
if m.input.Styles().Cursor.Blink {
t.Fatal("embedded composer cursor blink must stay disabled")
}
}
func TestComposerBlinkHiddenWhileUnfocused(t *testing.T) {
base := time.Date(2026, 7, 7, 12, 0, 0, 0, time.UTC)
m := model{
@@ -3004,10 +2920,10 @@ func TestComposerBlinkResumesAfterRefocusAndIdle(t *testing.T) {
t.Fatal("expected terminalFocused to be true after FocusMsg")
}
updated, _ = m.Update(composerBlinkMsg{seq: m.composerBlinkSeq})
updated, _ = m.Update(composerBlinkMsg{})
m = updated.(model)
first := m.composerCursorVisible
updated, _ = m.Update(composerBlinkMsg{seq: m.composerBlinkSeq})
updated, _ = m.Update(composerBlinkMsg{})
m = updated.(model)
second := m.composerCursorVisible
if first == second {
@@ -3015,34 +2931,6 @@ func TestComposerBlinkResumesAfterRefocusAndIdle(t *testing.T) {
}
}
func TestComposerBlinkRejectsTickFromBeforeRefocus(t *testing.T) {
base := time.Date(2026, 7, 7, 12, 0, 0, 0, time.UTC)
m := model{
now: func() time.Time { return base },
terminalFocused: true,
lastCharTime: base.Add(-time.Hour),
composerCursorVisible: true,
composerBlinkSeq: 5,
composerBlinkTicking: true,
}
staleSeq := m.composerBlinkSeq
updated, _ := m.Update(tea.BlurMsg{})
m = updated.(model)
updated, refocusCmd := m.Update(tea.FocusMsg{})
m = updated.(model)
if refocusCmd == nil || m.composerBlinkSeq == staleSeq {
t.Fatal("refocus did not re-arm composer blinking with a new sequence")
}
visible := m.composerCursorVisible
updated, staleCmd := m.Update(composerBlinkMsg{seq: staleSeq})
m = updated.(model)
if staleCmd != nil || m.composerCursorVisible != visible {
t.Fatalf("stale tick changed composer state: cmd=%v visible=%v want=%v", staleCmd != nil, m.composerCursorVisible, visible)
}
}
func TestComposerBlinkTogglesWhenIdleAndFocused(t *testing.T) {
base := time.Date(2026, 7, 7, 12, 0, 0, 0, time.UTC)
m := model{
@@ -3050,53 +2938,20 @@ func TestComposerBlinkTogglesWhenIdleAndFocused(t *testing.T) {
terminalFocused: true,
lastCharTime: base.Add(-time.Hour),
composerCursorVisible: true,
composerBlinkSeq: 1,
composerBlinkTicking: true,
}
updated, _ := m.Update(composerBlinkMsg{seq: m.composerBlinkSeq})
updated, _ := m.Update(composerBlinkMsg{})
m = updated.(model)
if m.composerCursorVisible {
t.Fatal("expected cursor to toggle off on first idle+focused tick")
}
updated, _ = m.Update(composerBlinkMsg{seq: m.composerBlinkSeq})
updated, _ = m.Update(composerBlinkMsg{})
m = updated.(model)
if !m.composerCursorVisible {
t.Fatal("expected cursor to toggle back on on second idle+focused tick")
}
}
func TestComposerBlinkSettlesVisibleWithoutRescheduling(t *testing.T) {
m := model{
now: func() time.Time { return time.Unix(100, 0) },
terminalFocused: true,
lastCharTime: time.Unix(0, 0),
composerCursorVisible: true,
composerBlinkSeq: 1,
composerBlinkTicking: true,
}
var cmd tea.Cmd
for range composerBlinkIdleTicksBeforeSettle {
var updated tea.Model
updated, cmd = m.Update(composerBlinkMsg{seq: m.composerBlinkSeq})
m = updated.(model)
}
if m.composerBlinkTicking || cmd != nil || !m.composerCursorVisible {
t.Fatalf("settled cursor = ticking:%v cmd:%v visible:%v", m.composerBlinkTicking, cmd != nil, m.composerCursorVisible)
}
}
func TestComposerInputRestartsSettledBlink(t *testing.T) {
m := newModel(t.Context(), Options{})
m.composerBlinkTicking = false
m.composerCursorVisible = true
updated, cmd := m.Update(tea.KeyPressMsg(tea.Key{Code: 'a', Text: "a"}))
m = updated.(model)
if !m.composerBlinkTicking || cmd == nil || !m.composerCursorVisible {
t.Fatalf("restarted cursor = ticking:%v cmd:%v visible:%v", m.composerBlinkTicking, cmd != nil, m.composerCursorVisible)
}
}
func TestComposerCursorShowsImmediatelyOnKeypress(t *testing.T) {
// The blink phase may have just hidden the caret when the user starts
// typing; the typed character must render with a caret immediately, not
+11 -16
View File
@@ -36,7 +36,6 @@ type Options struct {
RecapsEnabled bool
Provider zeroruntime.Provider
NewProvider func(config.ProviderProfile) (zeroruntime.Provider, error)
NewTurnSessionProvider func(config.ProviderProfile, zeroruntime.Provider) zeroruntime.TurnSessionProvider
ProbeProviderHealth func(context.Context, providerhealth.Options) providerhealth.Result
DiscoverProviderModels func(context.Context, config.ProviderProfile) ([]providermodeldiscovery.Model, error)
DiscoverOllamaContextWindow func(ctx context.Context, baseURL string, model string) (int, error)
@@ -44,21 +43,17 @@ type Options struct {
PrepareRunCompletionWarning func()
RunCompletionWarning func() string
Registry *tools.Registry
// AwaitToolReadiness gives prompt-critical integration startup a bounded
// chance to publish its tools before this turn snapshots the registry. The
// wait runs inside the asynchronous agent command, so the TUI stays usable.
AwaitToolReadiness func(context.Context)
SessionStore *sessions.Store
SandboxStore *sandbox.GrantStore
MCPConfig config.MCPConfig
MCPPermissionStore *mcp.PermissionStore
MCPTokenStore *mcp.TokenStore
MCPCommand func(context.Context, []string) MCPCommandResult
SandboxSetupCommand func(context.Context) SandboxSetupCommandResult
UsageTracker *usage.Tracker
SessionCompactor SessionCompactor
PrService *PrService
PeerService *peermsg.Service
SessionStore *sessions.Store
SandboxStore *sandbox.GrantStore
MCPConfig config.MCPConfig
MCPPermissionStore *mcp.PermissionStore
MCPTokenStore *mcp.TokenStore
MCPCommand func(context.Context, []string) MCPCommandResult
SandboxSetupCommand func(context.Context) SandboxSetupCommandResult
UsageTracker *usage.Tracker
SessionCompactor SessionCompactor
PrService *PrService
PeerService *peermsg.Service
AgentOptions agent.Options
// LoadSkills returns the installed skills (default skills dir merged with any
+4 -4
View File
@@ -521,9 +521,9 @@ func TestModelPickerAppliesActiveProviderCatalogModelID(t *testing.T) {
})
m.input.SetValue("/model openai/gpt-4.1")
updated, _ := m.Update(testKey(tea.KeyEnter))
updated, cmd := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if next.pending {
if cmd != nil {
t.Fatal("expected /model to be handled without starting a run")
}
if captured.Model != "openai/gpt-4.1" {
@@ -867,9 +867,9 @@ func TestModelCommandAcceptsManualModelForCustomProvider(t *testing.T) {
})
m.input.SetValue("/model qwen-custom-latest")
updated, _ := m.Update(testKey(tea.KeyEnter))
updated, cmd := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if next.pending {
if cmd != nil {
t.Fatal("expected /model to be handled without starting a run")
}
if captured.Model != "qwen-custom-latest" {
+3 -3
View File
@@ -1075,8 +1075,8 @@ func TestEditedDiffCardHandlesBareBlankHunkContext(t *testing.T) {
func TestReadCardBodyShowsExploredSummary(t *testing.T) {
m := limeTestModel()
// Mirrors the real read_file output shape: "N→<text>".
detail := "File: internal/agent/loop.go\n\n12→func Run() {\n13→}\n"
// Mirrors the real read_file output shape: "<right-aligned N> | <text>".
detail := "File: internal/agent/loop.go\n\n 12 | func Run() {\n 13 | }\n"
row := transcriptRow{kind: rowToolResult, id: "call_1", tool: "read_file", status: tools.StatusOK, detail: detail}
rc := buildRowContext([]transcriptRow{{kind: rowToolCall, id: "call_1", tool: "read_file", detail: "internal/agent/loop.go"}})
got := plainRender(t, m.renderRow(row, 80, rc))
@@ -1097,7 +1097,7 @@ func TestReadCardBodyShowsExploredSummary(t *testing.T) {
row.expanded = true
expanded := plainRender(t, m.renderRow(row, 80, rc))
for _, want := range []string{"Explored", "└ Read", "internal/agent/loop.go", "func Run()", "13→"} {
for _, want := range []string{"Explored", "└ Read", "internal/agent/loop.go", "func Run()", "13 |"} {
if !strings.Contains(expanded, want) {
t.Fatalf("expanded read card = %q, missing %q", expanded, want)
}
+6 -6
View File
@@ -46,10 +46,10 @@ func TestEffortCommandListsAndSetsSupportedEffort(t *testing.T) {
}
next.input.SetValue("/effort high")
updated, _ = next.Update(testKey(tea.KeyEnter))
updated, cmd = next.Update(testKey(tea.KeyEnter))
next = updated.(model)
if next.pending {
if cmd != nil {
t.Fatal("expected /effort high to be handled without starting an agent run")
}
if next.reasoningEffort != modelregistry.ReasoningEffortHigh {
@@ -153,10 +153,10 @@ func TestStyleCommandListsAndSetsSessionPreference(t *testing.T) {
}
next.input.SetValue("/style explanatory")
updated, _ = next.Update(testKey(tea.KeyEnter))
updated, cmd = next.Update(testKey(tea.KeyEnter))
next = updated.(model)
if next.pending {
if cmd != nil {
t.Fatal("expected /style explanatory to be handled without starting an agent run")
}
if next.responseStyle != "explanatory" {
@@ -915,10 +915,10 @@ func TestModelSwitchClearsUnsupportedEffortPreference(t *testing.T) {
})
m.input.SetValue("/model gpt-4.1")
updated, _ := m.Update(testKey(tea.KeyEnter))
updated, cmd := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if next.pending {
if cmd != nil {
t.Fatal("expected /model to be handled without starting an agent run")
}
if next.reasoningEffort != "" {
+8 -1
View File
@@ -247,7 +247,14 @@ func (m model) cancelSpecReview() (tea.Model, tea.Cmd) {
}
func cloneToolRegistry(registry *tools.Registry) *tools.Registry {
return registry.Clone()
clone := tools.NewRegistry()
if registry == nil {
return clone
}
for _, tool := range registry.All() {
clone.Register(tool)
}
return clone
}
// renderFocusedSpecReviewPrompt draws the spec-review gate in the shared card
+1 -1
View File
@@ -41,7 +41,7 @@ func TestDefaultToolBodyRegistrySelectsCoreRenderers(t *testing.T) {
{
name: "read_file",
hint: "README.md",
detail: "File: README.md\n\n7→# Zero",
detail: "File: README.md\n\n 7 | # Zero",
want: []string{"Read", "README.md"},
},
{
Loaded 100 of 105 files, more files were not shown because too many files have changed in this diff. Show more