mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-08 01:09:46 -10:00
Compare commits
15
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
221fb93041 | ||
|
|
ac3958f860 | ||
|
|
63767f3e42 | ||
|
|
f30f550e60 | ||
|
|
aadb4a27e9 | ||
|
|
a6ec0435e3 | ||
|
|
1b5db17656 | ||
|
|
eeea3308fe | ||
|
|
27b319ca88 | ||
|
|
1fde941803 | ||
|
|
7acac7acee | ||
|
|
1649008506 | ||
|
|
6fe0d1edcd | ||
|
|
e59bf78cab | ||
|
|
6645f9744b |
No files matched your search
@@ -273,6 +273,19 @@ zero sandbox grants list
|
||||
Zero includes local file/search/edit/shell tools, `web_fetch` for public URLs,
|
||||
and MCP support for additional tools.
|
||||
|
||||
`web_fetch` refuses loopback, private and other special-use addresses: it checks
|
||||
the URL before asking permission, resolves the host, and dials the address it
|
||||
validated so a name cannot resolve to something else in between.
|
||||
|
||||
If `HTTP_PROXY`/`HTTPS_PROXY` is set, `web_fetch` and the provider connectivity
|
||||
probe use it, and that last step changes: the proxy is dialed and the target
|
||||
hostname is sent to it, so the proxy decides which address the request actually
|
||||
reaches. The URL is still checked and resolved locally first, but a proxy that
|
||||
answers differently can reach a private service. A forward proxy already sees
|
||||
and can rewrite every request through it, so this is the trust you accept by
|
||||
configuring one. Leave the variables unset for the checks to be enforced end to
|
||||
end.
|
||||
|
||||
For local dev servers, use shell commands such as `curl` through `exec_command`
|
||||
so the normal sandbox and permission policy applies. Long-running commands stay
|
||||
attached to a background terminal session and can be listed or stopped from the
|
||||
@@ -371,9 +384,13 @@ 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`
|
||||
disable. A/B-benchmark it by running the same `zero-perf-bench` suite with the
|
||||
variable unset and set.
|
||||
`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.
|
||||
|
||||
### Code Quality and Security Checks
|
||||
|
||||
|
||||
@@ -627,7 +627,12 @@ Core tool groups include:
|
||||
skills, `ask_user`, permission requests.
|
||||
- **Write tools**: `write_file`, `edit_file`, `apply_patch`, `update_plan`.
|
||||
- **Shell tools**: `exec_command`, `write_stdin`, legacy `bash`.
|
||||
- **Network tools**: `web_fetch` and optional web search backend.
|
||||
- **Network tools**: `web_fetch` and optional web search backend. `web_fetch`
|
||||
validates the URL before the permission prompt, refuses loopback/private/
|
||||
special-use addresses, and dials the address it resolved. With
|
||||
`HTTP_PROXY`/`HTTPS_PROXY` set, the proxy is dialed instead and resolves the
|
||||
target itself, so destination enforcement ends at the proxy; see the note in
|
||||
the README.
|
||||
- **Local-control tools**: config-gated wrappers for browser, terminal, desktop,
|
||||
and artifact helpers backed by packages such as `internal/localcontrol` and
|
||||
`internal/browser`.
|
||||
|
||||
@@ -4,7 +4,7 @@ go 1.26.6
|
||||
|
||||
require (
|
||||
charm.land/bubbles/v2 v2.1.1
|
||||
charm.land/bubbletea/v2 v2.0.8
|
||||
charm.land/bubbletea/v2 v2.0.9
|
||||
charm.land/lipgloss/v2 v2.0.5
|
||||
github.com/Microsoft/go-winio v0.6.2
|
||||
github.com/alecthomas/chroma/v2 v2.27.0
|
||||
|
||||
@@ -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.8 h1:SxTJMhCAI3lbPmy4SgX5LWZ24AdINr4I6UEqzZvYJuY=
|
||||
charm.land/bubbletea/v2 v2.0.8/go.mod h1:2SkdgoTXluXJHOUwAoRlRXF/28vklb1rFl6GcgV1/ss=
|
||||
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/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,8 +64,6 @@ 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=
|
||||
|
||||
@@ -142,13 +142,15 @@ func actionOffered(optionID string, offered []PermissionOption) bool {
|
||||
// session/request_permission request from a ZERO permission request.
|
||||
func permissionToolCall(req agent.PermissionRequest) ToolCallUpdate {
|
||||
args := marshalArgs(req.Args)
|
||||
return ToolCallUpdate{
|
||||
upd := ToolCallUpdate{
|
||||
ToolCallID: req.ToolCallID,
|
||||
Title: toolTitle(req.ToolName, string(args)),
|
||||
Kind: toolKindFor(req.ToolName),
|
||||
Status: ToolStatusPending,
|
||||
RawInput: rawInputBytes(args),
|
||||
}
|
||||
attachBrowserToolDetails(&upd, req.ToolName)
|
||||
return upd
|
||||
}
|
||||
|
||||
func marshalArgs(args map[string]any) []byte {
|
||||
|
||||
@@ -82,3 +82,17 @@ func TestPermissionToolCall(t *testing.T) {
|
||||
t.Error("expected rawInput from args")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPermissionToolCallKeepsTheBrowserDescriptor(t *testing.T) {
|
||||
call := permissionToolCall(agent.PermissionRequest{
|
||||
ToolCallID: "browser-1",
|
||||
ToolName: "browser_connect",
|
||||
Args: map[string]any{"target": "127.0.0.1:9222"},
|
||||
})
|
||||
if got := browserDescriptor(t, call); got != (BrowserToolDetails{Version: 1, Command: "connect"}) {
|
||||
t.Fatalf("browser descriptor = %#v", got)
|
||||
}
|
||||
if call.Title != "browser connect" {
|
||||
t.Fatalf("title = %q", call.Title)
|
||||
}
|
||||
}
|
||||
+119
-1
@@ -2,7 +2,9 @@ package acp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/url"
|
||||
"strings"
|
||||
"unicode"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/agent"
|
||||
@@ -44,12 +46,125 @@ func toolKindFor(name string) string {
|
||||
|
||||
// toolTitle builds a concise human title, e.g. "read_file src/main.go".
|
||||
func toolTitle(name, rawArgs string) string {
|
||||
if browser, ok := browserToolDetails(name); ok {
|
||||
return browserToolTitle(browser.Command, rawArgs)
|
||||
}
|
||||
if hint := primaryArgHint(rawArgs); hint != "" {
|
||||
return name + " " + hint
|
||||
}
|
||||
return name
|
||||
}
|
||||
|
||||
// browserToolDetails identifies ZERO's local browser helpers without treating
|
||||
// similarly named MCP tools as browser automation. The descriptor intentionally
|
||||
// contains no request data: ACP tool input is already protocol-visible, but a
|
||||
// durable UI must not need to retain text, local CDP targets, or full URLs just
|
||||
// to recognise the browser operation.
|
||||
func browserToolDetails(name string) (*BrowserToolDetails, bool) {
|
||||
const prefix = "browser_"
|
||||
command, ok := strings.CutPrefix(name, prefix)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
switch command {
|
||||
case "install", "launch", "connect", "open", "snapshot", "click", "type", "press", "action":
|
||||
return &BrowserToolDetails{Version: 1, Command: command}, true
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
|
||||
const zeroBrowserMetaKey = "github.com/Gitlawb/zero/browser"
|
||||
|
||||
// attachBrowserToolDetails stores ZERO's browser descriptor in ACP's reserved
|
||||
// extension channel. Keeping this in one helper prevents start, result, and
|
||||
// permission payloads from drifting onto different wire shapes.
|
||||
func attachBrowserToolDetails(update *ToolCallUpdate, name string) {
|
||||
browser, ok := browserToolDetails(name)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
raw, err := json.Marshal(browser)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
update.Meta = map[string]json.RawMessage{zeroBrowserMetaKey: raw}
|
||||
}
|
||||
|
||||
// browserToolTitle avoids putting browser_type text, an attached DevTools
|
||||
// endpoint, or a URL query/fragment in a tool-card title. Those values can
|
||||
// carry credentials or session data; the UI only needs the operation and, for
|
||||
// navigation, a human-recognisable origin.
|
||||
func browserToolTitle(command, rawArgs string) string {
|
||||
switch command {
|
||||
case "action":
|
||||
action, ok := exactJSONStringArg(rawArgs, "command")
|
||||
if !ok {
|
||||
return "browser action"
|
||||
}
|
||||
if action, ok := tools.NormalizedBrowserActionCommand(action); ok {
|
||||
return "browser action " + action
|
||||
}
|
||||
return "browser action"
|
||||
case "open":
|
||||
rawURL, ok := exactJSONStringArg(rawArgs, "url")
|
||||
if !ok {
|
||||
return "browser open"
|
||||
}
|
||||
normalized, err := tools.NormalizeBrowserOpenURL(rawURL)
|
||||
if err != nil {
|
||||
return "browser open"
|
||||
}
|
||||
u, err := url.Parse(normalized)
|
||||
if err != nil || u.Scheme == "" || u.Host == "" {
|
||||
return "browser open"
|
||||
}
|
||||
origin := u.Scheme + "://" + u.Host
|
||||
if !browserTitleTextSafe(origin) {
|
||||
return "browser open"
|
||||
}
|
||||
return "browser open " + truncateHint(origin)
|
||||
default:
|
||||
return "browser " + command
|
||||
}
|
||||
}
|
||||
|
||||
// browserTitleTextSafe validates text after URL parsing has decoded escaped
|
||||
// UTF-8 in the host. Valid UTF-8 alone is not presentation-safe: control,
|
||||
// format/bidi, and line/paragraph separator runes can reorder or split the
|
||||
// permission label shown to a user. The execution URL remains unchanged.
|
||||
func browserTitleTextSafe(text string) bool {
|
||||
if !utf8.ValidString(text) {
|
||||
return false
|
||||
}
|
||||
for _, r := range text {
|
||||
if unicode.IsControl(r) || unicode.Is(unicode.Cf, r) || unicode.Is(unicode.Zl, r) || unicode.Is(unicode.Zp, r) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// exactJSONStringArg mirrors ZERO's map-based tool argument decoding: only the
|
||||
// exact JSON key is considered, and a non-string value is invalid. In
|
||||
// particular, an incidental "URL" key must not change a permission title when
|
||||
// browser_open will only read "url".
|
||||
func exactJSONStringArg(rawArgs, key string) (string, bool) {
|
||||
var args map[string]json.RawMessage
|
||||
if json.Unmarshal([]byte(rawArgs), &args) != nil {
|
||||
return "", false
|
||||
}
|
||||
raw, ok := args[key]
|
||||
if !ok {
|
||||
return "", false
|
||||
}
|
||||
var value string
|
||||
if json.Unmarshal(raw, &value) != nil {
|
||||
return "", false
|
||||
}
|
||||
return value, true
|
||||
}
|
||||
|
||||
// primaryArgHint extracts the most relevant argument (path/pattern/command) from
|
||||
// raw JSON arguments. Best-effort; returns "" when it can't parse.
|
||||
func primaryArgHint(rawArgs string) string {
|
||||
@@ -89,7 +204,7 @@ func rawInput(args string) json.RawMessage {
|
||||
// toolCallStart maps an advertised ZERO tool call to the initial ACP "tool_call"
|
||||
// update (status in_progress — ZERO executes immediately after advertising).
|
||||
func toolCallStart(call agent.ToolCall) ToolCallUpdate {
|
||||
return ToolCallUpdate{
|
||||
upd := ToolCallUpdate{
|
||||
SessionUpdate: UpdateToolCall,
|
||||
ToolCallID: call.ID,
|
||||
Title: toolTitle(call.Name, call.Arguments),
|
||||
@@ -97,6 +212,8 @@ func toolCallStart(call agent.ToolCall) ToolCallUpdate {
|
||||
Status: ToolStatusInProgress,
|
||||
RawInput: rawInput(call.Arguments),
|
||||
}
|
||||
attachBrowserToolDetails(&upd, call.Name)
|
||||
return upd
|
||||
}
|
||||
|
||||
// toolCallResult maps a finished ZERO tool result to a "tool_call_update".
|
||||
@@ -116,6 +233,7 @@ func toolCallResult(result agent.ToolResult) ToolCallUpdate {
|
||||
if locs := toolResultLocations(result); len(locs) > 0 {
|
||||
upd.Locations = locs
|
||||
}
|
||||
attachBrowserToolDetails(&upd, result.Name)
|
||||
return upd
|
||||
}
|
||||
|
||||
|
||||
@@ -1,14 +1,29 @@
|
||||
package acp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
"unicode"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/agent"
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
)
|
||||
|
||||
func browserDescriptor(t *testing.T, update ToolCallUpdate) BrowserToolDetails {
|
||||
t.Helper()
|
||||
raw, ok := update.Meta[zeroBrowserMetaKey]
|
||||
if !ok {
|
||||
t.Fatalf("browser metadata = %#v, want %q", update.Meta, zeroBrowserMetaKey)
|
||||
}
|
||||
var details BrowserToolDetails
|
||||
if err := json.Unmarshal(raw, &details); err != nil {
|
||||
t.Fatalf("decode browser metadata: %v", err)
|
||||
}
|
||||
return details
|
||||
}
|
||||
|
||||
func TestAgentMessageAndThoughtChunks(t *testing.T) {
|
||||
m := agentMessageChunk("hello")
|
||||
if m.SessionUpdate != UpdateAgentMessageChunk || m.Content.Type != "text" || m.Content.Text != "hello" {
|
||||
@@ -56,6 +71,217 @@ func TestToolTitleAndHint(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestBrowserToolUpdatesAreStructuredAndPresentationSafe(t *testing.T) {
|
||||
start := toolCallStart(agent.ToolCall{
|
||||
ID: "browser-1",
|
||||
Name: "browser_open",
|
||||
Arguments: `{"url":"https://example.com/settings?token=not-for-a-title#account"}`,
|
||||
})
|
||||
if got := browserDescriptor(t, start); got != (BrowserToolDetails{Version: 1, Command: "open"}) {
|
||||
t.Fatalf("browser descriptor = %#v, want open", got)
|
||||
}
|
||||
if start.Title != "browser open https://example.com" {
|
||||
t.Fatalf("browser title = %q", start.Title)
|
||||
}
|
||||
if strings.Contains(start.Title, "token=") || strings.Contains(start.Title, "#account") {
|
||||
t.Fatalf("browser title leaked URL-sensitive data: %q", start.Title)
|
||||
}
|
||||
encoded, err := json.Marshal(start)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var wire struct {
|
||||
Meta map[string]json.RawMessage `json:"_meta"`
|
||||
}
|
||||
if err := json.Unmarshal(encoded, &wire); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, ok := wire.Meta[zeroBrowserMetaKey]; !ok {
|
||||
t.Fatalf("browser wire metadata = %#v", wire.Meta)
|
||||
}
|
||||
|
||||
typed := toolCallStart(agent.ToolCall{
|
||||
ID: "browser-2",
|
||||
Name: "browser_type",
|
||||
Arguments: `{"ref":"email","text":"[email protected]"}`,
|
||||
})
|
||||
if got := browserDescriptor(t, typed); got.Command != "type" {
|
||||
t.Fatalf("browser type descriptor = %#v", got)
|
||||
}
|
||||
if typed.Title != "browser type" || strings.Contains(typed.Title, "[email protected]") {
|
||||
t.Fatalf("browser type title = %q", typed.Title)
|
||||
}
|
||||
|
||||
action := toolCallStart(agent.ToolCall{
|
||||
ID: "browser-3",
|
||||
Name: "browser_action",
|
||||
Arguments: `{"command":"keyboard_insert_text","args":["[email protected]"]}`,
|
||||
})
|
||||
if action.Title != "browser action keyboard_insert_text" {
|
||||
t.Fatalf("browser action title = %q", action.Title)
|
||||
}
|
||||
|
||||
result := toolCallResult(agent.ToolResult{
|
||||
ToolCallID: "browser-2",
|
||||
Name: "browser_type",
|
||||
Status: tools.StatusOK,
|
||||
})
|
||||
if got := browserDescriptor(t, result); got.Command != "type" {
|
||||
t.Fatalf("browser result descriptor = %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBrowserDescriptorSurvivesProtocolShapedRoundTrip(t *testing.T) {
|
||||
updates := []ToolCallUpdate{
|
||||
toolCallStart(agent.ToolCall{
|
||||
ID: "start",
|
||||
Name: "browser_open",
|
||||
Arguments: `{"url":"https://user:[email protected]/private?token=secret#fragment"}`,
|
||||
}),
|
||||
toolCallResult(agent.ToolResult{
|
||||
ToolCallID: "result",
|
||||
Name: "browser_type",
|
||||
Status: tools.StatusOK,
|
||||
}),
|
||||
permissionToolCall(agent.PermissionRequest{
|
||||
ToolCallID: "permission",
|
||||
ToolName: "browser_connect",
|
||||
Args: map[string]any{"target": "127.0.0.1:9222"},
|
||||
}),
|
||||
}
|
||||
|
||||
type protocolToolCallUpdate struct {
|
||||
SessionUpdate string `json:"sessionUpdate,omitempty"`
|
||||
ToolCallID string `json:"toolCallId"`
|
||||
Title string `json:"title,omitempty"`
|
||||
Kind string `json:"kind,omitempty"`
|
||||
Status string `json:"status,omitempty"`
|
||||
RawInput json.RawMessage `json:"rawInput,omitempty"`
|
||||
Content []ToolCallContent `json:"content,omitempty"`
|
||||
Locations []ToolCallLocation `json:"locations,omitempty"`
|
||||
Meta map[string]json.RawMessage `json:"_meta,omitempty"`
|
||||
}
|
||||
|
||||
for _, update := range updates {
|
||||
encoded, err := json.Marshal(update)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var root map[string]json.RawMessage
|
||||
if err := json.Unmarshal(encoded, &root); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, ok := root["browser"]; ok {
|
||||
t.Fatalf("browser descriptor escaped ACP _meta: %s", encoded)
|
||||
}
|
||||
|
||||
var protocol protocolToolCallUpdate
|
||||
if err := json.Unmarshal(encoded, &protocol); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
forwarded, err := json.Marshal(protocol)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var roundTripped ToolCallUpdate
|
||||
if err := json.Unmarshal(forwarded, &roundTripped); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
details := browserDescriptor(t, roundTripped)
|
||||
if details.Version != 1 || details.Command == "" {
|
||||
t.Fatalf("round-tripped browser descriptor = %#v", details)
|
||||
}
|
||||
descriptorJSON := string(roundTripped.Meta[zeroBrowserMetaKey])
|
||||
for _, secret := range []string{"password", "private", "token", "fragment", "127.0.0.1", "9222"} {
|
||||
if strings.Contains(descriptorJSON, secret) {
|
||||
t.Fatalf("browser descriptor leaked %q: %s", secret, descriptorJSON)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBrowserPermissionTitlesMirrorSafeToolArguments(t *testing.T) {
|
||||
if got := browserToolTitle("open", `{"url":"evil.example.test/pay?token=hidden#fragment"}`); got != "browser open https://evil.example.test" {
|
||||
t.Fatalf("bare-host title = %q", got)
|
||||
}
|
||||
if got := browserToolTitle("open", `{"URL":"https://different.example.test"}`); got != "browser open" {
|
||||
t.Fatalf("case-variant URL title = %q", got)
|
||||
}
|
||||
if got := browserToolTitle("open", `{"URL":"https://different.example.test","url":"https://actual.example.test/path"}`); got != "browser open https://actual.example.test" {
|
||||
t.Fatalf("exact URL key title = %q", got)
|
||||
}
|
||||
if got := browserToolTitle("action", `{"command":"not an action"}`); got != "browser action" {
|
||||
t.Fatalf("unknown browser action title = %q", got)
|
||||
}
|
||||
|
||||
longHost := "https://" + strings.Repeat("a", 200) + ".example.test/path?token=hidden"
|
||||
title := browserToolTitle("open", `{"url":"`+longHost+`"}`)
|
||||
if !utf8.ValidString(title) || utf8.RuneCountInString(title) > len("browser open ")+61 || strings.Contains(title, "token=") {
|
||||
t.Fatalf("bounded browser origin title = %q", title)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBrowserOpenTitlesRejectDecodedUnicodePresentationControls(t *testing.T) {
|
||||
for _, rawURL := range []string{
|
||||
"https://safe.example%E2%80%AEevil.test/path",
|
||||
"https://safe.example%E2%81%A6evil.test/path",
|
||||
"https://safe.example%C2%85evil.test/path",
|
||||
"https://safe.example%E2%80%A8evil.test/path",
|
||||
"https://safe.example%E2%80%A9evil.test/path",
|
||||
} {
|
||||
t.Run(rawURL, func(t *testing.T) {
|
||||
normalized, err := tools.NormalizeBrowserOpenURL(rawURL)
|
||||
if err != nil {
|
||||
t.Fatalf("execution URL rejected: %v", err)
|
||||
}
|
||||
if normalized != rawURL {
|
||||
t.Fatalf("execution URL = %q, want unchanged %q", normalized, rawURL)
|
||||
}
|
||||
|
||||
args, err := json.Marshal(map[string]any{"url": rawURL})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
updates := []ToolCallUpdate{
|
||||
toolCallStart(agent.ToolCall{ID: "start", Name: "browser_open", Arguments: string(args)}),
|
||||
permissionToolCall(agent.PermissionRequest{ToolCallID: "permission", ToolName: "browser_open", Args: map[string]any{"url": rawURL}}),
|
||||
}
|
||||
for _, update := range updates {
|
||||
encoded, err := json.Marshal(update)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var decoded ToolCallUpdate
|
||||
if err := json.Unmarshal(encoded, &decoded); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if decoded.Title != "browser open" {
|
||||
t.Fatalf("unsafe browser title survived wire round trip: %q", decoded.Title)
|
||||
}
|
||||
for _, r := range decoded.Title {
|
||||
if unicode.IsControl(r) || unicode.Is(unicode.Cf, r) || unicode.Is(unicode.Zl, r) || unicode.Is(unicode.Zp, r) {
|
||||
t.Fatalf("browser title contains unsafe presentation rune %U: %q", r, decoded.Title)
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBrowserDescriptorDoesNotClaimSimilarlyNamedMCPTools(t *testing.T) {
|
||||
start := toolCallStart(agent.ToolCall{ID: "mcp-1", Name: "browser_plugin_open", Arguments: `{}`})
|
||||
if len(start.Meta) != 0 {
|
||||
t.Fatalf("MCP-like tool received built-in browser metadata: %#v", start.Meta)
|
||||
}
|
||||
encoded, err := json.Marshal(start)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if strings.Contains(string(encoded), `"browser"`) {
|
||||
t.Fatalf("non-browser tool encoded browser field: %s", encoded)
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolCallStart(t *testing.T) {
|
||||
upd := toolCallStart(agent.ToolCall{ID: "tc1", Name: "read_file", Arguments: `{"path":"a.go"}`})
|
||||
if upd.SessionUpdate != UpdateToolCall {
|
||||
|
||||
@@ -210,6 +210,21 @@ type ToolCallUpdate struct {
|
||||
RawInput json.RawMessage `json:"rawInput,omitempty"`
|
||||
Content []ToolCallContent `json:"content,omitempty"`
|
||||
Locations []ToolCallLocation `json:"locations,omitempty"`
|
||||
// Meta is ACP's extension channel. ZERO-owned values must remain beneath a
|
||||
// namespaced key so protocol-shaped clients can preserve them while decoding
|
||||
// and re-encoding a tool call.
|
||||
Meta map[string]json.RawMessage `json:"_meta,omitempty"`
|
||||
}
|
||||
|
||||
// BrowserToolDetails identifies the browser helper operation behind a tool
|
||||
// call. Version is the schema version for this optional ZERO extension;
|
||||
// Command is one of install, launch, connect, open, snapshot, click, type,
|
||||
// press, or action. Future fields must remain display-safe and must not
|
||||
// include browser profile data, cookies, typed text, URL paths/queries, or
|
||||
// DevTools endpoints.
|
||||
type BrowserToolDetails struct {
|
||||
Version int `json:"version"`
|
||||
Command string `json:"command"`
|
||||
}
|
||||
|
||||
// ToolCallContent is a tool call's rendered output. ZERO emits "content" (a
|
||||
|
||||
+178
-11
@@ -4,7 +4,9 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strconv"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/trace"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
@@ -144,11 +146,17 @@ 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
|
||||
@@ -390,6 +398,27 @@ type compactionState struct {
|
||||
// starts at 1.0 and converges via an EMA as each turn reports actual usage, so
|
||||
// later turns compact nearer to real capacity. Zero is treated as 1.0.
|
||||
calibrationRatio float64
|
||||
|
||||
// summarizerFactory lazily builds the dedicated (cheap) summarizer provider
|
||||
// (Options.Summarizer) for the current main model. summarizerProvider
|
||||
// memoizes the built provider and summarizerResolved records that the
|
||||
// factory ran (it may legitimately return no provider: the main model is
|
||||
// already the cheap one). summarizerBroken is the sticky fallback: once a
|
||||
// build or call fails, the run stays on the main provider until the next
|
||||
// model switch — a misconfigured cheap model must never break (or
|
||||
// repeatedly delay) compaction.
|
||||
summarizerFactory func(ctx context.Context, mainModelID string) (Provider, error)
|
||||
summarizerProvider Provider
|
||||
summarizerResolved bool
|
||||
summarizerBroken bool
|
||||
// window is the context window the threshold was derived from. The
|
||||
// reactive free-prune stage compares against it (the limit the provider
|
||||
// enforces), not against the 70% proactive threshold — a pruned history
|
||||
// between the two would still succeed on retry.
|
||||
window int
|
||||
// mainModel is the model the run currently streams on; the summarizer
|
||||
// choice is made against it and remade after a mid-run switch.
|
||||
mainModel string
|
||||
}
|
||||
|
||||
// calibrate folds one turn's (rawEstimate, actualPromptTokens) sample into the
|
||||
@@ -424,17 +453,88 @@ func (state *compactionState) calibratedTokens(raw int) int {
|
||||
|
||||
func newCompactionState(options Options, task *taskState) *compactionState {
|
||||
state := &compactionState{
|
||||
enabled: options.ContextWindow > 0,
|
||||
threshold: compactionThreshold(options.ContextWindow),
|
||||
preserveLast: options.CompactionPreserveLast,
|
||||
onUsage: options.OnUsage,
|
||||
task: task,
|
||||
planner: newContextPlanner(contextPlannerConfig{contextWindow: options.ContextWindow, serviceTier: options.ServiceTier}),
|
||||
trace: options.Trace,
|
||||
enabled: options.ContextWindow > 0,
|
||||
threshold: compactionThreshold(options.ContextWindow),
|
||||
window: options.ContextWindow,
|
||||
mainModel: options.Model,
|
||||
preserveLast: options.CompactionPreserveLast,
|
||||
onUsage: options.OnUsage,
|
||||
task: task,
|
||||
planner: newContextPlanner(contextPlannerConfig{contextWindow: options.ContextWindow, serviceTier: options.ServiceTier}),
|
||||
trace: options.Trace,
|
||||
summarizerFactory: options.Summarizer,
|
||||
}
|
||||
return state
|
||||
}
|
||||
|
||||
// updateWindow re-derives the compaction threshold after a mid-run model
|
||||
// switch. The low-water mark is reset: it was measured against the old
|
||||
// threshold, and keeping it could suppress a compaction the smaller window
|
||||
// now requires. <= 0 (unknown model) leaves everything unchanged.
|
||||
func (state *compactionState) updateWindow(window int) {
|
||||
if window <= 0 {
|
||||
return
|
||||
}
|
||||
state.enabled = true
|
||||
state.window = window
|
||||
state.threshold = compactionThreshold(window)
|
||||
state.lowWaterMark = 0
|
||||
if state.planner != nil {
|
||||
state.planner.config.contextWindow = window
|
||||
}
|
||||
}
|
||||
|
||||
// switchModel records a mid-run model switch: the window is re-derived
|
||||
// (window <= 0 keeps the old one) and the summarizer choice is forgotten so
|
||||
// the next compaction re-resolves it against the new main model. A run that
|
||||
// started on the cheap model (no dedicated summarizer) and escalated to an
|
||||
// expensive one therefore regains cheap summaries — the exact workflow a
|
||||
// dedicated summarizer exists for.
|
||||
func (state *compactionState) switchModel(modelID string, window int) {
|
||||
state.mainModel = modelID
|
||||
state.updateWindow(window)
|
||||
state.summarizerProvider = nil
|
||||
state.summarizerResolved = false
|
||||
state.summarizerBroken = false
|
||||
}
|
||||
|
||||
// summarizeProvider returns the provider summarization calls should use and
|
||||
// whether it is the dedicated summarizer: that one when it applies to the
|
||||
// current main model and is healthy, else the main provider. The factory runs
|
||||
// at most once per main model; a build error marks the summarizer broken for
|
||||
// that model. Callers branch on the boolean rather than comparing Provider
|
||||
// values — interface comparison panics on a non-comparable dynamic type.
|
||||
func (state *compactionState) summarizeProvider(ctx context.Context, main Provider) (Provider, bool) {
|
||||
if state.summarizerFactory == nil || state.summarizerBroken {
|
||||
return main, false
|
||||
}
|
||||
if !state.summarizerResolved {
|
||||
state.summarizerResolved = true
|
||||
built, err := state.summarizerFactory(ctx, state.mainModel)
|
||||
if err != nil {
|
||||
state.summarizerBroken = true
|
||||
return main, false
|
||||
}
|
||||
state.summarizerProvider = built
|
||||
}
|
||||
if state.summarizerProvider == nil {
|
||||
return main, false
|
||||
}
|
||||
return state.summarizerProvider, true
|
||||
}
|
||||
|
||||
// fitsWindow reports whether a calibrated size is safely under the limit the
|
||||
// provider enforces. The proactive threshold (70% of the window) is a trigger
|
||||
// for the paid summarizer, not the hard limit; the reactive path compares
|
||||
// against the window itself, with a small margin for estimate error, so a
|
||||
// history pruned to, say, 85% of the window takes the free retry.
|
||||
func (state *compactionState) fitsWindow(size int) bool {
|
||||
if state.window <= 0 {
|
||||
return size <= state.threshold
|
||||
}
|
||||
return size <= state.window-state.window/20
|
||||
}
|
||||
|
||||
// maybeCompact runs proactive compaction at the top of a turn. It returns the
|
||||
// (possibly compacted) message slice. It is a no-op when compaction is disabled,
|
||||
// when the history is under threshold, or when the history has not grown past
|
||||
@@ -528,10 +628,34 @@ func (state *compactionState) recover(
|
||||
// context-limit error never triggers an unexpected summarization call.
|
||||
return messages, false, nil
|
||||
}
|
||||
if state.reactiveAttempted {
|
||||
if !isContextLimitError(errorMessage) {
|
||||
return messages, false, nil
|
||||
}
|
||||
if !isContextLimitError(errorMessage) {
|
||||
|
||||
// CHEAP FIRST STAGE (mirrors maybeCompact): a context-limit failure can
|
||||
// often be fixed for free by dropping superseded reads and pruning old
|
||||
// tool-result bodies. It runs regardless of the one-shot reactive budget
|
||||
// (it costs nothing and cannot loop: a second pass finds nothing left to
|
||||
// prune). The free retry is taken only when the pruned history fits the
|
||||
// window the provider enforces — otherwise a retry would fail the same
|
||||
// way, so fall through to the paid summarizer in this same call.
|
||||
toolTokens := estimateToolDefTokens(tools)
|
||||
prunedAny := false
|
||||
if pruned, reclaimed := pruneSupersededReadResults(messages); reclaimed > 0 {
|
||||
messages, prunedAny = pruned, true
|
||||
}
|
||||
if pruned, reclaimed := pruneStaleToolOutput(messages, state.preserveLast); reclaimed > 0 {
|
||||
messages, prunedAny = pruned, true
|
||||
}
|
||||
if prunedAny {
|
||||
size := state.calibratedTokens(estimateTokens(messages) + toolTokens)
|
||||
if state.fitsWindow(size) {
|
||||
state.lowWaterMark = size
|
||||
return messages, true, nil
|
||||
}
|
||||
}
|
||||
|
||||
if state.reactiveAttempted {
|
||||
return messages, false, nil
|
||||
}
|
||||
|
||||
@@ -571,11 +695,21 @@ func (state *compactionState) recover(
|
||||
// provider call. The summary stream intentionally does NOT forward OnText (so
|
||||
// compaction stays invisible on the user-facing surface), but it DOES forward
|
||||
// OnUsage so the summarizer's token cost is still counted by usage/budgeting.
|
||||
// The call goes to the dedicated (cheap) summarizer when one applies; if that
|
||||
// call fails the summarizer is marked broken for the current model and the
|
||||
// same slice is retried once on the main provider, so the caller sees at most
|
||||
// one (main-provider-quality) error.
|
||||
func (state *compactionState) summarizeClosure(ctx context.Context, provider Provider) func([]zeroruntime.Message) (string, error) {
|
||||
if state.planner == nil {
|
||||
state.planner = newContextPlanner(contextPlannerConfig{})
|
||||
}
|
||||
return func(toSummarize []zeroruntime.Message) (string, error) {
|
||||
chosen, dedicated := state.summarizeProvider(ctx, provider)
|
||||
summary, err := summarizeWithPlanner(ctx, chosen, toSummarize, state.onUsage, state.planner, state.trace)
|
||||
if err == nil || !dedicated {
|
||||
return summary, err
|
||||
}
|
||||
state.summarizerBroken = true
|
||||
return summarizeWithPlanner(ctx, provider, toSummarize, state.onUsage, state.planner, state.trace)
|
||||
}
|
||||
}
|
||||
@@ -659,8 +793,41 @@ func summarizeMessagesOnce(ctx context.Context, provider Provider, messages []ze
|
||||
return summary, nil
|
||||
}
|
||||
|
||||
// summaryToolResultClamp bounds each tool-result body in the summarizer input.
|
||||
// Tool outputs dominate the elided middle's bulk (build logs, file reads,
|
||||
// search results), yet the summary instructions only need their gist — the
|
||||
// model that ACTED on the full output already turned it into decisions the
|
||||
// surrounding messages record. 2,000 chars (~500 tokens) keeps the useful
|
||||
// head+tail of a result while cutting the summarizer input by 5-10x on
|
||||
// tool-heavy sessions; the summarizer's own cost is the price being cut.
|
||||
const summaryToolResultClamp = 2000
|
||||
|
||||
// summaryToolArgsClamp bounds each tool-call argument blob in the transcript:
|
||||
// a write_file call carries the whole file body in its arguments.
|
||||
const summaryToolArgsClamp = 500
|
||||
|
||||
// clampForSummary head+tail-truncates text for the summarizer transcript,
|
||||
// snapping to rune boundaries so the clamp never splits a UTF-8 sequence.
|
||||
func clampForSummary(text string, max int) string {
|
||||
if max <= 0 || len(text) <= max {
|
||||
return text
|
||||
}
|
||||
head := max * 3 / 4
|
||||
tail := max - head
|
||||
for head > 0 && !utf8.RuneStart(text[head]) {
|
||||
head--
|
||||
}
|
||||
start := len(text) - tail
|
||||
for start < len(text) && !utf8.RuneStart(text[start]) {
|
||||
start++
|
||||
}
|
||||
return text[:head] + "\n…[" + strconv.Itoa(len(text)-head-(len(text)-start)) + " chars omitted for summarization]…\n" + text[start:]
|
||||
}
|
||||
|
||||
// renderTranscript flattens messages into a plain-text transcript for the
|
||||
// summarizer. Secret scrubbing already happened upstream at the tool boundary.
|
||||
// Tool-result bodies and tool-call arguments are clamped (see the constants
|
||||
// above): the summary needs what was done and decided, not the raw payloads.
|
||||
func renderTranscript(messages []zeroruntime.Message) string {
|
||||
lines := make([]string, 0, len(messages))
|
||||
for _, message := range messages {
|
||||
@@ -670,13 +837,13 @@ func renderTranscript(messages []zeroruntime.Message) string {
|
||||
if len(message.ToolCalls) > 0 {
|
||||
calls := make([]string, 0, len(message.ToolCalls))
|
||||
for _, call := range message.ToolCalls {
|
||||
calls = append(calls, call.Name+"("+call.Arguments+")")
|
||||
calls = append(calls, call.Name+"("+clampForSummary(call.Arguments, summaryToolArgsClamp)+")")
|
||||
}
|
||||
line += "\n[tool calls: " + strings.Join(calls, "; ") + "]"
|
||||
}
|
||||
lines = append(lines, line)
|
||||
case zeroruntime.MessageRoleTool:
|
||||
lines = append(lines, "tool result: "+message.Content)
|
||||
lines = append(lines, "tool result: "+clampForSummary(message.Content, summaryToolResultClamp))
|
||||
default:
|
||||
lines = append(lines, string(message.Role)+": "+message.Content)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,469 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// compactableHistory returns a history whose middle is large enough to trip
|
||||
// the proactive threshold for a 1000-token window.
|
||||
func compactableHistory() []zeroruntime.Message {
|
||||
return []zeroruntime.Message{
|
||||
{Role: zeroruntime.MessageRoleSystem, Content: "sys"},
|
||||
{Role: zeroruntime.MessageRoleUser, Content: strings.Repeat("u", 4000)},
|
||||
{Role: zeroruntime.MessageRoleAssistant, Content: strings.Repeat("a", 4000)},
|
||||
{Role: zeroruntime.MessageRoleUser, Content: "u2"},
|
||||
{Role: zeroruntime.MessageRoleAssistant, Content: "a2"},
|
||||
{Role: zeroruntime.MessageRoleUser, Content: "u3"},
|
||||
}
|
||||
}
|
||||
|
||||
// Summarization must run on the dedicated (cheap) summarizer when one is
|
||||
// configured — the main provider sees no summarize traffic at all.
|
||||
func TestMaybeCompactUsesDedicatedSummarizer(t *testing.T) {
|
||||
cheap := &mockProvider{turns: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventText, Content: "CHEAP SUMMARY"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
}}}
|
||||
main := &mockProvider{}
|
||||
factoryCalls := 0
|
||||
state := newCompactionState(Options{
|
||||
ContextWindow: 1000,
|
||||
CompactionPreserveLast: 2,
|
||||
Summarizer: func(context.Context, string) (Provider, error) {
|
||||
factoryCalls++
|
||||
return cheap, nil
|
||||
},
|
||||
}, nil)
|
||||
|
||||
compacted := state.maybeCompact(context.Background(), main, compactableHistory(), nil)
|
||||
|
||||
if len(cheap.requests) == 0 {
|
||||
t.Fatal("dedicated summarizer never received the summarize call")
|
||||
}
|
||||
if len(main.requests) != 0 {
|
||||
t.Fatalf("main provider must see no summarize traffic, got %d requests", len(main.requests))
|
||||
}
|
||||
if factoryCalls != 1 {
|
||||
t.Fatalf("summarizer factory calls = %d, want 1 (lazy, memoized)", factoryCalls)
|
||||
}
|
||||
if estimateTokens(compacted) >= estimateTokens(compactableHistory()) {
|
||||
t.Fatal("compaction did not shrink the history")
|
||||
}
|
||||
found := false
|
||||
for _, message := range compacted {
|
||||
if strings.Contains(message.Content, "CHEAP SUMMARY") {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatal("cheap summarizer's summary missing from the compacted history")
|
||||
}
|
||||
}
|
||||
|
||||
// A failing summarizer falls back to the main provider for the SAME slice and
|
||||
// stays broken for the rest of the run — no repeated failed cheap calls.
|
||||
func TestSummarizerFailureFallsBackToMainAndSticks(t *testing.T) {
|
||||
broken := &mockProvider{turns: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventError, Error: "boom: model not found"},
|
||||
}}}
|
||||
main := &mockProvider{turns: [][]zeroruntime.StreamEvent{
|
||||
{{Type: zeroruntime.StreamEventText, Content: "MAIN SUMMARY"}, {Type: zeroruntime.StreamEventDone}},
|
||||
{{Type: zeroruntime.StreamEventText, Content: "MAIN SUMMARY 2"}, {Type: zeroruntime.StreamEventDone}},
|
||||
}}
|
||||
state := newCompactionState(Options{
|
||||
ContextWindow: 1000,
|
||||
CompactionPreserveLast: 2,
|
||||
Summarizer: func(context.Context, string) (Provider, error) { return broken, nil },
|
||||
}, nil)
|
||||
summarize := state.summarizeClosure(context.Background(), main)
|
||||
input := []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "history"}}
|
||||
|
||||
summary, err := summarize(input)
|
||||
if err != nil {
|
||||
t.Fatalf("fallback must succeed via the main provider: %v", err)
|
||||
}
|
||||
if summary != "MAIN SUMMARY" {
|
||||
t.Fatalf("summary = %q, want the main provider's", summary)
|
||||
}
|
||||
if len(broken.requests) != 1 {
|
||||
t.Fatalf("broken summarizer requests = %d, want exactly 1", len(broken.requests))
|
||||
}
|
||||
if !state.summarizerBroken {
|
||||
t.Fatal("a failed summarizer must be marked broken for the run")
|
||||
}
|
||||
|
||||
// Second summarization: straight to main, no cheap retry.
|
||||
if _, err := summarize(input); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(broken.requests) != 1 {
|
||||
t.Fatalf("broken summarizer was retried (%d requests)", len(broken.requests))
|
||||
}
|
||||
if len(main.requests) != 2 {
|
||||
t.Fatalf("main provider requests = %d, want 2", len(main.requests))
|
||||
}
|
||||
}
|
||||
|
||||
// A factory that errors (misconfigured model) is equivalent to no summarizer:
|
||||
// main provider used, factory not retried.
|
||||
func TestSummarizerFactoryErrorFallsBack(t *testing.T) {
|
||||
main := &mockProvider{turns: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventText, Content: "MAIN SUMMARY"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
}}}
|
||||
factoryCalls := 0
|
||||
state := newCompactionState(Options{
|
||||
ContextWindow: 1000,
|
||||
CompactionPreserveLast: 2,
|
||||
Summarizer: func(context.Context, string) (Provider, error) {
|
||||
factoryCalls++
|
||||
return nil, errors.New("unknown model")
|
||||
},
|
||||
}, nil)
|
||||
summarize := state.summarizeClosure(context.Background(), main)
|
||||
|
||||
summary, err := summarize([]zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "history"}})
|
||||
if err != nil || summary != "MAIN SUMMARY" {
|
||||
t.Fatalf("summary = %q, err = %v; want main provider fallback", summary, err)
|
||||
}
|
||||
if factoryCalls != 1 {
|
||||
t.Fatalf("factory calls = %d, want 1 (sticky failure)", factoryCalls)
|
||||
}
|
||||
if _, dedicated := state.summarizeProvider(context.Background(), main); dedicated {
|
||||
t.Fatal("after a factory failure the run must stay on the main provider")
|
||||
}
|
||||
if factoryCalls != 1 {
|
||||
t.Fatalf("factory retried after failure (%d calls)", factoryCalls)
|
||||
}
|
||||
}
|
||||
|
||||
// Tool payloads are clamped in the summarizer transcript: results to
|
||||
// summaryToolResultClamp, call arguments to summaryToolArgsClamp.
|
||||
func TestRenderTranscriptClampsToolPayloads(t *testing.T) {
|
||||
bigResult := strings.Repeat("r", 50_000)
|
||||
bigArgs := strings.Repeat("w", 20_000)
|
||||
transcript := renderTranscript([]zeroruntime.Message{
|
||||
{Role: zeroruntime.MessageRoleAssistant, Content: "editing", ToolCalls: []zeroruntime.ToolCall{
|
||||
{ID: "c1", Name: "write_file", Arguments: bigArgs},
|
||||
}},
|
||||
{Role: zeroruntime.MessageRoleTool, ToolCallID: "c1", Content: bigResult},
|
||||
{Role: zeroruntime.MessageRoleUser, Content: "keep going"},
|
||||
})
|
||||
|
||||
if len(transcript) > summaryToolResultClamp+summaryToolArgsClamp+1024 {
|
||||
t.Fatalf("transcript is %d bytes; clamps did not apply", len(transcript))
|
||||
}
|
||||
if !strings.Contains(transcript, "chars omitted for summarization") {
|
||||
t.Fatal("clamped transcript must carry the omission marker")
|
||||
}
|
||||
if !strings.Contains(transcript, "keep going") {
|
||||
t.Fatal("non-tool content must survive unclamped")
|
||||
}
|
||||
// Head AND tail of the clamped payloads survive.
|
||||
if !strings.HasPrefix(transcript[strings.Index(transcript, "tool result: ")+len("tool result: "):], "r") ||
|
||||
!strings.Contains(transcript, "rrr\n\nuser: keep going") {
|
||||
t.Fatalf("head/tail of the tool result lost:\n%s", transcript[:200])
|
||||
}
|
||||
}
|
||||
|
||||
// bigToolHistory returns a history with n large tool results (~40k estimated
|
||||
// tokens each): the trailing protection window keeps the newest verbatim,
|
||||
// the older ones are prunable.
|
||||
func bigToolHistory(n int) []zeroruntime.Message {
|
||||
bigBody := strings.Repeat("x", 160_000)
|
||||
messages := []zeroruntime.Message{{Role: zeroruntime.MessageRoleSystem, Content: "sys"}}
|
||||
for i := 0; i < n; i++ {
|
||||
id := string(rune('a' + i))
|
||||
messages = append(messages,
|
||||
zeroruntime.Message{Role: zeroruntime.MessageRoleAssistant, ToolCalls: []zeroruntime.ToolCall{{ID: id, Name: "read_file", Arguments: "{\"path\":\"" + id + "\"}"}}},
|
||||
zeroruntime.Message{Role: zeroruntime.MessageRoleTool, ToolCallID: id, Content: bigBody},
|
||||
)
|
||||
}
|
||||
return append(messages, zeroruntime.Message{Role: zeroruntime.MessageRoleUser, Content: "go"})
|
||||
}
|
||||
|
||||
func summaryProvider(text string) *mockProvider {
|
||||
return &mockProvider{turns: [][]zeroruntime.StreamEvent{
|
||||
{{Type: zeroruntime.StreamEventText, Content: text}, {Type: zeroruntime.StreamEventDone}},
|
||||
{{Type: zeroruntime.StreamEventText, Content: text}, {Type: zeroruntime.StreamEventDone}},
|
||||
}}
|
||||
}
|
||||
|
||||
// A context-limit error is first answered with the free prune stage; when the
|
||||
// pruned history fits the window the retry is free (no summarizer call) and
|
||||
// the one-shot reactive budget is untouched. The paid summarizer only runs
|
||||
// when a later error finds nothing left to prune.
|
||||
func TestRecoverPrunesBeforeSummarizing(t *testing.T) {
|
||||
provider := summaryProvider("SUMMARY")
|
||||
// The protected trailing tool output (~40k tokens) must fit the window.
|
||||
state := newCompactionState(Options{ContextWindow: 200_000, CompactionPreserveLast: 2}, nil)
|
||||
messages := bigToolHistory(5)
|
||||
|
||||
pruned, retried, err := state.recover(context.Background(), provider, messages, nil, "maximum context length exceeded")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !retried {
|
||||
t.Fatal("prunable history that fits must produce a free retry")
|
||||
}
|
||||
if len(provider.requests) != 0 {
|
||||
t.Fatalf("free prune stage must not call the summarizer (%d requests)", len(provider.requests))
|
||||
}
|
||||
if state.reactiveAttempted {
|
||||
t.Fatal("a free prune retry must not consume the reactive budget")
|
||||
}
|
||||
if estimateTokens(pruned) >= estimateTokens(messages) {
|
||||
t.Fatal("prune did not shrink the history")
|
||||
}
|
||||
|
||||
// Second context-limit error on the pruned history: nothing meaningfully
|
||||
// prunable remains, so the paid summarizer fires and consumes the budget.
|
||||
_, retried, err = state.recover(context.Background(), provider, pruned, nil, "maximum context length exceeded")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !retried {
|
||||
t.Fatal("second recover must compact and retry")
|
||||
}
|
||||
if len(provider.requests) == 0 {
|
||||
t.Fatal("second recover must reach the summarizer")
|
||||
}
|
||||
if !state.reactiveAttempted {
|
||||
t.Fatal("a paid recover must consume the reactive budget")
|
||||
}
|
||||
}
|
||||
|
||||
// When pruning reclaims something but the pruned history still exceeds the
|
||||
// window, recover falls through to the paid summarizer in the SAME call
|
||||
// instead of spending a retry that would fail the same way.
|
||||
func TestRecoverFallsThroughToSummarizerWhenPruneIsInsufficient(t *testing.T) {
|
||||
provider := summaryProvider("SUMMARY")
|
||||
// A 1,000-token window: even the protected trailing output overflows it.
|
||||
state := newCompactionState(Options{ContextWindow: 1000, CompactionPreserveLast: 2}, nil)
|
||||
messages := bigToolHistory(5)
|
||||
|
||||
compacted, retried, err := state.recover(context.Background(), provider, messages, nil, "maximum context length exceeded")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !retried {
|
||||
t.Fatal("insufficient prune must still end in a retry via summarization")
|
||||
}
|
||||
if len(provider.requests) == 0 {
|
||||
t.Fatal("insufficient prune must reach the summarizer in the same recover call")
|
||||
}
|
||||
if !state.reactiveAttempted {
|
||||
t.Fatal("the paid summarizer consumes the reactive budget")
|
||||
}
|
||||
if estimateTokens(compacted) >= estimateTokens(messages) {
|
||||
t.Fatal("recover did not shrink the history")
|
||||
}
|
||||
}
|
||||
|
||||
// The free prune stage is not gated behind the one-shot budget: after a paid
|
||||
// reactive compaction, a later context-limit error still gets a free retry
|
||||
// when pruning makes the history fit.
|
||||
func TestRecoverFreePruneRunsAfterBudgetIsSpent(t *testing.T) {
|
||||
provider := summaryProvider("SUMMARY")
|
||||
state := newCompactionState(Options{ContextWindow: 200_000, CompactionPreserveLast: 2}, nil)
|
||||
state.reactiveAttempted = true
|
||||
messages := bigToolHistory(5)
|
||||
|
||||
pruned, retried, err := state.recover(context.Background(), provider, messages, nil, "maximum context length exceeded")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !retried || len(provider.requests) != 0 {
|
||||
t.Fatalf("free prune must retry without the summarizer even when the budget is spent (retried=%v, requests=%d)", retried, len(provider.requests))
|
||||
}
|
||||
if estimateTokens(pruned) >= estimateTokens(messages) {
|
||||
t.Fatal("prune did not shrink the history")
|
||||
}
|
||||
// With nothing left to prune and the budget spent, the caller keeps its
|
||||
// original error — recover neither loops nor summarizes.
|
||||
_, retried, err = state.recover(context.Background(), provider, pruned, nil, "maximum context length exceeded")
|
||||
if err != nil || retried || len(provider.requests) != 0 {
|
||||
t.Fatalf("spent budget must not summarize again (retried=%v, requests=%d, err=%v)", retried, len(provider.requests), err)
|
||||
}
|
||||
}
|
||||
|
||||
// The free-retry gate compares against the window the provider enforces, not
|
||||
// the 70% proactive threshold: a history pruned to ~80% of the window still
|
||||
// takes the free retry.
|
||||
func TestRecoverFreeRetryGateUsesWindowNotThreshold(t *testing.T) {
|
||||
provider := summaryProvider("SUMMARY")
|
||||
messages := bigToolHistory(5)
|
||||
pruned, _ := pruneStaleToolOutput(messages, 2)
|
||||
prunedTokens := estimateTokens(pruned)
|
||||
// Choose a window so the pruned size sits between 70% and 95% of it.
|
||||
window := prunedTokens * 100 / 80
|
||||
state := newCompactionState(Options{ContextWindow: window, CompactionPreserveLast: 2}, nil)
|
||||
if prunedTokens <= state.threshold {
|
||||
t.Fatalf("test setup: pruned size %d must exceed the proactive threshold %d", prunedTokens, state.threshold)
|
||||
}
|
||||
|
||||
_, retried, err := state.recover(context.Background(), provider, messages, nil, "maximum context length exceeded")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !retried || len(provider.requests) != 0 {
|
||||
t.Fatalf("a history that fits the window must take the free retry (retried=%v, requests=%d)", retried, len(provider.requests))
|
||||
}
|
||||
}
|
||||
|
||||
// After a mid-run model switch the summarizer is re-resolved against the new
|
||||
// main model: a run that started on the cheap model (no dedicated summarizer)
|
||||
// and escalated gets a cheap summarizer, and a previously broken one is
|
||||
// retried for the new model.
|
||||
func TestSwitchModelReresolvesSummarizer(t *testing.T) {
|
||||
cheap := summaryProvider("CHEAP SUMMARY")
|
||||
main := summaryProvider("MAIN SUMMARY")
|
||||
var factoryModels []string
|
||||
state := newCompactionState(Options{
|
||||
ContextWindow: 1000,
|
||||
CompactionPreserveLast: 2,
|
||||
Model: "cheap-model",
|
||||
Summarizer: func(_ context.Context, mainModelID string) (Provider, error) {
|
||||
factoryModels = append(factoryModels, mainModelID)
|
||||
if mainModelID == "cheap-model" {
|
||||
return nil, nil // main already is the cheap model
|
||||
}
|
||||
return cheap, nil
|
||||
},
|
||||
}, nil)
|
||||
|
||||
if _, dedicated := state.summarizeProvider(context.Background(), main); dedicated {
|
||||
t.Fatal("on the cheap model the main provider must summarize")
|
||||
}
|
||||
if _, dedicated := state.summarizeProvider(context.Background(), main); dedicated || len(factoryModels) != 1 {
|
||||
t.Fatalf("factory must run once per main model, calls=%v", factoryModels)
|
||||
}
|
||||
|
||||
state.switchModel("big-model", 400_000)
|
||||
if state.threshold != compactionThreshold(400_000) || state.window != 400_000 {
|
||||
t.Fatalf("switchModel must re-derive the window, got threshold=%d window=%d", state.threshold, state.window)
|
||||
}
|
||||
if got, dedicated := state.summarizeProvider(context.Background(), main); !dedicated || got != Provider(cheap) {
|
||||
t.Fatal("after escalating to an expensive model the dedicated cheap summarizer must be used")
|
||||
}
|
||||
if len(factoryModels) != 2 || factoryModels[1] != "big-model" {
|
||||
t.Fatalf("factory must be re-invoked with the new main model, calls=%v", factoryModels)
|
||||
}
|
||||
|
||||
// A broken summarizer is retried after the next switch.
|
||||
state.summarizerBroken = true
|
||||
if _, dedicated := state.summarizeProvider(context.Background(), main); dedicated {
|
||||
t.Fatal("a broken summarizer must fall back to main")
|
||||
}
|
||||
state.switchModel("other-model", 0)
|
||||
if got, dedicated := state.summarizeProvider(context.Background(), main); !dedicated || got != Provider(cheap) {
|
||||
t.Fatal("a model switch must clear the broken flag and rebuild the summarizer")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactionUpdateWindow(t *testing.T) {
|
||||
state := newCompactionState(Options{ContextWindow: 100_000, CompactionPreserveLast: 2}, nil)
|
||||
state.lowWaterMark = 50_000
|
||||
before := state.threshold
|
||||
|
||||
state.updateWindow(0) // unknown model: unchanged
|
||||
if state.threshold != before || state.lowWaterMark != 50_000 {
|
||||
t.Fatal("updateWindow(0) must be a no-op")
|
||||
}
|
||||
|
||||
state.updateWindow(10_000)
|
||||
if state.threshold >= before {
|
||||
t.Fatalf("threshold = %d, want re-derived below %d", state.threshold, before)
|
||||
}
|
||||
if state.lowWaterMark != 0 {
|
||||
t.Fatal("low-water mark must reset on a window change")
|
||||
}
|
||||
if !state.enabled {
|
||||
t.Fatal("a positive window must keep compaction enabled")
|
||||
}
|
||||
}
|
||||
|
||||
// End-to-end: after escalate_model switches providers, the loop re-derives the
|
||||
// context window via ContextWindowFor, visible in the OnContext report.
|
||||
func TestRunEscalationRederivesContextWindow(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(escalatingTool{target: "big-model"})
|
||||
switched := &mockProvider{turns: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventText, Content: "done"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
}}}
|
||||
provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call-1", ToolName: "escalate"},
|
||||
{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call-1", ArgumentsFragment: `{}`},
|
||||
{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call-1"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
}}}
|
||||
|
||||
var windows []int
|
||||
result, err := Run(context.Background(), "go", provider, Options{
|
||||
Registry: registry,
|
||||
ContextWindow: 100_000,
|
||||
Model: "small-model",
|
||||
ModelSwitcher: func(context.Context, string) (Provider, error) { return switched, nil },
|
||||
ContextWindowFor: func(modelID string) int {
|
||||
if modelID == "big-model" {
|
||||
return 400_000
|
||||
}
|
||||
return 0
|
||||
},
|
||||
OnContext: func(breakdown ContextBreakdown) { windows = append(windows, breakdown.ContextWindow) },
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.FinalAnswer != "done" {
|
||||
t.Fatalf("final answer = %q", result.FinalAnswer)
|
||||
}
|
||||
if len(windows) < 2 {
|
||||
t.Fatalf("expected at least 2 OnContext reports, got %d", len(windows))
|
||||
}
|
||||
if windows[0] != 100_000 {
|
||||
t.Fatalf("turn 1 window = %d, want the original 100000", windows[0])
|
||||
}
|
||||
if windows[len(windows)-1] != 400_000 {
|
||||
t.Fatalf("post-escalation window = %d, want 400000", windows[len(windows)-1])
|
||||
}
|
||||
}
|
||||
|
||||
// nonComparableProvider has a slice field, so comparing two Provider interface
|
||||
// values holding it would panic; the summarizer path must never compare them.
|
||||
type nonComparableProvider struct {
|
||||
inner *mockProvider
|
||||
tags []string
|
||||
}
|
||||
|
||||
func (p nonComparableProvider) StreamCompletion(ctx context.Context, request zeroruntime.CompletionRequest) (<-chan zeroruntime.StreamEvent, error) {
|
||||
return p.inner.StreamCompletion(ctx, request)
|
||||
}
|
||||
|
||||
func TestSummarizeClosureNeverComparesProviderValues(t *testing.T) {
|
||||
broken := &mockProvider{turns: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventError, Error: "boom: model not found"},
|
||||
}}}
|
||||
main := nonComparableProvider{inner: summaryProvider("MAIN SUMMARY"), tags: []string{"x"}}
|
||||
state := newCompactionState(Options{
|
||||
ContextWindow: 1000,
|
||||
CompactionPreserveLast: 2,
|
||||
Summarizer: func(context.Context, string) (Provider, error) { return nonComparableProvider{inner: broken}, nil },
|
||||
}, nil)
|
||||
summary, err := state.summarizeClosure(context.Background(), main)([]zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "history"}})
|
||||
if err != nil || summary != "MAIN SUMMARY" {
|
||||
t.Fatalf("fallback with non-comparable providers must not panic and must use main, got %q err=%v", summary, err)
|
||||
}
|
||||
// And with no dedicated summarizer at all the closure still works.
|
||||
state = newCompactionState(Options{ContextWindow: 1000, CompactionPreserveLast: 2}, nil)
|
||||
if _, err := state.summarizeClosure(context.Background(), main)([]zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "history"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
@@ -126,8 +126,11 @@ 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 json.Unmarshal([]byte(strings.TrimSpace(call.Arguments)), &arguments) == nil {
|
||||
if detail == "" && 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
|
||||
|
||||
@@ -73,6 +73,7 @@ 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"}`},
|
||||
}},
|
||||
@@ -80,7 +81,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", "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", "internal/api.go", "internal/config.go", "https://example.com/spec", "Inspect the Windows implementation"} {
|
||||
if !strings.Contains(content, want) {
|
||||
t.Fatalf("projection missing %q: %q", want, content)
|
||||
}
|
||||
|
||||
@@ -275,6 +275,27 @@ 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.",
|
||||
|
||||
@@ -102,6 +102,10 @@ 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
|
||||
}
|
||||
|
||||
@@ -15,6 +15,7 @@ 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"},
|
||||
@@ -53,10 +54,12 @@ 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" {
|
||||
|
||||
@@ -0,0 +1,378 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -33,11 +33,11 @@ const (
|
||||
// knows before spending more tool turns.
|
||||
toolOnlyProgressReminderAt = 6
|
||||
|
||||
// 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
|
||||
// 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
|
||||
|
||||
// planToolName is the planning tool the loop watches for by name.
|
||||
planToolName = "update_plan"
|
||||
@@ -591,13 +591,12 @@ func (state *guardState) planReminder(turn int) string {
|
||||
return planStaleReminder(state.toolCallsSincePlanUpdate)
|
||||
}
|
||||
|
||||
// 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.
|
||||
// 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.
|
||||
if !state.notCalledReminderSent &&
|
||||
!state.planEverCalled &&
|
||||
turn >= planReminderTurn &&
|
||||
state.totalToolCalls >= 1 {
|
||||
state.totalToolCalls >= planReminderToolThreshold {
|
||||
state.notCalledReminderSent = true
|
||||
return planNotCalledReminder()
|
||||
}
|
||||
|
||||
@@ -257,19 +257,16 @@ 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: [][]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"),
|
||||
},
|
||||
}
|
||||
provider := &mockProvider{turns: append(repeatedReadTurns(expectedThreshold), textTurn("done"))}
|
||||
|
||||
result, err := Run(context.Background(), "go", provider, Options{
|
||||
Registry: registry,
|
||||
@@ -287,6 +284,44 @@ 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")
|
||||
|
||||
+56
-17
@@ -141,6 +141,7 @@ 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))
|
||||
}
|
||||
}
|
||||
@@ -697,11 +698,12 @@ 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,
|
||||
IsError: toolResult.Status == tools.StatusError,
|
||||
ChangedFiles: append([]string(nil), toolResult.ChangedFiles...),
|
||||
Role: zeroruntime.MessageRoleTool,
|
||||
Content: toolResult.ModelOutput(),
|
||||
ToolCallID: toolResult.ToolCallID,
|
||||
ToolCallProviderID: call.ProviderCallID,
|
||||
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 —
|
||||
@@ -844,11 +846,21 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
|
||||
options.ServiceTier = ""
|
||||
planner.config.serviceTier = ""
|
||||
options.Trace.Counter(trace.CounterModelSwitches, 1)
|
||||
// KNOWN LIMITATION (deferred): the compactor's context-window budget
|
||||
// is fixed at run start from options.ContextWindow and is NOT updated
|
||||
// here, so a switch to a model with a different window keeps compacting
|
||||
// against the original budget. Fixing it needs the switcher to also
|
||||
// report the new window — out of scope for this change.
|
||||
// Re-derive the context window for the new model so compaction and
|
||||
// the OnContext budget report track the model actually in force —
|
||||
// without this a switch to a smaller-window model can overflow the
|
||||
// target, and a switch to a larger one over-compacts. An unknown
|
||||
// model (<= 0) keeps the original window. The compactor also
|
||||
// re-resolves its dedicated summarizer against the new model.
|
||||
window := 0
|
||||
if options.ContextWindowFor != nil {
|
||||
window = options.ContextWindowFor(turnRequestedModel)
|
||||
}
|
||||
if window > 0 {
|
||||
options.ContextWindow = window
|
||||
planner.config.contextWindow = window
|
||||
}
|
||||
compactor.switchModel(turnRequestedModel, window)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1035,6 +1047,9 @@ 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 == "":
|
||||
@@ -1118,8 +1133,28 @@ 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.Arguments != "" {
|
||||
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 err := decodeToolArguments(call.Arguments, &args); err != nil {
|
||||
return ToolResult{
|
||||
ToolCallID: call.ID,
|
||||
@@ -1147,7 +1182,6 @@ 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{
|
||||
@@ -3210,11 +3244,15 @@ 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 {
|
||||
return zeroruntime.ToolDefinition{
|
||||
definition := 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 {
|
||||
@@ -3369,10 +3407,11 @@ 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,
|
||||
IsError: true,
|
||||
Role: zeroruntime.MessageRoleTool,
|
||||
Content: abortedToolResultNotice,
|
||||
ToolCallID: call.ID,
|
||||
ToolCallProviderID: call.ProviderCallID,
|
||||
IsError: true,
|
||||
})
|
||||
}
|
||||
return messages
|
||||
|
||||
@@ -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{InputTokens: 100, CachedInputTokens: 20, OutputTokens: 40}},
|
||||
{Type: zeroruntime.StreamEventUsage, Usage: zeroruntime.Usage{PromptTokens: 100, CachedInputTokens: 20, CacheWriteTokens: 10, CompletionTokens: 40}},
|
||||
{Type: zeroruntime.StreamEventText, Content: "done"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
}}}
|
||||
@@ -4135,6 +4135,9 @@ 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)
|
||||
}
|
||||
|
||||
@@ -31,6 +31,19 @@ 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) {
|
||||
|
||||
@@ -58,12 +58,18 @@ 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()
|
||||
}
|
||||
|
||||
@@ -73,6 +73,38 @@ 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
|
||||
|
||||
@@ -30,13 +30,10 @@ 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 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.
|
||||
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.
|
||||
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
|
||||
@@ -55,9 +52,9 @@ work.
|
||||
edit_file, apply_patch, write_file - over shelling out to
|
||||
cat/sed/awk/python for file operations.
|
||||
They are safer, reviewable, and produce clean diffs.
|
||||
- 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.
|
||||
- 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.
|
||||
- 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.
|
||||
@@ -76,21 +73,24 @@ work.
|
||||
- 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.
|
||||
over-built; if a 200-line solution could be 50, rewrite it. Once the required
|
||||
behavior passes validation, stop instead of adding optional cleanup.
|
||||
- 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 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.
|
||||
- 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.
|
||||
- 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
|
||||
|
||||
@@ -148,6 +148,7 @@ 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 as `file:line` so it is clickable.
|
||||
short snippets. Reference code with file paths and include line numbers when
|
||||
already known; do not look them up solely for the final summary.
|
||||
- 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.
|
||||
@@ -60,12 +60,11 @@ 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. read_file, edit_file and write_file take one file per call;
|
||||
a single apply_patch may span several files when they form one coherent
|
||||
change. Do not batch file writes into a script. Independent tool calls
|
||||
(several reads, several edit_file calls to different files) belong in the
|
||||
same turn.
|
||||
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.
|
||||
- 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>`
|
||||
|
||||
@@ -29,6 +29,9 @@ 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"})
|
||||
|
||||
@@ -29,7 +29,10 @@ func TestMain(m *testing.M) {
|
||||
}
|
||||
|
||||
func TestCoreSystemPromptIncludesCodingQualityRules(t *testing.T) {
|
||||
prompt := strings.ToLower(buildSystemPrompt(Options{}))
|
||||
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)
|
||||
}
|
||||
|
||||
for _, want := range []string{
|
||||
"read-before-edit",
|
||||
@@ -44,6 +47,13 @@ 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{}))
|
||||
|
||||
+20
-4
@@ -348,10 +348,26 @@ type Options struct {
|
||||
// CompactionPreserveLast is how many trailing messages compaction keeps
|
||||
// verbatim. <= 0 falls back to defaultCompactionPreserveLast.
|
||||
CompactionPreserveLast int
|
||||
Registry *tools.Registry
|
||||
PermissionMode PermissionMode
|
||||
Autonomy string
|
||||
Sandbox *sandbox.Engine
|
||||
// Summarizer, when set, lazily builds the provider used for compaction
|
||||
// summarization calls — typically a cheap/fast model, since summaries at
|
||||
// main-model prices are the single most expensive recurring event in a
|
||||
// long run. It receives the main model currently in force and may return
|
||||
// (nil, nil) when no dedicated summarizer applies to it (the main model is
|
||||
// already the cheap one). Built on the first paid compaction and again
|
||||
// after a mid-run model switch. Any failure (build or call) falls back to
|
||||
// the run's main provider until the next switch, so a misconfigured
|
||||
// summarizer can never break compaction. nil keeps today's behavior.
|
||||
Summarizer func(ctx context.Context, mainModelID string) (Provider, error)
|
||||
// ContextWindowFor, when set, resolves a model ID to its context window so
|
||||
// the compactor can re-derive its threshold after a mid-run model switch
|
||||
// (escalate_model). Without it a switch keeps compacting against the
|
||||
// original model's window — overflowing a smaller target or over-compacting
|
||||
// a larger one. Return <= 0 when the model is unknown (window unchanged).
|
||||
ContextWindowFor func(modelID string) int
|
||||
Registry *tools.Registry
|
||||
PermissionMode PermissionMode
|
||||
Autonomy string
|
||||
Sandbox *sandbox.Engine
|
||||
// FileTracker records per-session file read/write versions so the write tools
|
||||
// can detect a file changed on disk outside Zero since it was last read. nil
|
||||
// disables the check. Created once per session and threaded into every tool run.
|
||||
|
||||
+49
-12
@@ -839,18 +839,20 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
|
||||
mcpTokenStore = nil
|
||||
err = nil
|
||||
}
|
||||
criticalMCPConfig, optionalMCPConfig := splitMCPStartupConfig(mcpConfig)
|
||||
mcpRuntime := mcpToolRuntime(noopMCPRuntime{})
|
||||
if len(mcpConfig.Servers) > 0 {
|
||||
mcpRuntime, err = deps.registerMCPTools(context.Background(), registry, mcpConfig, mcp.RegisterOptions{
|
||||
if len(criticalMCPConfig.Servers) > 0 {
|
||||
runtime, registerErr := deps.registerMCPTools(context.Background(), registry, criticalMCPConfig, mcp.RegisterOptions{
|
||||
PermissionStore: mcpPermissionStore,
|
||||
Autonomy: mcp.AutonomyLow,
|
||||
Execution: executionRunner,
|
||||
WorkspaceRoot: workspaceRoot,
|
||||
})
|
||||
}
|
||||
if err != nil {
|
||||
closeMCPRuntime(stderr, mcpRuntime)
|
||||
return writeAppError(stderr, err.Error(), 1)
|
||||
if registerErr != nil {
|
||||
closeMCPRuntime(stderr, runtime)
|
||||
return writeAppError(stderr, redaction.ErrorMessage(registerErr, redaction.Options{}), 1)
|
||||
}
|
||||
mcpRuntime = runtime
|
||||
}
|
||||
defer closeMCPRuntime(stderr, mcpRuntime)
|
||||
// A server that could not be reached or validated is skipped, not fatal (one
|
||||
@@ -909,11 +911,35 @@ 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.
|
||||
// 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.
|
||||
// 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.
|
||||
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
|
||||
@@ -948,10 +974,20 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
|
||||
FavoriteModels: resolved.Preferences.FavoriteModels,
|
||||
RecentModels: resolved.Preferences.RecentModels,
|
||||
RecapsEnabled: resolved.Preferences.RecapsEnabled(),
|
||||
CompactionModel: resolved.Preferences.CompactionModel,
|
||||
Provider: provider,
|
||||
NewProvider: deps.newProvider,
|
||||
ProbeProviderHealth: deps.probeProviderHealth,
|
||||
UserAgent: userAgent(),
|
||||
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(),
|
||||
PrepareRunCompletionWarning: func() {
|
||||
scratchBaseline = scratchFileSnapshot(workspaceRoot)
|
||||
},
|
||||
@@ -959,6 +995,7 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
|
||||
return scratchFileWarning(workspaceRoot, scratchBaseline)
|
||||
},
|
||||
Registry: registry,
|
||||
AwaitToolReadiness: awaitToolReadiness,
|
||||
SessionStore: sessionStore,
|
||||
PeerService: peerService,
|
||||
SandboxStore: sandboxStore,
|
||||
|
||||
@@ -17,6 +17,7 @@ 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"
|
||||
@@ -486,6 +487,74 @@ 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
|
||||
@@ -565,6 +634,7 @@ 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) {
|
||||
@@ -591,7 +661,7 @@ func TestRunNoArgsClosesPartialMCPRuntimeWhenRegistrationFails(t *testing.T) {
|
||||
return closeFunc(func() error {
|
||||
runtimeClosed = true
|
||||
return nil
|
||||
}), errors.New("register mcp tools failed")
|
||||
}), fmt.Errorf("register mcp tools failed: %s", secret)
|
||||
},
|
||||
runTUI: func(ctx context.Context, options tui.Options) int {
|
||||
t.Fatal("TUI should not launch after MCP registration fails")
|
||||
@@ -605,6 +675,12 @@ 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")
|
||||
}
|
||||
|
||||
+17
-10
@@ -682,16 +682,20 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
|
||||
ModelSwitcher: modelSwitcher,
|
||||
TurnSessionProvider: turnSessions,
|
||||
ModelSessionSwitcher: modelSessionSwitcher,
|
||||
ReasoningEffort: forwardEffort,
|
||||
Trace: traceRecorder,
|
||||
Cwd: workspaceRoot,
|
||||
Images: images,
|
||||
Registry: registry,
|
||||
PermissionMode: permissionMode,
|
||||
Autonomy: options.autonomy,
|
||||
SelfCorrect: selfCorrector,
|
||||
FileDiagnostics: fileDiagnostics,
|
||||
Profile: execProfile.Policy(displacedMaxTurns, execProfileFilledEffort),
|
||||
Summarizer: summarizerFactory(resolved, deps.newProvider),
|
||||
ContextWindowFor: func(modelID string) int {
|
||||
return modelregistry.AgentContextWindow(modelContextWindow(modelRegistry, modelID))
|
||||
},
|
||||
ReasoningEffort: forwardEffort,
|
||||
Trace: traceRecorder,
|
||||
Cwd: workspaceRoot,
|
||||
Images: images,
|
||||
Registry: registry,
|
||||
PermissionMode: permissionMode,
|
||||
Autonomy: options.autonomy,
|
||||
SelfCorrect: selfCorrector,
|
||||
FileDiagnostics: fileDiagnostics,
|
||||
Profile: execProfile.Policy(displacedMaxTurns, execProfileFilledEffort),
|
||||
// Headless exec: don't accept a no-tool-call turn as "done" while work
|
||||
// clearly remains (pending plan items / a mid-step continuation cue) —
|
||||
// nudge to continue, and finalize as INCOMPLETE rather than false success
|
||||
@@ -939,6 +943,9 @@ 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
|
||||
}
|
||||
|
||||
@@ -267,23 +267,29 @@ 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,
|
||||
"total_tokens": usage.TotalTokens(),
|
||||
"type": "usage",
|
||||
"prompt_tokens": usage.PromptTokens,
|
||||
"completion_tokens": usage.CompletionTokens,
|
||||
"cached_input_tokens": usage.CachedInputTokens,
|
||||
"cache_write_tokens": usage.CacheWriteTokens,
|
||||
"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,
|
||||
TotalTokens: &totalTokens,
|
||||
Type: streamjson.EventUsage,
|
||||
RunID: writer.runID,
|
||||
PromptTokens: &promptTokens,
|
||||
CompletionTokens: &completionTokens,
|
||||
CachedInputTokens: &cachedInputTokens,
|
||||
CacheWriteTokens: &cacheWriteTokens,
|
||||
TotalTokens: &totalTokens,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -52,3 +52,39 @@ 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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,147 @@
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
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")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/agent"
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/providers"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// summarizerFactory adapts the resolved profile and the authenticated provider
|
||||
// builder into agent.Options.Summarizer; the selection rules live in
|
||||
// providers.CompactionSummarizerFactory, shared with the TUI.
|
||||
func summarizerFactory(resolved config.ResolvedConfig, newProvider func(config.ProviderProfile) (zeroruntime.Provider, error)) func(context.Context, string) (agent.Provider, error) {
|
||||
return providers.CompactionSummarizerFactory(resolved.Provider, resolved.Preferences.CompactionModel, newProvider)
|
||||
}
|
||||
@@ -119,6 +119,12 @@ type PreferencesConfig struct {
|
||||
// the user turned idle recaps off. A *bool is its own tri-state, so no
|
||||
// custom unmarshal is needed (unlike ToolsConfig.DeferThreshold's int).
|
||||
Recaps *bool `json:"recaps,omitempty"`
|
||||
// CompactionModel routes compaction summarization calls to this model
|
||||
// instead of the session's main model (summaries at main-model prices are
|
||||
// the most expensive recurring event in long runs). Empty = automatic: a
|
||||
// curated cheap model on official endpoints, the main model elsewhere.
|
||||
// "main" forces the main model. ZERO_COMPACTION_MODEL overrides.
|
||||
CompactionModel string `json:"compactionModel,omitempty"`
|
||||
}
|
||||
|
||||
// RecentModelEntry is one provider-qualified model selection recorded in
|
||||
|
||||
+9
-66
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -3,13 +3,93 @@
|
||||
package execution
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"os"
|
||||
"os/exec"
|
||||
"strconv"
|
||||
"strings"
|
||||
"syscall"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// requireProcessState refuses to let a broken ps pipeline masquerade as passing
|
||||
// coverage. Every test in this file proves something about zombie detection,
|
||||
// and every one of those proofs is read out of ps — so if ps cannot report
|
||||
// state, the tests verify nothing whatever they do next.
|
||||
//
|
||||
// The old code answered that with t.Skip, which is indistinguishable from
|
||||
// success in CI output. This probes BOTH ps pipelines against this test
|
||||
// process, which is by definition alive and not a zombie, so a working ps must
|
||||
// report a live state for it. Failing that:
|
||||
//
|
||||
// - In CI (or with ZERO_REQUIRE_PS=1) it is a hard failure. ps is present on
|
||||
// every image used today, so its absence means the pipeline regressed and
|
||||
// that has to surface immediately rather than as skipped-green.
|
||||
// - Elsewhere it stays a skip, so a contributor on an exotic box is not
|
||||
// blocked by an environment the project does not claim to support.
|
||||
//
|
||||
// This only covers "ps cannot report state at all". Once it passes, a later
|
||||
// failure to observe an EXPECTED state is a real defect, and the tests below
|
||||
// treat it as one.
|
||||
func requireProcessState(t *testing.T) {
|
||||
t.Helper()
|
||||
self := os.Getpid()
|
||||
var reasons []string
|
||||
// The production path: ps -A -o pid=,pgid=,stat=, parsed by
|
||||
// processTableStateMatches. This is what signalTargetRunning consults.
|
||||
if running, ok := processIsRunning(self); !ok || !running {
|
||||
reasons = append(reasons, "the process table read used by signalTargetRunning reported no live state for this test process")
|
||||
}
|
||||
// This file's independent oracle: ps -o stat= -p <pid>. It is deliberately
|
||||
// NOT the production parser, so a regression in that parser cannot hide by
|
||||
// breaking the test's own precondition check in the same way.
|
||||
if zombie, ok := processIsZombie(self); !ok || zombie {
|
||||
reasons = append(reasons, "the per-PID state read used by this test file reported no live state for this test process")
|
||||
}
|
||||
if len(reasons) == 0 {
|
||||
return
|
||||
}
|
||||
reason := "cannot read process state, so ps-based zombie detection is unverifiable here: " + strings.Join(reasons, "; ")
|
||||
if requireWorkingPS() {
|
||||
t.Fatal(reason + " — failing rather than skipping because CI (or ZERO_REQUIRE_PS) is set, where ps is expected to work")
|
||||
}
|
||||
t.Skip(reason)
|
||||
}
|
||||
|
||||
// requireWorkingPS reports whether an unusable ps must fail the run instead of
|
||||
// skipping it. ZERO_REQUIRE_PS lets a developer opt into the strict behavior
|
||||
// locally to reproduce what CI does.
|
||||
func requireWorkingPS() bool {
|
||||
if strings.TrimSpace(os.Getenv("ZERO_REQUIRE_PS")) != "" {
|
||||
return true
|
||||
}
|
||||
return strings.TrimSpace(os.Getenv("CI")) != ""
|
||||
}
|
||||
|
||||
// awaitZombie blocks until pid is an unreaped zombie. The caller must not have
|
||||
// waited on the child, so the zombie is a stable state rather than a race: it
|
||||
// persists until something reaps it. Failing to reach it therefore means the
|
||||
// child never exited or ps stopped reporting its state, and both are defects
|
||||
// rather than reasons to skip.
|
||||
func awaitZombie(t *testing.T, pid int) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for {
|
||||
zombie, ok := processIsZombie(pid)
|
||||
if ok && zombie {
|
||||
return
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
if !ok {
|
||||
t.Fatalf("pid %d left the process table without being reaped; this test never waited on it, so its zombie row must persist", pid)
|
||||
}
|
||||
t.Fatalf("pid %d did not reach the zombie state within 5s", pid)
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTerminateProcessTreeTreatsZombieGroupAsExited is the regression test for
|
||||
// counting unreaped processes as alive: kill(2) with signal 0 succeeds for a
|
||||
// zombie, so a group whose only remaining member is waiting to be reaped used to
|
||||
@@ -17,6 +97,8 @@ import (
|
||||
// as "did not exit after SIGKILL" — a spurious cleanup failure whose timing
|
||||
// depended on when the reaper ran.
|
||||
func TestTerminateProcessTreeTreatsZombieGroupAsExited(t *testing.T) {
|
||||
requireProcessState(t)
|
||||
|
||||
cmd := exec.Command("sh", "-c", "exit 0")
|
||||
ConfigureProcessGroup(cmd)
|
||||
if err := cmd.Start(); err != nil {
|
||||
@@ -27,23 +109,59 @@ func TestTerminateProcessTreeTreatsZombieGroupAsExited(t *testing.T) {
|
||||
// is exactly the state a terminated tree passes through.
|
||||
t.Cleanup(func() { _ = cmd.Wait() })
|
||||
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for {
|
||||
if zombie, ok := processIsZombie(pid); ok && zombie {
|
||||
break
|
||||
awaitZombie(t, pid)
|
||||
|
||||
// Whether a zombie-only group can be signalled is a KERNEL difference, not a
|
||||
// ps one. On Linux the zombie still occupies the process table and
|
||||
// kill(-pid, 0) succeeds, so signalTargetRunning goes on to consult ps —
|
||||
// which is the branch this test was written for. Darwin returns EPERM for
|
||||
// the same call, so signalTargetRunning short-circuits on the kill and never
|
||||
// reaches the ps check.
|
||||
//
|
||||
// The OUTCOME is identical and is asserted on every platform below. Only the
|
||||
// branch that produces it differs, so record that instead of skipping the
|
||||
// whole test (which would drop Darwin's coverage of the outcome) or failing
|
||||
// it (which would report a supported platform as broken). Any errno other
|
||||
// than EPERM means the group is not in the state this test needs.
|
||||
psPathReachable := true
|
||||
if err := syscall.Kill(-pid, syscall.Signal(0)); err != nil {
|
||||
if !errors.Is(err, syscall.EPERM) {
|
||||
t.Fatalf("zombie group %d is unsignalable with an unexpected error (%v); the zombie case is not present, so nothing below would be meaningful", pid, err)
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Skip("child did not reach the zombie state; ps output is unavailable in this environment")
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
if syscall.Kill(-pid, syscall.Signal(0)) != nil {
|
||||
t.Skip("the zombie group is no longer signalable; cannot exercise the case")
|
||||
psPathReachable = false
|
||||
t.Logf("kernel refuses to signal zombie-only group %d (EPERM); signalTargetRunning answers from kill(2) rather than ps here, so the ps-based zombie check is asserted directly on platforms that allow the signal", pid)
|
||||
}
|
||||
|
||||
// Where the signal IS permitted, pin the ps-based check itself. This is the
|
||||
// regression #774 was about, and asserting only signalTargetRunning would
|
||||
// let a ps-side regression hide behind the kill result.
|
||||
if psPathReachable {
|
||||
running, ok := processGroupHasRunningMember(pid)
|
||||
if !ok {
|
||||
t.Fatalf("process group %d has no rows while its zombie is unreaped; ps stopped reporting the group rather than the group being gone", pid)
|
||||
}
|
||||
if running {
|
||||
t.Fatal("the ps-based group check must not report a zombie-only group as running")
|
||||
}
|
||||
}
|
||||
|
||||
// Valid on every platform: whichever branch answers, a zombie-only group is
|
||||
// not running. On a kernel that refuses the signal this is the whole of the
|
||||
// coverage, and it is more than the previous skip provided.
|
||||
if signalTargetRunning(-pid) {
|
||||
t.Fatal("a group holding only a zombie must not count as running")
|
||||
}
|
||||
|
||||
if !psPathReachable {
|
||||
// terminateTarget's own kill(2) would hit the same EPERM and be returned
|
||||
// as a cleanup error, so this assertion cannot be made where the kernel
|
||||
// refuses the signal. That may be a real gap in TerminateProcessTree on
|
||||
// such a platform rather than a test limitation — a zombie-only tree is
|
||||
// already exited and should terminate cleanly — but confirming and
|
||||
// changing that is a production fix, not this test's business.
|
||||
t.Log("skipping the TerminateProcessTree assertion: its kill(2) would hit the same EPERM as the probe above")
|
||||
return
|
||||
}
|
||||
if err := TerminateProcessTree(pid, 50*time.Millisecond, 10*time.Millisecond); err != nil {
|
||||
t.Fatalf("TerminateProcessTree on an already-exited tree: %v", err)
|
||||
}
|
||||
@@ -53,6 +171,8 @@ func TestTerminateProcessTreeTreatsZombieGroupAsExited(t *testing.T) {
|
||||
// a live member must still be seen as running and then actually stopped, so the
|
||||
// zombie tolerance above cannot degrade into ignoring real processes.
|
||||
func TestTerminateProcessTreeStopsRunningGroup(t *testing.T) {
|
||||
requireProcessState(t)
|
||||
|
||||
cmd := exec.Command("sh", "-c", "sleep 30")
|
||||
ConfigureProcessGroup(cmd)
|
||||
if err := cmd.Start(); err != nil {
|
||||
@@ -67,7 +187,15 @@ func TestTerminateProcessTreeStopsRunningGroup(t *testing.T) {
|
||||
if err := TerminateProcessTree(pid, time.Second, 10*time.Millisecond); err != nil {
|
||||
t.Fatalf("TerminateProcessTree: %v", err)
|
||||
}
|
||||
if running, ok := processGroupHasRunningMember(pid); ok && running {
|
||||
// ok reports whether ANY row for this group was found. This test never waits
|
||||
// on the child, so its zombie row must still be there — an unknown state
|
||||
// means ps stopped answering, not that the group is gone. Treating that as a
|
||||
// pass is what let this assertion succeed without checking anything.
|
||||
running, ok := processGroupHasRunningMember(pid)
|
||||
if !ok {
|
||||
t.Fatalf("process group %d has no rows after termination; the unreaped child's zombie row must persist, so the group's state is undeterminable rather than clean", pid)
|
||||
}
|
||||
if running {
|
||||
t.Fatal("the group still has a running member after termination")
|
||||
}
|
||||
}
|
||||
@@ -82,6 +210,8 @@ func TestTerminateProcessTreeStopsRunningGroup(t *testing.T) {
|
||||
// "unknown" resulted every time, and signalTargetRunning conservatively (and
|
||||
// incorrectly) treated a zombie individual-PID target as still running.
|
||||
func TestSignalTargetRunningTreatsZombieIndividualPIDAsExited(t *testing.T) {
|
||||
requireProcessState(t)
|
||||
|
||||
cmd := exec.Command("sh", "-c", "exit 0")
|
||||
// Deliberately do NOT call ConfigureProcessGroup: the child inherits this
|
||||
// test process's group, so it is not its own leader and processSignalTarget
|
||||
@@ -92,40 +222,38 @@ func TestSignalTargetRunningTreatsZombieIndividualPIDAsExited(t *testing.T) {
|
||||
pid := cmd.Process.Pid
|
||||
t.Cleanup(func() { _ = cmd.Wait() })
|
||||
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for {
|
||||
if zombie, ok := processIsZombie(pid); ok && zombie {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Skip("child did not reach the zombie state; ps output is unavailable in this environment")
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
awaitZombie(t, pid)
|
||||
|
||||
target, err := processSignalTarget(pid)
|
||||
if err != nil {
|
||||
t.Fatalf("processSignalTarget: %v", err)
|
||||
}
|
||||
// The child inherited this process's group, so its PGID is this process's
|
||||
// PID and can never equal its own PID. A negative target here would mean
|
||||
// processSignalTarget misread the group, which is the very thing the
|
||||
// individual-PID fallback exists to handle.
|
||||
if target != pid {
|
||||
t.Skip("child unexpectedly became its own group leader; cannot exercise the individual-PID fallback")
|
||||
t.Fatalf("processSignalTarget(%d) = %d, want the individual PID: the child inherited this process's group, so it is not its own leader", pid, target)
|
||||
}
|
||||
if signalTargetRunning(target) {
|
||||
t.Fatal("a zombie individual-PID target must not count as running")
|
||||
}
|
||||
}
|
||||
|
||||
// processIsZombie reports a process's zombie state via ps. ok is false when the
|
||||
// state could not be read.
|
||||
// processIsZombie reports a process's zombie state via ps. It is deliberately a
|
||||
// separate, per-PID ps invocation rather than a reuse of the production
|
||||
// process-table parser: these tests assert what that parser concludes, so
|
||||
// borrowing it to establish their own preconditions would let one regression
|
||||
// silence both sides at once.
|
||||
//
|
||||
// ok is false when no state could be read, which after requireProcessState has
|
||||
// passed means the process left the table rather than that ps is broken.
|
||||
func processIsZombie(pid int) (zombie bool, ok bool) {
|
||||
output, err := exec.Command("ps", "-o", "stat=", "-p", strconv.Itoa(pid)).Output()
|
||||
if err != nil {
|
||||
return false, false
|
||||
}
|
||||
state := string(output)
|
||||
for len(state) > 0 && (state[0] == ' ' || state[0] == '\n' || state[0] == '\t') {
|
||||
state = state[1:]
|
||||
}
|
||||
state := strings.TrimLeft(string(output), " \n\t")
|
||||
if state == "" {
|
||||
return false, false
|
||||
}
|
||||
|
||||
+9
-63
@@ -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
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
+100
-14
@@ -63,8 +63,9 @@ func tokenToStored(t oauth.Token) StoredToken {
|
||||
}
|
||||
|
||||
// OAuthConfig describes how to authenticate to a remote MCP server using OAuth.
|
||||
// Endpoints may be discovered from the server's metadata document; explicit
|
||||
// values here override or fill in anything discovery cannot provide.
|
||||
// Endpoints may be discovered through RFC 9728 protected-resource metadata and
|
||||
// RFC 8414 authorization-server metadata; explicit values override or fill in
|
||||
// anything discovery cannot provide.
|
||||
type OAuthConfig struct {
|
||||
ClientID string
|
||||
ClientSecret string
|
||||
@@ -85,6 +86,11 @@ type authServerMetadata struct {
|
||||
TokenEndpoint string `json:"token_endpoint"`
|
||||
RegistrationEndpoint string `json:"registration_endpoint"`
|
||||
ScopesSupported []string `json:"scopes_supported"`
|
||||
|
||||
protectedResourceURL string
|
||||
protectAuthorizationEndpoint bool
|
||||
protectTokenEndpoint bool
|
||||
protectRegistrationEndpoint bool
|
||||
}
|
||||
|
||||
// pkceParams holds a PKCE verifier/challenge pair.
|
||||
@@ -120,7 +126,7 @@ type authorizationFlow struct {
|
||||
// discoverAuthorizationServer fetches the RFC 8414 authorization server metadata
|
||||
// at the well-known path under baseURL, via the shared engine.
|
||||
func discoverAuthorizationServer(ctx context.Context, client *http.Client, baseURL string) (authServerMetadata, error) {
|
||||
meta, err := oauth.DiscoverAuthorizationServer(ctx, client, baseURL)
|
||||
meta, err := oauth.DiscoverAuthorizationServer(ctx, withoutDiscoveryRedirects(client), baseURL)
|
||||
if err != nil {
|
||||
return authServerMetadata{}, err
|
||||
}
|
||||
@@ -133,10 +139,9 @@ func discoverAuthorizationServer(ctx context.Context, client *http.Client, baseU
|
||||
}, nil
|
||||
}
|
||||
|
||||
// validateResolvedEndpoints applies the shared https/loopback endpoint rule to
|
||||
// every non-empty endpoint in the resolved metadata — configured or discovered
|
||||
// alike — so discovery metadata can never downgrade an MCP login to an insecure
|
||||
// or attacker-controlled authorization, token, or registration endpoint.
|
||||
// validateResolvedEndpoints applies the shared https/loopback scheme rule to
|
||||
// every non-empty endpoint in the resolved metadata. Protected-resource
|
||||
// discovery adds its stricter provenance-aware network rule below.
|
||||
func validateResolvedEndpoints(metadata authServerMetadata) error {
|
||||
for _, ep := range []struct{ kind, url string }{
|
||||
{"authorization", metadata.AuthorizationEndpoint},
|
||||
@@ -153,9 +158,43 @@ func validateResolvedEndpoints(metadata authServerMetadata) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// resolveAuthorizationServer discovers metadata and applies explicit config
|
||||
// overrides. Configured endpoints take precedence over discovered ones, and act
|
||||
// as a fallback when discovery fails or omits a value.
|
||||
func validateProtectedDiscoveredEndpoints(ctx context.Context, metadata authServerMetadata) error {
|
||||
if metadata.protectedResourceURL == "" {
|
||||
return nil
|
||||
}
|
||||
policy, err := newOAuthDiscoveryNetworkPolicy(metadata.protectedResourceURL)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, endpoint := range []struct {
|
||||
kind string
|
||||
url string
|
||||
protected bool
|
||||
}{
|
||||
{kind: "authorization", url: metadata.AuthorizationEndpoint, protected: metadata.protectAuthorizationEndpoint},
|
||||
{kind: "token", url: metadata.TokenEndpoint, protected: metadata.protectTokenEndpoint},
|
||||
{kind: "registration", url: metadata.RegistrationEndpoint, protected: metadata.protectRegistrationEndpoint},
|
||||
} {
|
||||
if !endpoint.protected || strings.TrimSpace(endpoint.url) == "" {
|
||||
continue
|
||||
}
|
||||
parsed, parseErr := url.Parse(endpoint.url)
|
||||
if parseErr != nil || parsed.Hostname() == "" {
|
||||
return fmt.Errorf("mcp oauth: %s endpoint: %w", endpoint.kind, errUnsafeOAuthDiscoveryTarget)
|
||||
}
|
||||
if err := policy.validateLiteralHost(parsed.Hostname()); err != nil {
|
||||
return fmt.Errorf("mcp oauth: %s endpoint: %w", endpoint.kind, err)
|
||||
}
|
||||
if _, err := policy.resolve(ctx, parsed.Hostname()); err != nil {
|
||||
return fmt.Errorf("mcp oauth: %s endpoint: %w", endpoint.kind, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// resolveAuthorizationServer follows protected-resource metadata when the MCP
|
||||
// endpoint advertises it, falls back to direct authorization-server discovery
|
||||
// for legacy providers, then applies explicit config overrides.
|
||||
func resolveAuthorizationServer(ctx context.Context, client *http.Client, baseURL string, cfg OAuthConfig) (authServerMetadata, error) {
|
||||
// When the config supplies both the authorization and token endpoints
|
||||
// directly, skip network discovery entirely: there is nothing to discover,
|
||||
@@ -174,17 +213,40 @@ func resolveAuthorizationServer(ctx context.Context, client *http.Client, baseUR
|
||||
return metadata, nil
|
||||
}
|
||||
|
||||
discoveryClient := withoutDiscoveryRedirects(client)
|
||||
discoveryBase := strings.TrimSpace(cfg.IssuerURL)
|
||||
protectedResourceDiscovery := false
|
||||
if discoveryBase == "" {
|
||||
discoveryBase = baseURL
|
||||
issuer, found, protectedErr := discoverProtectedResourceAuthorizationServer(ctx, discoveryClient, baseURL)
|
||||
if protectedErr != nil {
|
||||
return authServerMetadata{}, protectedErr
|
||||
}
|
||||
if found {
|
||||
discoveryBase = issuer
|
||||
protectedResourceDiscovery = true
|
||||
discoveryClient, protectedErr = newAdvertisedOAuthDiscoveryClient(discoveryClient, baseURL)
|
||||
if protectedErr != nil {
|
||||
return authServerMetadata{}, protectedErr
|
||||
}
|
||||
} else {
|
||||
discoveryBase = baseURL
|
||||
}
|
||||
}
|
||||
|
||||
metadata, err := discoverAuthorizationServer(ctx, client, discoveryBase)
|
||||
metadata, err := discoverAuthorizationServer(ctx, discoveryClient, discoveryBase)
|
||||
if err != nil {
|
||||
// Discovery failures are non-fatal when the config supplies the endpoints
|
||||
// directly; otherwise surface the discovery error.
|
||||
metadata = authServerMetadata{}
|
||||
}
|
||||
if protectedResourceDiscovery {
|
||||
if strings.TrimSpace(metadata.Issuer) == "" {
|
||||
return authServerMetadata{}, errors.New("mcp oauth: authorization server metadata is missing its required issuer")
|
||||
}
|
||||
if metadata.Issuer != discoveryBase {
|
||||
return authServerMetadata{}, errors.New("mcp oauth: authorization server metadata issuer does not match protected resource metadata")
|
||||
}
|
||||
}
|
||||
|
||||
if endpoint := strings.TrimSpace(cfg.AuthorizationEndpoint); endpoint != "" {
|
||||
metadata.AuthorizationEndpoint = endpoint
|
||||
@@ -195,6 +257,12 @@ func resolveAuthorizationServer(ctx context.Context, client *http.Client, baseUR
|
||||
if endpoint := strings.TrimSpace(cfg.RegistrationEndpoint); endpoint != "" {
|
||||
metadata.RegistrationEndpoint = endpoint
|
||||
}
|
||||
if protectedResourceDiscovery {
|
||||
metadata.protectedResourceURL = strings.TrimSpace(baseURL)
|
||||
metadata.protectAuthorizationEndpoint = strings.TrimSpace(cfg.AuthorizationEndpoint) == ""
|
||||
metadata.protectTokenEndpoint = strings.TrimSpace(cfg.TokenEndpoint) == ""
|
||||
metadata.protectRegistrationEndpoint = strings.TrimSpace(cfg.RegistrationEndpoint) == ""
|
||||
}
|
||||
|
||||
if strings.TrimSpace(metadata.AuthorizationEndpoint) == "" {
|
||||
return authServerMetadata{}, errors.New("no authorization endpoint discovered or configured")
|
||||
@@ -205,6 +273,9 @@ func resolveAuthorizationServer(ctx context.Context, client *http.Client, baseUR
|
||||
if err := validateResolvedEndpoints(metadata); err != nil {
|
||||
return authServerMetadata{}, err
|
||||
}
|
||||
if err := validateProtectedDiscoveredEndpoints(ctx, metadata); err != nil {
|
||||
return authServerMetadata{}, err
|
||||
}
|
||||
return metadata, nil
|
||||
}
|
||||
|
||||
@@ -363,10 +434,25 @@ func Login(ctx context.Context, options LoginOptions) (StoredToken, error) {
|
||||
}
|
||||
defer listener.Close()
|
||||
redirectURI := fmt.Sprintf("http://%s/callback", listener.Addr().String())
|
||||
protectedEndpointClient := httpClient
|
||||
if metadata.protectRegistrationEndpoint || metadata.protectTokenEndpoint {
|
||||
protectedEndpointClient, err = newAdvertisedOAuthDiscoveryClient(httpClient, metadata.protectedResourceURL)
|
||||
if err != nil {
|
||||
return StoredToken{}, err
|
||||
}
|
||||
}
|
||||
tokenClient := httpClient
|
||||
if metadata.protectTokenEndpoint {
|
||||
tokenClient = protectedEndpointClient
|
||||
}
|
||||
|
||||
if strings.TrimSpace(cfg.ClientID) == "" {
|
||||
if registration := strings.TrimSpace(metadata.RegistrationEndpoint); registration != "" {
|
||||
clientID, clientSecret, regErr := registerClient(loginCtx, httpClient, registration, redirectURI, cfg.Scopes)
|
||||
registrationClient := httpClient
|
||||
if metadata.protectRegistrationEndpoint {
|
||||
registrationClient = protectedEndpointClient
|
||||
}
|
||||
clientID, clientSecret, regErr := registerClient(loginCtx, registrationClient, registration, redirectURI, cfg.Scopes)
|
||||
if regErr != nil {
|
||||
return StoredToken{}, regErr
|
||||
}
|
||||
@@ -390,7 +476,7 @@ func Login(ctx context.Context, options LoginOptions) (StoredToken, error) {
|
||||
}
|
||||
|
||||
flow := &authorizationFlow{
|
||||
httpClient: httpClient,
|
||||
httpClient: tokenClient,
|
||||
metadata: metadata,
|
||||
config: cfg,
|
||||
pkce: pkce,
|
||||
|
||||
@@ -0,0 +1,290 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
var errUnsafeOAuthDiscoveryTarget = errors.New("mcp oauth: refused non-public server-advertised discovery target")
|
||||
|
||||
var nonPublicOAuthDiscoveryPrefixes = []netip.Prefix{
|
||||
netip.MustParsePrefix("0.0.0.0/8"),
|
||||
netip.MustParsePrefix("10.0.0.0/8"),
|
||||
netip.MustParsePrefix("100.64.0.0/10"),
|
||||
netip.MustParsePrefix("127.0.0.0/8"),
|
||||
netip.MustParsePrefix("169.254.0.0/16"),
|
||||
netip.MustParsePrefix("172.16.0.0/12"),
|
||||
netip.MustParsePrefix("192.0.0.0/24"),
|
||||
netip.MustParsePrefix("192.0.2.0/24"),
|
||||
netip.MustParsePrefix("192.88.99.0/24"),
|
||||
netip.MustParsePrefix("192.168.0.0/16"),
|
||||
netip.MustParsePrefix("198.18.0.0/15"),
|
||||
netip.MustParsePrefix("198.51.100.0/24"),
|
||||
netip.MustParsePrefix("203.0.113.0/24"),
|
||||
netip.MustParsePrefix("224.0.0.0/4"),
|
||||
netip.MustParsePrefix("240.0.0.0/4"),
|
||||
netip.MustParsePrefix("::/128"),
|
||||
netip.MustParsePrefix("::1/128"),
|
||||
netip.MustParsePrefix("100::/64"),
|
||||
netip.MustParsePrefix("2001:2::/48"),
|
||||
netip.MustParsePrefix("2001:db8::/32"),
|
||||
netip.MustParsePrefix("fc00::/7"),
|
||||
netip.MustParsePrefix("fe80::/10"),
|
||||
netip.MustParsePrefix("ff00::/8"),
|
||||
}
|
||||
|
||||
type oauthDiscoveryNetworkPolicy struct {
|
||||
allowLoopback bool
|
||||
lookupNetIP func(context.Context, string, string) ([]netip.Addr, error)
|
||||
}
|
||||
|
||||
type advertisedOAuthDiscoveryTransport struct {
|
||||
base http.RoundTripper
|
||||
policy oauthDiscoveryNetworkPolicy
|
||||
}
|
||||
|
||||
// newAdvertisedOAuthDiscoveryClient binds server-advertised discovery to the
|
||||
// public network. Loopback targets remain available only when the operator
|
||||
// explicitly configured a loopback MCP resource.
|
||||
func newAdvertisedOAuthDiscoveryClient(client *http.Client, resourceURL string) (*http.Client, error) {
|
||||
policy, err := newOAuthDiscoveryNetworkPolicy(resourceURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
copied := withoutDiscoveryRedirects(client)
|
||||
transport := copied.Transport
|
||||
if transport == nil {
|
||||
transport = http.DefaultTransport
|
||||
}
|
||||
copied.Transport = advertisedOAuthDiscoveryTransport{base: transport, policy: policy}
|
||||
return copied, nil
|
||||
}
|
||||
|
||||
func newOAuthDiscoveryNetworkPolicy(resourceURL string) (oauthDiscoveryNetworkPolicy, error) {
|
||||
parsed, err := url.Parse(strings.TrimSpace(resourceURL))
|
||||
if err != nil || parsed.Hostname() == "" {
|
||||
return oauthDiscoveryNetworkPolicy{}, errors.New("mcp oauth: invalid MCP resource URL for discovery policy")
|
||||
}
|
||||
return oauthDiscoveryNetworkPolicy{allowLoopback: isExplicitOAuthLoopbackHost(parsed.Hostname())}, nil
|
||||
}
|
||||
|
||||
func validateAdvertisedOAuthDiscoveryURL(resourceURL string, targetURL string) error {
|
||||
policy, err := newOAuthDiscoveryNetworkPolicy(resourceURL)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
parsed, err := url.Parse(targetURL)
|
||||
if err != nil || parsed.Hostname() == "" {
|
||||
return errUnsafeOAuthDiscoveryTarget
|
||||
}
|
||||
return policy.validateLiteralHost(parsed.Hostname())
|
||||
}
|
||||
|
||||
func (transport advertisedOAuthDiscoveryTransport) RoundTrip(request *http.Request) (*http.Response, error) {
|
||||
if request == nil || request.URL == nil {
|
||||
return nil, errors.New("mcp oauth: discovery request has no URL")
|
||||
}
|
||||
if err := transport.policy.validateLiteralHost(request.URL.Hostname()); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
baseTransport, ok := transport.base.(*http.Transport)
|
||||
if !ok {
|
||||
return nil, errors.New("mcp oauth: cannot enforce discovery network policy with a custom HTTP transport")
|
||||
}
|
||||
addresses, err := transport.policy.resolve(request.Context(), request.URL.Hostname())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var lastErr error
|
||||
for index, address := range addresses {
|
||||
attempt := request
|
||||
if index > 0 && request.Body != nil {
|
||||
if request.GetBody == nil {
|
||||
break
|
||||
}
|
||||
body, bodyErr := request.GetBody()
|
||||
if bodyErr != nil {
|
||||
return nil, bodyErr
|
||||
}
|
||||
attempt = request.Clone(request.Context())
|
||||
attempt.Body = body
|
||||
}
|
||||
response, roundTripErr := roundTripPinnedOAuthDiscovery(attempt, baseTransport, address)
|
||||
if roundTripErr == nil {
|
||||
return response, nil
|
||||
}
|
||||
lastErr = roundTripErr
|
||||
}
|
||||
if lastErr == nil {
|
||||
lastErr = errUnsafeOAuthDiscoveryTarget
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
func roundTripPinnedOAuthDiscovery(request *http.Request, base *http.Transport, address netip.Addr) (*http.Response, error) {
|
||||
pinnedRequest := request.Clone(request.Context())
|
||||
pinnedURL := *request.URL
|
||||
originalHost := request.URL.Hostname()
|
||||
port := request.URL.Port()
|
||||
if port == "" {
|
||||
switch strings.ToLower(request.URL.Scheme) {
|
||||
case "http":
|
||||
port = "80"
|
||||
case "https":
|
||||
port = "443"
|
||||
default:
|
||||
return nil, errUnsafeOAuthDiscoveryTarget
|
||||
}
|
||||
}
|
||||
pinnedURL.Host = net.JoinHostPort(address.String(), port)
|
||||
pinnedRequest.URL = &pinnedURL
|
||||
if pinnedRequest.Host == "" {
|
||||
pinnedRequest.Host = request.URL.Host
|
||||
}
|
||||
|
||||
pinnedTransport := base.Clone()
|
||||
pinnedTransport.DisableKeepAlives = true
|
||||
var proxyURL *url.URL
|
||||
if base.Proxy != nil {
|
||||
var err error
|
||||
proxyURL, err = base.Proxy(request)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("mcp oauth: resolve discovery proxy: %w", err)
|
||||
}
|
||||
if proxyURL == nil {
|
||||
pinnedTransport.Proxy = nil
|
||||
} else {
|
||||
pinnedTransport.Proxy = http.ProxyURL(proxyURL)
|
||||
}
|
||||
}
|
||||
if pinnedTransport.TLSClientConfig == nil {
|
||||
pinnedTransport.TLSClientConfig = &tls.Config{ServerName: originalHost} //nolint:gosec // standard verification remains enabled
|
||||
} else {
|
||||
pinnedTransport.TLSClientConfig = pinnedTransport.TLSClientConfig.Clone()
|
||||
pinnedTransport.TLSClientConfig.ServerName = originalHost
|
||||
}
|
||||
// A custom TLS dialer must not resolve the advertised hostname again after
|
||||
// the policy has pinned it. For an HTTPS proxy, however, that proxy remains
|
||||
// the TLS first hop and the transport sends CONNECT to the pinned target.
|
||||
dialContext := pinnedTransport.DialContext
|
||||
if dialContext == nil {
|
||||
dialContext = (&net.Dialer{}).DialContext
|
||||
}
|
||||
pinnedAddress := net.JoinHostPort(address.String(), port)
|
||||
targetTLSConfig := pinnedTransport.TLSClientConfig
|
||||
pinnedTransport.DialTLSContext = func(ctx context.Context, network string, address string) (net.Conn, error) {
|
||||
if proxyURL != nil && strings.EqualFold(proxyURL.Scheme, "https") {
|
||||
proxyTLSConfig := targetTLSConfig.Clone()
|
||||
proxyTLSConfig.ServerName = proxyURL.Hostname()
|
||||
return dialOAuthDiscoveryTLS(ctx, dialContext, network, address, proxyTLSConfig)
|
||||
}
|
||||
return dialOAuthDiscoveryTLS(ctx, dialContext, network, pinnedAddress, targetTLSConfig)
|
||||
}
|
||||
return pinnedTransport.RoundTrip(pinnedRequest)
|
||||
}
|
||||
|
||||
func dialOAuthDiscoveryTLS(ctx context.Context, dialContext func(context.Context, string, string) (net.Conn, error), network string, address string, config *tls.Config) (net.Conn, error) {
|
||||
connection, err := dialContext(ctx, network, address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tlsConnection := tls.Client(connection, config)
|
||||
if err := tlsConnection.HandshakeContext(ctx); err != nil {
|
||||
connection.Close()
|
||||
return nil, err
|
||||
}
|
||||
return tlsConnection, nil
|
||||
}
|
||||
|
||||
func (policy oauthDiscoveryNetworkPolicy) resolve(ctx context.Context, host string) ([]netip.Addr, error) {
|
||||
if address, err := netip.ParseAddr(strings.Trim(host, "[]")); err == nil {
|
||||
address = address.Unmap()
|
||||
if err := policy.validateAddress(address); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []netip.Addr{address}, nil
|
||||
}
|
||||
|
||||
lookupNetIP := policy.lookupNetIP
|
||||
if lookupNetIP == nil {
|
||||
lookupNetIP = net.DefaultResolver.LookupNetIP
|
||||
}
|
||||
addresses, err := lookupNetIP(ctx, "ip", host)
|
||||
if err != nil {
|
||||
return nil, errors.New("mcp oauth: resolve server-advertised discovery target failed")
|
||||
}
|
||||
if len(addresses) == 0 {
|
||||
return nil, errors.New("mcp oauth: server-advertised discovery target resolved to no addresses")
|
||||
}
|
||||
explicitLoopback := isExplicitOAuthLoopbackHost(host)
|
||||
resolved := make([]netip.Addr, 0, len(addresses))
|
||||
for _, address := range addresses {
|
||||
address = address.Unmap()
|
||||
if explicitLoopback && !address.IsLoopback() {
|
||||
return nil, errUnsafeOAuthDiscoveryTarget
|
||||
}
|
||||
if err := policy.validateAddress(address); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resolved = append(resolved, address)
|
||||
}
|
||||
return resolved, nil
|
||||
}
|
||||
|
||||
func (policy oauthDiscoveryNetworkPolicy) validateLiteralHost(host string) error {
|
||||
normalized := strings.ToLower(strings.TrimSuffix(strings.Trim(host, "[]"), "."))
|
||||
if normalized == "" {
|
||||
return errUnsafeOAuthDiscoveryTarget
|
||||
}
|
||||
if isCloudMetadataHostname(normalized) {
|
||||
return errUnsafeOAuthDiscoveryTarget
|
||||
}
|
||||
if isExplicitOAuthLoopbackHost(normalized) && !policy.allowLoopback {
|
||||
return errUnsafeOAuthDiscoveryTarget
|
||||
}
|
||||
if address, err := netip.ParseAddr(normalized); err == nil {
|
||||
return policy.validateAddress(address.Unmap())
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (policy oauthDiscoveryNetworkPolicy) validateAddress(address netip.Addr) error {
|
||||
if policy.allowLoopback && address.IsLoopback() {
|
||||
return nil
|
||||
}
|
||||
if !address.IsValid() || !address.IsGlobalUnicast() || address.IsPrivate() {
|
||||
return errUnsafeOAuthDiscoveryTarget
|
||||
}
|
||||
for _, prefix := range nonPublicOAuthDiscoveryPrefixes {
|
||||
if prefix.Contains(address) {
|
||||
return errUnsafeOAuthDiscoveryTarget
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func isExplicitOAuthLoopbackHost(host string) bool {
|
||||
normalized := strings.ToLower(strings.TrimSuffix(strings.Trim(host, "[]"), "."))
|
||||
if normalized == "localhost" || strings.HasSuffix(normalized, ".localhost") {
|
||||
return true
|
||||
}
|
||||
address, err := netip.ParseAddr(normalized)
|
||||
return err == nil && address.Unmap().IsLoopback()
|
||||
}
|
||||
|
||||
func isCloudMetadataHostname(host string) bool {
|
||||
switch strings.ToLower(strings.TrimSuffix(host, ".")) {
|
||||
case "metadata", "metadata.google.internal", "metadata.azure.internal":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,180 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestOAuthDiscoveryNetworkPolicyRejectsNonPublicAddresses(t *testing.T) {
|
||||
policy := oauthDiscoveryNetworkPolicy{}
|
||||
for _, value := range []string{
|
||||
"127.0.0.1",
|
||||
"10.0.0.1",
|
||||
"169.254.169.254",
|
||||
"192.0.2.10",
|
||||
"224.0.0.1",
|
||||
"::1",
|
||||
"fc00::1",
|
||||
"fe80::1",
|
||||
"2001:db8::1",
|
||||
} {
|
||||
t.Run(value, func(t *testing.T) {
|
||||
if err := policy.validateAddress(netip.MustParseAddr(value)); err == nil {
|
||||
t.Fatalf("address %s should be rejected", value)
|
||||
}
|
||||
})
|
||||
}
|
||||
for _, value := range []string{"8.8.8.8", "2606:4700:4700::1111"} {
|
||||
if err := policy.validateAddress(netip.MustParseAddr(value)); err != nil {
|
||||
t.Fatalf("public address %s rejected: %v", value, err)
|
||||
}
|
||||
}
|
||||
loopbackPolicy := oauthDiscoveryNetworkPolicy{allowLoopback: true}
|
||||
if err := loopbackPolicy.validateAddress(netip.MustParseAddr("127.0.0.1")); err != nil {
|
||||
t.Fatalf("explicit loopback resource should allow loopback discovery: %v", err)
|
||||
}
|
||||
if err := loopbackPolicy.validateAddress(netip.MustParseAddr("10.0.0.1")); err == nil {
|
||||
t.Fatal("loopback development must not permit arbitrary private targets")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProtectedDiscoveryRejectsNonPublicOAuthEndpoints(t *testing.T) {
|
||||
for _, endpoint := range []struct {
|
||||
name string
|
||||
metadata authServerMetadata
|
||||
}{
|
||||
{
|
||||
name: "authorization",
|
||||
metadata: authServerMetadata{
|
||||
AuthorizationEndpoint: "http://127.0.0.1/authorize",
|
||||
protectAuthorizationEndpoint: true,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "token",
|
||||
metadata: authServerMetadata{
|
||||
TokenEndpoint: "https://10.0.0.1/token",
|
||||
protectTokenEndpoint: true,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "registration",
|
||||
metadata: authServerMetadata{
|
||||
RegistrationEndpoint: "https://169.254.169.254/register",
|
||||
protectRegistrationEndpoint: true,
|
||||
},
|
||||
},
|
||||
} {
|
||||
t.Run(endpoint.name, func(t *testing.T) {
|
||||
endpoint.metadata.protectedResourceURL = "https://mcp.example/mcp"
|
||||
if err := validateProtectedDiscoveredEndpoints(context.Background(), endpoint.metadata); err == nil {
|
||||
t.Fatalf("protected %s endpoint should be rejected", endpoint.name)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdvertisedOAuthDiscoveryRejectsHostnameResolvingPrivateAtUseTime(t *testing.T) {
|
||||
var dialed atomic.Bool
|
||||
policy := oauthDiscoveryNetworkPolicy{
|
||||
lookupNetIP: func(context.Context, string, string) ([]netip.Addr, error) {
|
||||
return []netip.Addr{netip.MustParseAddr("10.0.0.12")}, nil
|
||||
},
|
||||
}
|
||||
client := &http.Client{Transport: advertisedOAuthDiscoveryTransport{
|
||||
base: &http.Transport{DialContext: func(context.Context, string, string) (net.Conn, error) {
|
||||
dialed.Store(true)
|
||||
return nil, io.EOF
|
||||
}},
|
||||
policy: policy,
|
||||
}}
|
||||
|
||||
_, err := client.Get("https://auth.example/.well-known/oauth-authorization-server")
|
||||
if err == nil || !strings.Contains(err.Error(), errUnsafeOAuthDiscoveryTarget.Error()) {
|
||||
t.Fatalf("error = %v, want private-resolution rejection", err)
|
||||
}
|
||||
if dialed.Load() {
|
||||
t.Fatal("private resolved address reached the dialer")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdvertisedOAuthDiscoveryPinsResolvedAddress(t *testing.T) {
|
||||
const publicAddress = "93.184.216.34"
|
||||
var dialedAddress string
|
||||
policy := oauthDiscoveryNetworkPolicy{
|
||||
lookupNetIP: func(context.Context, string, string) ([]netip.Addr, error) {
|
||||
return []netip.Addr{netip.MustParseAddr(publicAddress)}, nil
|
||||
},
|
||||
}
|
||||
baseTransport := &http.Transport{DialContext: func(_ context.Context, _, address string) (net.Conn, error) {
|
||||
dialedAddress = address
|
||||
clientConn, serverConn := net.Pipe()
|
||||
go func() {
|
||||
defer serverConn.Close()
|
||||
reader := bufio.NewReader(serverConn)
|
||||
for {
|
||||
line, err := reader.ReadString('\n')
|
||||
if err != nil || line == "\r\n" {
|
||||
break
|
||||
}
|
||||
}
|
||||
_, _ = io.WriteString(serverConn, "HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}")
|
||||
}()
|
||||
return clientConn, nil
|
||||
}}
|
||||
client := &http.Client{Transport: advertisedOAuthDiscoveryTransport{base: baseTransport, policy: policy}}
|
||||
|
||||
response, err := client.Get("http://auth.example/metadata")
|
||||
if err != nil {
|
||||
t.Fatalf("GET pinned discovery target: %v", err)
|
||||
}
|
||||
response.Body.Close()
|
||||
if dialedAddress != net.JoinHostPort(publicAddress, "80") {
|
||||
t.Fatalf("dialed %q, want resolved address", dialedAddress)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPinnedOAuthDiscoveryPreservesHTTPSProxyFirstHop(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
proxy string
|
||||
firstHop string
|
||||
}{
|
||||
{name: "HTTP", proxy: "http://proxy.example:8080", firstHop: "proxy.example:8080"},
|
||||
{name: "HTTPS", proxy: "https://proxy.example:8443", firstHop: "proxy.example:8443"},
|
||||
{name: "SOCKS", proxy: "socks5://proxy.example:1080", firstHop: "proxy.example:1080"},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
proxyURL, err := url.Parse(test.proxy)
|
||||
if err != nil {
|
||||
t.Fatalf("parse proxy URL: %v", err)
|
||||
}
|
||||
var dialedAddress string
|
||||
baseTransport := &http.Transport{
|
||||
Proxy: http.ProxyURL(proxyURL),
|
||||
DialContext: func(_ context.Context, _, address string) (net.Conn, error) {
|
||||
dialedAddress = address
|
||||
return nil, errors.New("stop after observing first hop")
|
||||
},
|
||||
}
|
||||
request, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://auth.example/metadata", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("create request: %v", err)
|
||||
}
|
||||
|
||||
_, _ = roundTripPinnedOAuthDiscovery(request, baseTransport, netip.MustParseAddr("93.184.216.34"))
|
||||
if dialedAddress != test.firstHop {
|
||||
t.Fatalf("dialed %q, want %s proxy first hop", dialedAddress, test.name)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,233 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/oauth"
|
||||
)
|
||||
|
||||
const protectedResourceMetadataLimit = 1 << 20
|
||||
|
||||
type protectedResourceMetadata struct {
|
||||
Resource string `json:"resource"`
|
||||
AuthorizationServers []string `json:"authorization_servers"`
|
||||
}
|
||||
|
||||
// withoutDiscoveryRedirects returns a shallow client copy that exposes 3xx
|
||||
// responses to the caller. Every discovered URL is validated before use, so a
|
||||
// redirect must never replace that validated target with an unchecked one.
|
||||
func withoutDiscoveryRedirects(client *http.Client) *http.Client {
|
||||
if client == nil {
|
||||
client = http.DefaultClient
|
||||
}
|
||||
copied := *client
|
||||
// Discovery requests are deliberately unauthenticated. A caller-supplied
|
||||
// jar must not turn a plain metadata GET into a credential-bearing request.
|
||||
copied.Jar = nil
|
||||
copied.CheckRedirect = func(*http.Request, []*http.Request) error {
|
||||
return http.ErrUseLastResponse
|
||||
}
|
||||
return &copied
|
||||
}
|
||||
|
||||
// discoverProtectedResourceAuthorizationServer follows the RFC 9728 URL
|
||||
// advertised by a protected MCP resource. found is false when the endpoint does
|
||||
// not advertise protected-resource metadata, allowing legacy direct discovery.
|
||||
func discoverProtectedResourceAuthorizationServer(ctx context.Context, client *http.Client, resourceURL string) (issuer string, found bool, err error) {
|
||||
resourceURL = strings.TrimSpace(resourceURL)
|
||||
client = withoutDiscoveryRedirects(client)
|
||||
metadataURL, found, err := probeProtectedResourceMetadataURL(ctx, client, resourceURL)
|
||||
if err != nil || !found {
|
||||
return "", found, err
|
||||
}
|
||||
if err := validateResourceMetadataOrigin(resourceURL, metadataURL); err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
client, err = newAdvertisedOAuthDiscoveryClient(client, resourceURL)
|
||||
if err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
metadata, err := fetchProtectedResourceMetadata(ctx, client, metadataURL)
|
||||
if err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
if metadata.Resource != resourceURL {
|
||||
return "", true, errors.New("mcp oauth: protected resource metadata resource does not match the MCP server URL")
|
||||
}
|
||||
if len(metadata.AuthorizationServers) == 0 {
|
||||
return "", true, errors.New("mcp oauth: protected resource metadata has no authorization_servers")
|
||||
}
|
||||
issuer = metadata.AuthorizationServers[0]
|
||||
if strings.TrimSpace(issuer) == "" {
|
||||
return "", true, errors.New("mcp oauth: protected resource metadata has an empty authorization server")
|
||||
}
|
||||
if issuer != strings.TrimSpace(issuer) {
|
||||
return "", true, errors.New("mcp oauth: protected resource authorization server must not contain surrounding whitespace")
|
||||
}
|
||||
if err := oauth.ValidateEndpointURL(issuer); err != nil {
|
||||
return "", true, fmt.Errorf("mcp oauth: protected resource authorization server: %w", err)
|
||||
}
|
||||
if err := validateAdvertisedOAuthDiscoveryURL(resourceURL, issuer); err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
return issuer, true, nil
|
||||
}
|
||||
|
||||
func validateResourceMetadataOrigin(resourceURL string, metadataURL string) error {
|
||||
resource, resourceErr := url.Parse(resourceURL)
|
||||
metadata, metadataErr := url.Parse(metadataURL)
|
||||
if resourceErr != nil || metadataErr != nil || !sameMCPOrigin(resource, metadata) {
|
||||
return errors.New("mcp oauth: protected resource metadata URL must be same-origin with the MCP server URL")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func probeProtectedResourceMetadataURL(ctx context.Context, client *http.Client, resourceURL string) (string, bool, error) {
|
||||
if client == nil {
|
||||
client = http.DefaultClient
|
||||
}
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.TrimSpace(resourceURL), nil)
|
||||
if err != nil {
|
||||
return "", false, errors.New("mcp oauth: invalid MCP resource URL")
|
||||
}
|
||||
request.Header.Set("Accept", "application/json, text/event-stream")
|
||||
request.Header.Set("MCP-Protocol-Version", defaultProtocolVersion)
|
||||
response, err := client.Do(request)
|
||||
if err != nil {
|
||||
// A resource that cannot be probed may still support the legacy direct
|
||||
// authorization-server discovery flow.
|
||||
return "", false, nil
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode != http.StatusUnauthorized {
|
||||
return "", false, nil
|
||||
}
|
||||
metadataURL, found, err := parseResourceMetadataURL(response.Header.Values("WWW-Authenticate"))
|
||||
if err != nil {
|
||||
return "", false, fmt.Errorf("mcp oauth: invalid WWW-Authenticate resource_metadata parameter: %w", err)
|
||||
}
|
||||
if !found {
|
||||
return "", false, nil
|
||||
}
|
||||
if err := oauth.ValidateEndpointURL(metadataURL); err != nil {
|
||||
return "", false, fmt.Errorf("mcp oauth: protected resource metadata URL: %w", err)
|
||||
}
|
||||
return metadataURL, true, nil
|
||||
}
|
||||
|
||||
func fetchProtectedResourceMetadata(ctx context.Context, client *http.Client, metadataURL string) (protectedResourceMetadata, error) {
|
||||
if client == nil {
|
||||
client = http.DefaultClient
|
||||
}
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodGet, metadataURL, nil)
|
||||
if err != nil {
|
||||
return protectedResourceMetadata{}, errors.New("mcp oauth: invalid protected resource metadata URL")
|
||||
}
|
||||
request.Header.Set("Accept", "application/json")
|
||||
response, err := client.Do(request)
|
||||
if err != nil {
|
||||
return protectedResourceMetadata{}, errors.New("mcp oauth: fetch protected resource metadata failed")
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
|
||||
return protectedResourceMetadata{}, fmt.Errorf("mcp oauth: protected resource metadata returned HTTP %d", response.StatusCode)
|
||||
}
|
||||
var metadata protectedResourceMetadata
|
||||
if err := json.NewDecoder(io.LimitReader(response.Body, protectedResourceMetadataLimit)).Decode(&metadata); err != nil {
|
||||
return protectedResourceMetadata{}, errors.New("mcp oauth: decode protected resource metadata failed")
|
||||
}
|
||||
return metadata, nil
|
||||
}
|
||||
|
||||
func parseResourceMetadataURL(headerValues []string) (string, bool, error) {
|
||||
var selected string
|
||||
for _, headerValue := range headerValues {
|
||||
for index := 0; index < len(headerValue); {
|
||||
if headerValue[index] == '"' {
|
||||
_, next, err := readAuthQuotedString(headerValue, index)
|
||||
if err != nil {
|
||||
// Malformed parameters for another authentication scheme do
|
||||
// not suppress a valid challenge in a separate header value.
|
||||
break
|
||||
}
|
||||
index = next
|
||||
continue
|
||||
}
|
||||
if !isAuthTokenByte(headerValue[index]) {
|
||||
index++
|
||||
continue
|
||||
}
|
||||
start := index
|
||||
for index < len(headerValue) && isAuthTokenByte(headerValue[index]) {
|
||||
index++
|
||||
}
|
||||
name := headerValue[start:index]
|
||||
for index < len(headerValue) && (headerValue[index] == ' ' || headerValue[index] == '\t') {
|
||||
index++
|
||||
}
|
||||
if index >= len(headerValue) || headerValue[index] != '=' {
|
||||
continue
|
||||
}
|
||||
index++
|
||||
for index < len(headerValue) && (headerValue[index] == ' ' || headerValue[index] == '\t') {
|
||||
index++
|
||||
}
|
||||
if !strings.EqualFold(name, "resource_metadata") {
|
||||
continue
|
||||
}
|
||||
if index >= len(headerValue) || headerValue[index] != '"' {
|
||||
return "", false, errors.New("resource_metadata must be a quoted URL")
|
||||
}
|
||||
value, next, err := readAuthQuotedString(headerValue, index)
|
||||
if err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return "", false, errors.New("resource_metadata URL is empty")
|
||||
}
|
||||
if selected != "" && selected != value {
|
||||
return "", false, errors.New("conflicting resource_metadata URLs")
|
||||
}
|
||||
selected = value
|
||||
index = next
|
||||
}
|
||||
}
|
||||
return selected, selected != "", nil
|
||||
}
|
||||
|
||||
func readAuthQuotedString(value string, quote int) (string, int, error) {
|
||||
var decoded strings.Builder
|
||||
for index := quote + 1; index < len(value); index++ {
|
||||
switch value[index] {
|
||||
case '"':
|
||||
return decoded.String(), index + 1, nil
|
||||
case '\\':
|
||||
index++
|
||||
if index >= len(value) {
|
||||
return "", 0, errors.New("unterminated quoted resource_metadata URL")
|
||||
}
|
||||
decoded.WriteByte(value[index])
|
||||
default:
|
||||
if value[index] < 0x20 || value[index] == 0x7f {
|
||||
return "", 0, errors.New("resource_metadata URL contains a control character")
|
||||
}
|
||||
decoded.WriteByte(value[index])
|
||||
}
|
||||
}
|
||||
return "", 0, errors.New("unterminated quoted resource_metadata URL")
|
||||
}
|
||||
|
||||
func isAuthTokenByte(value byte) bool {
|
||||
if value >= 'a' && value <= 'z' || value >= 'A' && value <= 'Z' || value >= '0' && value <= '9' {
|
||||
return true
|
||||
}
|
||||
return strings.ContainsRune("!#$%&'*+-.^_`|~", rune(value))
|
||||
}
|
||||
@@ -0,0 +1,742 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/cookiejar"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestParseResourceMetadataURL(t *testing.T) {
|
||||
got, found, err := parseResourceMetadataURL([]string{
|
||||
`Basic realm="resource_metadata=not-a-parameter"`,
|
||||
`Bearer realm="OAuth", scope="read", Resource_Metadata="https://mcp.example/resource-metadata"`,
|
||||
})
|
||||
if err != nil || !found {
|
||||
t.Fatalf("parseResourceMetadataURL() found=%v err=%v", found, err)
|
||||
}
|
||||
if got != "https://mcp.example/resource-metadata" {
|
||||
t.Fatalf("resource metadata URL = %q", got)
|
||||
}
|
||||
|
||||
got, found, err = parseResourceMetadataURL([]string{
|
||||
`Basic realm="unterminated`,
|
||||
`Bearer resource_metadata="https://mcp.example/fallback-metadata"`,
|
||||
})
|
||||
if err != nil || !found || got != "https://mcp.example/fallback-metadata" {
|
||||
t.Fatalf("unrelated malformed challenge blocked discovery: got=%q found=%v err=%v", got, found, err)
|
||||
}
|
||||
|
||||
for name, headers := range map[string][]string{
|
||||
"unquoted": {`Bearer resource_metadata=https://mcp.example/metadata`},
|
||||
"unterminated": {`Bearer resource_metadata="https://mcp.example/metadata`},
|
||||
"conflicting": {
|
||||
`Bearer resource_metadata="https://mcp.example/one"`,
|
||||
`DPoP resource_metadata="https://mcp.example/two"`,
|
||||
},
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
if _, _, err := parseResourceMetadataURL(headers); err == nil {
|
||||
t.Fatal("expected malformed resource_metadata rejection")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestProtectedResourceDiscoveryRejectsInvalidMetadata(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
challenge func(base string) string
|
||||
metadata func(base string) string
|
||||
wantError string
|
||||
}{
|
||||
{
|
||||
name: "malformed challenge",
|
||||
challenge: func(string) string { return `Bearer resource_metadata="unterminated` },
|
||||
wantError: "WWW-Authenticate",
|
||||
},
|
||||
{
|
||||
name: "insecure metadata URL",
|
||||
challenge: func(string) string { return `Bearer resource_metadata="http://metadata.example/resource"` },
|
||||
wantError: "metadata URL",
|
||||
},
|
||||
{
|
||||
name: "invalid JSON",
|
||||
challenge: func(base string) string { return `Bearer resource_metadata="` + base + `/metadata"` },
|
||||
metadata: func(string) string { return `{` },
|
||||
wantError: "decode protected resource metadata",
|
||||
},
|
||||
{
|
||||
name: "resource mismatch",
|
||||
challenge: func(base string) string { return `Bearer resource_metadata="` + base + `/metadata"` },
|
||||
metadata: func(base string) string {
|
||||
return `{"resource":"` + base + `/other","authorization_servers":["https://auth.example"]}`
|
||||
},
|
||||
wantError: "does not match",
|
||||
},
|
||||
{
|
||||
name: "resource leading whitespace mismatch",
|
||||
challenge: func(base string) string { return `Bearer resource_metadata="` + base + `/metadata"` },
|
||||
metadata: func(base string) string {
|
||||
return `{"resource":" ` + base + `/mcp","authorization_servers":["https://auth.example"]}`
|
||||
},
|
||||
wantError: "does not match",
|
||||
},
|
||||
{
|
||||
name: "resource trailing whitespace mismatch",
|
||||
challenge: func(base string) string { return `Bearer resource_metadata="` + base + `/metadata"` },
|
||||
metadata: func(base string) string {
|
||||
return `{"resource":"` + base + `/mcp ","authorization_servers":["https://auth.example"]}`
|
||||
},
|
||||
wantError: "does not match",
|
||||
},
|
||||
{
|
||||
name: "authorization servers missing",
|
||||
challenge: func(base string) string { return `Bearer resource_metadata="` + base + `/metadata"` },
|
||||
metadata: func(base string) string {
|
||||
return `{"resource":"` + base + `/mcp"}`
|
||||
},
|
||||
wantError: "authorization_servers",
|
||||
},
|
||||
{
|
||||
name: "insecure authorization server",
|
||||
challenge: func(base string) string { return `Bearer resource_metadata="` + base + `/metadata"` },
|
||||
metadata: func(base string) string {
|
||||
return `{"resource":"` + base + `/mcp","authorization_servers":["http://auth.example"]}`
|
||||
},
|
||||
wantError: "authorization server",
|
||||
},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
base := "http://" + r.Host
|
||||
switch r.URL.Path {
|
||||
case "/mcp":
|
||||
w.Header().Set("WWW-Authenticate", test.challenge(base))
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
case "/metadata":
|
||||
if test.metadata == nil {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(test.metadata(base)))
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
_, _, err := discoverProtectedResourceAuthorizationServer(context.Background(), server.Client(), server.URL+"/mcp")
|
||||
if err == nil || !strings.Contains(err.Error(), test.wantError) {
|
||||
t.Fatalf("error = %v, want %q", err, test.wantError)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestProtectedResourceProbeSendsNoCredentials(t *testing.T) {
|
||||
var authorization, cookie string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
authorization = r.Header.Get("Authorization")
|
||||
cookie = r.Header.Get("Cookie")
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}))
|
||||
defer server.Close()
|
||||
jar, err := cookiejar.New(nil)
|
||||
if err != nil {
|
||||
t.Fatalf("create cookie jar: %v", err)
|
||||
}
|
||||
serverURL, err := url.Parse(server.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("parse server URL: %v", err)
|
||||
}
|
||||
jar.SetCookies(serverURL, []*http.Cookie{{Name: "session", Value: "secret"}})
|
||||
client := server.Client()
|
||||
client.Jar = jar
|
||||
|
||||
_, found, err := discoverProtectedResourceAuthorizationServer(context.Background(), client, server.URL+"/mcp")
|
||||
if err != nil || found {
|
||||
t.Fatalf("discovery found=%v err=%v", found, err)
|
||||
}
|
||||
if authorization != "" || cookie != "" {
|
||||
t.Fatalf("probe sent credentials: authorization=%q cookie=%q", authorization, cookie)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDirectAuthorizationServerDiscoverySendsNoCookies(t *testing.T) {
|
||||
var cookie string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
cookie = r.Header.Get("Cookie")
|
||||
base := "http://" + r.Host
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"issuer": base,
|
||||
"authorization_endpoint": base + "/authorize",
|
||||
"token_endpoint": base + "/token",
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
jar, err := cookiejar.New(nil)
|
||||
if err != nil {
|
||||
t.Fatalf("create cookie jar: %v", err)
|
||||
}
|
||||
serverURL, err := url.Parse(server.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("parse server URL: %v", err)
|
||||
}
|
||||
jar.SetCookies(serverURL, []*http.Cookie{{Name: "session", Value: "secret"}})
|
||||
client := server.Client()
|
||||
client.Jar = jar
|
||||
|
||||
if _, err := resolveAuthorizationServer(context.Background(), client, server.URL, OAuthConfig{}); err != nil {
|
||||
t.Fatalf("resolveAuthorizationServer() error = %v", err)
|
||||
}
|
||||
if cookie != "" {
|
||||
t.Fatalf("direct discovery sent cookie %q", cookie)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProtectedResourceDiscoveryRejectsAdvertisedLoopbackMetadata(t *testing.T) {
|
||||
var metadataHits atomic.Int64
|
||||
client := &http.Client{Transport: roundTripperFunc(func(request *http.Request) (*http.Response, error) {
|
||||
response := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader("")),
|
||||
Request: request,
|
||||
}
|
||||
switch request.URL.String() {
|
||||
case "https://mcp.example/mcp":
|
||||
response.StatusCode = http.StatusUnauthorized
|
||||
response.Header.Set("WWW-Authenticate", `Bearer resource_metadata="http://127.0.0.1/metadata"`)
|
||||
case "http://127.0.0.1/metadata":
|
||||
metadataHits.Add(1)
|
||||
response.Body = io.NopCloser(strings.NewReader(`{"resource":"https://mcp.example/mcp","authorization_servers":["https://auth.example"]}`))
|
||||
default:
|
||||
response.StatusCode = http.StatusNotFound
|
||||
}
|
||||
return response, nil
|
||||
})}
|
||||
|
||||
_, _, err := discoverProtectedResourceAuthorizationServer(context.Background(), client, "https://mcp.example/mcp")
|
||||
if err == nil {
|
||||
t.Fatal("advertised loopback metadata URL should be rejected")
|
||||
}
|
||||
if metadataHits.Load() != 0 {
|
||||
t.Fatalf("loopback metadata target was fetched %d time(s)", metadataHits.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestProtectedResourceDiscoveryRejectsAdvertisedLoopbackIssuer(t *testing.T) {
|
||||
var issuerHits atomic.Int64
|
||||
client := &http.Client{Transport: roundTripperFunc(func(request *http.Request) (*http.Response, error) {
|
||||
response := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader("")),
|
||||
Request: request,
|
||||
}
|
||||
switch request.URL.String() {
|
||||
case "https://mcp.example/mcp":
|
||||
response.StatusCode = http.StatusUnauthorized
|
||||
response.Header.Set("WWW-Authenticate", `Bearer resource_metadata="https://mcp.example/metadata"`)
|
||||
case "https://mcp.example/metadata":
|
||||
response.Body = io.NopCloser(strings.NewReader(`{"resource":"https://mcp.example/mcp","authorization_servers":["http://127.0.0.1"]}`))
|
||||
case "http://127.0.0.1/.well-known/oauth-authorization-server":
|
||||
issuerHits.Add(1)
|
||||
response.Body = io.NopCloser(strings.NewReader(`{"issuer":"http://127.0.0.1","authorization_endpoint":"http://127.0.0.1/authorize","token_endpoint":"http://127.0.0.1/token"}`))
|
||||
default:
|
||||
response.StatusCode = http.StatusNotFound
|
||||
}
|
||||
return response, nil
|
||||
})}
|
||||
|
||||
_, err := resolveAuthorizationServer(context.Background(), client, "https://mcp.example/mcp", OAuthConfig{})
|
||||
if err == nil {
|
||||
t.Fatal("advertised loopback issuer should be rejected")
|
||||
}
|
||||
if issuerHits.Load() != 0 {
|
||||
t.Fatalf("loopback issuer target was fetched %d time(s)", issuerHits.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestProtectedResourceDiscoveryRejectsAdvertisedPrivateTargets(t *testing.T) {
|
||||
for _, target := range []string{
|
||||
"http://10.0.0.8/metadata",
|
||||
"https://169.254.169.254/metadata",
|
||||
"https://metadata.google.internal/metadata",
|
||||
} {
|
||||
t.Run(target, func(t *testing.T) {
|
||||
if err := validateAdvertisedOAuthDiscoveryURL("https://mcp.example/mcp", target); err == nil {
|
||||
t.Fatalf("advertised target %q should be rejected", target)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestProtectedDiscoveryPathsSendNoCookies(t *testing.T) {
|
||||
var probeCookie, resourceCookie, issuerCookie string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
base := "http://" + r.Host
|
||||
switch r.URL.Path {
|
||||
case "/mcp":
|
||||
probeCookie = r.Header.Get("Cookie")
|
||||
w.Header().Set("WWW-Authenticate", `Bearer resource_metadata="`+base+`/metadata"`)
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
case "/metadata":
|
||||
resourceCookie = r.Header.Get("Cookie")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"resource": base + "/mcp",
|
||||
"authorization_servers": []string{base},
|
||||
})
|
||||
case "/.well-known/oauth-authorization-server":
|
||||
issuerCookie = r.Header.Get("Cookie")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"issuer": base,
|
||||
"authorization_endpoint": base + "/authorize",
|
||||
"token_endpoint": base + "/token",
|
||||
})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
jar, err := cookiejar.New(nil)
|
||||
if err != nil {
|
||||
t.Fatalf("create cookie jar: %v", err)
|
||||
}
|
||||
serverURL, err := url.Parse(server.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("parse server URL: %v", err)
|
||||
}
|
||||
jar.SetCookies(serverURL, []*http.Cookie{{Name: "session", Value: "secret"}})
|
||||
client := server.Client()
|
||||
client.Jar = jar
|
||||
|
||||
if _, err := resolveAuthorizationServer(context.Background(), client, server.URL+"/mcp", OAuthConfig{}); err != nil {
|
||||
t.Fatalf("resolveAuthorizationServer() error = %v", err)
|
||||
}
|
||||
if probeCookie != "" || resourceCookie != "" || issuerCookie != "" {
|
||||
t.Fatalf("discovery sent cookies: probe=%q resource=%q issuer=%q", probeCookie, resourceCookie, issuerCookie)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProtectedDiscoveryPreservesExplicitEndpointOverride(t *testing.T) {
|
||||
fixture := newPublicProtectedOAuthFixture(t, "", "https://93.184.216.35/token", "", "")
|
||||
explicitToken := "http://127.0.0.1:19432/token"
|
||||
metadata, err := resolveAuthorizationServer(context.Background(), fixture.client, fixture.resourceURL, OAuthConfig{TokenEndpoint: explicitToken})
|
||||
if err != nil {
|
||||
t.Fatalf("resolve with explicit endpoint override: %v", err)
|
||||
}
|
||||
if metadata.TokenEndpoint != explicitToken || metadata.protectTokenEndpoint {
|
||||
t.Fatalf("explicit token override lost provenance: endpoint=%q protected=%v", metadata.TokenEndpoint, metadata.protectTokenEndpoint)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProtectedDiscoveryRejectsLoopbackRegistrationEndpointBeforeUse(t *testing.T) {
|
||||
var registrationHits atomic.Int64
|
||||
registration := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
registrationHits.Add(1)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"client_id": "registered-client"})
|
||||
}))
|
||||
defer registration.Close()
|
||||
fixture := newPublicProtectedOAuthFixture(t, registration.URL+"/register", "https://93.184.216.35/token", "", "")
|
||||
|
||||
_, err := Login(context.Background(), LoginOptions{
|
||||
ServerName: "registration-ssrf",
|
||||
ServerURL: fixture.resourceURL,
|
||||
HTTPClient: fixture.client,
|
||||
OpenBrowser: func(string) error {
|
||||
return errors.New("authorization browser should not open")
|
||||
},
|
||||
Timeout: 5 * time.Second,
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("protected discovery should reject a loopback registration endpoint")
|
||||
}
|
||||
if registrationHits.Load() != 0 {
|
||||
t.Fatalf("loopback registration endpoint received %d request(s)", registrationHits.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestProtectedDiscoveryRejectsLoopbackTokenEndpointBeforeUse(t *testing.T) {
|
||||
var tokenHits atomic.Int64
|
||||
tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
tokenHits.Add(1)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"access_token": "unexpected",
|
||||
"token_type": "Bearer",
|
||||
})
|
||||
}))
|
||||
defer tokenServer.Close()
|
||||
fixture := newPublicProtectedOAuthFixture(t, "", tokenServer.URL+"/token", "", "")
|
||||
|
||||
_, err := Login(context.Background(), LoginOptions{
|
||||
ServerName: "token-ssrf",
|
||||
ServerURL: fixture.resourceURL,
|
||||
Config: OAuthConfig{ClientID: "configured-client"},
|
||||
HTTPClient: fixture.client,
|
||||
OpenBrowser: func(authURL string) error {
|
||||
parsed, parseErr := url.Parse(authURL)
|
||||
if parseErr != nil {
|
||||
return parseErr
|
||||
}
|
||||
callbackURL := parsed.Query().Get("redirect_uri") + "?code=code&state=" + url.QueryEscape(parsed.Query().Get("state"))
|
||||
go func() {
|
||||
for attempt := 0; attempt < 20; attempt++ {
|
||||
response, requestErr := http.Get(callbackURL)
|
||||
if requestErr == nil {
|
||||
response.Body.Close()
|
||||
return
|
||||
}
|
||||
time.Sleep(25 * time.Millisecond)
|
||||
}
|
||||
}()
|
||||
return nil
|
||||
},
|
||||
Timeout: 5 * time.Second,
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("protected discovery should reject a loopback token endpoint")
|
||||
}
|
||||
if tokenHits.Load() != 0 {
|
||||
t.Fatalf("loopback token endpoint received %d request(s)", tokenHits.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestProtectedDiscoveryRegistrationRefusesPublicToPrivateRedirect(t *testing.T) {
|
||||
var victimHits atomic.Int64
|
||||
victim := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
victimHits.Add(1)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"client_id": "redirected-client"})
|
||||
}))
|
||||
defer victim.Close()
|
||||
fixture := newPublicProtectedOAuthFixture(t, "", "https://93.184.216.35/token", victim.URL+"/register", "")
|
||||
|
||||
_, err := Login(context.Background(), LoginOptions{
|
||||
ServerName: "registration-redirect",
|
||||
ServerURL: fixture.resourceURL,
|
||||
HTTPClient: fixture.client,
|
||||
OpenBrowser: func(string) error {
|
||||
return errors.New("authorization browser should not open")
|
||||
},
|
||||
Timeout: 5 * time.Second,
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("protected registration redirect should be refused")
|
||||
}
|
||||
if victimHits.Load() != 0 {
|
||||
t.Fatalf("registration redirect reached private victim %d time(s)", victimHits.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestProtectedDiscoveryTokenRefusesPublicToPrivateRedirect(t *testing.T) {
|
||||
var victimHits atomic.Int64
|
||||
victim := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
victimHits.Add(1)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"access_token": "redirected-token", "token_type": "Bearer"})
|
||||
}))
|
||||
defer victim.Close()
|
||||
fixture := newPublicProtectedOAuthFixture(t, "", "", "", victim.URL+"/token")
|
||||
|
||||
_, err := Login(context.Background(), LoginOptions{
|
||||
ServerName: "token-redirect",
|
||||
ServerURL: fixture.resourceURL,
|
||||
Config: OAuthConfig{ClientID: "configured-client"},
|
||||
HTTPClient: fixture.client,
|
||||
OpenBrowser: func(authURL string) error {
|
||||
parsed, parseErr := url.Parse(authURL)
|
||||
if parseErr != nil {
|
||||
return parseErr
|
||||
}
|
||||
callbackURL := parsed.Query().Get("redirect_uri") + "?code=code&state=" + url.QueryEscape(parsed.Query().Get("state"))
|
||||
go func() {
|
||||
for attempt := 0; attempt < 20; attempt++ {
|
||||
response, requestErr := http.Get(callbackURL)
|
||||
if requestErr == nil {
|
||||
response.Body.Close()
|
||||
return
|
||||
}
|
||||
time.Sleep(25 * time.Millisecond)
|
||||
}
|
||||
}()
|
||||
return nil
|
||||
},
|
||||
Timeout: 5 * time.Second,
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("protected token redirect should be refused")
|
||||
}
|
||||
if victimHits.Load() != 0 {
|
||||
t.Fatalf("token redirect reached private victim %d time(s)", victimHits.Load())
|
||||
}
|
||||
}
|
||||
|
||||
type publicProtectedOAuthFixture struct {
|
||||
client *http.Client
|
||||
resourceURL string
|
||||
}
|
||||
|
||||
func newPublicProtectedOAuthFixture(t *testing.T, registrationEndpoint string, tokenEndpoint string, registrationRedirect string, tokenRedirect string) publicProtectedOAuthFixture {
|
||||
t.Helper()
|
||||
var resourceOrigin, issuerOrigin string
|
||||
gateway := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
resourceHost := strings.TrimPrefix(resourceOrigin, "https://")
|
||||
issuerHost := strings.TrimPrefix(issuerOrigin, "https://")
|
||||
switch {
|
||||
case r.Host == resourceHost && r.URL.Path == "/mcp":
|
||||
w.Header().Set("WWW-Authenticate", `Bearer resource_metadata="`+resourceOrigin+`/metadata"`)
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
case r.Host == resourceHost && r.URL.Path == "/metadata":
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"resource": resourceOrigin + "/mcp",
|
||||
"authorization_servers": []string{issuerOrigin},
|
||||
})
|
||||
case r.Host == issuerHost && r.URL.Path == "/.well-known/oauth-authorization-server":
|
||||
metadataRegistration := registrationEndpoint
|
||||
metadataToken := tokenEndpoint
|
||||
if registrationRedirect != "" {
|
||||
metadataRegistration = issuerOrigin + "/register"
|
||||
}
|
||||
if tokenRedirect != "" {
|
||||
metadataToken = issuerOrigin + "/token"
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"issuer": issuerOrigin,
|
||||
"authorization_endpoint": issuerOrigin + "/authorize",
|
||||
"token_endpoint": metadataToken,
|
||||
"registration_endpoint": metadataRegistration,
|
||||
})
|
||||
case r.Host == issuerHost && r.URL.Path == "/register" && registrationRedirect != "":
|
||||
http.Redirect(w, r, registrationRedirect, http.StatusTemporaryRedirect)
|
||||
case r.Host == issuerHost && r.URL.Path == "/token" && tokenRedirect != "":
|
||||
http.Redirect(w, r, tokenRedirect, http.StatusTemporaryRedirect)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
t.Cleanup(gateway.Close)
|
||||
port := gateway.Listener.Addr().(*net.TCPAddr).Port
|
||||
resourceOrigin = fmt.Sprintf("https://93.184.216.34:%d", port)
|
||||
issuerOrigin = fmt.Sprintf("https://93.184.216.35:%d", port)
|
||||
|
||||
transport := gateway.Client().Transport.(*http.Transport).Clone()
|
||||
transport.TLSClientConfig = transport.TLSClientConfig.Clone()
|
||||
transport.TLSClientConfig.InsecureSkipVerify = true //nolint:gosec // hermetic public-address routing fixture
|
||||
dialer := &net.Dialer{}
|
||||
gatewayAddress := gateway.Listener.Addr().String()
|
||||
transport.DialContext = func(ctx context.Context, network string, address string) (net.Conn, error) {
|
||||
host, _, err := net.SplitHostPort(address)
|
||||
if err == nil && (host == "93.184.216.34" || host == "93.184.216.35") {
|
||||
return dialer.DialContext(ctx, network, gatewayAddress)
|
||||
}
|
||||
return dialer.DialContext(ctx, network, address)
|
||||
}
|
||||
return publicProtectedOAuthFixture{
|
||||
client: &http.Client{Transport: transport},
|
||||
resourceURL: resourceOrigin + "/mcp",
|
||||
}
|
||||
}
|
||||
|
||||
func TestProtectedResourceProbeDoesNotFollowRedirects(t *testing.T) {
|
||||
var redirectedHits atomic.Int64
|
||||
redirected := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
redirectedHits.Add(1)
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
}))
|
||||
defer redirected.Close()
|
||||
source := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Redirect(w, r, redirected.URL+"/mcp", http.StatusFound)
|
||||
}))
|
||||
defer source.Close()
|
||||
|
||||
_, found, err := discoverProtectedResourceAuthorizationServer(context.Background(), source.Client(), source.URL+"/mcp")
|
||||
if err != nil || found {
|
||||
t.Fatalf("redirected probe found=%v err=%v", found, err)
|
||||
}
|
||||
if redirectedHits.Load() != 0 {
|
||||
t.Fatalf("probe followed redirect %d time(s)", redirectedHits.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestProtectedResourceMetadataDoesNotFollowRedirects(t *testing.T) {
|
||||
var redirectedHits atomic.Int64
|
||||
redirected := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
redirectedHits.Add(1)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"resource": "https://attacker.example/mcp",
|
||||
"authorization_servers": []string{"https://attacker.example"},
|
||||
})
|
||||
}))
|
||||
defer redirected.Close()
|
||||
source := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
base := "http://" + r.Host
|
||||
switch r.URL.Path {
|
||||
case "/mcp":
|
||||
w.Header().Set("WWW-Authenticate", `Bearer resource_metadata="`+base+`/metadata"`)
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
case "/metadata":
|
||||
http.Redirect(w, r, redirected.URL+"/metadata", http.StatusFound)
|
||||
}
|
||||
}))
|
||||
defer source.Close()
|
||||
|
||||
_, _, err := discoverProtectedResourceAuthorizationServer(context.Background(), source.Client(), source.URL+"/mcp")
|
||||
if err == nil || !strings.Contains(err.Error(), "HTTP 302") {
|
||||
t.Fatalf("error = %v, want redirect refusal", err)
|
||||
}
|
||||
if redirectedHits.Load() != 0 {
|
||||
t.Fatalf("metadata fetch followed redirect %d time(s)", redirectedHits.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAuthorizationServerRejectsProtectedIssuerMismatch(t *testing.T) {
|
||||
authorizationServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
base := "http://" + r.Host
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"issuer": "https://unexpected.example",
|
||||
"authorization_endpoint": base + "/authorize",
|
||||
"token_endpoint": base + "/token",
|
||||
})
|
||||
}))
|
||||
defer authorizationServer.Close()
|
||||
resourceServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
base := "http://" + r.Host
|
||||
switch r.URL.Path {
|
||||
case "/mcp":
|
||||
w.Header().Set("WWW-Authenticate", `Bearer resource_metadata="`+base+`/metadata"`)
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
case "/metadata":
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"resource": base + "/mcp",
|
||||
"authorization_servers": []string{authorizationServer.URL},
|
||||
})
|
||||
}
|
||||
}))
|
||||
defer resourceServer.Close()
|
||||
|
||||
_, err := resolveAuthorizationServer(context.Background(), resourceServer.Client(), resourceServer.URL+"/mcp", OAuthConfig{})
|
||||
if err == nil || !strings.Contains(err.Error(), "issuer does not match") {
|
||||
t.Fatalf("error = %v, want issuer mismatch", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAuthorizationServerRejectsMissingProtectedIssuer(t *testing.T) {
|
||||
authorizationServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
base := "http://" + r.Host
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"authorization_endpoint": base + "/authorize",
|
||||
"token_endpoint": base + "/token",
|
||||
})
|
||||
}))
|
||||
defer authorizationServer.Close()
|
||||
resourceServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
base := "http://" + r.Host
|
||||
switch r.URL.Path {
|
||||
case "/mcp":
|
||||
w.Header().Set("WWW-Authenticate", `Bearer resource_metadata="`+base+`/metadata"`)
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
case "/metadata":
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"resource": base + "/mcp",
|
||||
"authorization_servers": []string{authorizationServer.URL},
|
||||
})
|
||||
}
|
||||
}))
|
||||
defer resourceServer.Close()
|
||||
|
||||
_, err := resolveAuthorizationServer(context.Background(), resourceServer.Client(), resourceServer.URL+"/mcp", OAuthConfig{})
|
||||
if err == nil || !strings.Contains(err.Error(), "issuer") {
|
||||
t.Fatalf("error = %v, want missing issuer rejection", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAuthorizationServerRejectsWhitespaceAlteredProtectedIssuer(t *testing.T) {
|
||||
for _, alter := range []struct {
|
||||
name string
|
||||
apply func(string) string
|
||||
}{
|
||||
{name: "leading", apply: func(value string) string { return " " + value }},
|
||||
{name: "trailing", apply: func(value string) string { return value + " " }},
|
||||
} {
|
||||
t.Run(alter.name, func(t *testing.T) {
|
||||
authorizationServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
base := "http://" + r.Host
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"issuer": alter.apply(base),
|
||||
"authorization_endpoint": base + "/authorize",
|
||||
"token_endpoint": base + "/token",
|
||||
})
|
||||
}))
|
||||
defer authorizationServer.Close()
|
||||
resourceServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
base := "http://" + r.Host
|
||||
switch r.URL.Path {
|
||||
case "/mcp":
|
||||
w.Header().Set("WWW-Authenticate", `Bearer resource_metadata="`+base+`/metadata"`)
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
case "/metadata":
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"resource": base + "/mcp",
|
||||
"authorization_servers": []string{authorizationServer.URL},
|
||||
})
|
||||
}
|
||||
}))
|
||||
defer resourceServer.Close()
|
||||
|
||||
_, err := resolveAuthorizationServer(context.Background(), resourceServer.Client(), resourceServer.URL+"/mcp", OAuthConfig{})
|
||||
if err == nil || !strings.Contains(err.Error(), "issuer does not match") {
|
||||
t.Fatalf("error = %v, want whitespace issuer mismatch", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthorizationServerMetadataDoesNotFollowRedirects(t *testing.T) {
|
||||
var redirectedHits atomic.Int64
|
||||
redirected := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
redirectedHits.Add(1)
|
||||
base := "http://" + r.Host
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"issuer": base,
|
||||
"authorization_endpoint": base + "/authorize",
|
||||
"token_endpoint": base + "/token",
|
||||
})
|
||||
}))
|
||||
defer redirected.Close()
|
||||
authorizationServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Redirect(w, r, redirected.URL+"/.well-known/oauth-authorization-server", http.StatusFound)
|
||||
}))
|
||||
defer authorizationServer.Close()
|
||||
resourceServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
base := "http://" + r.Host
|
||||
switch r.URL.Path {
|
||||
case "/mcp":
|
||||
w.Header().Set("WWW-Authenticate", `Bearer resource_metadata="`+base+`/metadata"`)
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
case "/metadata":
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"resource": base + "/mcp",
|
||||
"authorization_servers": []string{authorizationServer.URL},
|
||||
})
|
||||
}
|
||||
}))
|
||||
defer resourceServer.Close()
|
||||
|
||||
_, err := resolveAuthorizationServer(context.Background(), resourceServer.Client(), resourceServer.URL+"/mcp", OAuthConfig{})
|
||||
if err == nil {
|
||||
t.Fatal("authorization metadata redirect should fail")
|
||||
}
|
||||
if redirectedHits.Load() != 0 {
|
||||
t.Fatalf("authorization metadata fetch followed redirect %d time(s)", redirectedHits.Load())
|
||||
}
|
||||
}
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -145,6 +146,130 @@ func TestResolveEndpointsConfigOverridesDiscovery(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAuthorizationServerUsesProtectedResourceMetadata(t *testing.T) {
|
||||
var authorizationMetadataHits atomic.Int64
|
||||
authorizationServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/.well-known/oauth-authorization-server" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
authorizationMetadataHits.Add(1)
|
||||
base := "http://" + r.Host
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"issuer": base,
|
||||
"authorization_endpoint": base + "/authorize",
|
||||
"token_endpoint": base + "/token",
|
||||
"registration_endpoint": base + "/register",
|
||||
})
|
||||
}))
|
||||
defer authorizationServer.Close()
|
||||
|
||||
var resourceProbeHits, resourceMetadataHits atomic.Int64
|
||||
resourceServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/mcp":
|
||||
resourceProbeHits.Add(1)
|
||||
metadataURL := "http://" + r.Host + "/resource-metadata"
|
||||
w.Header().Set("WWW-Authenticate", `Bearer realm="OAuth", resource_metadata="`+metadataURL+`"`)
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
case "/resource-metadata":
|
||||
resourceMetadataHits.Add(1)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"resource": "http://" + r.Host + "/mcp",
|
||||
"authorization_servers": []string{authorizationServer.URL},
|
||||
})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer resourceServer.Close()
|
||||
|
||||
metadata, err := resolveAuthorizationServer(context.Background(), resourceServer.Client(), resourceServer.URL+"/mcp", OAuthConfig{})
|
||||
if err != nil {
|
||||
t.Fatalf("resolveAuthorizationServer() error = %v", err)
|
||||
}
|
||||
if metadata.Issuer != authorizationServer.URL {
|
||||
t.Fatalf("issuer = %q, want %q", metadata.Issuer, authorizationServer.URL)
|
||||
}
|
||||
if metadata.AuthorizationEndpoint != authorizationServer.URL+"/authorize" || metadata.TokenEndpoint != authorizationServer.URL+"/token" {
|
||||
t.Fatalf("resolved metadata = %#v", metadata)
|
||||
}
|
||||
if resourceProbeHits.Load() != 1 || resourceMetadataHits.Load() != 1 || authorizationMetadataHits.Load() != 1 {
|
||||
t.Fatalf("discovery hits = probe:%d resource:%d authorization:%d", resourceProbeHits.Load(), resourceMetadataHits.Load(), authorizationMetadataHits.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAuthorizationServerFallsBackToDirectDiscovery(t *testing.T) {
|
||||
var resourceProbeHits, directDiscoveryHits atomic.Int64
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/mcp":
|
||||
resourceProbeHits.Add(1)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
if flusher, ok := w.(http.Flusher); ok {
|
||||
flusher.Flush()
|
||||
}
|
||||
<-r.Context().Done()
|
||||
case "/.well-known/oauth-authorization-server/mcp":
|
||||
directDiscoveryHits.Add(1)
|
||||
base := "http://" + r.Host
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"issuer": base + "/mcp",
|
||||
"authorization_endpoint": base + "/authorize",
|
||||
"token_endpoint": base + "/token",
|
||||
})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
metadata, err := resolveAuthorizationServer(context.Background(), server.Client(), server.URL+"/mcp", OAuthConfig{})
|
||||
if err != nil {
|
||||
t.Fatalf("resolveAuthorizationServer() error = %v", err)
|
||||
}
|
||||
if metadata.AuthorizationEndpoint != server.URL+"/authorize" || metadata.TokenEndpoint != server.URL+"/token" {
|
||||
t.Fatalf("resolved metadata = %#v", metadata)
|
||||
}
|
||||
if resourceProbeHits.Load() != 1 || directDiscoveryHits.Load() != 1 {
|
||||
t.Fatalf("discovery hits = probe:%d direct:%d", resourceProbeHits.Load(), directDiscoveryHits.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAuthorizationServerIssuerOverrideSkipsResourceProbe(t *testing.T) {
|
||||
var resourceProbeHits, issuerDiscoveryHits atomic.Int64
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
base := "http://" + r.Host
|
||||
switch r.URL.Path {
|
||||
case "/mcp":
|
||||
resourceProbeHits.Add(1)
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
case "/.well-known/oauth-authorization-server/issuer":
|
||||
issuerDiscoveryHits.Add(1)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"issuer": base + "/issuer",
|
||||
"authorization_endpoint": base + "/authorize",
|
||||
"token_endpoint": base + "/token",
|
||||
})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
metadata, err := resolveAuthorizationServer(context.Background(), server.Client(), server.URL+"/mcp", OAuthConfig{IssuerURL: server.URL + "/issuer"})
|
||||
if err != nil {
|
||||
t.Fatalf("resolveAuthorizationServer() error = %v", err)
|
||||
}
|
||||
if metadata.AuthorizationEndpoint != server.URL+"/authorize" || metadata.TokenEndpoint != server.URL+"/token" {
|
||||
t.Fatalf("resolved metadata = %#v", metadata)
|
||||
}
|
||||
if resourceProbeHits.Load() != 0 || issuerDiscoveryHits.Load() != 1 {
|
||||
t.Fatalf("discovery hits = resource:%d issuer:%d", resourceProbeHits.Load(), issuerDiscoveryHits.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestPKCEChallengeIsS256OfVerifier(t *testing.T) {
|
||||
pkce, err := newPKCE()
|
||||
if err != nil {
|
||||
@@ -448,3 +573,98 @@ func TestFullFlowStoresTokens(t *testing.T) {
|
||||
t.Fatalf("login token = %#v", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFullFlowUsesProtectedResourceDiscovery(t *testing.T) {
|
||||
var resourceProbeHits, resourceMetadataHits, authorizationMetadataHits, tokenHits atomic.Int64
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
base := "http://" + r.Host
|
||||
switch r.URL.Path {
|
||||
case "/mcp":
|
||||
resourceProbeHits.Add(1)
|
||||
if r.Header.Get("Authorization") != "" || r.Header.Get("Cookie") != "" {
|
||||
t.Error("resource probe sent credentials")
|
||||
}
|
||||
w.Header().Set("WWW-Authenticate", `Bearer resource_metadata="`+base+`/resource-metadata"`)
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
case "/resource-metadata":
|
||||
resourceMetadataHits.Add(1)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"resource": base + "/mcp",
|
||||
"authorization_servers": []string{base},
|
||||
})
|
||||
case "/.well-known/oauth-authorization-server":
|
||||
authorizationMetadataHits.Add(1)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"issuer": base,
|
||||
"authorization_endpoint": base + "/authorize",
|
||||
"token_endpoint": base + "/token",
|
||||
})
|
||||
case "/token":
|
||||
tokenHits.Add(1)
|
||||
if err := r.ParseForm(); err != nil {
|
||||
t.Errorf("parse token form: %v", err)
|
||||
}
|
||||
if r.Form.Get("code") != "authorization-code-secret" || r.Form.Get("code_verifier") == "" {
|
||||
t.Error("token exchange omitted the authorization code or PKCE verifier")
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"access_token": "access-token-secret",
|
||||
"refresh_token": "refresh-token-secret",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600,
|
||||
})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
opener := func(authURL string) error {
|
||||
parsed, err := url.Parse(authURL)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if parsed.Path != "/authorize" || parsed.Query().Get("code_challenge_method") != "S256" || parsed.Query().Get("code_challenge") == "" {
|
||||
return errors.New("authorization request omitted the discovered endpoint or PKCE challenge")
|
||||
}
|
||||
if parsed.Query().Get("code") != "" || parsed.Query().Get("code_verifier") != "" {
|
||||
return errors.New("authorization request exposed an authorization code or PKCE verifier")
|
||||
}
|
||||
state := parsed.Query().Get("state")
|
||||
redirect := parsed.Query().Get("redirect_uri")
|
||||
go func() {
|
||||
callbackURL := redirect + "?code=authorization-code-secret&state=" + url.QueryEscape(state)
|
||||
for attempt := 0; attempt < 20; attempt++ {
|
||||
response, requestErr := http.Get(callbackURL)
|
||||
if requestErr == nil {
|
||||
response.Body.Close()
|
||||
return
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
}()
|
||||
return nil
|
||||
}
|
||||
|
||||
token, err := Login(context.Background(), LoginOptions{
|
||||
ServerName: "notion-style",
|
||||
ServerURL: server.URL + "/mcp",
|
||||
Config: OAuthConfig{
|
||||
ClientID: "client-123",
|
||||
Scopes: []string{"read"},
|
||||
},
|
||||
HTTPClient: server.Client(),
|
||||
OpenBrowser: opener,
|
||||
Timeout: 5 * time.Second,
|
||||
Now: time.Now,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Login() error = %v", err)
|
||||
}
|
||||
if token.AccessToken != "access-token-secret" || token.RefreshToken != "refresh-token-secret" {
|
||||
t.Fatal("login did not return the issued tokens")
|
||||
}
|
||||
if resourceProbeHits.Load() != 1 || resourceMetadataHits.Load() != 1 || authorizationMetadataHits.Load() != 1 || tokenHits.Load() != 1 {
|
||||
t.Fatalf("flow hits = probe:%d resource:%d authorization:%d token:%d", resourceProbeHits.Load(), resourceMetadataHits.Load(), authorizationMetadataHits.Load(), tokenHits.Load())
|
||||
}
|
||||
}
|
||||
@@ -176,9 +176,11 @@ func RegisterTools(ctx context.Context, registry *tools.Registry, cfg config.MCP
|
||||
staged = append(staged, tool)
|
||||
}
|
||||
}
|
||||
for _, tool := range staged {
|
||||
registry.Register(tool)
|
||||
registered := make([]tools.Tool, 0, len(staged))
|
||||
for index := range staged {
|
||||
registered = append(registered, staged[index])
|
||||
}
|
||||
registry.RegisterBatch(registered)
|
||||
return runtime, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -70,6 +70,62 @@ 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
@@ -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))
|
||||
}
|
||||
|
||||
@@ -191,6 +191,7 @@ 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"`
|
||||
@@ -479,6 +480,7 @@ 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)
|
||||
@@ -531,8 +533,9 @@ 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 (cached %d) out=%d | requests=%d tools=%d retries=%d reconnects=%d compactions=%d",
|
||||
result.Totals.InputTokens, result.Totals.CachedInputTokens, result.Totals.OutputTokens,
|
||||
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,
|
||||
result.Totals.ModelRequests, result.Totals.ToolCalls, result.Totals.Retries,
|
||||
result.Totals.Reconnects, result.Totals.Compactions))
|
||||
for _, class := range sortedClasses(result.PerClass) {
|
||||
@@ -545,6 +548,23 @@ 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.
|
||||
|
||||
@@ -74,6 +74,7 @@ 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,
|
||||
@@ -155,6 +156,9 @@ 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)
|
||||
}
|
||||
@@ -334,6 +338,35 @@ 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)
|
||||
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
"github.com/Gitlawb/zero/internal/providercatalog"
|
||||
"github.com/Gitlawb/zero/internal/providers"
|
||||
"github.com/Gitlawb/zero/internal/providers/providerio"
|
||||
"github.com/Gitlawb/zero/internal/proxydial"
|
||||
"github.com/Gitlawb/zero/internal/redaction"
|
||||
)
|
||||
|
||||
@@ -518,9 +519,15 @@ func newConnectivityClient(timeout time.Duration, resolver Resolver, sensitiveHe
|
||||
transport = &http.Transport{}
|
||||
}
|
||||
transport.DialContext = safeDialContext(resolver, allowLoopbackOrPrivate)
|
||||
// WRAPPED, SO THE DIALER LEARNS THE ROUTE THIS TRANSPORT CHOSE. The wrapper
|
||||
// records the proxy selected for each request on that request, and the
|
||||
// dialer reads it back on the dial that request causes; installing the bare
|
||||
// transport instead would leave a proxied dial looking like a loopback
|
||||
// target and fail the probe (#569). The cloned default transport already
|
||||
// reads HTTPS_PROXY and friends; nothing new is enabled.
|
||||
return &http.Client{
|
||||
Timeout: timeout,
|
||||
Transport: transport,
|
||||
Transport: proxydial.Wrap(transport),
|
||||
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
||||
if len(via) >= maxConnectivityRedirects {
|
||||
return fmt.Errorf("provider connectivity exceeded %d redirects", maxConnectivityRedirects)
|
||||
@@ -583,6 +590,17 @@ func sensitiveAuthHeaderNames(profile config.ProviderProfile, kind config.Provid
|
||||
func safeDialContext(resolver Resolver, allowLoopbackOrPrivate bool) func(context.Context, string, string) (net.Conn, error) {
|
||||
dialer := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}
|
||||
return func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
// THE PROXY IS NOT THE TARGET. With HTTPS_PROXY set to a local forward
|
||||
// proxy, the transport dials the proxy first and tunnels the request
|
||||
// through it, and this dialer used to refuse that dial as a loopback
|
||||
// address: "proxyconnect tcp: ... loopback hosts are blocked" for a
|
||||
// request whose real target had already been validated (#569). The
|
||||
// target validation is unchanged; only the dial the transport opens to
|
||||
// the proxy it picked for this very request is let past the guard, and a
|
||||
// direct dial is checked as it always was.
|
||||
if proxydial.IsProxyDial(ctx) {
|
||||
return dialer.DialContext(ctx, network, address)
|
||||
}
|
||||
host, port, err := net.SplitHostPort(address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -0,0 +1,218 @@
|
||||
package providerhealth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/proxydial"
|
||||
)
|
||||
|
||||
// connectivityClientWithProxy builds the real connectivity client and swaps in
|
||||
// a proxy function, which is how the transport learns its route. The
|
||||
// environment is left alone deliberately: http.ProxyFromEnvironment reads it
|
||||
// once per process, so an environment-driven test would depend on test order.
|
||||
func connectivityClientWithProxy(t *testing.T, resolver Resolver, proxyFor func(*http.Request) (*url.URL, error)) *http.Client {
|
||||
t.Helper()
|
||||
client := newConnectivityClient(3*time.Second, resolver, nil, false)
|
||||
wrapped, ok := client.Transport.(*proxydial.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("transport = %T, want the route-recording wrapper; without it no dial can be recognised as the proxy dial", client.Transport)
|
||||
}
|
||||
wrapped.Proxy = proxyFor
|
||||
return client
|
||||
}
|
||||
|
||||
// recordingProxy is a listener standing in for a local forward proxy. It
|
||||
// records the first line it receives and hangs up.
|
||||
func recordingProxy(t *testing.T) (*url.URL, <-chan string) {
|
||||
t.Helper()
|
||||
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = listener.Close() })
|
||||
arrived := make(chan string, 1)
|
||||
go func() {
|
||||
conn, err := listener.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
buffer := make([]byte, 256)
|
||||
n, _ := conn.Read(buffer)
|
||||
arrived <- string(buffer[:n])
|
||||
}()
|
||||
proxyURL, err := url.Parse("http://" + listener.Addr().String())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return proxyURL, arrived
|
||||
}
|
||||
|
||||
// THE PROXY IS NOT THE TARGET, AND THE GUARD HAS TO KNOW THE DIFFERENCE.
|
||||
//
|
||||
// With HTTPS_PROXY set to a local forward proxy, the transport dials the proxy
|
||||
// first and tunnels the request through it. This dialer refused that dial as a
|
||||
// loopback address, so `zero doctor --connectivity` failed with "proxyconnect
|
||||
// tcp: ... loopback hosts are blocked" for a target that had already been
|
||||
// validated and was public (#569). Driven through the real client, so the route
|
||||
// is the transport's own decision.
|
||||
func TestConnectivityClientTunnelsThroughTheConfiguredProxy(t *testing.T) {
|
||||
proxyURL, arrived := recordingProxy(t)
|
||||
client := connectivityClientWithProxy(t,
|
||||
staticResolver{err: errors.New("resolver must not be called for a proxied request")},
|
||||
func(*http.Request) (*url.URL, error) { return proxyURL, nil })
|
||||
|
||||
request, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://api.example.com/v1/models", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = client.Do(request)
|
||||
select {
|
||||
case first := <-arrived:
|
||||
if line := strings.SplitN(first, "\r\n", 2)[0]; line != "CONNECT api.example.com:443 HTTP/1.1" {
|
||||
t.Fatalf("proxy received %q, want a CONNECT for the validated target", line)
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatalf("the proxy never received a connection; request error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// AND A DIRECT REQUEST STAYS GUARDED, EVEN WHEN ITS ORIGIN IS THE PROXY'S OWN
|
||||
// ADDRESS.
|
||||
//
|
||||
// This is the NO_PROXY shape: the proxy function sends one host direct. An
|
||||
// exemption derived from asking that function about a request of our own
|
||||
// invention would answer "proxied" for the excluded host too, and this dial,
|
||||
// whose address equals the configured proxy endpoint, would skip the resolver
|
||||
// and reach whatever the name resolves to at connect time.
|
||||
func TestDirectRequestToTheProxyEndpointIsStillGuarded(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://rebind.example:80")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
lookups := 0
|
||||
resolver := countingResolver{
|
||||
onLookup: func() { lookups++ },
|
||||
addr: netip.MustParseAddr("10.0.0.5"),
|
||||
}
|
||||
client := connectivityClientWithProxy(t, resolver, func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Hostname() == "rebind.example" {
|
||||
return nil, nil
|
||||
}
|
||||
return proxyURL, nil
|
||||
})
|
||||
|
||||
request, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "http://rebind.example/data", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = client.Do(request)
|
||||
|
||||
if lookups == 0 {
|
||||
t.Fatal("the direct dial skipped the safety resolver because its address matched the configured proxy endpoint")
|
||||
}
|
||||
var safety endpointSafetyError
|
||||
if !errors.As(err, &safety) {
|
||||
t.Fatalf("err = %v, want the private answer to be refused", err)
|
||||
}
|
||||
}
|
||||
|
||||
// A proxy configured with a Unicode hostname is dialed by its IDNA-encoded
|
||||
// name, and it is still the proxy: the request must reach it rather than be
|
||||
// refused as an ordinary private address.
|
||||
func TestUnicodeProxyHostnameIsStillTheProxy(t *testing.T) {
|
||||
dialed := ""
|
||||
client := connectivityClientWithProxy(t,
|
||||
staticResolver{err: errors.New("resolver must not be called for a proxied request")},
|
||||
func(*http.Request) (*url.URL, error) { return url.Parse("http://büro.example:3067") })
|
||||
wrapped := client.Transport.(*proxydial.Transport)
|
||||
guarded := wrapped.DialContext
|
||||
wrapped.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
dialed = address
|
||||
return guarded(ctx, network, address)
|
||||
}
|
||||
|
||||
request, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://api.example.com/v1/models", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = client.Do(request)
|
||||
|
||||
if dialed != "xn--bro-hoa.example:3067" {
|
||||
t.Fatalf("SETUP INVALID: dialed %q, want the IDNA-encoded proxy endpoint", dialed)
|
||||
}
|
||||
var safety endpointSafetyError
|
||||
if errors.As(err, &safety) {
|
||||
t.Fatalf("the configured proxy was refused by the guard: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Without a proxy the guard is exactly what it was: loopback is refused.
|
||||
func TestSafeDialContextStillRefusesLoopbackWithoutAProxy(t *testing.T) {
|
||||
dial := safeDialContext(staticResolver{err: errors.New("resolver must not be called")}, false)
|
||||
conn, err := dial(context.Background(), "tcp", "127.0.0.1:3067")
|
||||
if conn != nil {
|
||||
_ = conn.Close()
|
||||
t.Fatal("loopback was dialed with no proxy configured")
|
||||
}
|
||||
var safety endpointSafetyError
|
||||
if !errors.As(err, &safety) {
|
||||
t.Fatalf("err = %v, want endpointSafetyError", err)
|
||||
}
|
||||
}
|
||||
|
||||
// A sibling loopback port is not the proxy, and the guard is unchanged for it.
|
||||
func TestSafeDialContextRefusesALoopbackPortThatIsNotTheProxy(t *testing.T) {
|
||||
proxyURL, arrived := recordingProxy(t)
|
||||
client := connectivityClientWithProxy(t,
|
||||
staticResolver{err: errors.New("resolver must not be called")},
|
||||
func(request *http.Request) (*url.URL, error) {
|
||||
// Only the https probe is proxied; the http one below is direct and
|
||||
// happens to point at the neighbouring port.
|
||||
if request.URL.Scheme == "https" {
|
||||
return proxyURL, nil
|
||||
}
|
||||
return nil, nil
|
||||
})
|
||||
|
||||
_, port, _ := net.SplitHostPort(proxyURL.Host)
|
||||
number, _ := strconv.Atoi(port)
|
||||
sibling := "http://" + net.JoinHostPort("127.0.0.1", strconv.Itoa(number+1)) + "/"
|
||||
request, err := http.NewRequestWithContext(context.Background(), http.MethodGet, sibling, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = client.Do(request)
|
||||
var safety endpointSafetyError
|
||||
if !errors.As(err, &safety) {
|
||||
t.Fatalf("err = %v, want a loopback port beside the proxy to stay blocked", err)
|
||||
}
|
||||
select {
|
||||
case first := <-arrived:
|
||||
t.Fatalf("the proxy received a connection it should not have: %q", first)
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// countingResolver answers with one address and reports each lookup, so a test
|
||||
// can tell a skipped guard from a guard that ran and allowed the address.
|
||||
type countingResolver struct {
|
||||
onLookup func()
|
||||
addr netip.Addr
|
||||
}
|
||||
|
||||
func (resolver countingResolver) LookupNetIP(context.Context, string, string) ([]netip.Addr, error) {
|
||||
if resolver.onLookup != nil {
|
||||
resolver.onLookup()
|
||||
}
|
||||
return []netip.Addr{resolver.addr}, nil
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// Compaction summarization model selection.
|
||||
//
|
||||
// Compaction summaries used to run on the session's main model — near a 200k
|
||||
// window one summarization is a ~100k-token input call at full price, the
|
||||
// single most expensive recurring event in a long run. A cheap model handles
|
||||
// the "dense factual summary" task fine, so summarization is routed to one
|
||||
// whenever the target is knowable.
|
||||
//
|
||||
// Resolution order: ZERO_COMPACTION_MODEL env > preferences.compactionModel
|
||||
// config > a curated per-kind default (OFFICIAL endpoints only) > "" (use the
|
||||
// main model). The literal value "main" at any level forces the main model.
|
||||
// A wrong choice is never fatal: the agent falls back to the main provider on
|
||||
// the first summarizer failure for the rest of the run.
|
||||
|
||||
// CompactionModelEnv overrides the compaction summarization model.
|
||||
const CompactionModelEnv = "ZERO_COMPACTION_MODEL"
|
||||
|
||||
// compactionMainSentinel disables the dedicated summarizer explicitly.
|
||||
const compactionMainSentinel = "main"
|
||||
|
||||
// Curated cheap summarizer defaults, applied only for provider kinds whose
|
||||
// model catalog is guaranteed (official endpoints). Custom / *-compatible
|
||||
// endpoints get NO default — their catalogs are unknowable, and a guaranteed
|
||||
// failed call per run is worse than main-model prices.
|
||||
const (
|
||||
defaultAnthropicCompactionModel = "claude-haiku-4-5-20251001"
|
||||
defaultGoogleCompactionModel = "gemini-2.5-flash-lite"
|
||||
defaultOpenAICompactionModel = "gpt-4.1-mini"
|
||||
)
|
||||
|
||||
// CompactionModelID returns the model that compaction summarization calls
|
||||
// should use for the given profile, or "" to use the session's main model.
|
||||
// preference is the resolved config value (preferences.compactionModel).
|
||||
func CompactionModelID(profile config.ProviderProfile, preference string) string {
|
||||
value := strings.TrimSpace(os.Getenv(CompactionModelEnv))
|
||||
if value == "" {
|
||||
value = strings.TrimSpace(preference)
|
||||
}
|
||||
if value != "" {
|
||||
if strings.EqualFold(value, compactionMainSentinel) {
|
||||
return ""
|
||||
}
|
||||
return value
|
||||
}
|
||||
return defaultCompactionModel(profile)
|
||||
}
|
||||
|
||||
func defaultCompactionModel(profile config.ProviderProfile) string {
|
||||
metadata, err := ResolveRuntimeMetadata(profile, Options{})
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
var candidate string
|
||||
switch metadata.ProviderKind {
|
||||
case config.ProviderKindAnthropic:
|
||||
candidate = defaultAnthropicCompactionModel
|
||||
case config.ProviderKindGoogle:
|
||||
candidate = defaultGoogleCompactionModel
|
||||
case config.ProviderKindOpenAI:
|
||||
candidate = defaultOpenAICompactionModel
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
// Session already runs on the cheap model: a dedicated summarizer buys
|
||||
// nothing and would only add a second provider instance.
|
||||
if strings.EqualFold(strings.TrimSpace(metadata.APIModel), candidate) {
|
||||
return ""
|
||||
}
|
||||
return candidate
|
||||
}
|
||||
|
||||
// CompactionSummarizerFactory returns the lazy constructor the agent loop
|
||||
// calls for its dedicated compaction summarizer (agent.Options.Summarizer).
|
||||
// The compactor passes the main model in force at that moment — on the first
|
||||
// paid compaction and again after a mid-run model switch — so the choice
|
||||
// tracks escalations: a run that starts on the cheap model gets no dedicated
|
||||
// summarizer, and gains one once it escalates to a pricier main model.
|
||||
// Returning (nil, nil) means "summarize on the main provider". The factory
|
||||
// itself is nil only when no provider builder exists. Shared by the exec and
|
||||
// TUI front ends so the selection rules cannot drift between them.
|
||||
func CompactionSummarizerFactory(base config.ProviderProfile, preference string, newProvider func(config.ProviderProfile) (zeroruntime.Provider, error)) func(context.Context, string) (zeroruntime.Provider, error) {
|
||||
if newProvider == nil {
|
||||
return nil
|
||||
}
|
||||
return func(_ context.Context, mainModelID string) (zeroruntime.Provider, error) {
|
||||
profile := base
|
||||
if strings.TrimSpace(mainModelID) != "" {
|
||||
profile.Model = mainModelID
|
||||
}
|
||||
modelID := CompactionModelID(profile, preference)
|
||||
if modelID == "" {
|
||||
return nil, nil
|
||||
}
|
||||
profile.Model = modelID
|
||||
return newProvider(profile)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
func TestCompactionModelIDResolutionOrder(t *testing.T) {
|
||||
anthropic := config.ProviderProfile{
|
||||
Name: "anthropic",
|
||||
ProviderKind: config.ProviderKindAnthropic,
|
||||
Model: "claude-sonnet-4.5",
|
||||
}
|
||||
|
||||
t.Setenv(CompactionModelEnv, "my-env-model")
|
||||
if got := CompactionModelID(anthropic, "my-config-model"); got != "my-env-model" {
|
||||
t.Fatalf("env must win, got %q", got)
|
||||
}
|
||||
|
||||
t.Setenv(CompactionModelEnv, "main")
|
||||
if got := CompactionModelID(anthropic, "my-config-model"); got != "" {
|
||||
t.Fatalf("env 'main' must force the main model, got %q", got)
|
||||
}
|
||||
|
||||
t.Setenv(CompactionModelEnv, "")
|
||||
if got := CompactionModelID(anthropic, "my-config-model"); got != "my-config-model" {
|
||||
t.Fatalf("config preference must apply when env unset, got %q", got)
|
||||
}
|
||||
if got := CompactionModelID(anthropic, "main"); got != "" {
|
||||
t.Fatalf("config 'main' must force the main model, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactionModelIDCuratedDefaults(t *testing.T) {
|
||||
t.Setenv(CompactionModelEnv, "")
|
||||
|
||||
anthropic := config.ProviderProfile{Name: "anthropic", ProviderKind: config.ProviderKindAnthropic, Model: "claude-sonnet-4.5"}
|
||||
if got := CompactionModelID(anthropic, ""); got != defaultAnthropicCompactionModel {
|
||||
t.Fatalf("anthropic default = %q, want %q", got, defaultAnthropicCompactionModel)
|
||||
}
|
||||
|
||||
google := config.ProviderProfile{Name: "google", ProviderKind: config.ProviderKindGoogle, Model: "gemini-2.5-pro"}
|
||||
if got := CompactionModelID(google, ""); got != defaultGoogleCompactionModel {
|
||||
t.Fatalf("google default = %q, want %q", got, defaultGoogleCompactionModel)
|
||||
}
|
||||
|
||||
openai := config.ProviderProfile{Name: "openai", ProviderKind: config.ProviderKindOpenAI, Model: "gpt-4.1"}
|
||||
if got := CompactionModelID(openai, ""); got != defaultOpenAICompactionModel {
|
||||
t.Fatalf("openai default = %q, want %q", got, defaultOpenAICompactionModel)
|
||||
}
|
||||
|
||||
// Custom/compatible endpoints: catalog unknowable, no default.
|
||||
compatible := config.ProviderProfile{
|
||||
Name: "opengateway",
|
||||
ProviderKind: config.ProviderKindOpenAICompatible,
|
||||
BaseURL: "https://opengateway.gitlawb.com/v1",
|
||||
Model: "some-model",
|
||||
}
|
||||
if got := CompactionModelID(compatible, ""); got != "" {
|
||||
t.Fatalf("openai-compatible must have no default, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactionModelIDSkipsWhenMainIsAlreadyCheap(t *testing.T) {
|
||||
t.Setenv(CompactionModelEnv, "")
|
||||
haiku := config.ProviderProfile{
|
||||
Name: "anthropic",
|
||||
ProviderKind: config.ProviderKindAnthropic,
|
||||
Model: defaultAnthropicCompactionModel,
|
||||
}
|
||||
if got := CompactionModelID(haiku, ""); got != "" {
|
||||
t.Fatalf("already-cheap session must not get a dedicated summarizer, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactionSummarizerFactoryTracksMainModel(t *testing.T) {
|
||||
t.Setenv(CompactionModelEnv, "")
|
||||
base := config.ProviderProfile{Name: "anthropic", ProviderKind: config.ProviderKindAnthropic, Model: defaultAnthropicCompactionModel}
|
||||
var built []string
|
||||
newProvider := func(profile config.ProviderProfile) (zeroruntime.Provider, error) {
|
||||
built = append(built, profile.Model)
|
||||
return nil, errors.New("not needed")
|
||||
}
|
||||
factory := CompactionSummarizerFactory(base, "", newProvider)
|
||||
if factory == nil {
|
||||
t.Fatal("factory must exist when a provider builder exists")
|
||||
}
|
||||
// Main model is the cheap default: no dedicated summarizer, nothing built.
|
||||
if provider, err := factory(context.Background(), defaultAnthropicCompactionModel); provider != nil || err != nil || len(built) != 0 {
|
||||
t.Fatalf("cheap main model must yield (nil, nil) without building, got %v %v built=%v", provider, err, built)
|
||||
}
|
||||
// After escalating to an expensive model the cheap default is built.
|
||||
if _, err := factory(context.Background(), "claude-sonnet-4-5"); err == nil || len(built) != 1 || built[0] != defaultAnthropicCompactionModel {
|
||||
t.Fatalf("escalated main model must build the cheap summarizer, built=%v err=%v", built, err)
|
||||
}
|
||||
// An empty main model falls back to the base profile's model — the cheap
|
||||
// default here, so again nothing is built.
|
||||
built = nil
|
||||
if provider, err := factory(context.Background(), ""); provider != nil || err != nil || len(built) != 0 {
|
||||
t.Fatalf("empty main model must use the base profile, got %v %v built=%v", provider, err, built)
|
||||
}
|
||||
// The explicit "main" preference disables the dedicated summarizer.
|
||||
if provider, err := CompactionSummarizerFactory(base, "main", newProvider)(context.Background(), "claude-sonnet-4-5"); provider != nil || err != nil {
|
||||
t.Fatalf("preference main must disable the summarizer, got %v %v", provider, err)
|
||||
}
|
||||
if CompactionSummarizerFactory(base, "", nil) != nil {
|
||||
t.Fatal("no provider builder means no factory")
|
||||
}
|
||||
}
|
||||
@@ -62,6 +62,7 @@ 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"
|
||||
@@ -70,17 +71,37 @@ 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 {
|
||||
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"`
|
||||
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"`
|
||||
}
|
||||
|
||||
// responsesReasoning carries the reasoning controls for the Responses API. The
|
||||
@@ -105,6 +126,7 @@ 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"`
|
||||
}
|
||||
@@ -118,10 +140,11 @@ 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"`
|
||||
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"`
|
||||
}
|
||||
|
||||
// responsesEvent is the decoded form of one Responses SSE data payload.
|
||||
@@ -132,6 +155,7 @@ 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"`
|
||||
@@ -150,6 +174,7 @@ 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"`
|
||||
}
|
||||
|
||||
@@ -173,7 +198,8 @@ type outputDetails struct {
|
||||
}
|
||||
|
||||
type inputDetails struct {
|
||||
CachedTokens int `json:"cached_tokens"`
|
||||
CachedTokens int `json:"cached_tokens"`
|
||||
CacheWriteTokens int `json:"cache_write_tokens"`
|
||||
}
|
||||
|
||||
type errorPayload struct {
|
||||
@@ -187,8 +213,10 @@ 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
|
||||
}
|
||||
@@ -197,15 +225,57 @@ 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
|
||||
usage *usagePayload
|
||||
done bool
|
||||
toolCalls map[string]*toolCallBuilder
|
||||
toolOrder []string
|
||||
text strings.Builder
|
||||
responseID string
|
||||
usage *usagePayload
|
||||
emitted bool
|
||||
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
|
||||
@@ -215,11 +285,14 @@ func (p *CodexProvider) buildResponsesRequest(request zeroruntime.CompletionRequ
|
||||
return nil, errors.New("codex provider: model is required")
|
||||
}
|
||||
req := &responsesRequest{
|
||||
Model: p.inner.model,
|
||||
Stream: true,
|
||||
MaxOutputTokens: p.inner.maxTokens,
|
||||
Model: p.inner.model,
|
||||
Stream: true,
|
||||
ParallelToolCalls: true,
|
||||
MaxOutputTokens: p.inner.maxTokens,
|
||||
PromptCacheKey: request.PromptCacheKey,
|
||||
}
|
||||
instructions := []string{}
|
||||
customCallIDs := map[string]string{}
|
||||
for _, msg := range request.Messages {
|
||||
switch msg.Role {
|
||||
case zeroruntime.MessageRoleSystem:
|
||||
@@ -234,13 +307,31 @@ 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: "function_call_output",
|
||||
CallID: msg.ToolCallID,
|
||||
Type: outputType,
|
||||
CallID: callID,
|
||||
Output: msg.Content,
|
||||
})
|
||||
default:
|
||||
@@ -249,12 +340,25 @@ func (p *CodexProvider) buildResponsesRequest(request zeroruntime.CompletionRequ
|
||||
}
|
||||
req.Instructions = strings.Join(instructions, "\n\n")
|
||||
for _, tool := range request.Tools {
|
||||
req.Tools = append(req.Tools, responsesTool{
|
||||
Type: "function",
|
||||
Name: tool.Name,
|
||||
Description: tool.Description,
|
||||
Parameters: tool.Parameters,
|
||||
})
|
||||
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)
|
||||
}
|
||||
// Codex / o-series reasoning models take the effort tier nested under
|
||||
// `reasoning` on the Responses API (the chat-completions `reasoning_effort`
|
||||
@@ -312,13 +416,26 @@ func (p *CodexProvider) assistantInputItems(msg zeroruntime.Message) ([]inputIte
|
||||
if tc.Name == "" {
|
||||
return nil, errors.New("codex provider: assistant tool call missing name")
|
||||
}
|
||||
items = append(items, inputItem{
|
||||
Type: "function_call",
|
||||
ID: tc.ID,
|
||||
CallID: tc.ID,
|
||||
Name: tc.Name,
|
||||
Arguments: tc.Arguments,
|
||||
})
|
||||
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)
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
@@ -479,6 +596,18 @@ 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
|
||||
@@ -500,12 +629,17 @@ func (p *CodexProvider) emitResponsesEvent(
|
||||
// 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,
|
||||
@@ -517,6 +651,7 @@ func (p *CodexProvider) emitResponsesEvent(
|
||||
// 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,
|
||||
@@ -524,13 +659,16 @@ func (p *CodexProvider) emitResponsesEvent(
|
||||
}
|
||||
return true
|
||||
case responsesEventFunctionArgsDelta:
|
||||
p.handleFunctionArgsDelta(ctx, &event, state, events)
|
||||
p.handleToolInputDelta(ctx, event, state, events, false)
|
||||
return true
|
||||
case responsesEventCustomInputDelta:
|
||||
p.handleToolInputDelta(ctx, event, state, events, true)
|
||||
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 != "" {
|
||||
@@ -570,22 +708,38 @@ func (p *CodexProvider) handleOutputItemAdded(
|
||||
return
|
||||
}
|
||||
switch event.Item.Type {
|
||||
case "function_call":
|
||||
case "function_call", "custom_tool_call":
|
||||
key := p.toolCallKey(event)
|
||||
if key == "" {
|
||||
return
|
||||
}
|
||||
builder := &toolCallBuilder{
|
||||
ID: key,
|
||||
Name: event.Item.Name,
|
||||
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
|
||||
}
|
||||
state.toolCalls[key] = builder
|
||||
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
|
||||
Type: zeroruntime.StreamEventToolCallStart,
|
||||
ToolCallID: key,
|
||||
ToolName: event.Item.Name,
|
||||
})
|
||||
builder.started = true
|
||||
}
|
||||
}
|
||||
|
||||
@@ -593,11 +747,12 @@ 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) handleFunctionArgsDelta(
|
||||
func (p *CodexProvider) handleToolInputDelta(
|
||||
ctx context.Context,
|
||||
event *responsesEvent,
|
||||
state *responsesState,
|
||||
events chan<- zeroruntime.StreamEvent,
|
||||
freeform bool,
|
||||
) {
|
||||
key := p.toolCallKey(event)
|
||||
if key == "" {
|
||||
@@ -610,10 +765,16 @@ func (p *CodexProvider) handleFunctionArgsDelta(
|
||||
// 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}
|
||||
builder = &toolCallBuilder{ID: key, Freeform: freeform}
|
||||
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,
|
||||
@@ -632,7 +793,7 @@ func (p *CodexProvider) handleOutputItemDone(
|
||||
state *responsesState,
|
||||
events chan<- zeroruntime.StreamEvent,
|
||||
) {
|
||||
if event.Item == nil || event.Item.Type != "function_call" {
|
||||
if event.Item == nil || (event.Item.Type != "function_call" && event.Item.Type != "custom_tool_call") {
|
||||
return
|
||||
}
|
||||
key := p.toolCallKey(event)
|
||||
@@ -641,20 +802,31 @@ func (p *CodexProvider) handleOutputItemDone(
|
||||
}
|
||||
builder, ok := state.toolCalls[key]
|
||||
if !ok {
|
||||
builder = &toolCallBuilder{ID: key}
|
||||
builder = &toolCallBuilder{ID: key, Freeform: event.Item.Type == "custom_tool_call"}
|
||||
state.toolCalls[key] = builder
|
||||
state.toolOrder = append(state.toolOrder, key)
|
||||
}
|
||||
if event.Item.Name != "" {
|
||||
builder.Name = event.Item.Name
|
||||
}
|
||||
if event.Item.Arguments != "" && builder.Arguments.Len() == 0 {
|
||||
builder.Arguments.WriteString(event.Item.Arguments)
|
||||
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 !builder.started {
|
||||
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
|
||||
Type: zeroruntime.StreamEventToolCallStart,
|
||||
ToolCallID: key,
|
||||
ToolName: builder.Name,
|
||||
Type: zeroruntime.StreamEventToolCallStart,
|
||||
ToolCallID: key,
|
||||
ToolCallProviderID: builder.CallID,
|
||||
ToolName: builder.Name,
|
||||
ToolCallFreeform: builder.Freeform,
|
||||
})
|
||||
builder.started = true
|
||||
}
|
||||
@@ -664,6 +836,7 @@ func (p *CodexProvider) handleOutputItemDone(
|
||||
ToolName: builder.Name,
|
||||
})
|
||||
builder.ended = true
|
||||
state.emitted = true
|
||||
}
|
||||
|
||||
// handleTerminalResponse emits the final usage + done event for a
|
||||
@@ -693,6 +866,7 @@ 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,
|
||||
@@ -700,6 +874,7 @@ 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
|
||||
@@ -735,7 +910,7 @@ func (p *CodexProvider) handleTerminalResponse(
|
||||
state.done = true
|
||||
return false
|
||||
}
|
||||
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone})
|
||||
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone, ResponseID: event.Response.ID})
|
||||
state.done = true
|
||||
return false
|
||||
}
|
||||
@@ -754,6 +929,9 @@ 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
|
||||
|
||||
@@ -0,0 +1,421 @@
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,590 @@
|
||||
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,6 +2,8 @@ package openai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
@@ -25,6 +27,126 @@ 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{}
|
||||
|
||||
@@ -90,11 +212,14 @@ 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{Status: "completed"}}, &responsesState{}, ok)
|
||||
&responsesEvent{Type: responsesEventCompleted, Response: &responsePayload{ID: "resp-1", 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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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}}}}`,
|
||||
`{"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}}}}`,
|
||||
)
|
||||
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.ReasoningTokens != 1 {
|
||||
t.Fatalf("usage = %+v, want input=7 output=3 cached=2 reasoning=1", *usage)
|
||||
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 !gotDone {
|
||||
t.Fatal("expected a done event, got none")
|
||||
|
||||
@@ -357,6 +357,7 @@ 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,
|
||||
},
|
||||
})
|
||||
|
||||
@@ -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}}}`)
|
||||
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, `[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 {
|
||||
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 {
|
||||
t.Fatalf("unexpected usage event: %#v", events[2])
|
||||
}
|
||||
if events[3].Type != zeroruntime.StreamEventDone {
|
||||
|
||||
@@ -195,7 +195,12 @@ func (s *turnSession) computeFingerprint(request zeroruntime.CompletionRequest)
|
||||
func wireToolsDigest(tools []zeroruntime.ToolDefinition) string {
|
||||
rendered := make([]string, 0, len(tools))
|
||||
for _, tool := range tools {
|
||||
schema, err := json.Marshal(tool.Parameters)
|
||||
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)
|
||||
if err != nil {
|
||||
rendered = append(rendered, tool.Name+"\n"+tool.Description+"\n__non_json:"+tool.Name)
|
||||
continue
|
||||
|
||||
@@ -109,7 +109,8 @@ type usage struct {
|
||||
}
|
||||
|
||||
type promptTokenDetails struct {
|
||||
CachedTokens int `json:"cached_tokens"`
|
||||
CachedTokens int `json:"cached_tokens"`
|
||||
CacheWriteTokens int `json:"cache_write_tokens"`
|
||||
}
|
||||
|
||||
type completionTokenDetails struct {
|
||||
|
||||
@@ -15,9 +15,27 @@ 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 != "" && value != "0" && !strings.EqualFold(value, "false")
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
// OptimizedTurnSessions returns the gated optimized TurnSessionProvider for an
|
||||
@@ -25,20 +43,33 @@ func openaiTurnSessionEnabled() 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: 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.
|
||||
// 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.
|
||||
func OptimizedTurnSessions(profile config.ProviderProfile, provider zeroruntime.Provider, options Options) (zeroruntime.TurnSessionProvider, bool) {
|
||||
if !openaiTurnSessionEnabled() {
|
||||
resolved, err := resolveProfile(profile, options)
|
||||
if err != nil {
|
||||
return nil, false
|
||||
}
|
||||
resolved, err := resolveProfile(profile, options)
|
||||
if err != nil || resolved.providerKind != config.ProviderKindOpenAI || isCodexCatalog(profile, resolved) {
|
||||
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
|
||||
}
|
||||
concrete, ok := provider.(*openai.Provider)
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/providers/openai"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
@@ -45,7 +46,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", " ", ""} {
|
||||
for _, value := range []string{"0", "false", "FALSE", "False", "off", "OFF", " ", ""} {
|
||||
t.Setenv(openaiTurnSessionEnv, value)
|
||||
if _, ok := OptimizedTurnSessions(profile, provider, Options{}); ok {
|
||||
t.Fatalf("gate enabled for falsey value %q", value)
|
||||
@@ -59,6 +60,21 @@ 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()
|
||||
@@ -88,13 +104,64 @@ func TestOptimizedTurnSessionsRejectsCompatible(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestOptimizedTurnSessionsRejectsCodexCatalog(t *testing.T) {
|
||||
t.Setenv(openaiTurnSessionEnv, "1")
|
||||
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")
|
||||
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("the ChatGPT Codex catalog must not get the optimized session")
|
||||
t.Fatal("a ChatGPT catalog profile with a non-Codex provider must use the default session")
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
// Package proxydial answers one question for an SSRF-guarded dialer: is this
|
||||
// dial the connection dial of a request the transport routed through a forward
|
||||
// proxy?
|
||||
//
|
||||
// A guarded transport validates the TARGET of a request before dialing, and its
|
||||
// dialer refuses loopback, private and link-local addresses so a redirect or a
|
||||
// DNS rebind cannot walk the request onto the local network. When a forward
|
||||
// proxy is configured, the transport dials the PROXY first and tunnels the
|
||||
// request through it, and a local proxy is by far the common shape:
|
||||
// HTTPS_PROXY=http://127.0.0.1:3067. The dialer then sees 127.0.0.1, refuses it,
|
||||
// and the user gets "loopback hosts are blocked" for a request whose real target
|
||||
// was validated and public (#569).
|
||||
//
|
||||
// The proxy address is not the request target, and the guard was never about it.
|
||||
// It is configuration the user set on purpose, and the transport is going to dial
|
||||
// it for every request whether the guard likes it or not. So that dial is let
|
||||
// through, and only that dial.
|
||||
//
|
||||
// THE EXEMPTION FOLLOWS THE TRANSPORT'S OWN ROUTE DECISION, NOT A GUESS AT IT.
|
||||
// Asking the Proxy function a question of one's own invention answers for a
|
||||
// request the transport never made: with NO_PROXY excluding a host, the real
|
||||
// request to that host goes direct while a synthetic probe still comes back
|
||||
// proxied, and a direct dial whose address happens to equal the proxy endpoint
|
||||
// would then skip the guard entirely. Host-dependent proxy callbacks make the
|
||||
// mismatch general, and no amount of care over the probe's spelling closes it,
|
||||
// because the two calls are about different requests.
|
||||
//
|
||||
// So the decision is recorded where it is made, on the request, and read back on
|
||||
// the dial that request causes. Transport.RoundTrip asks the wrapped transport's
|
||||
// Proxy function about the actual request, stores the answer in that request's
|
||||
// context, and hands it on; the transport derives the dial context from the
|
||||
// request context, so the dialer reads back the route this very connection is
|
||||
// being opened for. Per-request, so concurrent requests cannot overwrite each
|
||||
// other, and unforgeable from outside the package: the context key is private.
|
||||
package proxydial
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/url"
|
||||
)
|
||||
|
||||
// selectedProxyKey addresses the route decision inside a request context. It is
|
||||
// unexported so nothing outside this package can plant one.
|
||||
type selectedProxyKey struct{}
|
||||
|
||||
// Transport records the proxy its wrapped transport selects for each request in
|
||||
// that request's context, so a guarded DialContext installed on the same
|
||||
// transport can tell the dial to a proxy from a direct dial.
|
||||
//
|
||||
// The embedded transport stays reachable: callers configure Proxy, DialContext
|
||||
// and the rest on it exactly as before, and tests can read them back.
|
||||
type Transport struct {
|
||||
*http.Transport
|
||||
}
|
||||
|
||||
// Wrap returns a Transport around transport. The result is the http.RoundTripper
|
||||
// to install on the client; installing the bare transport instead leaves every
|
||||
// dial looking direct, which is safe (the guard applies) but reopens #569.
|
||||
func Wrap(transport *http.Transport) *Transport {
|
||||
return &Transport{Transport: transport}
|
||||
}
|
||||
|
||||
// RoundTrip resolves the route for this request and passes the decision down to
|
||||
// the dialer through the request context.
|
||||
func (transport *Transport) RoundTrip(request *http.Request) (*http.Response, error) {
|
||||
if transport.Transport.Proxy != nil {
|
||||
// A Proxy error is left for the transport to report: it fails the
|
||||
// request before dialing, so there is no dial to classify. The function
|
||||
// is called once here and once by the transport, which is what a
|
||||
// host-dependent callback already expects from redirects and retries.
|
||||
if proxyURL, err := transport.Transport.Proxy(request); err == nil && proxyURL != nil {
|
||||
request = request.WithContext(context.WithValue(request.Context(), selectedProxyKey{}, proxyURL))
|
||||
}
|
||||
}
|
||||
return transport.Transport.RoundTrip(request)
|
||||
}
|
||||
|
||||
// IsProxyDial reports whether ctx belongs to a request this package's Transport
|
||||
// routed through a proxy, which makes the dial about to happen the dial to that
|
||||
// proxy.
|
||||
//
|
||||
// THE ADDRESS IS NOT COMPARED, AND DOES NOT NEED TO BE. A transport opens
|
||||
// exactly one connection dial per connection, and when a proxy is selected that
|
||||
// dial goes to the proxy; the target is reached through it, never dialed. So the
|
||||
// route decision on the context already identifies the dial, and matching the
|
||||
// address against the proxy URL on top of it could only ever disagree with the
|
||||
// transport, which is what a punycoded Unicode proxy hostname does: Go dials
|
||||
// xn--bro-hoa.example while the URL still spells buero as the user wrote it.
|
||||
func IsProxyDial(ctx context.Context) bool {
|
||||
proxyURL, _ := ctx.Value(selectedProxyKey{}).(*url.URL)
|
||||
return proxyURL != nil
|
||||
}
|
||||
@@ -0,0 +1,224 @@
|
||||
package proxydial
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// errDialed ends the request at the dial, after the route decision has been
|
||||
// observed. Nothing here needs a live server.
|
||||
var errDialed = errors.New("dial recorded")
|
||||
|
||||
type dialRecord struct {
|
||||
address string
|
||||
proxied bool
|
||||
}
|
||||
|
||||
// driveRequest sends one request through a wrapped transport whose dialer
|
||||
// records what it was asked for and whether the context says this dial belongs
|
||||
// to a proxied request.
|
||||
func driveRequest(t *testing.T, target string, proxyFor func(*http.Request) (*url.URL, error)) dialRecord {
|
||||
t.Helper()
|
||||
var record dialRecord
|
||||
transport := &http.Transport{Proxy: proxyFor}
|
||||
transport.DialContext = func(ctx context.Context, _, address string) (net.Conn, error) {
|
||||
record = dialRecord{address: address, proxied: IsProxyDial(ctx)}
|
||||
return nil, errDialed
|
||||
}
|
||||
client := &http.Client{Transport: Wrap(transport)}
|
||||
request, err := http.NewRequest(http.MethodGet, target, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := client.Do(request); !errors.Is(err, errDialed) {
|
||||
t.Fatalf("%s: request ended before the dial: %v", target, err)
|
||||
}
|
||||
if record.address == "" {
|
||||
t.Fatalf("%s: the dialer was never called", target)
|
||||
}
|
||||
return record
|
||||
}
|
||||
|
||||
func fixedProxy(raw string) func(*http.Request) (*url.URL, error) {
|
||||
parsed, err := url.Parse(raw)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return func(*http.Request) (*url.URL, error) { return parsed, nil }
|
||||
}
|
||||
|
||||
// THE DIAL TO A CONFIGURED PROXY IS THE ONE THAT IS EXEMPT. A proxied request
|
||||
// dials the proxy, never the target, so the route on the context is what
|
||||
// identifies the dial.
|
||||
func TestProxiedRequestDialsTheProxyAndSaysSo(t *testing.T) {
|
||||
record := driveRequest(t, "https://api.example.com/v1/models", fixedProxy("http://127.0.0.1:3067"))
|
||||
if !record.proxied {
|
||||
t.Error("a request routed through a proxy did not reach the dialer as proxied")
|
||||
}
|
||||
if record.address != "127.0.0.1:3067" {
|
||||
t.Errorf("dial address = %q, want the proxy endpoint", record.address)
|
||||
}
|
||||
}
|
||||
|
||||
// AND A DIRECT REQUEST IS NOT, EVEN WHEN ITS ORIGIN IS THE PROXY'S OWN ADDRESS.
|
||||
//
|
||||
// This is the NO_PROXY shape: the proxy function excludes one host, so a request
|
||||
// to that host goes direct. Asking that same function about a request of our own
|
||||
// invention would answer "proxied" (the invented host is not excluded) and the
|
||||
// direct dial, whose address equals the proxy endpoint, would skip the guard
|
||||
// while resolving to whatever the name resolves to at connect time.
|
||||
func TestDirectRequestToTheProxyEndpointIsNotAProxyDial(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://rebind.example:80")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
noProxyForRebind := func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Hostname() == "rebind.example" {
|
||||
return nil, nil
|
||||
}
|
||||
return proxyURL, nil
|
||||
}
|
||||
|
||||
record := driveRequest(t, "http://rebind.example/data", noProxyForRebind)
|
||||
if record.address != "rebind.example:80" {
|
||||
t.Fatalf("SETUP INVALID: dial address = %q, want the direct origin, which is also the proxy endpoint", record.address)
|
||||
}
|
||||
if record.proxied {
|
||||
t.Error("a direct dial was reported as the proxy dial because its address matched the configured proxy")
|
||||
}
|
||||
|
||||
// The same function still routes another host through that proxy, so the
|
||||
// exclusion above is a route difference and not a disabled proxy.
|
||||
if record := driveRequest(t, "http://elsewhere.example/data", noProxyForRebind); !record.proxied {
|
||||
t.Error("the configured proxy was not used for a host outside the exclusion")
|
||||
}
|
||||
}
|
||||
|
||||
// The inverse mismatch: a callback that routes by host must not have its answer
|
||||
// taken from some other host. A proxy selected only for the real target is
|
||||
// honoured.
|
||||
func TestHostDependentProxySelectionIsHonoured(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://127.0.0.1:3067")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
onlyForOneHost := func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Hostname() == "api.example.com" {
|
||||
return proxyURL, nil
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
if record := driveRequest(t, "https://api.example.com/v1", onlyForOneHost); !record.proxied || record.address != "127.0.0.1:3067" {
|
||||
t.Errorf("selected proxy not reached: %+v", record)
|
||||
}
|
||||
if record := driveRequest(t, "https://other.example/v1", onlyForOneHost); record.proxied {
|
||||
t.Error("an unproxied host was reported as proxied")
|
||||
}
|
||||
}
|
||||
|
||||
// Scheme separation: an https-only proxy leaves an http request direct, and the
|
||||
// decision recorded is the one made for that request's own scheme.
|
||||
func TestSchemeSeparation(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://127.0.0.1:3067")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
httpsOnly := func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Scheme == "https" {
|
||||
return proxyURL, nil
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
if record := driveRequest(t, "https://api.example.com/v1", httpsOnly); !record.proxied {
|
||||
t.Error("the https proxy was not applied to an https request")
|
||||
}
|
||||
if record := driveRequest(t, "http://api.example.com/v1", httpsOnly); record.proxied {
|
||||
t.Error("an http request was treated as proxied by an https-only proxy")
|
||||
}
|
||||
}
|
||||
|
||||
// A UNICODE PROXY HOSTNAME IS STILL THE PROXY. Go dials proxies by their
|
||||
// IDNA-encoded name, so a matcher comparing the URL's own spelling would call
|
||||
// this dial direct and refuse it as an ordinary private address. Nothing is
|
||||
// compared here, so both spellings of one host are one host.
|
||||
func TestUnicodeProxyHostnameIsRecognised(t *testing.T) {
|
||||
unicodeRecord := driveRequest(t, "https://api.example.com/v1", fixedProxy("http://büro.example:3067"))
|
||||
if !unicodeRecord.proxied {
|
||||
t.Error("a proxy configured with a Unicode hostname was not recognised as the proxy")
|
||||
}
|
||||
asciiRecord := driveRequest(t, "https://api.example.com/v1", fixedProxy("http://xn--bro-hoa.example:3067"))
|
||||
if !asciiRecord.proxied {
|
||||
t.Error("a proxy configured with the ASCII spelling was not recognised as the proxy")
|
||||
}
|
||||
// The two spellings name one endpoint, and it is the ASCII one that is
|
||||
// dialed. This is the mismatch a string comparison would have had to close.
|
||||
if unicodeRecord.address != asciiRecord.address {
|
||||
t.Errorf("dial addresses differ across spellings of one host: %q vs %q", unicodeRecord.address, asciiRecord.address)
|
||||
}
|
||||
}
|
||||
|
||||
// A transport with no proxy leaves every dial direct.
|
||||
func TestNoProxyMeansNoExemption(t *testing.T) {
|
||||
if record := driveRequest(t, "http://127.0.0.1:3067/", nil); record.proxied {
|
||||
t.Error("a dial was reported as proxied by a transport with no Proxy function")
|
||||
}
|
||||
if IsProxyDial(context.Background()) {
|
||||
t.Error("a context that never went through the wrapper reported a proxy dial")
|
||||
}
|
||||
}
|
||||
|
||||
// THE DECISION IS PER REQUEST. A shared record of the last selected proxy would
|
||||
// let a proxied request hand its exemption to a direct one running beside it.
|
||||
func TestConcurrentRequestsKeepTheirOwnRoute(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://127.0.0.1:3067")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
byHost := func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Hostname() == "proxied.example" {
|
||||
return proxyURL, nil
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var mutex sync.Mutex
|
||||
seen := map[string]bool{}
|
||||
transport := &http.Transport{Proxy: byHost}
|
||||
transport.DialContext = func(ctx context.Context, _, address string) (net.Conn, error) {
|
||||
mutex.Lock()
|
||||
seen[address] = IsProxyDial(ctx)
|
||||
mutex.Unlock()
|
||||
return nil, errDialed
|
||||
}
|
||||
client := &http.Client{Transport: Wrap(transport)}
|
||||
|
||||
var group sync.WaitGroup
|
||||
for range 40 {
|
||||
for _, target := range []string{"http://proxied.example/", "http://direct.example/"} {
|
||||
group.Add(1)
|
||||
go func() {
|
||||
defer group.Done()
|
||||
request, _ := http.NewRequest(http.MethodGet, target, nil)
|
||||
_, _ = client.Do(request)
|
||||
}()
|
||||
}
|
||||
}
|
||||
group.Wait()
|
||||
|
||||
mutex.Lock()
|
||||
defer mutex.Unlock()
|
||||
if !seen["127.0.0.1:3067"] {
|
||||
t.Error("the proxy dial was not recorded as proxied under concurrency")
|
||||
}
|
||||
if seen["direct.example:80"] {
|
||||
t.Error("a direct dial picked up another request's proxy exemption")
|
||||
}
|
||||
if len(seen) != 2 {
|
||||
t.Fatalf("SETUP INVALID: dialed %v, want exactly the proxy and the direct origin", seen)
|
||||
}
|
||||
}
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
type BackendOptions struct {
|
||||
GOOS string
|
||||
LookupExecutable func(string) (string, error)
|
||||
DetectWSL func() WSLInfo
|
||||
}
|
||||
|
||||
type Backend struct {
|
||||
@@ -53,6 +54,7 @@ func SelectBackend(options BackendOptions) Backend {
|
||||
return NewSandboxManager(SandboxManagerOptions{
|
||||
GOOS: options.GOOS,
|
||||
LookupExecutable: options.LookupExecutable,
|
||||
DetectWSL: options.DetectWSL,
|
||||
}).Backend()
|
||||
}
|
||||
|
||||
|
||||
@@ -45,6 +45,7 @@ func TestSelectBackendChoosesPlatformAdapterWithFallback(t *testing.T) {
|
||||
}
|
||||
return "", errors.New("missing")
|
||||
},
|
||||
DetectWSL: func() WSLInfo { return WSLInfo{} },
|
||||
})
|
||||
if backend.Name != BackendUnavailable || backend.Available {
|
||||
t.Fatalf("linux backend = %#v, want native sandbox unavailable without Linux helper", backend)
|
||||
@@ -54,6 +55,20 @@ func TestSelectBackendChoosesPlatformAdapterWithFallback(t *testing.T) {
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("linux helper missing under WSL reports WSL fallback", func(t *testing.T) {
|
||||
backend := SelectBackend(BackendOptions{
|
||||
GOOS: "linux",
|
||||
LookupExecutable: func(string) (string, error) { return "", errors.New("missing") },
|
||||
DetectWSL: func() WSLInfo { return WSLInfo{IsWSL: true, IsWSL2: true} },
|
||||
})
|
||||
if backend.Name != BackendWSL || backend.Available || !backend.Fallback {
|
||||
t.Fatalf("linux backend = %#v, want explicit WSL fallback", backend)
|
||||
}
|
||||
if !strings.Contains(backend.Message, "WSL2") {
|
||||
t.Fatalf("linux fallback message = %q, want WSL2 diagnostic", backend.Message)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("darwin sandbox exec available", func(t *testing.T) {
|
||||
backend := SelectBackend(BackendOptions{
|
||||
GOOS: "darwin",
|
||||
|
||||
@@ -51,6 +51,7 @@ const (
|
||||
type SandboxManagerOptions struct {
|
||||
GOOS string
|
||||
LookupExecutable func(string) (string, error)
|
||||
DetectWSL func() WSLInfo
|
||||
Backend Backend
|
||||
}
|
||||
|
||||
@@ -96,7 +97,7 @@ func NewSandboxManager(options SandboxManagerOptions) SandboxManager {
|
||||
goos = runtime.GOOS
|
||||
}
|
||||
if backend.Name == "" {
|
||||
backend = selectPlatformBackend(goos, options.LookupExecutable)
|
||||
backend = selectPlatformBackend(goos, options.LookupExecutable, options.DetectWSL)
|
||||
}
|
||||
if backend.Platform == "" {
|
||||
backend.Platform = goos
|
||||
@@ -168,10 +169,13 @@ func lookupExecutable(name string) (string, error) {
|
||||
return "", errors.New("executable file not found")
|
||||
}
|
||||
|
||||
func selectPlatformBackend(goos string, lookup func(string) (string, error)) Backend {
|
||||
func selectPlatformBackend(goos string, lookup func(string) (string, error), detect func() WSLInfo) Backend {
|
||||
if lookup == nil {
|
||||
lookup = lookupExecutable
|
||||
}
|
||||
if detect == nil {
|
||||
detect = detectWSL
|
||||
}
|
||||
switch goos {
|
||||
case "linux":
|
||||
if helper, err := lookup(LinuxSandboxHelperName); err == nil && helper != "" {
|
||||
@@ -180,7 +184,7 @@ func selectPlatformBackend(goos string, lookup func(string) (string, error)) Bac
|
||||
}
|
||||
return nativeBackend(goos, BackendLinuxBwrap, helper, "Linux sandbox helper available")
|
||||
}
|
||||
if info := detectWSL(); info.IsWSL {
|
||||
if info := detect(); info.IsWSL {
|
||||
return wslBackend(goos, info)
|
||||
}
|
||||
return unavailableBackend(goos, "Linux sandbox helper is not available")
|
||||
|
||||
@@ -96,15 +96,28 @@ var protectedMetadataNames = []string{".git", ".zero", ".agents"}
|
||||
// gitMetadataWriteCarveouts below.
|
||||
var sandboxFullyProtectedMetadataNames = []string{".zero", ".agents"}
|
||||
|
||||
// gitMetadataWriteCarveouts returns the .git subpaths that stay write-denied
|
||||
// under the OS-level sandbox even though the rest of .git is writable to git
|
||||
// subprocesses. Nonexistent paths are harmless no-ops in every backend's
|
||||
// enforcement (seatbelt regex, bwrap ro-bind, Windows ACL deny entry).
|
||||
// gitMetadataWriteCarveouts returns the .git paths that stay write-denied under
|
||||
// the OS-level sandbox even though the rest of .git is writable to git
|
||||
// subprocesses. Worktrees and submodules store .git as a regular pointer file,
|
||||
// so they protect that file itself: constructing .git/hooks below it makes
|
||||
// bwrap fail before the command starts. All other entry and error states retain
|
||||
// the legacy child paths. The pointer is deliberately not parsed here;
|
||||
// repository-controlled metadata must not redirect sandbox rules elsewhere.
|
||||
func gitMetadataWriteCarveouts(root string) []string {
|
||||
return []string{
|
||||
filepath.Join(root, ".git", "hooks"),
|
||||
filepath.Join(root, ".git", "config"),
|
||||
return gitMetadataWriteCarveoutsWithLstat(root, os.Lstat)
|
||||
}
|
||||
|
||||
func gitMetadataWriteCarveoutsWithLstat(root string, lstat func(string) (os.FileInfo, error)) []string {
|
||||
gitPath := filepath.Join(root, ".git")
|
||||
children := []string{
|
||||
filepath.Join(gitPath, "hooks"),
|
||||
filepath.Join(gitPath, "config"),
|
||||
}
|
||||
info, err := lstat(gitPath)
|
||||
if err != nil || !info.Mode().IsRegular() {
|
||||
return children
|
||||
}
|
||||
return []string{gitPath}
|
||||
}
|
||||
|
||||
func PermissionProfileFromPolicy(workspaceRoot string, policy Policy, scope *Scope) PermissionProfile {
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
package sandbox
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func mustEqualCarveouts(t *testing.T, got, want []string) {
|
||||
t.Helper()
|
||||
if len(got) != len(want) {
|
||||
t.Fatalf("carveouts = %v, want %v", got, want)
|
||||
}
|
||||
for i := range want {
|
||||
if got[i] != want[i] {
|
||||
t.Fatalf("carveout[%d] = %q, want %q", i, got[i], want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Worktree/submodule checkouts store .git as a pointer file. Protect the file
|
||||
// itself and never create a bwrap mount target beneath it.
|
||||
func TestGitMetadataWriteCarveoutsProtectsWorktreePointer(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
gitPath := filepath.Join(root, ".git")
|
||||
if err := os.WriteFile(gitPath, []byte("gitdir: ../outside\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
got := gitMetadataWriteCarveouts(root)
|
||||
mustEqualCarveouts(t, got, []string{gitPath})
|
||||
for _, carveout := range got {
|
||||
if carveout == filepath.Join(gitPath, "hooks") || carveout == filepath.Join(gitPath, "config") {
|
||||
t.Fatalf("carveout points beneath .git file: %q", carveout)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLinuxBwrapPlanDoesNotMountBelowWorktreePointer(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
gitPath := filepath.Join(root, ".git")
|
||||
if err := os.WriteFile(gitPath, []byte("gitdir: ../outside\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
profile := PermissionProfile{FileSystem: FileSystemPolicy{
|
||||
Kind: FileSystemRestricted,
|
||||
WriteRoots: []WritableRoot{{
|
||||
Root: root,
|
||||
ReadOnlySubpaths: gitMetadataWriteCarveouts(root),
|
||||
}},
|
||||
}}
|
||||
args := linuxBwrapFilesystemArgs(profile)
|
||||
assertArgsContainSequence(t, args, "--ro-bind", gitPath, gitPath)
|
||||
for _, child := range []string{"hooks", "config"} {
|
||||
bogus := filepath.Join(gitPath, child)
|
||||
if argsContainSequence(args, "--tmpfs", bogus) || argsContainSequence(args, "--ro-bind", bogus, bogus) {
|
||||
t.Fatalf("bwrap args contain mount below .git file %q: %v", bogus, args)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGitMetadataWriteCarveoutsPlainCheckout(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
want := []string{
|
||||
filepath.Join(root, ".git", "hooks"),
|
||||
filepath.Join(root, ".git", "config"),
|
||||
}
|
||||
mustEqualCarveouts(t, gitMetadataWriteCarveouts(root), want)
|
||||
|
||||
if err := os.MkdirAll(filepath.Join(root, ".git"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
mustEqualCarveouts(t, gitMetadataWriteCarveouts(root), want)
|
||||
}
|
||||
|
||||
func TestGitMetadataWriteCarveoutsPreservesSymlinkBehavior(t *testing.T) {
|
||||
for _, targetKind := range []string{"directory", "file", "dangling"} {
|
||||
t.Run(targetKind, func(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
target := filepath.Join(root, "real.git")
|
||||
switch targetKind {
|
||||
case "directory":
|
||||
if err := os.MkdirAll(target, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
case "file":
|
||||
if err := os.WriteFile(target, []byte("gitdir: ../outside\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := os.Symlink(target, filepath.Join(root, ".git")); err != nil {
|
||||
t.Skipf("symlink unavailable: %v", err)
|
||||
}
|
||||
want := []string{
|
||||
filepath.Join(root, ".git", "hooks"),
|
||||
filepath.Join(root, ".git", "config"),
|
||||
}
|
||||
mustEqualCarveouts(t, gitMetadataWriteCarveouts(root), want)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGitMetadataWriteCarveoutsPreservesInspectionErrorBehavior(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
gitPath := filepath.Join(root, ".git")
|
||||
want := []string{
|
||||
filepath.Join(gitPath, "hooks"),
|
||||
filepath.Join(gitPath, "config"),
|
||||
}
|
||||
got := gitMetadataWriteCarveoutsWithLstat(root, func(string) (os.FileInfo, error) {
|
||||
return nil, errors.New("inspection failed")
|
||||
})
|
||||
mustEqualCarveouts(t, got, want)
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
//go:build unix
|
||||
|
||||
package sandbox
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"syscall"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestGitMetadataWriteCarveoutsPreservesSpecialGitEntryBehavior(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
gitPath := filepath.Join(root, ".git")
|
||||
if err := syscall.Mkfifo(gitPath, 0o644); err != nil {
|
||||
t.Skip(err)
|
||||
}
|
||||
mustEqualCarveouts(t, gitMetadataWriteCarveouts(root), []string{
|
||||
filepath.Join(gitPath, "hooks"),
|
||||
filepath.Join(gitPath, "config"),
|
||||
})
|
||||
}
|
||||
+239
-25
@@ -19,6 +19,17 @@ type Scope struct {
|
||||
workspaceRoot string
|
||||
readRoots []string
|
||||
extraRoots []string
|
||||
// tempReads / tempWrites count how many LIVE temporary grants depend on a
|
||||
// root, so one holder's cleanup cannot revoke another's access.
|
||||
//
|
||||
// Without them a temporary grant was add-then-remove with no notion of who
|
||||
// still needed it: the second caller to ask for a root already present got a
|
||||
// NO-OP undo, and the first caller's cleanup removed the root out from under
|
||||
// it. Two read-only tools in the same parallel batch, both blocked on the
|
||||
// same directory, is exactly that shape — and read-only tools are precisely
|
||||
// the ones the batch runs concurrently.
|
||||
tempReads map[string]int
|
||||
tempWrites map[string]int
|
||||
}
|
||||
|
||||
// NewScope builds a scope for workspaceRoot plus the given extra roots. The
|
||||
@@ -81,15 +92,88 @@ func (s *Scope) Add(path string) (string, error) {
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for _, existing := range append([]string{s.workspaceRoot}, s.extraRoots...) {
|
||||
if pathWithinRoot(existing, root) {
|
||||
// COVERED BY WHAT. extraRoots holds temporary write roots alongside
|
||||
// permanent ones, so "something already covers this" was not the question
|
||||
// worth asking: a permanent grant made over a path a temporary holder
|
||||
// happened to cover recorded nothing, and vanished the moment that holder
|
||||
// released. A session-scoped grant outliving the request that prompted it is
|
||||
// the entire difference between Add and AddTemporaryWrite.
|
||||
//
|
||||
// Permanent coverage is looked for FIRST and across every root, because a
|
||||
// path can be covered twice and the temporary cover must not decide the
|
||||
// answer when a permanent one is also present.
|
||||
if s.permanentWriteRootCoversLocked(root) {
|
||||
return root, nil
|
||||
}
|
||||
for _, existing := range s.extraRoots {
|
||||
// The same root, held temporarily: promote it in place. Its holders'
|
||||
// releases become no-ops, which is what permanence means here.
|
||||
if existing == root && s.temporaryWriteLocked(existing) {
|
||||
delete(s.tempWrites, root)
|
||||
return root, nil
|
||||
}
|
||||
}
|
||||
// Either uncovered, or covered only by a BROADER temporary root that is
|
||||
// going to be released. Recording the narrower root in its own right is what
|
||||
// survives that release.
|
||||
s.extraRoots = append(s.extraRoots, root)
|
||||
return root, nil
|
||||
}
|
||||
|
||||
// temporaryWriteLocked reports whether root is held by a temporary write grant
|
||||
// rather than a session-scoped one. Callers must hold the lock.
|
||||
func (s *Scope) temporaryWriteLocked(root string) bool {
|
||||
_, temporary := s.tempWrites[root]
|
||||
return temporary
|
||||
}
|
||||
|
||||
// permanentWriteRootCoversLocked reports whether root sits under write authority
|
||||
// that outlives every temporary holder: the workspace root, or a session-scoped
|
||||
// grant. Coverage by a TEMPORARY write root deliberately does not count —
|
||||
// it is going to be released, so nothing recorded on the strength of it
|
||||
// survives. Callers must hold the lock.
|
||||
//
|
||||
// One helper rather than a copy per caller because Add, AddRead and
|
||||
// AddTemporaryRead all turn on the same question, and three copies of "covered,
|
||||
// but by what" is how one of them ends up answering it differently.
|
||||
// permanentReadRootCoversLocked reports whether root is already readable by a
|
||||
// grant that outlives every temporary holder.
|
||||
//
|
||||
// SLICE ORDER MUST NOT DECIDE THIS. The old check walked readRoots and stopped at
|
||||
// the first entry covering the path, so a BROAD TEMPORARY root sitting earlier in
|
||||
// the slice shadowed a narrower PERMANENT one later — and the caller was handed a
|
||||
// reference on somebody else's grant when it needed no reference at all. Asking
|
||||
// the whole question first, over both lists, makes the answer independent of
|
||||
// insertion order.
|
||||
//
|
||||
// Write roots count: a grant that permits writing permits reading the same tree.
|
||||
func (s *Scope) permanentReadRootCoversLocked(root string) bool {
|
||||
if s.permanentWriteRootCoversLocked(root) {
|
||||
return true
|
||||
}
|
||||
for _, existing := range s.readRoots {
|
||||
if pathWithinRoot(existing, root) && !s.temporaryReadLocked(existing) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (s *Scope) permanentWriteRootCoversLocked(root string) bool {
|
||||
for _, existing := range append([]string{s.workspaceRoot}, s.extraRoots...) {
|
||||
if pathWithinRoot(existing, root) && !s.temporaryWriteLocked(existing) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// temporaryReadLocked is temporaryWriteLocked for read roots.
|
||||
func (s *Scope) temporaryReadLocked(root string) bool {
|
||||
_, temporary := s.tempReads[root]
|
||||
return temporary
|
||||
}
|
||||
|
||||
// AddRead grants read-only access under path. If the path is already covered by
|
||||
// a writable root, no separate read root is stored.
|
||||
func (s *Scope) AddRead(path string) (string, error) {
|
||||
@@ -99,11 +183,19 @@ func (s *Scope) AddRead(path string) (string, error) {
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.writeRootCoversLocked(root) {
|
||||
// Same rule as Add: coverage that is going to be released is not coverage a
|
||||
// permanent grant can rely on, whichever side of the read/write boundary it
|
||||
// sits on. A permanent read covered only by a temporary WRITE root died with
|
||||
// that root just as surely.
|
||||
if s.permanentWriteRootCoversLocked(root) {
|
||||
return root, nil
|
||||
}
|
||||
if s.permanentReadRootCoversLocked(root) {
|
||||
return root, nil
|
||||
}
|
||||
for _, existing := range s.readRoots {
|
||||
if pathWithinRoot(existing, root) {
|
||||
if existing == root && s.temporaryReadLocked(existing) {
|
||||
delete(s.tempReads, root)
|
||||
return root, nil
|
||||
}
|
||||
}
|
||||
@@ -118,16 +210,63 @@ func (s *Scope) AddTemporaryRead(path string) (string, func(), error) {
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.writeRootCoversLocked(root) {
|
||||
// Genuinely permanent — the workspace root, or a session-scoped grant. It
|
||||
// outlives this reader, so there is nothing to release.
|
||||
if s.permanentWriteRootCoversLocked(root) {
|
||||
return root, func() {}, nil
|
||||
}
|
||||
for _, existing := range s.readRoots {
|
||||
if pathWithinRoot(existing, root) {
|
||||
return root, func() {}, nil
|
||||
}
|
||||
// A TEMPORARY write root covering this path is NOT borrowed, and this is the
|
||||
// one place the distinction bites hardest.
|
||||
//
|
||||
// The dependency on that root is real — the write holder's cleanup must not
|
||||
// revoke this reader, which is the defect the refcount exists to close,
|
||||
// reached across the read/write boundary. But taking a reference on the
|
||||
// WRITE root pays for the lifetime with the authority: the reference keeps
|
||||
// the root in extraRoots, extraRoots is what Roots() feeds validate(), and
|
||||
// validate() is WRITE authorization. So a reader outliving its writer went on
|
||||
// writing anywhere under a write grant that had already ended — not merely
|
||||
// under the path it asked to read.
|
||||
//
|
||||
// One reference cannot be both the lifetime and the capability. The lifetime
|
||||
// is kept below, as a temporary READ root for the path this caller actually
|
||||
// asked for. readRoots feeds ReadRoots(), which confers no write authority —
|
||||
// validate() authorises writes from workspaceRoot plus extraRoots, and this
|
||||
// path never enters that list. So the reader survives its writer with read
|
||||
// authority and only read authority.
|
||||
//
|
||||
// That claim is the written justification for the whole design, so it has to
|
||||
// stay exhaustive as the accessors change. It has already been wrong once, in
|
||||
// both directions: an earlier draft said readRoots fed ReadRoots() and nothing
|
||||
// else while ExtraReadRoots() also read it, and that accessor has since been
|
||||
// removed as belonging to a different change.
|
||||
if s.permanentReadRootCoversLocked(root) {
|
||||
return root, func() {}, nil
|
||||
}
|
||||
if s.tempReads == nil {
|
||||
s.tempReads = map[string]int{}
|
||||
}
|
||||
// THE SAME ROOT, ANOTHER HOLDER: one entry, two references. Only an EXACT
|
||||
// match shares — a narrower request must not, which is the whole point below.
|
||||
if _, held := s.tempReads[root]; held {
|
||||
s.tempReads[root]++
|
||||
return root, oncePerHolder(func() { s.releaseTemporaryRead(root) }), nil
|
||||
}
|
||||
// A NARROWER REQUEST RECORDS ITS OWN ROOT, even under a broader temporary
|
||||
// read that currently covers it.
|
||||
//
|
||||
// Borrowing the covering root's reference kept this reader alive by keeping
|
||||
// the BROAD ROOT alive, and readRoots feeds validateRead and the native
|
||||
// sandbox profile. So once the broad reader released, a caller that had asked
|
||||
// to read one subdirectory was left able to read the whole tree its neighbour
|
||||
// had asked for, siblings included. Reported by @jatmn, and it is the read
|
||||
// twin of the write escalation fixed above: the refcount key decides not only
|
||||
// WHEN an entry disappears but WHICH path stays authorised while it lives.
|
||||
s.readRoots = append(s.readRoots, root)
|
||||
return root, func() { s.removeReadRoot(root) }, nil
|
||||
if s.tempReads == nil {
|
||||
s.tempReads = map[string]int{}
|
||||
}
|
||||
s.tempReads[root] = 1
|
||||
return root, oncePerHolder(func() { s.releaseTemporaryRead(root) }), nil
|
||||
}
|
||||
|
||||
func (s *Scope) AddTemporaryWrite(path string) (string, func(), error) {
|
||||
@@ -137,32 +276,107 @@ func (s *Scope) AddTemporaryWrite(path string) (string, func(), error) {
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.writeRootCoversLocked(root) {
|
||||
// A PERMANENT grant is the only cover that needs no bookkeeping: it outlives
|
||||
// every temporary holder, so there is nothing to keep alive and nothing to
|
||||
// release.
|
||||
if s.permanentWriteRootCoversLocked(root) {
|
||||
return root, func() {}, nil
|
||||
}
|
||||
s.extraRoots = append(s.extraRoots, root)
|
||||
return root, func() { s.removeWriteRoot(root) }, nil
|
||||
}
|
||||
|
||||
func (s *Scope) writeRootCoversLocked(root string) bool {
|
||||
for _, existing := range append([]string{s.workspaceRoot}, s.extraRoots...) {
|
||||
if pathWithinRoot(existing, root) {
|
||||
return true
|
||||
}
|
||||
if s.tempWrites == nil {
|
||||
s.tempWrites = map[string]int{}
|
||||
}
|
||||
return false
|
||||
// THE SAME ROOT, ANOTHER HOLDER: one entry, two references. Only an exact
|
||||
// match shares an entry — a NARROWER request must not, which is the whole
|
||||
// point below.
|
||||
if _, held := s.tempWrites[root]; held {
|
||||
s.tempWrites[root]++
|
||||
return root, oncePerHolder(func() { s.releaseTemporaryWrite(root) }), nil
|
||||
}
|
||||
// A NARROWER REQUEST RECORDS ITS OWN ROOT, even when a broader temporary
|
||||
// grant currently covers it.
|
||||
//
|
||||
// Borrowing a reference on the covering root kept this holder alive, which
|
||||
// was the point, but it kept the COVERING ROOT alive to do it — and
|
||||
// extraRoots is what validate() authorises writes from. So once the broad
|
||||
// holder released, a caller that had asked to write one subdirectory was
|
||||
// left holding write authority over the whole tree its neighbour had asked
|
||||
// for, including siblings it never named. Reported by CodeRabbit, and it is
|
||||
// the write-side twin of the read-side escalation @jatmn found: one
|
||||
// reference cannot be both a lifetime and a capability, whichever side of
|
||||
// the read/write boundary it sits on.
|
||||
//
|
||||
// Recording the requested root separately gives each holder exactly the
|
||||
// authority it asked for and a lifetime of its own. The nesting costs an
|
||||
// extra entry in extraRoots while both are live, which is redundant but not
|
||||
// wrong — the broader root already permits everything the narrower one does.
|
||||
s.extraRoots = append(s.extraRoots, root)
|
||||
s.tempWrites[root] = 1
|
||||
return root, oncePerHolder(func() { s.releaseTemporaryWrite(root) }), nil
|
||||
}
|
||||
|
||||
func (s *Scope) removeReadRoot(root string) {
|
||||
// releaseTemporaryRead drops ONE holder's reference and removes the root only
|
||||
// when the last one is gone.
|
||||
//
|
||||
// The count alone cannot make this safe, and an earlier comment here claimed it
|
||||
// could — that each undo is called exactly once, and flooring at zero handled
|
||||
// the rest. Flooring stops the count going negative; it does not stop one
|
||||
// holder's second call consuming a DIFFERENT holder's reference. With two
|
||||
// readers, calling the first's undo twice took the count 2 -> 1 -> 0 and removed
|
||||
// a root the second was still using. Idempotency belongs to the closure, which
|
||||
// is the thing that knows whose reference it is, so every undo handed out is
|
||||
// wrapped in oncePerHolder.
|
||||
func (s *Scope) releaseTemporaryRead(root string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
remaining, tracked := s.tempReads[root]
|
||||
if !tracked {
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
remaining--
|
||||
if remaining > 0 {
|
||||
s.tempReads[root] = remaining
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
// Both mutations under ONE hold. Dropping the lock between them opened a
|
||||
// window where the root was still in readRoots but no longer in tempReads:
|
||||
// a concurrent AddTemporaryRead landing there reads it as a PERMANENT root,
|
||||
// hands its caller a no-op undo, and then this call strips the root — so
|
||||
// that caller believes it holds access it has already silently lost.
|
||||
delete(s.tempReads, root)
|
||||
s.readRoots = removeScopeRoot(s.readRoots, root)
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func (s *Scope) removeWriteRoot(root string) {
|
||||
// releaseTemporaryWrite is releaseTemporaryRead for write roots. Two functions
|
||||
// rather than one generic helper because they guard different slices and the
|
||||
// generic version would take the slice by name — which is how the wrong one
|
||||
// gets passed.
|
||||
func (s *Scope) releaseTemporaryWrite(root string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
remaining, tracked := s.tempWrites[root]
|
||||
if !tracked {
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
remaining--
|
||||
if remaining > 0 {
|
||||
s.tempWrites[root] = remaining
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
// Same single-hold rule as releaseTemporaryRead above.
|
||||
delete(s.tempWrites, root)
|
||||
s.extraRoots = removeScopeRoot(s.extraRoots, root)
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// oncePerHolder makes one holder's undo safe to call more than once. The
|
||||
// reference it drops is that holder's own, so a second call must do nothing
|
||||
// rather than reach into the shared count and take somebody else's.
|
||||
func oncePerHolder(release func()) func() {
|
||||
var once sync.Once
|
||||
return func() { once.Do(release) }
|
||||
}
|
||||
|
||||
func removeScopeRoot(roots []string, root string) []string {
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
package sandbox
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// THE ANTI-ESCALATION PROPERTY. A read grant (AddRead, where --add-read-dir lands)
|
||||
// is READABLE but NOT WRITABLE. This is what makes propagating a parent's
|
||||
// request_permissions read grant to a child safe: the child can audit the granted
|
||||
// path but cannot modify it. Emitting the grant as a WRITE root (--add-dir) would
|
||||
// break exactly this.
|
||||
func TestAReadGrantIsReadableButNotWritable(t *testing.T) {
|
||||
scope, granted := grantOutsideDefaults(t)
|
||||
root, err := scope.AddRead(granted)
|
||||
if err != nil {
|
||||
t.Fatalf("AddRead: %v", err)
|
||||
}
|
||||
target := filepath.Join(root, "audit-me.go")
|
||||
|
||||
if block := scope.validateRead(target); block != nil {
|
||||
t.Fatalf("a read grant did not allow READING %q: %v — the audit would fail 'outside the workspace'", target, block)
|
||||
}
|
||||
if block := scope.validate(target); block == nil {
|
||||
t.Fatalf("a read grant ALLOWED WRITING %q — a read grant escalated to write", target)
|
||||
}
|
||||
}
|
||||
|
||||
// grantOutsideDefaults returns a scope and a directory that scope does not
|
||||
// already cover, both under test-owned storage.
|
||||
//
|
||||
// THE CAUSE, NOT THE SYMPTOM. NewScope seeds the system temporary directory as a
|
||||
// permanent write root, so a plain t.TempDir() grant is writable before any
|
||||
// grant exists and the anti-escalation assertions below prove nothing. The first
|
||||
// fix placed the fixture under os.UserHomeDir() to find ground the defaults do
|
||||
// not own. That worked and bought two problems: a hermetic runner may expose
|
||||
// HOME read-only, so these tests failed at MkdirTemp before asserting anything,
|
||||
// and on an ordinary machine they wrote into the developer's real home and left
|
||||
// zero-scope-* debris whenever a run was interrupted. Reported by @jatmn.
|
||||
//
|
||||
// Building the Scope directly removes the seeded defaults, which is what made
|
||||
// HOME necessary in the first place. t.TempDir() is then outside by
|
||||
// construction, the fixture is hermetic, and nothing touches HOME.
|
||||
func grantOutsideDefaults(t *testing.T) (*Scope, string) {
|
||||
t.Helper()
|
||||
workspace := resolvedFixturePath(t, t.TempDir())
|
||||
granted := resolvedFixturePath(t, t.TempDir())
|
||||
scope := &Scope{workspaceRoot: workspace}
|
||||
// PROVED AGAINST THIS SCOPE. The point of the helper is that the grant is
|
||||
// what changes access, so the target has to begin unauthorised — asserted
|
||||
// here once rather than trusted in every caller.
|
||||
if scope.validate(filepath.Join(granted, "probe.txt")) == nil {
|
||||
t.Fatalf("%s is writable before any grant exists; the assertions below would prove nothing", granted)
|
||||
}
|
||||
return scope, granted
|
||||
}
|
||||
@@ -0,0 +1,595 @@
|
||||
package sandbox
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// scopeOutsideDefaults builds a Scope directly instead of through NewScope.
|
||||
//
|
||||
// NewScope adds the system temp directory as a write root, so a target under
|
||||
// t.TempDir() is already covered and AddTemporaryRead returns before it ever
|
||||
// touches the refcount. A test built that way exercises none of this and passes
|
||||
// against the bug — which is how the first version of this test passed.
|
||||
func scopeOutsideDefaults(t *testing.T) (*Scope, string) {
|
||||
t.Helper()
|
||||
workspace := t.TempDir()
|
||||
target := filepath.Join(t.TempDir(), "shared")
|
||||
if err := os.MkdirAll(target, 0o700); err != nil {
|
||||
t.Fatalf("mkdir target: %v", err)
|
||||
}
|
||||
// Built directly, NOT via NewScope: real directories are needed because
|
||||
// normalizeScopeRoot requires them to exist, but NewScope's default temp
|
||||
// write root would then cover the target and short-circuit the path.
|
||||
return &Scope{workspaceRoot: workspace}, target
|
||||
}
|
||||
|
||||
// The refcount and the root list have to move together.
|
||||
//
|
||||
// Release used to delete the refcount, drop the lock, then strip the root in a
|
||||
// second acquisition. In that window the root was still in readRoots but no
|
||||
// longer in tempReads, so a concurrent AddTemporaryRead read it as a PERMANENT
|
||||
// root and handed its caller a no-op undo — then the release stripped it and
|
||||
// that caller silently lost access it believed it held.
|
||||
func TestReleasingATemporaryReadCannotStripALiveGrant(t *testing.T) {
|
||||
scope, target := scopeOutsideDefaults(t)
|
||||
|
||||
// First holder establishes the root.
|
||||
first, releaseFirst, err := scope.AddTemporaryRead(target)
|
||||
if err != nil {
|
||||
t.Fatalf("AddTemporaryRead: %v", err)
|
||||
}
|
||||
if len(scope.tempReads) == 0 {
|
||||
t.Fatal("precondition: the grant should be refcounted, not covered by a default root")
|
||||
}
|
||||
|
||||
// Second holder takes a reference on the same root.
|
||||
second, releaseSecond, err := scope.AddTemporaryRead(target)
|
||||
if err != nil {
|
||||
t.Fatalf("AddTemporaryRead (second): %v", err)
|
||||
}
|
||||
|
||||
// The first holder leaving must not revoke the second's access.
|
||||
releaseFirst()
|
||||
if block := scope.validateRead(second); block != nil {
|
||||
t.Fatalf("the second holder lost its grant when the first released: %v", block)
|
||||
}
|
||||
|
||||
// Only the last release retires the root.
|
||||
releaseSecond()
|
||||
if block := scope.validateRead(first); block == nil {
|
||||
t.Error("the root outlived its last holder")
|
||||
}
|
||||
}
|
||||
|
||||
// Under contention the same invariant has to hold: while a grant is live, its
|
||||
// root is readable.
|
||||
func TestTemporaryReadGrantsSurviveConcurrentReleases(t *testing.T) {
|
||||
scope, target := scopeOutsideDefaults(t)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
var mu sync.Mutex
|
||||
var failures []string
|
||||
|
||||
for i := 0; i < 24; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < 200; j++ {
|
||||
granted, undo, err := scope.AddTemporaryRead(target)
|
||||
if err != nil {
|
||||
mu.Lock()
|
||||
failures = append(failures, "AddTemporaryRead: "+err.Error())
|
||||
mu.Unlock()
|
||||
return
|
||||
}
|
||||
if block := scope.validateRead(granted); block != nil {
|
||||
mu.Lock()
|
||||
failures = append(failures, "root not readable while a grant was live")
|
||||
mu.Unlock()
|
||||
}
|
||||
undo()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if len(failures) > 0 {
|
||||
t.Fatalf("%d failures, first: %s", len(failures), failures[0])
|
||||
}
|
||||
}
|
||||
|
||||
// A read covered by a TEMPORARY WRITE root must take a reference too.
|
||||
//
|
||||
// permanentWriteRootCoversLocked scans extraRoots, and AddTemporaryWrite appends
|
||||
// temporary write roots there — so "covered by a write root" does not mean
|
||||
// "covered permanently". Without a reference the write holder's cleanup silently
|
||||
// revoked the reader's access: the same defect the refcount closes, reached
|
||||
// across the read/write boundary instead of within one side of it.
|
||||
func TestAReadCoveredByATemporaryWriteSurvivesItsRelease(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
outer := filepath.Join(t.TempDir(), "outer")
|
||||
inner := filepath.Join(outer, "inner")
|
||||
if err := os.MkdirAll(inner, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Built directly for the same reason scopeOutsideDefaults is: NewScope's
|
||||
// default temp write root would cover the fixture and short-circuit this.
|
||||
scope := &Scope{workspaceRoot: workspace}
|
||||
|
||||
_, releaseWrite, err := scope.AddTemporaryWrite(outer)
|
||||
if err != nil {
|
||||
t.Fatalf("temporary write: %v", err)
|
||||
}
|
||||
readRoot, releaseRead, err := scope.AddTemporaryRead(inner)
|
||||
if err != nil {
|
||||
t.Fatalf("temporary read: %v", err)
|
||||
}
|
||||
|
||||
covered := func() bool {
|
||||
for _, root := range scope.ReadRoots() {
|
||||
if pathWithinRoot(root, readRoot) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
if !covered() {
|
||||
t.Fatal("the reader did not hold access before any release")
|
||||
}
|
||||
|
||||
// The WRITE holder finishes first; the reader has not released.
|
||||
releaseWrite()
|
||||
if !covered() {
|
||||
t.Errorf("the write holder's cleanup revoked a live read grant on %s", readRoot)
|
||||
}
|
||||
|
||||
// And once the reader releases too, the root is genuinely gone.
|
||||
releaseRead()
|
||||
if covered() {
|
||||
t.Errorf("the root outlived its last holder: %v", scope.ReadRoots())
|
||||
}
|
||||
}
|
||||
|
||||
func hasExtraRoot(scope *Scope, root string) bool {
|
||||
for _, existing := range writeRootsBeyondWorkspace(scope) {
|
||||
if existing == root {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// A PERMANENT GRANT MUST OUTLIVE THE TEMPORARY ONE IT LANDED ON TOP OF.
|
||||
//
|
||||
// extraRoots holds temporary write roots alongside permanent ones, so Add asked
|
||||
// only "is this already covered" — and a session-scoped grant made while a
|
||||
// temporary holder happened to cover the path recorded nothing, then vanished
|
||||
// when that holder released. Outliving the request that prompted it is the
|
||||
// entire difference between Add and AddTemporaryWrite.
|
||||
func TestAPermanentGrantSurvivesTheTemporaryOneItCovered(t *testing.T) {
|
||||
t.Run("write, same root", func(t *testing.T) {
|
||||
scope, _, outside := scopeOutsideRoots(t)
|
||||
_, releaseTemp, err := scope.AddTemporaryWrite(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := scope.Add(outside); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
releaseTemp()
|
||||
if !hasExtraRoot(scope, outside) {
|
||||
t.Error("the temporary holder's release revoked a permanent write grant")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("read, same root", func(t *testing.T) {
|
||||
scope, _, outside := scopeOutsideRoots(t)
|
||||
_, releaseTemp, err := scope.AddTemporaryRead(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := scope.AddRead(outside); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
releaseTemp()
|
||||
if !hasReadRoot(scope, outside) {
|
||||
t.Error("the temporary holder's release revoked a permanent read grant")
|
||||
}
|
||||
})
|
||||
|
||||
// The nested shape: a BROAD temporary root covering a NARROW permanent one.
|
||||
// Promoting in place cannot help here — the narrow root has to be recorded
|
||||
// in its own right, or it goes when the broad one does.
|
||||
t.Run("broad temporary over narrow permanent", func(t *testing.T) {
|
||||
scope, _, outside := scopeOutsideRoots(t)
|
||||
inner := filepath.Join(outside, "inner")
|
||||
if err := os.MkdirAll(inner, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, releaseBroad, err := scope.AddTemporaryWrite(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := scope.Add(inner); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
releaseBroad()
|
||||
covered := false
|
||||
for _, existing := range writeRootsBeyondWorkspace(scope) {
|
||||
if pathWithinRoot(existing, inner) {
|
||||
covered = true
|
||||
}
|
||||
}
|
||||
if !covered {
|
||||
t.Error("a narrower permanent grant died with the broader temporary root it sat under")
|
||||
}
|
||||
})
|
||||
|
||||
// THE READ MIRROR OF BOTH SHAPES ABOVE. AddRead makes two claims that the
|
||||
// subtests so far only made for Add: the loop over extraRoots skips coverage
|
||||
// that is a temporary WRITE, and the loop over readRoots skips coverage that
|
||||
// is a temporary READ. Each claim is a separate branch, and neither was
|
||||
// exercised from the read side.
|
||||
t.Run("permanent read under a temporary write", func(t *testing.T) {
|
||||
scope, _, outside := scopeOutsideRoots(t)
|
||||
_, releaseWrite, err := scope.AddTemporaryWrite(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := scope.AddRead(outside); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
releaseWrite()
|
||||
if !hasReadRoot(scope, outside) {
|
||||
t.Error("a temporary write holder's release revoked a permanent read grant")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("broad temporary read over narrow permanent read", func(t *testing.T) {
|
||||
scope, _, outside := scopeOutsideRoots(t)
|
||||
inner := filepath.Join(outside, "inner")
|
||||
if err := os.MkdirAll(inner, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, releaseBroad, err := scope.AddTemporaryRead(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := scope.AddRead(inner); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
releaseBroad()
|
||||
if block := scope.validateRead(filepath.Join(inner, "audit-me.go")); block != nil {
|
||||
t.Errorf("a narrower permanent read died with the broader temporary read it sat under: %v", block)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ONE HOLDER'S UNDO DROPS ONE HOLDER'S REFERENCE, HOWEVER OFTEN IT IS CALLED.
|
||||
//
|
||||
// The count flooring at zero stops it going negative; it does not stop a
|
||||
// holder's second call consuming somebody else's reference. With two readers,
|
||||
// calling the first's undo twice took the count 2 -> 1 -> 0 and removed a root
|
||||
// the second was still using — the exact revocation this refcount exists to
|
||||
// prevent, reached through a duplicate call rather than a sibling's cleanup.
|
||||
// Deferred cleanups in a retry path are how a real caller does this by accident.
|
||||
func TestOneHoldersUndoIsIdempotent(t *testing.T) {
|
||||
scope, workspace, outside := scopeOutsideRoots(t)
|
||||
_, undoFirst, err := scope.AddTemporaryRead(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, undoSecond, err := scope.AddTemporaryRead(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
undoFirst()
|
||||
undoFirst()
|
||||
undoFirst()
|
||||
if !hasReadRoot(scope, outside) {
|
||||
t.Fatal("repeated calls to one holder's undo revoked a live holder's access")
|
||||
}
|
||||
undoSecond()
|
||||
if hasReadRoot(scope, outside) {
|
||||
t.Error("the root outlived its last holder")
|
||||
}
|
||||
|
||||
// The write side takes the same rule. Built directly for the same reason the
|
||||
// helper is: NewScope would seed the system temp directory as a permanent
|
||||
// write root, which already covers this fixture and makes the grant vacuous.
|
||||
writeScope := &Scope{workspaceRoot: workspace}
|
||||
_, undoWriteFirst, err := writeScope.AddTemporaryWrite(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, undoWriteSecond, err := writeScope.AddTemporaryWrite(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
undoWriteFirst()
|
||||
undoWriteFirst()
|
||||
if !hasExtraRoot(writeScope, outside) {
|
||||
t.Fatal("repeated calls to one write holder's undo revoked a live holder's access")
|
||||
}
|
||||
undoWriteSecond()
|
||||
if hasExtraRoot(writeScope, outside) {
|
||||
t.Error("the write root outlived its last holder")
|
||||
}
|
||||
}
|
||||
|
||||
// temporaryWriteFixture builds a nested pair of real directories and a scope
|
||||
// that covers neither: an OUTER directory a temporary write grant will take,
|
||||
// and an INNER one below it that a temporary read grant will ask for.
|
||||
//
|
||||
// Built directly rather than through NewScope, for the reason
|
||||
// scopeOutsideDefaults documents: NewScope seeds the system temp directory as a
|
||||
// PERMANENT write root, which would cover the whole fixture and send every
|
||||
// branch under test down the "genuinely permanent" path instead. Both paths are
|
||||
// returned symlink-resolved because normalizeScopeRoot resolves what it stores,
|
||||
// and on macOS the fixture's /var spelling is a symlink to /private/var.
|
||||
func temporaryWriteFixture(t *testing.T) (scope *Scope, workspace, outer, inner string) {
|
||||
t.Helper()
|
||||
outer = filepath.Join(t.TempDir(), "outer")
|
||||
inner = filepath.Join(outer, "inner")
|
||||
if err := os.MkdirAll(inner, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
workspace = resolvedFixturePath(t, t.TempDir())
|
||||
return &Scope{workspaceRoot: workspace}, workspace, resolvedFixturePath(t, outer), resolvedFixturePath(t, inner)
|
||||
}
|
||||
|
||||
func resolvedFixturePath(t *testing.T, path string) string {
|
||||
t.Helper()
|
||||
resolved, err := filepath.EvalSymlinks(path)
|
||||
if err != nil {
|
||||
t.Fatalf("resolve fixture %s: %v", path, err)
|
||||
}
|
||||
return resolved
|
||||
}
|
||||
|
||||
// A READER MUST NOT INHERIT ITS WRITER'S AUTHORITY.
|
||||
//
|
||||
// The reader's dependency on a covering TEMPORARY WRITE root was recorded as a
|
||||
// reference on that root — one reference standing for both the lifetime and the
|
||||
// capability. Release the WRITER first and the reference kept the path in
|
||||
// extraRoots, which is the list Roots() feeds validate(), so the surviving
|
||||
// read-only holder could WRITE anywhere below a write grant that had already
|
||||
// ended — including paths it never asked to read.
|
||||
//
|
||||
// This is TestAReadGrantIsReadableButNotWritable's property for temporary
|
||||
// grants: read authority stays read authority, whichever holder leaves first.
|
||||
func TestAReaderOutlivingItsWriterKeepsReadAuthorityOnly(t *testing.T) {
|
||||
t.Run("the writer releases first", func(t *testing.T) {
|
||||
scope, workspace, outer, inner := temporaryWriteFixture(t)
|
||||
|
||||
writeRoot, releaseWrite, err := scope.AddTemporaryWrite(outer)
|
||||
if err != nil {
|
||||
t.Fatalf("temporary write: %v", err)
|
||||
}
|
||||
if writeRoot != outer {
|
||||
t.Fatalf("temporary write granted %q, want %q", writeRoot, outer)
|
||||
}
|
||||
readRoot, releaseRead, err := scope.AddTemporaryRead(inner)
|
||||
if err != nil {
|
||||
t.Fatalf("temporary read: %v", err)
|
||||
}
|
||||
if readRoot != inner {
|
||||
t.Fatalf("temporary read granted %q, want %q", readRoot, inner)
|
||||
}
|
||||
// The reader holds a read root of its own and NO reference on the write
|
||||
// root. Both halves matter: the first is what survives the writer, the
|
||||
// second is what stops the write root surviving with it.
|
||||
//
|
||||
// Errorf, not Fatalf: these say WHY the behavior below breaks, and a
|
||||
// bookkeeping check that aborts the run hides the behavior it explains.
|
||||
scope.mu.RLock()
|
||||
writeHolders := scope.tempWrites[outer]
|
||||
readHolders := scope.tempReads[inner]
|
||||
scope.mu.RUnlock()
|
||||
if writeHolders != 1 {
|
||||
t.Errorf("tempWrites[%s] = %d, want 1: the reader took a reference on the WRITE root", outer, writeHolders)
|
||||
}
|
||||
if readHolders != 1 {
|
||||
t.Errorf("tempReads[%s] = %d, want 1: the reader kept no read root of its own", inner, readHolders)
|
||||
}
|
||||
|
||||
// The WRITE holder finishes; only the read-only holder is left.
|
||||
releaseWrite()
|
||||
|
||||
if got := writeRootsBeyondWorkspace(scope); len(got) != 0 {
|
||||
t.Errorf("ExtraRoots() = %v, want none: the reader held the write root open", got)
|
||||
}
|
||||
if got, want := scope.ReadRoots(), []string{workspace, inner}; !slices.Equal(got, want) {
|
||||
t.Errorf("ReadRoots() = %v, want %v", got, want)
|
||||
}
|
||||
|
||||
target := filepath.Join(inner, "audit-me.go")
|
||||
if block := scope.validateRead(target); block != nil {
|
||||
t.Errorf("the writer's cleanup revoked a live read grant on %q: %v", target, block)
|
||||
}
|
||||
block := scope.validate(target)
|
||||
if block == nil {
|
||||
t.Fatalf("a read-only holder may still WRITE %q after the write grant ended", target)
|
||||
}
|
||||
if block.Code != BlockOutsideWorkspace {
|
||||
t.Errorf("write block code = %q, want %q", block.Code, BlockOutsideWorkspace)
|
||||
}
|
||||
if block.Path != target {
|
||||
t.Errorf("write block path = %q, want %q", block.Path, target)
|
||||
}
|
||||
// The escalation was never confined to the path the reader asked for —
|
||||
// a borrowed reference carries the whole covering root — so a sibling
|
||||
// the reader never named has to be denied too.
|
||||
sibling := filepath.Join(outer, "sibling.txt")
|
||||
if block := scope.validate(sibling); block == nil {
|
||||
t.Errorf("a read-only holder of %q may still WRITE %q, which it never asked to read", inner, sibling)
|
||||
}
|
||||
|
||||
releaseRead()
|
||||
if got, want := scope.ReadRoots(), []string{workspace}; !slices.Equal(got, want) {
|
||||
t.Errorf("ReadRoots() = %v, want %v: the read root outlived its last holder", got, want)
|
||||
}
|
||||
})
|
||||
|
||||
// A CONTROL, NOT COVERAGE FOR THE FINDING. This subtest passes on the
|
||||
// unfixed tree too — the escalation needs the WRITER to go first, which the
|
||||
// subtest above exercises. It earns its place by pinning that the opposite
|
||||
// release order was not broken while fixing the dangerous one, and it is
|
||||
// labelled so nobody counts it as evidence the finding is closed.
|
||||
t.Run("the reader releases first (control: passes unfixed)", func(t *testing.T) {
|
||||
scope, _, outer, inner := temporaryWriteFixture(t)
|
||||
|
||||
_, releaseWrite, err := scope.AddTemporaryWrite(outer)
|
||||
if err != nil {
|
||||
t.Fatalf("temporary write: %v", err)
|
||||
}
|
||||
_, releaseRead, err := scope.AddTemporaryRead(inner)
|
||||
if err != nil {
|
||||
t.Fatalf("temporary read: %v", err)
|
||||
}
|
||||
|
||||
// The other order, which the fix must not break: the writer keeps FULL
|
||||
// authority for as long as it holds the grant.
|
||||
releaseRead()
|
||||
buildLog := filepath.Join(outer, "build.log")
|
||||
if block := scope.validate(buildLog); block != nil {
|
||||
t.Fatalf("the reader's cleanup revoked the live write grant on %q: %v", buildLog, block)
|
||||
}
|
||||
if got, want := writeRootsBeyondWorkspace(scope), []string{outer}; !slices.Equal(got, want) {
|
||||
t.Errorf("ExtraRoots() = %v, want %v", got, want)
|
||||
}
|
||||
|
||||
releaseWrite()
|
||||
if block := scope.validate(buildLog); block == nil {
|
||||
t.Errorf("the write root %q outlived its only holder", outer)
|
||||
}
|
||||
if block := scope.validateRead(filepath.Join(inner, "audit-me.go")); block == nil {
|
||||
t.Errorf("the read root %q outlived its only holder", inner)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// A NARROWER WRITE HOLDER MUST NOT INHERIT ITS NEIGHBOUR'S TREE.
|
||||
//
|
||||
// The write-side twin of the read-side escalation above, and the same root
|
||||
// cause: a nested holder used to take a reference on the BROADER temporary root
|
||||
// that happened to cover it. That kept the nested holder alive, which was the
|
||||
// point, but it kept the covering root alive to do it — and extraRoots is what
|
||||
// validate() authorises writes from. Once the broad holder released, a caller
|
||||
// that had asked for one subdirectory held write authority over the whole tree,
|
||||
// siblings included. Reported by CodeRabbit.
|
||||
func TestANestedWriteHolderKeepsOnlyItsOwnSubtree(t *testing.T) {
|
||||
scope, _, outsideBase := scopeOutsideRoots(t)
|
||||
outer := filepath.Join(outsideBase, "outer")
|
||||
inner := filepath.Join(outer, "inner")
|
||||
sibling := filepath.Join(outer, "sibling")
|
||||
for _, dir := range []string{outer, inner, sibling} {
|
||||
if err := os.MkdirAll(dir, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
_, releaseOuter, err := scope.AddTemporaryWrite(outer)
|
||||
if err != nil {
|
||||
t.Fatalf("temporary write on the broad root: %v", err)
|
||||
}
|
||||
_, releaseInner, err := scope.AddTemporaryWrite(inner)
|
||||
if err != nil {
|
||||
t.Fatalf("temporary write on the nested root: %v", err)
|
||||
}
|
||||
|
||||
// THE BROAD HOLDER GOES FIRST — the nested holder outliving it is what
|
||||
// exposes whose authority it is actually holding.
|
||||
releaseOuter()
|
||||
|
||||
if block := scope.validate(filepath.Join(inner, "mine.txt")); block != nil {
|
||||
t.Errorf("the nested holder lost the subtree it asked for: %v", block.Reason)
|
||||
}
|
||||
if block := scope.validate(filepath.Join(sibling, "not-mine.txt")); block == nil {
|
||||
t.Error("the nested holder can write a SIBLING it never asked for: the broad holder's authority outlived the broad holder")
|
||||
}
|
||||
if block := scope.validate(filepath.Join(outer, "not-mine.txt")); block == nil {
|
||||
t.Error("the nested holder can write the covering root itself after its holder released")
|
||||
}
|
||||
|
||||
releaseInner()
|
||||
if block := scope.validate(filepath.Join(inner, "mine.txt")); block == nil {
|
||||
t.Error("the nested root outlived its only holder")
|
||||
}
|
||||
}
|
||||
|
||||
// The opposite order, as a control: the nested holder releasing first must not
|
||||
// disturb the broad holder that is still live.
|
||||
func TestANestedWriteHolderReleasingFirstLeavesTheBroadGrant(t *testing.T) {
|
||||
scope, _, outsideBase := scopeOutsideRoots(t)
|
||||
outer := filepath.Join(outsideBase, "outer")
|
||||
inner := filepath.Join(outer, "inner")
|
||||
for _, dir := range []string{outer, inner} {
|
||||
if err := os.MkdirAll(dir, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
_, releaseOuter, err := scope.AddTemporaryWrite(outer)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, releaseInner, err := scope.AddTemporaryWrite(inner)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
releaseInner()
|
||||
if block := scope.validate(filepath.Join(inner, "still-ok.txt")); block != nil {
|
||||
t.Errorf("the nested holder's release revoked the broad holder's coverage: %v", block.Reason)
|
||||
}
|
||||
releaseOuter()
|
||||
if block := scope.validate(filepath.Join(inner, "gone.txt")); block == nil {
|
||||
t.Error("the tree outlived every holder")
|
||||
}
|
||||
}
|
||||
|
||||
// TWO HOLDERS OF THE SAME ROOT SHARE ONE ENTRY. Only an exact match shares —
|
||||
// a narrower request must not, which is what the escalation test above pins.
|
||||
func TestTwoHoldersOfTheSameWriteRootShareOneEntry(t *testing.T) {
|
||||
scope, _, outsideBase := scopeOutsideRoots(t)
|
||||
target := filepath.Join(outsideBase, "shared")
|
||||
if err := os.MkdirAll(target, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, releaseFirst, err := scope.AddTemporaryWrite(target)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, releaseSecond, err := scope.AddTemporaryWrite(target)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
releaseFirst()
|
||||
if block := scope.validate(filepath.Join(target, "still-held.txt")); block != nil {
|
||||
t.Errorf("one holder's release revoked the other's grant on the same root: %v", block.Reason)
|
||||
}
|
||||
releaseSecond()
|
||||
if block := scope.validate(filepath.Join(target, "gone.txt")); block == nil {
|
||||
t.Error("the shared root outlived both holders")
|
||||
}
|
||||
}
|
||||
|
||||
// writeRootsBeyondWorkspace is Roots() without the workspace entry.
|
||||
//
|
||||
// These tests need to see which BEYOND-WORKSPACE write roots a scope currently
|
||||
// holds. There was an ExtraRoots() accessor that answered exactly this, but it
|
||||
// had no production caller — it belongs to #829's child-grant propagation, whose
|
||||
// consumer side is not on this branch, and @jatmn was right that carrying it here
|
||||
// leaves unrelated production surface in a PR scoped to refcounting. The
|
||||
// inspection is a test concern, so it lives in the tests.
|
||||
func writeRootsBeyondWorkspace(scope *Scope) []string {
|
||||
all := scope.Roots()
|
||||
if len(all) == 0 {
|
||||
return nil
|
||||
}
|
||||
return append([]string(nil), all[1:]...)
|
||||
}
|
||||
@@ -0,0 +1,366 @@
|
||||
package sandbox
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// scopeOutsideRoots builds a workspace and an unrelated directory that no
|
||||
// default write root already covers.
|
||||
//
|
||||
// NOT t.TempDir(). /tmp and $TMPDIR are default write roots, so a temporary
|
||||
// grant for a path under one is a no-op — every assertion here would pass
|
||||
// against code that grants nothing at all. That is RULES §2.3 in its exact
|
||||
// form, and it caught the first version of this test.
|
||||
// scopeOutsideRoots builds a workspace and a directory outside it, both under
|
||||
// test-owned storage, and returns a Scope that treats the second as genuinely
|
||||
// outside.
|
||||
//
|
||||
// THE FIXTURE PROBLEM AND ITS ACTUAL CAUSE. NewScope seeds the system temporary
|
||||
// directory as a permanent write root, so an ordinary t.TempDir() grant is
|
||||
// covered before any temporary grant exists and every assertion built on it is
|
||||
// vacuous. The first fix for that reached into os.UserHomeDir() to find ground
|
||||
// the defaults do not already own — which worked, and bought two new problems:
|
||||
// a hermetic runner may expose HOME read-only, so the tests failed at MkdirTemp
|
||||
// before asserting anything; and on an ordinary machine they wrote fixtures into
|
||||
// the developer's real home and left zeromax-scope-* debris behind whenever a
|
||||
// run was interrupted. Reported by @jatmn.
|
||||
//
|
||||
// The cause is the seeded defaults, not the location, so the Scope is built
|
||||
// DIRECTLY here — the same pattern scopeOutsideDefaults and temporaryWriteFixture
|
||||
// already use. With no defaults seeded, t.TempDir() is outside by construction,
|
||||
// the fixture is hermetic, and nothing touches HOME.
|
||||
//
|
||||
// Paths are symlink-resolved because the scope stores resolved roots and macOS
|
||||
// hands out /var/... temp directories that resolve to /private/var/...; an
|
||||
// unresolved fixture path compares unequal to the root the scope just recorded.
|
||||
func scopeOutsideRoots(t *testing.T) (scope *Scope, workspace string, outside string) {
|
||||
t.Helper()
|
||||
workspace = resolvedFixturePath(t, t.TempDir())
|
||||
outside = filepath.Join(resolvedFixturePath(t, t.TempDir()), "outside")
|
||||
// FATAL, NOT SKIP. Every caller of this helper asserts an AUTHORIZATION
|
||||
// boundary, so a skipped fixture means that assertion never ran and the build
|
||||
// reported green anyway. t.TempDir() already failed the test if it could not
|
||||
// provide storage, so a subdirectory failing under it is a broken machine.
|
||||
if err := os.MkdirAll(outside, 0o700); err != nil {
|
||||
t.Fatalf("cannot prepare %s: %v", outside, err)
|
||||
}
|
||||
scope = &Scope{workspaceRoot: workspace}
|
||||
// PROVED AGAINST THIS SCOPE, not against production's default list. The point
|
||||
// of the helper is that a grant here changes something, and the only thing
|
||||
// that can answer that is the scope the test will actually use. Checking
|
||||
// defaultTempWriteRoots() instead asks about a scope nobody builds here — and
|
||||
// says "prove nothing" for a t.TempDir() path that this scope, seeded with no
|
||||
// defaults, genuinely does not cover.
|
||||
//
|
||||
// Asserting the NEGATIVE up front is also what keeps every caller honest: it
|
||||
// establishes that the target begins unauthorised, so a later success is the
|
||||
// grant's doing and not the fixture's placement.
|
||||
if scope.validate(filepath.Join(outside, "probe.txt")) == nil {
|
||||
t.Fatalf("%s is writable before any grant exists, so a temporary grant here would prove nothing", outside)
|
||||
}
|
||||
if scope.validateRead(filepath.Join(outside, "probe.txt")) == nil {
|
||||
t.Fatalf("%s is readable before any grant exists, so a temporary read grant here would prove nothing", outside)
|
||||
}
|
||||
return scope, workspace, outside
|
||||
}
|
||||
|
||||
func hasReadRoot(scope *Scope, root string) bool {
|
||||
for _, existing := range scope.ReadRoots() {
|
||||
if existing == root {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// THE DEFECT: one holder's cleanup revoked another's access. Two read-only
|
||||
// tools in the same parallel batch, both blocked on the same directory, is
|
||||
// exactly this — and read-only tools are the ones the batch runs concurrently.
|
||||
func TestATemporaryReadSurvivesASiblingsCleanup(t *testing.T) {
|
||||
scope, _, outside := scopeOutsideRoots(t)
|
||||
// The probe that the fix must make unnecessary: without it this reads
|
||||
// true -> false.
|
||||
if hasReadRoot(scope, outside) {
|
||||
t.Fatal("the outside root is already covered; this test proves nothing")
|
||||
}
|
||||
|
||||
_, undoA, err := scope.AddTemporaryRead(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, undoB, err := scope.AddTemporaryRead(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !hasReadRoot(scope, outside) {
|
||||
t.Fatal("the grant did not take effect")
|
||||
}
|
||||
|
||||
undoA()
|
||||
if !hasReadRoot(scope, outside) {
|
||||
t.Fatal("A's cleanup revoked the root while B still held a grant")
|
||||
}
|
||||
undoB()
|
||||
if hasReadRoot(scope, outside) {
|
||||
t.Fatal("the root outlived its last holder")
|
||||
}
|
||||
}
|
||||
|
||||
// A NARROWER REQUEST KEEPS ITS OWN SUBTREE AND NOTHING WIDER.
|
||||
//
|
||||
// THIS TEST USED TO ASSERT THE OPPOSITE, and the earlier version is worth
|
||||
// recording because it is how the escalation survived a review round. It was
|
||||
// written to prove a LIFETIME property — the broad holder's cleanup must not
|
||||
// revoke a live narrow holder — and it proved that by asserting the BROAD root
|
||||
// was still present afterwards. That is the defect stated as a requirement: the
|
||||
// narrow holder was keeping its neighbour's whole tree alive in order to keep
|
||||
// itself alive, so the surviving reader could read siblings it never asked for.
|
||||
//
|
||||
// Lifetime and capability are two properties and one reference cannot carry
|
||||
// both. The assertions below name them separately: the requested path stays
|
||||
// readable while its holder lives, AND a sibling under the released broad root
|
||||
// is denied as soon as the holder that authorised it exits. Reported by @jatmn.
|
||||
func TestANarrowerReaderKeepsOnlyItsOwnSubtree(t *testing.T) {
|
||||
scope, _, outside := scopeOutsideRoots(t)
|
||||
nested := filepath.Join(outside, "nested")
|
||||
sibling := filepath.Join(outside, "sibling")
|
||||
for _, dir := range []string{nested, sibling} {
|
||||
// FATAL, NOT SKIP. These directories are the fixture, not a precondition
|
||||
// the environment might reasonably withhold — skipping on them reports a
|
||||
// broken setup as a passing run.
|
||||
if err := os.MkdirAll(dir, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
_, undoBroad, err := scope.AddTemporaryRead(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, undoNarrow, err := scope.AddTemporaryRead(nested)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// THE BROAD HOLDER GOES FIRST. The narrow holder outliving it is what exposes
|
||||
// whose authority it is actually holding.
|
||||
undoBroad()
|
||||
if block := scope.validateRead(filepath.Join(nested, "mine.txt")); block != nil {
|
||||
t.Errorf("the broad holder's cleanup revoked the subtree the narrow holder asked for: %v", block.Reason)
|
||||
}
|
||||
if block := scope.validateRead(filepath.Join(sibling, "not-mine.txt")); block == nil {
|
||||
t.Error("the narrow reader can read a SIBLING it never asked for: the broad holder's authority outlived the broad holder")
|
||||
}
|
||||
if block := scope.validateRead(filepath.Join(outside, "not-mine.txt")); block == nil {
|
||||
t.Error("the narrow reader can read the covering root itself after its holder released")
|
||||
}
|
||||
|
||||
undoNarrow()
|
||||
if block := scope.validateRead(filepath.Join(nested, "mine.txt")); block == nil {
|
||||
t.Error("the nested root outlived its only holder")
|
||||
}
|
||||
}
|
||||
|
||||
// The opposite order, as a control: the narrow reader releasing first must not
|
||||
// disturb the broad reader that is still live.
|
||||
func TestANarrowerReaderReleasingFirstLeavesTheBroadGrant(t *testing.T) {
|
||||
scope, _, outside := scopeOutsideRoots(t)
|
||||
nested := filepath.Join(outside, "nested")
|
||||
if err := os.MkdirAll(nested, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, undoBroad, err := scope.AddTemporaryRead(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, undoNarrow, err := scope.AddTemporaryRead(nested)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
undoNarrow()
|
||||
if block := scope.validateRead(filepath.Join(nested, "still-ok.txt")); block != nil {
|
||||
t.Errorf("the narrow holder's release revoked the broad holder's coverage: %v", block.Reason)
|
||||
}
|
||||
undoBroad()
|
||||
if block := scope.validateRead(filepath.Join(nested, "gone.txt")); block == nil {
|
||||
t.Error("the tree outlived every holder")
|
||||
}
|
||||
}
|
||||
|
||||
// TWO READERS OF THE SAME ROOT SHARE ONE ENTRY. Only an exact match shares; a
|
||||
// narrower one never does, which is what the escalation test above pins.
|
||||
func TestTwoReadersOfTheSameRootShareOneEntry(t *testing.T) {
|
||||
scope, _, outside := scopeOutsideRoots(t)
|
||||
_, undoFirst, err := scope.AddTemporaryRead(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, undoSecond, err := scope.AddTemporaryRead(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
undoFirst()
|
||||
if block := scope.validateRead(filepath.Join(outside, "still-held.txt")); block != nil {
|
||||
t.Errorf("one reader's release revoked the other's grant on the same root: %v", block.Reason)
|
||||
}
|
||||
undoSecond()
|
||||
if block := scope.validateRead(filepath.Join(outside, "gone.txt")); block == nil {
|
||||
t.Error("the shared root outlived both holders")
|
||||
}
|
||||
}
|
||||
|
||||
// A NARROW PERMANENT READ IS NOT BORROWED FROM A BROAD TEMPORARY ONE. The old
|
||||
// scan walked readRoots in order and stopped at the first entry covering the
|
||||
// path, so a broad TEMPORARY root earlier in the slice shadowed a narrower
|
||||
// PERMANENT one later — and the caller took a reference it did not need on a
|
||||
// grant that was about to end.
|
||||
func TestAPermanentReadIsNotShadowedByABroadTemporaryOne(t *testing.T) {
|
||||
scope, _, outside := scopeOutsideRoots(t)
|
||||
nested := filepath.Join(outside, "nested")
|
||||
if err := os.MkdirAll(nested, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// The broad TEMPORARY grant is taken first, so it sits earlier in readRoots.
|
||||
_, undoBroad, err := scope.AddTemporaryRead(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := scope.AddRead(nested); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
undoBroad()
|
||||
if block := scope.validateRead(filepath.Join(nested, "mine.txt")); block != nil {
|
||||
t.Errorf("a PERMANENT read grant died with a broad temporary one that merely preceded it: %v", block.Reason)
|
||||
}
|
||||
}
|
||||
|
||||
// A PERMANENT root is not refcounted and must never be removed by a temporary
|
||||
// holder's cleanup — the undo for a request it already covers is genuinely
|
||||
// nothing.
|
||||
func TestATemporaryGrantNeverRevokesAPermanentRoot(t *testing.T) {
|
||||
scope, _, outside := scopeOutsideRoots(t)
|
||||
if _, err := scope.AddRead(outside); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, undo, err := scope.AddTemporaryRead(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
undo()
|
||||
if !hasReadRoot(scope, outside) {
|
||||
t.Fatal("a temporary holder's cleanup removed a permanent read root")
|
||||
}
|
||||
}
|
||||
|
||||
// The same rule for WRITE roots, or a write grant keeps the bug reads no longer
|
||||
// have.
|
||||
func TestATemporaryWriteSurvivesASiblingsCleanup(t *testing.T) {
|
||||
scope, _, outside := scopeOutsideRoots(t)
|
||||
covered := func() bool {
|
||||
for _, existing := range scope.Roots() {
|
||||
if existing == outside {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
_, undoA, err := scope.AddTemporaryWrite(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, undoB, err := scope.AddTemporaryWrite(outside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !covered() {
|
||||
t.Fatal("the write grant did not take effect")
|
||||
}
|
||||
undoA()
|
||||
if !covered() {
|
||||
t.Fatal("A's cleanup revoked the write root while B still held a grant")
|
||||
}
|
||||
undoB()
|
||||
if covered() {
|
||||
t.Fatal("the write root outlived its last holder")
|
||||
}
|
||||
}
|
||||
|
||||
// UNDER REAL CONCURRENCY, which is the shape the parallel tool batch produces:
|
||||
// eight holders taking and releasing the same root in overlapping windows. The
|
||||
// root must be present the whole time at least one holder has it, and gone once
|
||||
// the last releases.
|
||||
func TestConcurrentHoldersOfOneRoot(t *testing.T) {
|
||||
scope, _, outside := scopeOutsideRoots(t)
|
||||
|
||||
const holders = 8
|
||||
var wg sync.WaitGroup
|
||||
start := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
// TWO PER HOLDER, because each one can report at both hasReadRoot checks. At
|
||||
// one slot per holder the channel fills after eight reports, the ninth send
|
||||
// blocks forever, wg.Done() is never reached and wg.Wait() hangs — so the
|
||||
// test would DEADLOCK rather than fail in exactly the case it exists to
|
||||
// catch, since nothing drains failures until after the wait.
|
||||
failures := make(chan string, 2*holders)
|
||||
// Signalled by EVERY holder once its own AddTemporaryRead has returned,
|
||||
// failure included: the gate below waits for all of them, so a holder that
|
||||
// returned early without signalling would hang this test rather than fail it.
|
||||
acquired := make(chan struct{}, holders)
|
||||
for i := 0; i < holders; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
_, undo, err := scope.AddTemporaryRead(outside)
|
||||
acquired <- struct{}{}
|
||||
if err != nil {
|
||||
failures <- err.Error()
|
||||
return
|
||||
}
|
||||
// Every holder must SEE its own grant for as long as it holds it.
|
||||
if !hasReadRoot(scope, outside) {
|
||||
failures <- "a holder could not see the root it had just been granted"
|
||||
}
|
||||
<-release
|
||||
if !hasReadRoot(scope, outside) {
|
||||
failures <- "a holder lost the root while still holding it"
|
||||
}
|
||||
undo()
|
||||
}()
|
||||
}
|
||||
close(start)
|
||||
// EVERY holder, not the first one. Spinning on hasReadRoot only waited for
|
||||
// somebody to hold the root, and the spin is tight enough to win that race
|
||||
// against the goroutines still being scheduled: measured over 200 runs, 194
|
||||
// of them peaked at a single simultaneous holder and none ever reached
|
||||
// eight. A test named for concurrent holders was testing one holder at a
|
||||
// time, and would have passed against an implementation with no refcount at
|
||||
// all. Counting the acquisitions makes the overlap real instead of hoped for.
|
||||
for i := 0; i < holders; i++ {
|
||||
<-acquired
|
||||
}
|
||||
close(release)
|
||||
wg.Wait()
|
||||
close(failures)
|
||||
for failure := range failures {
|
||||
t.Fatal(failure)
|
||||
}
|
||||
if hasReadRoot(scope, outside) {
|
||||
t.Fatal("the root outlived every holder")
|
||||
}
|
||||
// AND THE BOOKKEEPING IS EMPTY, not merely invisible. hasReadRoot answers
|
||||
// what a caller can see; a refcount entry left behind at zero is not visible
|
||||
// that way and would still leak, one map entry per acquire/release cycle.
|
||||
scope.mu.Lock()
|
||||
leftover, stillCounted := scope.tempReads[outside]
|
||||
scope.mu.Unlock()
|
||||
if stillCounted {
|
||||
t.Fatalf("the refcount table kept an entry for %q after every holder released: %v", outside, leftover)
|
||||
}
|
||||
}
|
||||
+31
-12
@@ -76,12 +76,17 @@ type BuildArgsInput struct {
|
||||
}
|
||||
|
||||
type BuildResumeArgsInput struct {
|
||||
SessionID string
|
||||
Prompt string
|
||||
CurrentDepth int
|
||||
Manifest Manifest
|
||||
Cwd string
|
||||
PermissionMode string
|
||||
SessionID string
|
||||
Prompt string
|
||||
CurrentDepth int
|
||||
Manifest Manifest
|
||||
Cwd string
|
||||
// ParentModel and ParentReasoningEffort are the same fallbacks the fresh
|
||||
// path takes, and they are here for the same reason: appendModelArgs uses
|
||||
// them only when the manifest pins nothing of its own.
|
||||
ParentModel string
|
||||
ParentReasoningEffort string
|
||||
PermissionMode string
|
||||
}
|
||||
|
||||
type BuildArgsResult struct {
|
||||
@@ -344,6 +349,18 @@ func (executor Executor) BuildResumeArgs(input BuildResumeArgsInput) (BuildArgsR
|
||||
}
|
||||
args := []string{"exec", "--resume", sessionID}
|
||||
args = append(args, promptArgs...)
|
||||
// A RESUMED SPECIALIST KEEPS THE MODEL ITS MANIFEST PINNED.
|
||||
//
|
||||
// The fresh path appends this and the resume path did not, so a specialist
|
||||
// pinned to a cheap model ran on that model once and then silently reverted
|
||||
// to the parent's configured model on every resume. Nothing reports it: the
|
||||
// child starts normally and the only symptom is the bill.
|
||||
//
|
||||
// Resuming does not restore the recorded model on its own. sessions.PrepareExec
|
||||
// records the model it ran under but never feeds it back into provider
|
||||
// construction, so without this the child takes whatever the config default
|
||||
// resolves to now.
|
||||
args = appendModelArgs(args, input.Manifest, input.ParentModel, input.ParentReasoningEffort)
|
||||
args = append(args, "--auto", specialistAutonomy(input.PermissionMode), "--output-format", "stream-json")
|
||||
// See BuildArgs: only plan/spec-draft propagate --permission-mode so --auto
|
||||
// remains the authority for member/auto/ask/unsafe child resolution.
|
||||
@@ -424,12 +441,14 @@ func (executor Executor) runResume(ctx context.Context, params TaskParameters, o
|
||||
return ExecResult{}, err
|
||||
}
|
||||
built, err := executor.BuildResumeArgs(BuildResumeArgsInput{
|
||||
SessionID: params.Resume,
|
||||
Prompt: params.Prompt,
|
||||
CurrentDepth: options.CurrentDepth,
|
||||
Manifest: manifest,
|
||||
Cwd: options.Cwd,
|
||||
PermissionMode: options.PermissionMode,
|
||||
SessionID: params.Resume,
|
||||
Prompt: params.Prompt,
|
||||
CurrentDepth: options.CurrentDepth,
|
||||
Manifest: manifest,
|
||||
Cwd: options.Cwd,
|
||||
ParentModel: options.ParentModel,
|
||||
ParentReasoningEffort: options.ParentReasoningEffort,
|
||||
PermissionMode: options.PermissionMode,
|
||||
})
|
||||
if err != nil {
|
||||
return ExecResult{}, err
|
||||
|
||||
@@ -0,0 +1,303 @@
|
||||
package specialist
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/sessions"
|
||||
"github.com/Gitlawb/zero/internal/streamjson"
|
||||
)
|
||||
|
||||
// argValue returns the value following flag in argv, and whether it was present.
|
||||
func argValue(args []string, flag string) (string, bool) {
|
||||
for index, arg := range args {
|
||||
if arg == flag && index+1 < len(args) {
|
||||
return args[index+1], true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func pinnedManifest(model string, effort string) Manifest {
|
||||
return Manifest{
|
||||
Metadata: Metadata{
|
||||
Name: "skim",
|
||||
Description: "Bounded read-only lookups.",
|
||||
Model: model,
|
||||
ReasoningEffort: effort,
|
||||
Tools: []string{"read_file"},
|
||||
},
|
||||
SystemPrompt: "Find the thing and stop.",
|
||||
ResolvedTools: []string{"read_file"},
|
||||
}
|
||||
}
|
||||
|
||||
// A RESUMED SPECIALIST KEEPS THE MODEL ITS MANIFEST PINNED.
|
||||
//
|
||||
// Metadata.Model exists so a bounded, delegated task can run on a cheaper model
|
||||
// than its parent. BuildArgs appended it; BuildResumeArgs did not. So the
|
||||
// specialist ran on the cheap model once, and the moment the orchestrator
|
||||
// resumed it the child fell back to whatever the parent's configured model
|
||||
// resolved to. Nothing surfaced it: the resumed child starts normally, does the
|
||||
// work, and the only symptom is the bill.
|
||||
//
|
||||
// Resuming does not restore the recorded model on its own, which is why the flag
|
||||
// has to be passed again rather than relied upon.
|
||||
func TestAResumedSpecialistKeepsItsPinnedModel(t *testing.T) {
|
||||
manifest := pinnedManifest("claude-haiku-4.5", "")
|
||||
|
||||
fresh, err := (Executor{}).BuildArgs(BuildArgsInput{
|
||||
Manifest: manifest,
|
||||
Prompt: "find the thing",
|
||||
CurrentDepth: 0,
|
||||
ParentModel: "claude-opus-4.1",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildArgs: %v", err)
|
||||
}
|
||||
// SETUP: the fresh path really does pin it, or the comparison below is
|
||||
// asserting against a path that never worked either.
|
||||
freshModel, ok := argValue(fresh.Args, "--model")
|
||||
if !ok || freshModel != "claude-haiku-4.5" {
|
||||
t.Fatalf("SETUP INVALID: the fresh path passed --model %q (present=%t), want the manifest's model", freshModel, ok)
|
||||
}
|
||||
|
||||
resumed, err := (Executor{}).BuildResumeArgs(BuildResumeArgsInput{
|
||||
SessionID: "01HZZZZZZZZZZZZZZZZZZZZZZZ",
|
||||
Prompt: "keep going",
|
||||
CurrentDepth: 0,
|
||||
Manifest: manifest,
|
||||
ParentModel: "claude-opus-4.1",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildResumeArgs: %v", err)
|
||||
}
|
||||
resumedModel, ok := argValue(resumed.Args, "--model")
|
||||
if !ok {
|
||||
t.Fatalf("a resumed specialist carries no --model at all, so it reverts to the parent's configured model: %v", resumed.Args)
|
||||
}
|
||||
if resumedModel != "claude-haiku-4.5" {
|
||||
t.Fatalf("a resumed specialist runs on %q, want the manifest's pinned %q", resumedModel, "claude-haiku-4.5")
|
||||
}
|
||||
}
|
||||
|
||||
// And a manifest that pins nothing still inherits the parent's model, which is
|
||||
// what makes the fallback in appendModelArgs meaningful rather than the pinned
|
||||
// case being special-cased.
|
||||
func TestAResumedSpecialistWithNoPinInheritsTheParentModel(t *testing.T) {
|
||||
resumed, err := (Executor{}).BuildResumeArgs(BuildResumeArgsInput{
|
||||
SessionID: "01HZZZZZZZZZZZZZZZZZZZZZZZ",
|
||||
Prompt: "keep going",
|
||||
CurrentDepth: 0,
|
||||
Manifest: pinnedManifest("", ""),
|
||||
ParentModel: "claude-opus-4.1",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildResumeArgs: %v", err)
|
||||
}
|
||||
model, ok := argValue(resumed.Args, "--model")
|
||||
if !ok || model != "claude-opus-4.1" {
|
||||
t.Fatalf("a resumed specialist with no pinned model passed --model %q (present=%t), want the parent's", model, ok)
|
||||
}
|
||||
}
|
||||
|
||||
// The reasoning-effort rule travels with the model, or the two paths disagree
|
||||
// about what a pinned model implies. appendModelArgs inherits the parent's
|
||||
// effort ONLY when the manifest pins no model of its own: a manifest that chose
|
||||
// a different model has not agreed to the parent's effort for it.
|
||||
func TestAResumedSpecialistFollowsTheSameReasoningEffortRule(t *testing.T) {
|
||||
t.Run("pinned model does not inherit parent effort", func(t *testing.T) {
|
||||
resumed, err := (Executor{}).BuildResumeArgs(BuildResumeArgsInput{
|
||||
SessionID: "01HZZZZZZZZZZZZZZZZZZZZZZZ",
|
||||
Prompt: "keep going",
|
||||
CurrentDepth: 0,
|
||||
Manifest: pinnedManifest("claude-haiku-4.5", ""),
|
||||
ParentModel: "claude-opus-4.1",
|
||||
ParentReasoningEffort: "high",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildResumeArgs: %v", err)
|
||||
}
|
||||
if effort, ok := argValue(resumed.Args, "--reasoning-effort"); ok {
|
||||
t.Fatalf("a manifest that pinned its own model inherited the parent's effort %q", effort)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("no pinned model inherits parent effort", func(t *testing.T) {
|
||||
resumed, err := (Executor{}).BuildResumeArgs(BuildResumeArgsInput{
|
||||
SessionID: "01HZZZZZZZZZZZZZZZZZZZZZZZ",
|
||||
Prompt: "keep going",
|
||||
CurrentDepth: 0,
|
||||
Manifest: pinnedManifest("", ""),
|
||||
ParentModel: "claude-opus-4.1",
|
||||
ParentReasoningEffort: "high",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildResumeArgs: %v", err)
|
||||
}
|
||||
effort, ok := argValue(resumed.Args, "--reasoning-effort")
|
||||
if !ok || effort != "high" {
|
||||
t.Fatalf("a manifest pinning nothing passed --reasoning-effort %q (present=%t), want the parent's", effort, ok)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// The two builders agree on where the flag sits relative to the rest of the
|
||||
// argv, so a future change to one does not silently reorder only the other.
|
||||
func TestBothBuildersPlaceTheModelBeforeTheAutonomyFlag(t *testing.T) {
|
||||
manifest := pinnedManifest("claude-haiku-4.5", "")
|
||||
fresh, err := (Executor{}).BuildArgs(BuildArgsInput{Manifest: manifest, Prompt: "go", CurrentDepth: 0})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildArgs: %v", err)
|
||||
}
|
||||
resumed, err := (Executor{}).BuildResumeArgs(BuildResumeArgsInput{
|
||||
SessionID: "01HZZZZZZZZZZZZZZZZZZZZZZZ", Prompt: "go", CurrentDepth: 0, Manifest: manifest,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildResumeArgs: %v", err)
|
||||
}
|
||||
for _, argv := range [][]string{fresh.Args, resumed.Args} {
|
||||
model := indexOf(argv, "--model")
|
||||
auto := indexOf(argv, "--auto")
|
||||
if model < 0 || auto < 0 || model > auto {
|
||||
t.Fatalf("--model at %d, --auto at %d in %s", model, auto, strings.Join(argv, " "))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func indexOf(args []string, want string) int {
|
||||
for index, arg := range args {
|
||||
if arg == want {
|
||||
return index
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
// AND THE CALL SITE PASSES IT, WHICH THE BUILDER TESTS ABOVE CANNOT SEE.
|
||||
//
|
||||
// Every test above calls BuildResumeArgs directly. Dropping the ParentModel
|
||||
// field from the runResume call site still compiles and still passes all of
|
||||
// them, because the builder is doing its job with whatever it is handed. The
|
||||
// defect this fix repairs lived at the call site, not in the builder, so it
|
||||
// needs a test that goes through Run.
|
||||
//
|
||||
// Driven through the real Run dispatch with the RunChild seam capturing argv.
|
||||
func TestRunResumePassesTheParentModelThroughToTheChild(t *testing.T) {
|
||||
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
|
||||
parent, err := store.Create(sessions.CreateInput{SessionID: "parent_session"})
|
||||
if err != nil {
|
||||
t.Fatalf("create parent: %v", err)
|
||||
}
|
||||
if _, err := store.Create(sessions.CreateInput{
|
||||
SessionID: "child_task",
|
||||
SessionKind: sessions.SessionKindChild,
|
||||
Tag: SessionTagSpecialist,
|
||||
Depth: 1,
|
||||
ParentSessionID: parent.SessionID,
|
||||
AgentName: "skim",
|
||||
TaskID: "child_task",
|
||||
}); err != nil {
|
||||
t.Fatalf("create child: %v", err)
|
||||
}
|
||||
|
||||
zero := 0
|
||||
var captured []string
|
||||
executor := Executor{
|
||||
BinaryPath: "/usr/local/bin/zero",
|
||||
SessionStore: store,
|
||||
NewSessionID: func() (string, error) { return "child_task", nil },
|
||||
Load: func(LoadOptions) (LoadResult, error) {
|
||||
// No pinned model, so the parent's is the only thing that can supply one.
|
||||
return LoadResult{Specialists: []Manifest{pinnedManifest("", "")}}, nil
|
||||
},
|
||||
RunChild: func(_ context.Context, _ string, args []string, _ func(streamjson.Event)) (ChildRunResult, error) {
|
||||
captured = append([]string(nil), args...)
|
||||
return ChildRunResult{
|
||||
Events: []streamjson.Event{
|
||||
{Type: streamjson.EventRunStart, RunID: "run_1", SessionID: "child_task"},
|
||||
{Type: streamjson.EventFinal, RunID: "run_1", Text: "done"},
|
||||
{Type: streamjson.EventRunEnd, RunID: "run_1", Status: "success", ExitCode: &zero},
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
|
||||
if _, err := executor.Run(context.Background(), TaskParameters{
|
||||
Name: "skim",
|
||||
Prompt: "keep going",
|
||||
Resume: "child_task",
|
||||
}, TaskRunOptions{
|
||||
ParentSessionID: parent.SessionID,
|
||||
ParentModel: "claude-opus-4.1",
|
||||
}); err != nil {
|
||||
t.Fatalf("Run(resume): %v", err)
|
||||
}
|
||||
|
||||
// SETUP: this really was the resume path, not a fresh spawn.
|
||||
if index := indexOf(captured, "--resume"); index < 0 {
|
||||
t.Fatalf("SETUP INVALID: the child was not resumed: %v", captured)
|
||||
}
|
||||
model, ok := argValue(captured, "--model")
|
||||
if !ok {
|
||||
t.Fatalf("the resumed child was launched with no --model, so it runs on whatever the config default resolves to: %v", captured)
|
||||
}
|
||||
if model != "claude-opus-4.1" {
|
||||
t.Fatalf("the resumed child was launched with --model %q, want the parent's %q", model, "claude-opus-4.1")
|
||||
}
|
||||
}
|
||||
|
||||
// THE FRESH CALL SITE NEEDS THE SAME GUARD, FOR THE SAME REASON.
|
||||
//
|
||||
// runFresh and runResume both construct their builder input by hand, and both
|
||||
// have a byte-identical ParentModel line. Deleting either one compiles. The
|
||||
// resume half is what this change repairs; this covers the other half so the
|
||||
// pair cannot drift again in the direction nobody was looking.
|
||||
func TestRunFreshPassesTheParentModelThroughToTheChild(t *testing.T) {
|
||||
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
|
||||
parent, err := store.Create(sessions.CreateInput{SessionID: "parent_session"})
|
||||
if err != nil {
|
||||
t.Fatalf("create parent: %v", err)
|
||||
}
|
||||
|
||||
zero := 0
|
||||
var captured []string
|
||||
executor := Executor{
|
||||
BinaryPath: "/usr/local/bin/zero",
|
||||
SessionStore: store,
|
||||
NewSessionID: func() (string, error) { return "child_task", nil },
|
||||
Load: func(LoadOptions) (LoadResult, error) {
|
||||
return LoadResult{Specialists: []Manifest{pinnedManifest("", "")}}, nil
|
||||
},
|
||||
RunChild: func(_ context.Context, _ string, args []string, _ func(streamjson.Event)) (ChildRunResult, error) {
|
||||
captured = append([]string(nil), args...)
|
||||
return ChildRunResult{
|
||||
Events: []streamjson.Event{
|
||||
{Type: streamjson.EventRunStart, RunID: "run_1", SessionID: "child_task"},
|
||||
{Type: streamjson.EventFinal, RunID: "run_1", Text: "done"},
|
||||
{Type: streamjson.EventRunEnd, RunID: "run_1", Status: "success", ExitCode: &zero},
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
|
||||
if _, err := executor.Run(context.Background(), TaskParameters{
|
||||
Name: "skim",
|
||||
Prompt: "find the thing",
|
||||
}, TaskRunOptions{
|
||||
ParentSessionID: parent.SessionID,
|
||||
ParentModel: "claude-opus-4.1",
|
||||
}); err != nil {
|
||||
t.Fatalf("Run(fresh): %v", err)
|
||||
}
|
||||
|
||||
// SETUP: a fresh spawn, not a resume, or this covers the wrong site.
|
||||
if indexOf(captured, "--resume") >= 0 {
|
||||
t.Fatalf("SETUP INVALID: the child was resumed rather than freshly spawned: %v", captured)
|
||||
}
|
||||
model, ok := argValue(captured, "--model")
|
||||
if !ok || model != "claude-opus-4.1" {
|
||||
t.Fatalf("the fresh child was launched with --model %q (present=%t), want the parent's", model, ok)
|
||||
}
|
||||
}
|
||||
@@ -94,6 +94,8 @@ 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"`
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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
@@ -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))
|
||||
}
|
||||
|
||||
Loaded 100 of 139 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user