mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-09 13:31:26 -10:00
Compare commits
12
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
29a5253013 | ||
|
|
789bd389aa | ||
|
|
a67ea8735a | ||
|
|
e27174f6f4 | ||
|
|
c35412bd01 | ||
|
|
dab4691b21 | ||
|
|
23b1e5f92f | ||
|
|
1e52a3625d | ||
|
|
139cffa6c3 | ||
|
|
1e4db7c664 | ||
|
|
f9092668a8 | ||
|
|
7f5e5af7bb |
No files matched your search
@@ -273,6 +273,19 @@ zero sandbox grants list
|
||||
Zero includes local file/search/edit/shell tools, `web_fetch` for public URLs,
|
||||
and MCP support for additional tools.
|
||||
|
||||
`web_fetch` refuses loopback, private and other special-use addresses: it checks
|
||||
the URL before asking permission, resolves the host, and dials the address it
|
||||
validated so a name cannot resolve to something else in between.
|
||||
|
||||
If `HTTP_PROXY`/`HTTPS_PROXY` is set, `web_fetch` and the provider connectivity
|
||||
probe use it, and that last step changes: the proxy is dialed and the target
|
||||
hostname is sent to it, so the proxy decides which address the request actually
|
||||
reaches. The URL is still checked and resolved locally first, but a proxy that
|
||||
answers differently can reach a private service. A forward proxy already sees
|
||||
and can rewrite every request through it, so this is the trust you accept by
|
||||
configuring one. Leave the variables unset for the checks to be enforced end to
|
||||
end.
|
||||
|
||||
For local dev servers, use shell commands such as `curl` through `exec_command`
|
||||
so the normal sandbox and permission policy applies. Long-running commands stay
|
||||
attached to a background terminal session and can be listed or stopped from the
|
||||
|
||||
@@ -627,7 +627,12 @@ Core tool groups include:
|
||||
skills, `ask_user`, permission requests.
|
||||
- **Write tools**: `write_file`, `edit_file`, `apply_patch`, `update_plan`.
|
||||
- **Shell tools**: `exec_command`, `write_stdin`, legacy `bash`.
|
||||
- **Network tools**: `web_fetch` and optional web search backend.
|
||||
- **Network tools**: `web_fetch` and optional web search backend. `web_fetch`
|
||||
validates the URL before the permission prompt, refuses loopback/private/
|
||||
special-use addresses, and dials the address it resolved. With
|
||||
`HTTP_PROXY`/`HTTPS_PROXY` set, the proxy is dialed instead and resolves the
|
||||
target itself, so destination enforcement ends at the proxy; see the note in
|
||||
the README.
|
||||
- **Local-control tools**: config-gated wrappers for browser, terminal, desktop,
|
||||
and artifact helpers backed by packages such as `internal/localcontrol` and
|
||||
`internal/browser`.
|
||||
|
||||
+52
-6
@@ -14,6 +14,7 @@ import (
|
||||
"github.com/Gitlawb/zero/internal/providercatalog"
|
||||
"github.com/Gitlawb/zero/internal/providermodelcatalog"
|
||||
"github.com/Gitlawb/zero/internal/providermodeldiscovery"
|
||||
"github.com/Gitlawb/zero/internal/redaction"
|
||||
"github.com/Gitlawb/zero/internal/sandbox"
|
||||
"github.com/Gitlawb/zero/internal/sessions"
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
@@ -30,10 +31,12 @@ type Deps struct {
|
||||
DiscoverModels func(context.Context, config.ProviderProfile) ([]providermodeldiscovery.Model, error)
|
||||
NewProvider func(profile config.ProviderProfile) (zeroruntime.Provider, error)
|
||||
RunAgent func(ctx context.Context, prompt string, provider zeroruntime.Provider, opts agent.Options) (agent.Result, error)
|
||||
// BuildWorkspace builds the SCOPED tool registry and the sandbox engine for a
|
||||
// validated workspace root, so ACP shell tools (bash/exec_command) are confined
|
||||
// exactly like the exec surface — never run unconfined on the host.
|
||||
BuildWorkspace func(workspaceRoot string, resolved config.ResolvedConfig) (*tools.Registry, *sandbox.Engine, error)
|
||||
// BuildWorkspace creates the per-turn scoped tool workspace for a validated
|
||||
// workspace root. Its Close method releases any resources the registry owns
|
||||
// (notably MCP server connections) after the turn completes or is cancelled.
|
||||
// Shell tools remain confined exactly like the exec surface — never unconfined
|
||||
// on the host.
|
||||
BuildWorkspace func(ctx context.Context, workspaceRoot string, resolved config.ResolvedConfig, mode agent.PermissionMode) (*Workspace, error)
|
||||
// ResolveWorkspaceRoot validates + normalizes a client-supplied cwd (must be an
|
||||
// existing directory; never the bare root). It is the file-tool confinement root.
|
||||
ResolveWorkspaceRoot func(cwd string) (string, error)
|
||||
@@ -41,6 +44,27 @@ type Deps struct {
|
||||
AgentInfo Implementation
|
||||
}
|
||||
|
||||
// Workspace is the per-turn execution environment passed to the agent. ACP does
|
||||
// not retain it across turns because MCP connections belong to the registry that
|
||||
// advertised their tools; keeping a stale registry after a cancelled turn would
|
||||
// leak its server process and its permission state into the next turn.
|
||||
type Workspace struct {
|
||||
Registry *tools.Registry
|
||||
Sandbox *sandbox.Engine
|
||||
DeferThreshold int
|
||||
Notices []string
|
||||
Cleanup func() error
|
||||
}
|
||||
|
||||
// Close releases resources created with the workspace. A nil cleanup is valid
|
||||
// for core-only workspaces.
|
||||
func (w *Workspace) Close() error {
|
||||
if w == nil || w.Cleanup == nil {
|
||||
return nil
|
||||
}
|
||||
return w.Cleanup()
|
||||
}
|
||||
|
||||
// Agent is the ACP agent server bound to one JSON-RPC connection (one editor).
|
||||
type Agent struct {
|
||||
conn *Conn
|
||||
@@ -241,13 +265,34 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string,
|
||||
if err != nil {
|
||||
return "", RPCError(codeInternalError, "provider: "+err.Error())
|
||||
}
|
||||
mode := sess.currentMode()
|
||||
// Build the SCOPED registry + sandbox engine for this session's workspace so
|
||||
// shell/file tools are confined to the workspace exactly like the exec surface.
|
||||
registry, sandboxEngine, err := a.deps.BuildWorkspace(sess.cwd, resolved)
|
||||
// The workspace can also own MCP connections, which must be closed after every
|
||||
// turn even if the agent run returns an error or the client cancels it.
|
||||
workspace, err := a.deps.BuildWorkspace(ctx, sess.cwd, resolved, mode)
|
||||
if err != nil {
|
||||
return "", RPCError(codeInternalError, "workspace: "+err.Error())
|
||||
}
|
||||
if workspace == nil || workspace.Registry == nil {
|
||||
if workspace != nil {
|
||||
_ = workspace.Close()
|
||||
}
|
||||
return "", RPCError(codeInternalError, "workspace: missing tool registry")
|
||||
}
|
||||
defer func() {
|
||||
if err := workspace.Close(); err != nil {
|
||||
log.Printf("acp: close workspace: %v", err)
|
||||
}
|
||||
}()
|
||||
registry := workspace.Registry
|
||||
sandboxEngine := workspace.Sandbox
|
||||
note := ¬ifier{conn: a.conn, sessionID: sess.id}
|
||||
for _, notice := range workspace.Notices {
|
||||
if text := strings.TrimSpace(redaction.RedactString(notice, redaction.Options{})); text != "" {
|
||||
note.text("\n\n[zero warning] " + text + "\n")
|
||||
}
|
||||
}
|
||||
|
||||
opts := agent.Options{
|
||||
Cwd: sess.cwd,
|
||||
@@ -256,7 +301,8 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string,
|
||||
Model: resolved.Provider.Model,
|
||||
Registry: registry,
|
||||
Sandbox: sandboxEngine,
|
||||
PermissionMode: sess.currentMode(),
|
||||
PermissionMode: mode,
|
||||
DeferThreshold: workspace.DeferThreshold,
|
||||
MaxTurns: resolved.MaxTurns,
|
||||
Images: images,
|
||||
OnText: note.text,
|
||||
|
||||
+152
-4
@@ -8,6 +8,7 @@ import (
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -35,6 +36,35 @@ func (f fakeProvider) StreamCompletion(_ context.Context, _ zeroruntime.Completi
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
type acpDeferredPromptTool struct{ name string }
|
||||
|
||||
func (t acpDeferredPromptTool) Name() string { return t.name }
|
||||
func (acpDeferredPromptTool) Description() string { return "deferred MCP test tool" }
|
||||
func (acpDeferredPromptTool) Parameters() tools.Schema { return tools.Schema{Type: "object"} }
|
||||
func (acpDeferredPromptTool) Deferred() bool { return true }
|
||||
func (acpDeferredPromptTool) Run(context.Context, map[string]any) tools.Result {
|
||||
return tools.Result{Status: tools.StatusOK, Output: "ok"}
|
||||
}
|
||||
func (acpDeferredPromptTool) Safety() tools.Safety {
|
||||
return tools.Safety{
|
||||
SideEffect: tools.SideEffectNetwork, Permission: tools.PermissionPrompt,
|
||||
AdvertiseInAuto: true,
|
||||
}
|
||||
}
|
||||
|
||||
type captureToolsProvider struct {
|
||||
requests chan zeroruntime.CompletionRequest
|
||||
}
|
||||
|
||||
func (p captureToolsProvider) StreamCompletion(_ context.Context, request zeroruntime.CompletionRequest) (<-chan zeroruntime.StreamEvent, error) {
|
||||
p.requests <- request
|
||||
ch := make(chan zeroruntime.StreamEvent, 2)
|
||||
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventText, Content: "done"}
|
||||
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}
|
||||
close(ch)
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
func testDeps(t *testing.T) Deps {
|
||||
t.Helper()
|
||||
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
|
||||
@@ -53,10 +83,10 @@ func testDeps(t *testing.T) Deps {
|
||||
return fakeProvider{text: "Hello from ZERO"}, nil
|
||||
},
|
||||
RunAgent: agent.Run,
|
||||
BuildWorkspace: func(string, config.ResolvedConfig) (*tools.Registry, *sandbox.Engine, error) {
|
||||
BuildWorkspace: func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
|
||||
r := tools.NewRegistry()
|
||||
r.Register(tools.NewUpdatePlanTool())
|
||||
return r, nil, nil
|
||||
return &Workspace{Registry: r}, nil
|
||||
},
|
||||
ResolveWorkspaceRoot: func(cwd string) (string, error) { return cwd, nil },
|
||||
Store: store,
|
||||
@@ -468,8 +498,8 @@ func TestACPRunTurnWiresSandboxAndScopedRegistry(t *testing.T) {
|
||||
reg := tools.NewRegistry()
|
||||
reg.Register(tools.NewUpdatePlanTool())
|
||||
engine := sandbox.NewEngine(sandbox.EngineOptions{WorkspaceRoot: t.TempDir()})
|
||||
deps.BuildWorkspace = func(string, config.ResolvedConfig) (*tools.Registry, *sandbox.Engine, error) {
|
||||
return reg, engine, nil
|
||||
deps.BuildWorkspace = func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
|
||||
return &Workspace{Registry: reg, Sandbox: engine}, nil
|
||||
}
|
||||
var captured agent.Options
|
||||
deps.RunAgent = func(_ context.Context, _ string, _ zeroruntime.Provider, opts agent.Options) (agent.Result, error) {
|
||||
@@ -497,6 +527,124 @@ func TestACPRunTurnWiresSandboxAndScopedRegistry(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestACPRunTurnUsesWorkspaceDeferThresholdAtProviderBoundary(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
threshold int
|
||||
want []string
|
||||
}{
|
||||
{name: "explicit zero disables", threshold: 0, want: []string{"mcp_docs_alpha", "mcp_docs_beta"}},
|
||||
{name: "below threshold stays eager", threshold: 3, want: []string{"mcp_docs_alpha", "mcp_docs_beta"}},
|
||||
{name: "at threshold defers", threshold: 2, want: []string{tools.ToolSearchToolName}},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
requests := make(chan zeroruntime.CompletionRequest, 4)
|
||||
deps := testDeps(t)
|
||||
deps.NewProvider = func(config.ProviderProfile) (zeroruntime.Provider, error) {
|
||||
return captureToolsProvider{requests: requests}, nil
|
||||
}
|
||||
deps.BuildWorkspace = func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(acpDeferredPromptTool{name: "mcp_docs_alpha"})
|
||||
registry.Register(acpDeferredPromptTool{name: "mcp_docs_beta"})
|
||||
registry.Register(tools.NewToolSearchTool(registry))
|
||||
return &Workspace{Registry: registry, DeferThreshold: tc.threshold}, nil
|
||||
}
|
||||
|
||||
h := newHarness(t, deps)
|
||||
defer h.stop()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
var created NewSessionResult
|
||||
if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir()}, &created); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{SessionID: created.SessionID, Prompt: []ContentBlock{TextBlock("hi")}}, &PromptResult{}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
request := <-requests
|
||||
got := make([]string, 0, len(request.Tools))
|
||||
for _, definition := range request.Tools {
|
||||
got = append(got, definition.Name)
|
||||
}
|
||||
if !slices.Equal(got, tc.want) {
|
||||
t.Fatalf("provider tools = %#v, want %#v", got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestACPRunTurnEmitsWorkspaceSetupNotices(t *testing.T) {
|
||||
deps := testDeps(t)
|
||||
deps.BuildWorkspace = func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.NewUpdatePlanTool())
|
||||
return &Workspace{Registry: registry, Notices: []string{"MCP server docs unavailable, skipped: timeout"}}, nil
|
||||
}
|
||||
h := newHarness(t, deps)
|
||||
defer h.stop()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
var created NewSessionResult
|
||||
if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir()}, &created); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{SessionID: created.SessionID, Prompt: []ContentBlock{TextBlock("hi")}}, &PromptResult{}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := drainText(t, h.updates)
|
||||
if !strings.Contains(got, "[zero warning] MCP server docs unavailable, skipped: timeout") || !strings.Contains(got, "Hello from ZERO") {
|
||||
t.Fatalf("streamed text = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestACPRunTurnClosesWorkspace proves a per-turn workspace never outlives the
|
||||
// agent run. This is particularly important for MCP: its registry owns live
|
||||
// client connections and, for stdio servers, child processes.
|
||||
func TestACPRunTurnClosesWorkspaceAfterSuccessAndFailure(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
runErr error
|
||||
}{
|
||||
{name: "success"},
|
||||
{name: "agent failure", runErr: errors.New("provider interrupted")},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
deps := testDeps(t)
|
||||
closed := 0
|
||||
deps.BuildWorkspace = func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.NewUpdatePlanTool())
|
||||
return &Workspace{Registry: registry, Cleanup: func() error {
|
||||
closed++
|
||||
return nil
|
||||
}}, nil
|
||||
}
|
||||
deps.RunAgent = func(context.Context, string, zeroruntime.Provider, agent.Options) (agent.Result, error) {
|
||||
return agent.Result{FinalAnswer: "ok"}, tc.runErr
|
||||
}
|
||||
h := newHarness(t, deps)
|
||||
defer h.stop()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
var created NewSessionResult
|
||||
if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir()}, &created); err != nil {
|
||||
t.Fatalf("session/new: %v", err)
|
||||
}
|
||||
err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{SessionID: created.SessionID, Prompt: []ContentBlock{TextBlock("hello")}}, &PromptResult{})
|
||||
if tc.runErr == nil && err != nil {
|
||||
t.Fatalf("session/prompt: %v", err)
|
||||
}
|
||||
if tc.runErr != nil && err == nil {
|
||||
t.Fatal("session/prompt succeeded after agent failure")
|
||||
}
|
||||
if closed != 1 {
|
||||
t.Fatalf("workspace cleanup calls = %d, want 1", closed)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestACPRejectsInvalidCwd confirms session/new fails when the workspace root
|
||||
// resolver rejects the client cwd (e.g. filesystem root).
|
||||
func TestACPRejectsInvalidCwd(t *testing.T) {
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// These ceilings are a deliberate ratchet on Zero's fixed per-turn overhead — the
|
||||
@@ -33,7 +35,24 @@ const (
|
||||
// The ratchet did its job: it caught the creep and forced a decision instead
|
||||
// of a drift. That is the point of it, so keep raising it deliberately rather
|
||||
// than reflexively.
|
||||
//
|
||||
// THIS COUNTS THE SCHEMAS ONLY, NOT THE HOST'S SHELL GUIDANCE. Until
|
||||
// 2026-09-07 it counted both, and that made it a different measurement on
|
||||
// every machine: exec_command appends shell guidance whose size depends on
|
||||
// which shell was detected, so the same tree measured 3498 on Linux, ~3600 on
|
||||
// a Windows host with pwsh 7, and 3653 on Windows PowerShell 5.1, which is
|
||||
// stock Windows. The 5.1 case has been over this ceiling since view_image
|
||||
// landed, so `go test ./...` failed for anyone on a stock Windows box while
|
||||
// CI stayed green, because the runners have pwsh 7 and take the cheaper
|
||||
// branch. A budget that passes or fails on which PowerShell happens to be
|
||||
// installed is not ratcheting anything.
|
||||
maxEagerToolSchemaTokens = 3650
|
||||
// The host-specific half, ratcheted on its own so it cannot creep either.
|
||||
// Measured 2026-09-07: Windows PowerShell 5.1 is the most expensive host at
|
||||
// 155 tokens (722 chars), pwsh 7 is 110, cmd.exe is 87, and every POSIX host
|
||||
// is 0 because no guidance is appended there. Set above the worst case, not
|
||||
// above whatever this machine happens to be.
|
||||
maxExecCommandShellGuidanceTokens = 200
|
||||
)
|
||||
|
||||
func TestSystemPromptTokenBudget(t *testing.T) {
|
||||
@@ -56,8 +75,35 @@ func TestEagerToolSchemaTokenBudget(t *testing.T) {
|
||||
// Options{} keeps DeferThreshold at 0, so deferral is inactive and every core
|
||||
// tool is exposed eagerly — exactly what a plugin-free session sends each turn.
|
||||
exposed, _ := partitionTools(registry, PermissionModeAuto, Options{}, map[string]bool{})
|
||||
got := estimateToolDefTokens(exposed)
|
||||
t.Logf("eager core tool schemas: %d tokens across %d tools", got, len(exposed))
|
||||
total := estimateToolDefTokens(exposed)
|
||||
// exec_command's description carries this host's shell guidance, and that
|
||||
// guidance is the one part of the eager set that is not the same everywhere.
|
||||
// Charged separately below rather than folded in here, or this ceiling would
|
||||
// mean a different thing on every machine. See the constants above.
|
||||
//
|
||||
// Removed from the description BEFORE estimating rather than subtracted from
|
||||
// the total afterwards. ApproxTextTokens is an integer division, so
|
||||
// floor((schema+guidance)/4) - floor(guidance/4) comes out one token high
|
||||
// whenever the two remainders add past four, and which side of that line a
|
||||
// host lands on depends on the length of its guidance: exactly the
|
||||
// host-dependence this number is supposed to be rid of.
|
||||
guidance := tools.HostExecCommandShellGuidance()
|
||||
schemas := make([]zeroruntime.ToolDefinition, len(exposed))
|
||||
copy(schemas, exposed)
|
||||
stripped := 0
|
||||
for index := range schemas {
|
||||
without := strings.TrimSuffix(schemas[index].Description, guidance)
|
||||
if without != schemas[index].Description {
|
||||
stripped++
|
||||
}
|
||||
schemas[index].Description = without
|
||||
}
|
||||
if guidance != "" && stripped != 1 {
|
||||
t.Fatalf("SETUP INVALID: this host's shell guidance came off %d descriptions, want exactly exec_command's; internal/tools pins it as that description's tail", stripped)
|
||||
}
|
||||
got := estimateToolDefTokens(schemas)
|
||||
guidanceTokens := ApproxTextTokens(guidance)
|
||||
t.Logf("eager core tool schemas: %d tokens across %d tools (%d total, %d host shell guidance)", got, len(exposed), total, guidanceTokens)
|
||||
if got > maxEagerToolSchemaTokens {
|
||||
// NOT "defer a tool": this test pins DeferThreshold at 0, so marking a
|
||||
// tool deferred leaves it exposed here and changes nothing. Deferral is
|
||||
@@ -66,6 +112,9 @@ func TestEagerToolSchemaTokenBudget(t *testing.T) {
|
||||
// deliberate raise.
|
||||
t.Fatalf("eager tool schemas are %d tokens, over the %d ceiling — trim a schema, drop a core tool, or raise the ceiling deliberately (deferring will NOT help: this test disables deferral)", got, maxEagerToolSchemaTokens)
|
||||
}
|
||||
if guidanceTokens > maxExecCommandShellGuidanceTokens {
|
||||
t.Fatalf("this host's exec_command shell guidance is %d tokens, over the %d ceiling — trim the guidance in internal/tools/shell_runtime.go or raise the ceiling deliberately", guidanceTokens, maxExecCommandShellGuidanceTokens)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentAdvertisesBothEditTools(t *testing.T) {
|
||||
|
||||
+75
-15
@@ -10,9 +10,11 @@ import (
|
||||
"github.com/Gitlawb/zero/internal/acp"
|
||||
"github.com/Gitlawb/zero/internal/agent"
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/execution"
|
||||
"github.com/Gitlawb/zero/internal/mcp"
|
||||
"github.com/Gitlawb/zero/internal/providermodeldiscovery"
|
||||
"github.com/Gitlawb/zero/internal/redaction"
|
||||
"github.com/Gitlawb/zero/internal/sandbox"
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
)
|
||||
|
||||
const acpUsage = `zero acp — serve the Agent Client Protocol (ACP) over stdio
|
||||
@@ -57,20 +59,8 @@ func runACP(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int
|
||||
// surface — no ACP-specific credential handling needed.
|
||||
NewProvider: deps.newProvider,
|
||||
RunAgent: agent.Run,
|
||||
// Build the SCOPED registry + sandbox engine per workspace, exactly like the
|
||||
// exec surface, so ACP shell/file tools are confined — never run unconfined.
|
||||
BuildWorkspace: func(workspaceRoot string, resolved config.ResolvedConfig) (*tools.Registry, *sandbox.Engine, error) {
|
||||
scope, err := sandbox.NewScope(workspaceRoot, resolved.Sandbox.AdditionalWriteRoots)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
engine, err := buildExecSandboxEngine(workspaceRoot, resolved, deps, scope)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
registry := newCoreRegistryScoped(workspaceRoot, scope)
|
||||
registerLocalControlTools(registry, workspaceRoot, resolved.LocalControl)
|
||||
return registry, engine, nil
|
||||
BuildWorkspace: func(ctx context.Context, workspaceRoot string, resolved config.ResolvedConfig, mode agent.PermissionMode) (*acp.Workspace, error) {
|
||||
return buildACPWorkspace(ctx, workspaceRoot, resolved, mode, deps)
|
||||
},
|
||||
ResolveWorkspaceRoot: acpWorkspaceRootResolver(deps),
|
||||
Store: deps.newSessionStore(),
|
||||
@@ -85,6 +75,76 @@ func runACP(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int
|
||||
return exitSuccess
|
||||
}
|
||||
|
||||
// buildACPWorkspace matches exec's registry construction for one ACP turn. MCP
|
||||
// servers use the same sandbox-prepared execution runner, project configuration is
|
||||
// gated by the validated workspace's trust state, and their runtime is released
|
||||
// when acp.Agent completes the turn. ACP deliberately stays at low autonomy: an
|
||||
// editor connection never upgrades MCP permissions beyond an interactive prompt.
|
||||
func buildACPWorkspace(ctx context.Context, workspaceRoot string, resolved config.ResolvedConfig, mode agent.PermissionMode, deps appDeps) (*acp.Workspace, error) {
|
||||
scope, err := sandbox.NewScope(workspaceRoot, resolved.Sandbox.AdditionalWriteRoots)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
engine, err := buildExecSandboxEngine(workspaceRoot, resolved, deps, scope)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
registry := newCoreRegistryScoped(workspaceRoot, scope)
|
||||
|
||||
workspace := &acp.Workspace{
|
||||
Registry: registry, Sandbox: engine,
|
||||
DeferThreshold: resolved.Tools.DeferThreshold,
|
||||
}
|
||||
if mode != agent.PermissionModePlan {
|
||||
// MCP stdio servers are subprocesses. Passing the engine to their runner
|
||||
// keeps them inside the exact sandbox / lifecycle path used by zero exec.
|
||||
runtime, trustSkip, err := registerMCPToolsForWorkspaceWithOptions(
|
||||
ctx, workspaceRoot, registry, deps, mcp.AutonomyLow, workspaceRoot,
|
||||
mcp.RegisterOptions{Execution: execution.NewRunner(engine), AdvertiseInAuto: true},
|
||||
)
|
||||
if err != nil {
|
||||
// RegisterTools may have connected an earlier server before reporting a
|
||||
// later failure, so do not orphan a partial runtime on this error path.
|
||||
if runtime != nil {
|
||||
_ = runtime.Close()
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
workspace.Cleanup = runtime.Close
|
||||
workspace.Notices = acpMCPSetupNotices(trustSkip, runtime)
|
||||
}
|
||||
registerLocalControlTools(registry, workspaceRoot, resolved.LocalControl)
|
||||
// MCP tools are deferred-eligible. Register their loader only after every
|
||||
// ACP-visible tool is present, using the same mode the agent receives.
|
||||
registerToolSearchIfEligible(registry, resolved.Tools.DeferThreshold, mode, nil, nil)
|
||||
return workspace, nil
|
||||
}
|
||||
|
||||
func acpMCPSetupNotices(skip trustSkip, runtime mcpToolRuntime) []string {
|
||||
var notices []string
|
||||
if skip.excludedProjectConfig {
|
||||
if skip.trustCheckErrored {
|
||||
notices = append(notices, "The workspace-trust store could not be read; project MCP servers were ignored (fail-closed). Run 'zero trust' to enable them.")
|
||||
} else {
|
||||
notices = append(notices, "Project MCP servers were ignored in this untrusted workspace. Run 'zero trust' to enable them.")
|
||||
}
|
||||
}
|
||||
if runtime == nil {
|
||||
return notices
|
||||
}
|
||||
for _, skipped := range runtime.Skipped() {
|
||||
if skipped.UnconfiguredDefault {
|
||||
continue
|
||||
}
|
||||
message := fmt.Sprintf("MCP server %s unavailable, skipped", skipped.Name)
|
||||
if skipped.Err != nil {
|
||||
message += ": " + redaction.ErrorMessage(skipped.Err, redaction.Options{})
|
||||
}
|
||||
notices = append(notices, redaction.RedactString(message, redaction.Options{}))
|
||||
}
|
||||
return notices
|
||||
}
|
||||
|
||||
// acpWorkspaceRootResolver validates a client-supplied cwd into a confinement
|
||||
// root. It reuses exec's resolveWorkspaceRoot (abs+clean, must be an existing
|
||||
// dir) and additionally rejects the filesystem root and the home directory — an
|
||||
|
||||
@@ -5,10 +5,19 @@ import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/agent"
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/mcp"
|
||||
"github.com/Gitlawb/zero/internal/sandbox"
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
"github.com/Gitlawb/zero/internal/workspacetrust"
|
||||
)
|
||||
|
||||
type acpTestReader func([]byte) (int, error)
|
||||
@@ -108,6 +117,198 @@ func TestRunACPIdleCancellationExitsCleanly(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildACPWorkspaceRegistersSandboxedMCPToolsAndClosesThem(t *testing.T) {
|
||||
setTrustConfigRoot(t)
|
||||
workspaceRoot := t.TempDir()
|
||||
if err := workspacetrust.Trust(workspaceRoot); err != nil {
|
||||
t.Fatalf("trust workspace: %v", err)
|
||||
}
|
||||
grantStore, err := sandbox.NewGrantStore(sandbox.StoreOptions{FilePath: t.TempDir() + "/grants.json"})
|
||||
if err != nil {
|
||||
t.Fatalf("new grant store: %v", err)
|
||||
}
|
||||
|
||||
var gotExclude, registered, closed bool
|
||||
deps := fillAppDeps(appDeps{
|
||||
resolveMCPConfig: func(root string, excludeProject bool) (config.MCPConfig, error) {
|
||||
if root != workspaceRoot {
|
||||
t.Fatalf("MCP workspace root = %q, want %q", root, workspaceRoot)
|
||||
}
|
||||
gotExclude = excludeProject
|
||||
return config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"docs": {Type: "stdio", Command: "fake-docs"},
|
||||
}}, nil
|
||||
},
|
||||
newMCPStore: func() (*mcp.PermissionStore, error) { return nil, nil },
|
||||
newSandboxStore: func() (*sandbox.GrantStore, error) { return grantStore, nil },
|
||||
registerMCPTools: func(ctx context.Context, registry *tools.Registry, cfg config.MCPConfig, options mcp.RegisterOptions) (mcpToolRuntime, error) {
|
||||
if ctx == nil {
|
||||
t.Fatal("MCP registration received nil context")
|
||||
}
|
||||
if len(cfg.Servers) != 1 || cfg.Servers["docs"].Command != "fake-docs" {
|
||||
t.Fatalf("MCP config = %+v", cfg)
|
||||
}
|
||||
if options.Autonomy != mcp.AutonomyLow {
|
||||
t.Fatalf("MCP autonomy = %q, want low", options.Autonomy)
|
||||
}
|
||||
if options.Execution == nil {
|
||||
t.Fatal("MCP stdio server was not given the sandbox execution runner")
|
||||
}
|
||||
if options.WorkspaceRoot != workspaceRoot {
|
||||
t.Fatalf("MCP execution workspace = %q, want %q", options.WorkspaceRoot, workspaceRoot)
|
||||
}
|
||||
if !options.AdvertiseInAuto {
|
||||
t.Fatal("ACP MCP tools must be advertised in auto mode without being pre-approved")
|
||||
}
|
||||
registered = true
|
||||
registry.Register(cliFakeDeferredTool{
|
||||
name: "mcp_docs_search", permission: tools.PermissionPrompt,
|
||||
advertiseInAuto: options.AdvertiseInAuto,
|
||||
})
|
||||
return closeFunc(func() error {
|
||||
closed = true
|
||||
return nil
|
||||
}), nil
|
||||
},
|
||||
})
|
||||
|
||||
workspace, err := buildACPWorkspace(context.Background(), workspaceRoot, config.ResolvedConfig{
|
||||
Tools: config.ToolsConfig{DeferThreshold: 1},
|
||||
}, agent.PermissionModeAuto, deps)
|
||||
if err != nil {
|
||||
t.Fatalf("buildACPWorkspace: %v", err)
|
||||
}
|
||||
if gotExclude {
|
||||
t.Fatal("trusted ACP workspace excluded project MCP configuration")
|
||||
}
|
||||
if !registered {
|
||||
t.Fatal("ACP workspace did not register configured MCP tools")
|
||||
}
|
||||
if _, ok := workspace.Registry.Get("mcp_docs_search"); !ok {
|
||||
t.Fatal("MCP tool is absent from ACP agent registry")
|
||||
}
|
||||
if _, ok := workspace.Registry.Get(tools.ToolSearchToolName); !ok {
|
||||
t.Fatal("ACP registry is missing tool_search for deferred MCP tools")
|
||||
}
|
||||
if workspace.DeferThreshold != 1 {
|
||||
t.Fatalf("workspace defer threshold = %d, want 1", workspace.DeferThreshold)
|
||||
}
|
||||
if err := workspace.Close(); err != nil {
|
||||
t.Fatalf("close ACP workspace: %v", err)
|
||||
}
|
||||
if !closed {
|
||||
t.Fatal("ACP workspace did not close its MCP runtime")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildACPWorkspacePreservesRedactedMCPSetupNotices(t *testing.T) {
|
||||
setTrustConfigRoot(t)
|
||||
workspaceRoot := t.TempDir()
|
||||
if err := os.MkdirAll(filepath.Join(workspaceRoot, ".zero"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
projectConfig := `{"mcp":{"servers":{"project-docs":{"type":"stdio","command":"project-docs"}}}}`
|
||||
if err := os.WriteFile(filepath.Join(workspaceRoot, ".zero", "config.json"), []byte(projectConfig), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
secret := "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
|
||||
grantStore, err := sandbox.NewGrantStore(sandbox.StoreOptions{FilePath: filepath.Join(t.TempDir(), "grants.json")})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
deps := fillAppDeps(appDeps{
|
||||
resolveMCPConfig: func(string, bool) (config.MCPConfig, error) {
|
||||
return config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"docs": {Type: "stdio", Command: "fake-docs"},
|
||||
}}, nil
|
||||
},
|
||||
newMCPStore: func() (*mcp.PermissionStore, error) { return nil, nil },
|
||||
newSandboxStore: func() (*sandbox.GrantStore, error) { return grantStore, nil },
|
||||
registerMCPTools: func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
|
||||
return fakeMCPRuntimeWithSkips{skipped: []mcp.SkippedServer{{
|
||||
Name: "docs", Err: errors.New("token=" + secret),
|
||||
}}}, nil
|
||||
},
|
||||
})
|
||||
|
||||
workspace, err := buildACPWorkspace(context.Background(), workspaceRoot, config.ResolvedConfig{}, agent.PermissionModeAuto, deps)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer workspace.Close()
|
||||
joined := strings.Join(workspace.Notices, "\n")
|
||||
if !strings.Contains(joined, "untrusted workspace") || !strings.Contains(joined, "MCP server docs unavailable") {
|
||||
t.Fatalf("workspace notices = %#v", workspace.Notices)
|
||||
}
|
||||
if strings.Contains(joined, secret) || !strings.Contains(joined, "[REDACTED]") {
|
||||
t.Fatalf("workspace notices leaked secret: %q", joined)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildACPWorkspaceSkipsMCPInPlanMode(t *testing.T) {
|
||||
grantStore, err := sandbox.NewGrantStore(sandbox.StoreOptions{FilePath: t.TempDir() + "/grants.json"})
|
||||
if err != nil {
|
||||
t.Fatalf("new grant store: %v", err)
|
||||
}
|
||||
called := false
|
||||
deps := fillAppDeps(appDeps{
|
||||
resolveMCPConfig: func(string, bool) (config.MCPConfig, error) {
|
||||
called = true
|
||||
return config.MCPConfig{}, nil
|
||||
},
|
||||
newSandboxStore: func() (*sandbox.GrantStore, error) { return grantStore, nil },
|
||||
registerMCPTools: func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
|
||||
called = true
|
||||
return nil, nil
|
||||
},
|
||||
})
|
||||
|
||||
workspace, err := buildACPWorkspace(context.Background(), t.TempDir(), config.ResolvedConfig{}, agent.PermissionModePlan, deps)
|
||||
if err != nil {
|
||||
t.Fatalf("buildACPWorkspace: %v", err)
|
||||
}
|
||||
if called {
|
||||
t.Fatal("plan-mode ACP workspace must not resolve or start MCP servers")
|
||||
}
|
||||
if err := workspace.Close(); err != nil {
|
||||
t.Fatalf("close plan workspace: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildACPWorkspaceClosesPartialMCPRuntimeOnRegistrationError(t *testing.T) {
|
||||
grantStore, err := sandbox.NewGrantStore(sandbox.StoreOptions{FilePath: t.TempDir() + "/grants.json"})
|
||||
if err != nil {
|
||||
t.Fatalf("new grant store: %v", err)
|
||||
}
|
||||
closed := false
|
||||
deps := fillAppDeps(appDeps{
|
||||
resolveMCPConfig: func(string, bool) (config.MCPConfig, error) {
|
||||
return config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"partial": {Type: "stdio", Command: "fake-partial"},
|
||||
}}, nil
|
||||
},
|
||||
newMCPStore: func() (*mcp.PermissionStore, error) { return nil, nil },
|
||||
newSandboxStore: func() (*sandbox.GrantStore, error) { return grantStore, nil },
|
||||
registerMCPTools: func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
|
||||
return closeFunc(func() error {
|
||||
closed = true
|
||||
return nil
|
||||
}), errors.New("second MCP server failed")
|
||||
},
|
||||
})
|
||||
|
||||
workspace, err := buildACPWorkspace(context.Background(), t.TempDir(), config.ResolvedConfig{}, agent.PermissionModeAuto, deps)
|
||||
if err == nil {
|
||||
t.Fatal("buildACPWorkspace succeeded after MCP registration error")
|
||||
}
|
||||
if workspace != nil {
|
||||
t.Fatal("buildACPWorkspace returned a workspace after MCP registration error")
|
||||
}
|
||||
if !closed {
|
||||
t.Fatal("partial MCP runtime was not closed after registration error")
|
||||
}
|
||||
}
|
||||
|
||||
type acpNotifyingReadCloser struct {
|
||||
io.ReadCloser
|
||||
readStarted chan<- struct{}
|
||||
|
||||
@@ -18,14 +18,20 @@ import (
|
||||
// cliFakeDeferredTool is deferred-eligible (implements Deferred() bool), mirroring
|
||||
// an MCP registry tool, so it counts toward the deferral threshold.
|
||||
type cliFakeDeferredTool struct {
|
||||
name string
|
||||
name string
|
||||
permission tools.Permission
|
||||
advertiseInAuto bool
|
||||
}
|
||||
|
||||
func (t cliFakeDeferredTool) Name() string { return t.name }
|
||||
func (t cliFakeDeferredTool) Description() string { return "fake deferred tool" }
|
||||
func (t cliFakeDeferredTool) Parameters() tools.Schema { return tools.Schema{Type: "object"} }
|
||||
func (t cliFakeDeferredTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectNetwork, Permission: tools.PermissionAllow}
|
||||
permission := t.permission
|
||||
if permission == "" {
|
||||
permission = tools.PermissionAllow
|
||||
}
|
||||
return tools.Safety{SideEffect: tools.SideEffectNetwork, Permission: permission, AdvertiseInAuto: t.advertiseInAuto}
|
||||
}
|
||||
func (t cliFakeDeferredTool) Run(context.Context, map[string]any) tools.Result {
|
||||
return tools.Result{Status: tools.StatusOK, Output: "ok"}
|
||||
|
||||
+12
-10
@@ -31,6 +31,14 @@ type mcpToolListItem struct {
|
||||
// into the one-line trust notice (mirroring the hooks and plugins chokepoints);
|
||||
// otherwise a workspace whose only project config is MCP would be gated silently.
|
||||
func registerMCPToolsForWorkspace(ctx context.Context, workspaceRoot string, registry *tools.Registry, deps appDeps, autonomy mcp.PermissionAutonomy, trustRoot string, runners ...*execution.Runner) (mcpToolRuntime, trustSkip, error) {
|
||||
options := mcp.RegisterOptions{}
|
||||
if len(runners) > 0 {
|
||||
options.Execution = runners[0]
|
||||
}
|
||||
return registerMCPToolsForWorkspaceWithOptions(ctx, workspaceRoot, registry, deps, autonomy, trustRoot, options)
|
||||
}
|
||||
|
||||
func registerMCPToolsForWorkspaceWithOptions(ctx context.Context, workspaceRoot string, registry *tools.Registry, deps appDeps, autonomy mcp.PermissionAutonomy, trustRoot string, options mcp.RegisterOptions) (mcpToolRuntime, trustSkip, error) {
|
||||
excludeProject, trustCheckErrored := resolveTrust(trustRoot)
|
||||
skip := trustSkip{
|
||||
excludedProjectConfig: excludeProject && projectMCPConfigExists(workspaceRoot),
|
||||
@@ -47,16 +55,10 @@ func registerMCPToolsForWorkspace(ctx context.Context, workspaceRoot string, reg
|
||||
if err != nil {
|
||||
return nil, skip, err
|
||||
}
|
||||
var runner *execution.Runner
|
||||
if len(runners) > 0 {
|
||||
runner = runners[0]
|
||||
}
|
||||
runtime, err := deps.registerMCPTools(ctx, registry, cfg, mcp.RegisterOptions{
|
||||
PermissionStore: store,
|
||||
Autonomy: autonomy,
|
||||
Execution: runner,
|
||||
WorkspaceRoot: workspaceRoot,
|
||||
})
|
||||
options.PermissionStore = store
|
||||
options.Autonomy = autonomy
|
||||
options.WorkspaceRoot = workspaceRoot
|
||||
runtime, err := deps.registerMCPTools(ctx, registry, cfg, options)
|
||||
return runtime, skip, err
|
||||
}
|
||||
|
||||
|
||||
@@ -23,6 +23,9 @@ const defaultConnectTimeout = 8 * time.Second
|
||||
type RegisterOptions struct {
|
||||
PermissionStore *PermissionStore
|
||||
Autonomy PermissionAutonomy
|
||||
// AdvertiseInAuto exposes prompt-gated MCP tools to model surfaces whose
|
||||
// default mode is auto while preserving the execution-time prompt.
|
||||
AdvertiseInAuto bool
|
||||
ClientFactory func(context.Context, Server) (ToolClient, error)
|
||||
// ConnectTimeout bounds the per-server connect+list at startup. Zero uses
|
||||
// defaultConnectTimeout.
|
||||
@@ -271,9 +274,10 @@ func newRegistryTool(server Server, remote RemoteTool, client ToolClient, option
|
||||
client: client,
|
||||
parameters: SchemaFromMCP(remote.InputSchema),
|
||||
safety: tools.Safety{
|
||||
SideEffect: tools.SideEffectNetwork,
|
||||
Permission: permission,
|
||||
Reason: fmt.Sprintf("MCP tool %s/%s runs through the configured %s server.", server.Name, remote.Name, server.Type),
|
||||
SideEffect: tools.SideEffectNetwork,
|
||||
Permission: permission,
|
||||
Reason: fmt.Sprintf("MCP tool %s/%s runs through the configured %s server.", server.Name, remote.Name, server.Type),
|
||||
AdvertiseInAuto: options.AdvertiseInAuto,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -342,6 +342,19 @@ func TestRegistryToolIsDeferredEligible(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistryToolCanAdvertiseInAutoWithoutPreapproval(t *testing.T) {
|
||||
tool := newRegistryTool(
|
||||
Server{Name: "docs", Type: "stdio"},
|
||||
RemoteTool{Name: "search"},
|
||||
&fakeToolClient{},
|
||||
RegisterOptions{Autonomy: AutonomyLow, AdvertiseInAuto: true},
|
||||
)
|
||||
safety := tool.Safety()
|
||||
if safety.Permission != tools.PermissionPrompt || !safety.AdvertiseInAuto {
|
||||
t.Fatalf("MCP safety = %#v", safety)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistryToolReportsMCPServerName verifies the registryTool reports its true
|
||||
// configured server name (not the sanitized tool-name token) so the deferred-tools
|
||||
// discovery label names a multi-token server correctly via tools.DeferredSource.
|
||||
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
"github.com/Gitlawb/zero/internal/providercatalog"
|
||||
"github.com/Gitlawb/zero/internal/providers"
|
||||
"github.com/Gitlawb/zero/internal/providers/providerio"
|
||||
"github.com/Gitlawb/zero/internal/proxydial"
|
||||
"github.com/Gitlawb/zero/internal/redaction"
|
||||
)
|
||||
|
||||
@@ -518,9 +519,15 @@ func newConnectivityClient(timeout time.Duration, resolver Resolver, sensitiveHe
|
||||
transport = &http.Transport{}
|
||||
}
|
||||
transport.DialContext = safeDialContext(resolver, allowLoopbackOrPrivate)
|
||||
// WRAPPED, SO THE DIALER LEARNS THE ROUTE THIS TRANSPORT CHOSE. The wrapper
|
||||
// records the proxy selected for each request on that request, and the
|
||||
// dialer reads it back on the dial that request causes; installing the bare
|
||||
// transport instead would leave a proxied dial looking like a loopback
|
||||
// target and fail the probe (#569). The cloned default transport already
|
||||
// reads HTTPS_PROXY and friends; nothing new is enabled.
|
||||
return &http.Client{
|
||||
Timeout: timeout,
|
||||
Transport: transport,
|
||||
Transport: proxydial.Wrap(transport),
|
||||
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
||||
if len(via) >= maxConnectivityRedirects {
|
||||
return fmt.Errorf("provider connectivity exceeded %d redirects", maxConnectivityRedirects)
|
||||
@@ -583,6 +590,17 @@ func sensitiveAuthHeaderNames(profile config.ProviderProfile, kind config.Provid
|
||||
func safeDialContext(resolver Resolver, allowLoopbackOrPrivate bool) func(context.Context, string, string) (net.Conn, error) {
|
||||
dialer := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}
|
||||
return func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
// THE PROXY IS NOT THE TARGET. With HTTPS_PROXY set to a local forward
|
||||
// proxy, the transport dials the proxy first and tunnels the request
|
||||
// through it, and this dialer used to refuse that dial as a loopback
|
||||
// address: "proxyconnect tcp: ... loopback hosts are blocked" for a
|
||||
// request whose real target had already been validated (#569). The
|
||||
// target validation is unchanged; only the dial the transport opens to
|
||||
// the proxy it picked for this very request is let past the guard, and a
|
||||
// direct dial is checked as it always was.
|
||||
if proxydial.IsProxyDial(ctx) {
|
||||
return dialer.DialContext(ctx, network, address)
|
||||
}
|
||||
host, port, err := net.SplitHostPort(address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -0,0 +1,218 @@
|
||||
package providerhealth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/proxydial"
|
||||
)
|
||||
|
||||
// connectivityClientWithProxy builds the real connectivity client and swaps in
|
||||
// a proxy function, which is how the transport learns its route. The
|
||||
// environment is left alone deliberately: http.ProxyFromEnvironment reads it
|
||||
// once per process, so an environment-driven test would depend on test order.
|
||||
func connectivityClientWithProxy(t *testing.T, resolver Resolver, proxyFor func(*http.Request) (*url.URL, error)) *http.Client {
|
||||
t.Helper()
|
||||
client := newConnectivityClient(3*time.Second, resolver, nil, false)
|
||||
wrapped, ok := client.Transport.(*proxydial.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("transport = %T, want the route-recording wrapper; without it no dial can be recognised as the proxy dial", client.Transport)
|
||||
}
|
||||
wrapped.Proxy = proxyFor
|
||||
return client
|
||||
}
|
||||
|
||||
// recordingProxy is a listener standing in for a local forward proxy. It
|
||||
// records the first line it receives and hangs up.
|
||||
func recordingProxy(t *testing.T) (*url.URL, <-chan string) {
|
||||
t.Helper()
|
||||
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = listener.Close() })
|
||||
arrived := make(chan string, 1)
|
||||
go func() {
|
||||
conn, err := listener.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
buffer := make([]byte, 256)
|
||||
n, _ := conn.Read(buffer)
|
||||
arrived <- string(buffer[:n])
|
||||
}()
|
||||
proxyURL, err := url.Parse("http://" + listener.Addr().String())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return proxyURL, arrived
|
||||
}
|
||||
|
||||
// THE PROXY IS NOT THE TARGET, AND THE GUARD HAS TO KNOW THE DIFFERENCE.
|
||||
//
|
||||
// With HTTPS_PROXY set to a local forward proxy, the transport dials the proxy
|
||||
// first and tunnels the request through it. This dialer refused that dial as a
|
||||
// loopback address, so `zero doctor --connectivity` failed with "proxyconnect
|
||||
// tcp: ... loopback hosts are blocked" for a target that had already been
|
||||
// validated and was public (#569). Driven through the real client, so the route
|
||||
// is the transport's own decision.
|
||||
func TestConnectivityClientTunnelsThroughTheConfiguredProxy(t *testing.T) {
|
||||
proxyURL, arrived := recordingProxy(t)
|
||||
client := connectivityClientWithProxy(t,
|
||||
staticResolver{err: errors.New("resolver must not be called for a proxied request")},
|
||||
func(*http.Request) (*url.URL, error) { return proxyURL, nil })
|
||||
|
||||
request, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://api.example.com/v1/models", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = client.Do(request)
|
||||
select {
|
||||
case first := <-arrived:
|
||||
if line := strings.SplitN(first, "\r\n", 2)[0]; line != "CONNECT api.example.com:443 HTTP/1.1" {
|
||||
t.Fatalf("proxy received %q, want a CONNECT for the validated target", line)
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatalf("the proxy never received a connection; request error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// AND A DIRECT REQUEST STAYS GUARDED, EVEN WHEN ITS ORIGIN IS THE PROXY'S OWN
|
||||
// ADDRESS.
|
||||
//
|
||||
// This is the NO_PROXY shape: the proxy function sends one host direct. An
|
||||
// exemption derived from asking that function about a request of our own
|
||||
// invention would answer "proxied" for the excluded host too, and this dial,
|
||||
// whose address equals the configured proxy endpoint, would skip the resolver
|
||||
// and reach whatever the name resolves to at connect time.
|
||||
func TestDirectRequestToTheProxyEndpointIsStillGuarded(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://rebind.example:80")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
lookups := 0
|
||||
resolver := countingResolver{
|
||||
onLookup: func() { lookups++ },
|
||||
addr: netip.MustParseAddr("10.0.0.5"),
|
||||
}
|
||||
client := connectivityClientWithProxy(t, resolver, func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Hostname() == "rebind.example" {
|
||||
return nil, nil
|
||||
}
|
||||
return proxyURL, nil
|
||||
})
|
||||
|
||||
request, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "http://rebind.example/data", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = client.Do(request)
|
||||
|
||||
if lookups == 0 {
|
||||
t.Fatal("the direct dial skipped the safety resolver because its address matched the configured proxy endpoint")
|
||||
}
|
||||
var safety endpointSafetyError
|
||||
if !errors.As(err, &safety) {
|
||||
t.Fatalf("err = %v, want the private answer to be refused", err)
|
||||
}
|
||||
}
|
||||
|
||||
// A proxy configured with a Unicode hostname is dialed by its IDNA-encoded
|
||||
// name, and it is still the proxy: the request must reach it rather than be
|
||||
// refused as an ordinary private address.
|
||||
func TestUnicodeProxyHostnameIsStillTheProxy(t *testing.T) {
|
||||
dialed := ""
|
||||
client := connectivityClientWithProxy(t,
|
||||
staticResolver{err: errors.New("resolver must not be called for a proxied request")},
|
||||
func(*http.Request) (*url.URL, error) { return url.Parse("http://büro.example:3067") })
|
||||
wrapped := client.Transport.(*proxydial.Transport)
|
||||
guarded := wrapped.DialContext
|
||||
wrapped.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
dialed = address
|
||||
return guarded(ctx, network, address)
|
||||
}
|
||||
|
||||
request, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://api.example.com/v1/models", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = client.Do(request)
|
||||
|
||||
if dialed != "xn--bro-hoa.example:3067" {
|
||||
t.Fatalf("SETUP INVALID: dialed %q, want the IDNA-encoded proxy endpoint", dialed)
|
||||
}
|
||||
var safety endpointSafetyError
|
||||
if errors.As(err, &safety) {
|
||||
t.Fatalf("the configured proxy was refused by the guard: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Without a proxy the guard is exactly what it was: loopback is refused.
|
||||
func TestSafeDialContextStillRefusesLoopbackWithoutAProxy(t *testing.T) {
|
||||
dial := safeDialContext(staticResolver{err: errors.New("resolver must not be called")}, false)
|
||||
conn, err := dial(context.Background(), "tcp", "127.0.0.1:3067")
|
||||
if conn != nil {
|
||||
_ = conn.Close()
|
||||
t.Fatal("loopback was dialed with no proxy configured")
|
||||
}
|
||||
var safety endpointSafetyError
|
||||
if !errors.As(err, &safety) {
|
||||
t.Fatalf("err = %v, want endpointSafetyError", err)
|
||||
}
|
||||
}
|
||||
|
||||
// A sibling loopback port is not the proxy, and the guard is unchanged for it.
|
||||
func TestSafeDialContextRefusesALoopbackPortThatIsNotTheProxy(t *testing.T) {
|
||||
proxyURL, arrived := recordingProxy(t)
|
||||
client := connectivityClientWithProxy(t,
|
||||
staticResolver{err: errors.New("resolver must not be called")},
|
||||
func(request *http.Request) (*url.URL, error) {
|
||||
// Only the https probe is proxied; the http one below is direct and
|
||||
// happens to point at the neighbouring port.
|
||||
if request.URL.Scheme == "https" {
|
||||
return proxyURL, nil
|
||||
}
|
||||
return nil, nil
|
||||
})
|
||||
|
||||
_, port, _ := net.SplitHostPort(proxyURL.Host)
|
||||
number, _ := strconv.Atoi(port)
|
||||
sibling := "http://" + net.JoinHostPort("127.0.0.1", strconv.Itoa(number+1)) + "/"
|
||||
request, err := http.NewRequestWithContext(context.Background(), http.MethodGet, sibling, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = client.Do(request)
|
||||
var safety endpointSafetyError
|
||||
if !errors.As(err, &safety) {
|
||||
t.Fatalf("err = %v, want a loopback port beside the proxy to stay blocked", err)
|
||||
}
|
||||
select {
|
||||
case first := <-arrived:
|
||||
t.Fatalf("the proxy received a connection it should not have: %q", first)
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// countingResolver answers with one address and reports each lookup, so a test
|
||||
// can tell a skipped guard from a guard that ran and allowed the address.
|
||||
type countingResolver struct {
|
||||
onLookup func()
|
||||
addr netip.Addr
|
||||
}
|
||||
|
||||
func (resolver countingResolver) LookupNetIP(context.Context, string, string) ([]netip.Addr, error) {
|
||||
if resolver.onLookup != nil {
|
||||
resolver.onLookup()
|
||||
}
|
||||
return []netip.Addr{resolver.addr}, nil
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
// Package proxydial answers one question for an SSRF-guarded dialer: is this
|
||||
// dial the connection dial of a request the transport routed through a forward
|
||||
// proxy?
|
||||
//
|
||||
// A guarded transport validates the TARGET of a request before dialing, and its
|
||||
// dialer refuses loopback, private and link-local addresses so a redirect or a
|
||||
// DNS rebind cannot walk the request onto the local network. When a forward
|
||||
// proxy is configured, the transport dials the PROXY first and tunnels the
|
||||
// request through it, and a local proxy is by far the common shape:
|
||||
// HTTPS_PROXY=http://127.0.0.1:3067. The dialer then sees 127.0.0.1, refuses it,
|
||||
// and the user gets "loopback hosts are blocked" for a request whose real target
|
||||
// was validated and public (#569).
|
||||
//
|
||||
// The proxy address is not the request target, and the guard was never about it.
|
||||
// It is configuration the user set on purpose, and the transport is going to dial
|
||||
// it for every request whether the guard likes it or not. So that dial is let
|
||||
// through, and only that dial.
|
||||
//
|
||||
// THE EXEMPTION FOLLOWS THE TRANSPORT'S OWN ROUTE DECISION, NOT A GUESS AT IT.
|
||||
// Asking the Proxy function a question of one's own invention answers for a
|
||||
// request the transport never made: with NO_PROXY excluding a host, the real
|
||||
// request to that host goes direct while a synthetic probe still comes back
|
||||
// proxied, and a direct dial whose address happens to equal the proxy endpoint
|
||||
// would then skip the guard entirely. Host-dependent proxy callbacks make the
|
||||
// mismatch general, and no amount of care over the probe's spelling closes it,
|
||||
// because the two calls are about different requests.
|
||||
//
|
||||
// So the decision is recorded where it is made, on the request, and read back on
|
||||
// the dial that request causes. Transport.RoundTrip asks the wrapped transport's
|
||||
// Proxy function about the actual request, stores the answer in that request's
|
||||
// context, and hands it on; the transport derives the dial context from the
|
||||
// request context, so the dialer reads back the route this very connection is
|
||||
// being opened for. Per-request, so concurrent requests cannot overwrite each
|
||||
// other, and unforgeable from outside the package: the context key is private.
|
||||
package proxydial
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/url"
|
||||
)
|
||||
|
||||
// selectedProxyKey addresses the route decision inside a request context. It is
|
||||
// unexported so nothing outside this package can plant one.
|
||||
type selectedProxyKey struct{}
|
||||
|
||||
// Transport records the proxy its wrapped transport selects for each request in
|
||||
// that request's context, so a guarded DialContext installed on the same
|
||||
// transport can tell the dial to a proxy from a direct dial.
|
||||
//
|
||||
// The embedded transport stays reachable: callers configure Proxy, DialContext
|
||||
// and the rest on it exactly as before, and tests can read them back.
|
||||
type Transport struct {
|
||||
*http.Transport
|
||||
}
|
||||
|
||||
// Wrap returns a Transport around transport. The result is the http.RoundTripper
|
||||
// to install on the client; installing the bare transport instead leaves every
|
||||
// dial looking direct, which is safe (the guard applies) but reopens #569.
|
||||
func Wrap(transport *http.Transport) *Transport {
|
||||
return &Transport{Transport: transport}
|
||||
}
|
||||
|
||||
// RoundTrip resolves the route for this request and passes the decision down to
|
||||
// the dialer through the request context.
|
||||
func (transport *Transport) RoundTrip(request *http.Request) (*http.Response, error) {
|
||||
if transport.Transport.Proxy != nil {
|
||||
// A Proxy error is left for the transport to report: it fails the
|
||||
// request before dialing, so there is no dial to classify. The function
|
||||
// is called once here and once by the transport, which is what a
|
||||
// host-dependent callback already expects from redirects and retries.
|
||||
if proxyURL, err := transport.Transport.Proxy(request); err == nil && proxyURL != nil {
|
||||
request = request.WithContext(context.WithValue(request.Context(), selectedProxyKey{}, proxyURL))
|
||||
}
|
||||
}
|
||||
return transport.Transport.RoundTrip(request)
|
||||
}
|
||||
|
||||
// IsProxyDial reports whether ctx belongs to a request this package's Transport
|
||||
// routed through a proxy, which makes the dial about to happen the dial to that
|
||||
// proxy.
|
||||
//
|
||||
// THE ADDRESS IS NOT COMPARED, AND DOES NOT NEED TO BE. A transport opens
|
||||
// exactly one connection dial per connection, and when a proxy is selected that
|
||||
// dial goes to the proxy; the target is reached through it, never dialed. So the
|
||||
// route decision on the context already identifies the dial, and matching the
|
||||
// address against the proxy URL on top of it could only ever disagree with the
|
||||
// transport, which is what a punycoded Unicode proxy hostname does: Go dials
|
||||
// xn--bro-hoa.example while the URL still spells buero as the user wrote it.
|
||||
func IsProxyDial(ctx context.Context) bool {
|
||||
proxyURL, _ := ctx.Value(selectedProxyKey{}).(*url.URL)
|
||||
return proxyURL != nil
|
||||
}
|
||||
@@ -0,0 +1,224 @@
|
||||
package proxydial
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// errDialed ends the request at the dial, after the route decision has been
|
||||
// observed. Nothing here needs a live server.
|
||||
var errDialed = errors.New("dial recorded")
|
||||
|
||||
type dialRecord struct {
|
||||
address string
|
||||
proxied bool
|
||||
}
|
||||
|
||||
// driveRequest sends one request through a wrapped transport whose dialer
|
||||
// records what it was asked for and whether the context says this dial belongs
|
||||
// to a proxied request.
|
||||
func driveRequest(t *testing.T, target string, proxyFor func(*http.Request) (*url.URL, error)) dialRecord {
|
||||
t.Helper()
|
||||
var record dialRecord
|
||||
transport := &http.Transport{Proxy: proxyFor}
|
||||
transport.DialContext = func(ctx context.Context, _, address string) (net.Conn, error) {
|
||||
record = dialRecord{address: address, proxied: IsProxyDial(ctx)}
|
||||
return nil, errDialed
|
||||
}
|
||||
client := &http.Client{Transport: Wrap(transport)}
|
||||
request, err := http.NewRequest(http.MethodGet, target, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := client.Do(request); !errors.Is(err, errDialed) {
|
||||
t.Fatalf("%s: request ended before the dial: %v", target, err)
|
||||
}
|
||||
if record.address == "" {
|
||||
t.Fatalf("%s: the dialer was never called", target)
|
||||
}
|
||||
return record
|
||||
}
|
||||
|
||||
func fixedProxy(raw string) func(*http.Request) (*url.URL, error) {
|
||||
parsed, err := url.Parse(raw)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return func(*http.Request) (*url.URL, error) { return parsed, nil }
|
||||
}
|
||||
|
||||
// THE DIAL TO A CONFIGURED PROXY IS THE ONE THAT IS EXEMPT. A proxied request
|
||||
// dials the proxy, never the target, so the route on the context is what
|
||||
// identifies the dial.
|
||||
func TestProxiedRequestDialsTheProxyAndSaysSo(t *testing.T) {
|
||||
record := driveRequest(t, "https://api.example.com/v1/models", fixedProxy("http://127.0.0.1:3067"))
|
||||
if !record.proxied {
|
||||
t.Error("a request routed through a proxy did not reach the dialer as proxied")
|
||||
}
|
||||
if record.address != "127.0.0.1:3067" {
|
||||
t.Errorf("dial address = %q, want the proxy endpoint", record.address)
|
||||
}
|
||||
}
|
||||
|
||||
// AND A DIRECT REQUEST IS NOT, EVEN WHEN ITS ORIGIN IS THE PROXY'S OWN ADDRESS.
|
||||
//
|
||||
// This is the NO_PROXY shape: the proxy function excludes one host, so a request
|
||||
// to that host goes direct. Asking that same function about a request of our own
|
||||
// invention would answer "proxied" (the invented host is not excluded) and the
|
||||
// direct dial, whose address equals the proxy endpoint, would skip the guard
|
||||
// while resolving to whatever the name resolves to at connect time.
|
||||
func TestDirectRequestToTheProxyEndpointIsNotAProxyDial(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://rebind.example:80")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
noProxyForRebind := func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Hostname() == "rebind.example" {
|
||||
return nil, nil
|
||||
}
|
||||
return proxyURL, nil
|
||||
}
|
||||
|
||||
record := driveRequest(t, "http://rebind.example/data", noProxyForRebind)
|
||||
if record.address != "rebind.example:80" {
|
||||
t.Fatalf("SETUP INVALID: dial address = %q, want the direct origin, which is also the proxy endpoint", record.address)
|
||||
}
|
||||
if record.proxied {
|
||||
t.Error("a direct dial was reported as the proxy dial because its address matched the configured proxy")
|
||||
}
|
||||
|
||||
// The same function still routes another host through that proxy, so the
|
||||
// exclusion above is a route difference and not a disabled proxy.
|
||||
if record := driveRequest(t, "http://elsewhere.example/data", noProxyForRebind); !record.proxied {
|
||||
t.Error("the configured proxy was not used for a host outside the exclusion")
|
||||
}
|
||||
}
|
||||
|
||||
// The inverse mismatch: a callback that routes by host must not have its answer
|
||||
// taken from some other host. A proxy selected only for the real target is
|
||||
// honoured.
|
||||
func TestHostDependentProxySelectionIsHonoured(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://127.0.0.1:3067")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
onlyForOneHost := func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Hostname() == "api.example.com" {
|
||||
return proxyURL, nil
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
if record := driveRequest(t, "https://api.example.com/v1", onlyForOneHost); !record.proxied || record.address != "127.0.0.1:3067" {
|
||||
t.Errorf("selected proxy not reached: %+v", record)
|
||||
}
|
||||
if record := driveRequest(t, "https://other.example/v1", onlyForOneHost); record.proxied {
|
||||
t.Error("an unproxied host was reported as proxied")
|
||||
}
|
||||
}
|
||||
|
||||
// Scheme separation: an https-only proxy leaves an http request direct, and the
|
||||
// decision recorded is the one made for that request's own scheme.
|
||||
func TestSchemeSeparation(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://127.0.0.1:3067")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
httpsOnly := func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Scheme == "https" {
|
||||
return proxyURL, nil
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
if record := driveRequest(t, "https://api.example.com/v1", httpsOnly); !record.proxied {
|
||||
t.Error("the https proxy was not applied to an https request")
|
||||
}
|
||||
if record := driveRequest(t, "http://api.example.com/v1", httpsOnly); record.proxied {
|
||||
t.Error("an http request was treated as proxied by an https-only proxy")
|
||||
}
|
||||
}
|
||||
|
||||
// A UNICODE PROXY HOSTNAME IS STILL THE PROXY. Go dials proxies by their
|
||||
// IDNA-encoded name, so a matcher comparing the URL's own spelling would call
|
||||
// this dial direct and refuse it as an ordinary private address. Nothing is
|
||||
// compared here, so both spellings of one host are one host.
|
||||
func TestUnicodeProxyHostnameIsRecognised(t *testing.T) {
|
||||
unicodeRecord := driveRequest(t, "https://api.example.com/v1", fixedProxy("http://büro.example:3067"))
|
||||
if !unicodeRecord.proxied {
|
||||
t.Error("a proxy configured with a Unicode hostname was not recognised as the proxy")
|
||||
}
|
||||
asciiRecord := driveRequest(t, "https://api.example.com/v1", fixedProxy("http://xn--bro-hoa.example:3067"))
|
||||
if !asciiRecord.proxied {
|
||||
t.Error("a proxy configured with the ASCII spelling was not recognised as the proxy")
|
||||
}
|
||||
// The two spellings name one endpoint, and it is the ASCII one that is
|
||||
// dialed. This is the mismatch a string comparison would have had to close.
|
||||
if unicodeRecord.address != asciiRecord.address {
|
||||
t.Errorf("dial addresses differ across spellings of one host: %q vs %q", unicodeRecord.address, asciiRecord.address)
|
||||
}
|
||||
}
|
||||
|
||||
// A transport with no proxy leaves every dial direct.
|
||||
func TestNoProxyMeansNoExemption(t *testing.T) {
|
||||
if record := driveRequest(t, "http://127.0.0.1:3067/", nil); record.proxied {
|
||||
t.Error("a dial was reported as proxied by a transport with no Proxy function")
|
||||
}
|
||||
if IsProxyDial(context.Background()) {
|
||||
t.Error("a context that never went through the wrapper reported a proxy dial")
|
||||
}
|
||||
}
|
||||
|
||||
// THE DECISION IS PER REQUEST. A shared record of the last selected proxy would
|
||||
// let a proxied request hand its exemption to a direct one running beside it.
|
||||
func TestConcurrentRequestsKeepTheirOwnRoute(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://127.0.0.1:3067")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
byHost := func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Hostname() == "proxied.example" {
|
||||
return proxyURL, nil
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var mutex sync.Mutex
|
||||
seen := map[string]bool{}
|
||||
transport := &http.Transport{Proxy: byHost}
|
||||
transport.DialContext = func(ctx context.Context, _, address string) (net.Conn, error) {
|
||||
mutex.Lock()
|
||||
seen[address] = IsProxyDial(ctx)
|
||||
mutex.Unlock()
|
||||
return nil, errDialed
|
||||
}
|
||||
client := &http.Client{Transport: Wrap(transport)}
|
||||
|
||||
var group sync.WaitGroup
|
||||
for range 40 {
|
||||
for _, target := range []string{"http://proxied.example/", "http://direct.example/"} {
|
||||
group.Add(1)
|
||||
go func() {
|
||||
defer group.Done()
|
||||
request, _ := http.NewRequest(http.MethodGet, target, nil)
|
||||
_, _ = client.Do(request)
|
||||
}()
|
||||
}
|
||||
}
|
||||
group.Wait()
|
||||
|
||||
mutex.Lock()
|
||||
defer mutex.Unlock()
|
||||
if !seen["127.0.0.1:3067"] {
|
||||
t.Error("the proxy dial was not recorded as proxied under concurrency")
|
||||
}
|
||||
if seen["direct.example:80"] {
|
||||
t.Error("a direct dial picked up another request's proxy exemption")
|
||||
}
|
||||
if len(seen) != 2 {
|
||||
t.Fatalf("SETUP INVALID: dialed %v, want exactly the proxy and the direct origin", seen)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,463 @@
|
||||
package sessions
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// A CALL'S PAYLOAD MUST NOT RIDE INTO A LATER PROMPT ANY MORE THAN A RESULT'S.
|
||||
//
|
||||
// toolResultOutcome drops result bodies, and calls were left verbatim, so an
|
||||
// interrupted write_file put its content into the next turn's prompt while the
|
||||
// matching result body did not. Same secret, one door left open. The identity
|
||||
// still has to survive, since a resumed turn knowing WHICH file was written is
|
||||
// the whole point of admitting calls at all.
|
||||
//
|
||||
// Pinned as whole lines, compared as fields, for the same reason the result test
|
||||
// is: a substring search for one fixture secret passes when a truncated or
|
||||
// reworded copy leaks, and the renderer walks a map so field order varies.
|
||||
func TestResumePromptCarriesToolCallIdentityWithoutPayload(t *testing.T) {
|
||||
const secret = "AWS_SECRET_ACCESS_KEY=wJalrXUtnFEMI/K7MDENG"
|
||||
const token = "Bearer sk-live-4eC39HqLyjWDarjtT1zdp7dc"
|
||||
events := []Event{
|
||||
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "rotate the deploy key"})},
|
||||
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c1", "name": "write_file", "arguments": `{"path":"deploy/prod.env","content":"` + secret + `"}`})},
|
||||
{Sequence: 3, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "write_file", "status": "ok", "output": "written"})},
|
||||
{Sequence: 4, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c2", "name": "edit_file", "arguments": `{"path":"deploy/prod.env","old_string":"` + secret + `","new_string":"rotated"}`})},
|
||||
{Sequence: 5, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "edit_file", "status": "ok"})},
|
||||
{Sequence: 6, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c3", "name": "exec_command", "arguments": `{"cmd":"curl -H '` + token + `' https://api.example/rotate","workdir":"deploy"}`})},
|
||||
{Sequence: 7, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "exec_command", "status": "ok"})},
|
||||
{Sequence: 8, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c4", "name": "apply_patch", "arguments": `{"patch":"*** Begin Patch\n*** Update File: deploy/prod.env\n-` + secret + `\n+rotated\n*** End Patch"}`})},
|
||||
{Sequence: 9, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "apply_patch", "status": "ok"})},
|
||||
{Sequence: 10, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c5", "name": "read_file", "arguments": `{"path":"deploy/prod.env","offset":1,"limit":40}`})},
|
||||
{Sequence: 11, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "read_file", "status": "ok", "output": secret})},
|
||||
{Sequence: 12, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error: upstream timeout"})},
|
||||
}
|
||||
out := resumePrompt(t, events)
|
||||
|
||||
// The payloads: none of them, in any form.
|
||||
for _, leaked := range []string{secret, secret[:16], token, token[:14], "rotated", "Begin Patch", "curl -H"} {
|
||||
if strings.Contains(out, leaked) {
|
||||
t.Errorf("tool call payload %q reached the resume prompt:\n%s", leaked, out)
|
||||
}
|
||||
}
|
||||
|
||||
// The identities: every one of them, from the call side.
|
||||
want := map[int]struct {
|
||||
kind string
|
||||
fields []string
|
||||
}{
|
||||
2: {"tool_call", []string{"c1", "write_file", `{"path":"deploy/prod.env"}`}},
|
||||
4: {"tool_call", []string{"c2", "edit_file", `{"path":"deploy/prod.env"}`}},
|
||||
6: {"tool_call", []string{"c3", "exec_command", `{"workdir":"deploy"}`}},
|
||||
8: {"tool_call", []string{"c4", "apply_patch"}},
|
||||
10: {"tool_call", []string{"c5", "read_file"}},
|
||||
}
|
||||
seen := map[int]bool{}
|
||||
for _, line := range strings.Split(out, "\n") {
|
||||
if !strings.HasPrefix(line, "- #") || !strings.Contains(line, "tool_call:") {
|
||||
continue
|
||||
}
|
||||
var sequence int
|
||||
var kind string
|
||||
rest, _ := strings.CutPrefix(line, "- #")
|
||||
if _, err := fmt.Sscanf(rest, "%d %s", &sequence, &kind); err != nil {
|
||||
t.Fatalf("unparsable context line %q: %v", line, err)
|
||||
}
|
||||
expected, wanted := want[sequence]
|
||||
if !wanted {
|
||||
t.Errorf("unexpected tool_call line %q", line)
|
||||
continue
|
||||
}
|
||||
seen[sequence] = true
|
||||
_, payload, _ := strings.Cut(rest, ": ")
|
||||
fields := strings.Fields(payload)
|
||||
for _, field := range expected.fields {
|
||||
if !containsField(fields, field) {
|
||||
t.Errorf("line #%d lost identity field %q: %q", sequence, field, line)
|
||||
}
|
||||
}
|
||||
}
|
||||
for sequence := range want {
|
||||
if !seen[sequence] {
|
||||
t.Errorf("tool_call #%d did not reach the resume prompt at all, so the interrupted work is invisible again", sequence)
|
||||
}
|
||||
}
|
||||
// read_file keeps its window too, or a resumed turn re-reads a part it had.
|
||||
if !strings.Contains(out, `"offset":1`) || !strings.Contains(out, `"limit":40`) {
|
||||
t.Errorf("read_file lost its offset/limit window:\n%s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func containsField(fields []string, want string) bool {
|
||||
for _, field := range fields {
|
||||
if field == want {
|
||||
return true
|
||||
}
|
||||
}
|
||||
// The reduced arguments object may carry more than one identity key, in map
|
||||
// order, so a single-key expectation is also satisfied by an object that
|
||||
// contains that key.
|
||||
if strings.HasPrefix(want, "{") {
|
||||
key := strings.TrimSuffix(strings.TrimPrefix(want, "{"), "}")
|
||||
for _, field := range fields {
|
||||
if strings.HasPrefix(field, "{") && strings.Contains(field, key) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ARGUMENTS THAT CANNOT BE READ ARE NOT PASSED THROUGH AS TEXT.
|
||||
//
|
||||
// An allow-list only protects what it can see. Arguments that do not decode as
|
||||
// an object are dropped rather than rendered, since text this file could not
|
||||
// inspect is text it cannot vouch for.
|
||||
func TestResumePromptDropsUnreadableToolCallArguments(t *testing.T) {
|
||||
const secret = "sk-live-unparseable-9f8e7d6c"
|
||||
events := []Event{
|
||||
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "go"})},
|
||||
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c1", "name": "bash", "arguments": `not json at all ` + secret})},
|
||||
{Sequence: 3, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error"})},
|
||||
}
|
||||
out := resumePrompt(t, events)
|
||||
if strings.Contains(out, secret) {
|
||||
t.Errorf("unparseable arguments were rendered as text:\n%s", out)
|
||||
}
|
||||
if !strings.Contains(out, "bash") {
|
||||
t.Errorf("the call's identity was lost along with its unreadable arguments:\n%s", out)
|
||||
}
|
||||
}
|
||||
|
||||
// AND A TOOL THIS FILE HAS NEVER HEARD OF STILL KEEPS ITS PATH.
|
||||
//
|
||||
// The allow-list is by key, not by tool name, so an MCP tool or a future core
|
||||
// tool whose argument is a path or a url carries it into the resume prompt
|
||||
// without being enumerated here, while any body field it has is dropped.
|
||||
func TestResumePromptKeepsIdentityForUnknownTools(t *testing.T) {
|
||||
const body = "large opaque payload that must not be replayed"
|
||||
events := []Event{
|
||||
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "go"})},
|
||||
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c1", "name": "mcp_uploader", "arguments": `{"url":"https://files.example/x","blob":"` + body + `"}`})},
|
||||
{Sequence: 3, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error"})},
|
||||
}
|
||||
out := resumePrompt(t, events)
|
||||
if strings.Contains(out, body) {
|
||||
t.Errorf("an unlisted body field was replayed:\n%s", out)
|
||||
}
|
||||
if !strings.Contains(out, "https://files.example/x") {
|
||||
t.Errorf("an unlisted tool lost its url identity:\n%s", out)
|
||||
}
|
||||
}
|
||||
|
||||
// AND AN ALLOW-LISTED KEY IS NOT A SAFE VALUE.
|
||||
//
|
||||
// web_fetch accepts a credential in the query string and redacts the URL it
|
||||
// reports back, so an interrupted fetch had its token dropped from the result
|
||||
// and kept verbatim in the call. This projection is what admits the call into a
|
||||
// later turn's prompt, so the token rode into the next resume or fork. The
|
||||
// identity the resumed turn actually needs is the host and path, and those
|
||||
// survive.
|
||||
func TestResumePromptRedactsCredentialsInsideRetainedValues(t *testing.T) {
|
||||
const token = "secret-value-9f8e7d6c5b4a"
|
||||
events := []Event{
|
||||
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "check the feed"})},
|
||||
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c1", "name": "web_fetch", "arguments": `{"url":"https://api.example/data?access_token=` + token + `&page=2"}`})},
|
||||
{Sequence: 3, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c2", "name": "web_fetch", "arguments": `{"url":"https://reader:[email protected]/feed"}`})},
|
||||
{Sequence: 4, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error: upstream timeout"})},
|
||||
}
|
||||
out := resumePrompt(t, events)
|
||||
|
||||
for _, leaked := range []string{token, token[:12], "hunter2"} {
|
||||
if strings.Contains(out, leaked) {
|
||||
t.Errorf("credential %q reached the resume prompt:\n%s", leaked, out)
|
||||
}
|
||||
}
|
||||
// The non-secret identity is the whole reason calls are admitted at all.
|
||||
for _, kept := range []string{"api.example", "/data", "page=2", "/feed"} {
|
||||
if !strings.Contains(out, kept) {
|
||||
t.Errorf("URL identity %q was lost to the scrub:\n%s", kept, out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The scrub must not eat ordinary identities. A path, a glob and a query are
|
||||
// what the resumed turn navigates by, and none of them is a credential.
|
||||
func TestResumePromptKeepsOrdinaryIdentitiesThroughTheScrub(t *testing.T) {
|
||||
events := []Event{
|
||||
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "go"})},
|
||||
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c1", "name": "read_file", "arguments": `{"path":"internal/sessions/exec_session.go","offset":1,"limit":40}`})},
|
||||
{Sequence: 3, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c2", "name": "grep", "arguments": `{"pattern":"func toolCallIdentity","glob":"src/**/*.tsx"}`})},
|
||||
{Sequence: 4, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error"})},
|
||||
}
|
||||
out := resumePrompt(t, events)
|
||||
for _, kept := range []string{"internal/sessions/exec_session.go", "func toolCallIdentity", "src/**/*.tsx"} {
|
||||
if !strings.Contains(out, kept) {
|
||||
t.Errorf("identity %q was lost to the scrub:\n%s", kept, out)
|
||||
}
|
||||
}
|
||||
// Numbers are not strings and are nothing to scrub; they must survive whole.
|
||||
if !strings.Contains(out, "40") {
|
||||
t.Errorf("a numeric read window did not survive the scrub:\n%s", out)
|
||||
}
|
||||
}
|
||||
|
||||
// AND A FIELD THIS FILE HAS NEVER HEARD OF DOES NOT RIDE ALONG.
|
||||
//
|
||||
// toolResultOutcome builds its output from the fields it keeps, so an unknown
|
||||
// one cannot reach a prompt through it. The call side edited arguments in place
|
||||
// and returned every sibling key, which is the opposite contract: a producer
|
||||
// recording one more top-level field, or the same payload with no arguments at
|
||||
// all, put that field straight into the next turn. The projection names what
|
||||
// survives, so both halves now drop what they do not name.
|
||||
func TestResumePromptProjectsTopLevelToolCallFields(t *testing.T) {
|
||||
const body = "opaque top-level payload that must not be replayed"
|
||||
events := []Event{
|
||||
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "go"})},
|
||||
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
|
||||
"id": "c1", "name": "write_file", "arguments": `{"path":"deploy/prod.env"}`, "rawInput": body,
|
||||
})},
|
||||
// The same shape with no arguments at all, which used to return the whole
|
||||
// payload untouched.
|
||||
{Sequence: 3, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
|
||||
"id": "c2", "name": "read_file", "stashedContent": body,
|
||||
})},
|
||||
{Sequence: 4, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error"})},
|
||||
}
|
||||
out := resumePrompt(t, events)
|
||||
|
||||
for _, leaked := range []string{body, "rawInput", "stashedContent"} {
|
||||
if strings.Contains(out, leaked) {
|
||||
t.Errorf("an unnamed top-level field %q reached the resume prompt:\n%s", leaked, out)
|
||||
}
|
||||
}
|
||||
// The identity itself still survives, from both calls.
|
||||
for _, kept := range []string{"c1", "write_file", "deploy/prod.env", "c2", "read_file"} {
|
||||
if !strings.Contains(out, kept) {
|
||||
t.Errorf("call identity %q was lost to the projection:\n%s", kept, out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// AND A PERMITTED KEY IS NOT PERMISSION FOR WHATEVER HANGS UNDER IT.
|
||||
//
|
||||
// The projection scrubbed strings and passed everything else through, which
|
||||
// reads as "a number is nothing to scrub" and is true of a number and not of an
|
||||
// object. MCP schemas allow object and array properties, so a whole payload can
|
||||
// arrive under a permitted key. Redacting a serialized object would not help:
|
||||
// ordinary private text has no credential shape to match.
|
||||
func TestResumePromptDropsNestedArgumentContainers(t *testing.T) {
|
||||
const body = "private body text with no credential shape at all"
|
||||
const token = "secret-value-9f8e7d6c5b4a"
|
||||
events := []Event{
|
||||
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "go"})},
|
||||
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
|
||||
"id": "c1", "name": "mcp_search",
|
||||
"arguments": `{"query":{"api_key":"` + token + `","content":"` + body + `"}}`,
|
||||
})},
|
||||
{Sequence: 3, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
|
||||
"id": "c2", "name": "web_fetch",
|
||||
"arguments": `{"url":["https://api.example/data?access_token=` + token + `"]}`,
|
||||
})},
|
||||
{Sequence: 4, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error"})},
|
||||
}
|
||||
out := resumePrompt(t, events)
|
||||
|
||||
for _, leaked := range []string{body, body[:20], token, token[:12], "api_key"} {
|
||||
if strings.Contains(out, leaked) {
|
||||
t.Errorf("nested argument content %q reached the resume prompt:\n%s", leaked, out)
|
||||
}
|
||||
}
|
||||
// The call still says which tool ran, so a resumed turn can ask again.
|
||||
for _, kept := range []string{"c1", "mcp_search", "c2", "web_fetch"} {
|
||||
if !strings.Contains(out, kept) {
|
||||
t.Errorf("call identity %q was lost along with the container:\n%s", kept, out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// THE SEARCH ALIASES ARE TOOL-SPECIFIC, AND ONE OF THEM IS A TRAP.
|
||||
//
|
||||
// grep accepts `search` for its pattern and edit_file accepts the same word for
|
||||
// the text being replaced. Admitting it globally for the sake of the first would
|
||||
// replay file contents through the second.
|
||||
func TestResumePromptKeepsGrepSearchAndDropsEditSearch(t *testing.T) {
|
||||
const oldText = "OLD FILE CONTENTS THAT MUST NOT BE REPLAYED"
|
||||
events := []Event{
|
||||
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "go"})},
|
||||
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
|
||||
"id": "c1", "name": "grep", "arguments": `{"search":"func toolCallIdentity","path":"internal/sessions"}`,
|
||||
})},
|
||||
{Sequence: 3, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
|
||||
"id": "c2", "name": "edit_file", "arguments": `{"path":"a.go","search":"` + oldText + `","replace":"new"}`,
|
||||
})},
|
||||
{Sequence: 4, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error"})},
|
||||
}
|
||||
out := resumePrompt(t, events)
|
||||
|
||||
if !strings.Contains(out, "func toolCallIdentity") {
|
||||
t.Errorf("grep lost the pattern it was searching for:\n%s", out)
|
||||
}
|
||||
if strings.Contains(out, oldText) || strings.Contains(out, oldText[:20]) {
|
||||
t.Errorf("edit_file replayed the text it was replacing:\n%s", out)
|
||||
}
|
||||
// Both calls keep the path, which is identity for either tool.
|
||||
for _, kept := range []string{"internal/sessions", "a.go"} {
|
||||
if !strings.Contains(out, kept) {
|
||||
t.Errorf("path identity %q was lost:\n%s", kept, out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// AND THE SUPPORTED SEARCH AND WINDOW FORMS SURVIVE, not only the canonical
|
||||
// ones. A resumed turn that knows the path but not which slice was asked for
|
||||
// has half the identity and will read again.
|
||||
func TestResumePromptKeepsSupportedAliasesAndWindows(t *testing.T) {
|
||||
events := []Event{
|
||||
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "go"})},
|
||||
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
|
||||
"id": "c1", "name": "glob", "arguments": `{"match":"src/**/*.go"}`,
|
||||
})},
|
||||
{Sequence: 3, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
|
||||
"id": "c2", "name": "read_file", "arguments": `{"path":"large.go","byte_offset":4096,"byte_limit":512}`,
|
||||
})},
|
||||
{Sequence: 4, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
|
||||
"id": "c3", "name": "read_file", "arguments": `{"path":"big.go","start_line":40,"end_line":80,"max_lines":41}`,
|
||||
})},
|
||||
{Sequence: 5, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error"})},
|
||||
}
|
||||
out := resumePrompt(t, events)
|
||||
|
||||
// KEY AND VALUE TOGETHER, NOT THE NUMBER ALONE. Searching for "40" finds
|
||||
// it inside "4096", so the start_line assertion passed whether or not
|
||||
// start_line survived at all, which is the opposite of what it claims.
|
||||
for _, kept := range []string{
|
||||
"src/**/*.go",
|
||||
`"byte_offset":4096`, `"byte_limit":512`,
|
||||
`"start_line":40`, `"end_line":80`, `"max_lines":41`,
|
||||
} {
|
||||
if !strings.Contains(out, kept) {
|
||||
t.Errorf("supported identity or window %s was lost:\n%s", kept, out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// THE SHARED TABLE MUST NEVER NAME A MUTATING TOOL'S BODY ARGUMENT.
|
||||
//
|
||||
// This is the rule the grep/edit_file collision taught, kept as a test so the
|
||||
// next alias added for a search tool cannot quietly open the same door. The
|
||||
// aliases below are the ones the mutating tools accept for the content they
|
||||
// write; they belong in the per-tool table or nowhere.
|
||||
func TestIdentityKeysNeverNameAMutatingToolBodyArgument(t *testing.T) {
|
||||
bodyAliases := map[string][]string{
|
||||
"edit_file": {"old_string", "old", "search", "find", "old_str", "new_string", "new", "replace", "replacement", "new_str"},
|
||||
"write_file": {"content", "contents", "text", "body", "data", "file_content"},
|
||||
"apply_patch": {"patch"},
|
||||
}
|
||||
for tool, aliases := range bodyAliases {
|
||||
for _, alias := range aliases {
|
||||
if toolCallIdentityKeys[alias] {
|
||||
t.Errorf("%s accepts %q for content it writes, and the shared identity table admits it: a resumed prompt would replay the body", tool, alias)
|
||||
}
|
||||
}
|
||||
}
|
||||
// The per-tool table is where an ambiguous name belongs, and grep's use of
|
||||
// `search` is the case that proves the mechanism is wired.
|
||||
if !toolCallIdentityKeysByTool["grep"]["search"] {
|
||||
t.Error("grep lost its search alias, so an interrupted grep resumes without the pattern")
|
||||
}
|
||||
}
|
||||
|
||||
// THROUGH THE STORE, NOT JUST THE RENDERER.
|
||||
//
|
||||
// Every test above hands FormatExecPrompt an events slice it built in memory,
|
||||
// so none of them shows the projection still applies to events that were
|
||||
// written, reloaded and prepared by the real resume and fork paths. That is
|
||||
// where it has to hold: the leak this guards against is a later turn reading a
|
||||
// recorded call, and the recording is what makes it later.
|
||||
//
|
||||
// Includes a call the tool would have REJECTED. The agent records OnToolCall
|
||||
// before execution, so a read_file whose path is an object is persisted even
|
||||
// though the tool refuses it, and a projection that only considered
|
||||
// well-formed calls would carry it through.
|
||||
func TestPersistedToolCallsAreProjectedOnResumeAndFork(t *testing.T) {
|
||||
const body = "private body text with no credential shape at all"
|
||||
const token = "secret-value-9f8e7d6c5b4a"
|
||||
const oldText = "OLD FILE CONTENTS THAT MUST NOT BE REPLAYED"
|
||||
|
||||
store := NewStore(StoreOptions{RootDir: t.TempDir()})
|
||||
session, err := store.Create(CreateInput{SessionID: "projection_persisted"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
appends := []AppendEventInput{
|
||||
{Type: EventMessage, Payload: map[string]any{"role": "user", "content": "go"}},
|
||||
{Type: EventToolCall, Payload: map[string]any{
|
||||
"id": "c1", "name": "mcp_search",
|
||||
"arguments": `{"query":{"api_key":"` + token + `","content":"` + body + `"}}`,
|
||||
}},
|
||||
{Type: EventToolCall, Payload: map[string]any{
|
||||
"id": "c2", "name": "edit_file",
|
||||
"arguments": `{"path":"a.go","search":"` + oldText + `","replace":"new"}`,
|
||||
}},
|
||||
// Rejected by read_file, recorded anyway.
|
||||
{Type: EventToolCall, Payload: map[string]any{
|
||||
"id": "c3", "name": "read_file",
|
||||
"arguments": `{"path":{"nested":"` + body + `"},"start_line":40}`,
|
||||
}},
|
||||
{Type: EventToolCall, Payload: map[string]any{
|
||||
"id": "c4", "name": "grep",
|
||||
"arguments": `{"search":"func toolCallIdentity","path":"internal/sessions"}`,
|
||||
}},
|
||||
{Type: EventError, Payload: map[string]any{"message": "provider error: upstream timeout"}},
|
||||
}
|
||||
if _, err := store.AppendEvents(session.SessionID, appends); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
for _, mode := range []struct {
|
||||
name string
|
||||
options PrepareExecOptions
|
||||
}{
|
||||
{"resume", PrepareExecOptions{Store: store, Resume: session.SessionID}},
|
||||
{"fork", PrepareExecOptions{Store: store, Fork: session.SessionID, SessionID: "projection_persisted_fork"}},
|
||||
} {
|
||||
t.Run(mode.name, func(t *testing.T) {
|
||||
prepared, err := PrepareExec(mode.options)
|
||||
if err != nil {
|
||||
t.Fatalf("PrepareExec: %v", err)
|
||||
}
|
||||
prompt := FormatExecPrompt("continue", prepared)
|
||||
|
||||
for _, leaked := range []string{body, body[:20], token, token[:12], oldText, oldText[:20], "api_key"} {
|
||||
if strings.Contains(prompt, leaked) {
|
||||
t.Errorf("%q survived persistence into the %s prompt:\n%s", leaked, mode.name, prompt)
|
||||
}
|
||||
}
|
||||
// The identities that make the tail worth carrying at all.
|
||||
for _, kept := range []string{"mcp_search", "edit_file", "read_file", "grep", "func toolCallIdentity", "internal/sessions", "a.go"} {
|
||||
if !strings.Contains(prompt, kept) {
|
||||
t.Errorf("identity %q was lost from the %s prompt:\n%s", kept, mode.name, prompt)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// AND THE STORE STILL HOLDS THE ORIGINALS. The projection is a view for the
|
||||
// prompt, not an edit to the record: /rewind and an audit read the events.
|
||||
events, err := store.ReadEvents(session.SessionID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
stored := ""
|
||||
for _, event := range events {
|
||||
stored += string(event.Payload)
|
||||
}
|
||||
for _, want := range []string{body, token, oldText} {
|
||||
if !strings.Contains(stored, want) {
|
||||
t.Errorf("the projection modified the persisted events: %q is no longer in the store", want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,227 @@
|
||||
package sessions
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func toolContextPayload(t *testing.T, value map[string]any) json.RawMessage {
|
||||
t.Helper()
|
||||
raw, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return json.RawMessage(raw)
|
||||
}
|
||||
|
||||
// interruptedTurnEvents is a turn that read two files and then died on a
|
||||
// provider error, which is the shape #913 reports: no assistant answer was ever
|
||||
// produced, so nothing in the conversation describes the work.
|
||||
func interruptedTurnEvents(t *testing.T) []Event {
|
||||
t.Helper()
|
||||
return []Event{
|
||||
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "add retries to the http client"})},
|
||||
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c1", "name": "read_file", "arguments": `{"path":"internal/http/client.go"}`})},
|
||||
{Sequence: 3, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "read_file", "content": "package http"})},
|
||||
{Sequence: 4, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c2", "name": "read_file", "arguments": `{"path":"internal/http/retry.go"}`})},
|
||||
{Sequence: 5, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "read_file", "content": "var defaultBackoff = 250ms"})},
|
||||
{Sequence: 6, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error: upstream timeout"})},
|
||||
}
|
||||
}
|
||||
|
||||
func resumePrompt(t *testing.T, events []Event) string {
|
||||
t.Helper()
|
||||
return FormatExecPrompt("continue", PreparedExec{
|
||||
Mode: ModeResume,
|
||||
Session: Metadata{SessionID: "s1"},
|
||||
ContextEvents: events,
|
||||
})
|
||||
}
|
||||
|
||||
// AN INTERRUPTED TURN HAS TO SAY WHAT IT DID.
|
||||
//
|
||||
// A turn that ends normally describes its own work in the assistant's answer,
|
||||
// so the next turn inherits a prose record. A turn killed by a provider error
|
||||
// produces no answer, and the tool events that were the only record of the work
|
||||
// were filtered out of the resume context. The next turn was told a request had
|
||||
// been made and an error had happened, and nothing else, so it re-read from
|
||||
// scratch (#913).
|
||||
func TestResumePromptCarriesInterruptedToolWork(t *testing.T) {
|
||||
out := resumePrompt(t, interruptedTurnEvents(t))
|
||||
|
||||
for _, want := range []string{"internal/http/client.go", "internal/http/retry.go"} {
|
||||
if !strings.Contains(out, want) {
|
||||
t.Errorf("the resume prompt does not name %s, so the next turn cannot know it was already read:\n%s", want, out)
|
||||
}
|
||||
}
|
||||
// The conversation spine is still there.
|
||||
for _, want := range []string{"add retries to the http client", "upstream timeout"} {
|
||||
if !strings.Contains(out, want) {
|
||||
t.Errorf("the resume prompt lost conversation context %q:\n%s", want, out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// AND THE ORDER STAYS THE ORDER IT HAPPENED IN.
|
||||
//
|
||||
// The sequence number is rendered into each line, so a list that jumps
|
||||
// backwards reads as a corrupted history rather than a merge artifact.
|
||||
func TestResumePromptKeepsEventsInSequenceOrder(t *testing.T) {
|
||||
out := resumePrompt(t, interruptedTurnEvents(t))
|
||||
|
||||
last := -1
|
||||
for _, line := range strings.Split(out, "\n") {
|
||||
if !strings.HasPrefix(line, "- #") {
|
||||
continue
|
||||
}
|
||||
var seq int
|
||||
if _, err := fmt.Sscanf(line, "- #%d", &seq); err != nil {
|
||||
t.Fatalf("unparsable context line %q", line)
|
||||
}
|
||||
if seq <= last {
|
||||
t.Fatalf("context is out of order at %q (previous #%d):\n%s", line, last, out)
|
||||
}
|
||||
last = seq
|
||||
}
|
||||
if last < 0 {
|
||||
t.Fatal("SETUP INVALID: the prompt rendered no context lines")
|
||||
}
|
||||
}
|
||||
|
||||
// TOOL EVENTS MUST NOT EVICT CONVERSATION.
|
||||
//
|
||||
// This is what the original filter was added for (#460): tool events vastly
|
||||
// outnumber messages, so admitting them into the same trailing budget would let
|
||||
// one tool-heavy turn push every earlier message out of the context. A separate
|
||||
// allowance is what keeps both properties at once.
|
||||
func TestToolEventsNeverEvictConversation(t *testing.T) {
|
||||
var events []Event
|
||||
seq := 0
|
||||
next := func() int { seq++; return seq }
|
||||
// An early message that must survive, then a flood of tool work.
|
||||
events = append(events, Event{Sequence: next(), Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "EARLIEST-REQUEST-MARKER"})})
|
||||
for i := 0; i < 500; i++ {
|
||||
arguments := fmt.Sprintf(`{"path":"noise/%03d.go"}`, i)
|
||||
events = append(events, Event{Sequence: next(), Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"name": "read_file", "arguments": arguments})})
|
||||
events = append(events, Event{Sequence: next(), Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "read_file", "content": "noise"})})
|
||||
}
|
||||
events = append(events, Event{Sequence: next(), Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "LATEST-REQUEST-MARKER"})})
|
||||
|
||||
out := resumePrompt(t, events)
|
||||
for _, want := range []string{"EARLIEST-REQUEST-MARKER", "LATEST-REQUEST-MARKER"} {
|
||||
if !strings.Contains(out, want) {
|
||||
t.Errorf("a tool-heavy turn evicted conversation event %q, which is the regression the filter exists to prevent", want)
|
||||
}
|
||||
}
|
||||
|
||||
selected := promptContextEvents(events)
|
||||
tools := 0
|
||||
for _, event := range selected {
|
||||
if event.Type == EventToolCall || event.Type == EventToolResult {
|
||||
tools++
|
||||
}
|
||||
}
|
||||
if tools == 0 {
|
||||
t.Error("no tool work survived at all, so an interrupted tool-heavy turn still says nothing about what it did")
|
||||
}
|
||||
if len(selected) > 80 {
|
||||
t.Errorf("selected %d events, over the 80 the prompt budget allows", len(selected))
|
||||
}
|
||||
}
|
||||
|
||||
// A session with no tool work renders exactly as it did before.
|
||||
func TestResumePromptUnchangedWithoutToolEvents(t *testing.T) {
|
||||
events := []Event{
|
||||
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "hello"})},
|
||||
{Sequence: 2, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "assistant", "content": "hi"})},
|
||||
}
|
||||
out := resumePrompt(t, events)
|
||||
if strings.Count(out, "- #") != 2 {
|
||||
t.Fatalf("a tool-free session rendered %d context lines, want 2:\n%s", strings.Count(out, "- #"), out)
|
||||
}
|
||||
}
|
||||
|
||||
// TOOL OUTPUT MUST NOT RIDE ALONG INTO A LATER PROMPT.
|
||||
//
|
||||
// The tail is there to say what the interrupted turn did, and the call already
|
||||
// says that. Carrying the result body would put up to 500 bytes of raw tool
|
||||
// output into a prompt on a later turn, and nothing redacts on the way in.
|
||||
func TestResumePromptCarriesToolOutcomeWithoutOutput(t *testing.T) {
|
||||
secret := "AKIAIOSFODNN7EXAMPLE-and-more-file-contents"
|
||||
events := []Event{
|
||||
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "look at the config"})},
|
||||
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c1", "name": "read_file", "arguments": `{"path":"deploy/prod.env"}`})},
|
||||
{Sequence: 3, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "read_file", "status": "ok", "output": secret})},
|
||||
{Sequence: 4, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c2", "name": "read_file", "arguments": `{"path":"deploy/missing.env"}`})},
|
||||
{Sequence: 5, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "read_file", "status": "error", "output": "no such file"})},
|
||||
{Sequence: 6, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error: upstream timeout"})},
|
||||
}
|
||||
out := resumePrompt(t, events)
|
||||
|
||||
// THE TOOL LINES ARE PINNED WHOLE, NOT SEARCHED.
|
||||
//
|
||||
// Two reasons the obvious assertions are too weak to say what the trim
|
||||
// promises. Searching the prompt for the fixture's full secret only rejects
|
||||
// that exact string, so a change that carried a truncated or partly
|
||||
// redacted prefix of the output would pass it. And asking whether "ok" and
|
||||
// "error" each appear somewhere among the result lines passes just as well
|
||||
// when the two statuses are swapped, or both attached to the wrong call.
|
||||
//
|
||||
// Pinning the whole of every tool line covers both: nothing but the tool
|
||||
// name and the outcome survives a result, each outcome sits on its own
|
||||
// line, and the line order pairs each one with the call above it.
|
||||
//
|
||||
// Compared as a set of fields per line rather than as a string, because the
|
||||
// renderer walks the payload map and Go randomizes that order, so the same
|
||||
// events render their fields in a different order on every run.
|
||||
want := []struct {
|
||||
sequence int
|
||||
kind string
|
||||
fields []string
|
||||
}{
|
||||
{2, "tool_call", []string{"c1", "read_file", `{"path":"deploy/prod.env"}`}},
|
||||
{3, "tool_result", []string{"read_file", "ok"}},
|
||||
{4, "tool_call", []string{"c2", "read_file", `{"path":"deploy/missing.env"}`}},
|
||||
{5, "tool_result", []string{"read_file", "error"}},
|
||||
}
|
||||
var toolLines []string
|
||||
for _, line := range strings.Split(out, "\n") {
|
||||
if strings.HasPrefix(line, "- #") && (strings.Contains(line, "tool_call:") || strings.Contains(line, "tool_result:")) {
|
||||
toolLines = append(toolLines, line)
|
||||
}
|
||||
}
|
||||
if len(toolLines) != len(want) {
|
||||
t.Fatalf("the resume prompt carries %d tool lines, want %d:\n%s", len(toolLines), len(want), out)
|
||||
}
|
||||
for index, line := range toolLines {
|
||||
var sequence int
|
||||
var kind string
|
||||
rest, found := strings.CutPrefix(line, "- #")
|
||||
if !found {
|
||||
t.Fatalf("unparsable context line %q", line)
|
||||
}
|
||||
if _, err := fmt.Sscanf(rest, "%d %s", &sequence, &kind); err != nil {
|
||||
t.Fatalf("unparsable context line %q: %v", line, err)
|
||||
}
|
||||
kind = strings.TrimSuffix(kind, ":")
|
||||
_, payload, _ := strings.Cut(rest, ": ")
|
||||
fields := strings.Fields(payload)
|
||||
slices.Sort(fields)
|
||||
expected := slices.Clone(want[index].fields)
|
||||
slices.Sort(expected)
|
||||
if sequence != want[index].sequence || kind != want[index].kind || !slices.Equal(fields, expected) {
|
||||
t.Errorf("tool line %d is not what the trim promises.\ngot: #%d %s %v\nwant: #%d %s %v\nfull prompt:\n%s",
|
||||
index, sequence, kind, fields, want[index].sequence, want[index].kind, expected, out)
|
||||
}
|
||||
}
|
||||
// Named on its own so a leak reports as a leak rather than as the
|
||||
// formatting drift the comparison above would also catch. A PREFIX, because
|
||||
// the danger is output reaching a later prompt at all, not this exact
|
||||
// fixture string reaching it.
|
||||
if strings.Contains(out, secret[:16]) {
|
||||
t.Errorf("tool output reached the resume prompt:\n%s", out)
|
||||
}
|
||||
}
|
||||
@@ -4,8 +4,11 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"sort"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/redaction"
|
||||
)
|
||||
|
||||
type ExecMode string
|
||||
@@ -188,23 +191,307 @@ func FormatExecPrompt(prompt string, prepared PreparedExec) string {
|
||||
}, "\n")
|
||||
}
|
||||
|
||||
// promptContextEvents chooses what a resumed turn is told about the session so
|
||||
// far.
|
||||
//
|
||||
// A RESUMED TURN NEEDS TO KNOW WHAT THE PREVIOUS ONE DID, NOT ONLY WHAT IT SAID.
|
||||
//
|
||||
// Conversation events alone were selected here, so a turn that read six files
|
||||
// and then died before answering left this behind:
|
||||
//
|
||||
// - #1 message: add retries to the http client
|
||||
// - #6 error: provider error: upstream timeout
|
||||
//
|
||||
// Nothing names the files, so the next turn re-reads them from scratch (#913).
|
||||
// On a turn that ENDS NORMALLY the assistant's answer describes the work, which
|
||||
// is why this was survivable; an interrupted turn produces no such answer, and
|
||||
// the record of the work goes with it.
|
||||
//
|
||||
// ONLY THE UNSUMMARIZED TAIL, NOT EVERY TOOL EVENT. Filtering tool events out
|
||||
// was deliberate (#460) and is right for work the assistant has already
|
||||
// described: forty read_file results add length and no information next to the
|
||||
// answer that explains them. It is wrong only for work with no answer after it,
|
||||
// which is precisely what an interrupted turn leaves behind. So the events after
|
||||
// the last assistant message come along, and everything the assistant already
|
||||
// spoke for stays filtered.
|
||||
//
|
||||
// The tail also gets its own allowance rather than sharing the conversation
|
||||
// budget, so a tool-heavy interrupted turn can never push an earlier message out
|
||||
// of the context. That is the property #460 added this filter for and it still
|
||||
// holds.
|
||||
func promptContextEvents(events []Event) []Event {
|
||||
const maxPromptContextEvents = 80
|
||||
// Enough to describe an interrupted turn's work, small enough that the
|
||||
// conversation still dominates the prompt.
|
||||
const maxPromptContextTailEvents = 24
|
||||
|
||||
filtered := make([]Event, 0, len(events))
|
||||
for _, event := range events {
|
||||
switch event.Type {
|
||||
case EventMessage, EventCompaction, EventSessionFork, EventSessionChild, EventSpecialistStart, EventSpecialistStop, EventError:
|
||||
filtered = append(filtered, event)
|
||||
lastSpoken := -1
|
||||
for index, event := range events {
|
||||
if event.Type == EventMessage && payloadRole(event.Payload) == "assistant" {
|
||||
lastSpoken = index
|
||||
}
|
||||
}
|
||||
if len(filtered) == 0 {
|
||||
filtered = append(filtered, events...)
|
||||
|
||||
conversation := make([]Event, 0, len(events))
|
||||
tail := make([]Event, 0, maxPromptContextTailEvents)
|
||||
for index, event := range events {
|
||||
switch event.Type {
|
||||
case EventMessage, EventCompaction, EventSessionFork, EventSessionChild, EventSpecialistStart, EventSpecialistStop, EventError:
|
||||
conversation = append(conversation, event)
|
||||
case EventToolCall:
|
||||
if index > lastSpoken {
|
||||
tail = append(tail, toolCallIdentity(event))
|
||||
}
|
||||
case EventToolResult:
|
||||
if index > lastSpoken {
|
||||
tail = append(tail, toolResultOutcome(event))
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(filtered) > maxPromptContextEvents {
|
||||
filtered = filtered[len(filtered)-maxPromptContextEvents:]
|
||||
if len(conversation) == 0 && len(tail) == 0 {
|
||||
// An unrecognized event stream still says more than nothing.
|
||||
conversation = append(conversation, events...)
|
||||
}
|
||||
return filtered
|
||||
if len(conversation) > maxPromptContextEvents {
|
||||
conversation = conversation[len(conversation)-maxPromptContextEvents:]
|
||||
}
|
||||
// Tail slots exist only when the conversation uses fewer than the 80-event
|
||||
// cap. A session whose conversation already fills it carries no interrupted
|
||||
// tool work, which is what every session did before tool events were admitted
|
||||
// at all, and what the #460 cap is there to hold. Reserving a minimum tail by
|
||||
// trimming conversation further is a product decision, not this change.
|
||||
tailBudget := min(maxPromptContextTailEvents, maxPromptContextEvents-len(conversation))
|
||||
if tailBudget < 0 {
|
||||
tailBudget = 0
|
||||
}
|
||||
if len(tail) > tailBudget {
|
||||
tail = tail[len(tail)-tailBudget:]
|
||||
}
|
||||
if len(tail) == 0 {
|
||||
return conversation
|
||||
}
|
||||
// Merged back into execution order: the sequence numbers are rendered, so a
|
||||
// list that jumps backwards would read as a corrupted history.
|
||||
merged := make([]Event, 0, len(conversation)+len(tail))
|
||||
merged = append(merged, conversation...)
|
||||
merged = append(merged, tail...)
|
||||
sort.SliceStable(merged, func(i, j int) bool { return merged[i].Sequence < merged[j].Sequence })
|
||||
return merged
|
||||
}
|
||||
|
||||
// toolCallIdentityKeys are the argument fields that say WHAT a call was about
|
||||
// without carrying what it was about to write, run, or search for.
|
||||
//
|
||||
// Path, directory, URL, name and pattern fields identify the work: the file that
|
||||
// was read, the tree that was listed, the expression that was searched. Body
|
||||
// fields carry payload: write_file's content, edit_file's old and new strings,
|
||||
// apply_patch's hunks, a shell command with whatever credential was on its
|
||||
// line. The result side already drops payload through toolResultOutcome, and
|
||||
// admitting calls into resume context without the same care put an interrupted
|
||||
// write_file's content into the next turn's prompt while its result body did
|
||||
// not. Same fact, one door left open.
|
||||
//
|
||||
// AN ALLOW-LIST, NOT A DENY-LIST, so an argument this file has never heard of is
|
||||
// dropped rather than replayed. A new tool with a new body field is then a
|
||||
// missing path in a resume prompt, which is visible, instead of a new leak, which
|
||||
// is not. The keys cover every alias the tools accept for the identity fields.
|
||||
var toolCallIdentityKeys = map[string]bool{
|
||||
// files and directories
|
||||
"path": true, "file": true, "file_path": true, "filepath": true, "filename": true,
|
||||
"dir": true, "directory": true, "cwd": true, "workdir": true,
|
||||
// what a search was for; these are what the interrupted turn was looking at
|
||||
"pattern": true, "glob": true, "query": true, "regex": true, "expression": true,
|
||||
"match": true,
|
||||
// fetches and named resources
|
||||
"url": true, "name": true,
|
||||
// read windows, so a resumed turn knows which part it already had. Every
|
||||
// form read_file accepts, because a resumed turn that knows the path but not
|
||||
// which slice was asked for has half the identity and will read again.
|
||||
"offset": true, "limit": true,
|
||||
"start_line": true, "end_line": true, "max_lines": true,
|
||||
"byte_offset": true, "byte_limit": true,
|
||||
}
|
||||
|
||||
// toolCallIdentityKeysByTool holds the argument names whose meaning depends on
|
||||
// which tool was called, so they cannot live in the table above.
|
||||
//
|
||||
// grep accepts `search` for its pattern, and edit_file accepts the same word
|
||||
// for the text being replaced. Admitting it globally for the sake of the first
|
||||
// would replay file contents through the second, which is the whole thing this
|
||||
// projection exists to stop. The rule this encodes: an argument name that any
|
||||
// MUTATING tool accepts for body content is scoped to the tools that mean
|
||||
// something else by it, never added to the shared table.
|
||||
var toolCallIdentityKeysByTool = map[string]map[string]bool{
|
||||
"grep": {"search": true},
|
||||
}
|
||||
|
||||
// retainedIdentityValue decides whether a value under a permitted key may be
|
||||
// kept, and what it looks like if so.
|
||||
//
|
||||
// AN ALLOW-LISTED KEY IS NOT A SAFE VALUE. A url is a valid place for a
|
||||
// credential to appear: web_fetch accepts a query token and redacts the URL it
|
||||
// reports back, so an interrupted fetch of
|
||||
// https://api.example/data?access_token=... had its token dropped from the
|
||||
// result and kept verbatim in the call. This projection is what admits call
|
||||
// arguments into a later turn's prompt, so the same scrub belongs here or the
|
||||
// credential is replayed on resume. Host and path survive it, which is the
|
||||
// identity the resumed turn needs.
|
||||
//
|
||||
// AND A KEY IS NOT PERMISSION FOR WHATEVER HANGS UNDER IT. Scrubbing strings
|
||||
// and passing everything else through read as "numbers are nothing to scrub",
|
||||
// which is true of a number and not of an object: MCP schemas allow object and
|
||||
// array properties, so {"query":{"api_key":"...","content":"private body"}}
|
||||
// arrives under a permitted key and carries a whole payload with it. Redacting
|
||||
// a serialized object would not help either, since ordinary private text has no
|
||||
// credential shape to match.
|
||||
//
|
||||
// So the retained shapes are named: a string, scrubbed, and the scalars a read
|
||||
// window is made of. A container is dropped. The call keeps its name and id, so
|
||||
// a resumed turn still knows which tool ran and can ask again; what it does not
|
||||
// get is arbitrary nested content it was never meant to see.
|
||||
func retainedIdentityValue(value any) (any, bool) {
|
||||
switch typed := value.(type) {
|
||||
case string:
|
||||
return redaction.RedactString(typed, redaction.Options{}), true
|
||||
case float64, int, int64, bool, json.Number:
|
||||
return typed, true
|
||||
default:
|
||||
// Objects, arrays, null, and anything a future decoder invents.
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
|
||||
// toolCallIdentityKeyAllowed reports whether key is an identity field for the
|
||||
// tool that was called.
|
||||
func toolCallIdentityKeyAllowed(toolName, key string) bool {
|
||||
lowered := strings.ToLower(key)
|
||||
if toolCallIdentityKeys[lowered] {
|
||||
return true
|
||||
}
|
||||
return toolCallIdentityKeysByTool[strings.ToLower(strings.TrimSpace(toolName))][lowered]
|
||||
}
|
||||
|
||||
// toolCallIdentityFields are the top-level payload fields a projected call
|
||||
// keeps: which call it was, and which tool. Everything else on the payload is
|
||||
// dropped, including a field added by a future producer that this file has
|
||||
// never seen.
|
||||
var toolCallIdentityFields = []string{"id", "name"}
|
||||
|
||||
// toolCallIdentity keeps a tool call's identity and drops its payload, the
|
||||
// symmetric half of toolResultOutcome.
|
||||
//
|
||||
// The arguments travel as a JSON string inside the payload. They are decoded,
|
||||
// reduced to the identity keys, and re-encoded; anything that does not decode as
|
||||
// an object is removed outright rather than passed through as text, since text
|
||||
// that could not be read is text that cannot be checked.
|
||||
//
|
||||
// REBUILT, NOT PATCHED, WHICH IS WHAT MAKES IT THE SYMMETRIC HALF.
|
||||
// toolResultOutcome constructs its output from the two fields it keeps, so a
|
||||
// field it has never heard of cannot reach a prompt through it. Editing
|
||||
// arguments in place and returning the rest of the payload is the opposite
|
||||
// contract: it drops what it recognises as a body and forwards every sibling
|
||||
// key, so a producer recording one more top-level field would put that field
|
||||
// into the next turn. The projection names what survives.
|
||||
func toolCallIdentity(event Event) Event {
|
||||
var decoded map[string]json.RawMessage
|
||||
if err := json.Unmarshal(event.Payload, &decoded); err != nil {
|
||||
event.Payload = json.RawMessage(`{}`)
|
||||
return event
|
||||
}
|
||||
projected := map[string]json.RawMessage{}
|
||||
for _, field := range toolCallIdentityFields {
|
||||
if value, present := decoded[field]; present {
|
||||
projected[field] = value
|
||||
}
|
||||
}
|
||||
// The tool decides what some argument names mean, so the projection has to
|
||||
// know which tool this was before it can say which keys are identity.
|
||||
toolName := ""
|
||||
if raw, present := decoded["name"]; present {
|
||||
_ = json.Unmarshal(raw, &toolName)
|
||||
}
|
||||
kept := map[string]any{}
|
||||
if raw, present := decoded["arguments"]; present {
|
||||
var argumentsText string
|
||||
if err := json.Unmarshal(raw, &argumentsText); err == nil {
|
||||
var arguments map[string]any
|
||||
if err := json.Unmarshal([]byte(argumentsText), &arguments); err == nil {
|
||||
for key, value := range arguments {
|
||||
if !toolCallIdentityKeyAllowed(toolName, key) {
|
||||
continue
|
||||
}
|
||||
if retained, ok := retainedIdentityValue(value); ok {
|
||||
kept[key] = retained
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(kept) > 0 {
|
||||
if reduced, err := json.Marshal(kept); err == nil {
|
||||
if quoted, err := json.Marshal(string(reduced)); err == nil {
|
||||
projected["arguments"] = quoted
|
||||
}
|
||||
}
|
||||
}
|
||||
rebuilt, err := json.Marshal(projected)
|
||||
if err != nil {
|
||||
event.Payload = json.RawMessage(`{}`)
|
||||
return event
|
||||
}
|
||||
event.Payload = json.RawMessage(rebuilt)
|
||||
return event
|
||||
}
|
||||
|
||||
// toolResultOutcome strips a tool result down to WHICH tool ran and HOW IT
|
||||
// ENDED, dropping the output body.
|
||||
//
|
||||
// The tail exists so a resumed turn knows what the interrupted one did, and the
|
||||
// CALL already carries that: the tool name and its arguments, which is the path
|
||||
// for a read. The result adds only a 500-byte prefix of the output, which is
|
||||
// worth little beside the call and is the one part of an event that can carry
|
||||
// file contents. Nothing redacts on the way into a prompt, so keeping it would
|
||||
// re-emit raw tool output into a later turn on the strength of a truncation
|
||||
// limit alone.
|
||||
//
|
||||
// The status stays, because dropping it would be worse than dropping the whole
|
||||
// result: a bare call reads as work that succeeded, so a failed read would come
|
||||
// back as a file the next turn believes it already has.
|
||||
func toolResultOutcome(event Event) Event {
|
||||
var decoded struct {
|
||||
Name string `json:"name"`
|
||||
Status string `json:"status"`
|
||||
}
|
||||
if err := json.Unmarshal(event.Payload, &decoded); err != nil {
|
||||
// Undecodable: drop the payload rather than pass an unknown shape through.
|
||||
event.Payload = json.RawMessage(`{}`)
|
||||
return event
|
||||
}
|
||||
trimmed, err := json.Marshal(map[string]string{
|
||||
"name": decoded.Name,
|
||||
"status": decoded.Status,
|
||||
})
|
||||
if err != nil {
|
||||
event.Payload = json.RawMessage(`{}`)
|
||||
return event
|
||||
}
|
||||
event.Payload = json.RawMessage(trimmed)
|
||||
return event
|
||||
}
|
||||
|
||||
// payloadRole reads the "role" field of a message payload, returning "" when the
|
||||
// payload is not a decodable object or carries no role.
|
||||
func payloadRole(payload json.RawMessage) string {
|
||||
if len(payload) == 0 {
|
||||
return ""
|
||||
}
|
||||
var decoded struct {
|
||||
Role string `json:"role"`
|
||||
}
|
||||
if err := json.Unmarshal(payload, &decoded); err != nil {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(decoded.Role)
|
||||
}
|
||||
|
||||
func forkTitle(title string) string {
|
||||
|
||||
@@ -610,8 +610,60 @@ func TestFormatExecPromptKeepsConversationMessagesWhenNoisyEventsFollow(t *testi
|
||||
t.Fatalf("expected prompt to contain %q, got %q", want, prompt)
|
||||
}
|
||||
}
|
||||
// The tool results here follow the last assistant answer, so they are work
|
||||
// nothing has spoken for yet. They are carried on purpose: that is the whole
|
||||
// of what an interrupted turn leaves behind (#913). The property this test
|
||||
// was written for (#460) is that CONVERSATION survives a noisy turn, and the
|
||||
// assertions above still hold it.
|
||||
// The tool results here follow the last assistant answer, so they are work
|
||||
// nothing has spoken for yet, and an interrupted turn leaves exactly this
|
||||
// behind (#913). Their IDENTITY is carried; their OUTPUT is not, because the
|
||||
// call beside them already names what was touched and the body is the part
|
||||
// that can carry file contents into a later prompt.
|
||||
if !strings.Contains(prompt, "tool_result: read_file") {
|
||||
t.Fatalf("expected unanswered tool work to be named, got %q", prompt)
|
||||
}
|
||||
if strings.Contains(prompt, "noisy tool result") {
|
||||
t.Fatalf("expected prompt to omit noisy non-conversation events, got %q", prompt)
|
||||
t.Fatalf("tool output body reached the prompt, got %q", prompt)
|
||||
}
|
||||
if count := strings.Count(prompt, "tool_result:"); count > 24 {
|
||||
t.Fatalf("carried %d tool events, over the tail allowance of 24: %q", count, prompt)
|
||||
}
|
||||
}
|
||||
|
||||
// AND WORK THE ASSISTANT ALREADY SPOKE FOR STAYS FILTERED.
|
||||
//
|
||||
// The counterpart to the case above, and what keeps that one from being
|
||||
// satisfied by carrying every tool event ever recorded. Once an answer describes
|
||||
// the work, repeating the raw results adds length and no information.
|
||||
func TestFormatExecPromptOmitsToolWorkAnAnswerAlreadyCovers(t *testing.T) {
|
||||
events := []Event{
|
||||
{Sequence: 1, Type: EventMessage, Payload: json.RawMessage(`{"role":"user","content":"first user request"}`)},
|
||||
}
|
||||
for sequence := 2; sequence <= 20; sequence++ {
|
||||
events = append(events, Event{Sequence: sequence, Type: EventToolResult, Payload: json.RawMessage(`{"name":"read_file","output":"already summarized tool result"}`)})
|
||||
}
|
||||
events = append(events,
|
||||
Event{Sequence: 21, Type: EventMessage, Payload: json.RawMessage(`{"role":"assistant","content":"I read the files and here is the summary"}`)},
|
||||
Event{Sequence: 22, Type: EventMessage, Payload: json.RawMessage(`{"role":"user","content":"latest user request"}`)},
|
||||
)
|
||||
|
||||
prompt := FormatExecPrompt("continue", PreparedExec{
|
||||
Mode: ModeResume,
|
||||
Session: Metadata{SessionID: "session-already-answered"},
|
||||
ContextEvents: events,
|
||||
})
|
||||
|
||||
if strings.Contains(prompt, "already summarized tool result") {
|
||||
t.Fatalf("tool work the assistant already described was repeated verbatim, got %q", prompt)
|
||||
}
|
||||
if strings.Contains(prompt, "tool_result:") {
|
||||
t.Fatalf("an answered turn still listed its tool events, got %q", prompt)
|
||||
}
|
||||
for _, want := range []string{"first user request", "here is the summary", "latest user request"} {
|
||||
if !strings.Contains(prompt, want) {
|
||||
t.Fatalf("expected prompt to contain %q, got %q", want, prompt)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -75,14 +75,10 @@ func NewScopedExecCommandTool(workspaceRoot string, scope PathScope, manager *ex
|
||||
if manager == nil {
|
||||
manager = defaultExecSessionManager
|
||||
}
|
||||
description := "Runs a command in a PTY, returning output or a session ID for ongoing interaction."
|
||||
if runtimeGOOS() == "windows" {
|
||||
description += "\n\n" + shellGuidanceForGOOS(runtimeGOOS())
|
||||
}
|
||||
return execCommandTool{
|
||||
baseTool: baseTool{
|
||||
name: ExecCommandToolName,
|
||||
description: description,
|
||||
description: execCommandDescription(runtimeGOOS()),
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]PropertySchema{
|
||||
@@ -674,3 +670,19 @@ func execDisplaySummary(commandText string, sessionID int, exited bool, exitCode
|
||||
func runtimeGOOS() string {
|
||||
return runtime.GOOS
|
||||
}
|
||||
|
||||
// execCommandDescription is exec_command's description as a host running goos
|
||||
// would see it. The shell guidance is appended on Windows only.
|
||||
//
|
||||
// Taking goos as an argument rather than reading it is what lets the guidance
|
||||
// contract be checked for every platform from any one of them. internal/agent's
|
||||
// token ratchet removes exactly this appended text before charging the schemas,
|
||||
// so the condition here and hostExecCommandShellGuidance have to agree on all
|
||||
// three platforms, not only on whichever one the test happens to run on.
|
||||
func execCommandDescription(goos string) string {
|
||||
description := "Runs a command in a PTY, returning output or a session ID for ongoing interaction."
|
||||
if goos == "windows" {
|
||||
description += "\n\n" + shellGuidanceForGOOS(goos)
|
||||
}
|
||||
return description
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// THE BUDGET TEST REMOVES THIS TEXT, SO IT HAS TO BE EXACTLY WHAT IS ADDED.
|
||||
//
|
||||
// internal/agent's eager-schema ratchet charges the tool schemas and this
|
||||
// host's shell guidance separately, and it separates them by stripping
|
||||
// HostExecCommandShellGuidance off exec_command's description before
|
||||
// estimating. That is only valid while the guidance is appended verbatim as the
|
||||
// tail of that description, and only while a platform that reports no guidance
|
||||
// really appends none. If the two drift apart the ratchet keeps reporting a
|
||||
// number, just not the one it claims.
|
||||
//
|
||||
// Checked for all three platforms rather than for the running one: the
|
||||
// appending is itself conditional on GOOS, so a host-only check is the shape of
|
||||
// bug this pins. Off Windows it would pass while reporting nothing at all.
|
||||
func TestExecCommandDescriptionMatchesReportedGuidance(t *testing.T) {
|
||||
for _, goos := range []string{"windows", "linux", "darwin"} {
|
||||
t.Run(goos, func(t *testing.T) {
|
||||
description := execCommandDescription(goos)
|
||||
guidance := hostExecCommandShellGuidance(goos)
|
||||
|
||||
if goos == "windows" {
|
||||
if guidance == "" {
|
||||
t.Fatal("exec_command appends shell guidance on Windows, but the accessor reports none; the budget would charge it to the schemas")
|
||||
}
|
||||
if !strings.HasSuffix(description, guidance) {
|
||||
t.Fatalf("the guidance is not the tail of the description, so removing it from the schema total is wrong.\nguidance: %q\ndescription tail: %q",
|
||||
guidance, description[max(0, len(description)-len(guidance)-40):])
|
||||
}
|
||||
} else if guidance != "" {
|
||||
t.Fatalf("no shell guidance is appended off Windows, but the accessor reported %q", guidance)
|
||||
}
|
||||
|
||||
// Whatever guidance THIS platform would produce, none of it may be
|
||||
// left in the description once the reported guidance is removed.
|
||||
// Off Windows the accessor reports nothing, so this is the check
|
||||
// that the description really carries nothing: drop the GOOS
|
||||
// condition in execCommandDescription and a POSIX host starts
|
||||
// appending its own guidance (about 150 characters on darwin) with
|
||||
// the ratchet still charging every byte of it to the schemas.
|
||||
platformGuidance := shellGuidanceForGOOS(goos)
|
||||
if platformGuidance == "" {
|
||||
t.Fatalf("SETUP INVALID: %s produces no shell guidance at all, so the check below cannot fail", goos)
|
||||
}
|
||||
if remainder := strings.TrimSuffix(description, guidance); strings.Contains(remainder, platformGuidance) {
|
||||
t.Fatalf("exec_command's description still carries %s shell guidance after removing the %d characters the accessor reports, so the ratchet charges the rest to the tool schemas.\nremainder: %q",
|
||||
goos, len(guidance), remainder)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// And the running host agrees with the table above, so the exported accessor
|
||||
// and the real constructor cannot drift away from the platform contract.
|
||||
func TestHostExecCommandShellGuidanceMatchesTheDescription(t *testing.T) {
|
||||
tool := NewExecCommandTool(t.TempDir(), nil)
|
||||
if got, want := tool.Description(), execCommandDescription(runtime.GOOS); got != want {
|
||||
t.Fatalf("the constructed tool's description is not what the platform contract says.\ngot: %q\nwant: %q", got, want)
|
||||
}
|
||||
if got, want := HostExecCommandShellGuidance(), hostExecCommandShellGuidance(runtime.GOOS); got != want {
|
||||
t.Fatalf("the exported accessor disagrees with the platform contract.\ngot: %q\nwant: %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// And the guidance a Windows host gets is never empty, whichever shell was
|
||||
// detected, since an empty one would mean the model gets no syntax rule at all.
|
||||
func TestWindowsShellGuidanceIsNeverEmpty(t *testing.T) {
|
||||
for _, shell := range []shellRuntime{
|
||||
{GOOS: "windows", Executable: `C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe`, Kind: shellKindPowerShell, Syntax: "PowerShell"},
|
||||
{GOOS: "windows", Executable: `C:\Program Files\PowerShell\7\pwsh.exe`, Kind: shellKindPowerShell, Syntax: "PowerShell"},
|
||||
{GOOS: "windows", Executable: `C:\Windows\System32\cmd.exe`, Kind: shellKindCmd, Syntax: "cmd.exe"},
|
||||
} {
|
||||
if guidance := shellGuidanceForRuntime(shell); strings.TrimSpace(guidance) == "" {
|
||||
t.Errorf("%s produced no shell guidance", shell.Executable)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -181,6 +181,28 @@ func shellGuidanceForRuntime(shell shellRuntime) string {
|
||||
return guidance
|
||||
}
|
||||
|
||||
// HostExecCommandShellGuidance returns the shell guidance appended to
|
||||
// exec_command's description on this host, or "" where none is appended.
|
||||
//
|
||||
// Exported for the per-turn token ratchet in internal/agent. That budget has to
|
||||
// separate what every host pays for the tool schemas from what THIS host pays
|
||||
// for its shell guidance, because the guidance varies by shell and by
|
||||
// PowerShell version, and a single number covering both silently means a
|
||||
// different thing on each machine.
|
||||
func HostExecCommandShellGuidance() string {
|
||||
return hostExecCommandShellGuidance(runtimeGOOS())
|
||||
}
|
||||
|
||||
// hostExecCommandShellGuidance is the guidance a host running goos appends. It
|
||||
// is the pair of execCommandDescription, so the two can be checked against each
|
||||
// other for every platform rather than only for the running one.
|
||||
func hostExecCommandShellGuidance(goos string) string {
|
||||
if goos != "windows" {
|
||||
return ""
|
||||
}
|
||||
return shellGuidanceForGOOS(goos)
|
||||
}
|
||||
|
||||
// HostShellEnvironmentGuidance returns the concise, model-facing shell rule
|
||||
// for the current host.
|
||||
func HostShellEnvironmentGuidance() string {
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/proxydial"
|
||||
"github.com/Gitlawb/zero/internal/redaction"
|
||||
zeroSandbox "github.com/Gitlawb/zero/internal/sandbox"
|
||||
)
|
||||
@@ -308,19 +309,48 @@ func webFetchSafeTransport(roundTripper http.RoundTripper, resolver webFetchReso
|
||||
transport = http.DefaultTransport.(*http.Transport).Clone()
|
||||
case *http.Transport:
|
||||
transport = typed.Clone()
|
||||
case *proxydial.Transport:
|
||||
// Already guarded: rebuild from the inner transport so a second pass
|
||||
// installs one dialer rather than nesting wrappers.
|
||||
transport = typed.Transport.Clone()
|
||||
default:
|
||||
return roundTripper
|
||||
}
|
||||
|
||||
dialer := &net.Dialer{Timeout: webFetchTimeout, KeepAlive: 30 * time.Second}
|
||||
transport.Proxy = nil
|
||||
// THE PROXY IS HONOURED, AND THAT IS A DELIBERATE REVERSAL. The original
|
||||
// guarded fetch set Proxy to nil as part of tightening the SSRF guardrails,
|
||||
// and the reasoning was sound as far as it went: a proxied request is dialed
|
||||
// to the proxy, so the dial-time address pin below never sees the target and
|
||||
// the DNS-rebind protection it provides does not apply. What that traded away
|
||||
// was every fetch on a machine that reaches the network only through a
|
||||
// configured proxy, which fails outright rather than degrades (#569).
|
||||
//
|
||||
// The target is still validated before the request is made, lexically and
|
||||
// through the resolver, with the same block list the dialer applies. What a
|
||||
// forward proxy removes is only the re-check at dial time, and that is
|
||||
// inherent to using any forward proxy; a user who sets HTTPS_PROXY has chosen
|
||||
// it.
|
||||
//
|
||||
// The returned round tripper is wrapped so the dialer learns which route
|
||||
// the transport chose for each request: the one dial it lets past the guard
|
||||
// is the one this transport opens to the proxy it picked for that request.
|
||||
transport.Proxy = http.ProxyFromEnvironment
|
||||
transport.DialContext = webFetchSafeDialContext(resolver, dialer)
|
||||
transport.DialTLSContext = nil
|
||||
return transport
|
||||
return proxydial.Wrap(transport)
|
||||
}
|
||||
|
||||
func webFetchSafeDialContext(resolver webFetchResolver, dialer webFetchDialer) func(context.Context, string, string) (net.Conn, error) {
|
||||
return func(ctx context.Context, network string, address string) (net.Conn, error) {
|
||||
// The proxy is not the target. See webFetchSafeTransport for why the
|
||||
// proxy is honoured at all; here, the dial the transport opens to the
|
||||
// proxy it selected for this request skips the pin, and nothing else
|
||||
// does. A direct dial is pinned as it always was, including one whose
|
||||
// address happens to equal a configured proxy endpoint.
|
||||
if proxydial.IsProxyDial(ctx) {
|
||||
return dialer.DialContext(ctx, network, address)
|
||||
}
|
||||
pinnedAddress, err := webFetchSafeDialAddress(ctx, resolver, address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/proxydial"
|
||||
)
|
||||
|
||||
// webFetchTransportWithProxy builds the guarded web_fetch round tripper and
|
||||
// swaps in a proxy function, which is how the transport learns its route. The
|
||||
// environment is left alone deliberately: http.ProxyFromEnvironment reads it
|
||||
// once per process, so an environment-driven test would depend on test order.
|
||||
func webFetchTransportWithProxy(t *testing.T, resolver webFetchResolver, proxyFor func(*http.Request) (*url.URL, error)) *proxydial.Transport {
|
||||
t.Helper()
|
||||
roundTripper := webFetchSafeTransport(nil, resolver)
|
||||
wrapped, ok := roundTripper.(*proxydial.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("transport = %T, want the route-recording wrapper; without it no dial can be recognised as the proxy dial", roundTripper)
|
||||
}
|
||||
wrapped.Proxy = proxyFor
|
||||
return wrapped
|
||||
}
|
||||
|
||||
// THE PROXY IS NOT THE TARGET. With HTTPS_PROXY pointing at a local forward
|
||||
// proxy, the transport dials the proxy and tunnels the request through it. The
|
||||
// dialer used to refuse that dial as a loopback address, so every fetch on a
|
||||
// machine behind such a proxy failed with "loopback hosts are blocked" for a
|
||||
// public target that had already been validated (#569).
|
||||
func TestWebFetchTunnelsThroughTheConfiguredProxy(t *testing.T) {
|
||||
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer listener.Close()
|
||||
arrived := make(chan string, 1)
|
||||
go func() {
|
||||
conn, err := listener.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
buffer := make([]byte, 256)
|
||||
n, _ := conn.Read(buffer)
|
||||
arrived <- string(buffer[:n])
|
||||
}()
|
||||
proxyURL, err := url.Parse("http://" + listener.Addr().String())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// A proxied fetch never resolves the target on this side; the CONNECT
|
||||
// carries the hostname and the proxy resolves it.
|
||||
transport := webFetchTransportWithProxy(t,
|
||||
webFetchResolverFunc(func(context.Context, string, string) ([]netip.Addr, error) {
|
||||
return nil, errors.New("resolver must not be called for a proxied fetch")
|
||||
}),
|
||||
func(*http.Request) (*url.URL, error) { return proxyURL, nil })
|
||||
client := &http.Client{Transport: transport, Timeout: 3 * time.Second}
|
||||
|
||||
_, err = client.Get("https://public.example/resource")
|
||||
select {
|
||||
case first := <-arrived:
|
||||
if line := strings.SplitN(first, "\r\n", 2)[0]; line != "CONNECT public.example:443 HTTP/1.1" {
|
||||
t.Fatalf("proxy received %q, want a CONNECT for the target", line)
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatalf("the proxy never received a connection; request error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// AND A DIRECT FETCH STAYS PINNED, EVEN WHEN ITS ORIGIN IS THE PROXY'S OWN
|
||||
// ADDRESS.
|
||||
//
|
||||
// This is the NO_PROXY shape: the proxy function sends one host direct. An
|
||||
// exemption derived from asking that function about a request of our own
|
||||
// invention would answer "proxied" for the excluded host too, and this dial
|
||||
// would skip the pin and reach whatever the name resolves to at connect time.
|
||||
func TestWebFetchDirectRequestToTheProxyEndpointIsStillPinned(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://rebind.example:80")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
lookups := 0
|
||||
transport := webFetchTransportWithProxy(t,
|
||||
webFetchResolverFunc(func(_ context.Context, _ string, host string) ([]netip.Addr, error) {
|
||||
lookups++
|
||||
return []netip.Addr{netip.MustParseAddr("127.0.0.1")}, nil
|
||||
}),
|
||||
func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Hostname() == "rebind.example" {
|
||||
return nil, nil
|
||||
}
|
||||
return proxyURL, nil
|
||||
})
|
||||
client := &http.Client{Transport: transport, Timeout: 3 * time.Second}
|
||||
|
||||
_, err = client.Get("http://rebind.example/data")
|
||||
|
||||
if lookups == 0 {
|
||||
t.Fatal("the direct dial skipped the address pin because its address matched the configured proxy endpoint")
|
||||
}
|
||||
if err == nil || !strings.Contains(err.Error(), "loopback hosts are blocked") {
|
||||
t.Fatalf("err = %v, want the private answer to be refused", err)
|
||||
}
|
||||
}
|
||||
|
||||
// A proxy configured with a Unicode hostname is dialed by its IDNA-encoded
|
||||
// name, and it is still the proxy.
|
||||
func TestWebFetchUnicodeProxyHostnameIsStillTheProxy(t *testing.T) {
|
||||
dialed := ""
|
||||
transport := webFetchTransportWithProxy(t,
|
||||
webFetchResolverFunc(func(context.Context, string, string) ([]netip.Addr, error) {
|
||||
return nil, errors.New("resolver must not be called for a proxied fetch")
|
||||
}),
|
||||
func(*http.Request) (*url.URL, error) { return url.Parse("http://büro.example:3067") })
|
||||
guarded := transport.DialContext
|
||||
transport.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
dialed = address
|
||||
return guarded(ctx, network, address)
|
||||
}
|
||||
client := &http.Client{Transport: transport, Timeout: 3 * time.Second}
|
||||
|
||||
_, err := client.Get("https://public.example/resource")
|
||||
|
||||
if dialed != "xn--bro-hoa.example:3067" {
|
||||
t.Fatalf("SETUP INVALID: dialed %q, want the IDNA-encoded proxy endpoint", dialed)
|
||||
}
|
||||
if err != nil && strings.Contains(err.Error(), "hosts are blocked") {
|
||||
t.Fatalf("the configured proxy was refused by the guard: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Without a proxy the guard is exactly what it was: loopback is refused and
|
||||
// the dialer never runs.
|
||||
func TestWebFetchSafeDialStillRefusesLoopbackWithoutAProxy(t *testing.T) {
|
||||
dialCalled := false
|
||||
dial := webFetchSafeDialContext(
|
||||
webFetchResolverFunc(func(_ context.Context, network string, host string) ([]netip.Addr, error) {
|
||||
t.Fatalf("resolver consulted for a loopback literal: network=%q host=%q", network, host)
|
||||
return nil, nil
|
||||
}),
|
||||
webFetchDialFunc(func(context.Context, string, string) (net.Conn, error) {
|
||||
dialCalled = true
|
||||
return nil, errors.New("dial should not run")
|
||||
}),
|
||||
)
|
||||
|
||||
_, err := dial(context.Background(), "tcp", "127.0.0.1:3067")
|
||||
if err == nil || !strings.Contains(err.Error(), "loopback hosts are blocked") {
|
||||
t.Fatalf("expected loopback rejection, got %v", err)
|
||||
}
|
||||
if dialCalled {
|
||||
t.Fatal("loopback was dialed with no proxy configured")
|
||||
}
|
||||
}
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/proxydial"
|
||||
zeroSandbox "github.com/Gitlawb/zero/internal/sandbox"
|
||||
)
|
||||
|
||||
@@ -325,15 +326,26 @@ func TestWebFetchToolConfiguresDialTimeSafetyForDefaultTransport(t *testing.T) {
|
||||
}
|
||||
|
||||
client := tool.clientForRun()
|
||||
transport, ok := client.Transport.(*http.Transport)
|
||||
// The round tripper is the route-recording wrapper, not the bare transport:
|
||||
// it is what tells the dialer which dial is the one to the proxy the
|
||||
// transport picked for a given request. The guarded transport it wraps
|
||||
// stays reachable, and is what the rest of this test inspects.
|
||||
transport, ok := client.Transport.(*proxydial.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("client transport = %T, want *http.Transport", client.Transport)
|
||||
t.Fatalf("client transport = %T, want *proxydial.Transport", client.Transport)
|
||||
}
|
||||
if transport.DialContext == nil {
|
||||
t.Fatal("expected web_fetch transport to install a safe DialContext")
|
||||
}
|
||||
if transport.Proxy != nil {
|
||||
t.Fatal("expected web_fetch transport to disable proxy resolution")
|
||||
// REVERSED ON PURPOSE. The guarded fetch used to pin Proxy to nil as part of
|
||||
// the SSRF hardening, and this assertion held it there. #569 is what that
|
||||
// cost: on a machine that reaches the network only through a configured
|
||||
// forward proxy, every fetch failed outright. The proxy is honoured now, and
|
||||
// the wrapper hands the dialer the route chosen for each request so the one
|
||||
// dial it lets past the guard is the one to that proxy. The target is still
|
||||
// validated before the request, which the tests below pin.
|
||||
if transport.Proxy == nil {
|
||||
t.Fatal("expected web_fetch transport to honour the environment proxy; a machine behind a forward proxy cannot fetch at all without it (#569)")
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user