Compare commits

...
Author SHA1 Message Date
Vasanthdev2004 9371890bf4 fix(tools): put the written bytes back when the formatter fails
These formatters edit in place, so one killed by the deadline, killed by
the caller, or exiting partway through its own rewrite can leave the target
truncated: neither the input nor the output. Returning the written bytes on
top of that left the tracker baseline and the diff preview describing a
file that is not on disk, so the next edit would compare against content the
file does not have. That is a worse failure than the missing formatting the
previous commit set out to disclose.

The failure paths write the bytes back, and when even that fails the user
hears about it, since it is their file. Written back unconditionally on that
path rather than only when the bytes differ: comparing first means reading
the file to find out, and a read that fails leaves the same ambiguity this
exists to remove.

The fixture is proven to damage the file on a throwaway copy, because a
successful restore erases the evidence from the real one. My first attempt
asserted that on the target itself, where it could only pass vacuously or
report the opposite of what it saw.
2026-09-10 21:26:22 +05:30
Vasanthdev2004 94695aaad7 fix(tools): say when format-on-write ran out of time
The 10s bound exists so a wedged formatter cannot hang a tool call, which
is right, and on expiry the unformatted write stood with nothing said. Every
other fallback here is a standing fact about the environment (toggle off, no
formatter for the extension, binary not installed) and is silent because
nothing is wrong. A deadline firing is not that: formatting was configured,
available and expected, and the file was written unformatted because the
machine was slow. Left silent, the caller believes it wrote canonical style
and finds out from a CI format check it cannot see, which is the failure
this feature exists to prevent.

The timeout is now reported on the write and edit summaries, naming the file
and the formatter. The caller cancelling and the formatter exiting non-zero
are deliberately not reported: the first is already surfaced as a cancelled
run, and the second usually means content the formatter could not parse,
which a write does not promise to fix.

The bound stays 10s rather than being raised. Raising it makes a wedged
formatter block longer, which is the thing the bound is for; what was
missing was knowing when it fired.

Closes #1022
2026-09-09 22:49:08 +05:30
f30f550e60 fix(sandbox): protect worktree git pointers (#805)
* fix(sandbox): resolve git metadata carveouts for worktrees

.git is a file (gitdir: ...) in worktrees/submodules, not a dir. The
hooks/config write-deny carveout assumed a directory and pointed at
<root>/.git/hooks, which has no mkdir-able parent under the .git file.
bwrap then failed 'Can't mkdir parents ... Not a directory', blocking
every sandboxed tool in worktree checkouts.

Resolve the real (common) git dir via the gitdir pointer + commondir so
worktrees deny the shared hooks/config like a plain clone. Plain
checkouts (.git dir) and .git-absent roots keep the literal paths.

* fix(sandbox): bound worktree git metadata carveouts to the write root

Refuse gitdir/commondir pointers that escape the workspace instead of
carving outside confinement. Invalid, stale, or non-directory pointers
return no carveouts so bwrap does not mkdir under a .git file.

* fix(sandbox): evaluate gitdir/commondir through symlinks

Containment must use the real path so an in-root symlink cannot punch
hooks/config carveouts outside the write root.

* fix(sandbox): fail closed on gitdir eval and protect the pointer file

Refuse carveouts when EvalSymlinks cannot produce a real in-root path,
compare evaluated roots and targets, and write-deny the .git pointer so
it cannot redirect hooks/config between commands.

* fix(sandbox): deny worktree-local gitdir metadata and parse first gitdir line

Keep hooks/config write-denied on both the common dir and the in-root
worktree gitdir, emit those paths under the lexical write root, and
accept a gitdir: pointer that has trailing comment lines.

* fix(sandbox): skip special git metadata files and re-check carveouts at apply

Read .git and commondir only when they are regular files so a FIFO cannot
block profile construction. At sandbox apply, drop ReadOnlySubpaths that
no longer sit inside the write root after a symlink swap.

* fix(sandbox): follow directory .git symlinks and unix-tag FIFO tests

A .git symlink to a directory is a normal checkout, so keep lexical
hooks/config carveouts. FIFO coverage stays on unix so Windows still
compiles.

* fix(sandbox): protect worktree git pointer directly

* test(sandbox): skip unavailable symlink fixture

* fix(sandbox): preserve git symlink semantics

* fix(sandbox): limit git pointer carveout to regular files

---------

Co-authored-by: jatmn <[email protected]>
Co-authored-by: jatmn <[email protected]>
2026-09-08 07:54:22 +08:00
Vasanth T aadb4a27e9 fix(specialist): keep the pinned model when a specialist is resumed (#1009)
* fix(specialist): keep the pinned model when a specialist is resumed

Metadata.Model exists so a bounded, delegated task can run on a cheaper model
than its parent, and BuildArgs appends it to the child argv through
appendModelArgs. BuildResumeArgs never did.

So a specialist pinned to a cheap model ran on that model exactly once. 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 and does the work, so the only symptom is the bill.
Cost-motivated delegation quietly stopped saving anything.

Resuming does not restore the recorded model on its own. sessions.PrepareExec
records the model a run used but never feeds it back into provider
construction, so the flag has to be passed again rather than relied upon.

BuildResumeArgsInput now carries ParentModel and ParentReasoningEffort, the
same fallbacks the fresh path takes, and runResume passes what TaskRunOptions
already held. The reasoning-effort rule travels with the model unchanged: the
parent's effort is inherited only when the manifest pins no model of its own,
because a manifest that chose a different model has not agreed to the parent's
effort for it.

Regressions drive both builders and compare them, so the two paths cannot
drift again: a pinned model survives resume, an unpinned one still inherits the
parent's, both halves of the effort rule hold, and the flag keeps its position
relative to --auto in both.

Refs #554

* test(specialist): cover the resume call site, not just its builder

The builder tests all call BuildResumeArgs directly, so dropping the
ParentModel field from the runResume call site still compiled and still passed
every one of them. The defect this fixes lived at the call site, so it needs a
test that goes through Run.

Driven through the real dispatch with the RunChild seam capturing argv.

* test(specialist): guard the fresh call site as well

runFresh and runResume each construct their builder input by hand and carry a
byte-identical ParentModel line. Deleting either compiles and, until now,
deleting the fresh one was silent.

The resume half is what this branch repairs. This covers the other half so the
pair cannot drift again in the direction nobody was looking.
2026-09-07 16:02:39 +05:30
KRATOS a6ec0435e3 feat(acp): surface safe browser tool metadata (#1000)
* feat(acp): surface safe browser tool metadata

* fix(acp): align browser permission titles

* fix(acp): namespace browser metadata

* fix(acp): reject unsafe browser title text
2026-09-07 15:50:49 +05:30
Anandan 1b5db17656 fix(mcp): discover OAuth through protected resource metadata (#991)
* fix(mcp): discover OAuth through protected resource metadata

* fix(mcp): refuse OAuth discovery redirects

* fix(mcp): harden OAuth metadata discovery

* fix(mcp): protect discovered OAuth endpoints
2026-08-28 21:17:08 +05:30
Anandan eeea3308fe fix(tui): move voice capture to Ctrl+Space (#968)
* fix(tui): keep voice mode from swallowing spaces

* fix(tui): clarify voice capture controls

* fix(tui): align active voice status with capture key
2026-08-28 10:27:33 +05:30
29 changed files with 3298 additions and 106 deletions

No files matched your search

+3 -1
View File
@@ -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 {
+14
View File
@@ -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
View File
@@ -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
}
+226
View File
@@ -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 {
+15
View File
@@ -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
+100 -14
View File
@@ -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,
+290
View File
@@ -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)
}
})
}
}
+233
View File
@@ -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())
}
}
+220
View File
@@ -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())
}
}
+20 -7
View File
@@ -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 {
+115
View File
@@ -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"),
})
}
+31 -12
View File
@@ -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
+303
View File
@@ -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)
}
}
+3 -1
View File
@@ -160,7 +160,8 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
// Optional format-on-write (ZERO_FORMAT_ON_WRITE). Must run BEFORE the
// FileTracker re-baseline: recording pre-format content would make the very
// next edit look like an external modification and trip the conflict guard.
updated = maybeFormatWrittenFile(ctx, absolutePath, updated)
formatting := maybeFormatWrittenFile(ctx, absolutePath, updated)
updated = formatting.Content
// Re-baseline to the content we just wrote so subsequent edits in this session
// compare against the current on-disk state, not the pre-edit version.
newInfo, _ := os.Stat(absolutePath)
@@ -193,6 +194,7 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
suffix = "s"
}
summary := fmt.Sprintf("Successfully edited %s (replaced %d occurrence%s).", relativePath, replacedCount, suffix)
summary += formatting.notice(relativePath)
summary += inlineDiagnostics(ctx, options, absolutePath, relativePath)
result := okResult(summary)
result.ChangedFiles = []string{relativePath}
+88 -10
View File
@@ -2,6 +2,7 @@ package tools
import (
"context"
"errors"
"os"
"os/exec"
"path/filepath"
@@ -22,8 +23,59 @@ import (
// the conflict guard.
// formatOnWriteTimeout bounds one formatter run; a wedged formatter must never
// hang a tool call. On timeout the unformatted write stands.
const formatOnWriteTimeout = 10 * time.Second
// hang a tool call. On timeout the unformatted write stands, and the caller
// says so: see formatOnWriteResult.
//
// A var rather than a const so a test can shorten it. Nothing outside a test
// assigns to it, and the deadline path was previously unreachable in a test at
// any speed, which is part of why it went unnoticed that it reported nothing.
var formatOnWriteTimeout = 10 * time.Second
// formatOnWriteResult is the content after formatting, plus whether the
// formatting that was supposed to happen actually did.
//
// A TIMEOUT IS NOT THE SAME KIND OF MISS AS THE OTHERS. Every other way this
// falls back is a standing fact about the environment: the toggle is off, the
// extension has no formatter, the binary is not installed. Those are silent on
// purpose, because nothing is wrong and saying so on every write would be
// noise. A deadline firing is different: formatting was configured, available
// and expected, and the file was written unformatted anyway, on a machine that
// was merely slow. Left silent, the caller believes it wrote canonical style
// and finds out from a CI format check it cannot see, which is the thing this
// feature exists to prevent.
type formatOnWriteResult struct {
Content string
// Formatter is the binary that was run, named in the notice so the user can
// tell a slow gofmt from a slow prettier.
Formatter string
TimedOut bool
// RestoreFailed means the file on disk is not known to hold Content.
//
// These formatters edit in place, so one that is killed or fails partway
// can leave the target truncated or half-rewritten: what a dead
// `prettier --write` leaves behind is not the input and not the output.
// Returning the written bytes while disk holds something else would put the
// tracker baseline, the diff preview and the file itself into three
// different states, so the failure paths write the bytes back. When even
// that fails the user has to hear about it: it is their file.
RestoreFailed bool
}
// notice is the line appended to the tool summary when formatting was expected
// and did not happen, and empty in every other case.
func (result formatOnWriteResult) notice(relativePath string) string {
if result.RestoreFailed {
return "\n\nWARNING: " + relativePath + " may not hold what was written. " +
result.Formatter + " was interrupted while rewriting it in place and the " +
"content could not be written back. Re-read the file before trusting it."
}
if !result.TimedOut {
return ""
}
return "\n\nNote: " + relativePath + " was written but not formatted: " +
result.Formatter + " did not finish within " + formatOnWriteTimeout.String() +
". The file holds exactly what was written, so a project format check may still flag it."
}
// formatterCommands maps a file extension to the formatter argv; the file path
// is appended as the final argument. Only in-place, config-respecting,
@@ -69,18 +121,20 @@ func formatOnWriteEnabled() bool {
// enabled and on PATH) and returns the file's content afterwards. Best-effort
// throughout: any failure — no formatter, formatter error, timeout, unreadable
// result — returns writtenContent so the caller's state matches the last write
// it performed itself.
func maybeFormatWrittenFile(ctx context.Context, absolutePath string, writtenContent string) string {
// it performed itself. Only the timeout is reported back, for the reason on
// formatOnWriteResult.
func maybeFormatWrittenFile(ctx context.Context, absolutePath string, writtenContent string) formatOnWriteResult {
unformatted := formatOnWriteResult{Content: writtenContent}
if !formatOnWriteEnabled() {
return writtenContent
return unformatted
}
command, ok := formatterCommands[strings.ToLower(filepath.Ext(absolutePath))]
if !ok {
return writtenContent
return unformatted
}
binaryPath, err := exec.LookPath(command[0])
if err != nil {
return writtenContent
return unformatted
}
formatCtx, cancel := context.WithTimeout(ctx, formatOnWriteTimeout)
defer cancel()
@@ -89,11 +143,35 @@ func maybeFormatWrittenFile(ctx context.Context, absolutePath string, writtenCon
formatter.Dir = filepath.Dir(absolutePath)
formatter.Stdin = strings.NewReader("")
if err := formatter.Run(); err != nil {
return writtenContent
unformatted.Formatter = command[0]
// THE FORMATTER EDITS IN PLACE, SO A FAILED RUN CAN LEAVE THE FILE
// NEITHER FORMATTED NOR AS WRITTEN. Killed by the deadline or by the
// caller, or exiting partway through its own rewrite, the target can hold
// a truncation. Returning the written bytes on top of that would leave the
// tracker baseline and the diff preview describing a file that is not on
// disk, which is a worse failure than the missing formatting: the next
// edit compares against content the file does not have.
//
// Written back unconditionally on this path rather than only when the
// bytes differ. Comparing first means reading the file to find out, and a
// read that fails leaves the same ambiguity this exists to remove.
if restoreErr := os.WriteFile(absolutePath, []byte(writtenContent), 0o644); restoreErr != nil {
unformatted.RestoreFailed = true
}
// OUR deadline, not the caller's cancellation and not the formatter's own
// exit status. A cancelled tool call is already being reported as
// cancelled, and a formatter that ran and refused the file usually means
// content it could not parse, which the write itself does not promise to
// fix. Neither is this notice's business; the restore above is, for all
// three.
if errors.Is(formatCtx.Err(), context.DeadlineExceeded) && ctx.Err() == nil {
unformatted.TimedOut = true
}
return unformatted
}
formatted, err := os.ReadFile(absolutePath)
if err != nil {
return writtenContent
return unformatted
}
return string(formatted)
return formatOnWriteResult{Content: string(formatted), Formatter: command[0]}
}
+12 -6
View File
@@ -97,9 +97,12 @@ func TestFormatOnWriteFormatsAndKeepsTrackerConsistent(t *testing.T) {
func TestFormatOnWriteSkipsUnknownExtensions(t *testing.T) {
t.Setenv("ZERO_FORMAT_ON_WRITE", "1")
content := maybeFormatWrittenFile(context.Background(), filepath.Join(t.TempDir(), "notes.xyz"), "raw text")
if content != "raw text" {
t.Fatalf("unknown extension must pass through: %q", content)
formatting := maybeFormatWrittenFile(context.Background(), filepath.Join(t.TempDir(), "notes.xyz"), "raw text")
if formatting.Content != "raw text" {
t.Fatalf("unknown extension must pass through: %q", formatting.Content)
}
if notice := formatting.notice("notes.xyz"); notice != "" {
t.Fatalf("an extension with no formatter is not a miss worth reporting, got %q", notice)
}
}
@@ -111,8 +114,11 @@ func TestFormatOnWriteFormatterLookupFailure(t *testing.T) {
if err := os.WriteFile(targetPath, []byte(uglyContent), 0o644); err != nil {
t.Fatal(err)
}
content := maybeFormatWrittenFile(context.Background(), targetPath, uglyContent)
if content != uglyContent {
t.Fatalf("missing formatter must return written content, got %q", content)
formatting := maybeFormatWrittenFile(context.Background(), targetPath, uglyContent)
if formatting.Content != uglyContent {
t.Fatalf("missing formatter must return written content, got %q", formatting.Content)
}
if notice := formatting.notice("a.go"); notice != "" {
t.Fatalf("an uninstalled formatter is a standing fact, not a miss worth reporting, got %q", notice)
}
}
@@ -0,0 +1,280 @@
package tools
import (
"context"
"os"
"os/exec"
"path/filepath"
"runtime"
"strings"
"testing"
"time"
)
// A FORMATTER THAT RUNS OUT OF TIME MUST NOT LOOK LIKE ONE THAT SUCCEEDED.
//
// Every other way format-on-write falls back is a standing fact about the
// environment: the toggle is off, the extension has no formatter, the binary is
// not installed. Those are silent because nothing is wrong. A deadline firing is
// different: formatting was configured, available and expected, and the file was
// written unformatted anyway because the machine was slow. Silent, the caller
// believes it wrote canonical style and learns otherwise from a CI format check
// it cannot see, which is the failure this feature exists to prevent.
//
// THE DEADLINE IS SHRUNK RATHER THAN THE FORMATTER SLOWED. A formatter that
// really sleeps has to be killed, and on Windows the batch file that hosts the
// sleep leaves the sleeping grandchild holding the working directory, so the
// test then fails in TempDir cleanup rather than on anything it asserts. An
// expired budget against an ordinary fast formatter reaches the same branch by
// the same route, deterministically and in microseconds.
func TestFormatOnWriteReportsATimeout(t *testing.T) {
formatter := installFakeFormatter(t, ".fastfmt", "fastfmt", succeedingFormatterScript())
shortenFormatOnWriteTimeout(t, time.Nanosecond)
t.Setenv("ZERO_FORMAT_ON_WRITE", "1")
target := filepath.Join(t.TempDir(), "subject.fastfmt")
const written = "written but not formatted\n"
if err := os.WriteFile(target, []byte(written), 0o644); err != nil {
t.Fatal(err)
}
formatting := maybeFormatWrittenFile(context.Background(), target, written)
if formatting.Content != written {
t.Fatalf("content = %q, want the bytes that were written", formatting.Content)
}
if !formatting.TimedOut {
t.Fatal("a formatter cut off by the deadline was reported as an ordinary miss, so the write claims a formatting that never happened")
}
notice := formatting.notice("subject.fastfmt")
for _, want := range []string{"subject.fastfmt", "not formatted", formatter} {
if !strings.Contains(notice, want) {
t.Errorf("notice %q does not mention %q", notice, want)
}
}
}
// And the same formatter inside its budget says nothing at all, or the notice
// above would be reporting the deadline rather than the miss.
func TestFormatOnWriteStaysQuietWhenTheFormatterFinishes(t *testing.T) {
installFakeFormatter(t, ".fastfmt", "fastfmt", succeedingFormatterScript())
t.Setenv("ZERO_FORMAT_ON_WRITE", "1")
target := filepath.Join(t.TempDir(), "subject.fastfmt")
const written = "written\n"
if err := os.WriteFile(target, []byte(written), 0o644); err != nil {
t.Fatal(err)
}
formatting := maybeFormatWrittenFile(context.Background(), target, written)
if formatting.TimedOut {
t.Error("a formatter that finished was reported as timed out")
}
if notice := formatting.notice("subject.fastfmt"); notice != "" {
t.Errorf("a successful format produced a notice: %q", notice)
}
}
// The caller cancelling is not this notice's business: that run is already
// being reported as cancelled, and saying the formatter was too slow on top of
// it would be wrong about why.
func TestFormatOnWriteStaysQuietWhenTheCallerCancels(t *testing.T) {
installFakeFormatter(t, ".fastfmt", "fastfmt", succeedingFormatterScript())
t.Setenv("ZERO_FORMAT_ON_WRITE", "1")
target := filepath.Join(t.TempDir(), "subject.fastfmt")
const written = "written\n"
if err := os.WriteFile(target, []byte(written), 0o644); err != nil {
t.Fatal(err)
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
formatting := maybeFormatWrittenFile(ctx, target, written)
if formatting.Content != written {
t.Fatalf("content = %q, want the bytes that were written", formatting.Content)
}
if formatting.TimedOut {
t.Error("a cancelled run was reported as a formatter timeout")
}
if notice := formatting.notice("subject.fastfmt"); notice != "" {
t.Errorf("a cancelled run produced a notice: %q", notice)
}
}
// A formatter that runs and refuses the file is also not a timeout. It usually
// means content it could not parse, which the write does not promise to fix,
// and reporting it as a slow machine would be wrong about the cause.
func TestFormatOnWriteStaysQuietWhenTheFormatterFails(t *testing.T) {
installFakeFormatter(t, ".failfmt", "failfmt", failingFormatterScript())
t.Setenv("ZERO_FORMAT_ON_WRITE", "1")
target := filepath.Join(t.TempDir(), "subject.failfmt")
const written = "unparseable\n"
if err := os.WriteFile(target, []byte(written), 0o644); err != nil {
t.Fatal(err)
}
formatting := maybeFormatWrittenFile(context.Background(), target, written)
if formatting.TimedOut {
t.Error("a formatter that exited non-zero was reported as a timeout")
}
if notice := formatting.notice("subject.failfmt"); notice != "" {
t.Errorf("a failing formatter produced a timeout notice: %q", notice)
}
}
// installFakeFormatter puts a formatter on PATH and registers it in the command
// table for the test's duration, returning the binary name the notice carries.
func installFakeFormatter(t *testing.T, extension, name, script string) string {
t.Helper()
binaryName := name + formatterScriptExtension()
directory := t.TempDir()
if err := os.WriteFile(filepath.Join(directory, binaryName), []byte(script), 0o755); err != nil {
t.Fatal(err)
}
t.Setenv("PATH", directory+string(os.PathListSeparator)+os.Getenv("PATH"))
previous, existed := formatterCommands[extension]
formatterCommands[extension] = []string{binaryName}
t.Cleanup(func() {
if existed {
formatterCommands[extension] = previous
return
}
delete(formatterCommands, extension)
})
return binaryName
}
func formatterScriptExtension() string {
if runtime.GOOS == "windows" {
return ".bat"
}
return ""
}
func succeedingFormatterScript() string {
if runtime.GOOS == "windows" {
return "@echo off\r\nexit /b 0\r\n"
}
return "#!/bin/sh\nexit 0\n"
}
func failingFormatterScript() string {
if runtime.GOOS == "windows" {
return "@echo off\r\nexit /b 3\r\n"
}
return "#!/bin/sh\nexit 3\n"
}
// shortenFormatOnWriteTimeout narrows the production deadline for one test.
func shortenFormatOnWriteTimeout(t *testing.T, timeout time.Duration) {
t.Helper()
previous := formatOnWriteTimeout
formatOnWriteTimeout = timeout
t.Cleanup(func() { formatOnWriteTimeout = previous })
}
// A FAILED FORMATTER MUST NOT LEAVE THE FILE HALF-REWRITTEN.
//
// These commands edit in place, so one killed by the deadline, killed by the
// caller, or exiting partway through its own rewrite can leave the target
// truncated: neither the input nor the output. Returning the written bytes on
// top of that would leave the tracker baseline and the diff preview describing
// a file that is not on disk, and the next edit would compare against content
// the file does not have.
func TestFormatOnWriteRestoresTheFileWhenTheFormatterFails(t *testing.T) {
installFakeFormatter(t, ".clobberfmt", "clobberfmt", clobberingFormatterScript())
t.Setenv("ZERO_FORMAT_ON_WRITE", "1")
target := filepath.Join(t.TempDir(), "subject.clobberfmt")
const written = "the bytes the caller wrote\n"
if err := os.WriteFile(target, []byte(written), 0o644); err != nil {
t.Fatal(err)
}
requireFormatterClobbers(t, written)
formatting := maybeFormatWrittenFile(context.Background(), target, written)
onDisk, err := os.ReadFile(target)
if err != nil {
t.Fatal(err)
}
if string(onDisk) != written {
t.Errorf("file on disk = %q, want the bytes that were written back", onDisk)
}
if formatting.Content != written {
t.Errorf("content = %q, want the bytes that were written", formatting.Content)
}
if formatting.RestoreFailed {
t.Error("restoration was reported as failed on a writable file")
}
}
// And the same for a run cut off by the deadline, which is the case the notice
// already covers: the disclosure and the file have to agree.
func TestFormatOnWriteRestoresTheFileOnTimeout(t *testing.T) {
installFakeFormatter(t, ".clobberfmt", "clobberfmt", clobberingFormatterScript())
shortenFormatOnWriteTimeout(t, time.Nanosecond)
t.Setenv("ZERO_FORMAT_ON_WRITE", "1")
target := filepath.Join(t.TempDir(), "subject.clobberfmt")
const written = "the bytes the caller wrote\n"
if err := os.WriteFile(target, []byte(written), 0o644); err != nil {
t.Fatal(err)
}
requireFormatterClobbers(t, written)
formatting := maybeFormatWrittenFile(context.Background(), target, written)
onDisk, err := os.ReadFile(target)
if err != nil {
t.Fatal(err)
}
if string(onDisk) != written {
t.Errorf("file on disk = %q, want the bytes that were written back", onDisk)
}
if !formatting.TimedOut {
t.Error("the deadline path stopped being reported once restoration was added")
}
if notice := formatting.notice("subject.clobberfmt"); !strings.Contains(notice, "not formatted") {
t.Errorf("notice = %q, want the timeout note", notice)
}
}
// clobberingFormatterScript truncates the file it is handed and then fails, the
// way an interrupted in-place formatter leaves a partial rewrite.
func clobberingFormatterScript() string {
if runtime.GOOS == "windows" {
return "@echo off\r\necho CLOBBERED> %1\r\nexit /b 3\r\n"
}
return "#!/bin/sh\necho CLOBBERED > \"$1\"\nexit 3\n"
}
// requireFormatterClobbers proves the fixture really does damage the file it is
// handed, on a throwaway copy.
//
// It cannot be checked on the real target: restoration is the behaviour under
// test, so when it works the evidence is gone, and asserting on the target
// afterwards would either pass vacuously or report the opposite of what it saw.
func requireFormatterClobbers(t *testing.T, written string) {
t.Helper()
probe := filepath.Join(t.TempDir(), "probe.clobberfmt")
if err := os.WriteFile(probe, []byte(written), 0o644); err != nil {
t.Fatal(err)
}
binary, err := exec.LookPath("clobberfmt" + formatterScriptExtension())
if err != nil {
t.Fatalf("SETUP INVALID: the fake formatter is not on PATH: %v", err)
}
_ = exec.Command(binary, probe).Run()
after, err := os.ReadFile(probe)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(string(after), "CLOBBERED") {
t.Fatalf("SETUP INVALID: the fake formatter left %q, so it does not damage the file and restoration is not under test", after)
}
}
+19 -2
View File
@@ -548,11 +548,11 @@ func browserActionArgs(args map[string]any) (string, []string, error) {
if err != nil {
return "", nil, err
}
command = strings.ToLower(strings.TrimSpace(command))
spec, ok := browserActionSpecs[command]
command, ok := NormalizedBrowserActionCommand(command)
if !ok {
return "", nil, fmt.Errorf("command must be one of: %s", strings.Join(browserActionCommandNames(), ", "))
}
spec := browserActionSpecs[command]
values, err := stringArrayArg(args, "args")
if err != nil {
return "", nil, err
@@ -570,6 +570,16 @@ func browserActionArgs(args map[string]any) (string, []string, error) {
return command, commandArgs, nil
}
// NormalizedBrowserActionCommand returns the exact action that browser_action
// will execute after normalizing its command argument. ACP uses it only for a
// permission title, so an unrecognised command is never reflected as text that
// the tool would reject.
func NormalizedBrowserActionCommand(command string) (string, bool) {
command = strings.ToLower(strings.TrimSpace(command))
_, ok := browserActionSpecs[command]
return command, ok
}
func browserActionCommandArgs(command string, spec browserActionSpec, values []string) ([]string, error) {
switch command {
case "connect":
@@ -657,6 +667,13 @@ func browserOpenURLArg(args map[string]any) (string, error) {
if err != nil {
return "", err
}
return NormalizeBrowserOpenURL(rawURL)
}
// NormalizeBrowserOpenURL applies the browser_open URL rules before execution.
// Keeping this exported within the internal package lets permission displays
// describe the same destination the browser helper will open.
func NormalizeBrowserOpenURL(rawURL string) (string, error) {
normalized := strings.TrimSpace(rawURL)
if !strings.Contains(normalized, "://") {
normalized = "https://" + normalized
+3 -1
View File
@@ -114,7 +114,8 @@ func (tool writeFileTool) RunWithOptions(ctx context.Context, args map[string]an
// Optional format-on-write (ZERO_FORMAT_ON_WRITE). Must run BEFORE the
// FileTracker baseline: recording pre-format content would make the very
// next edit look like an external modification and trip the conflict guard.
content = maybeFormatWrittenFile(ctx, absolutePath, content)
formatting := maybeFormatWrittenFile(ctx, absolutePath, content)
content = formatting.Content
// Baseline the freshly written content so a later edit/overwrite in this
// session compares against what is now on disk.
newInfo, _ := os.Stat(absolutePath)
@@ -137,6 +138,7 @@ func (tool writeFileTool) RunWithOptions(ctx context.Context, args map[string]an
lines++
}
summary := fmt.Sprintf("%s %s (%d lines).", verb, relativePath, lines)
summary += formatting.notice(relativePath)
summary += inlineDiagnostics(ctx, options, absolutePath, relativePath)
result := okResult(summary)
result.ChangedFiles = []string{relativePath}
+1 -1
View File
@@ -109,7 +109,7 @@ var commandDefinitions = []commandDefinition{
name: "/voice",
usage: "/voice",
group: commandGroupRuntime,
description: "Toggle voice mode (hold Space to dictate).",
description: "Toggle voice mode (" + voiceCaptureUsage + ").",
kind: commandVoice,
},
{
+9 -14
View File
@@ -88,7 +88,7 @@ type dictationController struct {
waveBars []int
waveTick int
// voiceModeEnabled repurposes Space into hold-to-record (§13.9).
// voiceModeEnabled repurposes Ctrl+Space into hold-to-record (§13.9).
voiceModeEnabled bool
// eventTypesSupported records whether the terminal confirmed key-release
// reporting (Kitty protocol); eventTypesKnown gates it until the terminal has
@@ -173,7 +173,7 @@ func firstNonEmptyStr(a, b string) string {
func (d dictationController) active() bool { return d.phase != dictIdle }
// toggleDictation starts a recording when idle and stops-and-transcribes when
// recording. Invoked by the voice-mode Space gesture; a call during
// recording. Invoked by the voice-mode Ctrl+Space gesture; a call during
// startup/transcription is ignored (the machine is mid-transition).
func (m model) toggleDictation() (model, tea.Cmd) {
if !m.dictation.available() {
@@ -193,15 +193,15 @@ func (m model) toggleDictation() (model, tea.Cmd) {
}
}
// toggleVoiceMode flips the /voice hold-to-record gesture — the dictation
// trigger. While on, Space records; run /voice again to type spaces normally.
// toggleVoiceMode flips the /voice hold-to-record gesture. While on,
// Ctrl+Space records and ordinary Space continues to type normally.
func (m model) toggleVoiceMode() (model, tea.Cmd) {
if !m.dictation.available() {
return m.appendSystemNotice("Dictation is not configured. See docs/dictation.md to set up a local engine or a Groq/OpenAI key."), nil
}
m.dictation.voiceModeEnabled = !m.dictation.voiceModeEnabled
if m.dictation.voiceModeEnabled {
return m.showTransientNoticeInline("Voice mode on — hold Space to dictate; run /voice again to turn it off.", transientNoticeSuccess), nil
return m.showTransientNoticeInline("Voice mode on — "+voiceCaptureUsage+". Run /voice again to turn it off.", transientNoticeSuccess), nil
}
// Turning voice off is the "done dictating" signal, so release the warm sherpa
// streaming server — otherwise a loaded model keeps idling in RAM (and holding
@@ -374,7 +374,7 @@ func (m model) handleDictationStarted(msg dictationStartedMsg) (model, tea.Cmd)
if m.dictation.phase == dictStarting {
m.dictation.phase = dictRecording
}
// A voice-mode Space release that arrived mid-startup asked us to stop as soon
// A voice-mode Ctrl+Space release that arrived mid-startup asked us to stop as soon
// as recording began.
if m.dictation.voiceStopPending && m.dictation.phase == dictRecording {
m.dictation.voiceStopPending = false
@@ -480,12 +480,8 @@ func needsLeadingSpace(state composerState) bool {
func (m model) dictationStatusChip() string {
switch m.dictation.phase {
case dictStarting, dictRecording:
stop := "release Space to stop"
if !m.dictation.eventTypesSupported {
stop = "press Space to stop" // press-to-toggle fallback
}
wave := zeroTheme.amber.Render("● " + renderWaveBars(m.dictation.waveBars) + " REC")
return wave + zeroTheme.muted.Render(" · "+stop+", Esc to cancel")
return wave + zeroTheme.muted.Render(" · "+m.voiceCaptureStatus()+", Esc to cancel")
case dictTranscribing:
return zeroTheme.accent.Render("●") + " " + zeroTheme.muted.Render("transcribing…")
}
@@ -601,13 +597,12 @@ func (m model) handleRecTick() (model, tea.Cmd) {
return m, nil
}
// voiceModeIndicator renders the persistent "voice mode on" hint shown while
// idle so the user knows Space is repurposed to hold-to-record (§10).
// voiceModeIndicator renders the current voice gesture or transcription state.
func (m model) voiceModeIndicator() string {
if !m.dictation.voiceModeEnabled {
return ""
}
return zeroTheme.accent.Render("🎙 voice") + zeroTheme.muted.Render(" · "+m.dictation.currentModelLabel())
return zeroTheme.accent.Render("🎙 voice") + zeroTheme.muted.Render(" · "+m.voiceCaptureStatus()+" · "+m.dictation.currentModelLabel())
}
// dictationErrorText renders a dictation error for the transcript: a missing-
+41 -10
View File
@@ -4,17 +4,19 @@ import (
tea "charm.land/bubbletea/v2"
)
// Voice mode's Space-hold gesture — the (only) dictation trigger. Two terminal
const voiceCaptureUsage = "hold Ctrl+Space to record; press it to toggle where key releases are unavailable"
// Voice mode's Ctrl+Space gesture — the (only) dictation trigger. Two terminal
// tiers:
//
// 1. The terminal confirms key-release reporting (Ghostty/Kitty/WezTerm, …):
// Space press starts recording, Space release stops it — true hold-to-record,
// no ambiguity.
// 2. The terminal does not confirm release events: Space falls back to
// Ctrl+Space press starts recording and release stops it — true
// hold-to-record, with ordinary Space left available for typing.
// 2. The terminal does not confirm release events: Ctrl+Space falls back to
// press-to-toggle (press to start, press again to stop) — a deliberately
// simpler, robust fallback than inferring release from key-repeat timing
// (racy in a terminal). This works on ANY terminal, so voice mode is always
// usable once enabled.
// (racy in a terminal). Legacy NUL input is decoded as Ctrl+Space by the
// terminal input layer, so the same matcher covers both tiers.
//
// Only active while voice mode is on; the rest of dispatch is built on
// KeyPressMsg and is untouched.
@@ -27,8 +29,37 @@ func (m model) handleKeyboardEnhancements(msg tea.KeyboardEnhancementsMsg) model
return m
}
// handleVoiceSpacePress handles a Space press while voice mode is on.
func (m model) handleVoiceSpacePress(msg tea.KeyMsg) (model, tea.Cmd) {
func voiceCaptureKey(msg tea.KeyMsg) bool {
mod := msg.Key().Mod &^ (tea.ModCapsLock | tea.ModNumLock)
return keyIs(msg, tea.KeySpace) && mod == tea.ModCtrl
}
// voiceCaptureReleaseKey matches either half of the held Ctrl+Space chord. A
// terminal may report Ctrl releasing first, or report the later Space release
// without ModCtrl because Ctrl is no longer down.
func voiceCaptureReleaseKey(msg tea.KeyMsg) bool {
return keyIs(msg, tea.KeySpace) || keyIs(msg, tea.KeyLeftCtrl) || keyIs(msg, tea.KeyRightCtrl)
}
func (m model) voiceCaptureStatus() string {
switch m.dictation.phase {
case dictStarting, dictRecording:
if m.dictation.eventTypesSupported {
return "release Ctrl+Space to stop"
}
return "press Ctrl+Space to stop"
case dictTranscribing:
return "transcribing…"
default:
if m.dictation.eventTypesSupported {
return "hold Ctrl+Space to record"
}
return "press Ctrl+Space to record"
}
}
// handleVoiceCapturePress handles Ctrl+Space while voice mode is on.
func (m model) handleVoiceCapturePress(msg tea.KeyMsg) (model, tea.Cmd) {
if !m.dictation.eventTypesSupported {
// Tier 2: no release events — press-to-toggle.
return m.toggleDictation()
@@ -45,8 +76,8 @@ func (m model) handleVoiceSpacePress(msg tea.KeyMsg) (model, tea.Cmd) {
return m, nil // already recording; the release will stop it
}
// handleVoiceSpaceRelease stops a hold-to-record session when Space is released.
func (m model) handleVoiceSpaceRelease() (model, tea.Cmd) {
// handleVoiceCaptureRelease stops a hold-to-record session when Ctrl+Space is released.
func (m model) handleVoiceCaptureRelease() (model, tea.Cmd) {
if !m.dictation.spaceHeld {
return m, nil
}
+200 -15
View File
@@ -2,6 +2,7 @@ package tui
import (
"context"
"strings"
"testing"
tea "charm.land/bubbletea/v2"
@@ -36,7 +37,7 @@ func TestToggleVoiceModeFlips(t *testing.T) {
if !next.dictation.voiceModeEnabled {
t.Fatal("first /voice should enable voice mode")
}
if next.transientNotice.text != "Voice mode on — hold Space to dictate; run /voice again to turn it off." {
if next.transientNotice.text != "Voice mode on — hold Ctrl+Space to record; press it to toggle where key releases are unavailable. Run /voice again to turn it off." {
t.Errorf("voice-mode-on notice = %q", next.transientNotice.text)
}
if transcriptHasText(next, "Voice mode on") {
@@ -68,15 +69,27 @@ func TestKeyboardEnhancementsRecorded(t *testing.T) {
}
}
func TestVoiceSpaceHoldStartsRecording(t *testing.T) {
func TestVoiceCaptureKeyIgnoresLockModifiers(t *testing.T) {
for _, mod := range []tea.KeyMod{
tea.ModCtrl | tea.ModCapsLock,
tea.ModCtrl | tea.ModNumLock,
tea.ModCtrl | tea.ModCapsLock | tea.ModNumLock,
} {
if !voiceCaptureKey(tea.KeyPressMsg(tea.Key{Code: tea.KeySpace, Mod: mod})) {
t.Fatalf("Ctrl+Space with lock modifiers %v was not recognized", mod)
}
}
}
func TestVoiceCaptureHoldStartsRecording(t *testing.T) {
m := model{dictation: batchOnlyController()}
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = true
press := tea.KeyPressMsg(tea.Key{Code: tea.KeySpace})
next, cmd := m.handleVoiceSpacePress(press)
press := tea.KeyPressMsg(tea.Key{Code: tea.KeySpace, Mod: tea.ModCtrl})
next, cmd := m.handleVoiceCapturePress(press)
if next.dictation.phase != dictStarting {
t.Fatalf("Space press should start recording (phase=%d)", next.dictation.phase)
t.Fatalf("Ctrl+Space press should start recording (phase=%d)", next.dictation.phase)
}
if !next.dictation.spaceHeld {
t.Error("spaceHeld should be set in hold mode")
@@ -86,28 +99,28 @@ func TestVoiceSpaceHoldStartsRecording(t *testing.T) {
}
}
func TestVoiceSpaceHoldIgnoresRepeat(t *testing.T) {
func TestVoiceCaptureHoldIgnoresRepeat(t *testing.T) {
m := model{dictation: batchOnlyController()}
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = true
m.dictation.phase = dictRecording
m.dictation.spaceHeld = true
repeat := tea.KeyPressMsg(tea.Key{Code: tea.KeySpace, IsRepeat: true})
next, _ := m.handleVoiceSpacePress(repeat)
repeat := tea.KeyPressMsg(tea.Key{Code: tea.KeySpace, Mod: tea.ModCtrl, IsRepeat: true})
next, _ := m.handleVoiceCapturePress(repeat)
if next.dictation.phase != dictRecording {
t.Error("auto-repeat while held must not restart or change phase")
}
}
func TestVoiceSpaceReleaseStops(t *testing.T) {
func TestVoiceCaptureReleaseStops(t *testing.T) {
m := model{dictation: batchOnlyController()}
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = true
m.dictation.phase = dictRecording
m.dictation.spaceHeld = true
next, _ := m.handleVoiceSpaceRelease()
next, _ := m.handleVoiceCaptureRelease()
if next.dictation.phase != dictTranscribing {
t.Errorf("release should stop recording → transcribing (phase=%d)", next.dictation.phase)
}
@@ -123,7 +136,7 @@ func TestVoiceReleaseDuringStartupDefersStop(t *testing.T) {
m.dictation.phase = dictStarting
m.dictation.spaceHeld = true
next, _ := m.handleVoiceSpaceRelease()
next, _ := m.handleVoiceCaptureRelease()
if !next.dictation.voiceStopPending {
t.Error("a release during startup should defer the stop")
}
@@ -134,14 +147,186 @@ func TestVoiceReleaseDuringStartupDefersStop(t *testing.T) {
}
}
func TestVoiceSpaceToggleFallback(t *testing.T) {
func TestVoiceCaptureToggleFallback(t *testing.T) {
m := model{dictation: batchOnlyController()}
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = false // no release events → toggle fallback
press := tea.KeyPressMsg(tea.Key{Code: tea.KeySpace})
next, _ := m.handleVoiceSpacePress(press)
press := tea.KeyPressMsg(tea.Key{Code: tea.KeySpace, Mod: tea.ModCtrl})
next, _ := m.handleVoiceCapturePress(press)
if next.dictation.phase != dictStarting {
t.Errorf("toggle-mode Space should start recording (phase=%d)", next.dictation.phase)
t.Errorf("toggle-mode Ctrl+Space should start recording (phase=%d)", next.dictation.phase)
}
}
func TestVoiceModeCtrlSpaceReleaseStopsHoldRecording(t *testing.T) {
m := newModel(t.Context(), Options{})
m.dictation = batchOnlyController()
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = true
m.dictation.phase = dictRecording
m.dictation.spaceHeld = true
updated, _ := m.Update(tea.KeyReleaseMsg(tea.Key{Code: tea.KeySpace, Mod: tea.ModCtrl}))
next := updated.(model)
if next.dictation.phase != dictTranscribing || next.dictation.spaceHeld {
t.Fatalf("Ctrl+Space release did not stop hold recording: phase=%d held=%v", next.dictation.phase, next.dictation.spaceHeld)
}
}
func TestVoiceModeCaptureStopsWhenEitherChordKeyIsReleased(t *testing.T) {
for _, test := range []struct {
name string
key tea.Key
}{
{name: "left Ctrl first", key: tea.Key{Code: tea.KeyLeftCtrl}},
{name: "right Ctrl first", key: tea.Key{Code: tea.KeyRightCtrl}},
{name: "Space after Ctrl", key: tea.Key{Code: tea.KeySpace}},
} {
t.Run(test.name, func(t *testing.T) {
m := newModel(t.Context(), Options{})
m.dictation = batchOnlyController()
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = true
m.dictation.phase = dictRecording
m.dictation.spaceHeld = true
updated, _ := m.Update(tea.KeyReleaseMsg(test.key))
next := updated.(model)
if next.dictation.phase != dictTranscribing || next.dictation.spaceHeld {
t.Fatalf("release did not stop hold recording: phase=%d held=%v", next.dictation.phase, next.dictation.spaceHeld)
}
})
}
}
func TestVoiceModeIndicatorDescribesTierAndPhase(t *testing.T) {
for _, test := range []struct {
name string
releaseEvents bool
phase dictationPhase
spaceHeld bool
want string
}{
{name: "toggle idle", want: "press Ctrl+Space to record"},
{name: "hold idle", releaseEvents: true, want: "hold Ctrl+Space to record"},
{name: "toggle recording", phase: dictRecording, want: "press Ctrl+Space to stop"},
{name: "hold recording", releaseEvents: true, phase: dictRecording, spaceHeld: true, want: "release Ctrl+Space to stop"},
{name: "transcribing", phase: dictTranscribing, want: "transcribing"},
} {
t.Run(test.name, func(t *testing.T) {
m := model{dictation: batchOnlyController()}
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = test.releaseEvents
m.dictation.phase = test.phase
m.dictation.spaceHeld = test.spaceHeld
if got := m.voiceModeIndicator(); !strings.Contains(strings.ToLower(got), strings.ToLower(test.want)) {
t.Fatalf("voice indicator = %q, want %q", got, test.want)
}
})
}
}
func TestDictationStatusChipUsesVoiceCaptureShortcut(t *testing.T) {
for _, test := range []struct {
name string
releaseEvents bool
want string
}{
{name: "toggle", want: "press Ctrl+Space to stop"},
{name: "hold", releaseEvents: true, want: "release Ctrl+Space to stop"},
} {
t.Run(test.name, func(t *testing.T) {
m := model{dictation: batchOnlyController()}
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = test.releaseEvents
m.dictation.phase = dictRecording
if got := m.dictationStatusChip(); !strings.Contains(got, test.want) {
t.Fatalf("dictation status = %q, want %q", got, test.want)
}
})
}
}
func TestVoiceCommandDescriptionExplainsBothTerminalTiers(t *testing.T) {
for _, command := range commandDefinitions {
if command.kind == commandVoice {
for _, want := range []string{"hold Ctrl+Space", "press it to toggle", "key releases"} {
if !strings.Contains(command.description, want) {
t.Fatalf("voice command description = %q, want %q", command.description, want)
}
}
return
}
}
t.Fatal("voice command definition not found")
}
func TestVoiceModePlainSpaceTypesNormally(t *testing.T) {
for _, releaseEvents := range []bool{false, true} {
t.Run(map[bool]string{false: "toggle", true: "hold"}[releaseEvents], func(t *testing.T) {
m := newModel(t.Context(), Options{})
m.dictation = batchOnlyController()
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = releaseEvents
updated, cmd := m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeySpace}))
next := updated.(model)
if got := next.composerValue(); got != " " {
t.Fatalf("plain Space composer = %q, want one space", got)
}
if next.dictation.phase != dictIdle || cmd != nil {
t.Fatalf("plain Space started dictation: phase=%d cmd=%v", next.dictation.phase, cmd != nil)
}
})
}
}
func TestVoiceModeCtrlSpaceStartsAcrossTerminalTiers(t *testing.T) {
for _, releaseEvents := range []bool{false, true} {
t.Run(map[bool]string{false: "toggle", true: "hold"}[releaseEvents], func(t *testing.T) {
m := newModel(t.Context(), Options{})
m.dictation = batchOnlyController()
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = releaseEvents
updated, cmd := m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeySpace, Mod: tea.ModCtrl}))
next := updated.(model)
if next.dictation.phase != dictStarting || cmd == nil {
t.Fatalf("Ctrl+Space did not start dictation: phase=%d cmd=%v", next.dictation.phase, cmd != nil)
}
if got := next.composerValue(); got != "" {
t.Fatalf("Ctrl+Space typed into composer: %q", got)
}
})
}
}
func TestVoiceModeTypingPreservesSpacesAcrossTerminalTiers(t *testing.T) {
for _, releaseEvents := range []bool{false, true} {
t.Run(map[bool]string{false: "toggle", true: "hold"}[releaseEvents], func(t *testing.T) {
m := newModel(t.Context(), Options{})
m.dictation = batchOnlyController()
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = releaseEvents
for _, key := range []tea.Key{
{Code: 'f', Text: "f"},
{Code: 'i', Text: "i"},
{Code: 'x', Text: "x"},
{Code: tea.KeySpace},
{Code: 't', Text: "t"},
{Code: 'h', Text: "h"},
{Code: 'e', Text: "e"},
} {
updated, _ := m.Update(tea.KeyPressMsg(key))
m = updated.(model)
}
if got := m.composerValue(); got != "fix the" {
t.Fatalf("composer = %q, want %q", got, "fix the")
}
if m.dictation.phase != dictIdle {
t.Fatalf("typing started dictation: phase=%d", m.dictation.phase)
}
})
}
}
+8 -9
View File
@@ -1539,10 +1539,10 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
case tea.KeyboardEnhancementsMsg:
return m.handleKeyboardEnhancements(msg), nil
case tea.KeyReleaseMsg:
// Voice mode's hold-to-record ends on Space release; every other release
// event is ignored (dispatch elsewhere is press-based).
if m.dictation.voiceModeEnabled && keyIs(msg, tea.KeySpace) {
return m.handleVoiceSpaceRelease()
// Voice mode's hold-to-record ends on Ctrl+Space release; every other
// release event is ignored (dispatch elsewhere is press-based).
if m.dictation.voiceModeEnabled && voiceCaptureReleaseKey(msg) {
return m.handleVoiceCaptureRelease()
}
return m, nil
case tea.KeyPressMsg:
@@ -1676,11 +1676,10 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
return m.appendSystemNotice(fmt.Sprintf("Mouse released — drag to select and copy text. Press %s again to re-enable mouse interaction (clicks, right-click paste).", mouseKey)), nil
}
return m.showTransientNoticeInline("Mouse interaction re-enabled.", transientNoticeSuccess), nil
case m.dictation.voiceModeEnabled && !m.transcriptDetailed && keyIs(msg, tea.KeySpace) && !keyHasMod(msg, tea.ModCtrl) && !keyAlt(msg) && m.noBlockingModal():
// Voice mode (/voice) repurposes Space into the record gesture — the only
// dictation trigger — so it must not also type a space. Turn voice mode
// off (/voice) to type normally.
return m.handleVoiceSpacePress(msg)
case m.dictation.voiceModeEnabled && !m.transcriptDetailed && voiceCaptureKey(msg) && m.noBlockingModal():
// Voice mode reserves Ctrl+Space for recording; ordinary Space continues
// through the composer path below.
return m.handleVoiceCapturePress(msg)
case keyIs(msg, tea.KeyEsc):
// Esc is heavily overloaded below (subchat exit, MCP cancel, ask-user,
// permission deny, wizard/picker/suggestions dismiss, ...) before ever
+1 -1
View File
@@ -123,7 +123,7 @@ func (m model) finalizeSTTModelFromKey(p *sttKeyPromptState, prefix string) (tea
if modelID != "" {
label += " " + modelID
}
return m.showTransientNoticeInline(prefix+"Dictation model: "+label+". Run /voice, then hold Space to dictate.", transientNoticeSuccess), nil
return m.showTransientNoticeInline(prefix+"Dictation model: "+label+". Run /voice; "+voiceCaptureUsage+".", transientNoticeSuccess), nil
}
// sttKeyPromptOverlay renders the API-key prompt as a centered modal, matching
+1 -1
View File
@@ -226,7 +226,7 @@ func (m model) handleSTTModelSelection(value string) (model, string) {
// instead); this is the manual-setup path (e.g. Termux).
hint += " Set stt.localModelPath to a model directory (see docs/dictation.md)."
} else {
hint += " Run /voice, then hold Space to dictate."
hint += " Run /voice; " + voiceCaptureUsage + "."
}
return m, hint
}