Compare commits

..
Author SHA1 Message Date
Vasanthdev2004 29a5253013 test(sessions): assert the window key with its value, not the number alone
strings.Contains(out, "40") finds it inside "4096", so the start_line
assertion passed whether or not start_line survived the projection, which
is the opposite of what it claimed. The rendered arguments carry the pairs,
so the assertions name them.
2026-09-11 13:17:19 +05:30
Vasanthdev2004 789bd389aa fix(sessions): name the value shapes a projected argument may keep, and scope the ambiguous aliases
Two halves of the same gap: the projection decided which KEYS survive and
said nothing about the values under them or about which tool gave them
meaning.

Values. Strings were scrubbed and everything else passed 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
{"query":{"api_key":"...","content":"private body"}} arrived under a
permitted key and carried a whole payload into the next prompt. Redacting a
serialized object would not have helped, since ordinary private text has no
credential shape to match. The retained shapes are named now: a scrubbed
string, and the scalars a read window is made of. A container is dropped and
the call keeps its name and id, so a resumed turn still knows which tool ran.

Meaning. grep accepts "search" for its pattern and edit_file accepts the
same word for the text being replaced, so admitting it globally would have
replayed file contents. It is scoped to grep, and the rule is kept as a
test: an argument name any mutating tool accepts for body content belongs in
the per-tool table or nowhere. The unambiguous supported forms are added to
the shared table: glob's match, and read_file's start_line, end_line,
max_lines, byte_offset and byte_limit.

Driven through the store on both resume and fork, including a call the tool
would have rejected, since the agent records OnToolCall before execution. The
persisted events are asserted unchanged: this is a view for the prompt, not
an edit to the record.
2026-09-11 10:11:51 +05:30
Vasanthdev2004 a67ea8735a fix(sessions): project a tool call's top-level fields instead of patching one
toolResultOutcome builds its output from the fields it keeps, so a field it
has never heard of cannot reach a prompt through it. The call side edited
arguments in place and returned every sibling key, and returned the payload
untouched when there were no arguments at all, so a producer recording one
more top-level field would put it into the next turn. Both halves now name
what survives and drop the rest.
2026-09-11 10:11:51 +05:30
Vasanthdev2004 e27174f6f4 fix(sessions): scrub credentials out of retained tool-call values
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,
and this projection then carried it into the next resume or fork prompt.
Retained string values now go through the same redaction web_fetch uses on
its own returned URL: host, path and ordinary query survive, the token does
not, and non-string values are untouched.
2026-09-11 10:11:51 +05:30
Vasanthdev2004 c35412bd01 fix(sessions): keep a tool call's identity and drop its payload in resume context
toolResultOutcome strips a result down to name and status, 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.
Admitting calls is the #913 fix and has to stay: a resumed turn knowing which
file was read is the whole point. What had to change is the symmetry.

toolCallIdentity is the call-side half of toolResultOutcome. It decodes the
arguments, keeps the fields that say what a call was about, and drops the ones
that carry what it was about to write, run or search for. Paths, directories,
urls, names, patterns and a read window survive; write_file's content, edit_file's
old and new strings, apply_patch's hunks and a shell command line do not.

An allow-list rather than a deny-list, so an argument this file has never seen is
dropped rather than replayed. A new tool with a new body field then shows up as a
missing path in a resume prompt, which is visible, instead of a new leak, which
is not. Keyed by argument name rather than tool name, so an MCP tool whose
argument is a url keeps it without being enumerated. Arguments that do not decode
as an object are removed outright: text that could not be read is text that
cannot be checked.

The tail budget gets a comment rather than a change. Tail slots exist only when
the conversation uses fewer than the 80-event cap, so a session already at the
cap carries no interrupted tool work, which is what every session did before
tool events were admitted and what the #460 cap is there to hold. Reserving a
minimum tail by trimming conversation further is a product decision, and it is
not this one.

Reported by jatmn.
2026-09-11 10:11:51 +05:30
Vasanthdev2004 dab4691b21 test(sessions): pin the whole of every tool line, not a search for one string
Two assertions here were weaker than what they claimed.

Searching the prompt for the fixture's whole secret only rejects that exact
string, so a change that carried a truncated or partly redacted prefix of the
output into the prompt would have passed. And asking whether "ok" and "error"
each appear somewhere among the result lines passes just as well with the two
statuses swapped, or both hung off the wrong call.

Both tool calls and both results are now pinned outright, which says what the
trim actually promises: nothing but the tool name and the outcome survives a
result, and the line order pairs each outcome with its own call. 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 come out in a
different field order on every run.

Reported by CodeRabbit on this PR.
2026-09-11 10:11:51 +05:30
Vasanthdev2004 23b1e5f92f test(sessions): scope the outcome assertion to the tool_result lines
The status check searched the whole prompt for "error", which the provider error
message in the same fixture also contains. Blanking the status out left the test
passing, so it was asserting the presence of an unrelated line.

Scoped to the tool_result lines, it fails on that mutation naming the outcome it
lost.
2026-09-11 10:11:51 +05:30
Vasanthdev2004 1e52a3625d fix(sessions): carry the tool outcome, not the tool output
The tail was carrying whole tool-result events, so up to 500 bytes of raw tool
output rode into a later turn prompt. Nothing redacts on the way in, and the
prompt renderer truncates rather than filters, so the only thing standing between
a file read and a later prompt was a length limit.

It bought almost nothing. The tail exists so a resumed turn knows what the
interrupted one did, and the CALL already says that: the tool name and its
arguments, which is the path for a read. The result body is the one part of the
pair that can carry file contents.

The status is kept rather than dropping the result entirely. 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.

Raised by CodeRabbit on the PR, and right.
2026-09-11 10:11:51 +05:30
Vasanthdev2004 139cffa6c3 fix(sessions): carry an interrupted turn's work into the resume prompt
A turn that read six files and then died on a provider error left this behind
for the next turn:

  - #1 message: add retries to the http client
  - #6 error: provider error: upstream timeout

Nothing named the files, so the next turn re-read them from scratch (#913).

The work was never lost from the session: tool calls and results are recorded as
they happen, and the error path carries them out with everything else.
promptContextEvents then dropped both types on the way into the prompt, so the
record existed and nothing read it.

Filtering them was deliberate (#460) and is right for work an answer already
describes: forty read_file results add length and no information next to the
assistant message explaining them. It is wrong only for work with no answer
after it, and that is exactly what an interruption leaves. A turn that ends
normally was survivable for that reason, which is why this went unnoticed: the
answer was doing the remembering.

So the events after the last assistant message come along and the rest stays
filtered. The tail gets its own allowance rather than sharing the conversation
budget, so a tool-heavy interrupted turn still cannot push an earlier message
out of the context. That is the property #460 added the filter for and it is
asserted directly.

This changes a tested contract. TestFormatExecPromptKeepsConversationMessagesWhenNoisyEventsFollow
asserted that tool results are omitted, and its fixture is an interrupted turn:
43 results after an assistant answer with nothing speaking for them. That is the
#913 shape, so no rule satisfies both and the contract had to move. The test
keeps the guarantee it was written for, that conversation survives a noisy turn,
and a counterpart now pins the other half: work an answer already covers is
still filtered.

Not fixed here, and worth saying plainly: this carries what was DONE, not what
was READ. Zero has no conversation-history parameter at all, agent.Run seeds
from the system prompt and the user prompt every time, so file contents cannot
survive a turn boundary today. The next turn learns it already read a path and
can decide, rather than starting blind.
2026-09-11 10:11:50 +05:30
Vasanth Tandbaoyu0 1e4db7c664 fix: let the configured forward proxy through the loopback guard (#1025)
* 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]>

* 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.

* 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.

---------

Co-authored-by: baoyu0 <[email protected]>
2026-09-10 21:50:26 +05:30
Vasanth T f9092668a8 fix(agent): stop the tool-schema ratchet meaning a different thing per host (#1017)
* fix(agent): stop the tool-schema ratchet meaning a different thing per host

TestEagerToolSchemaTokenBudget counted exec_command shell guidance along with
the schemas, and that guidance is host-specific. The same tree measured 3498 on
Linux, about 3600 on a Windows host with pwsh 7, and 3653 on Windows PowerShell
5.1, because shellGuidanceForRuntime appends a longer block for 5.1 (it has to
explain that && and || do not work there).

5.1 is stock Windows, and that case has been over the 3650 ceiling since
view_image landed. So `go test ./...` failed for anyone on a stock Windows box
while CI stayed green: the runners have pwsh 7 and take the cheaper branch. The
workflow does run the tests on windows-latest, so this was never a coverage gap,
just a ceiling sized against whichever shell the runner happened to have.

The two halves are now charged separately. The schema ceiling measures the
schemas alone and reads 3498 on Linux and on Windows PowerShell 5.1, so it
ratchets the same thing everywhere. The host guidance gets its own ceiling at
200, above the worst case rather than above this machine: 5.1 is 155 tokens,
pwsh 7 is 110, cmd.exe is 87, POSIX is 0.

No guidance text was trimmed. It is what stops the model writing bash on Windows,
and the 5.1 note is the reason a chained command does not silently do half its
work. Paying 155 tokens for it is the right trade; measuring it as if it were a
schema was not.

HostExecCommandShellGuidance is exported for the budget, which is production
surface added for a test. It is one accessor over existing logic, and the
alternative is a ratchet that cannot tell the two costs apart. A guard in
internal/tools pins that the guidance really is the tail of the description, so
the subtraction cannot silently start measuring the wrong thing.

* test(tools,agent): make the guidance contract checkable off Windows

The ratchet separates the fixed tool schemas from this host's exec_command
shell guidance, and both halves of that separation were only being checked on
Windows.

The description test returned early off Windows once the accessor reported no
guidance, so it never looked at the description. Dropping the GOOS condition in
the constructor would have a POSIX host append its own guidance (about 150
characters on darwin) with the accessor still reporting none and the ratchet
charging every byte of it to the schemas, and nothing would fail. Both the
description and the accessor now take goos, so the contract is checked for all
three platforms from any one of them.

The budget also went from subtracting a separately rounded guidance estimate to
removing the guidance from the description before estimating. ApproxTextTokens
is an integer division, so floor((schema+guidance)/4) - floor(guidance/4) is one
token high whenever the two remainders add past four, and which side of that a
host lands on depends on the length of its guidance. That is the same
host-dependence the ratchet was being fixed for, one token smaller.

No behaviour change: exec_command's description is byte for byte what it was on
every platform.
2026-09-10 21:43:58 +05:30
KRATOS 7f5e5af7bb fix(acp): make configured MCP tools callable (#1003)
* fix(acp): register configured MCP tools per turn

* fix(acp): address MCP review findings
2026-09-10 20:17:26 +05:30
16 changed files with 1733 additions and 58 deletions

No files matched your search

+52 -6
View File
@@ -14,6 +14,7 @@ import (
"github.com/Gitlawb/zero/internal/providercatalog"
"github.com/Gitlawb/zero/internal/providermodelcatalog"
"github.com/Gitlawb/zero/internal/providermodeldiscovery"
"github.com/Gitlawb/zero/internal/redaction"
"github.com/Gitlawb/zero/internal/sandbox"
"github.com/Gitlawb/zero/internal/sessions"
"github.com/Gitlawb/zero/internal/tools"
@@ -30,10 +31,12 @@ type Deps struct {
DiscoverModels func(context.Context, config.ProviderProfile) ([]providermodeldiscovery.Model, error)
NewProvider func(profile config.ProviderProfile) (zeroruntime.Provider, error)
RunAgent func(ctx context.Context, prompt string, provider zeroruntime.Provider, opts agent.Options) (agent.Result, error)
// BuildWorkspace builds the SCOPED tool registry and the sandbox engine for a
// validated workspace root, so ACP shell tools (bash/exec_command) are confined
// exactly like the exec surface — never run unconfined on the host.
BuildWorkspace func(workspaceRoot string, resolved config.ResolvedConfig) (*tools.Registry, *sandbox.Engine, error)
// BuildWorkspace creates the per-turn scoped tool workspace for a validated
// workspace root. Its Close method releases any resources the registry owns
// (notably MCP server connections) after the turn completes or is cancelled.
// Shell tools remain confined exactly like the exec surface — never unconfined
// on the host.
BuildWorkspace func(ctx context.Context, workspaceRoot string, resolved config.ResolvedConfig, mode agent.PermissionMode) (*Workspace, error)
// ResolveWorkspaceRoot validates + normalizes a client-supplied cwd (must be an
// existing directory; never the bare root). It is the file-tool confinement root.
ResolveWorkspaceRoot func(cwd string) (string, error)
@@ -41,6 +44,27 @@ type Deps struct {
AgentInfo Implementation
}
// Workspace is the per-turn execution environment passed to the agent. ACP does
// not retain it across turns because MCP connections belong to the registry that
// advertised their tools; keeping a stale registry after a cancelled turn would
// leak its server process and its permission state into the next turn.
type Workspace struct {
Registry *tools.Registry
Sandbox *sandbox.Engine
DeferThreshold int
Notices []string
Cleanup func() error
}
// Close releases resources created with the workspace. A nil cleanup is valid
// for core-only workspaces.
func (w *Workspace) Close() error {
if w == nil || w.Cleanup == nil {
return nil
}
return w.Cleanup()
}
// Agent is the ACP agent server bound to one JSON-RPC connection (one editor).
type Agent struct {
conn *Conn
@@ -241,13 +265,34 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string,
if err != nil {
return "", RPCError(codeInternalError, "provider: "+err.Error())
}
mode := sess.currentMode()
// Build the SCOPED registry + sandbox engine for this session's workspace so
// shell/file tools are confined to the workspace exactly like the exec surface.
registry, sandboxEngine, err := a.deps.BuildWorkspace(sess.cwd, resolved)
// The workspace can also own MCP connections, which must be closed after every
// turn even if the agent run returns an error or the client cancels it.
workspace, err := a.deps.BuildWorkspace(ctx, sess.cwd, resolved, mode)
if err != nil {
return "", RPCError(codeInternalError, "workspace: "+err.Error())
}
if workspace == nil || workspace.Registry == nil {
if workspace != nil {
_ = workspace.Close()
}
return "", RPCError(codeInternalError, "workspace: missing tool registry")
}
defer func() {
if err := workspace.Close(); err != nil {
log.Printf("acp: close workspace: %v", err)
}
}()
registry := workspace.Registry
sandboxEngine := workspace.Sandbox
note := &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,
@@ -256,7 +301,8 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string,
Model: resolved.Provider.Model,
Registry: registry,
Sandbox: sandboxEngine,
PermissionMode: sess.currentMode(),
PermissionMode: mode,
DeferThreshold: workspace.DeferThreshold,
MaxTurns: resolved.MaxTurns,
Images: images,
OnText: note.text,
+152 -4
View File
@@ -8,6 +8,7 @@ import (
"io"
"os"
"path/filepath"
"slices"
"strings"
"testing"
"time"
@@ -35,6 +36,35 @@ func (f fakeProvider) StreamCompletion(_ context.Context, _ zeroruntime.Completi
return ch, nil
}
type acpDeferredPromptTool struct{ name string }
func (t acpDeferredPromptTool) Name() string { return t.name }
func (acpDeferredPromptTool) Description() string { return "deferred MCP test tool" }
func (acpDeferredPromptTool) Parameters() tools.Schema { return tools.Schema{Type: "object"} }
func (acpDeferredPromptTool) Deferred() bool { return true }
func (acpDeferredPromptTool) Run(context.Context, map[string]any) tools.Result {
return tools.Result{Status: tools.StatusOK, Output: "ok"}
}
func (acpDeferredPromptTool) Safety() tools.Safety {
return tools.Safety{
SideEffect: tools.SideEffectNetwork, Permission: tools.PermissionPrompt,
AdvertiseInAuto: true,
}
}
type captureToolsProvider struct {
requests chan zeroruntime.CompletionRequest
}
func (p captureToolsProvider) StreamCompletion(_ context.Context, request zeroruntime.CompletionRequest) (<-chan zeroruntime.StreamEvent, error) {
p.requests <- request
ch := make(chan zeroruntime.StreamEvent, 2)
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventText, Content: "done"}
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}
close(ch)
return ch, nil
}
func testDeps(t *testing.T) Deps {
t.Helper()
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
@@ -53,10 +83,10 @@ func testDeps(t *testing.T) Deps {
return fakeProvider{text: "Hello from ZERO"}, nil
},
RunAgent: agent.Run,
BuildWorkspace: func(string, config.ResolvedConfig) (*tools.Registry, *sandbox.Engine, error) {
BuildWorkspace: func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
r := tools.NewRegistry()
r.Register(tools.NewUpdatePlanTool())
return r, nil, nil
return &Workspace{Registry: r}, nil
},
ResolveWorkspaceRoot: func(cwd string) (string, error) { return cwd, nil },
Store: store,
@@ -468,8 +498,8 @@ func TestACPRunTurnWiresSandboxAndScopedRegistry(t *testing.T) {
reg := tools.NewRegistry()
reg.Register(tools.NewUpdatePlanTool())
engine := sandbox.NewEngine(sandbox.EngineOptions{WorkspaceRoot: t.TempDir()})
deps.BuildWorkspace = func(string, config.ResolvedConfig) (*tools.Registry, *sandbox.Engine, error) {
return reg, engine, nil
deps.BuildWorkspace = func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
return &Workspace{Registry: reg, Sandbox: engine}, nil
}
var captured agent.Options
deps.RunAgent = func(_ context.Context, _ string, _ zeroruntime.Provider, opts agent.Options) (agent.Result, error) {
@@ -497,6 +527,124 @@ func TestACPRunTurnWiresSandboxAndScopedRegistry(t *testing.T) {
}
}
func TestACPRunTurnUsesWorkspaceDeferThresholdAtProviderBoundary(t *testing.T) {
for _, tc := range []struct {
name string
threshold int
want []string
}{
{name: "explicit zero disables", threshold: 0, want: []string{"mcp_docs_alpha", "mcp_docs_beta"}},
{name: "below threshold stays eager", threshold: 3, want: []string{"mcp_docs_alpha", "mcp_docs_beta"}},
{name: "at threshold defers", threshold: 2, want: []string{tools.ToolSearchToolName}},
} {
t.Run(tc.name, func(t *testing.T) {
requests := make(chan zeroruntime.CompletionRequest, 4)
deps := testDeps(t)
deps.NewProvider = func(config.ProviderProfile) (zeroruntime.Provider, error) {
return captureToolsProvider{requests: requests}, nil
}
deps.BuildWorkspace = func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
registry := tools.NewRegistry()
registry.Register(acpDeferredPromptTool{name: "mcp_docs_alpha"})
registry.Register(acpDeferredPromptTool{name: "mcp_docs_beta"})
registry.Register(tools.NewToolSearchTool(registry))
return &Workspace{Registry: registry, DeferThreshold: tc.threshold}, nil
}
h := newHarness(t, deps)
defer h.stop()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
var created NewSessionResult
if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir()}, &created); err != nil {
t.Fatal(err)
}
if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{SessionID: created.SessionID, Prompt: []ContentBlock{TextBlock("hi")}}, &PromptResult{}); err != nil {
t.Fatal(err)
}
request := <-requests
got := make([]string, 0, len(request.Tools))
for _, definition := range request.Tools {
got = append(got, definition.Name)
}
if !slices.Equal(got, tc.want) {
t.Fatalf("provider tools = %#v, want %#v", got, tc.want)
}
})
}
}
func TestACPRunTurnEmitsWorkspaceSetupNotices(t *testing.T) {
deps := testDeps(t)
deps.BuildWorkspace = func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
registry := tools.NewRegistry()
registry.Register(tools.NewUpdatePlanTool())
return &Workspace{Registry: registry, Notices: []string{"MCP server docs unavailable, skipped: timeout"}}, nil
}
h := newHarness(t, deps)
defer h.stop()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
var created NewSessionResult
if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir()}, &created); err != nil {
t.Fatal(err)
}
if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{SessionID: created.SessionID, Prompt: []ContentBlock{TextBlock("hi")}}, &PromptResult{}); err != nil {
t.Fatal(err)
}
got := drainText(t, h.updates)
if !strings.Contains(got, "[zero warning] MCP server docs unavailable, skipped: timeout") || !strings.Contains(got, "Hello from ZERO") {
t.Fatalf("streamed text = %q", got)
}
}
// TestACPRunTurnClosesWorkspace proves a per-turn workspace never outlives the
// agent run. This is particularly important for MCP: its registry owns live
// client connections and, for stdio servers, child processes.
func TestACPRunTurnClosesWorkspaceAfterSuccessAndFailure(t *testing.T) {
for _, tc := range []struct {
name string
runErr error
}{
{name: "success"},
{name: "agent failure", runErr: errors.New("provider interrupted")},
} {
t.Run(tc.name, func(t *testing.T) {
deps := testDeps(t)
closed := 0
deps.BuildWorkspace = func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
registry := tools.NewRegistry()
registry.Register(tools.NewUpdatePlanTool())
return &Workspace{Registry: registry, Cleanup: func() error {
closed++
return nil
}}, nil
}
deps.RunAgent = func(context.Context, string, zeroruntime.Provider, agent.Options) (agent.Result, error) {
return agent.Result{FinalAnswer: "ok"}, tc.runErr
}
h := newHarness(t, deps)
defer h.stop()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
var created NewSessionResult
if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir()}, &created); err != nil {
t.Fatalf("session/new: %v", err)
}
err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{SessionID: created.SessionID, Prompt: []ContentBlock{TextBlock("hello")}}, &PromptResult{})
if tc.runErr == nil && err != nil {
t.Fatalf("session/prompt: %v", err)
}
if tc.runErr != nil && err == nil {
t.Fatal("session/prompt succeeded after agent failure")
}
if closed != 1 {
t.Fatalf("workspace cleanup calls = %d, want 1", closed)
}
})
}
}
// TestACPRejectsInvalidCwd confirms session/new fails when the workspace root
// resolver rejects the client cwd (e.g. filesystem root).
func TestACPRejectsInvalidCwd(t *testing.T) {
+51 -2
View File
@@ -1,9 +1,11 @@
package agent
import (
"strings"
"testing"
"github.com/Gitlawb/zero/internal/tools"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
// These ceilings are a deliberate ratchet on Zero's fixed per-turn overhead — the
@@ -33,7 +35,24 @@ const (
// The ratchet did its job: it caught the creep and forced a decision instead
// of a drift. That is the point of it, so keep raising it deliberately rather
// than reflexively.
//
// THIS COUNTS THE SCHEMAS ONLY, NOT THE HOST'S SHELL GUIDANCE. Until
// 2026-09-07 it counted both, and that made it a different measurement on
// every machine: exec_command appends shell guidance whose size depends on
// which shell was detected, so the same tree measured 3498 on Linux, ~3600 on
// a Windows host with pwsh 7, and 3653 on Windows PowerShell 5.1, which is
// stock Windows. The 5.1 case has been over this ceiling since view_image
// landed, so `go test ./...` failed for anyone on a stock Windows box while
// CI stayed green, because the runners have pwsh 7 and take the cheaper
// branch. A budget that passes or fails on which PowerShell happens to be
// installed is not ratcheting anything.
maxEagerToolSchemaTokens = 3650
// The host-specific half, ratcheted on its own so it cannot creep either.
// Measured 2026-09-07: Windows PowerShell 5.1 is the most expensive host at
// 155 tokens (722 chars), pwsh 7 is 110, cmd.exe is 87, and every POSIX host
// is 0 because no guidance is appended there. Set above the worst case, not
// above whatever this machine happens to be.
maxExecCommandShellGuidanceTokens = 200
)
func TestSystemPromptTokenBudget(t *testing.T) {
@@ -56,8 +75,35 @@ func TestEagerToolSchemaTokenBudget(t *testing.T) {
// Options{} keeps DeferThreshold at 0, so deferral is inactive and every core
// tool is exposed eagerly — exactly what a plugin-free session sends each turn.
exposed, _ := partitionTools(registry, PermissionModeAuto, Options{}, map[string]bool{})
got := estimateToolDefTokens(exposed)
t.Logf("eager core tool schemas: %d tokens across %d tools", got, len(exposed))
total := estimateToolDefTokens(exposed)
// exec_command's description carries this host's shell guidance, and that
// guidance is the one part of the eager set that is not the same everywhere.
// Charged separately below rather than folded in here, or this ceiling would
// mean a different thing on every machine. See the constants above.
//
// Removed from the description BEFORE estimating rather than subtracted from
// the total afterwards. ApproxTextTokens is an integer division, so
// floor((schema+guidance)/4) - floor(guidance/4) comes out one token high
// whenever the two remainders add past four, and which side of that line a
// host lands on depends on the length of its guidance: exactly the
// host-dependence this number is supposed to be rid of.
guidance := tools.HostExecCommandShellGuidance()
schemas := make([]zeroruntime.ToolDefinition, len(exposed))
copy(schemas, exposed)
stripped := 0
for index := range schemas {
without := strings.TrimSuffix(schemas[index].Description, guidance)
if without != schemas[index].Description {
stripped++
}
schemas[index].Description = without
}
if guidance != "" && stripped != 1 {
t.Fatalf("SETUP INVALID: this host's shell guidance came off %d descriptions, want exactly exec_command's; internal/tools pins it as that description's tail", stripped)
}
got := estimateToolDefTokens(schemas)
guidanceTokens := ApproxTextTokens(guidance)
t.Logf("eager core tool schemas: %d tokens across %d tools (%d total, %d host shell guidance)", got, len(exposed), total, guidanceTokens)
if got > maxEagerToolSchemaTokens {
// NOT "defer a tool": this test pins DeferThreshold at 0, so marking a
// tool deferred leaves it exposed here and changes nothing. Deferral is
@@ -66,6 +112,9 @@ func TestEagerToolSchemaTokenBudget(t *testing.T) {
// deliberate raise.
t.Fatalf("eager tool schemas are %d tokens, over the %d ceiling — trim a schema, drop a core tool, or raise the ceiling deliberately (deferring will NOT help: this test disables deferral)", got, maxEagerToolSchemaTokens)
}
if guidanceTokens > maxExecCommandShellGuidanceTokens {
t.Fatalf("this host's exec_command shell guidance is %d tokens, over the %d ceiling — trim the guidance in internal/tools/shell_runtime.go or raise the ceiling deliberately", guidanceTokens, maxExecCommandShellGuidanceTokens)
}
}
func TestAgentAdvertisesBothEditTools(t *testing.T) {
+75 -15
View File
@@ -10,9 +10,11 @@ import (
"github.com/Gitlawb/zero/internal/acp"
"github.com/Gitlawb/zero/internal/agent"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/execution"
"github.com/Gitlawb/zero/internal/mcp"
"github.com/Gitlawb/zero/internal/providermodeldiscovery"
"github.com/Gitlawb/zero/internal/redaction"
"github.com/Gitlawb/zero/internal/sandbox"
"github.com/Gitlawb/zero/internal/tools"
)
const acpUsage = `zero acp — serve the Agent Client Protocol (ACP) over stdio
@@ -57,20 +59,8 @@ func runACP(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int
// surface — no ACP-specific credential handling needed.
NewProvider: deps.newProvider,
RunAgent: agent.Run,
// Build the SCOPED registry + sandbox engine per workspace, exactly like the
// exec surface, so ACP shell/file tools are confined — never run unconfined.
BuildWorkspace: func(workspaceRoot string, resolved config.ResolvedConfig) (*tools.Registry, *sandbox.Engine, error) {
scope, err := sandbox.NewScope(workspaceRoot, resolved.Sandbox.AdditionalWriteRoots)
if err != nil {
return nil, nil, err
}
engine, err := buildExecSandboxEngine(workspaceRoot, resolved, deps, scope)
if err != nil {
return nil, nil, err
}
registry := newCoreRegistryScoped(workspaceRoot, scope)
registerLocalControlTools(registry, workspaceRoot, resolved.LocalControl)
return registry, engine, nil
BuildWorkspace: func(ctx context.Context, workspaceRoot string, resolved config.ResolvedConfig, mode agent.PermissionMode) (*acp.Workspace, error) {
return buildACPWorkspace(ctx, workspaceRoot, resolved, mode, deps)
},
ResolveWorkspaceRoot: acpWorkspaceRootResolver(deps),
Store: deps.newSessionStore(),
@@ -85,6 +75,76 @@ func runACP(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int
return exitSuccess
}
// buildACPWorkspace matches exec's registry construction for one ACP turn. MCP
// servers use the same sandbox-prepared execution runner, project configuration is
// gated by the validated workspace's trust state, and their runtime is released
// when acp.Agent completes the turn. ACP deliberately stays at low autonomy: an
// editor connection never upgrades MCP permissions beyond an interactive prompt.
func buildACPWorkspace(ctx context.Context, workspaceRoot string, resolved config.ResolvedConfig, mode agent.PermissionMode, deps appDeps) (*acp.Workspace, error) {
scope, err := sandbox.NewScope(workspaceRoot, resolved.Sandbox.AdditionalWriteRoots)
if err != nil {
return nil, err
}
engine, err := buildExecSandboxEngine(workspaceRoot, resolved, deps, scope)
if err != nil {
return nil, err
}
registry := newCoreRegistryScoped(workspaceRoot, scope)
workspace := &acp.Workspace{
Registry: registry, Sandbox: engine,
DeferThreshold: resolved.Tools.DeferThreshold,
}
if mode != agent.PermissionModePlan {
// MCP stdio servers are subprocesses. Passing the engine to their runner
// keeps them inside the exact sandbox / lifecycle path used by zero exec.
runtime, trustSkip, err := registerMCPToolsForWorkspaceWithOptions(
ctx, workspaceRoot, registry, deps, mcp.AutonomyLow, workspaceRoot,
mcp.RegisterOptions{Execution: execution.NewRunner(engine), AdvertiseInAuto: true},
)
if err != nil {
// RegisterTools may have connected an earlier server before reporting a
// later failure, so do not orphan a partial runtime on this error path.
if runtime != nil {
_ = runtime.Close()
}
return nil, err
}
workspace.Cleanup = runtime.Close
workspace.Notices = acpMCPSetupNotices(trustSkip, runtime)
}
registerLocalControlTools(registry, workspaceRoot, resolved.LocalControl)
// MCP tools are deferred-eligible. Register their loader only after every
// ACP-visible tool is present, using the same mode the agent receives.
registerToolSearchIfEligible(registry, resolved.Tools.DeferThreshold, mode, nil, nil)
return workspace, nil
}
func acpMCPSetupNotices(skip trustSkip, runtime mcpToolRuntime) []string {
var notices []string
if skip.excludedProjectConfig {
if skip.trustCheckErrored {
notices = append(notices, "The workspace-trust store could not be read; project MCP servers were ignored (fail-closed). Run 'zero trust' to enable them.")
} else {
notices = append(notices, "Project MCP servers were ignored in this untrusted workspace. Run 'zero trust' to enable them.")
}
}
if runtime == nil {
return notices
}
for _, skipped := range runtime.Skipped() {
if skipped.UnconfiguredDefault {
continue
}
message := fmt.Sprintf("MCP server %s unavailable, skipped", skipped.Name)
if skipped.Err != nil {
message += ": " + redaction.ErrorMessage(skipped.Err, redaction.Options{})
}
notices = append(notices, redaction.RedactString(message, redaction.Options{}))
}
return notices
}
// acpWorkspaceRootResolver validates a client-supplied cwd into a confinement
// root. It reuses exec's resolveWorkspaceRoot (abs+clean, must be an existing
// dir) and additionally rejects the filesystem root and the home directory — an
+201
View File
@@ -5,10 +5,19 @@ import (
"context"
"errors"
"io"
"os"
"path/filepath"
"strings"
"sync"
"testing"
"time"
"github.com/Gitlawb/zero/internal/agent"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/mcp"
"github.com/Gitlawb/zero/internal/sandbox"
"github.com/Gitlawb/zero/internal/tools"
"github.com/Gitlawb/zero/internal/workspacetrust"
)
type acpTestReader func([]byte) (int, error)
@@ -108,6 +117,198 @@ func TestRunACPIdleCancellationExitsCleanly(t *testing.T) {
}
}
func TestBuildACPWorkspaceRegistersSandboxedMCPToolsAndClosesThem(t *testing.T) {
setTrustConfigRoot(t)
workspaceRoot := t.TempDir()
if err := workspacetrust.Trust(workspaceRoot); err != nil {
t.Fatalf("trust workspace: %v", err)
}
grantStore, err := sandbox.NewGrantStore(sandbox.StoreOptions{FilePath: t.TempDir() + "/grants.json"})
if err != nil {
t.Fatalf("new grant store: %v", err)
}
var gotExclude, registered, closed bool
deps := fillAppDeps(appDeps{
resolveMCPConfig: func(root string, excludeProject bool) (config.MCPConfig, error) {
if root != workspaceRoot {
t.Fatalf("MCP workspace root = %q, want %q", root, workspaceRoot)
}
gotExclude = excludeProject
return config.MCPConfig{Servers: map[string]config.MCPServerConfig{
"docs": {Type: "stdio", Command: "fake-docs"},
}}, nil
},
newMCPStore: func() (*mcp.PermissionStore, error) { return nil, nil },
newSandboxStore: func() (*sandbox.GrantStore, error) { return grantStore, nil },
registerMCPTools: func(ctx context.Context, registry *tools.Registry, cfg config.MCPConfig, options mcp.RegisterOptions) (mcpToolRuntime, error) {
if ctx == nil {
t.Fatal("MCP registration received nil context")
}
if len(cfg.Servers) != 1 || cfg.Servers["docs"].Command != "fake-docs" {
t.Fatalf("MCP config = %+v", cfg)
}
if options.Autonomy != mcp.AutonomyLow {
t.Fatalf("MCP autonomy = %q, want low", options.Autonomy)
}
if options.Execution == nil {
t.Fatal("MCP stdio server was not given the sandbox execution runner")
}
if options.WorkspaceRoot != workspaceRoot {
t.Fatalf("MCP execution workspace = %q, want %q", options.WorkspaceRoot, workspaceRoot)
}
if !options.AdvertiseInAuto {
t.Fatal("ACP MCP tools must be advertised in auto mode without being pre-approved")
}
registered = true
registry.Register(cliFakeDeferredTool{
name: "mcp_docs_search", permission: tools.PermissionPrompt,
advertiseInAuto: options.AdvertiseInAuto,
})
return closeFunc(func() error {
closed = true
return nil
}), nil
},
})
workspace, err := buildACPWorkspace(context.Background(), workspaceRoot, config.ResolvedConfig{
Tools: config.ToolsConfig{DeferThreshold: 1},
}, agent.PermissionModeAuto, deps)
if err != nil {
t.Fatalf("buildACPWorkspace: %v", err)
}
if gotExclude {
t.Fatal("trusted ACP workspace excluded project MCP configuration")
}
if !registered {
t.Fatal("ACP workspace did not register configured MCP tools")
}
if _, ok := workspace.Registry.Get("mcp_docs_search"); !ok {
t.Fatal("MCP tool is absent from ACP agent registry")
}
if _, ok := workspace.Registry.Get(tools.ToolSearchToolName); !ok {
t.Fatal("ACP registry is missing tool_search for deferred MCP tools")
}
if workspace.DeferThreshold != 1 {
t.Fatalf("workspace defer threshold = %d, want 1", workspace.DeferThreshold)
}
if err := workspace.Close(); err != nil {
t.Fatalf("close ACP workspace: %v", err)
}
if !closed {
t.Fatal("ACP workspace did not close its MCP runtime")
}
}
func TestBuildACPWorkspacePreservesRedactedMCPSetupNotices(t *testing.T) {
setTrustConfigRoot(t)
workspaceRoot := t.TempDir()
if err := os.MkdirAll(filepath.Join(workspaceRoot, ".zero"), 0o755); err != nil {
t.Fatal(err)
}
projectConfig := `{"mcp":{"servers":{"project-docs":{"type":"stdio","command":"project-docs"}}}}`
if err := os.WriteFile(filepath.Join(workspaceRoot, ".zero", "config.json"), []byte(projectConfig), 0o600); err != nil {
t.Fatal(err)
}
secret := "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
grantStore, err := sandbox.NewGrantStore(sandbox.StoreOptions{FilePath: filepath.Join(t.TempDir(), "grants.json")})
if err != nil {
t.Fatal(err)
}
deps := fillAppDeps(appDeps{
resolveMCPConfig: func(string, bool) (config.MCPConfig, error) {
return config.MCPConfig{Servers: map[string]config.MCPServerConfig{
"docs": {Type: "stdio", Command: "fake-docs"},
}}, nil
},
newMCPStore: func() (*mcp.PermissionStore, error) { return nil, nil },
newSandboxStore: func() (*sandbox.GrantStore, error) { return grantStore, nil },
registerMCPTools: func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
return fakeMCPRuntimeWithSkips{skipped: []mcp.SkippedServer{{
Name: "docs", Err: errors.New("token=" + secret),
}}}, nil
},
})
workspace, err := buildACPWorkspace(context.Background(), workspaceRoot, config.ResolvedConfig{}, agent.PermissionModeAuto, deps)
if err != nil {
t.Fatal(err)
}
defer workspace.Close()
joined := strings.Join(workspace.Notices, "\n")
if !strings.Contains(joined, "untrusted workspace") || !strings.Contains(joined, "MCP server docs unavailable") {
t.Fatalf("workspace notices = %#v", workspace.Notices)
}
if strings.Contains(joined, secret) || !strings.Contains(joined, "[REDACTED]") {
t.Fatalf("workspace notices leaked secret: %q", joined)
}
}
func TestBuildACPWorkspaceSkipsMCPInPlanMode(t *testing.T) {
grantStore, err := sandbox.NewGrantStore(sandbox.StoreOptions{FilePath: t.TempDir() + "/grants.json"})
if err != nil {
t.Fatalf("new grant store: %v", err)
}
called := false
deps := fillAppDeps(appDeps{
resolveMCPConfig: func(string, bool) (config.MCPConfig, error) {
called = true
return config.MCPConfig{}, nil
},
newSandboxStore: func() (*sandbox.GrantStore, error) { return grantStore, nil },
registerMCPTools: func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
called = true
return nil, nil
},
})
workspace, err := buildACPWorkspace(context.Background(), t.TempDir(), config.ResolvedConfig{}, agent.PermissionModePlan, deps)
if err != nil {
t.Fatalf("buildACPWorkspace: %v", err)
}
if called {
t.Fatal("plan-mode ACP workspace must not resolve or start MCP servers")
}
if err := workspace.Close(); err != nil {
t.Fatalf("close plan workspace: %v", err)
}
}
func TestBuildACPWorkspaceClosesPartialMCPRuntimeOnRegistrationError(t *testing.T) {
grantStore, err := sandbox.NewGrantStore(sandbox.StoreOptions{FilePath: t.TempDir() + "/grants.json"})
if err != nil {
t.Fatalf("new grant store: %v", err)
}
closed := false
deps := fillAppDeps(appDeps{
resolveMCPConfig: func(string, bool) (config.MCPConfig, error) {
return config.MCPConfig{Servers: map[string]config.MCPServerConfig{
"partial": {Type: "stdio", Command: "fake-partial"},
}}, nil
},
newMCPStore: func() (*mcp.PermissionStore, error) { return nil, nil },
newSandboxStore: func() (*sandbox.GrantStore, error) { return grantStore, nil },
registerMCPTools: func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
return closeFunc(func() error {
closed = true
return nil
}), errors.New("second MCP server failed")
},
})
workspace, err := buildACPWorkspace(context.Background(), t.TempDir(), config.ResolvedConfig{}, agent.PermissionModeAuto, deps)
if err == nil {
t.Fatal("buildACPWorkspace succeeded after MCP registration error")
}
if workspace != nil {
t.Fatal("buildACPWorkspace returned a workspace after MCP registration error")
}
if !closed {
t.Fatal("partial MCP runtime was not closed after registration error")
}
}
type acpNotifyingReadCloser struct {
io.ReadCloser
readStarted chan<- struct{}
+8 -2
View File
@@ -18,14 +18,20 @@ import (
// cliFakeDeferredTool is deferred-eligible (implements Deferred() bool), mirroring
// an MCP registry tool, so it counts toward the deferral threshold.
type cliFakeDeferredTool struct {
name string
name string
permission tools.Permission
advertiseInAuto bool
}
func (t cliFakeDeferredTool) Name() string { return t.name }
func (t cliFakeDeferredTool) Description() string { return "fake deferred tool" }
func (t cliFakeDeferredTool) Parameters() tools.Schema { return tools.Schema{Type: "object"} }
func (t cliFakeDeferredTool) Safety() tools.Safety {
return tools.Safety{SideEffect: tools.SideEffectNetwork, Permission: tools.PermissionAllow}
permission := t.permission
if permission == "" {
permission = tools.PermissionAllow
}
return tools.Safety{SideEffect: tools.SideEffectNetwork, Permission: permission, AdvertiseInAuto: t.advertiseInAuto}
}
func (t cliFakeDeferredTool) Run(context.Context, map[string]any) tools.Result {
return tools.Result{Status: tools.StatusOK, Output: "ok"}
+12 -10
View File
@@ -31,6 +31,14 @@ type mcpToolListItem struct {
// into the one-line trust notice (mirroring the hooks and plugins chokepoints);
// otherwise a workspace whose only project config is MCP would be gated silently.
func registerMCPToolsForWorkspace(ctx context.Context, workspaceRoot string, registry *tools.Registry, deps appDeps, autonomy mcp.PermissionAutonomy, trustRoot string, runners ...*execution.Runner) (mcpToolRuntime, trustSkip, error) {
options := mcp.RegisterOptions{}
if len(runners) > 0 {
options.Execution = runners[0]
}
return registerMCPToolsForWorkspaceWithOptions(ctx, workspaceRoot, registry, deps, autonomy, trustRoot, options)
}
func registerMCPToolsForWorkspaceWithOptions(ctx context.Context, workspaceRoot string, registry *tools.Registry, deps appDeps, autonomy mcp.PermissionAutonomy, trustRoot string, options mcp.RegisterOptions) (mcpToolRuntime, trustSkip, error) {
excludeProject, trustCheckErrored := resolveTrust(trustRoot)
skip := trustSkip{
excludedProjectConfig: excludeProject && projectMCPConfigExists(workspaceRoot),
@@ -47,16 +55,10 @@ func registerMCPToolsForWorkspace(ctx context.Context, workspaceRoot string, reg
if err != nil {
return nil, skip, err
}
var runner *execution.Runner
if len(runners) > 0 {
runner = runners[0]
}
runtime, err := deps.registerMCPTools(ctx, registry, cfg, mcp.RegisterOptions{
PermissionStore: store,
Autonomy: autonomy,
Execution: runner,
WorkspaceRoot: workspaceRoot,
})
options.PermissionStore = store
options.Autonomy = autonomy
options.WorkspaceRoot = workspaceRoot
runtime, err := deps.registerMCPTools(ctx, registry, cfg, options)
return runtime, skip, err
}
+7 -3
View File
@@ -23,6 +23,9 @@ const defaultConnectTimeout = 8 * time.Second
type RegisterOptions struct {
PermissionStore *PermissionStore
Autonomy PermissionAutonomy
// AdvertiseInAuto exposes prompt-gated MCP tools to model surfaces whose
// default mode is auto while preserving the execution-time prompt.
AdvertiseInAuto bool
ClientFactory func(context.Context, Server) (ToolClient, error)
// ConnectTimeout bounds the per-server connect+list at startup. Zero uses
// defaultConnectTimeout.
@@ -271,9 +274,10 @@ func newRegistryTool(server Server, remote RemoteTool, client ToolClient, option
client: client,
parameters: SchemaFromMCP(remote.InputSchema),
safety: tools.Safety{
SideEffect: tools.SideEffectNetwork,
Permission: permission,
Reason: fmt.Sprintf("MCP tool %s/%s runs through the configured %s server.", server.Name, remote.Name, server.Type),
SideEffect: tools.SideEffectNetwork,
Permission: permission,
Reason: fmt.Sprintf("MCP tool %s/%s runs through the configured %s server.", server.Name, remote.Name, server.Type),
AdvertiseInAuto: options.AdvertiseInAuto,
},
}
}
+13
View File
@@ -342,6 +342,19 @@ func TestRegistryToolIsDeferredEligible(t *testing.T) {
}
}
func TestRegistryToolCanAdvertiseInAutoWithoutPreapproval(t *testing.T) {
tool := newRegistryTool(
Server{Name: "docs", Type: "stdio"},
RemoteTool{Name: "search"},
&fakeToolClient{},
RegisterOptions{Autonomy: AutonomyLow, AdvertiseInAuto: true},
)
safety := tool.Safety()
if safety.Permission != tools.PermissionPrompt || !safety.AdvertiseInAuto {
t.Fatalf("MCP safety = %#v", safety)
}
}
// TestRegistryToolReportsMCPServerName verifies the registryTool reports its true
// configured server name (not the sanitized tool-name token) so the deferred-tools
// discovery label names a multi-token server correctly via tools.DeferredSource.
@@ -0,0 +1,463 @@
package sessions
import (
"fmt"
"strings"
"testing"
)
// A CALL'S PAYLOAD MUST NOT RIDE INTO A LATER PROMPT ANY MORE THAN A RESULT'S.
//
// toolResultOutcome drops result bodies, and calls were left verbatim, so an
// interrupted write_file put its content into the next turn's prompt while the
// matching result body did not. Same secret, one door left open. The identity
// still has to survive, since a resumed turn knowing WHICH file was written is
// the whole point of admitting calls at all.
//
// Pinned as whole lines, compared as fields, for the same reason the result test
// is: a substring search for one fixture secret passes when a truncated or
// reworded copy leaks, and the renderer walks a map so field order varies.
func TestResumePromptCarriesToolCallIdentityWithoutPayload(t *testing.T) {
const secret = "AWS_SECRET_ACCESS_KEY=wJalrXUtnFEMI/K7MDENG"
const token = "Bearer sk-live-4eC39HqLyjWDarjtT1zdp7dc"
events := []Event{
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "rotate the deploy key"})},
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c1", "name": "write_file", "arguments": `{"path":"deploy/prod.env","content":"` + secret + `"}`})},
{Sequence: 3, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "write_file", "status": "ok", "output": "written"})},
{Sequence: 4, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c2", "name": "edit_file", "arguments": `{"path":"deploy/prod.env","old_string":"` + secret + `","new_string":"rotated"}`})},
{Sequence: 5, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "edit_file", "status": "ok"})},
{Sequence: 6, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c3", "name": "exec_command", "arguments": `{"cmd":"curl -H '` + token + `' https://api.example/rotate","workdir":"deploy"}`})},
{Sequence: 7, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "exec_command", "status": "ok"})},
{Sequence: 8, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c4", "name": "apply_patch", "arguments": `{"patch":"*** Begin Patch\n*** Update File: deploy/prod.env\n-` + secret + `\n+rotated\n*** End Patch"}`})},
{Sequence: 9, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "apply_patch", "status": "ok"})},
{Sequence: 10, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c5", "name": "read_file", "arguments": `{"path":"deploy/prod.env","offset":1,"limit":40}`})},
{Sequence: 11, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "read_file", "status": "ok", "output": secret})},
{Sequence: 12, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error: upstream timeout"})},
}
out := resumePrompt(t, events)
// The payloads: none of them, in any form.
for _, leaked := range []string{secret, secret[:16], token, token[:14], "rotated", "Begin Patch", "curl -H"} {
if strings.Contains(out, leaked) {
t.Errorf("tool call payload %q reached the resume prompt:\n%s", leaked, out)
}
}
// The identities: every one of them, from the call side.
want := map[int]struct {
kind string
fields []string
}{
2: {"tool_call", []string{"c1", "write_file", `{"path":"deploy/prod.env"}`}},
4: {"tool_call", []string{"c2", "edit_file", `{"path":"deploy/prod.env"}`}},
6: {"tool_call", []string{"c3", "exec_command", `{"workdir":"deploy"}`}},
8: {"tool_call", []string{"c4", "apply_patch"}},
10: {"tool_call", []string{"c5", "read_file"}},
}
seen := map[int]bool{}
for _, line := range strings.Split(out, "\n") {
if !strings.HasPrefix(line, "- #") || !strings.Contains(line, "tool_call:") {
continue
}
var sequence int
var kind string
rest, _ := strings.CutPrefix(line, "- #")
if _, err := fmt.Sscanf(rest, "%d %s", &sequence, &kind); err != nil {
t.Fatalf("unparsable context line %q: %v", line, err)
}
expected, wanted := want[sequence]
if !wanted {
t.Errorf("unexpected tool_call line %q", line)
continue
}
seen[sequence] = true
_, payload, _ := strings.Cut(rest, ": ")
fields := strings.Fields(payload)
for _, field := range expected.fields {
if !containsField(fields, field) {
t.Errorf("line #%d lost identity field %q: %q", sequence, field, line)
}
}
}
for sequence := range want {
if !seen[sequence] {
t.Errorf("tool_call #%d did not reach the resume prompt at all, so the interrupted work is invisible again", sequence)
}
}
// read_file keeps its window too, or a resumed turn re-reads a part it had.
if !strings.Contains(out, `"offset":1`) || !strings.Contains(out, `"limit":40`) {
t.Errorf("read_file lost its offset/limit window:\n%s", out)
}
}
func containsField(fields []string, want string) bool {
for _, field := range fields {
if field == want {
return true
}
}
// The reduced arguments object may carry more than one identity key, in map
// order, so a single-key expectation is also satisfied by an object that
// contains that key.
if strings.HasPrefix(want, "{") {
key := strings.TrimSuffix(strings.TrimPrefix(want, "{"), "}")
for _, field := range fields {
if strings.HasPrefix(field, "{") && strings.Contains(field, key) {
return true
}
}
}
return false
}
// ARGUMENTS THAT CANNOT BE READ ARE NOT PASSED THROUGH AS TEXT.
//
// An allow-list only protects what it can see. Arguments that do not decode as
// an object are dropped rather than rendered, since text this file could not
// inspect is text it cannot vouch for.
func TestResumePromptDropsUnreadableToolCallArguments(t *testing.T) {
const secret = "sk-live-unparseable-9f8e7d6c"
events := []Event{
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "go"})},
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c1", "name": "bash", "arguments": `not json at all ` + secret})},
{Sequence: 3, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error"})},
}
out := resumePrompt(t, events)
if strings.Contains(out, secret) {
t.Errorf("unparseable arguments were rendered as text:\n%s", out)
}
if !strings.Contains(out, "bash") {
t.Errorf("the call's identity was lost along with its unreadable arguments:\n%s", out)
}
}
// AND A TOOL THIS FILE HAS NEVER HEARD OF STILL KEEPS ITS PATH.
//
// The allow-list is by key, not by tool name, so an MCP tool or a future core
// tool whose argument is a path or a url carries it into the resume prompt
// without being enumerated here, while any body field it has is dropped.
func TestResumePromptKeepsIdentityForUnknownTools(t *testing.T) {
const body = "large opaque payload that must not be replayed"
events := []Event{
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "go"})},
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c1", "name": "mcp_uploader", "arguments": `{"url":"https://files.example/x","blob":"` + body + `"}`})},
{Sequence: 3, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error"})},
}
out := resumePrompt(t, events)
if strings.Contains(out, body) {
t.Errorf("an unlisted body field was replayed:\n%s", out)
}
if !strings.Contains(out, "https://files.example/x") {
t.Errorf("an unlisted tool lost its url identity:\n%s", out)
}
}
// AND AN ALLOW-LISTED KEY IS NOT A SAFE VALUE.
//
// web_fetch accepts a credential in the query string and redacts the URL it
// reports back, so an interrupted fetch had its token dropped from the result
// and kept verbatim in the call. This projection is what admits the call into a
// later turn's prompt, so the token rode into the next resume or fork. The
// identity the resumed turn actually needs is the host and path, and those
// survive.
func TestResumePromptRedactsCredentialsInsideRetainedValues(t *testing.T) {
const token = "secret-value-9f8e7d6c5b4a"
events := []Event{
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "check the feed"})},
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c1", "name": "web_fetch", "arguments": `{"url":"https://api.example/data?access_token=` + token + `&page=2"}`})},
{Sequence: 3, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c2", "name": "web_fetch", "arguments": `{"url":"https://reader:[email protected]/feed"}`})},
{Sequence: 4, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error: upstream timeout"})},
}
out := resumePrompt(t, events)
for _, leaked := range []string{token, token[:12], "hunter2"} {
if strings.Contains(out, leaked) {
t.Errorf("credential %q reached the resume prompt:\n%s", leaked, out)
}
}
// The non-secret identity is the whole reason calls are admitted at all.
for _, kept := range []string{"api.example", "/data", "page=2", "/feed"} {
if !strings.Contains(out, kept) {
t.Errorf("URL identity %q was lost to the scrub:\n%s", kept, out)
}
}
}
// The scrub must not eat ordinary identities. A path, a glob and a query are
// what the resumed turn navigates by, and none of them is a credential.
func TestResumePromptKeepsOrdinaryIdentitiesThroughTheScrub(t *testing.T) {
events := []Event{
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "go"})},
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c1", "name": "read_file", "arguments": `{"path":"internal/sessions/exec_session.go","offset":1,"limit":40}`})},
{Sequence: 3, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c2", "name": "grep", "arguments": `{"pattern":"func toolCallIdentity","glob":"src/**/*.tsx"}`})},
{Sequence: 4, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error"})},
}
out := resumePrompt(t, events)
for _, kept := range []string{"internal/sessions/exec_session.go", "func toolCallIdentity", "src/**/*.tsx"} {
if !strings.Contains(out, kept) {
t.Errorf("identity %q was lost to the scrub:\n%s", kept, out)
}
}
// Numbers are not strings and are nothing to scrub; they must survive whole.
if !strings.Contains(out, "40") {
t.Errorf("a numeric read window did not survive the scrub:\n%s", out)
}
}
// AND A FIELD THIS FILE HAS NEVER HEARD OF DOES NOT RIDE ALONG.
//
// toolResultOutcome builds its output from the fields it keeps, so an unknown
// one cannot reach a prompt through it. The call side edited arguments in place
// and returned every sibling key, which is the opposite contract: a producer
// recording one more top-level field, or the same payload with no arguments at
// all, put that field straight into the next turn. The projection names what
// survives, so both halves now drop what they do not name.
func TestResumePromptProjectsTopLevelToolCallFields(t *testing.T) {
const body = "opaque top-level payload that must not be replayed"
events := []Event{
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "go"})},
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
"id": "c1", "name": "write_file", "arguments": `{"path":"deploy/prod.env"}`, "rawInput": body,
})},
// The same shape with no arguments at all, which used to return the whole
// payload untouched.
{Sequence: 3, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
"id": "c2", "name": "read_file", "stashedContent": body,
})},
{Sequence: 4, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error"})},
}
out := resumePrompt(t, events)
for _, leaked := range []string{body, "rawInput", "stashedContent"} {
if strings.Contains(out, leaked) {
t.Errorf("an unnamed top-level field %q reached the resume prompt:\n%s", leaked, out)
}
}
// The identity itself still survives, from both calls.
for _, kept := range []string{"c1", "write_file", "deploy/prod.env", "c2", "read_file"} {
if !strings.Contains(out, kept) {
t.Errorf("call identity %q was lost to the projection:\n%s", kept, out)
}
}
}
// AND A PERMITTED KEY IS NOT PERMISSION FOR WHATEVER HANGS UNDER IT.
//
// The projection scrubbed strings and passed everything else through, which
// reads as "a number is nothing to scrub" and is true of a number and not of an
// object. MCP schemas allow object and array properties, so a whole payload can
// arrive under a permitted key. Redacting a serialized object would not help:
// ordinary private text has no credential shape to match.
func TestResumePromptDropsNestedArgumentContainers(t *testing.T) {
const body = "private body text with no credential shape at all"
const token = "secret-value-9f8e7d6c5b4a"
events := []Event{
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "go"})},
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
"id": "c1", "name": "mcp_search",
"arguments": `{"query":{"api_key":"` + token + `","content":"` + body + `"}}`,
})},
{Sequence: 3, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
"id": "c2", "name": "web_fetch",
"arguments": `{"url":["https://api.example/data?access_token=` + token + `"]}`,
})},
{Sequence: 4, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error"})},
}
out := resumePrompt(t, events)
for _, leaked := range []string{body, body[:20], token, token[:12], "api_key"} {
if strings.Contains(out, leaked) {
t.Errorf("nested argument content %q reached the resume prompt:\n%s", leaked, out)
}
}
// The call still says which tool ran, so a resumed turn can ask again.
for _, kept := range []string{"c1", "mcp_search", "c2", "web_fetch"} {
if !strings.Contains(out, kept) {
t.Errorf("call identity %q was lost along with the container:\n%s", kept, out)
}
}
}
// THE SEARCH ALIASES ARE TOOL-SPECIFIC, AND ONE OF THEM IS A TRAP.
//
// grep accepts `search` for its pattern and edit_file accepts the same word for
// the text being replaced. Admitting it globally for the sake of the first would
// replay file contents through the second.
func TestResumePromptKeepsGrepSearchAndDropsEditSearch(t *testing.T) {
const oldText = "OLD FILE CONTENTS THAT MUST NOT BE REPLAYED"
events := []Event{
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "go"})},
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
"id": "c1", "name": "grep", "arguments": `{"search":"func toolCallIdentity","path":"internal/sessions"}`,
})},
{Sequence: 3, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
"id": "c2", "name": "edit_file", "arguments": `{"path":"a.go","search":"` + oldText + `","replace":"new"}`,
})},
{Sequence: 4, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error"})},
}
out := resumePrompt(t, events)
if !strings.Contains(out, "func toolCallIdentity") {
t.Errorf("grep lost the pattern it was searching for:\n%s", out)
}
if strings.Contains(out, oldText) || strings.Contains(out, oldText[:20]) {
t.Errorf("edit_file replayed the text it was replacing:\n%s", out)
}
// Both calls keep the path, which is identity for either tool.
for _, kept := range []string{"internal/sessions", "a.go"} {
if !strings.Contains(out, kept) {
t.Errorf("path identity %q was lost:\n%s", kept, out)
}
}
}
// AND THE SUPPORTED SEARCH AND WINDOW FORMS SURVIVE, not only the canonical
// ones. A resumed turn that knows the path but not which slice was asked for
// has half the identity and will read again.
func TestResumePromptKeepsSupportedAliasesAndWindows(t *testing.T) {
events := []Event{
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "go"})},
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
"id": "c1", "name": "glob", "arguments": `{"match":"src/**/*.go"}`,
})},
{Sequence: 3, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
"id": "c2", "name": "read_file", "arguments": `{"path":"large.go","byte_offset":4096,"byte_limit":512}`,
})},
{Sequence: 4, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{
"id": "c3", "name": "read_file", "arguments": `{"path":"big.go","start_line":40,"end_line":80,"max_lines":41}`,
})},
{Sequence: 5, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error"})},
}
out := resumePrompt(t, events)
// KEY AND VALUE TOGETHER, NOT THE NUMBER ALONE. Searching for "40" finds
// it inside "4096", so the start_line assertion passed whether or not
// start_line survived at all, which is the opposite of what it claims.
for _, kept := range []string{
"src/**/*.go",
`"byte_offset":4096`, `"byte_limit":512`,
`"start_line":40`, `"end_line":80`, `"max_lines":41`,
} {
if !strings.Contains(out, kept) {
t.Errorf("supported identity or window %s was lost:\n%s", kept, out)
}
}
}
// THE SHARED TABLE MUST NEVER NAME A MUTATING TOOL'S BODY ARGUMENT.
//
// This is the rule the grep/edit_file collision taught, kept as a test so the
// next alias added for a search tool cannot quietly open the same door. The
// aliases below are the ones the mutating tools accept for the content they
// write; they belong in the per-tool table or nowhere.
func TestIdentityKeysNeverNameAMutatingToolBodyArgument(t *testing.T) {
bodyAliases := map[string][]string{
"edit_file": {"old_string", "old", "search", "find", "old_str", "new_string", "new", "replace", "replacement", "new_str"},
"write_file": {"content", "contents", "text", "body", "data", "file_content"},
"apply_patch": {"patch"},
}
for tool, aliases := range bodyAliases {
for _, alias := range aliases {
if toolCallIdentityKeys[alias] {
t.Errorf("%s accepts %q for content it writes, and the shared identity table admits it: a resumed prompt would replay the body", tool, alias)
}
}
}
// The per-tool table is where an ambiguous name belongs, and grep's use of
// `search` is the case that proves the mechanism is wired.
if !toolCallIdentityKeysByTool["grep"]["search"] {
t.Error("grep lost its search alias, so an interrupted grep resumes without the pattern")
}
}
// THROUGH THE STORE, NOT JUST THE RENDERER.
//
// Every test above hands FormatExecPrompt an events slice it built in memory,
// so none of them shows the projection still applies to events that were
// written, reloaded and prepared by the real resume and fork paths. That is
// where it has to hold: the leak this guards against is a later turn reading a
// recorded call, and the recording is what makes it later.
//
// Includes a call the tool would have REJECTED. The agent records OnToolCall
// before execution, so a read_file whose path is an object is persisted even
// though the tool refuses it, and a projection that only considered
// well-formed calls would carry it through.
func TestPersistedToolCallsAreProjectedOnResumeAndFork(t *testing.T) {
const body = "private body text with no credential shape at all"
const token = "secret-value-9f8e7d6c5b4a"
const oldText = "OLD FILE CONTENTS THAT MUST NOT BE REPLAYED"
store := NewStore(StoreOptions{RootDir: t.TempDir()})
session, err := store.Create(CreateInput{SessionID: "projection_persisted"})
if err != nil {
t.Fatal(err)
}
appends := []AppendEventInput{
{Type: EventMessage, Payload: map[string]any{"role": "user", "content": "go"}},
{Type: EventToolCall, Payload: map[string]any{
"id": "c1", "name": "mcp_search",
"arguments": `{"query":{"api_key":"` + token + `","content":"` + body + `"}}`,
}},
{Type: EventToolCall, Payload: map[string]any{
"id": "c2", "name": "edit_file",
"arguments": `{"path":"a.go","search":"` + oldText + `","replace":"new"}`,
}},
// Rejected by read_file, recorded anyway.
{Type: EventToolCall, Payload: map[string]any{
"id": "c3", "name": "read_file",
"arguments": `{"path":{"nested":"` + body + `"},"start_line":40}`,
}},
{Type: EventToolCall, Payload: map[string]any{
"id": "c4", "name": "grep",
"arguments": `{"search":"func toolCallIdentity","path":"internal/sessions"}`,
}},
{Type: EventError, Payload: map[string]any{"message": "provider error: upstream timeout"}},
}
if _, err := store.AppendEvents(session.SessionID, appends); err != nil {
t.Fatal(err)
}
for _, mode := range []struct {
name string
options PrepareExecOptions
}{
{"resume", PrepareExecOptions{Store: store, Resume: session.SessionID}},
{"fork", PrepareExecOptions{Store: store, Fork: session.SessionID, SessionID: "projection_persisted_fork"}},
} {
t.Run(mode.name, func(t *testing.T) {
prepared, err := PrepareExec(mode.options)
if err != nil {
t.Fatalf("PrepareExec: %v", err)
}
prompt := FormatExecPrompt("continue", prepared)
for _, leaked := range []string{body, body[:20], token, token[:12], oldText, oldText[:20], "api_key"} {
if strings.Contains(prompt, leaked) {
t.Errorf("%q survived persistence into the %s prompt:\n%s", leaked, mode.name, prompt)
}
}
// The identities that make the tail worth carrying at all.
for _, kept := range []string{"mcp_search", "edit_file", "read_file", "grep", "func toolCallIdentity", "internal/sessions", "a.go"} {
if !strings.Contains(prompt, kept) {
t.Errorf("identity %q was lost from the %s prompt:\n%s", kept, mode.name, prompt)
}
}
})
}
// AND THE STORE STILL HOLDS THE ORIGINALS. The projection is a view for the
// prompt, not an edit to the record: /rewind and an audit read the events.
events, err := store.ReadEvents(session.SessionID)
if err != nil {
t.Fatal(err)
}
stored := ""
for _, event := range events {
stored += string(event.Payload)
}
for _, want := range []string{body, token, oldText} {
if !strings.Contains(stored, want) {
t.Errorf("the projection modified the persisted events: %q is no longer in the store", want)
}
}
}
@@ -0,0 +1,227 @@
package sessions
import (
"encoding/json"
"fmt"
"slices"
"strings"
"testing"
)
func toolContextPayload(t *testing.T, value map[string]any) json.RawMessage {
t.Helper()
raw, err := json.Marshal(value)
if err != nil {
t.Fatal(err)
}
return json.RawMessage(raw)
}
// interruptedTurnEvents is a turn that read two files and then died on a
// provider error, which is the shape #913 reports: no assistant answer was ever
// produced, so nothing in the conversation describes the work.
func interruptedTurnEvents(t *testing.T) []Event {
t.Helper()
return []Event{
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "add retries to the http client"})},
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c1", "name": "read_file", "arguments": `{"path":"internal/http/client.go"}`})},
{Sequence: 3, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "read_file", "content": "package http"})},
{Sequence: 4, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c2", "name": "read_file", "arguments": `{"path":"internal/http/retry.go"}`})},
{Sequence: 5, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "read_file", "content": "var defaultBackoff = 250ms"})},
{Sequence: 6, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error: upstream timeout"})},
}
}
func resumePrompt(t *testing.T, events []Event) string {
t.Helper()
return FormatExecPrompt("continue", PreparedExec{
Mode: ModeResume,
Session: Metadata{SessionID: "s1"},
ContextEvents: events,
})
}
// AN INTERRUPTED TURN HAS TO SAY WHAT IT DID.
//
// A turn that ends normally describes its own work in the assistant's answer,
// so the next turn inherits a prose record. A turn killed by a provider error
// produces no answer, and the tool events that were the only record of the work
// were filtered out of the resume context. The next turn was told a request had
// been made and an error had happened, and nothing else, so it re-read from
// scratch (#913).
func TestResumePromptCarriesInterruptedToolWork(t *testing.T) {
out := resumePrompt(t, interruptedTurnEvents(t))
for _, want := range []string{"internal/http/client.go", "internal/http/retry.go"} {
if !strings.Contains(out, want) {
t.Errorf("the resume prompt does not name %s, so the next turn cannot know it was already read:\n%s", want, out)
}
}
// The conversation spine is still there.
for _, want := range []string{"add retries to the http client", "upstream timeout"} {
if !strings.Contains(out, want) {
t.Errorf("the resume prompt lost conversation context %q:\n%s", want, out)
}
}
}
// AND THE ORDER STAYS THE ORDER IT HAPPENED IN.
//
// The sequence number is rendered into each line, so a list that jumps
// backwards reads as a corrupted history rather than a merge artifact.
func TestResumePromptKeepsEventsInSequenceOrder(t *testing.T) {
out := resumePrompt(t, interruptedTurnEvents(t))
last := -1
for _, line := range strings.Split(out, "\n") {
if !strings.HasPrefix(line, "- #") {
continue
}
var seq int
if _, err := fmt.Sscanf(line, "- #%d", &seq); err != nil {
t.Fatalf("unparsable context line %q", line)
}
if seq <= last {
t.Fatalf("context is out of order at %q (previous #%d):\n%s", line, last, out)
}
last = seq
}
if last < 0 {
t.Fatal("SETUP INVALID: the prompt rendered no context lines")
}
}
// TOOL EVENTS MUST NOT EVICT CONVERSATION.
//
// This is what the original filter was added for (#460): tool events vastly
// outnumber messages, so admitting them into the same trailing budget would let
// one tool-heavy turn push every earlier message out of the context. A separate
// allowance is what keeps both properties at once.
func TestToolEventsNeverEvictConversation(t *testing.T) {
var events []Event
seq := 0
next := func() int { seq++; return seq }
// An early message that must survive, then a flood of tool work.
events = append(events, Event{Sequence: next(), Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "EARLIEST-REQUEST-MARKER"})})
for i := 0; i < 500; i++ {
arguments := fmt.Sprintf(`{"path":"noise/%03d.go"}`, i)
events = append(events, Event{Sequence: next(), Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"name": "read_file", "arguments": arguments})})
events = append(events, Event{Sequence: next(), Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "read_file", "content": "noise"})})
}
events = append(events, Event{Sequence: next(), Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "LATEST-REQUEST-MARKER"})})
out := resumePrompt(t, events)
for _, want := range []string{"EARLIEST-REQUEST-MARKER", "LATEST-REQUEST-MARKER"} {
if !strings.Contains(out, want) {
t.Errorf("a tool-heavy turn evicted conversation event %q, which is the regression the filter exists to prevent", want)
}
}
selected := promptContextEvents(events)
tools := 0
for _, event := range selected {
if event.Type == EventToolCall || event.Type == EventToolResult {
tools++
}
}
if tools == 0 {
t.Error("no tool work survived at all, so an interrupted tool-heavy turn still says nothing about what it did")
}
if len(selected) > 80 {
t.Errorf("selected %d events, over the 80 the prompt budget allows", len(selected))
}
}
// A session with no tool work renders exactly as it did before.
func TestResumePromptUnchangedWithoutToolEvents(t *testing.T) {
events := []Event{
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "hello"})},
{Sequence: 2, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "assistant", "content": "hi"})},
}
out := resumePrompt(t, events)
if strings.Count(out, "- #") != 2 {
t.Fatalf("a tool-free session rendered %d context lines, want 2:\n%s", strings.Count(out, "- #"), out)
}
}
// TOOL OUTPUT MUST NOT RIDE ALONG INTO A LATER PROMPT.
//
// The tail is there to say what the interrupted turn did, and the call already
// says that. Carrying the result body would put up to 500 bytes of raw tool
// output into a prompt on a later turn, and nothing redacts on the way in.
func TestResumePromptCarriesToolOutcomeWithoutOutput(t *testing.T) {
secret := "AKIAIOSFODNN7EXAMPLE-and-more-file-contents"
events := []Event{
{Sequence: 1, Type: EventMessage, Payload: toolContextPayload(t, map[string]any{"role": "user", "content": "look at the config"})},
{Sequence: 2, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c1", "name": "read_file", "arguments": `{"path":"deploy/prod.env"}`})},
{Sequence: 3, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "read_file", "status": "ok", "output": secret})},
{Sequence: 4, Type: EventToolCall, Payload: toolContextPayload(t, map[string]any{"id": "c2", "name": "read_file", "arguments": `{"path":"deploy/missing.env"}`})},
{Sequence: 5, Type: EventToolResult, Payload: toolContextPayload(t, map[string]any{"name": "read_file", "status": "error", "output": "no such file"})},
{Sequence: 6, Type: EventError, Payload: toolContextPayload(t, map[string]any{"message": "provider error: upstream timeout"})},
}
out := resumePrompt(t, events)
// THE TOOL LINES ARE PINNED WHOLE, NOT SEARCHED.
//
// Two reasons the obvious assertions are too weak to say what the trim
// promises. Searching the prompt for the fixture's full secret only rejects
// that exact string, so a change that carried a truncated or partly
// redacted prefix of the output would pass it. And asking whether "ok" and
// "error" each appear somewhere among the result lines passes just as well
// when the two statuses are swapped, or both attached to the wrong call.
//
// Pinning the whole of every tool line covers both: nothing but the tool
// name and the outcome survives a result, each outcome sits on its own
// line, and the line order pairs each one with the call above it.
//
// Compared as a set of fields per line rather than as a string, because the
// renderer walks the payload map and Go randomizes that order, so the same
// events render their fields in a different order on every run.
want := []struct {
sequence int
kind string
fields []string
}{
{2, "tool_call", []string{"c1", "read_file", `{"path":"deploy/prod.env"}`}},
{3, "tool_result", []string{"read_file", "ok"}},
{4, "tool_call", []string{"c2", "read_file", `{"path":"deploy/missing.env"}`}},
{5, "tool_result", []string{"read_file", "error"}},
}
var toolLines []string
for _, line := range strings.Split(out, "\n") {
if strings.HasPrefix(line, "- #") && (strings.Contains(line, "tool_call:") || strings.Contains(line, "tool_result:")) {
toolLines = append(toolLines, line)
}
}
if len(toolLines) != len(want) {
t.Fatalf("the resume prompt carries %d tool lines, want %d:\n%s", len(toolLines), len(want), out)
}
for index, line := range toolLines {
var sequence int
var kind string
rest, found := strings.CutPrefix(line, "- #")
if !found {
t.Fatalf("unparsable context line %q", line)
}
if _, err := fmt.Sscanf(rest, "%d %s", &sequence, &kind); err != nil {
t.Fatalf("unparsable context line %q: %v", line, err)
}
kind = strings.TrimSuffix(kind, ":")
_, payload, _ := strings.Cut(rest, ": ")
fields := strings.Fields(payload)
slices.Sort(fields)
expected := slices.Clone(want[index].fields)
slices.Sort(expected)
if sequence != want[index].sequence || kind != want[index].kind || !slices.Equal(fields, expected) {
t.Errorf("tool line %d is not what the trim promises.\ngot: #%d %s %v\nwant: #%d %s %v\nfull prompt:\n%s",
index, sequence, kind, fields, want[index].sequence, want[index].kind, expected, out)
}
}
// Named on its own so a leak reports as a leak rather than as the
// formatting drift the comparison above would also catch. A PREFIX, because
// the danger is output reaching a later prompt at all, not this exact
// fixture string reaching it.
if strings.Contains(out, secret[:16]) {
t.Errorf("tool output reached the resume prompt:\n%s", out)
}
}
+297 -10
View File
@@ -4,8 +4,11 @@ import (
"encoding/json"
"fmt"
"log"
"sort"
"strings"
"unicode/utf8"
"github.com/Gitlawb/zero/internal/redaction"
)
type ExecMode string
@@ -188,23 +191,307 @@ func FormatExecPrompt(prompt string, prepared PreparedExec) string {
}, "\n")
}
// promptContextEvents chooses what a resumed turn is told about the session so
// far.
//
// A RESUMED TURN NEEDS TO KNOW WHAT THE PREVIOUS ONE DID, NOT ONLY WHAT IT SAID.
//
// Conversation events alone were selected here, so a turn that read six files
// and then died before answering left this behind:
//
// - #1 message: add retries to the http client
// - #6 error: provider error: upstream timeout
//
// Nothing names the files, so the next turn re-reads them from scratch (#913).
// On a turn that ENDS NORMALLY the assistant's answer describes the work, which
// is why this was survivable; an interrupted turn produces no such answer, and
// the record of the work goes with it.
//
// ONLY THE UNSUMMARIZED TAIL, NOT EVERY TOOL EVENT. Filtering tool events out
// was deliberate (#460) and is right for work the assistant has already
// described: forty read_file results add length and no information next to the
// answer that explains them. It is wrong only for work with no answer after it,
// which is precisely what an interrupted turn leaves behind. So the events after
// the last assistant message come along, and everything the assistant already
// spoke for stays filtered.
//
// The tail also gets its own allowance rather than sharing the conversation
// budget, so a tool-heavy interrupted turn can never push an earlier message out
// of the context. That is the property #460 added this filter for and it still
// holds.
func promptContextEvents(events []Event) []Event {
const maxPromptContextEvents = 80
// Enough to describe an interrupted turn's work, small enough that the
// conversation still dominates the prompt.
const maxPromptContextTailEvents = 24
filtered := make([]Event, 0, len(events))
for _, event := range events {
switch event.Type {
case EventMessage, EventCompaction, EventSessionFork, EventSessionChild, EventSpecialistStart, EventSpecialistStop, EventError:
filtered = append(filtered, event)
lastSpoken := -1
for index, event := range events {
if event.Type == EventMessage && payloadRole(event.Payload) == "assistant" {
lastSpoken = index
}
}
if len(filtered) == 0 {
filtered = append(filtered, events...)
conversation := make([]Event, 0, len(events))
tail := make([]Event, 0, maxPromptContextTailEvents)
for index, event := range events {
switch event.Type {
case EventMessage, EventCompaction, EventSessionFork, EventSessionChild, EventSpecialistStart, EventSpecialistStop, EventError:
conversation = append(conversation, event)
case EventToolCall:
if index > lastSpoken {
tail = append(tail, toolCallIdentity(event))
}
case EventToolResult:
if index > lastSpoken {
tail = append(tail, toolResultOutcome(event))
}
}
}
if len(filtered) > maxPromptContextEvents {
filtered = filtered[len(filtered)-maxPromptContextEvents:]
if len(conversation) == 0 && len(tail) == 0 {
// An unrecognized event stream still says more than nothing.
conversation = append(conversation, events...)
}
return filtered
if len(conversation) > maxPromptContextEvents {
conversation = conversation[len(conversation)-maxPromptContextEvents:]
}
// Tail slots exist only when the conversation uses fewer than the 80-event
// cap. A session whose conversation already fills it carries no interrupted
// tool work, which is what every session did before tool events were admitted
// at all, and what the #460 cap is there to hold. Reserving a minimum tail by
// trimming conversation further is a product decision, not this change.
tailBudget := min(maxPromptContextTailEvents, maxPromptContextEvents-len(conversation))
if tailBudget < 0 {
tailBudget = 0
}
if len(tail) > tailBudget {
tail = tail[len(tail)-tailBudget:]
}
if len(tail) == 0 {
return conversation
}
// Merged back into execution order: the sequence numbers are rendered, so a
// list that jumps backwards would read as a corrupted history.
merged := make([]Event, 0, len(conversation)+len(tail))
merged = append(merged, conversation...)
merged = append(merged, tail...)
sort.SliceStable(merged, func(i, j int) bool { return merged[i].Sequence < merged[j].Sequence })
return merged
}
// toolCallIdentityKeys are the argument fields that say WHAT a call was about
// without carrying what it was about to write, run, or search for.
//
// Path, directory, URL, name and pattern fields identify the work: the file that
// was read, the tree that was listed, the expression that was searched. Body
// fields carry payload: write_file's content, edit_file's old and new strings,
// apply_patch's hunks, a shell command with whatever credential was on its
// line. The result side already drops payload through toolResultOutcome, and
// admitting calls into resume context without the same care put an interrupted
// write_file's content into the next turn's prompt while its result body did
// not. Same fact, one door left open.
//
// AN ALLOW-LIST, NOT A DENY-LIST, so an argument this file has never heard of is
// dropped rather than replayed. A new tool with a new body field is then a
// missing path in a resume prompt, which is visible, instead of a new leak, which
// is not. The keys cover every alias the tools accept for the identity fields.
var toolCallIdentityKeys = map[string]bool{
// files and directories
"path": true, "file": true, "file_path": true, "filepath": true, "filename": true,
"dir": true, "directory": true, "cwd": true, "workdir": true,
// what a search was for; these are what the interrupted turn was looking at
"pattern": true, "glob": true, "query": true, "regex": true, "expression": true,
"match": true,
// fetches and named resources
"url": true, "name": true,
// read windows, so a resumed turn knows which part it already had. Every
// form read_file accepts, because a resumed turn that knows the path but not
// which slice was asked for has half the identity and will read again.
"offset": true, "limit": true,
"start_line": true, "end_line": true, "max_lines": true,
"byte_offset": true, "byte_limit": true,
}
// toolCallIdentityKeysByTool holds the argument names whose meaning depends on
// which tool was called, so they cannot live in the table above.
//
// grep accepts `search` for its pattern, and edit_file accepts the same word
// for the text being replaced. Admitting it globally for the sake of the first
// would replay file contents through the second, which is the whole thing this
// projection exists to stop. The rule this encodes: an argument name that any
// MUTATING tool accepts for body content is scoped to the tools that mean
// something else by it, never added to the shared table.
var toolCallIdentityKeysByTool = map[string]map[string]bool{
"grep": {"search": true},
}
// retainedIdentityValue decides whether a value under a permitted key may be
// kept, and what it looks like if so.
//
// AN ALLOW-LISTED KEY IS NOT A SAFE VALUE. A url is a valid place for a
// credential to appear: web_fetch accepts a query token and redacts the URL it
// reports back, so an interrupted fetch of
// https://api.example/data?access_token=... had its token dropped from the
// result and kept verbatim in the call. This projection is what admits call
// arguments into a later turn's prompt, so the same scrub belongs here or the
// credential is replayed on resume. Host and path survive it, which is the
// identity the resumed turn needs.
//
// AND A KEY IS NOT PERMISSION FOR WHATEVER HANGS UNDER IT. Scrubbing strings
// and passing everything else through read as "numbers are nothing to scrub",
// which is true of a number and not of an object: MCP schemas allow object and
// array properties, so {"query":{"api_key":"...","content":"private body"}}
// arrives under a permitted key and carries a whole payload with it. Redacting
// a serialized object would not help either, since ordinary private text has no
// credential shape to match.
//
// So the retained shapes are named: a string, scrubbed, and the scalars a read
// window is made of. A container is dropped. The call keeps its name and id, so
// a resumed turn still knows which tool ran and can ask again; what it does not
// get is arbitrary nested content it was never meant to see.
func retainedIdentityValue(value any) (any, bool) {
switch typed := value.(type) {
case string:
return redaction.RedactString(typed, redaction.Options{}), true
case float64, int, int64, bool, json.Number:
return typed, true
default:
// Objects, arrays, null, and anything a future decoder invents.
return nil, false
}
}
// toolCallIdentityKeyAllowed reports whether key is an identity field for the
// tool that was called.
func toolCallIdentityKeyAllowed(toolName, key string) bool {
lowered := strings.ToLower(key)
if toolCallIdentityKeys[lowered] {
return true
}
return toolCallIdentityKeysByTool[strings.ToLower(strings.TrimSpace(toolName))][lowered]
}
// toolCallIdentityFields are the top-level payload fields a projected call
// keeps: which call it was, and which tool. Everything else on the payload is
// dropped, including a field added by a future producer that this file has
// never seen.
var toolCallIdentityFields = []string{"id", "name"}
// toolCallIdentity keeps a tool call's identity and drops its payload, the
// symmetric half of toolResultOutcome.
//
// The arguments travel as a JSON string inside the payload. They are decoded,
// reduced to the identity keys, and re-encoded; anything that does not decode as
// an object is removed outright rather than passed through as text, since text
// that could not be read is text that cannot be checked.
//
// REBUILT, NOT PATCHED, WHICH IS WHAT MAKES IT THE SYMMETRIC HALF.
// toolResultOutcome constructs its output from the two fields it keeps, so a
// field it has never heard of cannot reach a prompt through it. Editing
// arguments in place and returning the rest of the payload is the opposite
// contract: it drops what it recognises as a body and forwards every sibling
// key, so a producer recording one more top-level field would put that field
// into the next turn. The projection names what survives.
func toolCallIdentity(event Event) Event {
var decoded map[string]json.RawMessage
if err := json.Unmarshal(event.Payload, &decoded); err != nil {
event.Payload = json.RawMessage(`{}`)
return event
}
projected := map[string]json.RawMessage{}
for _, field := range toolCallIdentityFields {
if value, present := decoded[field]; present {
projected[field] = value
}
}
// The tool decides what some argument names mean, so the projection has to
// know which tool this was before it can say which keys are identity.
toolName := ""
if raw, present := decoded["name"]; present {
_ = json.Unmarshal(raw, &toolName)
}
kept := map[string]any{}
if raw, present := decoded["arguments"]; present {
var argumentsText string
if err := json.Unmarshal(raw, &argumentsText); err == nil {
var arguments map[string]any
if err := json.Unmarshal([]byte(argumentsText), &arguments); err == nil {
for key, value := range arguments {
if !toolCallIdentityKeyAllowed(toolName, key) {
continue
}
if retained, ok := retainedIdentityValue(value); ok {
kept[key] = retained
}
}
}
}
}
if len(kept) > 0 {
if reduced, err := json.Marshal(kept); err == nil {
if quoted, err := json.Marshal(string(reduced)); err == nil {
projected["arguments"] = quoted
}
}
}
rebuilt, err := json.Marshal(projected)
if err != nil {
event.Payload = json.RawMessage(`{}`)
return event
}
event.Payload = json.RawMessage(rebuilt)
return event
}
// toolResultOutcome strips a tool result down to WHICH tool ran and HOW IT
// ENDED, dropping the output body.
//
// The tail exists so a resumed turn knows what the interrupted one did, and the
// CALL already carries that: the tool name and its arguments, which is the path
// for a read. The result adds only a 500-byte prefix of the output, which is
// worth little beside the call and is the one part of an event that can carry
// file contents. Nothing redacts on the way into a prompt, so keeping it would
// re-emit raw tool output into a later turn on the strength of a truncation
// limit alone.
//
// The status stays, because dropping it would be worse than dropping the whole
// result: a bare call reads as work that succeeded, so a failed read would come
// back as a file the next turn believes it already has.
func toolResultOutcome(event Event) Event {
var decoded struct {
Name string `json:"name"`
Status string `json:"status"`
}
if err := json.Unmarshal(event.Payload, &decoded); err != nil {
// Undecodable: drop the payload rather than pass an unknown shape through.
event.Payload = json.RawMessage(`{}`)
return event
}
trimmed, err := json.Marshal(map[string]string{
"name": decoded.Name,
"status": decoded.Status,
})
if err != nil {
event.Payload = json.RawMessage(`{}`)
return event
}
event.Payload = json.RawMessage(trimmed)
return event
}
// payloadRole reads the "role" field of a message payload, returning "" when the
// payload is not a decodable object or carries no role.
func payloadRole(payload json.RawMessage) string {
if len(payload) == 0 {
return ""
}
var decoded struct {
Role string `json:"role"`
}
if err := json.Unmarshal(payload, &decoded); err != nil {
return ""
}
return strings.TrimSpace(decoded.Role)
}
func forkTitle(title string) string {
+53 -1
View File
@@ -610,8 +610,60 @@ func TestFormatExecPromptKeepsConversationMessagesWhenNoisyEventsFollow(t *testi
t.Fatalf("expected prompt to contain %q, got %q", want, prompt)
}
}
// The tool results here follow the last assistant answer, so they are work
// nothing has spoken for yet. They are carried on purpose: that is the whole
// of what an interrupted turn leaves behind (#913). The property this test
// was written for (#460) is that CONVERSATION survives a noisy turn, and the
// assertions above still hold it.
// The tool results here follow the last assistant answer, so they are work
// nothing has spoken for yet, and an interrupted turn leaves exactly this
// behind (#913). Their IDENTITY is carried; their OUTPUT is not, because the
// call beside them already names what was touched and the body is the part
// that can carry file contents into a later prompt.
if !strings.Contains(prompt, "tool_result: read_file") {
t.Fatalf("expected unanswered tool work to be named, got %q", prompt)
}
if strings.Contains(prompt, "noisy tool result") {
t.Fatalf("expected prompt to omit noisy non-conversation events, got %q", prompt)
t.Fatalf("tool output body reached the prompt, got %q", prompt)
}
if count := strings.Count(prompt, "tool_result:"); count > 24 {
t.Fatalf("carried %d tool events, over the tail allowance of 24: %q", count, prompt)
}
}
// AND WORK THE ASSISTANT ALREADY SPOKE FOR STAYS FILTERED.
//
// The counterpart to the case above, and what keeps that one from being
// satisfied by carrying every tool event ever recorded. Once an answer describes
// the work, repeating the raw results adds length and no information.
func TestFormatExecPromptOmitsToolWorkAnAnswerAlreadyCovers(t *testing.T) {
events := []Event{
{Sequence: 1, Type: EventMessage, Payload: json.RawMessage(`{"role":"user","content":"first user request"}`)},
}
for sequence := 2; sequence <= 20; sequence++ {
events = append(events, Event{Sequence: sequence, Type: EventToolResult, Payload: json.RawMessage(`{"name":"read_file","output":"already summarized tool result"}`)})
}
events = append(events,
Event{Sequence: 21, Type: EventMessage, Payload: json.RawMessage(`{"role":"assistant","content":"I read the files and here is the summary"}`)},
Event{Sequence: 22, Type: EventMessage, Payload: json.RawMessage(`{"role":"user","content":"latest user request"}`)},
)
prompt := FormatExecPrompt("continue", PreparedExec{
Mode: ModeResume,
Session: Metadata{SessionID: "session-already-answered"},
ContextEvents: events,
})
if strings.Contains(prompt, "already summarized tool result") {
t.Fatalf("tool work the assistant already described was repeated verbatim, got %q", prompt)
}
if strings.Contains(prompt, "tool_result:") {
t.Fatalf("an answered turn still listed its tool events, got %q", prompt)
}
for _, want := range []string{"first user request", "here is the summary", "latest user request"} {
if !strings.Contains(prompt, want) {
t.Fatalf("expected prompt to contain %q, got %q", want, prompt)
}
}
}
+17 -5
View File
@@ -75,14 +75,10 @@ func NewScopedExecCommandTool(workspaceRoot string, scope PathScope, manager *ex
if manager == nil {
manager = defaultExecSessionManager
}
description := "Runs a command in a PTY, returning output or a session ID for ongoing interaction."
if runtimeGOOS() == "windows" {
description += "\n\n" + shellGuidanceForGOOS(runtimeGOOS())
}
return execCommandTool{
baseTool: baseTool{
name: ExecCommandToolName,
description: description,
description: execCommandDescription(runtimeGOOS()),
parameters: Schema{
Type: "object",
Properties: map[string]PropertySchema{
@@ -674,3 +670,19 @@ func execDisplaySummary(commandText string, sessionID int, exited bool, exitCode
func runtimeGOOS() string {
return runtime.GOOS
}
// execCommandDescription is exec_command's description as a host running goos
// would see it. The shell guidance is appended on Windows only.
//
// Taking goos as an argument rather than reading it is what lets the guidance
// contract be checked for every platform from any one of them. internal/agent's
// token ratchet removes exactly this appended text before charging the schemas,
// so the condition here and hostExecCommandShellGuidance have to agree on all
// three platforms, not only on whichever one the test happens to run on.
func execCommandDescription(goos string) string {
description := "Runs a command in a PTY, returning output or a session ID for ongoing interaction."
if goos == "windows" {
description += "\n\n" + shellGuidanceForGOOS(goos)
}
return description
}
@@ -0,0 +1,83 @@
package tools
import (
"runtime"
"strings"
"testing"
)
// THE BUDGET TEST REMOVES THIS TEXT, SO IT HAS TO BE EXACTLY WHAT IS ADDED.
//
// internal/agent's eager-schema ratchet charges the tool schemas and this
// host's shell guidance separately, and it separates them by stripping
// HostExecCommandShellGuidance off exec_command's description before
// estimating. That is only valid while the guidance is appended verbatim as the
// tail of that description, and only while a platform that reports no guidance
// really appends none. If the two drift apart the ratchet keeps reporting a
// number, just not the one it claims.
//
// Checked for all three platforms rather than for the running one: the
// appending is itself conditional on GOOS, so a host-only check is the shape of
// bug this pins. Off Windows it would pass while reporting nothing at all.
func TestExecCommandDescriptionMatchesReportedGuidance(t *testing.T) {
for _, goos := range []string{"windows", "linux", "darwin"} {
t.Run(goos, func(t *testing.T) {
description := execCommandDescription(goos)
guidance := hostExecCommandShellGuidance(goos)
if goos == "windows" {
if guidance == "" {
t.Fatal("exec_command appends shell guidance on Windows, but the accessor reports none; the budget would charge it to the schemas")
}
if !strings.HasSuffix(description, guidance) {
t.Fatalf("the guidance is not the tail of the description, so removing it from the schema total is wrong.\nguidance: %q\ndescription tail: %q",
guidance, description[max(0, len(description)-len(guidance)-40):])
}
} else if guidance != "" {
t.Fatalf("no shell guidance is appended off Windows, but the accessor reported %q", guidance)
}
// Whatever guidance THIS platform would produce, none of it may be
// left in the description once the reported guidance is removed.
// Off Windows the accessor reports nothing, so this is the check
// that the description really carries nothing: drop the GOOS
// condition in execCommandDescription and a POSIX host starts
// appending its own guidance (about 150 characters on darwin) with
// the ratchet still charging every byte of it to the schemas.
platformGuidance := shellGuidanceForGOOS(goos)
if platformGuidance == "" {
t.Fatalf("SETUP INVALID: %s produces no shell guidance at all, so the check below cannot fail", goos)
}
if remainder := strings.TrimSuffix(description, guidance); strings.Contains(remainder, platformGuidance) {
t.Fatalf("exec_command's description still carries %s shell guidance after removing the %d characters the accessor reports, so the ratchet charges the rest to the tool schemas.\nremainder: %q",
goos, len(guidance), remainder)
}
})
}
}
// And the running host agrees with the table above, so the exported accessor
// and the real constructor cannot drift away from the platform contract.
func TestHostExecCommandShellGuidanceMatchesTheDescription(t *testing.T) {
tool := NewExecCommandTool(t.TempDir(), nil)
if got, want := tool.Description(), execCommandDescription(runtime.GOOS); got != want {
t.Fatalf("the constructed tool's description is not what the platform contract says.\ngot: %q\nwant: %q", got, want)
}
if got, want := HostExecCommandShellGuidance(), hostExecCommandShellGuidance(runtime.GOOS); got != want {
t.Fatalf("the exported accessor disagrees with the platform contract.\ngot: %q\nwant: %q", got, want)
}
}
// And the guidance a Windows host gets is never empty, whichever shell was
// detected, since an empty one would mean the model gets no syntax rule at all.
func TestWindowsShellGuidanceIsNeverEmpty(t *testing.T) {
for _, shell := range []shellRuntime{
{GOOS: "windows", Executable: `C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe`, Kind: shellKindPowerShell, Syntax: "PowerShell"},
{GOOS: "windows", Executable: `C:\Program Files\PowerShell\7\pwsh.exe`, Kind: shellKindPowerShell, Syntax: "PowerShell"},
{GOOS: "windows", Executable: `C:\Windows\System32\cmd.exe`, Kind: shellKindCmd, Syntax: "cmd.exe"},
} {
if guidance := shellGuidanceForRuntime(shell); strings.TrimSpace(guidance) == "" {
t.Errorf("%s produced no shell guidance", shell.Executable)
}
}
}
+22
View File
@@ -181,6 +181,28 @@ func shellGuidanceForRuntime(shell shellRuntime) string {
return guidance
}
// HostExecCommandShellGuidance returns the shell guidance appended to
// exec_command's description on this host, or "" where none is appended.
//
// Exported for the per-turn token ratchet in internal/agent. That budget has to
// separate what every host pays for the tool schemas from what THIS host pays
// for its shell guidance, because the guidance varies by shell and by
// PowerShell version, and a single number covering both silently means a
// different thing on each machine.
func HostExecCommandShellGuidance() string {
return hostExecCommandShellGuidance(runtimeGOOS())
}
// hostExecCommandShellGuidance is the guidance a host running goos appends. It
// is the pair of execCommandDescription, so the two can be checked against each
// other for every platform rather than only for the running one.
func hostExecCommandShellGuidance(goos string) string {
if goos != "windows" {
return ""
}
return shellGuidanceForGOOS(goos)
}
// HostShellEnvironmentGuidance returns the concise, model-facing shell rule
// for the current host.
func HostShellEnvironmentGuidance() string {