Compare commits

..
Author SHA1 Message Date
Vasanthdev2004 221fb93041 fix(proxydial): bind the proxy exemption to the route the transport chose
Asking the Proxy function a question of our own invention answered 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, so a direct dial whose address happened to equal the configured
proxy endpoint skipped the guard and reached whatever the name resolved to
at connect time. A probe on the previous head dials rebind.example:80 with
zero safety lookups; the guard now runs and refuses its private answer.
Host-dependent proxy callbacks make the mismatch general, and no spelling
of the probe host closes it, because the two calls are about different
requests.

The decision is now recorded where it is made. proxydial.Transport wraps
the transport, asks its Proxy function about the actual request, and puts
the answer in that request's context; 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 the context key is unexported so nothing
outside the package can plant one.

Address matching is gone with it. A transport opens one connection dial,
and when a proxy is selected that dial goes to the proxy, so the recorded
route already identifies the dial. Comparing the address on top could only
disagree with the transport, which is what a Unicode proxy hostname did:
Go dials xn--bro-hoa.example while the URL still spells it as configured,
and the proxy was refused as an ordinary private address. Both spellings
are now one endpoint with no comparison to get wrong.

Documents where destination enforcement ends: with a proxy configured, the
target hostname is sent to the proxy and the proxy resolves it, so the
local pre-request checks no longer bind the address that is reached. A
forward proxy already sees and can rewrite every request through it, so
that trust comes with configuring one; the alternative is refusing to use
proxies at all, which is the outage in #569.
2026-09-09 11:03:14 +05:30
Vasanthdev2004 ac3958f860 fix(proxydial): have the dialer ask its transport for the proxy at dial time
Copying transport.Proxy into the dialer at construction left the wiring
between the two unpinned: a dialer handed nil would pass every unit test
while a proxied request still died as loopback. The dialer now consults
the transport it is installed on, and each package drives a real request
through a loopback listener acting as the proxy and checks that the
CONNECT arrives. The proxy is substituted on the transport rather than
through HTTPS_PROXY, since http.ProxyFromEnvironment reads the
environment once per process and an env-based test would depend on
test order.
2026-09-08 10:04:41 +05:30
Vasanthdev2004andbaoyu0 63767f3e42 fix(providerhealth,tools): let the configured forward proxy through the loopback guard
With HTTPS_PROXY pointing at a local forward proxy, the transport dials the
proxy first and tunnels the request through it. The SSRF-guarded dialer in
providerhealth refused that dial as a loopback address, so
`zero doctor --connectivity` failed with "proxyconnect tcp: provider
connectivity URL is unsafe: loopback hosts are blocked" for a public
provider that had already been validated, and web_fetch had no proxy
support at all (#569).

Each dialer is now handed its transport's own Proxy function and lets a dial
through unpinned only when the address is exactly the proxy endpoint that
function names for the request scheme: same host, same port, with a missing
proxy port taken as the scheme default. A sibling loopback port is still
refused. The shared check lives in internal/proxydial since providerhealth
and tools otherwise share nothing.

web_fetch honours http.ProxyFromEnvironment again. The guarded fetch had set
Proxy to nil as part of its SSRF hardening; the target is still validated
before the request, and what a forward proxy removes is only the dial-time
rebind pin, which is inherent to using any forward proxy. The pre-permission
URL validation stays in place.

Started from #594 by baoyu0. This version keeps the pre-permission check,
matches the proxy port exactly instead of any port when the proxy URL has
none, and injects the transport's Proxy function instead of reading the
environment a second time.

Closes #569

Co-authored-by: baoyu0 <[email protected]>
2026-09-08 09:59:50 +05:30
16 changed files with 56 additions and 1731 deletions

No files matched your search

+6 -52
View File
@@ -14,7 +14,6 @@ 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"
@@ -31,12 +30,10 @@ 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 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)
// 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)
// 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)
@@ -44,27 +41,6 @@ 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
@@ -265,34 +241,13 @@ 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.
// 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)
registry, sandboxEngine, err := a.deps.BuildWorkspace(sess.cwd, resolved)
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 := &notifier{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,
@@ -301,8 +256,7 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string,
Model: resolved.Provider.Model,
Registry: registry,
Sandbox: sandboxEngine,
PermissionMode: mode,
DeferThreshold: workspace.DeferThreshold,
PermissionMode: sess.currentMode(),
MaxTurns: resolved.MaxTurns,
Images: images,
OnText: note.text,
+4 -152
View File
@@ -8,7 +8,6 @@ import (
"io"
"os"
"path/filepath"
"slices"
"strings"
"testing"
"time"
@@ -36,35 +35,6 @@ 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()})
@@ -83,10 +53,10 @@ func testDeps(t *testing.T) Deps {
return fakeProvider{text: "Hello from ZERO"}, nil
},
RunAgent: agent.Run,
BuildWorkspace: func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
BuildWorkspace: func(string, config.ResolvedConfig) (*tools.Registry, *sandbox.Engine, error) {
r := tools.NewRegistry()
r.Register(tools.NewUpdatePlanTool())
return &Workspace{Registry: r}, nil
return r, nil, nil
},
ResolveWorkspaceRoot: func(cwd string) (string, error) { return cwd, nil },
Store: store,
@@ -498,8 +468,8 @@ func TestACPRunTurnWiresSandboxAndScopedRegistry(t *testing.T) {
reg := tools.NewRegistry()
reg.Register(tools.NewUpdatePlanTool())
engine := sandbox.NewEngine(sandbox.EngineOptions{WorkspaceRoot: t.TempDir()})
deps.BuildWorkspace = func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
return &Workspace{Registry: reg, Sandbox: engine}, nil
deps.BuildWorkspace = func(string, config.ResolvedConfig) (*tools.Registry, *sandbox.Engine, error) {
return reg, engine, nil
}
var captured agent.Options
deps.RunAgent = func(_ context.Context, _ string, _ zeroruntime.Provider, opts agent.Options) (agent.Result, error) {
@@ -527,124 +497,6 @@ 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) {
+2 -51
View File
@@ -1,11 +1,9 @@
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
@@ -35,24 +33,7 @@ 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) {
@@ -75,35 +56,8 @@ 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{})
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)
got := estimateToolDefTokens(exposed)
t.Logf("eager core tool schemas: %d tokens across %d tools", got, len(exposed))
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
@@ -112,9 +66,6 @@ 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) {
+15 -75
View File
@@ -10,11 +10,9 @@ 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
@@ -59,8 +57,20 @@ 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,
BuildWorkspace: func(ctx context.Context, workspaceRoot string, resolved config.ResolvedConfig, mode agent.PermissionMode) (*acp.Workspace, error) {
return buildACPWorkspace(ctx, workspaceRoot, resolved, mode, deps)
// 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
},
ResolveWorkspaceRoot: acpWorkspaceRootResolver(deps),
Store: deps.newSessionStore(),
@@ -75,76 +85,6 @@ 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
-201
View File
@@ -5,19 +5,10 @@ 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)
@@ -117,198 +108,6 @@ 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{}
+2 -8
View File
@@ -18,20 +18,14 @@ 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
permission tools.Permission
advertiseInAuto bool
name string
}
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 {
permission := t.permission
if permission == "" {
permission = tools.PermissionAllow
}
return tools.Safety{SideEffect: tools.SideEffectNetwork, Permission: permission, AdvertiseInAuto: t.advertiseInAuto}
return tools.Safety{SideEffect: tools.SideEffectNetwork, Permission: tools.PermissionAllow}
}
func (t cliFakeDeferredTool) Run(context.Context, map[string]any) tools.Result {
return tools.Result{Status: tools.StatusOK, Output: "ok"}
+10 -12
View File
@@ -31,14 +31,6 @@ 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),
@@ -55,10 +47,16 @@ func registerMCPToolsForWorkspaceWithOptions(ctx context.Context, workspaceRoot
if err != nil {
return nil, skip, err
}
options.PermissionStore = store
options.Autonomy = autonomy
options.WorkspaceRoot = workspaceRoot
runtime, err := deps.registerMCPTools(ctx, registry, cfg, options)
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,
})
return runtime, skip, err
}
+3 -7
View File
@@ -23,9 +23,6 @@ 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.
@@ -274,10 +271,9 @@ 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),
AdvertiseInAuto: options.AdvertiseInAuto,
SideEffect: tools.SideEffectNetwork,
Permission: permission,
Reason: fmt.Sprintf("MCP tool %s/%s runs through the configured %s server.", server.Name, remote.Name, server.Type),
},
}
}
-13
View File
@@ -342,19 +342,6 @@ 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.
@@ -1,463 +0,0 @@
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)
}
}
}
@@ -1,227 +0,0 @@
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)
}
}
+8 -295
View File
@@ -4,11 +4,8 @@ import (
"encoding/json"
"fmt"
"log"
"sort"
"strings"
"unicode/utf8"
"github.com/Gitlawb/zero/internal/redaction"
)
type ExecMode string
@@ -191,307 +188,23 @@ 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
lastSpoken := -1
for index, event := range events {
if event.Type == EventMessage && payloadRole(event.Payload) == "assistant" {
lastSpoken = index
}
}
conversation := make([]Event, 0, len(events))
tail := make([]Event, 0, maxPromptContextTailEvents)
for index, event := range events {
filtered := make([]Event, 0, len(events))
for _, 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))
}
filtered = append(filtered, event)
}
}
if len(conversation) == 0 && len(tail) == 0 {
// An unrecognized event stream still says more than nothing.
conversation = append(conversation, events...)
if len(filtered) == 0 {
filtered = append(filtered, events...)
}
if len(conversation) > maxPromptContextEvents {
conversation = conversation[len(conversation)-maxPromptContextEvents:]
if len(filtered) > maxPromptContextEvents {
filtered = filtered[len(filtered)-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)
return filtered
}
func forkTitle(title string) string {
+1 -53
View File
@@ -610,60 +610,8 @@ 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("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)
}
t.Fatalf("expected prompt to omit noisy non-conversation events, got %q", prompt)
}
}
+5 -17
View File
@@ -75,10 +75,14 @@ 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: execCommandDescription(runtimeGOOS()),
description: description,
parameters: Schema{
Type: "object",
Properties: map[string]PropertySchema{
@@ -670,19 +674,3 @@ 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
}
@@ -1,83 +0,0 @@
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)
}
}
}
-22
View File
@@ -181,28 +181,6 @@ 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 {