mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-08 01:09:46 -10:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
dba66b04fd | ||
|
|
185619e056 |
No files matched your search
+52
-6
@@ -14,6 +14,7 @@ import (
|
||||
"github.com/Gitlawb/zero/internal/providercatalog"
|
||||
"github.com/Gitlawb/zero/internal/providermodelcatalog"
|
||||
"github.com/Gitlawb/zero/internal/providermodeldiscovery"
|
||||
"github.com/Gitlawb/zero/internal/redaction"
|
||||
"github.com/Gitlawb/zero/internal/sandbox"
|
||||
"github.com/Gitlawb/zero/internal/sessions"
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
@@ -30,10 +31,12 @@ type Deps struct {
|
||||
DiscoverModels func(context.Context, config.ProviderProfile) ([]providermodeldiscovery.Model, error)
|
||||
NewProvider func(profile config.ProviderProfile) (zeroruntime.Provider, error)
|
||||
RunAgent func(ctx context.Context, prompt string, provider zeroruntime.Provider, opts agent.Options) (agent.Result, error)
|
||||
// BuildWorkspace builds the SCOPED tool registry and the sandbox engine for a
|
||||
// validated workspace root, so ACP shell tools (bash/exec_command) are confined
|
||||
// exactly like the exec surface — never run unconfined on the host.
|
||||
BuildWorkspace func(workspaceRoot string, resolved config.ResolvedConfig) (*tools.Registry, *sandbox.Engine, error)
|
||||
// BuildWorkspace creates the per-turn scoped tool workspace for a validated
|
||||
// workspace root. Its Close method releases any resources the registry owns
|
||||
// (notably MCP server connections) after the turn completes or is cancelled.
|
||||
// Shell tools remain confined exactly like the exec surface — never unconfined
|
||||
// on the host.
|
||||
BuildWorkspace func(ctx context.Context, workspaceRoot string, resolved config.ResolvedConfig, mode agent.PermissionMode) (*Workspace, error)
|
||||
// ResolveWorkspaceRoot validates + normalizes a client-supplied cwd (must be an
|
||||
// existing directory; never the bare root). It is the file-tool confinement root.
|
||||
ResolveWorkspaceRoot func(cwd string) (string, error)
|
||||
@@ -41,6 +44,27 @@ type Deps struct {
|
||||
AgentInfo Implementation
|
||||
}
|
||||
|
||||
// Workspace is the per-turn execution environment passed to the agent. ACP does
|
||||
// not retain it across turns because MCP connections belong to the registry that
|
||||
// advertised their tools; keeping a stale registry after a cancelled turn would
|
||||
// leak its server process and its permission state into the next turn.
|
||||
type Workspace struct {
|
||||
Registry *tools.Registry
|
||||
Sandbox *sandbox.Engine
|
||||
DeferThreshold int
|
||||
Notices []string
|
||||
Cleanup func() error
|
||||
}
|
||||
|
||||
// Close releases resources created with the workspace. A nil cleanup is valid
|
||||
// for core-only workspaces.
|
||||
func (w *Workspace) Close() error {
|
||||
if w == nil || w.Cleanup == nil {
|
||||
return nil
|
||||
}
|
||||
return w.Cleanup()
|
||||
}
|
||||
|
||||
// Agent is the ACP agent server bound to one JSON-RPC connection (one editor).
|
||||
type Agent struct {
|
||||
conn *Conn
|
||||
@@ -241,13 +265,34 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string,
|
||||
if err != nil {
|
||||
return "", RPCError(codeInternalError, "provider: "+err.Error())
|
||||
}
|
||||
mode := sess.currentMode()
|
||||
// Build the SCOPED registry + sandbox engine for this session's workspace so
|
||||
// shell/file tools are confined to the workspace exactly like the exec surface.
|
||||
registry, sandboxEngine, err := a.deps.BuildWorkspace(sess.cwd, resolved)
|
||||
// The workspace can also own MCP connections, which must be closed after every
|
||||
// turn even if the agent run returns an error or the client cancels it.
|
||||
workspace, err := a.deps.BuildWorkspace(ctx, sess.cwd, resolved, mode)
|
||||
if err != nil {
|
||||
return "", RPCError(codeInternalError, "workspace: "+err.Error())
|
||||
}
|
||||
if workspace == nil || workspace.Registry == nil {
|
||||
if workspace != nil {
|
||||
_ = workspace.Close()
|
||||
}
|
||||
return "", RPCError(codeInternalError, "workspace: missing tool registry")
|
||||
}
|
||||
defer func() {
|
||||
if err := workspace.Close(); err != nil {
|
||||
log.Printf("acp: close workspace: %v", err)
|
||||
}
|
||||
}()
|
||||
registry := workspace.Registry
|
||||
sandboxEngine := workspace.Sandbox
|
||||
note := ¬ifier{conn: a.conn, sessionID: sess.id}
|
||||
for _, notice := range workspace.Notices {
|
||||
if text := strings.TrimSpace(redaction.RedactString(notice, redaction.Options{})); text != "" {
|
||||
note.text("\n\n[zero warning] " + text + "\n")
|
||||
}
|
||||
}
|
||||
|
||||
opts := agent.Options{
|
||||
Cwd: sess.cwd,
|
||||
@@ -256,7 +301,8 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string,
|
||||
Model: resolved.Provider.Model,
|
||||
Registry: registry,
|
||||
Sandbox: sandboxEngine,
|
||||
PermissionMode: sess.currentMode(),
|
||||
PermissionMode: mode,
|
||||
DeferThreshold: workspace.DeferThreshold,
|
||||
MaxTurns: resolved.MaxTurns,
|
||||
Images: images,
|
||||
OnText: note.text,
|
||||
|
||||
+152
-4
@@ -8,6 +8,7 @@ import (
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -35,6 +36,35 @@ func (f fakeProvider) StreamCompletion(_ context.Context, _ zeroruntime.Completi
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
type acpDeferredPromptTool struct{ name string }
|
||||
|
||||
func (t acpDeferredPromptTool) Name() string { return t.name }
|
||||
func (acpDeferredPromptTool) Description() string { return "deferred MCP test tool" }
|
||||
func (acpDeferredPromptTool) Parameters() tools.Schema { return tools.Schema{Type: "object"} }
|
||||
func (acpDeferredPromptTool) Deferred() bool { return true }
|
||||
func (acpDeferredPromptTool) Run(context.Context, map[string]any) tools.Result {
|
||||
return tools.Result{Status: tools.StatusOK, Output: "ok"}
|
||||
}
|
||||
func (acpDeferredPromptTool) Safety() tools.Safety {
|
||||
return tools.Safety{
|
||||
SideEffect: tools.SideEffectNetwork, Permission: tools.PermissionPrompt,
|
||||
AdvertiseInAuto: true,
|
||||
}
|
||||
}
|
||||
|
||||
type captureToolsProvider struct {
|
||||
requests chan zeroruntime.CompletionRequest
|
||||
}
|
||||
|
||||
func (p captureToolsProvider) StreamCompletion(_ context.Context, request zeroruntime.CompletionRequest) (<-chan zeroruntime.StreamEvent, error) {
|
||||
p.requests <- request
|
||||
ch := make(chan zeroruntime.StreamEvent, 2)
|
||||
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventText, Content: "done"}
|
||||
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}
|
||||
close(ch)
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
func testDeps(t *testing.T) Deps {
|
||||
t.Helper()
|
||||
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
|
||||
@@ -53,10 +83,10 @@ func testDeps(t *testing.T) Deps {
|
||||
return fakeProvider{text: "Hello from ZERO"}, nil
|
||||
},
|
||||
RunAgent: agent.Run,
|
||||
BuildWorkspace: func(string, config.ResolvedConfig) (*tools.Registry, *sandbox.Engine, error) {
|
||||
BuildWorkspace: func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
|
||||
r := tools.NewRegistry()
|
||||
r.Register(tools.NewUpdatePlanTool())
|
||||
return r, nil, nil
|
||||
return &Workspace{Registry: r}, nil
|
||||
},
|
||||
ResolveWorkspaceRoot: func(cwd string) (string, error) { return cwd, nil },
|
||||
Store: store,
|
||||
@@ -468,8 +498,8 @@ func TestACPRunTurnWiresSandboxAndScopedRegistry(t *testing.T) {
|
||||
reg := tools.NewRegistry()
|
||||
reg.Register(tools.NewUpdatePlanTool())
|
||||
engine := sandbox.NewEngine(sandbox.EngineOptions{WorkspaceRoot: t.TempDir()})
|
||||
deps.BuildWorkspace = func(string, config.ResolvedConfig) (*tools.Registry, *sandbox.Engine, error) {
|
||||
return reg, engine, nil
|
||||
deps.BuildWorkspace = func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
|
||||
return &Workspace{Registry: reg, Sandbox: engine}, nil
|
||||
}
|
||||
var captured agent.Options
|
||||
deps.RunAgent = func(_ context.Context, _ string, _ zeroruntime.Provider, opts agent.Options) (agent.Result, error) {
|
||||
@@ -497,6 +527,124 @@ func TestACPRunTurnWiresSandboxAndScopedRegistry(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestACPRunTurnUsesWorkspaceDeferThresholdAtProviderBoundary(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
threshold int
|
||||
want []string
|
||||
}{
|
||||
{name: "explicit zero disables", threshold: 0, want: []string{"mcp_docs_alpha", "mcp_docs_beta"}},
|
||||
{name: "below threshold stays eager", threshold: 3, want: []string{"mcp_docs_alpha", "mcp_docs_beta"}},
|
||||
{name: "at threshold defers", threshold: 2, want: []string{tools.ToolSearchToolName}},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
requests := make(chan zeroruntime.CompletionRequest, 4)
|
||||
deps := testDeps(t)
|
||||
deps.NewProvider = func(config.ProviderProfile) (zeroruntime.Provider, error) {
|
||||
return captureToolsProvider{requests: requests}, nil
|
||||
}
|
||||
deps.BuildWorkspace = func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(acpDeferredPromptTool{name: "mcp_docs_alpha"})
|
||||
registry.Register(acpDeferredPromptTool{name: "mcp_docs_beta"})
|
||||
registry.Register(tools.NewToolSearchTool(registry))
|
||||
return &Workspace{Registry: registry, DeferThreshold: tc.threshold}, nil
|
||||
}
|
||||
|
||||
h := newHarness(t, deps)
|
||||
defer h.stop()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
var created NewSessionResult
|
||||
if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir()}, &created); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{SessionID: created.SessionID, Prompt: []ContentBlock{TextBlock("hi")}}, &PromptResult{}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
request := <-requests
|
||||
got := make([]string, 0, len(request.Tools))
|
||||
for _, definition := range request.Tools {
|
||||
got = append(got, definition.Name)
|
||||
}
|
||||
if !slices.Equal(got, tc.want) {
|
||||
t.Fatalf("provider tools = %#v, want %#v", got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestACPRunTurnEmitsWorkspaceSetupNotices(t *testing.T) {
|
||||
deps := testDeps(t)
|
||||
deps.BuildWorkspace = func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.NewUpdatePlanTool())
|
||||
return &Workspace{Registry: registry, Notices: []string{"MCP server docs unavailable, skipped: timeout"}}, nil
|
||||
}
|
||||
h := newHarness(t, deps)
|
||||
defer h.stop()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
var created NewSessionResult
|
||||
if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir()}, &created); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{SessionID: created.SessionID, Prompt: []ContentBlock{TextBlock("hi")}}, &PromptResult{}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := drainText(t, h.updates)
|
||||
if !strings.Contains(got, "[zero warning] MCP server docs unavailable, skipped: timeout") || !strings.Contains(got, "Hello from ZERO") {
|
||||
t.Fatalf("streamed text = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestACPRunTurnClosesWorkspace proves a per-turn workspace never outlives the
|
||||
// agent run. This is particularly important for MCP: its registry owns live
|
||||
// client connections and, for stdio servers, child processes.
|
||||
func TestACPRunTurnClosesWorkspaceAfterSuccessAndFailure(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
runErr error
|
||||
}{
|
||||
{name: "success"},
|
||||
{name: "agent failure", runErr: errors.New("provider interrupted")},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
deps := testDeps(t)
|
||||
closed := 0
|
||||
deps.BuildWorkspace = func(context.Context, string, config.ResolvedConfig, agent.PermissionMode) (*Workspace, error) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.NewUpdatePlanTool())
|
||||
return &Workspace{Registry: registry, Cleanup: func() error {
|
||||
closed++
|
||||
return nil
|
||||
}}, nil
|
||||
}
|
||||
deps.RunAgent = func(context.Context, string, zeroruntime.Provider, agent.Options) (agent.Result, error) {
|
||||
return agent.Result{FinalAnswer: "ok"}, tc.runErr
|
||||
}
|
||||
h := newHarness(t, deps)
|
||||
defer h.stop()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
var created NewSessionResult
|
||||
if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir()}, &created); err != nil {
|
||||
t.Fatalf("session/new: %v", err)
|
||||
}
|
||||
err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{SessionID: created.SessionID, Prompt: []ContentBlock{TextBlock("hello")}}, &PromptResult{})
|
||||
if tc.runErr == nil && err != nil {
|
||||
t.Fatalf("session/prompt: %v", err)
|
||||
}
|
||||
if tc.runErr != nil && err == nil {
|
||||
t.Fatal("session/prompt succeeded after agent failure")
|
||||
}
|
||||
if closed != 1 {
|
||||
t.Fatalf("workspace cleanup calls = %d, want 1", closed)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestACPRejectsInvalidCwd confirms session/new fails when the workspace root
|
||||
// resolver rejects the client cwd (e.g. filesystem root).
|
||||
func TestACPRejectsInvalidCwd(t *testing.T) {
|
||||
|
||||
+75
-15
@@ -10,9 +10,11 @@ import (
|
||||
"github.com/Gitlawb/zero/internal/acp"
|
||||
"github.com/Gitlawb/zero/internal/agent"
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/execution"
|
||||
"github.com/Gitlawb/zero/internal/mcp"
|
||||
"github.com/Gitlawb/zero/internal/providermodeldiscovery"
|
||||
"github.com/Gitlawb/zero/internal/redaction"
|
||||
"github.com/Gitlawb/zero/internal/sandbox"
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
)
|
||||
|
||||
const acpUsage = `zero acp — serve the Agent Client Protocol (ACP) over stdio
|
||||
@@ -57,20 +59,8 @@ func runACP(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int
|
||||
// surface — no ACP-specific credential handling needed.
|
||||
NewProvider: deps.newProvider,
|
||||
RunAgent: agent.Run,
|
||||
// Build the SCOPED registry + sandbox engine per workspace, exactly like the
|
||||
// exec surface, so ACP shell/file tools are confined — never run unconfined.
|
||||
BuildWorkspace: func(workspaceRoot string, resolved config.ResolvedConfig) (*tools.Registry, *sandbox.Engine, error) {
|
||||
scope, err := sandbox.NewScope(workspaceRoot, resolved.Sandbox.AdditionalWriteRoots)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
engine, err := buildExecSandboxEngine(workspaceRoot, resolved, deps, scope)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
registry := newCoreRegistryScoped(workspaceRoot, scope)
|
||||
registerLocalControlTools(registry, workspaceRoot, resolved.LocalControl)
|
||||
return registry, engine, nil
|
||||
BuildWorkspace: func(ctx context.Context, workspaceRoot string, resolved config.ResolvedConfig, mode agent.PermissionMode) (*acp.Workspace, error) {
|
||||
return buildACPWorkspace(ctx, workspaceRoot, resolved, mode, deps)
|
||||
},
|
||||
ResolveWorkspaceRoot: acpWorkspaceRootResolver(deps),
|
||||
Store: deps.newSessionStore(),
|
||||
@@ -85,6 +75,76 @@ func runACP(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int
|
||||
return exitSuccess
|
||||
}
|
||||
|
||||
// buildACPWorkspace matches exec's registry construction for one ACP turn. MCP
|
||||
// servers use the same sandbox-prepared execution runner, project configuration is
|
||||
// gated by the validated workspace's trust state, and their runtime is released
|
||||
// when acp.Agent completes the turn. ACP deliberately stays at low autonomy: an
|
||||
// editor connection never upgrades MCP permissions beyond an interactive prompt.
|
||||
func buildACPWorkspace(ctx context.Context, workspaceRoot string, resolved config.ResolvedConfig, mode agent.PermissionMode, deps appDeps) (*acp.Workspace, error) {
|
||||
scope, err := sandbox.NewScope(workspaceRoot, resolved.Sandbox.AdditionalWriteRoots)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
engine, err := buildExecSandboxEngine(workspaceRoot, resolved, deps, scope)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
registry := newCoreRegistryScoped(workspaceRoot, scope)
|
||||
|
||||
workspace := &acp.Workspace{
|
||||
Registry: registry, Sandbox: engine,
|
||||
DeferThreshold: resolved.Tools.DeferThreshold,
|
||||
}
|
||||
if mode != agent.PermissionModePlan {
|
||||
// MCP stdio servers are subprocesses. Passing the engine to their runner
|
||||
// keeps them inside the exact sandbox / lifecycle path used by zero exec.
|
||||
runtime, trustSkip, err := registerMCPToolsForWorkspaceWithOptions(
|
||||
ctx, workspaceRoot, registry, deps, mcp.AutonomyLow, workspaceRoot,
|
||||
mcp.RegisterOptions{Execution: execution.NewRunner(engine), AdvertiseInAuto: true},
|
||||
)
|
||||
if err != nil {
|
||||
// RegisterTools may have connected an earlier server before reporting a
|
||||
// later failure, so do not orphan a partial runtime on this error path.
|
||||
if runtime != nil {
|
||||
_ = runtime.Close()
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
workspace.Cleanup = runtime.Close
|
||||
workspace.Notices = acpMCPSetupNotices(trustSkip, runtime)
|
||||
}
|
||||
registerLocalControlTools(registry, workspaceRoot, resolved.LocalControl)
|
||||
// MCP tools are deferred-eligible. Register their loader only after every
|
||||
// ACP-visible tool is present, using the same mode the agent receives.
|
||||
registerToolSearchIfEligible(registry, resolved.Tools.DeferThreshold, mode, nil, nil)
|
||||
return workspace, nil
|
||||
}
|
||||
|
||||
func acpMCPSetupNotices(skip trustSkip, runtime mcpToolRuntime) []string {
|
||||
var notices []string
|
||||
if skip.excludedProjectConfig {
|
||||
if skip.trustCheckErrored {
|
||||
notices = append(notices, "The workspace-trust store could not be read; project MCP servers were ignored (fail-closed). Run 'zero trust' to enable them.")
|
||||
} else {
|
||||
notices = append(notices, "Project MCP servers were ignored in this untrusted workspace. Run 'zero trust' to enable them.")
|
||||
}
|
||||
}
|
||||
if runtime == nil {
|
||||
return notices
|
||||
}
|
||||
for _, skipped := range runtime.Skipped() {
|
||||
if skipped.UnconfiguredDefault {
|
||||
continue
|
||||
}
|
||||
message := fmt.Sprintf("MCP server %s unavailable, skipped", skipped.Name)
|
||||
if skipped.Err != nil {
|
||||
message += ": " + redaction.ErrorMessage(skipped.Err, redaction.Options{})
|
||||
}
|
||||
notices = append(notices, redaction.RedactString(message, redaction.Options{}))
|
||||
}
|
||||
return notices
|
||||
}
|
||||
|
||||
// acpWorkspaceRootResolver validates a client-supplied cwd into a confinement
|
||||
// root. It reuses exec's resolveWorkspaceRoot (abs+clean, must be an existing
|
||||
// dir) and additionally rejects the filesystem root and the home directory — an
|
||||
|
||||
@@ -5,10 +5,19 @@ import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/agent"
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/mcp"
|
||||
"github.com/Gitlawb/zero/internal/sandbox"
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
"github.com/Gitlawb/zero/internal/workspacetrust"
|
||||
)
|
||||
|
||||
type acpTestReader func([]byte) (int, error)
|
||||
@@ -108,6 +117,198 @@ func TestRunACPIdleCancellationExitsCleanly(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildACPWorkspaceRegistersSandboxedMCPToolsAndClosesThem(t *testing.T) {
|
||||
setTrustConfigRoot(t)
|
||||
workspaceRoot := t.TempDir()
|
||||
if err := workspacetrust.Trust(workspaceRoot); err != nil {
|
||||
t.Fatalf("trust workspace: %v", err)
|
||||
}
|
||||
grantStore, err := sandbox.NewGrantStore(sandbox.StoreOptions{FilePath: t.TempDir() + "/grants.json"})
|
||||
if err != nil {
|
||||
t.Fatalf("new grant store: %v", err)
|
||||
}
|
||||
|
||||
var gotExclude, registered, closed bool
|
||||
deps := fillAppDeps(appDeps{
|
||||
resolveMCPConfig: func(root string, excludeProject bool) (config.MCPConfig, error) {
|
||||
if root != workspaceRoot {
|
||||
t.Fatalf("MCP workspace root = %q, want %q", root, workspaceRoot)
|
||||
}
|
||||
gotExclude = excludeProject
|
||||
return config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"docs": {Type: "stdio", Command: "fake-docs"},
|
||||
}}, nil
|
||||
},
|
||||
newMCPStore: func() (*mcp.PermissionStore, error) { return nil, nil },
|
||||
newSandboxStore: func() (*sandbox.GrantStore, error) { return grantStore, nil },
|
||||
registerMCPTools: func(ctx context.Context, registry *tools.Registry, cfg config.MCPConfig, options mcp.RegisterOptions) (mcpToolRuntime, error) {
|
||||
if ctx == nil {
|
||||
t.Fatal("MCP registration received nil context")
|
||||
}
|
||||
if len(cfg.Servers) != 1 || cfg.Servers["docs"].Command != "fake-docs" {
|
||||
t.Fatalf("MCP config = %+v", cfg)
|
||||
}
|
||||
if options.Autonomy != mcp.AutonomyLow {
|
||||
t.Fatalf("MCP autonomy = %q, want low", options.Autonomy)
|
||||
}
|
||||
if options.Execution == nil {
|
||||
t.Fatal("MCP stdio server was not given the sandbox execution runner")
|
||||
}
|
||||
if options.WorkspaceRoot != workspaceRoot {
|
||||
t.Fatalf("MCP execution workspace = %q, want %q", options.WorkspaceRoot, workspaceRoot)
|
||||
}
|
||||
if !options.AdvertiseInAuto {
|
||||
t.Fatal("ACP MCP tools must be advertised in auto mode without being pre-approved")
|
||||
}
|
||||
registered = true
|
||||
registry.Register(cliFakeDeferredTool{
|
||||
name: "mcp_docs_search", permission: tools.PermissionPrompt,
|
||||
advertiseInAuto: options.AdvertiseInAuto,
|
||||
})
|
||||
return closeFunc(func() error {
|
||||
closed = true
|
||||
return nil
|
||||
}), nil
|
||||
},
|
||||
})
|
||||
|
||||
workspace, err := buildACPWorkspace(context.Background(), workspaceRoot, config.ResolvedConfig{
|
||||
Tools: config.ToolsConfig{DeferThreshold: 1},
|
||||
}, agent.PermissionModeAuto, deps)
|
||||
if err != nil {
|
||||
t.Fatalf("buildACPWorkspace: %v", err)
|
||||
}
|
||||
if gotExclude {
|
||||
t.Fatal("trusted ACP workspace excluded project MCP configuration")
|
||||
}
|
||||
if !registered {
|
||||
t.Fatal("ACP workspace did not register configured MCP tools")
|
||||
}
|
||||
if _, ok := workspace.Registry.Get("mcp_docs_search"); !ok {
|
||||
t.Fatal("MCP tool is absent from ACP agent registry")
|
||||
}
|
||||
if _, ok := workspace.Registry.Get(tools.ToolSearchToolName); !ok {
|
||||
t.Fatal("ACP registry is missing tool_search for deferred MCP tools")
|
||||
}
|
||||
if workspace.DeferThreshold != 1 {
|
||||
t.Fatalf("workspace defer threshold = %d, want 1", workspace.DeferThreshold)
|
||||
}
|
||||
if err := workspace.Close(); err != nil {
|
||||
t.Fatalf("close ACP workspace: %v", err)
|
||||
}
|
||||
if !closed {
|
||||
t.Fatal("ACP workspace did not close its MCP runtime")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildACPWorkspacePreservesRedactedMCPSetupNotices(t *testing.T) {
|
||||
setTrustConfigRoot(t)
|
||||
workspaceRoot := t.TempDir()
|
||||
if err := os.MkdirAll(filepath.Join(workspaceRoot, ".zero"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
projectConfig := `{"mcp":{"servers":{"project-docs":{"type":"stdio","command":"project-docs"}}}}`
|
||||
if err := os.WriteFile(filepath.Join(workspaceRoot, ".zero", "config.json"), []byte(projectConfig), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
secret := "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
|
||||
grantStore, err := sandbox.NewGrantStore(sandbox.StoreOptions{FilePath: filepath.Join(t.TempDir(), "grants.json")})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
deps := fillAppDeps(appDeps{
|
||||
resolveMCPConfig: func(string, bool) (config.MCPConfig, error) {
|
||||
return config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"docs": {Type: "stdio", Command: "fake-docs"},
|
||||
}}, nil
|
||||
},
|
||||
newMCPStore: func() (*mcp.PermissionStore, error) { return nil, nil },
|
||||
newSandboxStore: func() (*sandbox.GrantStore, error) { return grantStore, nil },
|
||||
registerMCPTools: func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
|
||||
return fakeMCPRuntimeWithSkips{skipped: []mcp.SkippedServer{{
|
||||
Name: "docs", Err: errors.New("token=" + secret),
|
||||
}}}, nil
|
||||
},
|
||||
})
|
||||
|
||||
workspace, err := buildACPWorkspace(context.Background(), workspaceRoot, config.ResolvedConfig{}, agent.PermissionModeAuto, deps)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer workspace.Close()
|
||||
joined := strings.Join(workspace.Notices, "\n")
|
||||
if !strings.Contains(joined, "untrusted workspace") || !strings.Contains(joined, "MCP server docs unavailable") {
|
||||
t.Fatalf("workspace notices = %#v", workspace.Notices)
|
||||
}
|
||||
if strings.Contains(joined, secret) || !strings.Contains(joined, "[REDACTED]") {
|
||||
t.Fatalf("workspace notices leaked secret: %q", joined)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildACPWorkspaceSkipsMCPInPlanMode(t *testing.T) {
|
||||
grantStore, err := sandbox.NewGrantStore(sandbox.StoreOptions{FilePath: t.TempDir() + "/grants.json"})
|
||||
if err != nil {
|
||||
t.Fatalf("new grant store: %v", err)
|
||||
}
|
||||
called := false
|
||||
deps := fillAppDeps(appDeps{
|
||||
resolveMCPConfig: func(string, bool) (config.MCPConfig, error) {
|
||||
called = true
|
||||
return config.MCPConfig{}, nil
|
||||
},
|
||||
newSandboxStore: func() (*sandbox.GrantStore, error) { return grantStore, nil },
|
||||
registerMCPTools: func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
|
||||
called = true
|
||||
return nil, nil
|
||||
},
|
||||
})
|
||||
|
||||
workspace, err := buildACPWorkspace(context.Background(), t.TempDir(), config.ResolvedConfig{}, agent.PermissionModePlan, deps)
|
||||
if err != nil {
|
||||
t.Fatalf("buildACPWorkspace: %v", err)
|
||||
}
|
||||
if called {
|
||||
t.Fatal("plan-mode ACP workspace must not resolve or start MCP servers")
|
||||
}
|
||||
if err := workspace.Close(); err != nil {
|
||||
t.Fatalf("close plan workspace: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildACPWorkspaceClosesPartialMCPRuntimeOnRegistrationError(t *testing.T) {
|
||||
grantStore, err := sandbox.NewGrantStore(sandbox.StoreOptions{FilePath: t.TempDir() + "/grants.json"})
|
||||
if err != nil {
|
||||
t.Fatalf("new grant store: %v", err)
|
||||
}
|
||||
closed := false
|
||||
deps := fillAppDeps(appDeps{
|
||||
resolveMCPConfig: func(string, bool) (config.MCPConfig, error) {
|
||||
return config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"partial": {Type: "stdio", Command: "fake-partial"},
|
||||
}}, nil
|
||||
},
|
||||
newMCPStore: func() (*mcp.PermissionStore, error) { return nil, nil },
|
||||
newSandboxStore: func() (*sandbox.GrantStore, error) { return grantStore, nil },
|
||||
registerMCPTools: func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
|
||||
return closeFunc(func() error {
|
||||
closed = true
|
||||
return nil
|
||||
}), errors.New("second MCP server failed")
|
||||
},
|
||||
})
|
||||
|
||||
workspace, err := buildACPWorkspace(context.Background(), t.TempDir(), config.ResolvedConfig{}, agent.PermissionModeAuto, deps)
|
||||
if err == nil {
|
||||
t.Fatal("buildACPWorkspace succeeded after MCP registration error")
|
||||
}
|
||||
if workspace != nil {
|
||||
t.Fatal("buildACPWorkspace returned a workspace after MCP registration error")
|
||||
}
|
||||
if !closed {
|
||||
t.Fatal("partial MCP runtime was not closed after registration error")
|
||||
}
|
||||
}
|
||||
|
||||
type acpNotifyingReadCloser struct {
|
||||
io.ReadCloser
|
||||
readStarted chan<- struct{}
|
||||
|
||||
@@ -18,14 +18,20 @@ import (
|
||||
// cliFakeDeferredTool is deferred-eligible (implements Deferred() bool), mirroring
|
||||
// an MCP registry tool, so it counts toward the deferral threshold.
|
||||
type cliFakeDeferredTool struct {
|
||||
name string
|
||||
name string
|
||||
permission tools.Permission
|
||||
advertiseInAuto bool
|
||||
}
|
||||
|
||||
func (t cliFakeDeferredTool) Name() string { return t.name }
|
||||
func (t cliFakeDeferredTool) Description() string { return "fake deferred tool" }
|
||||
func (t cliFakeDeferredTool) Parameters() tools.Schema { return tools.Schema{Type: "object"} }
|
||||
func (t cliFakeDeferredTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectNetwork, Permission: tools.PermissionAllow}
|
||||
permission := t.permission
|
||||
if permission == "" {
|
||||
permission = tools.PermissionAllow
|
||||
}
|
||||
return tools.Safety{SideEffect: tools.SideEffectNetwork, Permission: permission, AdvertiseInAuto: t.advertiseInAuto}
|
||||
}
|
||||
func (t cliFakeDeferredTool) Run(context.Context, map[string]any) tools.Result {
|
||||
return tools.Result{Status: tools.StatusOK, Output: "ok"}
|
||||
|
||||
+12
-10
@@ -31,6 +31,14 @@ type mcpToolListItem struct {
|
||||
// into the one-line trust notice (mirroring the hooks and plugins chokepoints);
|
||||
// otherwise a workspace whose only project config is MCP would be gated silently.
|
||||
func registerMCPToolsForWorkspace(ctx context.Context, workspaceRoot string, registry *tools.Registry, deps appDeps, autonomy mcp.PermissionAutonomy, trustRoot string, runners ...*execution.Runner) (mcpToolRuntime, trustSkip, error) {
|
||||
options := mcp.RegisterOptions{}
|
||||
if len(runners) > 0 {
|
||||
options.Execution = runners[0]
|
||||
}
|
||||
return registerMCPToolsForWorkspaceWithOptions(ctx, workspaceRoot, registry, deps, autonomy, trustRoot, options)
|
||||
}
|
||||
|
||||
func registerMCPToolsForWorkspaceWithOptions(ctx context.Context, workspaceRoot string, registry *tools.Registry, deps appDeps, autonomy mcp.PermissionAutonomy, trustRoot string, options mcp.RegisterOptions) (mcpToolRuntime, trustSkip, error) {
|
||||
excludeProject, trustCheckErrored := resolveTrust(trustRoot)
|
||||
skip := trustSkip{
|
||||
excludedProjectConfig: excludeProject && projectMCPConfigExists(workspaceRoot),
|
||||
@@ -47,16 +55,10 @@ func registerMCPToolsForWorkspace(ctx context.Context, workspaceRoot string, reg
|
||||
if err != nil {
|
||||
return nil, skip, err
|
||||
}
|
||||
var runner *execution.Runner
|
||||
if len(runners) > 0 {
|
||||
runner = runners[0]
|
||||
}
|
||||
runtime, err := deps.registerMCPTools(ctx, registry, cfg, mcp.RegisterOptions{
|
||||
PermissionStore: store,
|
||||
Autonomy: autonomy,
|
||||
Execution: runner,
|
||||
WorkspaceRoot: workspaceRoot,
|
||||
})
|
||||
options.PermissionStore = store
|
||||
options.Autonomy = autonomy
|
||||
options.WorkspaceRoot = workspaceRoot
|
||||
runtime, err := deps.registerMCPTools(ctx, registry, cfg, options)
|
||||
return runtime, skip, err
|
||||
}
|
||||
|
||||
|
||||
@@ -23,6 +23,9 @@ const defaultConnectTimeout = 8 * time.Second
|
||||
type RegisterOptions struct {
|
||||
PermissionStore *PermissionStore
|
||||
Autonomy PermissionAutonomy
|
||||
// AdvertiseInAuto exposes prompt-gated MCP tools to model surfaces whose
|
||||
// default mode is auto while preserving the execution-time prompt.
|
||||
AdvertiseInAuto bool
|
||||
ClientFactory func(context.Context, Server) (ToolClient, error)
|
||||
// ConnectTimeout bounds the per-server connect+list at startup. Zero uses
|
||||
// defaultConnectTimeout.
|
||||
@@ -271,9 +274,10 @@ func newRegistryTool(server Server, remote RemoteTool, client ToolClient, option
|
||||
client: client,
|
||||
parameters: SchemaFromMCP(remote.InputSchema),
|
||||
safety: tools.Safety{
|
||||
SideEffect: tools.SideEffectNetwork,
|
||||
Permission: permission,
|
||||
Reason: fmt.Sprintf("MCP tool %s/%s runs through the configured %s server.", server.Name, remote.Name, server.Type),
|
||||
SideEffect: tools.SideEffectNetwork,
|
||||
Permission: permission,
|
||||
Reason: fmt.Sprintf("MCP tool %s/%s runs through the configured %s server.", server.Name, remote.Name, server.Type),
|
||||
AdvertiseInAuto: options.AdvertiseInAuto,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -342,6 +342,19 @@ func TestRegistryToolIsDeferredEligible(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistryToolCanAdvertiseInAutoWithoutPreapproval(t *testing.T) {
|
||||
tool := newRegistryTool(
|
||||
Server{Name: "docs", Type: "stdio"},
|
||||
RemoteTool{Name: "search"},
|
||||
&fakeToolClient{},
|
||||
RegisterOptions{Autonomy: AutonomyLow, AdvertiseInAuto: true},
|
||||
)
|
||||
safety := tool.Safety()
|
||||
if safety.Permission != tools.PermissionPrompt || !safety.AdvertiseInAuto {
|
||||
t.Fatalf("MCP safety = %#v", safety)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistryToolReportsMCPServerName verifies the registryTool reports its true
|
||||
// configured server name (not the sanitized tool-name token) so the deferred-tools
|
||||
// discovery label names a multi-token server correctly via tools.DeferredSource.
|
||||
|
||||
Reference in new issue
Block a user