From 185619e056009649312fd1f7cc3eef0032eb073b Mon Sep 17 00:00:00 2001 From: KRATOS <84986124+gnanam1990@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:33:39 +0530 Subject: [PATCH] fix(acp): register configured MCP tools per turn --- internal/acp/agent.go | 49 +++++++++++-- internal/acp/agent_test.go | 55 +++++++++++++- internal/cli/acp.go | 57 +++++++++++---- internal/cli/acp_test.go | 146 +++++++++++++++++++++++++++++++++++++ 4 files changed, 282 insertions(+), 25 deletions(-) diff --git a/internal/acp/agent.go b/internal/acp/agent.go index 6050c7c8b..3a344c339 100644 --- a/internal/acp/agent.go +++ b/internal/acp/agent.go @@ -30,10 +30,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 +43,25 @@ 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 + 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,12 +262,28 @@ 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} opts := agent.Options{ @@ -256,7 +293,7 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string, Model: resolved.Provider.Model, Registry: registry, Sandbox: sandboxEngine, - PermissionMode: sess.currentMode(), + PermissionMode: mode, MaxTurns: resolved.MaxTurns, Images: images, OnText: note.text, diff --git a/internal/acp/agent_test.go b/internal/acp/agent_test.go index 4fa97a258..ba17841da 100644 --- a/internal/acp/agent_test.go +++ b/internal/acp/agent_test.go @@ -53,10 +53,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 +468,8 @@ func TestACPRunTurnWiresSandboxAndScopedRegistry(t *testing.T) { reg := tools.NewRegistry() reg.Register(tools.NewUpdatePlanTool()) engine := sandbox.NewEngine(sandbox.EngineOptions{WorkspaceRoot: t.TempDir()}) - deps.BuildWorkspace = func(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 +497,53 @@ func TestACPRunTurnWiresSandboxAndScopedRegistry(t *testing.T) { } } +// 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) { diff --git a/internal/cli/acp.go b/internal/cli/acp.go index c261c19ef..5c79fe1da 100644 --- a/internal/cli/acp.go +++ b/internal/cli/acp.go @@ -10,9 +10,10 @@ 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/sandbox" - "github.com/Gitlawb/zero/internal/tools" ) const acpUsage = `zero acp — serve the Agent Client Protocol (ACP) over stdio @@ -57,20 +58,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 +74,44 @@ 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} + 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, _, err := registerMCPToolsForWorkspace(ctx, workspaceRoot, registry, deps, mcp.AutonomyLow, workspaceRoot, execution.NewRunner(engine)) + 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 + } + 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 +} + // 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 diff --git a/internal/cli/acp_test.go b/internal/cli/acp_test.go index c957432b8..e183e8763 100644 --- a/internal/cli/acp_test.go +++ b/internal/cli/acp_test.go @@ -9,6 +9,13 @@ import ( "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 +115,145 @@ 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) + } + registered = true + registry.Register(cliFakeDeferredTool{name: "mcp_docs_search"}) + 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 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 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{}