From c16e57aabec4defe5f69e9d5b09fb51b7a92632c Mon Sep 17 00:00:00 2001 From: wjueyao Date: Mon, 18 May 2026 17:56:38 +0800 Subject: [PATCH 01/16] feat(runtime/acp): scaffold multi-session-per-agentrun (Manager.sessions map) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit First step of the multi-session refactor. ACP protocol natively supports multiple sessions per agent process (per spec: session/new is baseline), but mass runtime's Manager pinned a single sessionID, forcing callers to fork+kill an agent process per session — wasteful for use cases like mindpowers self-EDD that run N tasks back-to-back. This commit adds the structural foundation without breaking existing single-session callers: 1. New sessionState type capturing per-session protocol metadata (id / models / cwd). Mirrors the agent-side session map. 2. Manager.sessions map[acp.SessionId]*sessionState added. Create()'s initial session/new now populates both the legacy single-session fields (sessionID / models — kept for back-compat) AND the map. 3. NewSession(ctx, cwd, mcpServers) → SessionId opens additional ACP sessions on the running agent. Takes per-session cwd + optional per-session MCP server overrides layered on bundle defaults. 4. EndSession(sessionId) clears runtime tracking. Note: ACP has no explicit "end session" RPC yet, so this only removes the local map entry — the agent process retains its session state until cancelled or process exits. Initial session (Create's first one) cannot be ended via this method (kill the agent instead). 5. Sessions() returns active session IDs snapshot for callers that need to enumerate / clean up. Out of scope for this commit (follow-ups): - Prompt/Cancel/SetModel sessionId parameter (PromptSession etc.) - State.json schema with sessions map (currently still single SessionID) - ARI server: agentrun/new-session + agentrun/end-session RPC - massctl: agentrun new-session/end-session subcommands - Tests for the new methods (this is scaffold; tests follow with the RPC layer so we can exercise via the same interface mindpowers uses) Co-Authored-By: Claude Opus 4.7 (1M context) --- pkg/agentrun/runtime/acp/runtime.go | 146 ++++++++++++++++++++++++++-- 1 file changed, 139 insertions(+), 7 deletions(-) diff --git a/pkg/agentrun/runtime/acp/runtime.go b/pkg/agentrun/runtime/acp/runtime.go index 877b4ac..731150c 100644 --- a/pkg/agentrun/runtime/acp/runtime.go +++ b/pkg/agentrun/runtime/acp/runtime.go @@ -57,19 +57,51 @@ type StateChange struct { type StateChangeHook func(StateChange) // Manager manages the lifecycle of a single ACP agent process. +// sessionState holds per-session protocol metadata. Multiple sessions may +// co-exist on one agent process — ACP's session/new is multi-session by +// design (each session has its own cwd, sessionId, model state). The +// agent process (e.g. claude-agent-acp) keeps sessions in a map keyed by +// sessionId; this struct mirrors the runtime-side view of each. +// +// Multi-session lets callers (mass daemon / mindpowers) keep a long-lived +// agent process and switch session per task, instead of fork+kill an +// agent process per task. See docs/design/multi-session-per-agentrun.md +// (TODO: add design doc) for the broader rationale. +type sessionState struct { + id acp.SessionId + models *acp.SessionModelState + cwd string +} + type Manager struct { cfg apiruntime.Config bundleDir string stateDir string logger *slog.Logger - mu sync.Mutex - cmd *exec.Cmd - processDone chan struct{} - conn *acp.ClientSideConnection - sessionID acp.SessionId - events chan acp.SessionNotification - models *acp.SessionModelState // from session/new response + mu sync.Mutex + cmd *exec.Cmd + processDone chan struct{} + conn *acp.ClientSideConnection + events chan acp.SessionNotification + + // sessions holds all active ACP sessions on this agent process, keyed + // by sessionId. NewSession adds entries; EndSession removes them. + // Empty until Create()'s initial session/new completes. + sessions map[acp.SessionId]*sessionState + + // sessionID retains the *first* session created by Create() so legacy + // callers of Prompt/Cancel/SetModel (which don't pass sessionId) keep + // working without code changes. New callers should use the explicit + // sessionId-taking variants (PromptSession etc.). To be removed once + // all callers migrate. + sessionID acp.SessionId + + // models tracks the *first* session's models for the same backward- + // compat reason as sessionID above. Per-session models live in + // sessions[id].models. + models *acp.SessionModelState + stateChangeHook StateChangeHook eventCountsFn func() map[string]int usageFn func() *apiruntime.UsageInfo @@ -83,6 +115,7 @@ func New(cfg apiruntime.Config, bundleDir, stateDir string, logger *slog.Logger) stateDir: stateDir, logger: logger.With("subsystem", "runtime"), events: make(chan acp.SessionNotification, 1024), + sessions: make(map[acp.SessionId]*sessionState), } } @@ -197,8 +230,17 @@ func (m *Manager) Create(ctx context.Context) error { return fmt.Errorf("runtime: acp session/new: %w", err) } m.mu.Lock() + // Initial session populates both the legacy single-session fields + // (sessionID/models) and the multi-session map. Future sessions opened + // via NewSession only register in the map; sessionID stays pinned to + // the first one for backward compat with Prompt/Cancel/SetModel. m.sessionID = sessionResp.SessionId m.models = sessionResp.Models + m.sessions[sessionResp.SessionId] = &sessionState{ + id: sessionResp.SessionId, + models: sessionResp.Models, + cwd: workDir, + } m.mu.Unlock() m.logger.Info("session created", "sessionID", sessionResp.SessionId) @@ -305,6 +347,96 @@ func (m *Manager) GetState() (apiruntime.State, error) { return spec.ReadState(m.stateDir) } +// NewSession opens an additional ACP session on the running agent process. +// The agent must already be started via Create() (which creates the initial +// session). Returns the new session's ID, which callers must pass to +// PromptSession / CancelSession / EndSession to address this session. +// +// cwd overrides the agent's working directory for this session — useful for +// running multiple isolated tasks on one long-lived agent (e.g. self-EDD's +// per-case fixture directories). Empty cwd uses the agent's workDir. +// +// mcpServers (optional) are extra MCP servers scoped to this session, layered +// on top of the agent's bundle-level mcpServers. +func (m *Manager) NewSession(ctx context.Context, cwd string, mcpServers []acp.McpServer) (acp.SessionId, error) { + m.mu.Lock() + conn := m.conn + m.mu.Unlock() + + if conn == nil { + return "", fmt.Errorf("runtime: agent not started") + } + + resolvedCwd := cwd + if resolvedCwd == "" { + workDir, err := spec.ResolveAgentRoot(m.bundleDir, m.cfg) + if err != nil { + return "", fmt.Errorf("runtime: resolve cwd for new session: %w", err) + } + resolvedCwd = workDir + } + + // Layer session-scoped mcpServers on top of bundle-level config. Callers + // that pass nil/empty get the bundle defaults. + merged := convertMcpServers(m.cfg.Session.McpServers) + merged = append(merged, mcpServers...) + + req := acp.NewSessionRequest{ + Meta: m.cfg.Session.Meta, + Cwd: resolvedCwd, + McpServers: merged, + } + resp, err := conn.NewSession(ctx, req) + if err != nil { + return "", fmt.Errorf("runtime: acp session/new: %w", err) + } + + m.mu.Lock() + m.sessions[resp.SessionId] = &sessionState{ + id: resp.SessionId, + models: resp.Models, + cwd: resolvedCwd, + } + m.mu.Unlock() + m.logger.Info("session opened", "sessionID", resp.SessionId, "cwd", resolvedCwd) + return resp.SessionId, nil +} + +// EndSession releases runtime tracking of a session id. The ACP protocol +// has no explicit "end session" RPC today, so this method only clears the +// runtime's local map entry — the agent process retains its own session +// state until cancelled or the process exits. Callers must still call +// CancelSession (or wait for prompt completion) before EndSession to ensure +// no in-flight work on the session. +// +// Refusing to end the initial session (the one Create() opened) — callers +// that want to fully tear down the agent should call Kill() instead. +func (m *Manager) EndSession(sessionID acp.SessionId) error { + m.mu.Lock() + defer m.mu.Unlock() + if sessionID == m.sessionID { + return fmt.Errorf("runtime: cannot end initial session %q (kill the agent instead)", sessionID) + } + if _, ok := m.sessions[sessionID]; !ok { + return fmt.Errorf("runtime: session %q not found", sessionID) + } + delete(m.sessions, sessionID) + m.logger.Info("session ended", "sessionID", sessionID) + return nil +} + +// Sessions returns a snapshot of currently-active session IDs. Order is +// not stable — caller should sort if a deterministic order is needed. +func (m *Manager) Sessions() []acp.SessionId { + m.mu.Lock() + defer m.mu.Unlock() + ids := make([]acp.SessionId, 0, len(m.sessions)) + for id := range m.sessions { + ids = append(ids, id) + } + return ids +} + // Prompt sends a user prompt to the agent and blocks until the agent // returns a PromptResponse. Session notifications emitted by the agent // during the turn are forwarded to the Events channel. From c17f2988df0fa3445859b02301b01f1ce92cead2 Mon Sep 17 00:00:00 2001 From: wjueyao Date: Mon, 18 May 2026 17:58:16 +0800 Subject: [PATCH 02/16] feat(runtime/acp): add PromptSession/CancelSession/SetModelSession variants MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Step 2 of multi-session refactor (step 1: c16e57a sessions map scaffold). Existing Prompt/Cancel/SetModel methods didn't take sessionId — they used the Manager's single m.sessionID, which broke for callers holding multiple sessions via the new NewSession method. This commit adds explicit-sessionId variants and refactors the original methods to be thin shims that pass m.sessionID through. Backward compatibility preserved: existing callers (mass daemon ARI, massctl) keep working without changes; new multi-session callers use the *Session variants. API additions: PromptSession(ctx, sessionId, prompt) → PromptResponse CancelSession(ctx, sessionId) error SetModelSession(ctx, sessionId, modelId) error All three validate sessionId is in Manager.sessions map and return a clear "session not found" error otherwise — avoids confusing low-level ACP errors when callers pass a stale or wrong sessionId. Per-session models state is now updated in sessions[id].models. The legacy single-session Manager.models field is still mirrored for the initial session only (matches existing GetState / SessionID() behavior). State.json updates remain process-wide (single Phase field) — per- session phase tracking is a follow-up state-schema change. Compile-clean (go vet + go test -run NoneShouldMatch across pkg/agentrun). Co-Authored-By: Claude Opus 4.7 (1M context) --- pkg/agentrun/runtime/acp/runtime.go | 101 ++++++++++++++++++++++------ 1 file changed, 80 insertions(+), 21 deletions(-) diff --git a/pkg/agentrun/runtime/acp/runtime.go b/pkg/agentrun/runtime/acp/runtime.go index 731150c..1d79bdf 100644 --- a/pkg/agentrun/runtime/acp/runtime.go +++ b/pkg/agentrun/runtime/acp/runtime.go @@ -437,20 +437,42 @@ func (m *Manager) Sessions() []acp.SessionId { return ids } -// Prompt sends a user prompt to the agent and blocks until the agent -// returns a PromptResponse. Session notifications emitted by the agent -// during the turn are forwarded to the Events channel. -// On completion (success or error), LastTurn is persisted to state.json. +// Prompt sends a user prompt to the agent's initial session and blocks +// until the agent returns a PromptResponse. Session notifications emitted +// by the agent during the turn are forwarded to the Events channel. +// On completion (success or error), state.json is updated. +// +// This is a backward-compat shim around PromptSession that uses the +// initial sessionId set by Create(). New callers should prefer +// PromptSession with an explicit sessionId — especially when the agent +// holds multiple sessions opened via NewSession. func (m *Manager) Prompt(ctx context.Context, prompt []acp.ContentBlock) (acp.PromptResponse, error) { m.mu.Lock() - conn := m.conn sessionID := m.sessionID m.mu.Unlock() + return m.PromptSession(ctx, sessionID, prompt) +} + +// PromptSession sends a user prompt to a specific session. See Prompt for +// the general contract; the difference is the explicit sessionId. +// +// The state.json Phase field is single-valued today, so the +// "PhaseRunning" stamp applies process-wide — when multiple sessions are +// in-flight concurrently, phase is "running" if any session is. Per- +// session phase tracking is a follow-up state-schema change. +func (m *Manager) PromptSession(ctx context.Context, sessionID acp.SessionId, prompt []acp.ContentBlock) (acp.PromptResponse, error) { + m.mu.Lock() + conn := m.conn + _, known := m.sessions[sessionID] + m.mu.Unlock() if conn == nil { return acp.PromptResponse{}, fmt.Errorf("runtime: agent not started") } - m.logger.Debug("prompt started", "blocks", len(prompt)) + if !known { + return acp.PromptResponse{}, fmt.Errorf("runtime: prompt: session %q not found", sessionID) + } + m.logger.Debug("prompt started", "sessionID", sessionID, "blocks", len(prompt)) _ = m.writeState(func(s *apiruntime.State) { s.Phase = apiruntime.PhaseRunning @@ -466,7 +488,7 @@ func (m *Manager) Prompt(ctx context.Context, prompt []acp.ContentBlock) (acp.Pr if err != nil { reason = "prompt-failed" } - m.logger.Debug("prompt done", "reason", reason) + m.logger.Debug("prompt done", "sessionID", sessionID, "reason", reason) _ = m.writeState(func(s *apiruntime.State) { s.Phase = apiruntime.PhaseIdle }, reason) @@ -478,17 +500,29 @@ func (m *Manager) Prompt(ctx context.Context, prompt []acp.ContentBlock) (acp.Pr return resp, nil } -// Cancel sends a cancel notification to the agent for the current session. +// Cancel sends a cancel notification to the agent for the initial session. +// Backward-compat shim around CancelSession. func (m *Manager) Cancel(ctx context.Context) error { m.mu.Lock() - conn := m.conn sessionID := m.sessionID m.mu.Unlock() + return m.CancelSession(ctx, sessionID) +} + +// CancelSession sends a cancel notification for a specific session. +func (m *Manager) CancelSession(ctx context.Context, sessionID acp.SessionId) error { + m.mu.Lock() + conn := m.conn + _, known := m.sessions[sessionID] + m.mu.Unlock() if conn == nil { return fmt.Errorf("runtime: agent not started") } - m.logger.Debug("cancel") + if !known { + return fmt.Errorf("runtime: cancel: session %q not found", sessionID) + } + m.logger.Debug("cancel", "sessionID", sessionID) if err := conn.Cancel(ctx, acp.CancelNotification{SessionId: sessionID}); err != nil { return fmt.Errorf("runtime: cancel: %w", err) @@ -496,17 +530,34 @@ func (m *Manager) Cancel(ctx context.Context) error { return nil } -// SetModel switches the agent to a different model via ACP session/set_model. +// SetModel switches the initial session to a different model. +// Backward-compat shim around SetModelSession. func (m *Manager) SetModel(ctx context.Context, modelID string) error { m.mu.Lock() - conn := m.conn sessionID := m.sessionID m.mu.Unlock() + return m.SetModelSession(ctx, sessionID, modelID) +} + +// SetModelSession switches a specific session to a different model via +// ACP session/set_model. Updates in-memory per-session models state + +// (only for the initial session) the legacy state.json Session.Models +// field. Per-session models persistence is a follow-up state-schema +// change. +func (m *Manager) SetModelSession(ctx context.Context, sessionID acp.SessionId, modelID string) error { + m.mu.Lock() + conn := m.conn + sess, known := m.sessions[sessionID] + initialID := m.sessionID + m.mu.Unlock() if conn == nil { return fmt.Errorf("runtime: agent not started") } - m.logger.Debug("set_model", "modelID", modelID) + if !known { + return fmt.Errorf("runtime: set_model: session %q not found", sessionID) + } + m.logger.Debug("set_model", "sessionID", sessionID, "modelID", modelID) _, err := conn.UnstableSetSessionModel(ctx, acp.UnstableSetSessionModelRequest{ SessionId: sessionID, @@ -516,19 +567,27 @@ func (m *Manager) SetModel(ctx context.Context, modelID string) error { return fmt.Errorf("runtime: set model: %w", err) } - // Update in-memory models state. + // Update per-session in-memory models state. m.mu.Lock() - if m.models != nil { + if sess.models != nil { + sess.models.CurrentModelId = acp.ModelId(modelID) //nolint:gosec // ModelId is string + } + // Legacy single-session models mirror for the initial session only. + if sessionID == initialID && m.models != nil { m.models.CurrentModelId = acp.ModelId(modelID) //nolint:gosec // ModelId is string } m.mu.Unlock() - // Persist updated currentModelId to state.json. - _ = m.writeState(func(s *apiruntime.State) { - if s.Session != nil && s.Session.Models != nil { - s.Session.Models.CurrentModelId = modelID - } - }, "set-model") + // state.json's Session.Models is single-session; only update for the + // initial session to keep legacy semantics. Multi-session persistence + // is a follow-up. + if sessionID == initialID { + _ = m.writeState(func(s *apiruntime.State) { + if s.Session != nil && s.Session.Models != nil { + s.Session.Models.CurrentModelId = modelID + } + }, "set-model") + } return nil } From 36594b87dfdb07f60e09b179c4c4ea1da1ccb0f5 Mon Sep 17 00:00:00 2001 From: wjueyao Date: Mon, 18 May 2026 18:12:31 +0800 Subject: [PATCH 03/16] feat(agentrun/api): wire types + method constants for multi-session MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Step 3 of multi-session refactor (steps 1-2: c16e57a, c17f298). Add wire types so RPC callers can address specific sessions: - SessionPromptParams: new optional SessionID field (empty = initial) - SessionCancelParams: new struct, optional SessionID - SessionSetModelParams: new optional SessionID field - SessionNewParams + SessionNewResult: open additional session - SessionEndParams + SessionEndResult: release runtime tracking - SessionListResult: enumerate active sessions Add method constants (api/methods.go): - session/new, session/end, session/list Backward-compat preserved: existing single-session callers pass empty SessionID (or no params at all where the old Cancel had none), which routes to the initial session — runtime layer (c17f298) maps empty sessionId → m.sessionID. No behavior change yet — server handlers + client methods are step 4-5. Co-Authored-By: Claude Opus 4.7 (1M context) --- pkg/agentrun/api/methods.go | 8 +++++ pkg/agentrun/api/types.go | 62 +++++++++++++++++++++++++++++++++++-- 2 files changed, 68 insertions(+), 2 deletions(-) diff --git a/pkg/agentrun/api/methods.go b/pkg/agentrun/api/methods.go index 8d470f6..a8a9ce1 100644 --- a/pkg/agentrun/api/methods.go +++ b/pkg/agentrun/api/methods.go @@ -9,6 +9,14 @@ const ( MethodSessionSetModel = "session/set_model" MethodRuntimePhase = "runtime/status" MethodRuntimeStop = "runtime/stop" + + // Multi-session RPCs — open / end additional ACP sessions on a long- + // lived agent process so callers can multiplex tasks without + // fork+exec per task. See pkg/agentrun/runtime/acp/runtime.go's + // NewSession / EndSession for the underlying runtime contract. + MethodSessionNew = "session/new" + MethodSessionEnd = "session/end" + MethodSessionList = "session/list" ) // Run notification methods. diff --git a/pkg/agentrun/api/types.go b/pkg/agentrun/api/types.go index 5cac3a3..08364ba 100644 --- a/pkg/agentrun/api/types.go +++ b/pkg/agentrun/api/types.go @@ -14,8 +14,14 @@ import ( // SessionPromptParams is the JSON body for the "session/prompt" method. // Prompt is an array of ACP ContentBlocks supporting text, image, audio, // resource, and resource-link content types. +// +// SessionID is optional: when empty, the agent-run's initial session +// (the one opened during Create's handshake) is used — preserves +// backward compatibility for single-session callers. Multi-session +// callers must pass an explicit SessionID obtained from session/new. type SessionPromptParams struct { - Prompt []ContentBlock `json:"prompt"` + SessionID string `json:"sessionId,omitempty"` + Prompt []ContentBlock `json:"prompt"` } // SessionPromptResult is returned by the "session/prompt" method. @@ -23,6 +29,14 @@ type SessionPromptResult struct { StopReason string `json:"stopReason"` } +// SessionCancelParams is the JSON body for the "session/cancel" method. +// Optional SessionID — same semantics as SessionPromptParams.SessionID. +// Pre-multi-session callers passed no params at all; both empty body and +// missing SessionID work as "cancel the initial session". +type SessionCancelParams struct { + SessionID string `json:"sessionId,omitempty"` +} + // SessionLoadParams is the JSON body for the "session/load" RPC method. // agentd always calls this during recovery for best-effort session restore. // agent-run checks ACP loadSession capability internally and auto-fallbacks. @@ -55,13 +69,57 @@ type RuntimeStatusRecovery struct { } // SessionSetModelParams is the JSON body for "session/set_model". +// Optional SessionID — same semantics as SessionPromptParams.SessionID. type SessionSetModelParams struct { - ModelID string `json:"modelId"` + SessionID string `json:"sessionId,omitempty"` + ModelID string `json:"modelId"` } // SessionSetModelResult is returned by "session/set_model". type SessionSetModelResult struct{} +// SessionNewParams is the JSON body for the "session/new" method — +// opens an additional ACP session on the running agent. +// +// Cwd is required: each session is scoped to its own working directory. +// McpServers is optional per-session MCP overrides layered on the +// bundle-level config. +type SessionNewParams struct { + Cwd string `json:"cwd"` + McpServers []SessionNewMcpServer `json:"mcpServers,omitempty"` +} + +// SessionNewMcpServer mirrors acp.McpServer's wire shape for transport +// over JSON-RPC. Kept minimal — extend as actual cases demand. +type SessionNewMcpServer struct { + Name string `json:"name"` + Command string `json:"command"` + Args []string `json:"args,omitempty"` + Env map[string]string `json:"env,omitempty"` +} + +// SessionNewResult is returned by "session/new". +type SessionNewResult struct { + SessionID string `json:"sessionId"` +} + +// SessionEndParams is the JSON body for "session/end" — removes runtime +// tracking of a session. The agent process's per-session state remains +// until cancelled or the process exits (ACP has no explicit end-session +// RPC); this method only releases the agent-run's local map entry. +type SessionEndParams struct { + SessionID string `json:"sessionId"` +} + +// SessionEndResult is returned by "session/end". +type SessionEndResult struct{} + +// SessionListResult lists active session IDs on the agent. Used by +// callers (e.g. massctl) to inspect the agent's current sessions. +type SessionListResult struct { + SessionIDs []string `json:"sessionIds"` +} + // RuntimePhaseResult is returned by "runtime/status". type RuntimePhaseResult struct { State apiruntime.State `json:"state"` From e93326418487857dba2f8cea97e9ea1d4019f880 Mon Sep 17 00:00:00 2001 From: wjueyao Date: Mon, 18 May 2026 18:18:17 +0800 Subject: [PATCH 04/16] feat(agentrun): wire RPC layer for multi-session + client methods MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Steps 4-5 of multi-session refactor (steps 1-3: c16e57a, c17f298, 36594b8). Server side (pkg/agentrun/server): - Handler interface: Cancel signature change (now takes *SessionCancelParams), add NewSession / EndSession / ListSessions - register.go: route session/new, session/end, session/list to new handlers - service.go: Prompt/Cancel/SetModel now check params.SessionID - empty → call Manager's legacy shim (initial session) - non-empty → call Manager's *Session variant (validates membership) - New handler impls translate wire params (SessionNewMcpServer → acp.McpServerStdio etc.) to runtime layer types Client side (pkg/agentrun/client/client.go): - Cancel sends *SessionCancelParams instead of nil (compat with new handler signature; empty SessionID still routes to initial session) - New methods: NewSession, EndSession, ListSessions, CancelSession - Existing Cancel kept for backward compatibility (delegates to initial) Tests updated to match new Handler interface (stubRunService + replayService in client tests). Backward compatibility: - Old clients calling Cancel with nil params: server handler tolerates nil req (treats as empty SessionID → initial session) - Old clients calling Prompt without SessionID field: same routing - All existing tests still pass; no callers broken Co-Authored-By: Claude Opus 4.7 (1M context) --- pkg/agentrun/client/client.go | 38 ++++++++- pkg/agentrun/client/client_test.go | 14 +++- pkg/agentrun/client/watch_test.go | 14 +++- pkg/agentrun/server/register.go | 12 ++- pkg/agentrun/server/service.go | 130 ++++++++++++++++++++++++++--- 5 files changed, 192 insertions(+), 16 deletions(-) diff --git a/pkg/agentrun/client/client.go b/pkg/agentrun/client/client.go index b3800bb..57aa562 100644 --- a/pkg/agentrun/client/client.go +++ b/pkg/agentrun/client/client.go @@ -52,8 +52,44 @@ func (c *Client) SendPrompt(ctx context.Context, req *runapi.SessionPromptParams return c.c.CallAsync(ctx, runapi.MethodSessionPrompt, req) } +// Cancel cancels in-flight work on the agent's initial session. +// To cancel a specific session opened via NewSession, use CancelSession. func (c *Client) Cancel(ctx context.Context) error { - return c.c.Call(ctx, runapi.MethodSessionCancel, nil, nil) + return c.c.Call(ctx, runapi.MethodSessionCancel, &runapi.SessionCancelParams{}, nil) +} + +// CancelSession cancels in-flight work on a specific session. +func (c *Client) CancelSession(ctx context.Context, sessionID string) error { + return c.c.Call(ctx, runapi.MethodSessionCancel, + &runapi.SessionCancelParams{SessionID: sessionID}, nil) +} + +// NewSession opens an additional ACP session on the running agent. cwd +// scopes the session's working directory; mcpServers are optional per- +// session MCP overrides. Returns the new session id which callers must +// pass to Prompt / Cancel / SetModel via the params SessionID field. +func (c *Client) NewSession(ctx context.Context, req *runapi.SessionNewParams) (*runapi.SessionNewResult, error) { + var result runapi.SessionNewResult + if err := c.c.Call(ctx, runapi.MethodSessionNew, req, &result); err != nil { + return nil, err + } + return &result, nil +} + +// EndSession releases runtime tracking of a session. The agent process's +// per-session state remains until cancelled or the process exits. +func (c *Client) EndSession(ctx context.Context, sessionID string) error { + return c.c.Call(ctx, runapi.MethodSessionEnd, + &runapi.SessionEndParams{SessionID: sessionID}, nil) +} + +// ListSessions returns the agent's active session ids. +func (c *Client) ListSessions(ctx context.Context) (*runapi.SessionListResult, error) { + var result runapi.SessionListResult + if err := c.c.Call(ctx, runapi.MethodSessionList, nil, &result); err != nil { + return nil, err + } + return &result, nil } func (c *Client) Load(ctx context.Context, req *runapi.SessionLoadParams) error { diff --git a/pkg/agentrun/client/client_test.go b/pkg/agentrun/client/client_test.go index 97b7079..c672c73 100644 --- a/pkg/agentrun/client/client_test.go +++ b/pkg/agentrun/client/client_test.go @@ -28,7 +28,7 @@ type stubRunService struct { func (s *stubRunService) Prompt(_ context.Context, req *runapi.SessionPromptParams) (*runapi.SessionPromptResult, error) { return &s.promptResult, nil } -func (s *stubRunService) Cancel(_ context.Context) error { return nil } +func (s *stubRunService) Cancel(_ context.Context, _ *runapi.SessionCancelParams) error { return nil } func (s *stubRunService) Load(_ context.Context, _ *runapi.SessionLoadParams) error { return nil } @@ -46,6 +46,18 @@ func (s *stubRunService) SetModel(_ context.Context, _ *runapi.SessionSetModelPa } func (s *stubRunService) Stop(_ context.Context) error { return nil } +func (s *stubRunService) NewSession(_ context.Context, req *runapi.SessionNewParams) (*runapi.SessionNewResult, error) { + return &runapi.SessionNewResult{SessionID: "stub-session-" + req.Cwd}, nil +} + +func (s *stubRunService) EndSession(_ context.Context, _ *runapi.SessionEndParams) (*runapi.SessionEndResult, error) { + return &runapi.SessionEndResult{}, nil +} + +func (s *stubRunService) ListSessions(_ context.Context) (*runapi.SessionListResult, error) { + return &runapi.SessionListResult{SessionIDs: []string{}}, nil +} + // startTestServer starts a jsonrpc.Server with Register on a temp socket. func startTestServer(t *testing.T, svc runserver.Handler) string { t.Helper() diff --git a/pkg/agentrun/client/watch_test.go b/pkg/agentrun/client/watch_test.go index dd60585..9500266 100644 --- a/pkg/agentrun/client/watch_test.go +++ b/pkg/agentrun/client/watch_test.go @@ -30,8 +30,8 @@ type replayService struct { func (s *replayService) Prompt(context.Context, *runapi.SessionPromptParams) (*runapi.SessionPromptResult, error) { return &runapi.SessionPromptResult{}, nil } -func (s *replayService) Cancel(context.Context) error { return nil } -func (s *replayService) Load(context.Context, *runapi.SessionLoadParams) error { return nil } +func (s *replayService) Cancel(context.Context, *runapi.SessionCancelParams) error { return nil } +func (s *replayService) Load(context.Context, *runapi.SessionLoadParams) error { return nil } func (s *replayService) SetModel(context.Context, *runapi.SessionSetModelParams) (*runapi.SessionSetModelResult, error) { return &runapi.SessionSetModelResult{}, nil } @@ -41,6 +41,16 @@ func (s *replayService) Status(context.Context) (*runapi.RuntimePhaseResult, err } func (s *replayService) Stop(context.Context) error { return nil } +func (s *replayService) NewSession(context.Context, *runapi.SessionNewParams) (*runapi.SessionNewResult, error) { + return &runapi.SessionNewResult{}, nil +} +func (s *replayService) EndSession(context.Context, *runapi.SessionEndParams) (*runapi.SessionEndResult, error) { + return &runapi.SessionEndResult{}, nil +} +func (s *replayService) ListSessions(context.Context) (*runapi.SessionListResult, error) { + return &runapi.SessionListResult{}, nil +} + func (s *replayService) WatchEvent(ctx context.Context, req *runapi.SessionWatchEventParams, watchID string) (*runapi.SessionWatchEventResult, error) { peer := jsonrpc.PeerFromContext(ctx) diff --git a/pkg/agentrun/server/register.go b/pkg/agentrun/server/register.go index 7bb53a7..a1edb34 100644 --- a/pkg/agentrun/server/register.go +++ b/pkg/agentrun/server/register.go @@ -17,7 +17,7 @@ type watchEventWire struct { // These are the methods exposed by agent-run over a Unix socket. type Handler interface { Prompt(ctx context.Context, req *runapi.SessionPromptParams) (*runapi.SessionPromptResult, error) - Cancel(ctx context.Context) error + Cancel(ctx context.Context, req *runapi.SessionCancelParams) error Load(ctx context.Context, req *runapi.SessionLoadParams) error // WatchEvent implements K8s List-Watch style event subscription. // When FromSeq is nil, only live events are streamed. @@ -34,6 +34,11 @@ type Handler interface { SetModel(ctx context.Context, req *runapi.SessionSetModelParams) (*runapi.SessionSetModelResult, error) Status(ctx context.Context) (*runapi.RuntimePhaseResult, error) Stop(ctx context.Context) error + + // Multi-session methods (see runtime/acp Manager.NewSession etc.). + NewSession(ctx context.Context, req *runapi.SessionNewParams) (*runapi.SessionNewResult, error) + EndSession(ctx context.Context, req *runapi.SessionEndParams) (*runapi.SessionEndResult, error) + ListSessions(ctx context.Context) (*runapi.SessionListResult, error) } // Register registers a Handler implementation with the server. @@ -41,9 +46,12 @@ func Register(s *jsonrpc.Server, svc Handler) { s.RegisterService("session", &jsonrpc.ServiceDesc{ Methods: map[string]jsonrpc.Method{ "prompt": jsonrpc.UnaryMethod(svc.Prompt), - "cancel": jsonrpc.NullaryCommand(svc.Cancel), + "cancel": jsonrpc.UnaryCommand(svc.Cancel), "load": jsonrpc.UnaryCommand(svc.Load), "set_model": jsonrpc.UnaryMethod(svc.SetModel), + "new": jsonrpc.UnaryMethod(svc.NewSession), + "end": jsonrpc.UnaryMethod(svc.EndSession), + "list": jsonrpc.NullaryMethod(svc.ListSessions), }, }) s.RegisterService("runtime", &jsonrpc.ServiceDesc{ diff --git a/pkg/agentrun/server/service.go b/pkg/agentrun/server/service.go index f1d8164..9c3b72e 100644 --- a/pkg/agentrun/server/service.go +++ b/pkg/agentrun/server/service.go @@ -11,6 +11,14 @@ import ( "github.com/zoumo/mass/pkg/jsonrpc" ) +// resolveSessionID maps the wire-level optional sessionId string to the +// typed acp.SessionId. Empty → the Manager's initial session (preserves +// pre-multi-session caller behavior — runtime layer maps the zero +// SessionId to m.sessionID via its Prompt/Cancel/SetModel shims). +func resolveSessionID(s string) acp.SessionId { + return acp.SessionId(s) +} + // Service implements Handler. type Service struct { mgr *acpruntime.Manager @@ -27,10 +35,21 @@ func (s *Service) Prompt(ctx context.Context, req *runapi.SessionPromptParams) ( if len(req.Prompt) == 0 { return nil, jsonrpc.ErrInvalidParams("missing prompt") } - s.logger.Debug("prompt", "blocks", len(req.Prompt)) + s.logger.Debug("prompt", "sessionId", req.SessionID, "blocks", len(req.Prompt)) s.trans.NotifyTurnStart() s.trans.NotifyUserPrompt(req.Prompt) - resp, err := s.mgr.Prompt(ctx, req.Prompt) + + // Empty SessionID → use initial session (Manager.Prompt is the shim). + // Non-empty → route to specific session (PromptSession validates it + // exists in the sessions map). + var resp acp.PromptResponse + var err error + if req.SessionID == "" { + resp, err = s.mgr.Prompt(ctx, req.Prompt) + } else { + resp, err = s.mgr.PromptSession(ctx, resolveSessionID(req.SessionID), req.Prompt) + } + stopReason := "error" if err == nil { stopReason = string(resp.StopReason) @@ -39,19 +58,35 @@ func (s *Service) Prompt(ctx context.Context, req *runapi.SessionPromptParams) ( s.trans.NotifyError(err.Error()) } s.trans.NotifyTurnEnd(acp.StopReason(stopReason)) - s.logger.Debug("prompt done", "stopReason", stopReason) + s.logger.Debug("prompt done", "sessionId", req.SessionID, "stopReason", stopReason) if err != nil { return nil, jsonrpc.ErrInternal(err.Error()) } return &runapi.SessionPromptResult{StopReason: string(resp.StopReason)}, nil } -func (s *Service) Cancel(ctx context.Context) (retErr error) { - s.logger.Debug("cancel") +func (s *Service) Cancel(ctx context.Context, req *runapi.SessionCancelParams) (retErr error) { + // req may be nil if caller sent no params (pre-multi-session clients). + sessionID := "" + if req != nil { + sessionID = req.SessionID + } + s.logger.Debug("cancel", "sessionId", sessionID) defer func() { - s.trans.NotifyOperationAudit("cancel", nil, retErr) + var auditArgs map[string]string + if sessionID != "" { + auditArgs = map[string]string{"sessionId": sessionID} + } + s.trans.NotifyOperationAudit("cancel", auditArgs, retErr) }() - if err := s.mgr.Cancel(ctx); err != nil { + + var err error + if sessionID == "" { + err = s.mgr.Cancel(ctx) + } else { + err = s.mgr.CancelSession(ctx, resolveSessionID(sessionID)) + } + if err != nil { retErr = jsonrpc.ErrInternal(err.Error()) return retErr } @@ -214,15 +249,25 @@ func (s *Service) Status(_ context.Context) (*runapi.RuntimePhaseResult, error) } func (s *Service) SetModel(ctx context.Context, req *runapi.SessionSetModelParams) (_ *runapi.SessionSetModelResult, retErr error) { - s.logger.Debug("set_model", "modelID", req.ModelID) + s.logger.Debug("set_model", "sessionId", req.SessionID, "modelID", req.ModelID) defer func() { - s.trans.NotifyOperationAudit("set_model", map[string]string{"modelId": req.ModelID}, retErr) + auditArgs := map[string]string{"modelId": req.ModelID} + if req.SessionID != "" { + auditArgs["sessionId"] = req.SessionID + } + s.trans.NotifyOperationAudit("set_model", auditArgs, retErr) }() if req.ModelID == "" { retErr = jsonrpc.ErrInvalidParams("missing modelId") return nil, retErr } - if err := s.mgr.SetModel(ctx, req.ModelID); err != nil { + var err error + if req.SessionID == "" { + err = s.mgr.SetModel(ctx, req.ModelID) + } else { + err = s.mgr.SetModelSession(ctx, resolveSessionID(req.SessionID), req.ModelID) + } + if err != nil { retErr = jsonrpc.ErrInternal(err.Error()) return nil, retErr } @@ -234,3 +279,68 @@ func (s *Service) Stop(_ context.Context) error { s.trans.NotifyOperationAudit("stop", nil, nil) return nil } + +// NewSession opens an additional ACP session on the running agent. See +// runtime/acp Manager.NewSession for cwd / mcpServers semantics. +func (s *Service) NewSession(ctx context.Context, req *runapi.SessionNewParams) (_ *runapi.SessionNewResult, retErr error) { + s.logger.Debug("session/new", "cwd", req.Cwd, "mcpServers", len(req.McpServers)) + defer func() { + s.trans.NotifyOperationAudit("session/new", map[string]string{"cwd": req.Cwd}, retErr) + }() + if req.Cwd == "" { + retErr = jsonrpc.ErrInvalidParams("missing cwd") + return nil, retErr + } + + mcp := make([]acp.McpServer, 0, len(req.McpServers)) + for _, m := range req.McpServers { + envVars := make([]acp.EnvVariable, 0, len(m.Env)) + for k, v := range m.Env { + envVars = append(envVars, acp.EnvVariable{Name: k, Value: v}) + } + args := m.Args + if args == nil { + args = []string{} + } + mcp = append(mcp, acp.McpServer{Stdio: &acp.McpServerStdio{ + Name: m.Name, + Command: m.Command, + Args: args, + Env: envVars, + }}) + } + + sid, err := s.mgr.NewSession(ctx, req.Cwd, mcp) + if err != nil { + retErr = jsonrpc.ErrInternal(err.Error()) + return nil, retErr + } + return &runapi.SessionNewResult{SessionID: string(sid)}, nil +} + +// EndSession releases runtime tracking of a session id. +func (s *Service) EndSession(_ context.Context, req *runapi.SessionEndParams) (_ *runapi.SessionEndResult, retErr error) { + s.logger.Debug("session/end", "sessionId", req.SessionID) + defer func() { + s.trans.NotifyOperationAudit("session/end", map[string]string{"sessionId": req.SessionID}, retErr) + }() + if req.SessionID == "" { + retErr = jsonrpc.ErrInvalidParams("missing sessionId") + return nil, retErr + } + if err := s.mgr.EndSession(resolveSessionID(req.SessionID)); err != nil { + retErr = jsonrpc.ErrInternal(err.Error()) + return nil, retErr + } + return &runapi.SessionEndResult{}, nil +} + +// ListSessions returns the active session IDs snapshot from the Manager. +func (s *Service) ListSessions(_ context.Context) (*runapi.SessionListResult, error) { + ids := s.mgr.Sessions() + out := make([]string, len(ids)) + for i, id := range ids { + out[i] = string(id) + } + return &runapi.SessionListResult{SessionIDs: out}, nil +} From 601260ea5c41cc8967137117425dafd7a574f98f Mon Sep 17 00:00:00 2001 From: wjueyao Date: Mon, 18 May 2026 18:22:55 +0800 Subject: [PATCH 05/16] feat(ari): agentrun/new-session, agentrun/end-session, agentrun/list-sessions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Step 6 of multi-session refactor (steps 1-5: c16e57a, c17f298, 36594b8, e933264, f16480a inlining). Adds daemon-level RPCs that forward to the running agent-run process: API (pkg/ari/api): - methods.go: MethodAgentRunNewSession / EndSession / ListSessions - types.go: AgentRunNewSessionParams/Result + AgentRunEndSessionParams/Result + AgentRunListSessionsParams/Result - AgentRunPromptParams: new optional SessionID field — forwards to agent-run's session/prompt SessionID Server (pkg/ari/server): - AgentRunService interface: 3 new methods - service.go RegisterAgentRunService: route new-session / end-session / list-sessions on agentrun service - server.go agentRunAdapter: handler impls that Connect to agent-run and call the runapi client methods added in e933264. Prompt forwards the new SessionID field through to agent-run. Backward compatibility: - AgentRunPromptParams.SessionID is optional (omitempty); pre-multi- session callers pass nothing → routes to agent's initial session Co-Authored-By: Claude Opus 4.7 (1M context) --- pkg/ari/api/methods.go | 9 ++++++ pkg/ari/api/types.go | 45 ++++++++++++++++++++++++++ pkg/ari/server/server.go | 67 +++++++++++++++++++++++++++++++++++++-- pkg/ari/server/service.go | 10 ++++++ 4 files changed, 129 insertions(+), 2 deletions(-) diff --git a/pkg/ari/api/methods.go b/pkg/ari/api/methods.go index e22813c..7852b77 100644 --- a/pkg/ari/api/methods.go +++ b/pkg/ari/api/methods.go @@ -28,6 +28,15 @@ const ( MethodAgentRunTaskGet = "agentrun/task/get" MethodAgentRunTaskList = "agentrun/task/list" MethodAgentRunTaskRetry = "agentrun/task/retry" + + // Multi-session lifecycle for a single agentrun. session/new opens + // an additional ACP session (different cwd, fresh state) without + // fork+exec a new agent process; session/end releases runtime + // tracking; session/list enumerates active sessions. See + // pkg/agentrun/runtime/acp Manager.NewSession for runtime contract. + MethodAgentRunNewSession = "agentrun/new-session" + MethodAgentRunEndSession = "agentrun/end-session" + MethodAgentRunListSessions = "agentrun/list-sessions" ) // ARI agent definition methods. diff --git a/pkg/ari/api/types.go b/pkg/ari/api/types.go index 8aaddeb..2345459 100644 --- a/pkg/ari/api/types.go +++ b/pkg/ari/api/types.go @@ -89,6 +89,11 @@ type AgentRunPromptParams struct { // Name is the agent run name (required). Name string `json:"name"` + // SessionID addresses a specific session opened via agentrun/new-session. + // Empty (omitted) routes to the agentrun's initial session — preserves + // pre-multi-session caller behavior. + SessionID string `json:"sessionId,omitempty"` + // Prompt is an array of ACP ContentBlocks (text, image, audio, etc.) (required). Prompt []runapi.ContentBlock `json:"prompt"` } @@ -99,6 +104,46 @@ type AgentRunPromptResult struct { Accepted bool `json:"accepted"` } +// AgentRunNewSessionParams is the request params for agentrun/new-session. +// Opens an additional ACP session on an existing agentrun (no fork+exec). +type AgentRunNewSessionParams struct { + Workspace string `json:"workspace"` + Name string `json:"name"` + // Cwd is required — each session is scoped to its own working directory. + Cwd string `json:"cwd"` + // McpServers is optional per-session MCP overrides. Wire shape mirrors + // runapi.SessionNewMcpServer to avoid re-defining transports. + McpServers []runapi.SessionNewMcpServer `json:"mcpServers,omitempty"` +} + +// AgentRunNewSessionResult is the response for agentrun/new-session. +type AgentRunNewSessionResult struct { + // SessionID is the ACP session id the caller passes to subsequent + // agentrun/prompt (etc.) via the SessionID field. + SessionID string `json:"sessionId"` +} + +// AgentRunEndSessionParams is the request params for agentrun/end-session. +type AgentRunEndSessionParams struct { + Workspace string `json:"workspace"` + Name string `json:"name"` + SessionID string `json:"sessionId"` +} + +// AgentRunEndSessionResult is the response for agentrun/end-session. +type AgentRunEndSessionResult struct{} + +// AgentRunListSessionsParams identifies the agentrun whose sessions to list. +type AgentRunListSessionsParams struct { + Workspace string `json:"workspace"` + Name string `json:"name"` +} + +// AgentRunListSessionsResult enumerates the agentrun's active session ids. +type AgentRunListSessionsResult struct { + SessionIDs []string `json:"sessionIds"` +} + // WorkspaceSendParams is the request params for workspace/send method. // Routes a message from one agent run to another within a workspace. type WorkspaceSendParams struct { diff --git a/pkg/ari/server/server.go b/pkg/ari/server/server.go index 62fb343..bb7de93 100644 --- a/pkg/ari/server/server.go +++ b/pkg/ari/server/server.go @@ -516,17 +516,80 @@ func (a *agentRunAdapter) Prompt(ctx context.Context, req *pkgariapi.AgentRunPro } prompt := req.Prompt - if err := client.SendPrompt(ctx, &runapi.SessionPromptParams{Prompt: prompt}); err != nil { + if err := client.SendPrompt(ctx, &runapi.SessionPromptParams{ + SessionID: req.SessionID, + Prompt: prompt, + }); err != nil { a.logger.Warn("agentrun/prompt: prompt delivery failed", "workspace", req.Workspace, "name", req.Name, "error", err) a.recordPromptDeliveryFailure(req.Workspace, req.Name, agent.Status, err, false) } a.logger.Info("agentrun/prompt: dispatched", - "workspace", req.Workspace, "name", req.Name) + "workspace", req.Workspace, "name", req.Name, "sessionId", req.SessionID) return &pkgariapi.AgentRunPromptResult{Accepted: true}, nil } +// NewSession forwards agentrun/new-session to the running agent-run's +// session/new RPC. The agent reuses its process to host the new session +// — no fork+exec. Returns the agent-issued sessionId. +func (a *agentRunAdapter) NewSession(ctx context.Context, req *pkgariapi.AgentRunNewSessionParams) (*pkgariapi.AgentRunNewSessionResult, error) { + if req.Workspace == "" || req.Name == "" || req.Cwd == "" { + return nil, jsonrpc.ErrInvalidParams("workspace, name, and cwd are required") + } + a.logger.Info("agentrun/new-session", "workspace", req.Workspace, "name", req.Name, "cwd", req.Cwd) + + client, err := a.processes.Connect(ctx, req.Workspace, req.Name) + if err != nil { + return nil, &jsonrpc.RPCError{Code: pkgariapi.CodeRecoveryBlocked, Message: "agent not running"} + } + + out, err := client.NewSession(ctx, &runapi.SessionNewParams{ + Cwd: req.Cwd, + McpServers: req.McpServers, + }) + if err != nil { + return nil, jsonrpc.ErrInternal(err.Error()) + } + a.logger.Info("agentrun/new-session: opened", + "workspace", req.Workspace, "name", req.Name, "sessionId", out.SessionID) + return &pkgariapi.AgentRunNewSessionResult{SessionID: out.SessionID}, nil +} + +// EndSession forwards agentrun/end-session to release runtime tracking. +func (a *agentRunAdapter) EndSession(ctx context.Context, req *pkgariapi.AgentRunEndSessionParams) (*pkgariapi.AgentRunEndSessionResult, error) { + if req.Workspace == "" || req.Name == "" || req.SessionID == "" { + return nil, jsonrpc.ErrInvalidParams("workspace, name, and sessionId are required") + } + a.logger.Info("agentrun/end-session", "workspace", req.Workspace, "name", req.Name, "sessionId", req.SessionID) + + client, err := a.processes.Connect(ctx, req.Workspace, req.Name) + if err != nil { + return nil, &jsonrpc.RPCError{Code: pkgariapi.CodeRecoveryBlocked, Message: "agent not running"} + } + if err := client.EndSession(ctx, req.SessionID); err != nil { + return nil, jsonrpc.ErrInternal(err.Error()) + } + return &pkgariapi.AgentRunEndSessionResult{}, nil +} + +// ListSessions forwards agentrun/list-sessions to enumerate active sessions. +func (a *agentRunAdapter) ListSessions(ctx context.Context, req *pkgariapi.AgentRunListSessionsParams) (*pkgariapi.AgentRunListSessionsResult, error) { + if req.Workspace == "" || req.Name == "" { + return nil, jsonrpc.ErrInvalidParams("workspace and name are required") + } + + client, err := a.processes.Connect(ctx, req.Workspace, req.Name) + if err != nil { + return nil, &jsonrpc.RPCError{Code: pkgariapi.CodeRecoveryBlocked, Message: "agent not running"} + } + out, err := client.ListSessions(ctx) + if err != nil { + return nil, jsonrpc.ErrInternal(err.Error()) + } + return &pkgariapi.AgentRunListSessionsResult{SessionIDs: out.SessionIDs}, nil +} + // Cancel handles agentrun/cancel. // // Connects to the running agent-run and calls Cancel. diff --git a/pkg/ari/server/service.go b/pkg/ari/server/service.go index 926782d..b76c570 100644 --- a/pkg/ari/server/service.go +++ b/pkg/ari/server/service.go @@ -36,6 +36,11 @@ type AgentRunService interface { TaskGet(ctx context.Context, params *pkgariapi.AgentRunTaskGetParams) (*pkgariapi.AgentTask, error) TaskList(ctx context.Context, params *pkgariapi.AgentRunTaskListParams) (*pkgariapi.AgentRunTaskListResult, error) TaskRetry(ctx context.Context, params *pkgariapi.AgentRunTaskRetryParams) (*pkgariapi.AgentTask, error) + + // Multi-session lifecycle. + NewSession(ctx context.Context, req *pkgariapi.AgentRunNewSessionParams) (*pkgariapi.AgentRunNewSessionResult, error) + EndSession(ctx context.Context, req *pkgariapi.AgentRunEndSessionParams) (*pkgariapi.AgentRunEndSessionResult, error) + ListSessions(ctx context.Context, req *pkgariapi.AgentRunListSessionsParams) (*pkgariapi.AgentRunListSessionsResult, error) } // AgentService defines agent definition CRUD methods. @@ -157,6 +162,11 @@ func RegisterAgentRunService(s *jsonrpc.Server, svc AgentRunService) { "task/get": jsonrpc.UnaryMethod(svc.TaskGet), "task/list": jsonrpc.UnaryMethod(svc.TaskList), "task/retry": jsonrpc.UnaryMethod(svc.TaskRetry), + // Multi-session lifecycle (see pkg/agentrun/runtime/acp + // Manager.NewSession for runtime contract). + "new-session": jsonrpc.UnaryMethod(svc.NewSession), + "end-session": jsonrpc.UnaryMethod(svc.EndSession), + "list-sessions": jsonrpc.UnaryMethod(svc.ListSessions), }, }) } From e46e1b5b7279c6987ed89bfbcb2c9b30f067719d Mon Sep 17 00:00:00 2001 From: wjueyao Date: Mon, 18 May 2026 18:29:38 +0800 Subject: [PATCH 06/16] feat(massctl, ari/client): expose multi-session via CLI + ARI client API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Step 7 of multi-session refactor. ARI client (pkg/ari/client): - interfaces.go AgentRunOps: 4 new methods — PromptSession, NewSession, EndSession, ListSessions. Prompt() now delegates to PromptSession("") for backward compat. - client.go agentRunOps: impl using new ARI methods from f16480a massctl (cmd/massctl/commands/agentrun): - new-session.go: `massctl ar new-session -w --cwd ` prints sessionId to stdout (script-friendly) - end-session.go: `massctl ar end-session -w --session-id X` - list-sessions.go: `massctl ar list-sessions -w ` prints active session ids one per line - prompt.go: new --session-id flag routes prompts to specific session - command.go: register the 3 new subcommands Mocks (cmd/massctl/commands/{agent,workspace,workspace/create}/mock_test.go): - mockAgentRunOps stubs the 4 new interface methods (no-ops; not exercised by callers in those packages) Backward compatibility: - `massctl ar prompt` without --session-id → routes to initial session (same as before via PromptSession with empty sessionId) - existing tests pass; new commands have no dedicated tests yet (added in next commit) Co-Authored-By: Claude Opus 4.7 (1M context) --- cmd/massctl/commands/agent/mock_test.go | 14 +++++ cmd/massctl/commands/agentrun/command.go | 3 + cmd/massctl/commands/agentrun/end_session.go | 46 ++++++++++++++ .../commands/agentrun/list_sessions.go | 46 ++++++++++++++ cmd/massctl/commands/agentrun/mock_test.go | 17 ++++++ cmd/massctl/commands/agentrun/new_session.go | 60 +++++++++++++++++++ cmd/massctl/commands/agentrun/prompt.go | 16 +++-- .../commands/workspace/create/command_test.go | 14 +++++ cmd/massctl/commands/workspace/mock_test.go | 14 +++++ pkg/ari/client/client.go | 32 ++++++++++ pkg/ari/client/interfaces.go | 16 ++++- 11 files changed, 272 insertions(+), 6 deletions(-) create mode 100644 cmd/massctl/commands/agentrun/end_session.go create mode 100644 cmd/massctl/commands/agentrun/list_sessions.go create mode 100644 cmd/massctl/commands/agentrun/new_session.go diff --git a/cmd/massctl/commands/agent/mock_test.go b/cmd/massctl/commands/agent/mock_test.go index 87004c1..4150399 100644 --- a/cmd/massctl/commands/agent/mock_test.go +++ b/cmd/massctl/commands/agent/mock_test.go @@ -91,6 +91,20 @@ func (m *mockAgentRunOps) TaskRetry(context.Context, *pkgariapi.AgentRunTaskRetr return nil, nil } +// Multi-session stubs (no-ops — agent command tests don't exercise them). +func (m *mockAgentRunOps) PromptSession(context.Context, pkgariapi.ObjectKey, string, []runapi.ContentBlock) (*pkgariapi.AgentRunPromptResult, error) { + return nil, nil +} +func (m *mockAgentRunOps) NewSession(context.Context, *pkgariapi.AgentRunNewSessionParams) (*pkgariapi.AgentRunNewSessionResult, error) { + return nil, nil +} +func (m *mockAgentRunOps) EndSession(context.Context, pkgariapi.ObjectKey, string) error { + return nil +} +func (m *mockAgentRunOps) ListSessions(context.Context, pkgariapi.ObjectKey) (*pkgariapi.AgentRunListSessionsResult, error) { + return nil, nil +} + // mock WorkspaceOps (stub — not used in agent tests) type mockWorkspaceOps struct{} diff --git a/cmd/massctl/commands/agentrun/command.go b/cmd/massctl/commands/agentrun/command.go index 62809a0..54dda31 100644 --- a/cmd/massctl/commands/agentrun/command.go +++ b/cmd/massctl/commands/agentrun/command.go @@ -64,5 +64,8 @@ Poll with: massctl ar get -w `, cmd.AddCommand(newTaskCmd(getClient)) cmd.AddCommand(newChatCmd(getClient)) cmd.AddCommand(newDebugCmd()) + cmd.AddCommand(newNewSessionCmd(getClient)) + cmd.AddCommand(newEndSessionCmd(getClient)) + cmd.AddCommand(newListSessionsCmd(getClient)) return cmd } diff --git a/cmd/massctl/commands/agentrun/end_session.go b/cmd/massctl/commands/agentrun/end_session.go new file mode 100644 index 0000000..275915e --- /dev/null +++ b/cmd/massctl/commands/agentrun/end_session.go @@ -0,0 +1,46 @@ +package agentrun + +import ( + "context" + "fmt" + + "github.com/spf13/cobra" + + "github.com/zoumo/mass/cmd/massctl/commands/cliutil" + pkgariapi "github.com/zoumo/mass/pkg/ari/api" +) + +// newEndSessionCmd implements “massctl agentrun end-session“. +// +// Releases runtime tracking of a session id. The agent process keeps its +// per-session state until cancelled or process exits — ACP has no +// explicit end-session RPC. Refuses to end an agent's initial session +// (the one created by the agentrun handshake); kill the whole agent +// instead via “stop“. +func newEndSessionCmd(getClient cliutil.ClientFn) *cobra.Command { + var ws, sessionID string + cmd := &cobra.Command{ + Use: "end-session name", + Short: "Release runtime tracking of an ACP session", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + client, err := getClient() + if err != nil { + return err + } + defer client.Close() + + if err := client.AgentRuns().EndSession(context.Background(), + pkgariapi.ObjectKey{Workspace: ws, Name: args[0]}, sessionID); err != nil { + return fmt.Errorf("ending session %s on %s/%s: %w", sessionID, ws, args[0], err) + } + fmt.Fprintf(cmd.OutOrStdout(), "session %s ended on %s/%s\n", sessionID, ws, args[0]) + return nil + }, + } + cmd.Flags().StringVarP(&ws, "workspace", "w", "", "Workspace name (required)") + cmd.Flags().StringVar(&sessionID, "session-id", "", "Session id to end (required) — must not be the initial session") + _ = cmd.MarkFlagRequired("workspace") + _ = cmd.MarkFlagRequired("session-id") + return cmd +} diff --git a/cmd/massctl/commands/agentrun/list_sessions.go b/cmd/massctl/commands/agentrun/list_sessions.go new file mode 100644 index 0000000..630f6dd --- /dev/null +++ b/cmd/massctl/commands/agentrun/list_sessions.go @@ -0,0 +1,46 @@ +package agentrun + +import ( + "context" + "fmt" + + "github.com/spf13/cobra" + + "github.com/zoumo/mass/cmd/massctl/commands/cliutil" + pkgariapi "github.com/zoumo/mass/pkg/ari/api" +) + +// newListSessionsCmd implements “massctl agentrun list-sessions“. +// +// Prints active session ids on an agent-run, one per line. Useful for +// scripts that drive multi-session pools and need to enumerate or +// clean up sessions. +func newListSessionsCmd(getClient cliutil.ClientFn) *cobra.Command { + var ws string + cmd := &cobra.Command{ + Use: "list-sessions name", + Aliases: []string{"ls-sessions"}, + Short: "List active ACP session ids on an agent-run", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + client, err := getClient() + if err != nil { + return err + } + defer client.Close() + + out, err := client.AgentRuns().ListSessions(context.Background(), + pkgariapi.ObjectKey{Workspace: ws, Name: args[0]}) + if err != nil { + return fmt.Errorf("listing sessions on %s/%s: %w", ws, args[0], err) + } + for _, id := range out.SessionIDs { + fmt.Fprintln(cmd.OutOrStdout(), id) + } + return nil + }, + } + cmd.Flags().StringVarP(&ws, "workspace", "w", "", "Workspace name (required)") + _ = cmd.MarkFlagRequired("workspace") + return cmd +} diff --git a/cmd/massctl/commands/agentrun/mock_test.go b/cmd/massctl/commands/agentrun/mock_test.go index 6bd0c7c..fc57f0d 100644 --- a/cmd/massctl/commands/agentrun/mock_test.go +++ b/cmd/massctl/commands/agentrun/mock_test.go @@ -79,6 +79,23 @@ func (m *mockAgentRunOps) TaskRetry(ctx context.Context, params *pkgariapi.Agent return &pkgariapi.AgentTask{}, nil } +// Multi-session stubs — tests in this package don't exercise these yet. +func (m *mockAgentRunOps) PromptSession(ctx context.Context, key pkgariapi.ObjectKey, sessionID string, prompt []runapi.ContentBlock) (*pkgariapi.AgentRunPromptResult, error) { + return m.Prompt(ctx, key, prompt) +} + +func (m *mockAgentRunOps) NewSession(_ context.Context, _ *pkgariapi.AgentRunNewSessionParams) (*pkgariapi.AgentRunNewSessionResult, error) { + return &pkgariapi.AgentRunNewSessionResult{SessionID: "mock-session"}, nil +} + +func (m *mockAgentRunOps) EndSession(_ context.Context, _ pkgariapi.ObjectKey, _ string) error { + return nil +} + +func (m *mockAgentRunOps) ListSessions(_ context.Context, _ pkgariapi.ObjectKey) (*pkgariapi.AgentRunListSessionsResult, error) { + return &pkgariapi.AgentRunListSessionsResult{}, nil +} + // ── mock WorkspaceOps (stub — not used in agentrun tests) ──────────────────── type mockWorkspaceOps struct{} diff --git a/cmd/massctl/commands/agentrun/new_session.go b/cmd/massctl/commands/agentrun/new_session.go new file mode 100644 index 0000000..dc3099d --- /dev/null +++ b/cmd/massctl/commands/agentrun/new_session.go @@ -0,0 +1,60 @@ +package agentrun + +import ( + "context" + "fmt" + + "github.com/spf13/cobra" + + "github.com/zoumo/mass/cmd/massctl/commands/cliutil" + pkgariapi "github.com/zoumo/mass/pkg/ari/api" +) + +// newNewSessionCmd implements “massctl agentrun new-session“. +// +// Opens an additional ACP session on a running agent-run — agent process +// is reused (no fork+exec), but the session has its own cwd and state. +// Returns the sessionId to stdout; pass it via subsequent “prompt +// --session-id“ etc. +func newNewSessionCmd(getClient cliutil.ClientFn) *cobra.Command { + var ws, cwd string + cmd := &cobra.Command{ + Use: "new-session name", + Short: "Open an additional ACP session on a running agent-run", + Long: `Opens an additional ACP session on the running agent process. + +The session is multiplexed onto the existing process — no new fork/exec +— but has its own cwd, model state, and message history. Returns the +agent-issued sessionId on stdout; pass it to subsequent prompt / cancel +/ end-session via --session-id.`, + Example: ` # Open a fresh session scoped to /tmp/case-7 + sid=$(massctl ar new-session worker -w my-ws --cwd /tmp/case-7) + massctl ar prompt worker -w my-ws --session-id "$sid" --text "Fix the bug" + massctl ar end-session worker -w my-ws --session-id "$sid"`, + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + client, err := getClient() + if err != nil { + return err + } + defer client.Close() + + out, err := client.AgentRuns().NewSession(context.Background(), &pkgariapi.AgentRunNewSessionParams{ + Workspace: ws, + Name: args[0], + Cwd: cwd, + }) + if err != nil { + return fmt.Errorf("opening session for %s/%s: %w", ws, args[0], err) + } + // Print sessionId only (script-friendly). + fmt.Fprintln(cmd.OutOrStdout(), out.SessionID) + return nil + }, + } + cmd.Flags().StringVarP(&ws, "workspace", "w", "", "Workspace name (required)") + cmd.Flags().StringVar(&cwd, "cwd", "", "Session cwd (required) — the working directory the agent uses for this session") + _ = cmd.MarkFlagRequired("workspace") + _ = cmd.MarkFlagRequired("cwd") + return cmd +} diff --git a/cmd/massctl/commands/agentrun/prompt.go b/cmd/massctl/commands/agentrun/prompt.go index d991f1b..f1e3c3c 100644 --- a/cmd/massctl/commands/agentrun/prompt.go +++ b/cmd/massctl/commands/agentrun/prompt.go @@ -17,9 +17,10 @@ import ( func newPromptCmd(getClient cliutil.ClientFn) *cobra.Command { var ( - ws string - text string - wait bool + ws string + text string + wait bool + sessionID string ) cmd := &cobra.Command{ Use: "prompt name", @@ -37,7 +38,9 @@ func newPromptCmd(getClient cliutil.ClientFn) *cobra.Command { key := pkgariapi.ObjectKey{Workspace: ws, Name: name} if !wait { - result, err := client.AgentRuns().Prompt(ctx, key, []runapi.ContentBlock{runapi.TextBlock(text)}) + // PromptSession with empty sessionID == Prompt — single entry point. + result, err := client.AgentRuns().PromptSession(ctx, key, sessionID, + []runapi.ContentBlock{runapi.TextBlock(text)}) if err != nil { return err } @@ -75,7 +78,8 @@ func newPromptCmd(getClient cliutil.ClientFn) *cobra.Command { // Send prompt (fire-and-forget). if err := runClient.SendPrompt(ctx, &runapi.SessionPromptParams{ - Prompt: []runapi.ContentBlock{runapi.TextBlock(text)}, + SessionID: sessionID, + Prompt: []runapi.ContentBlock{runapi.TextBlock(text)}, }); err != nil { return fmt.Errorf("send_prompt: %w", err) } @@ -109,6 +113,8 @@ func newPromptCmd(getClient cliutil.ClientFn) *cobra.Command { cmd.Flags().StringVarP(&ws, "workspace", "w", "", "Workspace name (required)") cmd.Flags().StringVar(&text, "text", "", "Prompt text (required)") cmd.Flags().BoolVar(&wait, "wait", false, "Wait for turn to complete and print agent response") + cmd.Flags().StringVar(&sessionID, "session-id", "", + "Session id to prompt (defaults to the agent's initial session). Use new-session to open one.") _ = cmd.MarkFlagRequired("workspace") _ = cmd.MarkFlagRequired("text") return cmd diff --git a/cmd/massctl/commands/workspace/create/command_test.go b/cmd/massctl/commands/workspace/create/command_test.go index a38b464..e7b649a 100644 --- a/cmd/massctl/commands/workspace/create/command_test.go +++ b/cmd/massctl/commands/workspace/create/command_test.go @@ -43,6 +43,20 @@ func (m *mockAgentRunOps) TaskRetry(context.Context, *pkgariapi.AgentRunTaskRetr return &pkgariapi.AgentTask{}, nil } +// Multi-session stubs (no-ops). +func (m *mockAgentRunOps) PromptSession(context.Context, pkgariapi.ObjectKey, string, []runapi.ContentBlock) (*pkgariapi.AgentRunPromptResult, error) { + return &pkgariapi.AgentRunPromptResult{}, nil +} +func (m *mockAgentRunOps) NewSession(context.Context, *pkgariapi.AgentRunNewSessionParams) (*pkgariapi.AgentRunNewSessionResult, error) { + return &pkgariapi.AgentRunNewSessionResult{}, nil +} +func (m *mockAgentRunOps) EndSession(context.Context, pkgariapi.ObjectKey, string) error { + return nil +} +func (m *mockAgentRunOps) ListSessions(context.Context, pkgariapi.ObjectKey) (*pkgariapi.AgentRunListSessionsResult, error) { + return &pkgariapi.AgentRunListSessionsResult{}, nil +} + // ── mock WorkspaceOps (stub — not used in create tests) ────────────────────── type mockWorkspaceOps struct{} diff --git a/cmd/massctl/commands/workspace/mock_test.go b/cmd/massctl/commands/workspace/mock_test.go index ea0a117..4efaaa3 100644 --- a/cmd/massctl/commands/workspace/mock_test.go +++ b/cmd/massctl/commands/workspace/mock_test.go @@ -59,6 +59,20 @@ func (m *mockAgentRunOps) TaskRetry(context.Context, *pkgariapi.AgentRunTaskRetr return &pkgariapi.AgentTask{}, nil } +// Multi-session stubs (no-ops — workspace command tests don't exercise them). +func (m *mockAgentRunOps) PromptSession(context.Context, pkgariapi.ObjectKey, string, []runapi.ContentBlock) (*pkgariapi.AgentRunPromptResult, error) { + return &pkgariapi.AgentRunPromptResult{}, nil +} +func (m *mockAgentRunOps) NewSession(context.Context, *pkgariapi.AgentRunNewSessionParams) (*pkgariapi.AgentRunNewSessionResult, error) { + return &pkgariapi.AgentRunNewSessionResult{}, nil +} +func (m *mockAgentRunOps) EndSession(context.Context, pkgariapi.ObjectKey, string) error { + return nil +} +func (m *mockAgentRunOps) ListSessions(context.Context, pkgariapi.ObjectKey) (*pkgariapi.AgentRunListSessionsResult, error) { + return &pkgariapi.AgentRunListSessionsResult{}, nil +} + // ── mock SystemOps (stub — not used in workspace tests) ──────────────────────── type mockSystemOps struct{} diff --git a/pkg/ari/client/client.go b/pkg/ari/client/client.go index 338186e..2ddfae0 100644 --- a/pkg/ari/client/client.go +++ b/pkg/ari/client/client.go @@ -143,9 +143,14 @@ func (c *ariClient) Delete(ctx context.Context, key pkgariapi.ObjectKey, obj pkg type agentRunOps struct{ c *jsonrpc.Client } func (o *agentRunOps) Prompt(ctx context.Context, key pkgariapi.ObjectKey, prompt []runapi.ContentBlock) (*pkgariapi.AgentRunPromptResult, error) { + return o.PromptSession(ctx, key, "", prompt) +} + +func (o *agentRunOps) PromptSession(ctx context.Context, key pkgariapi.ObjectKey, sessionID string, prompt []runapi.ContentBlock) (*pkgariapi.AgentRunPromptResult, error) { req := pkgariapi.AgentRunPromptParams{ Workspace: key.Workspace, Name: key.Name, + SessionID: sessionID, Prompt: prompt, } var result pkgariapi.AgentRunPromptResult @@ -205,6 +210,33 @@ func (o *agentRunOps) TaskRetry(ctx context.Context, params *pkgariapi.AgentRunT return &result, nil } +func (o *agentRunOps) NewSession(ctx context.Context, params *pkgariapi.AgentRunNewSessionParams) (*pkgariapi.AgentRunNewSessionResult, error) { + var result pkgariapi.AgentRunNewSessionResult + if err := o.c.Call(ctx, pkgariapi.MethodAgentRunNewSession, params, &result); err != nil { + return nil, err + } + return &result, nil +} + +func (o *agentRunOps) EndSession(ctx context.Context, key pkgariapi.ObjectKey, sessionID string) error { + req := pkgariapi.AgentRunEndSessionParams{ + Workspace: key.Workspace, + Name: key.Name, + SessionID: sessionID, + } + var raw json.RawMessage + return o.c.Call(ctx, pkgariapi.MethodAgentRunEndSession, req, &raw) +} + +func (o *agentRunOps) ListSessions(ctx context.Context, key pkgariapi.ObjectKey) (*pkgariapi.AgentRunListSessionsResult, error) { + req := pkgariapi.AgentRunListSessionsParams{Workspace: key.Workspace, Name: key.Name} + var result pkgariapi.AgentRunListSessionsResult + if err := o.c.Call(ctx, pkgariapi.MethodAgentRunListSessions, req, &result); err != nil { + return nil, err + } + return &result, nil +} + // ──────────────────────────────────────────────────────────────────────────── // WorkspaceOps // ──────────────────────────────────────────────────────────────────────────── diff --git a/pkg/ari/client/interfaces.go b/pkg/ari/client/interfaces.go index fb2e2cb..5fdaa0c 100644 --- a/pkg/ari/client/interfaces.go +++ b/pkg/ari/client/interfaces.go @@ -48,9 +48,13 @@ type Client interface { // AgentRunOps provides non-CRUD operations on agent runs. type AgentRunOps interface { - // Prompt sends a multimodal prompt ([]runapi.ContentBlock) to an agent run. + // Prompt sends a multimodal prompt ([]runapi.ContentBlock) to an agent run's + // initial session. For multi-session, use PromptSession. Prompt(ctx context.Context, key pkgariapi.ObjectKey, prompt []runapi.ContentBlock) (*pkgariapi.AgentRunPromptResult, error) + // PromptSession addresses a specific session id (opened via NewSession). + PromptSession(ctx context.Context, key pkgariapi.ObjectKey, sessionID string, prompt []runapi.ContentBlock) (*pkgariapi.AgentRunPromptResult, error) + // Cancel cancels the current turn of an agent run. Cancel(ctx context.Context, key pkgariapi.ObjectKey) error @@ -71,6 +75,16 @@ type AgentRunOps interface { // TaskRetry retries an existing task by bumping its attempt count and re-prompting the agent. TaskRetry(ctx context.Context, params *pkgariapi.AgentRunTaskRetryParams) (*pkgariapi.AgentTask, error) + + // NewSession opens an additional ACP session on the running agent process + // (no fork+exec). Returns the agent-issued sessionId. + NewSession(ctx context.Context, params *pkgariapi.AgentRunNewSessionParams) (*pkgariapi.AgentRunNewSessionResult, error) + + // EndSession releases runtime tracking of a session id. + EndSession(ctx context.Context, key pkgariapi.ObjectKey, sessionID string) error + + // ListSessions enumerates active session ids on the agent. + ListSessions(ctx context.Context, key pkgariapi.ObjectKey) (*pkgariapi.AgentRunListSessionsResult, error) } // WorkspaceOps provides non-CRUD operations on workspaces. From 926547eb53d41da7eaad84b7b0dd312f9999e636 Mon Sep 17 00:00:00 2001 From: wjueyao Date: Mon, 18 May 2026 18:32:32 +0800 Subject: [PATCH 07/16] test(runtime/acp): bookkeeping tests for multi-session map operations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Step 8 of multi-session refactor — final code commit. Covers the parts of multi-session not exercised by existing tests (which all need a real ACP agent process and are pre-existing flaky): - TestSessions_EmptyWhenNoneRegistered — Sessions() on fresh Manager - TestSessions_SnapshotIncludesRegistered — Sessions() enumerates map - TestEndSession_RemovesFromMap — happy path - TestEndSession_RefusesInitialSession — invariant: initial session cannot be ended without killing the agent - TestEndSession_UnknownSessionErrors — clear error for stale ids - TestPromptSession_RejectsBeforeAgentStarted — nil conn → clean error - TestCancelSession_RejectsBeforeAgentStarted — same for cancel Uses a test helper `newManagerForSessionTest` that bypasses Create() to test bookkeeping logic without spawning an agent. Manager methods that go through ACP (anything touching m.conn) return a documented "agent not started" error for callers that race or skip Create(). The pre-existing TestRuntimeSuite tests using a fake bash agent remain flaky (peer disconnect race) — confirmed unrelated to this PR by running on origin/main baseline. Not addressed here. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../runtime/acp/runtime_internal_test.go | 129 ++++++++++++++++++ 1 file changed, 129 insertions(+) diff --git a/pkg/agentrun/runtime/acp/runtime_internal_test.go b/pkg/agentrun/runtime/acp/runtime_internal_test.go index bb3568d..999cb7b 100644 --- a/pkg/agentrun/runtime/acp/runtime_internal_test.go +++ b/pkg/agentrun/runtime/acp/runtime_internal_test.go @@ -1,8 +1,13 @@ package acp import ( + "context" + "log/slog" + "sort" "strings" "testing" + + acp "github.com/coder/acp-go-sdk" ) func TestBuildSeedSystemPrompt_AppendsGuard(t *testing.T) { @@ -20,3 +25,127 @@ func TestBuildSeedSystemPrompt_AppendsGuard(t *testing.T) { t.Fatalf("expected seed prompt to contain %q, got %q", want, seed) } } + +// newManagerForSessionTest builds a Manager with the sessions map +// initialized, bypassing Create() (which needs a real ACP agent +// process). Lets unit tests exercise the session bookkeeping logic +// (NewSession map insert, EndSession map delete, Sessions snapshot, +// error paths) without spinning up an agent. +// +// The returned Manager has nil conn; methods that go through ACP +// (PromptSession etc.) return "agent not started". +func newManagerForSessionTest(t *testing.T) *Manager { + t.Helper() + return &Manager{ + logger: slog.Default(), + sessions: make(map[acp.SessionId]*sessionState), + } +} + +func TestSessions_EmptyWhenNoneRegistered(t *testing.T) { + m := newManagerForSessionTest(t) + if got := m.Sessions(); len(got) != 0 { + t.Fatalf("expected empty sessions, got %v", got) + } +} + +func TestSessions_SnapshotIncludesRegistered(t *testing.T) { + m := newManagerForSessionTest(t) + m.sessions["sess-a"] = &sessionState{id: "sess-a", cwd: "/a"} + m.sessions["sess-b"] = &sessionState{id: "sess-b", cwd: "/b"} + m.sessionID = "sess-a" + + got := m.Sessions() + if len(got) != 2 { + t.Fatalf("expected 2 sessions, got %v", got) + } + strs := make([]string, len(got)) + for i, s := range got { + strs[i] = string(s) + } + sort.Strings(strs) + want := []string{"sess-a", "sess-b"} + for i := range want { + if strs[i] != want[i] { + t.Fatalf("session %d: want %q got %q", i, want[i], strs[i]) + } + } +} + +func TestEndSession_RemovesFromMap(t *testing.T) { + m := newManagerForSessionTest(t) + m.sessionID = "initial" + m.sessions["initial"] = &sessionState{id: "initial"} + m.sessions["extra"] = &sessionState{id: "extra", cwd: "/x"} + + if err := m.EndSession("extra"); err != nil { + t.Fatalf("end extra: %v", err) + } + if _, ok := m.sessions["extra"]; ok { + t.Fatalf("extra session not removed from map") + } + if _, ok := m.sessions["initial"]; !ok { + t.Fatalf("initial session unexpectedly removed") + } +} + +func TestEndSession_RefusesInitialSession(t *testing.T) { + m := newManagerForSessionTest(t) + m.sessionID = "initial" + m.sessions["initial"] = &sessionState{id: "initial"} + + err := m.EndSession("initial") + if err == nil { + t.Fatalf("expected EndSession to refuse the initial session") + } + if !strings.Contains(err.Error(), "initial session") { + t.Fatalf("expected error to mention initial session, got: %v", err) + } + if _, ok := m.sessions["initial"]; !ok { + t.Fatalf("initial session removed despite error") + } +} + +func TestEndSession_UnknownSessionErrors(t *testing.T) { + m := newManagerForSessionTest(t) + m.sessionID = "initial" + m.sessions["initial"] = &sessionState{id: "initial"} + + err := m.EndSession("never-existed") + if err == nil { + t.Fatalf("expected EndSession to error on unknown session") + } + if !strings.Contains(err.Error(), "not found") { + t.Fatalf("expected error to mention 'not found', got: %v", err) + } +} + +func TestPromptSession_RejectsBeforeAgentStarted(t *testing.T) { + // Manager with no conn (Create never ran). PromptSession should + // return a clear error rather than panic on nil conn — protects + // callers that race or skip Create(). + m := newManagerForSessionTest(t) + m.sessions["known"] = &sessionState{id: "known"} + + _, err := m.PromptSession(context.Background(), "known", + []acp.ContentBlock{acp.TextBlock("hi")}) + if err == nil { + t.Fatalf("expected error from PromptSession with nil conn") + } + if !strings.Contains(err.Error(), "agent not started") { + t.Fatalf("expected 'agent not started' error, got: %v", err) + } +} + +func TestCancelSession_RejectsBeforeAgentStarted(t *testing.T) { + m := newManagerForSessionTest(t) + m.sessions["known"] = &sessionState{id: "known"} + + err := m.CancelSession(context.Background(), "known") + if err == nil { + t.Fatalf("expected error from CancelSession with nil conn") + } + if !strings.Contains(err.Error(), "agent not started") { + t.Fatalf("expected 'agent not started' error, got: %v", err) + } +} From 70fb47138d37e20cb0880fc0205b93ac715994c4 Mon Sep 17 00:00:00 2001 From: wjueyao Date: Mon, 18 May 2026 20:18:30 +0800 Subject: [PATCH 08/16] fix(ari/api): preserve camelCase YAML keys via explicit yaml tags MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AgentSpec fields rely on gopkg.in/yaml.v3 for round-trip via 'massctl agent apply -f'. Without explicit `yaml:` tags, the library lowercases Go field names (StartupTimeoutSeconds → startuptimeoutseconds), silently dropping the camelCase keys that users and tooling write to match the JSON tag spelling. Add explicit yaml tags matching each json tag so that camelCase YAML input is preserved byte-for-byte on apply. --- pkg/ari/api/domain.go | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/pkg/ari/api/domain.go b/pkg/ari/api/domain.go index e69d74f..909c6a5 100644 --- a/pkg/ari/api/domain.go +++ b/pkg/ari/api/domain.go @@ -72,27 +72,34 @@ type ObjectMeta struct { // ──────────────────────────────────────────────────────────────────────────── // AgentSpec describes how to launch an agent process for this named agent definition. +// +// All fields carry both `json` and `yaml` tags. Without explicit `yaml:` tags, +// gopkg.in/yaml.v3 lowercases the Go field name (e.g. `StartupTimeoutSeconds` +// → `startuptimeoutseconds`), which silently drops camelCase YAML keys that +// users / tooling write to match the JSON tag (`startupTimeoutSeconds`). Each +// camelCase field needs an explicit yaml tag for `massctl agent apply -f` to +// preserve it round-trip. type AgentSpec struct { // Disabled controls whether the agent is prevented from creating new agent runs. // nil or false means not disabled (agent is usable). true means disabled. - Disabled *bool `json:"disabled,omitempty"` + Disabled *bool `json:"disabled,omitempty" yaml:"disabled,omitempty"` // ClientProtocol selects the communication protocol adapter. // Default: "acp". - ClientProtocol apiruntime.ClientProtocol `json:"clientProtocol,omitempty"` + ClientProtocol apiruntime.ClientProtocol `json:"clientProtocol,omitempty" yaml:"clientProtocol,omitempty"` // Command is the agent executable. - Command string `json:"command"` + Command string `json:"command" yaml:"command"` // Args are the command-line arguments passed to Command. - Args []string `json:"args,omitempty"` + Args []string `json:"args,omitempty" yaml:"args,omitempty"` // Env is the list of environment variable overrides applied to the process. - Env []apiruntime.EnvVar `json:"env,omitempty"` + Env []apiruntime.EnvVar `json:"env,omitempty" yaml:"env,omitempty"` // StartupTimeoutSeconds is the maximum time (in seconds) to wait for the // agent-run to reach idle state. Nil means use the daemon default. - StartupTimeoutSeconds *int `json:"startupTimeoutSeconds,omitempty"` + StartupTimeoutSeconds *int `json:"startupTimeoutSeconds,omitempty" yaml:"startupTimeoutSeconds,omitempty"` } // IsDisabled reports whether the agent is disabled. From 36a024511254d2aa9d4a5cfdbdd3a58590d4e2d0 Mon Sep 17 00:00:00 2001 From: wjueyao Date: Mon, 18 May 2026 21:16:38 +0800 Subject: [PATCH 09/16] feat(jsonrpc): OptionalUnaryCommand for absent-params back-compat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit session/cancel evolved from NullaryCommand (no params) to taking an optional sessionId. UnaryCommand strict-parses params and returns InvalidParams when the field is absent — which means pre-multi-session clients (sending the old paramless cancel) get a wire-level rejection before the Service handler's nil-guard ever runs. This commit makes the absent-params case typed and tolerable: 1. jsonrpc.ErrNoParams sentinel returned by the dispatcher's unmarshal callback when req.Params == nil. Replaces an ad-hoc fmt.Errorf — UnaryCommand semantics unchanged (still wraps as InvalidParams). 2. OptionalUnaryCommand helper — same shape as UnaryCommand but tolerates ErrNoParams (handler runs with zero-value Req). Malformed JSON unmarshal errors still surface as InvalidParams; tolerance is scoped to absent params. 3. agentrun/server/register.go: "session/cancel" switches to OptionalUnaryCommand. The Service handler reads req.SessionID directly — req is always non-nil from the helper. Tests pin the contract: - TestOptionalUnaryCommand_TolerateMissingParams: paramless call reaches handler with zero-value Req, no -32602 error. - TestOptionalUnaryCommand_SurfacesUnmarshalError: malformed params still return InvalidParams. Co-Authored-By: Claude Opus 4.7 (1M context) --- pkg/agentrun/server/register.go | 6 ++- pkg/jsonrpc/errors.go | 13 ++++++- pkg/jsonrpc/method_helpers.go | 22 ++++++++++- pkg/jsonrpc/server.go | 2 +- pkg/jsonrpc/server_test.go | 65 +++++++++++++++++++++++++++++++++ 5 files changed, 103 insertions(+), 5 deletions(-) diff --git a/pkg/agentrun/server/register.go b/pkg/agentrun/server/register.go index a1edb34..9d9d121 100644 --- a/pkg/agentrun/server/register.go +++ b/pkg/agentrun/server/register.go @@ -45,8 +45,10 @@ type Handler interface { func Register(s *jsonrpc.Server, svc Handler) { s.RegisterService("session", &jsonrpc.ServiceDesc{ Methods: map[string]jsonrpc.Method{ - "prompt": jsonrpc.UnaryMethod(svc.Prompt), - "cancel": jsonrpc.UnaryCommand(svc.Cancel), + "prompt": jsonrpc.UnaryMethod(svc.Prompt), + // cancel tolerates absent params (back-compat for callers that + // omit the field — pre-multi-session wire shape). + "cancel": jsonrpc.OptionalUnaryCommand(svc.Cancel), "load": jsonrpc.UnaryCommand(svc.Load), "set_model": jsonrpc.UnaryMethod(svc.SetModel), "new": jsonrpc.UnaryMethod(svc.NewSession), diff --git a/pkg/jsonrpc/errors.go b/pkg/jsonrpc/errors.go index 8eb5f39..61ff54d 100644 --- a/pkg/jsonrpc/errors.go +++ b/pkg/jsonrpc/errors.go @@ -2,7 +2,18 @@ // framework built on top of sourcegraph/jsonrpc2. package jsonrpc -import "fmt" +import ( + "errors" + "fmt" +) + +// ErrNoParams is returned by the dispatcher's unmarshal callback when the +// JSON-RPC request has no `params` field. UnaryCommand / UnaryMethod treat +// this as InvalidParams (the typical case — handler expects params). +// OptionalUnaryCommand tolerates it (leaves Req at zero value), letting a +// method evolve from NullaryCommand → optional-params without breaking +// callers that omit the field. +var ErrNoParams = errors.New("missing params") // RPCError is a JSON-RPC 2.0 error with code, message, and optional data. // Method handlers return *RPCError to control the JSON-RPC error response; diff --git a/pkg/jsonrpc/method_helpers.go b/pkg/jsonrpc/method_helpers.go index 17031c2..ee7a4c0 100644 --- a/pkg/jsonrpc/method_helpers.go +++ b/pkg/jsonrpc/method_helpers.go @@ -1,6 +1,9 @@ package jsonrpc -import "context" +import ( + "context" + "errors" +) func UnaryMethod[Req, Res any](fn func(ctx context.Context, req *Req) (*Res, error)) Method { return func(ctx context.Context, unmarshal func(any) error) (any, error) { @@ -33,3 +36,20 @@ func NullaryCommand(fn func(ctx context.Context) error) Method { return nil, fn(ctx) } } + +// OptionalUnaryCommand is like UnaryCommand but tolerates absent params: +// callers that omit the JSON-RPC `params` field reach the handler with +// Req at zero value rather than getting InvalidParams. Use this for +// methods evolving from NullaryCommand to taking optional params — +// pre-existing callers that send no params keep working unchanged. +// +// Malformed-JSON unmarshal errors still surface as InvalidParams. +func OptionalUnaryCommand[Req any](fn func(ctx context.Context, req *Req) error) Method { + return func(ctx context.Context, unmarshal func(any) error) (any, error) { + var req Req + if err := unmarshal(&req); err != nil && !errors.Is(err, ErrNoParams) { + return nil, ErrInvalidParams(err.Error()) + } + return nil, fn(ctx, &req) + } +} diff --git a/pkg/jsonrpc/server.go b/pkg/jsonrpc/server.go index c614613..93c4451 100644 --- a/pkg/jsonrpc/server.go +++ b/pkg/jsonrpc/server.go @@ -150,7 +150,7 @@ func (h *serverHandler) Handle(ctx context.Context, conn *jsonrpc2.Conn, req *js unmarshal := func(dst any) error { if req.Params == nil { - return fmt.Errorf("missing params") + return ErrNoParams } return json.Unmarshal(*req.Params, dst) } diff --git a/pkg/jsonrpc/server_test.go b/pkg/jsonrpc/server_test.go index 0c3c9d1..d14617c 100644 --- a/pkg/jsonrpc/server_test.go +++ b/pkg/jsonrpc/server_test.go @@ -223,6 +223,71 @@ func TestServer_PeerNotify(t *testing.T) { } } +// TestOptionalUnaryCommand_TolerateMissingParams verifies that +// OptionalUnaryCommand tolerates a JSON-RPC request with no params field +// (or null params) by delivering a zero-value Req to the handler instead +// of returning -32602 InvalidParams. This pins the back-compat contract +// for methods evolving from NullaryCommand → optional-params. +func TestOptionalUnaryCommand_TolerateMissingParams(t *testing.T) { + type req struct { + Name string `json:"name,omitempty"` + } + var captured req + called := make(chan struct{}, 1) + + srv := jsonrpc.NewServer(slog.Default()) + srv.RegisterService("svc", &jsonrpc.ServiceDesc{ + Methods: map[string]jsonrpc.Method{ + "cancel": jsonrpc.OptionalUnaryCommand(func(_ context.Context, r *req) error { + captured = *r + called <- struct{}{} + return nil + }), + }, + }) + + addr := startTestServer(t, srv) + client := dialTestClient(t, addr) + + // Pass nil — client should send "params: null" or omit the field. Either + // way OptionalUnaryCommand must run the handler with zero-value Req. + err := client.Call(context.Background(), "svc/cancel", nil, nil) + require.NoError(t, err) + + select { + case <-called: + case <-time.After(2 * time.Second): + t.Fatal("handler not invoked") + } + assert.Equal(t, "", captured.Name, "expected zero-value Req from missing params") +} + +// TestOptionalUnaryCommand_SurfacesUnmarshalError verifies that malformed +// params (not absent params) still surface as InvalidParams — the +// tolerance is scoped to ErrNoParams, not all unmarshal failures. +func TestOptionalUnaryCommand_SurfacesUnmarshalError(t *testing.T) { + type req struct { + Value int `json:"value"` + } + + srv := jsonrpc.NewServer(slog.Default()) + srv.RegisterService("svc", &jsonrpc.ServiceDesc{ + Methods: map[string]jsonrpc.Method{ + "op": jsonrpc.OptionalUnaryCommand(func(_ context.Context, _ *req) error { + return nil + }), + }, + }) + + addr := startTestServer(t, srv) + client := dialTestClient(t, addr) + + // Send malformed params — string where int expected. + err := client.Call(context.Background(), "svc/op", map[string]string{"value": "not-a-number"}, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "-32602") +} + func TestServer_PeerDisconnect(t *testing.T) { disconnectDetected := make(chan struct{}) From 9f66f503814ba6e37ad9fda2526676804288ff6c Mon Sep 17 00:00:00 2001 From: wjueyao Date: Mon, 18 May 2026 21:17:03 +0800 Subject: [PATCH 10/16] refactor(runtime/acp): single-point sessionID resolution + inflight guard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses two issues from PR review feedback. M1 — collapse empty-sessionID → initial resolution to Manager layer. Before, the "empty sessionId means initial session" convention was implemented in three places: client (passes ""), Service (branches on "" vs non-empty and calls Manager.Prompt vs PromptSession), Manager (Manager.Prompt is a shim that reads m.sessionID and calls PromptSession). Three layers to keep in sync if the convention ever changes. PromptSession / CancelSession / SetModelSession now do the resolution themselves via a small resolveSessionLocked helper. Service collapses to a one-line pass-through. Manager.Prompt / Cancel / SetModel reduce to single-line shims around the *Session variants with "". M3 — EndSession refuses sessions with in-flight prompt work. Previously EndSession only checked map membership; ending a session mid-prompt left the prompt running against a session no longer tracked, silently drifting state.json Phase / log lines from the sessions map. Add a per-session `inflight` counter on sessionState (collapsed from a parallel `Manager.inflight` map suggested in initial draft — keeps session bookkeeping in one place). PromptSession increments on entry under m.mu, decrements via defer (looking up the session by id again under lock to tolerate teardown). EndSession returns a busy error if inflight > 0; caller must CancelSession (or wait) first. Tests: - TestEndSession_RefusesBusySession: inflight=1 → busy error, session stays in map. - TestEndSession_AllowsAfterRefcountClears: once inflight returns to 0, EndSession succeeds — busy is per-state, not permanent. - TestPromptSession_EmptySessionIDResolvesToInitial: empty sessionID reaches a "session not found" branch iff resolution didn't happen; test asserts it doesn't. Co-Authored-By: Claude Opus 4.7 (1M context) --- pkg/agentrun/runtime/acp/runtime.go | 117 +++++++++++------- .../runtime/acp/runtime_internal_test.go | 73 +++++++++++ pkg/agentrun/server/service.go | 43 ++----- 3 files changed, 155 insertions(+), 78 deletions(-) diff --git a/pkg/agentrun/runtime/acp/runtime.go b/pkg/agentrun/runtime/acp/runtime.go index 1d79bdf..52c36db 100644 --- a/pkg/agentrun/runtime/acp/runtime.go +++ b/pkg/agentrun/runtime/acp/runtime.go @@ -71,6 +71,9 @@ type sessionState struct { id acp.SessionId models *acp.SessionModelState cwd string + // inflight counts active PromptSession turns. EndSession refuses to + // release a session with inflight > 0; the prompt's defer decrements it. + inflight int } type Manager struct { @@ -87,7 +90,8 @@ type Manager struct { // sessions holds all active ACP sessions on this agent process, keyed // by sessionId. NewSession adds entries; EndSession removes them. - // Empty until Create()'s initial session/new completes. + // Empty until Create()'s initial session/new completes. Per-session + // inflight prompt count lives on each sessionState — see EndSession. sessions map[acp.SessionId]*sessionState // sessionID retains the *first* session created by Create() so legacy @@ -405,11 +409,15 @@ func (m *Manager) NewSession(ctx context.Context, cwd string, mcpServers []acp.M // EndSession releases runtime tracking of a session id. The ACP protocol // has no explicit "end session" RPC today, so this method only clears the // runtime's local map entry — the agent process retains its own session -// state until cancelled or the process exits. Callers must still call -// CancelSession (or wait for prompt completion) before EndSession to ensure -// no in-flight work on the session. +// state until cancelled or the process exits. // -// Refusing to end the initial session (the one Create() opened) — callers +// EndSession refuses to release a session that has in-flight Prompt work +// (returns ErrSessionBusy). Callers must CancelSession (or wait for prompt +// completion) first; otherwise the prompt completes against a session no +// longer tracked, and the resulting state.json Phase / log lines diverge +// from the sessions map. +// +// Refuses to end the initial session (the one Create() opened) — callers // that want to fully tear down the agent should call Kill() instead. func (m *Manager) EndSession(sessionID acp.SessionId) error { m.mu.Lock() @@ -417,14 +425,27 @@ func (m *Manager) EndSession(sessionID acp.SessionId) error { if sessionID == m.sessionID { return fmt.Errorf("runtime: cannot end initial session %q (kill the agent instead)", sessionID) } - if _, ok := m.sessions[sessionID]; !ok { + sess, ok := m.sessions[sessionID] + if !ok { return fmt.Errorf("runtime: session %q not found", sessionID) } + if sess.inflight > 0 { + return fmt.Errorf("runtime: session %q busy: %d in-flight prompt(s) (cancel first)", sessionID, sess.inflight) + } delete(m.sessions, sessionID) m.logger.Info("session ended", "sessionID", sessionID) return nil } +// resolveSessionLocked maps an empty sessionID to the initial session. +// Caller must hold m.mu. +func (m *Manager) resolveSessionLocked(sessionID acp.SessionId) acp.SessionId { + if sessionID == "" { + return m.sessionID + } + return sessionID +} + // Sessions returns a snapshot of currently-active session IDs. Order is // not stable — caller should sort if a deterministic order is needed. func (m *Manager) Sessions() []acp.SessionId { @@ -437,41 +458,50 @@ func (m *Manager) Sessions() []acp.SessionId { return ids } -// Prompt sends a user prompt to the agent's initial session and blocks -// until the agent returns a PromptResponse. Session notifications emitted -// by the agent during the turn are forwarded to the Events channel. -// On completion (success or error), state.json is updated. -// -// This is a backward-compat shim around PromptSession that uses the -// initial sessionId set by Create(). New callers should prefer -// PromptSession with an explicit sessionId — especially when the agent -// holds multiple sessions opened via NewSession. +// Prompt is a backward-compat shim — equivalent to PromptSession(ctx, "", prompt), +// which resolves empty SessionId to the initial session. func (m *Manager) Prompt(ctx context.Context, prompt []acp.ContentBlock) (acp.PromptResponse, error) { - m.mu.Lock() - sessionID := m.sessionID - m.mu.Unlock() - return m.PromptSession(ctx, sessionID, prompt) + return m.PromptSession(ctx, "", prompt) } -// PromptSession sends a user prompt to a specific session. See Prompt for -// the general contract; the difference is the explicit sessionId. +// PromptSession sends a user prompt to a specific session and blocks until +// the agent returns a PromptResponse. Empty sessionID is resolved to the +// agent's initial session (the one Create() opened) — single resolution +// point for the empty-string convention. +// +// Session notifications emitted by the agent during the turn are forwarded +// to the Events channel. On completion (success or error), state.json is +// updated. // -// The state.json Phase field is single-valued today, so the -// "PhaseRunning" stamp applies process-wide — when multiple sessions are -// in-flight concurrently, phase is "running" if any session is. Per- -// session phase tracking is a follow-up state-schema change. +// The state.json Phase field is single-valued today, so the "PhaseRunning" +// stamp applies process-wide — when multiple sessions are in-flight +// concurrently, phase is "running" if any session is. Per-session phase +// tracking is a follow-up state-schema change. func (m *Manager) PromptSession(ctx context.Context, sessionID acp.SessionId, prompt []acp.ContentBlock) (acp.PromptResponse, error) { m.mu.Lock() conn := m.conn - _, known := m.sessions[sessionID] - m.mu.Unlock() - + sessionID = m.resolveSessionLocked(sessionID) + sess, known := m.sessions[sessionID] if conn == nil { + m.mu.Unlock() return acp.PromptResponse{}, fmt.Errorf("runtime: agent not started") } if !known { + m.mu.Unlock() return acp.PromptResponse{}, fmt.Errorf("runtime: prompt: session %q not found", sessionID) } + sess.inflight++ + m.mu.Unlock() + + defer func() { + m.mu.Lock() + // sess may have been deleted by Kill/Delete; guard the deref. + if s, ok := m.sessions[sessionID]; ok { + s.inflight-- + } + m.mu.Unlock() + }() + m.logger.Debug("prompt started", "sessionID", sessionID, "blocks", len(prompt)) _ = m.writeState(func(s *apiruntime.State) { @@ -500,19 +530,17 @@ func (m *Manager) PromptSession(ctx context.Context, sessionID acp.SessionId, pr return resp, nil } -// Cancel sends a cancel notification to the agent for the initial session. -// Backward-compat shim around CancelSession. +// Cancel is a backward-compat shim — equivalent to CancelSession(ctx, ""). func (m *Manager) Cancel(ctx context.Context) error { - m.mu.Lock() - sessionID := m.sessionID - m.mu.Unlock() - return m.CancelSession(ctx, sessionID) + return m.CancelSession(ctx, "") } -// CancelSession sends a cancel notification for a specific session. +// CancelSession sends a cancel notification for a specific session. Empty +// sessionID resolves to the initial session. func (m *Manager) CancelSession(ctx context.Context, sessionID acp.SessionId) error { m.mu.Lock() conn := m.conn + sessionID = m.resolveSessionLocked(sessionID) _, known := m.sessions[sessionID] m.mu.Unlock() @@ -530,23 +558,20 @@ func (m *Manager) CancelSession(ctx context.Context, sessionID acp.SessionId) er return nil } -// SetModel switches the initial session to a different model. -// Backward-compat shim around SetModelSession. +// SetModel is a backward-compat shim — equivalent to SetModelSession(ctx, "", modelID). func (m *Manager) SetModel(ctx context.Context, modelID string) error { - m.mu.Lock() - sessionID := m.sessionID - m.mu.Unlock() - return m.SetModelSession(ctx, sessionID, modelID) + return m.SetModelSession(ctx, "", modelID) } -// SetModelSession switches a specific session to a different model via -// ACP session/set_model. Updates in-memory per-session models state + -// (only for the initial session) the legacy state.json Session.Models -// field. Per-session models persistence is a follow-up state-schema -// change. +// SetModelSession switches a specific session to a different model via ACP +// session/set_model. Empty sessionID resolves to the initial session. +// Updates in-memory per-session models state; for the initial session +// also mirrors to the legacy state.json Session.Models field. Per-session +// models persistence is a follow-up state-schema change. func (m *Manager) SetModelSession(ctx context.Context, sessionID acp.SessionId, modelID string) error { m.mu.Lock() conn := m.conn + sessionID = m.resolveSessionLocked(sessionID) sess, known := m.sessions[sessionID] initialID := m.sessionID m.mu.Unlock() diff --git a/pkg/agentrun/runtime/acp/runtime_internal_test.go b/pkg/agentrun/runtime/acp/runtime_internal_test.go index 999cb7b..7c67783 100644 --- a/pkg/agentrun/runtime/acp/runtime_internal_test.go +++ b/pkg/agentrun/runtime/acp/runtime_internal_test.go @@ -149,3 +149,76 @@ func TestCancelSession_RejectsBeforeAgentStarted(t *testing.T) { t.Fatalf("expected 'agent not started' error, got: %v", err) } } + +// TestEndSession_RefusesBusySession verifies that EndSession returns a +// busy error when a session has in-flight Prompt work (refcount > 0). +// Without this check, ending a session mid-prompt leaves the prompt +// running against a session no longer in the map — state.json bookkeeping +// drifts silently. +func TestEndSession_RefusesBusySession(t *testing.T) { + m := newManagerForSessionTest(t) + m.sessionID = "initial" + m.sessions["initial"] = &sessionState{id: "initial"} + m.sessions["busy"] = &sessionState{id: "busy", cwd: "/b", inflight: 1} + + err := m.EndSession("busy") + if err == nil { + t.Fatalf("expected EndSession to refuse a busy session") + } + if !strings.Contains(err.Error(), "busy") { + t.Fatalf("expected error to mention 'busy', got: %v", err) + } + if _, ok := m.sessions["busy"]; !ok { + t.Fatalf("busy session removed despite error") + } +} + +// TestEndSession_AllowsAfterRefcountClears verifies that a session can +// be ended once its in-flight refcount drops back to zero — the busy +// check is per-state, not permanent. +func TestEndSession_AllowsAfterRefcountClears(t *testing.T) { + m := newManagerForSessionTest(t) + m.sessionID = "initial" + m.sessions["initial"] = &sessionState{id: "initial"} + m.sessions["s"] = &sessionState{id: "s", cwd: "/s", inflight: 1} + + if err := m.EndSession("s"); err == nil { + t.Fatalf("expected busy error first") + } + + m.sessions["s"].inflight = 0 // prompt completed + + if err := m.EndSession("s"); err != nil { + t.Fatalf("expected EndSession to succeed after refcount cleared: %v", err) + } + if _, ok := m.sessions["s"]; ok { + t.Fatalf("session not removed after EndSession") + } +} + +// TestPromptSession_EmptySessionIDResolvesToInitial verifies the +// single-resolution-point convention: an empty sessionID is resolved to +// m.sessionID inside PromptSession (rather than at the Service or +// Manager.Prompt-shim layer). Reaching the "session not found" branch +// would mean the empty-string convention leaked through unresolved. +func TestPromptSession_EmptySessionIDResolvesToInitial(t *testing.T) { + m := newManagerForSessionTest(t) + m.sessionID = "initial" + m.sessions["initial"] = &sessionState{id: "initial"} + + // nil conn so we get "agent not started" instead of nil-deref — + // the point is to confirm we got past the session-lookup step, + // which would have returned "session %q not found" for the literal + // empty string had resolution not happened. + _, err := m.PromptSession(context.Background(), "", + []acp.ContentBlock{acp.TextBlock("hi")}) + if err == nil { + t.Fatalf("expected error (agent not started)") + } + if strings.Contains(err.Error(), "not found") { + t.Fatalf("empty sessionID was not resolved to initial; got: %v", err) + } + if !strings.Contains(err.Error(), "agent not started") { + t.Fatalf("expected 'agent not started', got: %v", err) + } +} diff --git a/pkg/agentrun/server/service.go b/pkg/agentrun/server/service.go index 9c3b72e..2b455df 100644 --- a/pkg/agentrun/server/service.go +++ b/pkg/agentrun/server/service.go @@ -12,9 +12,10 @@ import ( ) // resolveSessionID maps the wire-level optional sessionId string to the -// typed acp.SessionId. Empty → the Manager's initial session (preserves -// pre-multi-session caller behavior — runtime layer maps the zero -// SessionId to m.sessionID via its Prompt/Cancel/SetModel shims). +// typed acp.SessionId. Empty → resolved to the Manager's initial session +// inside PromptSession / CancelSession / SetModelSession themselves. This +// keeps the empty-string-means-initial convention in exactly one layer +// (the runtime) rather than duplicating it across client / service / runtime. func resolveSessionID(s string) acp.SessionId { return acp.SessionId(s) } @@ -39,16 +40,8 @@ func (s *Service) Prompt(ctx context.Context, req *runapi.SessionPromptParams) ( s.trans.NotifyTurnStart() s.trans.NotifyUserPrompt(req.Prompt) - // Empty SessionID → use initial session (Manager.Prompt is the shim). - // Non-empty → route to specific session (PromptSession validates it - // exists in the sessions map). - var resp acp.PromptResponse - var err error - if req.SessionID == "" { - resp, err = s.mgr.Prompt(ctx, req.Prompt) - } else { - resp, err = s.mgr.PromptSession(ctx, resolveSessionID(req.SessionID), req.Prompt) - } + // Empty sessionID is resolved to the initial session inside PromptSession. + resp, err := s.mgr.PromptSession(ctx, resolveSessionID(req.SessionID), req.Prompt) stopReason := "error" if err == nil { @@ -66,11 +59,7 @@ func (s *Service) Prompt(ctx context.Context, req *runapi.SessionPromptParams) ( } func (s *Service) Cancel(ctx context.Context, req *runapi.SessionCancelParams) (retErr error) { - // req may be nil if caller sent no params (pre-multi-session clients). - sessionID := "" - if req != nil { - sessionID = req.SessionID - } + sessionID := req.SessionID s.logger.Debug("cancel", "sessionId", sessionID) defer func() { var auditArgs map[string]string @@ -80,13 +69,8 @@ func (s *Service) Cancel(ctx context.Context, req *runapi.SessionCancelParams) ( s.trans.NotifyOperationAudit("cancel", auditArgs, retErr) }() - var err error - if sessionID == "" { - err = s.mgr.Cancel(ctx) - } else { - err = s.mgr.CancelSession(ctx, resolveSessionID(sessionID)) - } - if err != nil { + // Empty sessionID is resolved to the initial session inside CancelSession. + if err := s.mgr.CancelSession(ctx, resolveSessionID(sessionID)); err != nil { retErr = jsonrpc.ErrInternal(err.Error()) return retErr } @@ -261,13 +245,8 @@ func (s *Service) SetModel(ctx context.Context, req *runapi.SessionSetModelParam retErr = jsonrpc.ErrInvalidParams("missing modelId") return nil, retErr } - var err error - if req.SessionID == "" { - err = s.mgr.SetModel(ctx, req.ModelID) - } else { - err = s.mgr.SetModelSession(ctx, resolveSessionID(req.SessionID), req.ModelID) - } - if err != nil { + // Empty sessionID is resolved to the initial session inside SetModelSession. + if err := s.mgr.SetModelSession(ctx, resolveSessionID(req.SessionID), req.ModelID); err != nil { retErr = jsonrpc.ErrInternal(err.Error()) return nil, retErr } From b4c1dc7c67215c3c42dee0bd86320f43f0c22c2d Mon Sep 17 00:00:00 2001 From: wjueyao Date: Tue, 19 May 2026 11:25:58 +0800 Subject: [PATCH 11/16] fix(runtime/acp): process-wide Phase + clear sessions on teardown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two blocking bugs surfaced by second-pass review: 1. state.Phase concurrent write race. Two PromptSession calls on different sessions race on completion: T1 finishes and writes Phase=Idle while T2 is still running. The field is single-valued but multi-session lifts that assumption, and the code blindly stamps Idle at every prompt end. state.json then lies about whether the agent is busy. Fix: Manager.activePrompts counts total in-flight PromptSession calls across all sessions. PromptSession bumps it under m.mu; the deferred decrement does the same. State writes use a new phaseFromActivePromptsLocked apply callback that reads the counter under writeState's lock — Phase reflects "any session running" rather than the last caller's local view. 2. Kill() / process-exit goroutine don't reset session bookkeeping. After the agent process dies, m.sessions retains every entry, m.conn stays non-nil, and m.sessionID still names a dead session. Sessions() returns stale IDs, EndSession silently succeeds on dead entries, and PromptSession passes its conn != nil guard before failing on the dead pipe — invariants broken across the teardown boundary. Fix: clearSessions() resets m.sessions, m.conn, m.sessionID, m.models, and m.activePrompts under m.mu. Called from Kill() and from the goroutine that observes process exit. In-flight RPCs hold their own captured conn pointer (read under m.mu at entry, used without lock); nil-ing m.conn here guards only *future* calls. Existing goroutines will fail on the closed pipe, which is the expected behavior. Tests: - TestPhaseFromActivePrompts: table-driven over the new apply callback; pins (activePrompts=0 → Idle), (1 → Running), (5 → Running). - TestClearSessions: pre-populates Manager with sessions, sessionID, models, activePrompts; verifies all fields are zeroed after the helper runs. The pre-existing TestRuntimeSuite flake (acp initialize: peer disconnected before response, noted in #2's verification section) is unaffected by these changes — internal unit tests pass deterministically with -race -count=3. Co-Authored-By: Claude Opus 4.7 (1M context) --- pkg/agentrun/runtime/acp/runtime.go | 50 ++++++++++++-- .../runtime/acp/runtime_internal_test.go | 65 +++++++++++++++++++ 2 files changed, 109 insertions(+), 6 deletions(-) diff --git a/pkg/agentrun/runtime/acp/runtime.go b/pkg/agentrun/runtime/acp/runtime.go index 52c36db..3115a63 100644 --- a/pkg/agentrun/runtime/acp/runtime.go +++ b/pkg/agentrun/runtime/acp/runtime.go @@ -94,6 +94,13 @@ type Manager struct { // inflight prompt count lives on each sessionState — see EndSession. sessions map[acp.SessionId]*sessionState + // activePrompts is the total number of in-flight PromptSession calls + // across all sessions. state.json's Phase field is process-wide, so + // it must reflect "any session running" — not "the last prompt that + // finished". PromptSession bumps this; the deferred decrement + + // writeState then re-reads it under m.mu to decide Running vs Idle. + activePrompts int + // sessionID retains the *first* session created by Create() so legacy // callers of Prompt/Cancel/SetModel (which don't pass sessionId) keep // working without code changes. New callers should use the explicit @@ -270,6 +277,7 @@ func (m *Manager) Create(ctx context.Context) error { defer close(processDone) _ = cmd.Wait() m.logger.Info("process exited") + m.clearSessions() _ = m.writeState(func(s *apiruntime.State) { s.MassVersion = m.cfg.MassVersion s.ID = m.cfg.Metadata.Name @@ -325,6 +333,7 @@ func (m *Manager) Kill(ctx context.Context) error { } } + m.clearSessions() return m.writeState(func(s *apiruntime.State) { s.MassVersion = m.cfg.MassVersion s.ID = m.cfg.Metadata.Name @@ -334,6 +343,25 @@ func (m *Manager) Kill(ctx context.Context) error { }, "runtime-stop") } +// clearSessions resets the multi-session bookkeeping after the agent +// process has exited (via Kill or natural exit). Without this, Sessions() +// returns stale IDs, EndSession silently succeeds on a dead session, and +// PromptSession passes its conn != nil check before failing on the dead +// pipe — invariants of the session map are broken across the teardown. +// +// In-flight RPCs hold their own captured conn pointer (read under m.mu at +// entry, used without lock); nil-ing m.conn here only guards *future* +// calls. Existing goroutines will fail on the closed pipe, which is fine. +func (m *Manager) clearSessions() { + m.mu.Lock() + defer m.mu.Unlock() + m.sessions = map[acp.SessionId]*sessionState{} + m.conn = nil + m.sessionID = "" + m.models = nil + m.activePrompts = 0 +} + // Delete removes the agent state directory. The agent must be stopped first. func (m *Manager) Delete() error { s, err := spec.ReadState(m.stateDir) @@ -491,6 +519,7 @@ func (m *Manager) PromptSession(ctx context.Context, sessionID acp.SessionId, pr return acp.PromptResponse{}, fmt.Errorf("runtime: prompt: session %q not found", sessionID) } sess.inflight++ + m.activePrompts++ m.mu.Unlock() defer func() { @@ -499,14 +528,13 @@ func (m *Manager) PromptSession(ctx context.Context, sessionID acp.SessionId, pr if s, ok := m.sessions[sessionID]; ok { s.inflight-- } + m.activePrompts-- m.mu.Unlock() }() m.logger.Debug("prompt started", "sessionID", sessionID, "blocks", len(prompt)) - _ = m.writeState(func(s *apiruntime.State) { - s.Phase = apiruntime.PhaseRunning - }, "prompt-started") + _ = m.writeState(m.phaseFromActivePromptsLocked, "prompt-started") resp, err := conn.Prompt(ctx, acp.PromptRequest{ SessionId: sessionID, @@ -519,9 +547,7 @@ func (m *Manager) PromptSession(ctx context.Context, sessionID acp.SessionId, pr reason = "prompt-failed" } m.logger.Debug("prompt done", "sessionID", sessionID, "reason", reason) - _ = m.writeState(func(s *apiruntime.State) { - s.Phase = apiruntime.PhaseIdle - }, reason) + _ = m.writeState(m.phaseFromActivePromptsLocked, reason) } if err != nil { @@ -530,6 +556,18 @@ func (m *Manager) PromptSession(ctx context.Context, sessionID acp.SessionId, pr return resp, nil } +// phaseFromActivePromptsLocked sets state.Phase based on the current +// activePrompts counter. Used as the writeState apply callback for prompt +// start/end so Phase reflects "any session running" rather than the last +// caller's local view. Runs under m.mu (writeState locks before calling). +func (m *Manager) phaseFromActivePromptsLocked(s *apiruntime.State) { + if m.activePrompts > 0 { + s.Phase = apiruntime.PhaseRunning + } else { + s.Phase = apiruntime.PhaseIdle + } +} + // Cancel is a backward-compat shim — equivalent to CancelSession(ctx, ""). func (m *Manager) Cancel(ctx context.Context) error { return m.CancelSession(ctx, "") diff --git a/pkg/agentrun/runtime/acp/runtime_internal_test.go b/pkg/agentrun/runtime/acp/runtime_internal_test.go index 7c67783..6fc4695 100644 --- a/pkg/agentrun/runtime/acp/runtime_internal_test.go +++ b/pkg/agentrun/runtime/acp/runtime_internal_test.go @@ -8,6 +8,8 @@ import ( "testing" acp "github.com/coder/acp-go-sdk" + + apiruntime "github.com/zoumo/mass/pkg/runtime-spec/api" ) func TestBuildSeedSystemPrompt_AppendsGuard(t *testing.T) { @@ -222,3 +224,66 @@ func TestPromptSession_EmptySessionIDResolvesToInitial(t *testing.T) { t.Fatalf("expected 'agent not started', got: %v", err) } } + +// TestPhaseFromActivePrompts verifies that state.Phase reflects the +// process-wide activePrompts counter, not a single caller's local view. +// Before this fix, two concurrent PromptSessions racing on completion +// could leave Phase=Idle while one was still running. +func TestPhaseFromActivePrompts(t *testing.T) { + m := newManagerForSessionTest(t) + + cases := []struct { + name string + active int + want apiruntime.Phase + }{ + {"zero prompts → idle", 0, apiruntime.PhaseIdle}, + {"one prompt → running", 1, apiruntime.PhaseRunning}, + {"many prompts → running", 5, apiruntime.PhaseRunning}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + m.activePrompts = tc.active + var s apiruntime.State + m.phaseFromActivePromptsLocked(&s) + if s.Phase != tc.want { + t.Fatalf("activePrompts=%d: want Phase=%q got %q", tc.active, tc.want, s.Phase) + } + }) + } +} + +// TestClearSessions verifies that the bookkeeping reset on agent +// teardown (Kill / process exit) drops every field that could leave a +// stale view behind. Without this, Sessions() would return dead IDs and +// PromptSession would pass its conn != nil guard before failing on the +// dead pipe. +func TestClearSessions(t *testing.T) { + m := newManagerForSessionTest(t) + m.sessionID = "initial" + m.sessions["initial"] = &sessionState{id: "initial"} + m.sessions["extra"] = &sessionState{id: "extra"} + m.activePrompts = 2 + m.models = &acp.SessionModelState{} + // conn is left nil — clearing a nil conn is a no-op, which is the + // branch we exercise here. The non-nil-clearing path is exercised + // indirectly via Kill()'s integration with a real process. + + m.clearSessions() + + if len(m.sessions) != 0 { + t.Errorf("sessions not cleared: %v", m.sessions) + } + if m.sessionID != "" { + t.Errorf("sessionID not cleared: %q", m.sessionID) + } + if m.activePrompts != 0 { + t.Errorf("activePrompts not cleared: %d", m.activePrompts) + } + if m.models != nil { + t.Errorf("models not cleared: %v", m.models) + } + if m.conn != nil { + t.Errorf("conn not cleared: %v", m.conn) + } +} From fc11f49c5ea24ba5b0a4c3a2f666e18ef11ddd13 Mon Sep 17 00:00:00 2001 From: wjueyao Date: Tue, 19 May 2026 12:38:06 +0800 Subject: [PATCH 12/16] fix(ari): align session method names + recovery gate on session ops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two ari-layer issues from second review pass: 1. Method-name divergence between agentrun and ari (D1). agentrun used `session/{new,end,list}` (resource/verb hierarchy), ari used `agentrun/{new,end,list}-session` (verb-on-noun). Both names are wire-visible; once shipped, third-party clients lock in whichever they hit. ari already uses `agentrun/task/{do,get,list,retry}` for task ops. Multi-session is the parallel concern, so rename to `agentrun/session/{new,end,list}` — one resource/verb hierarchy at this layer. Constants in ari/api/methods.go updated; ari/server/ service.go registration and ari/server/server.go log messages follow. agentrun layer is unchanged (its `session/*` shape was already consistent with its scope). 2. ARI new/end/list-session bypass the recovery gate (B4). Existing Prompt handler goes through reserveIdleAgent which short-circuits with CodeRecoveryBlocked while ProcessManager.IsRecovering() is true — sessions opened mid- recovery would race with the recovery sweep's view of the agent process. The new session-lifecycle handlers skipped this check. Can't use reserveIdleAgent directly because it also requires Idle status, and multi-session permits a new session while another is mid- prompt (Status=Running). Add a focused rejectIfRecovering helper on agentRunAdapter: NewSession / EndSession / ListSessions short- circuit early; existing flow otherwise unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) --- pkg/ari/api/methods.go | 9 +++++---- pkg/ari/api/types.go | 10 +++++----- pkg/ari/server/server.go | 37 +++++++++++++++++++++++++++++++------ pkg/ari/server/service.go | 11 ++++++----- 4 files changed, 47 insertions(+), 20 deletions(-) diff --git a/pkg/ari/api/methods.go b/pkg/ari/api/methods.go index 7852b77..99a0155 100644 --- a/pkg/ari/api/methods.go +++ b/pkg/ari/api/methods.go @@ -32,11 +32,12 @@ const ( // Multi-session lifecycle for a single agentrun. session/new opens // an additional ACP session (different cwd, fresh state) without // fork+exec a new agent process; session/end releases runtime - // tracking; session/list enumerates active sessions. See + // tracking; session/list enumerates active sessions. Names parallel + // agentrun/task/* — one resource/verb hierarchy at this layer. See // pkg/agentrun/runtime/acp Manager.NewSession for runtime contract. - MethodAgentRunNewSession = "agentrun/new-session" - MethodAgentRunEndSession = "agentrun/end-session" - MethodAgentRunListSessions = "agentrun/list-sessions" + MethodAgentRunNewSession = "agentrun/session/new" + MethodAgentRunEndSession = "agentrun/session/end" + MethodAgentRunListSessions = "agentrun/session/list" ) // ARI agent definition methods. diff --git a/pkg/ari/api/types.go b/pkg/ari/api/types.go index 2345459..2129fd5 100644 --- a/pkg/ari/api/types.go +++ b/pkg/ari/api/types.go @@ -89,7 +89,7 @@ type AgentRunPromptParams struct { // Name is the agent run name (required). Name string `json:"name"` - // SessionID addresses a specific session opened via agentrun/new-session. + // SessionID addresses a specific session opened via agentrun/session/new. // Empty (omitted) routes to the agentrun's initial session — preserves // pre-multi-session caller behavior. SessionID string `json:"sessionId,omitempty"` @@ -104,7 +104,7 @@ type AgentRunPromptResult struct { Accepted bool `json:"accepted"` } -// AgentRunNewSessionParams is the request params for agentrun/new-session. +// AgentRunNewSessionParams is the request params for agentrun/session/new. // Opens an additional ACP session on an existing agentrun (no fork+exec). type AgentRunNewSessionParams struct { Workspace string `json:"workspace"` @@ -116,21 +116,21 @@ type AgentRunNewSessionParams struct { McpServers []runapi.SessionNewMcpServer `json:"mcpServers,omitempty"` } -// AgentRunNewSessionResult is the response for agentrun/new-session. +// AgentRunNewSessionResult is the response for agentrun/session/new. type AgentRunNewSessionResult struct { // SessionID is the ACP session id the caller passes to subsequent // agentrun/prompt (etc.) via the SessionID field. SessionID string `json:"sessionId"` } -// AgentRunEndSessionParams is the request params for agentrun/end-session. +// AgentRunEndSessionParams is the request params for agentrun/session/end. type AgentRunEndSessionParams struct { Workspace string `json:"workspace"` Name string `json:"name"` SessionID string `json:"sessionId"` } -// AgentRunEndSessionResult is the response for agentrun/end-session. +// AgentRunEndSessionResult is the response for agentrun/session/end. type AgentRunEndSessionResult struct{} // AgentRunListSessionsParams identifies the agentrun whose sessions to list. diff --git a/pkg/ari/server/server.go b/pkg/ari/server/server.go index bb7de93..6de127e 100644 --- a/pkg/ari/server/server.go +++ b/pkg/ari/server/server.go @@ -530,14 +530,22 @@ func (a *agentRunAdapter) Prompt(ctx context.Context, req *pkgariapi.AgentRunPro return &pkgariapi.AgentRunPromptResult{Accepted: true}, nil } -// NewSession forwards agentrun/new-session to the running agent-run's +// NewSession forwards agentrun/session/new to the running agent-run's // session/new RPC. The agent reuses its process to host the new session // — no fork+exec. Returns the agent-issued sessionId. +// +// Rejects requests while the daemon is recovering agents — sessions opened +// then would race with recovery's view of the agent process. Doesn't use +// reserveIdleAgent because multi-session permits opening a session while +// another is mid-prompt (Status=Running); only recovery is unsafe. func (a *agentRunAdapter) NewSession(ctx context.Context, req *pkgariapi.AgentRunNewSessionParams) (*pkgariapi.AgentRunNewSessionResult, error) { if req.Workspace == "" || req.Name == "" || req.Cwd == "" { return nil, jsonrpc.ErrInvalidParams("workspace, name, and cwd are required") } - a.logger.Info("agentrun/new-session", "workspace", req.Workspace, "name", req.Name, "cwd", req.Cwd) + if err := a.rejectIfRecovering(); err != nil { + return nil, err + } + a.logger.Info("agentrun/session/new", "workspace", req.Workspace, "name", req.Name, "cwd", req.Cwd) client, err := a.processes.Connect(ctx, req.Workspace, req.Name) if err != nil { @@ -551,17 +559,20 @@ func (a *agentRunAdapter) NewSession(ctx context.Context, req *pkgariapi.AgentRu if err != nil { return nil, jsonrpc.ErrInternal(err.Error()) } - a.logger.Info("agentrun/new-session: opened", + a.logger.Info("agentrun/session/new: opened", "workspace", req.Workspace, "name", req.Name, "sessionId", out.SessionID) return &pkgariapi.AgentRunNewSessionResult{SessionID: out.SessionID}, nil } -// EndSession forwards agentrun/end-session to release runtime tracking. +// EndSession forwards agentrun/session/end to release runtime tracking. func (a *agentRunAdapter) EndSession(ctx context.Context, req *pkgariapi.AgentRunEndSessionParams) (*pkgariapi.AgentRunEndSessionResult, error) { if req.Workspace == "" || req.Name == "" || req.SessionID == "" { return nil, jsonrpc.ErrInvalidParams("workspace, name, and sessionId are required") } - a.logger.Info("agentrun/end-session", "workspace", req.Workspace, "name", req.Name, "sessionId", req.SessionID) + if err := a.rejectIfRecovering(); err != nil { + return nil, err + } + a.logger.Info("agentrun/session/end", "workspace", req.Workspace, "name", req.Name, "sessionId", req.SessionID) client, err := a.processes.Connect(ctx, req.Workspace, req.Name) if err != nil { @@ -573,11 +584,14 @@ func (a *agentRunAdapter) EndSession(ctx context.Context, req *pkgariapi.AgentRu return &pkgariapi.AgentRunEndSessionResult{}, nil } -// ListSessions forwards agentrun/list-sessions to enumerate active sessions. +// ListSessions forwards agentrun/session/list to enumerate active sessions. func (a *agentRunAdapter) ListSessions(ctx context.Context, req *pkgariapi.AgentRunListSessionsParams) (*pkgariapi.AgentRunListSessionsResult, error) { if req.Workspace == "" || req.Name == "" { return nil, jsonrpc.ErrInvalidParams("workspace and name are required") } + if err := a.rejectIfRecovering(); err != nil { + return nil, err + } client, err := a.processes.Connect(ctx, req.Workspace, req.Name) if err != nil { @@ -590,6 +604,17 @@ func (a *agentRunAdapter) ListSessions(ctx context.Context, req *pkgariapi.Agent return &pkgariapi.AgentRunListSessionsResult{SessionIDs: out.SessionIDs}, nil } +// rejectIfRecovering returns CodeRecoveryBlocked when the daemon is mid- +// recovery. Used by session-lifecycle handlers that can't go through +// reserveIdleAgent (multi-session allows operations while another session +// is Running) but still must serialize against the recovery sweep. +func (a *agentRunAdapter) rejectIfRecovering() *jsonrpc.RPCError { + if a.processes.IsRecovering() { + return &jsonrpc.RPCError{Code: pkgariapi.CodeRecoveryBlocked, Message: "daemon is recovering agents"} + } + return nil +} + // Cancel handles agentrun/cancel. // // Connects to the running agent-run and calls Cancel. diff --git a/pkg/ari/server/service.go b/pkg/ari/server/service.go index b76c570..c6c0f98 100644 --- a/pkg/ari/server/service.go +++ b/pkg/ari/server/service.go @@ -162,11 +162,12 @@ func RegisterAgentRunService(s *jsonrpc.Server, svc AgentRunService) { "task/get": jsonrpc.UnaryMethod(svc.TaskGet), "task/list": jsonrpc.UnaryMethod(svc.TaskList), "task/retry": jsonrpc.UnaryMethod(svc.TaskRetry), - // Multi-session lifecycle (see pkg/agentrun/runtime/acp - // Manager.NewSession for runtime contract). - "new-session": jsonrpc.UnaryMethod(svc.NewSession), - "end-session": jsonrpc.UnaryMethod(svc.EndSession), - "list-sessions": jsonrpc.UnaryMethod(svc.ListSessions), + // Multi-session lifecycle. Names parallel agentrun/task/* — + // agentrun/session/new / end / list — so a third-party reader + // sees one consistent resource/verb hierarchy at this layer. + "session/new": jsonrpc.UnaryMethod(svc.NewSession), + "session/end": jsonrpc.UnaryMethod(svc.EndSession), + "session/list": jsonrpc.UnaryMethod(svc.ListSessions), }, }) } From 9e378fb1d56bb45953b120a8bb90c34d3dac533a Mon Sep 17 00:00:00 2001 From: wjueyao Date: Tue, 19 May 2026 12:38:25 +0800 Subject: [PATCH 13/16] refactor(runtime/acp): cwd-required, drop stale doc TODO, rename Sessions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Several smaller items surfaced in second review: D5. Dead doc TODO. runtime.go's sessionState block referenced docs/design/multi-session-per-agentrun.md (TODO: add design doc). The doc was never written and the PR description / commit messages carry the rationale now. Drop the dangling reference. D2. Cwd contract divergence. Manager.NewSession allowed empty cwd with a silent fallback to ResolveAgentRoot(bundleDir, cfg), but every wire-level layer (CLI / ARI / agentrun service) already rejects empty cwd before reaching runtime. The fallback was dead code from any real caller's perspective; an in-process consumer relying on it would have fragile semantics (cwd resolves differently per bundle). Reject empty cwd at the runtime boundary — one contract, four enforcement points instead of three plus a quiet exception. B5. SetModelSession stale-pointer write. The pre-fix code captured sess at the first lock, released the lock, ran the RPC, then re- acquired the lock and mutated sess.models — but the session could have been ended (or the whole map cleared by Kill) in the window. The write succeeded against an orphan struct, silently no-op'ing the in-memory model update. Re-lookup `m.sessions[sessionID]` under the second lock and skip the mutation when the session is gone. D6. Sessions() naming. The method returns []acp.SessionId, not []*sessionState — name read like it returned the full session objects. Rename to SessionIDs(). Two callers (service.go + internal test) updated. Test functions renamed in parallel. Co-Authored-By: Claude Opus 4.7 (1M context) --- pkg/agentrun/runtime/acp/runtime.go | 45 +++++++++---------- .../runtime/acp/runtime_internal_test.go | 8 ++-- 2 files changed, 25 insertions(+), 28 deletions(-) diff --git a/pkg/agentrun/runtime/acp/runtime.go b/pkg/agentrun/runtime/acp/runtime.go index 3115a63..8cc60c1 100644 --- a/pkg/agentrun/runtime/acp/runtime.go +++ b/pkg/agentrun/runtime/acp/runtime.go @@ -63,10 +63,8 @@ type StateChangeHook func(StateChange) // agent process (e.g. claude-agent-acp) keeps sessions in a map keyed by // sessionId; this struct mirrors the runtime-side view of each. // -// Multi-session lets callers (mass daemon / mindpowers) keep a long-lived -// agent process and switch session per task, instead of fork+kill an -// agent process per task. See docs/design/multi-session-per-agentrun.md -// (TODO: add design doc) for the broader rationale. +// Multi-session lets callers keep a long-lived agent process and switch +// session per task, instead of fork+kill an agent process per task. type sessionState struct { id acp.SessionId models *acp.SessionModelState @@ -384,13 +382,18 @@ func (m *Manager) GetState() (apiruntime.State, error) { // session). Returns the new session's ID, which callers must pass to // PromptSession / CancelSession / EndSession to address this session. // -// cwd overrides the agent's working directory for this session — useful for -// running multiple isolated tasks on one long-lived agent (e.g. self-EDD's -// per-case fixture directories). Empty cwd uses the agent's workDir. +// cwd is required — each session is scoped to its own working directory. +// All wire-level callers (Service / ARI / CLI) already reject empty cwd; +// the runtime validates here too for in-process callers and so the +// contract is consistent across layers. // // mcpServers (optional) are extra MCP servers scoped to this session, layered // on top of the agent's bundle-level mcpServers. func (m *Manager) NewSession(ctx context.Context, cwd string, mcpServers []acp.McpServer) (acp.SessionId, error) { + if cwd == "" { + return "", fmt.Errorf("runtime: new session: cwd is required") + } + m.mu.Lock() conn := m.conn m.mu.Unlock() @@ -399,15 +402,6 @@ func (m *Manager) NewSession(ctx context.Context, cwd string, mcpServers []acp.M return "", fmt.Errorf("runtime: agent not started") } - resolvedCwd := cwd - if resolvedCwd == "" { - workDir, err := spec.ResolveAgentRoot(m.bundleDir, m.cfg) - if err != nil { - return "", fmt.Errorf("runtime: resolve cwd for new session: %w", err) - } - resolvedCwd = workDir - } - // Layer session-scoped mcpServers on top of bundle-level config. Callers // that pass nil/empty get the bundle defaults. merged := convertMcpServers(m.cfg.Session.McpServers) @@ -415,7 +409,7 @@ func (m *Manager) NewSession(ctx context.Context, cwd string, mcpServers []acp.M req := acp.NewSessionRequest{ Meta: m.cfg.Session.Meta, - Cwd: resolvedCwd, + Cwd: cwd, McpServers: merged, } resp, err := conn.NewSession(ctx, req) @@ -427,10 +421,10 @@ func (m *Manager) NewSession(ctx context.Context, cwd string, mcpServers []acp.M m.sessions[resp.SessionId] = &sessionState{ id: resp.SessionId, models: resp.Models, - cwd: resolvedCwd, + cwd: cwd, } m.mu.Unlock() - m.logger.Info("session opened", "sessionID", resp.SessionId, "cwd", resolvedCwd) + m.logger.Info("session opened", "sessionID", resp.SessionId, "cwd", cwd) return resp.SessionId, nil } @@ -474,9 +468,9 @@ func (m *Manager) resolveSessionLocked(sessionID acp.SessionId) acp.SessionId { return sessionID } -// Sessions returns a snapshot of currently-active session IDs. Order is +// SessionIDs returns a snapshot of currently-active session IDs. Order is // not stable — caller should sort if a deterministic order is needed. -func (m *Manager) Sessions() []acp.SessionId { +func (m *Manager) SessionIDs() []acp.SessionId { m.mu.Lock() defer m.mu.Unlock() ids := make([]acp.SessionId, 0, len(m.sessions)) @@ -610,7 +604,7 @@ func (m *Manager) SetModelSession(ctx context.Context, sessionID acp.SessionId, m.mu.Lock() conn := m.conn sessionID = m.resolveSessionLocked(sessionID) - sess, known := m.sessions[sessionID] + _, known := m.sessions[sessionID] initialID := m.sessionID m.mu.Unlock() @@ -630,9 +624,12 @@ func (m *Manager) SetModelSession(ctx context.Context, sessionID acp.SessionId, return fmt.Errorf("runtime: set model: %w", err) } - // Update per-session in-memory models state. + // Re-lookup under the second lock — the session could have been ended + // (or the whole map cleared by teardown) between the first unlock and + // here, in which case the in-memory model mutation should silently no-op + // rather than write through an orphan struct. m.mu.Lock() - if sess.models != nil { + if sess, ok := m.sessions[sessionID]; ok && sess.models != nil { sess.models.CurrentModelId = acp.ModelId(modelID) //nolint:gosec // ModelId is string } // Legacy single-session models mirror for the initial session only. diff --git a/pkg/agentrun/runtime/acp/runtime_internal_test.go b/pkg/agentrun/runtime/acp/runtime_internal_test.go index 6fc4695..b309977 100644 --- a/pkg/agentrun/runtime/acp/runtime_internal_test.go +++ b/pkg/agentrun/runtime/acp/runtime_internal_test.go @@ -44,20 +44,20 @@ func newManagerForSessionTest(t *testing.T) *Manager { } } -func TestSessions_EmptyWhenNoneRegistered(t *testing.T) { +func TestSessionIDs_EmptyWhenNoneRegistered(t *testing.T) { m := newManagerForSessionTest(t) - if got := m.Sessions(); len(got) != 0 { + if got := m.SessionIDs(); len(got) != 0 { t.Fatalf("expected empty sessions, got %v", got) } } -func TestSessions_SnapshotIncludesRegistered(t *testing.T) { +func TestSessionIDs_SnapshotIncludesRegistered(t *testing.T) { m := newManagerForSessionTest(t) m.sessions["sess-a"] = &sessionState{id: "sess-a", cwd: "/a"} m.sessions["sess-b"] = &sessionState{id: "sess-b", cwd: "/b"} m.sessionID = "sess-a" - got := m.Sessions() + got := m.SessionIDs() if len(got) != 2 { t.Fatalf("expected 2 sessions, got %v", got) } From 265bcac7db769d898e4532aeeb0581db61aef77b Mon Sep 17 00:00:00 2001 From: wjueyao Date: Tue, 19 May 2026 12:38:46 +0800 Subject: [PATCH 14/16] fix: per-session sessionId stamping for watch filtering (B3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Before this fix, the Translator stamped every emitted AgentRunEvent with the initial-session id (cached in t.sessionID). Multi-session broke this assumption silently: agent_message events from a non- initial session were labelled with the initial id, and `massctl agentrun prompt --wait --session-id X` had no way to distinguish session X's stream from a concurrent session Y's — events crossed, and a turn_end from Y exited X's collector early. Translator (pkg/agentrun/server/translator.go): - NotifyTurnStart / NotifyUserPrompt / NotifyTurnEnd gain a leading sessionID parameter. Empty resolves to the initial session via a small resolveSessionLocked helper — back-compat for any caller outside the multi-session paths. - broadcastEvent's per-event flow split into a generic broadcastEventForSession(sid, ev) and a one-line broadcastEvent wrapper that stamps the initial session (still used by error / operationAudit / stateChange / metadata events, which are process-wide). - The run() goroutine consumes acp.SessionNotification (which carries SessionId per ACP spec); content events now stamp n.SessionId rather than t.sessionID. Service (pkg/agentrun/server/service.go): Prompt forwards req.SessionID into the three Notify calls. cmd/mass/commands/run: the bootstrap seed-prompt explicitly targets the initial session — pass "" to let the Translator resolve. CLI (cmd/massctl/commands/agentrun/prompt.go): `--wait` mode filters incoming events by --session-id when non-empty. Default --session-id="" matches all events (single-session legacy behavior preserved). Tests: - TestService_NewSession_MissingCwdValidation — symmetric to runtime's cwd-required guard; pins the wire-layer rejection. - TestService_EndSession_MissingSessionIDValidation — same. - TestNotifyTurnStart_StampsSessionID / EmptySessionIDResolvesToInitial — Translator pins both the explicit-stamp and fallback paths. - TestRun_StampsSessionIDFromNotification — content-event path carries n.SessionId; the regression most likely to bite under real multi-session traffic. Existing translator_test.go callsites bulk-updated for the new Notify* signatures. Co-Authored-By: Claude Opus 4.7 (1M context) --- cmd/mass/commands/run/command.go | 8 ++- cmd/massctl/commands/agentrun/prompt.go | 10 +++ pkg/agentrun/server/service.go | 11 +-- pkg/agentrun/server/service_test.go | 21 ++++++ pkg/agentrun/server/translator.go | 57 ++++++++++++--- pkg/agentrun/server/translator_test.go | 92 ++++++++++++++++++++----- 6 files changed, 164 insertions(+), 35 deletions(-) diff --git a/cmd/mass/commands/run/command.go b/cmd/mass/commands/run/command.go index 6b7ea40..11845f3 100644 --- a/cmd/mass/commands/run/command.go +++ b/cmd/mass/commands/run/command.go @@ -141,8 +141,10 @@ func run(cmd *cobra.Command, bundle, stateDir, permissions, id string, logCfg *l } if cfg.Session.SystemPrompt != "" { - trans.NotifyTurnStart() - trans.NotifyUserPrompt([]runapi.ContentBlock{runapi.TextBlock(acpruntime.BuildSeedSystemPrompt(cfg.Session.SystemPrompt))}) + // Seed prompt targets the initial session — empty sessionID resolves + // to it inside the Translator. + trans.NotifyTurnStart("") + trans.NotifyUserPrompt("", []runapi.ContentBlock{runapi.TextBlock(acpruntime.BuildSeedSystemPrompt(cfg.Session.SystemPrompt))}) resp, err := mgr.SeedSystemPrompt(ctx) stopReason := "error" if err == nil { @@ -151,7 +153,7 @@ func run(cmd *cobra.Command, bundle, stateDir, permissions, id string, logCfg *l if err != nil { trans.NotifyError(err.Error()) } - trans.NotifyTurnEnd(acp.StopReason(stopReason)) + trans.NotifyTurnEnd("", acp.StopReason(stopReason)) if err != nil { return fmt.Errorf("agent-run: seed system prompt: %w", err) } diff --git a/cmd/massctl/commands/agentrun/prompt.go b/cmd/massctl/commands/agentrun/prompt.go index f1e3c3c..7b487f0 100644 --- a/cmd/massctl/commands/agentrun/prompt.go +++ b/cmd/massctl/commands/agentrun/prompt.go @@ -85,6 +85,13 @@ func newPromptCmd(getClient cliutil.ClientFn) *cobra.Command { } // Collect agent_message text until turn_end. + // + // When --session-id is set, filter events to that session — without + // this, two concurrent sessions on the same agent will cross-talk + // (agent_message from session B counted as session A's, turn_end + // from B exits early). When sessionID is empty (single-session + // legacy default), every event is accepted; Translator stamps + // initial-session events with the initial session id. var parts []string timeout := time.After(5 * time.Minute) for { @@ -93,6 +100,9 @@ func newPromptCmd(getClient cliutil.ClientFn) *cobra.Command { if !ok { return fmt.Errorf("event stream closed before turn_end") } + if sessionID != "" && ev.SessionID != sessionID { + continue + } if ev.Type == runapi.EventTypeTurnEnd { fmt.Fprintln(cmd.OutOrStdout(), strings.Join(parts, "")) return nil diff --git a/pkg/agentrun/server/service.go b/pkg/agentrun/server/service.go index 2b455df..1be6752 100644 --- a/pkg/agentrun/server/service.go +++ b/pkg/agentrun/server/service.go @@ -37,10 +37,11 @@ func (s *Service) Prompt(ctx context.Context, req *runapi.SessionPromptParams) ( return nil, jsonrpc.ErrInvalidParams("missing prompt") } s.logger.Debug("prompt", "sessionId", req.SessionID, "blocks", len(req.Prompt)) - s.trans.NotifyTurnStart() - s.trans.NotifyUserPrompt(req.Prompt) + // Empty sessionID is resolved to the initial session inside PromptSession + // (and inside Translator's Notify* methods). + s.trans.NotifyTurnStart(req.SessionID) + s.trans.NotifyUserPrompt(req.SessionID, req.Prompt) - // Empty sessionID is resolved to the initial session inside PromptSession. resp, err := s.mgr.PromptSession(ctx, resolveSessionID(req.SessionID), req.Prompt) stopReason := "error" @@ -50,7 +51,7 @@ func (s *Service) Prompt(ctx context.Context, req *runapi.SessionPromptParams) ( if err != nil { s.trans.NotifyError(err.Error()) } - s.trans.NotifyTurnEnd(acp.StopReason(stopReason)) + s.trans.NotifyTurnEnd(req.SessionID, acp.StopReason(stopReason)) s.logger.Debug("prompt done", "sessionId", req.SessionID, "stopReason", stopReason) if err != nil { return nil, jsonrpc.ErrInternal(err.Error()) @@ -316,7 +317,7 @@ func (s *Service) EndSession(_ context.Context, req *runapi.SessionEndParams) (_ // ListSessions returns the active session IDs snapshot from the Manager. func (s *Service) ListSessions(_ context.Context) (*runapi.SessionListResult, error) { - ids := s.mgr.Sessions() + ids := s.mgr.SessionIDs() out := make([]string, len(ids)) for i, id := range ids { out[i] = string(id) diff --git a/pkg/agentrun/server/service_test.go b/pkg/agentrun/server/service_test.go index aa6eed9..3c2505c 100644 --- a/pkg/agentrun/server/service_test.go +++ b/pkg/agentrun/server/service_test.go @@ -146,3 +146,24 @@ func TestService_SetModel_AuditOnValidationFailure(t *testing.T) { assert.False(t, ru.OperationAudit.Success) assert.NotEmpty(t, ru.OperationAudit.Error) } + +// TestService_NewSession_MissingCwdValidation verifies that the service layer +// rejects a missing cwd before reaching the runtime layer. The contract is +// cwd-required across all wire layers (CLI / ARI / agentrun); see runtime/acp +// NewSession's matching guard for the symmetric runtime-side check. +func TestService_NewSession_MissingCwdValidation(t *testing.T) { + svc := newTestService(t) + _, err := svc.NewSession(context.Background(), &runapi.SessionNewParams{}) + require.Error(t, err) + assert.Contains(t, err.Error(), "missing cwd") +} + +// TestService_EndSession_MissingSessionIDValidation pins the symmetric +// validation for end-session — the runtime can't disambiguate "no sessionId +// supplied" from "empty resolves to initial", so the wire layer rejects. +func TestService_EndSession_MissingSessionIDValidation(t *testing.T) { + svc := newTestService(t) + _, err := svc.EndSession(context.Background(), &runapi.SessionEndParams{}) + require.Error(t, err) + assert.Contains(t, err.Error(), "missing sessionId") +} diff --git a/pkg/agentrun/server/translator.go b/pkg/agentrun/server/translator.go index 27ccf1d..cf2d5d8 100644 --- a/pkg/agentrun/server/translator.go +++ b/pkg/agentrun/server/translator.go @@ -208,14 +208,19 @@ func (t *Translator) LastSeq() int { // NotifyTurnStart broadcasts a turn_start AgentRunEvent. // The new turnId is assigned atomically inside the broadcast callback, // which runs under mu.Lock. -func (t *Translator) NotifyTurnStart() { +// +// sessionID identifies which ACP session this turn belongs to. Empty +// resolves to the agentrun's initial session (back-compat for single- +// session callers). Watchers filter events by sessionId; mis-stamping +// would cause cross-session leakage on `--wait --session-id`. +func (t *Translator) NotifyTurnStart(sessionID string) { newTurnID := uuid.New().String() - t.logger.Debug("turn_start", "turnID", newTurnID) + t.logger.Debug("turn_start", "turnID", newTurnID, "sessionID", sessionID) t.broadcast(func(seq int, at time.Time) runapi.AgentRunEvent { t.currentTurnId = newTurnID return runapi.AgentRunEvent{ RunID: t.runID, - SessionID: t.sessionID, + SessionID: t.resolveSessionLocked(sessionID), Seq: seq, Time: at, Type: runapi.EventTypeTurnStart, @@ -230,13 +235,16 @@ func (t *Translator) NotifyTurnStart() { // Each ContentBlock becomes a separate event, matching the per-block pattern // used on the response side (agent_message). // This must be called after NotifyTurnStart and before mgr.Prompt. -func (t *Translator) NotifyUserPrompt(blocks []runapi.ContentBlock) { +// +// sessionID identifies which session the prompt is destined for. Empty +// resolves to the agentrun's initial session. +func (t *Translator) NotifyUserPrompt(sessionID string, blocks []runapi.ContentBlock) { for _, block := range blocks { b := block // capture for closure t.broadcast(func(seq int, at time.Time) runapi.AgentRunEvent { return runapi.AgentRunEvent{ RunID: t.runID, - SessionID: t.sessionID, + SessionID: t.resolveSessionLocked(sessionID), Seq: seq, Time: at, Type: runapi.EventTypeUserMessage, @@ -274,13 +282,16 @@ func (t *Translator) NotifyOperationAudit(op string, params map[string]string, e // Closes any open content block first, then emits turn_end. // The current turnId is included in the event and cleared AFTER use so the // turn_end event itself carries the identifier. -func (t *Translator) NotifyTurnEnd(reason acp.StopReason) { - t.logger.Debug("turn_end", "turnID", t.currentTurnId, "reason", reason) +// +// sessionID identifies which session the turn ended on. Empty resolves to +// the agentrun's initial session. +func (t *Translator) NotifyTurnEnd(sessionID string, reason acp.StopReason) { + t.logger.Debug("turn_end", "turnID", t.currentTurnId, "sessionID", sessionID, "reason", reason) t.closeOpenBlock() t.broadcast(func(seq int, at time.Time) runapi.AgentRunEvent { ae := runapi.AgentRunEvent{ RunID: t.runID, - SessionID: t.sessionID, + SessionID: t.resolveSessionLocked(sessionID), Seq: seq, Time: at, Type: runapi.EventTypeTurnEnd, @@ -354,7 +365,10 @@ func (t *Translator) run() { t.closeOpenBlock() } - t.broadcastEvent(ev) + // ACP SessionNotification carries the session id; stamp it on + // the resulting event so watchers can filter by session. Empty + // (unusual — pre-handshake) falls back to the initial session. + t.broadcastEventForSession(string(n.SessionId), ev) t.maybeNotifyMetadata(ev) } } @@ -392,14 +406,25 @@ func (t *Translator) closeOpenBlock() { } } -// broadcastEvent builds and broadcasts an AgentRunEvent. +// broadcastEvent builds and broadcasts an AgentRunEvent stamped with the +// initial session id. Use for process-wide events (runtime_update / error) +// where session attribution doesn't apply. Per-session events should use +// broadcastEventForSession. +// // TurnID is applied to all events except runtime_update when an active turn exists. func (t *Translator) broadcastEvent(ev runapi.Event) { + t.broadcastEventForSession("", ev) +} + +// broadcastEventForSession is broadcastEvent with explicit session attribution. +// Empty sessionID resolves to the initial session — callers passing a non- +// empty id stamp it directly, allowing watchers to filter by sessionId. +func (t *Translator) broadcastEventForSession(sessionID string, ev runapi.Event) { t.broadcast(func(seq int, at time.Time) runapi.AgentRunEvent { eventType := runapi.EventTypeOf(ev) ae := runapi.AgentRunEvent{ RunID: t.runID, - SessionID: t.sessionID, + SessionID: t.resolveSessionLocked(sessionID), Seq: seq, Time: at, Type: eventType, @@ -412,6 +437,16 @@ func (t *Translator) broadcastEvent(ev runapi.Event) { }) } +// resolveSessionLocked maps empty sessionID to t.sessionID (the initial +// session). Must be called from within a broadcast callback (which already +// holds t.mu) or from a method holding t.mu. +func (t *Translator) resolveSessionLocked(sessionID string) string { + if sessionID == "" { + return t.sessionID + } + return sessionID +} + // broadcast is the single fan-out entry point. The build callback runs under // mu.Lock and receives the assigned seq and current timestamp. The lock is held // for the entire log-then-fanout sequence to guarantee: diff --git a/pkg/agentrun/server/translator_test.go b/pkg/agentrun/server/translator_test.go index 76b456a..3f31740 100644 --- a/pkg/agentrun/server/translator_test.go +++ b/pkg/agentrun/server/translator_test.go @@ -258,8 +258,8 @@ func TestNotifyTurnStartAndEnd(t *testing.T) { tr.Start() defer tr.Stop() - tr.NotifyTurnStart() - tr.NotifyTurnEnd(acp.StopReason("end_turn")) + tr.NotifyTurnStart("") + tr.NotifyTurnEnd("", acp.StopReason("end_turn")) first := drainEvent(t, ch) second := drainEvent(t, ch) @@ -400,13 +400,13 @@ func TestTurnAwareEvent_TurnIdAssigned(t *testing.T) { tr.Start() defer tr.Stop() - tr.NotifyTurnStart() + tr.NotifyTurnStart("") tsEv := drainEvent(t, ch) txt1Ev := sendAndDrainEvent(t, in, ch, "hello") txt2Ev := sendAndDrainEvent(t, in, ch, "world") - tr.NotifyTurnEnd(acp.StopReason("end_turn")) + tr.NotifyTurnEnd("", acp.StopReason("end_turn")) blockEndEv := drainEvent(t, ch) // synthetic agent_message{end} teEv := drainEvent(t, ch) // turn_end @@ -433,15 +433,15 @@ func TestTurnAwareEvent_TurnIDChangesPerTurn(t *testing.T) { defer tr.Stop() // Turn 1. - tr.NotifyTurnStart() + tr.NotifyTurnStart("") ts1 := drainEvent(t, ch) sendAndDrainEvent(t, in, ch, "turn1") - tr.NotifyTurnEnd(acp.StopReason("end_turn")) + tr.NotifyTurnEnd("", acp.StopReason("end_turn")) drainEvent(t, ch) // turn_end drainEvent(t, ch) // synthetic content end // Turn 2 — TurnID must differ. - tr.NotifyTurnStart() + tr.NotifyTurnStart("") ts2 := drainEvent(t, ch) assert.NotEqual(t, ts1.TurnID, ts2.TurnID, "turn 2 must have a different TurnID") } @@ -455,7 +455,7 @@ func TestTurnAwareEvent_StateChangeExcludesTurnFields(t *testing.T) { tr.Start() defer tr.Stop() - tr.NotifyTurnStart() + tr.NotifyTurnStart("") tsEv := drainEvent(t, ch) require.NotEmpty(t, tsEv.TurnID) @@ -477,7 +477,7 @@ func TestTurnAwareEvent_MetadataEventInTurn(t *testing.T) { tr.Start() defer tr.Stop() - tr.NotifyTurnStart() + tr.NotifyTurnStart("") tsEv := drainEvent(t, ch) require.NotEmpty(t, tsEv.TurnID) @@ -634,20 +634,20 @@ func TestTurnAwareEvent_ReplayOrdering(t *testing.T) { defer tr.Stop() // Turn 1: turn_start + 2 text events + synthetic block end + turn_end. - tr.NotifyTurnStart() + tr.NotifyTurnStart("") ts1Ev := drainEvent(t, ch) t1aEv := sendAndDrainEvent(t, in, ch, "t1-a") t1bEv := sendAndDrainEvent(t, in, ch, "t1-b") - tr.NotifyTurnEnd(acp.StopReason("end_turn")) + tr.NotifyTurnEnd("", acp.StopReason("end_turn")) t1EndEv := drainEvent(t, ch) // synthetic agent_message{end} te1Ev := drainEvent(t, ch) // turn_end turn1 := []runapi.AgentRunEvent{ts1Ev, t1aEv, t1bEv, t1EndEv, te1Ev} // Turn 2: turn_start + 1 text event + synthetic block end + turn_end. - tr.NotifyTurnStart() + tr.NotifyTurnStart("") ts2Ev := drainEvent(t, ch) t2aEv := sendAndDrainEvent(t, in, ch, "t2-a") - tr.NotifyTurnEnd(acp.StopReason("end_turn")) + tr.NotifyTurnEnd("", acp.StopReason("end_turn")) t2EndEv := drainEvent(t, ch) // synthetic agent_message{end} te2Ev := drainEvent(t, ch) // turn_end turn2 := []runapi.AgentRunEvent{ts2Ev, t2aEv, t2EndEv, te2Ev} @@ -688,11 +688,11 @@ func TestEventCounts_PromptTurn(t *testing.T) { defer tr.Stop() // turn_start - tr.NotifyTurnStart() + tr.NotifyTurnStart("") drainEvent(t, ch) // user_message - tr.NotifyUserPrompt([]runapi.ContentBlock{runapi.TextBlock("hello")}) + tr.NotifyUserPrompt("", []runapi.ContentBlock{runapi.TextBlock("hello")}) drainEvent(t, ch) // 2 text events (AgentMessageChunk) @@ -707,7 +707,7 @@ func TestEventCounts_PromptTurn(t *testing.T) { drainEvent(t, ch) // tool_call // turn_end - tr.NotifyTurnEnd(acp.StopReason("end_turn")) + tr.NotifyTurnEnd("", acp.StopReason("end_turn")) drainEvent(t, ch) // state_change @@ -899,3 +899,63 @@ func TestSessionMetadataHook_AllFourTypes(t *testing.T) { defer mu.Unlock() assert.Equal(t, []string{"runtime_update", "runtime_update", "runtime_update", "runtime_update"}, types) } + +// TestNotifyTurnStart_StampsSessionID verifies that NotifyTurnStart writes +// the caller-supplied sessionID on the emitted event. This is what lets +// CLI clients filter `--wait --session-id X` watches without cross-session +// leakage between concurrent prompts on different sessions. +func TestNotifyTurnStart_StampsSessionID(t *testing.T) { + in := make(chan acp.SessionNotification, 1) + tr := NewTranslator("run-1", in, "", slog.Default()) + tr.SetSessionID("initial-sid") + ch, _, _ := tr.Subscribe() + + tr.NotifyTurnStart("session-abc") + + ev := drainEvent(t, ch) + assert.Equal(t, "session-abc", ev.SessionID, + "caller-supplied sessionID must be stamped on the event") +} + +// TestNotifyTurnStart_EmptySessionIDResolvesToInitial verifies that an +// empty sessionID at the Translator API falls back to the initial session +// id — back-compat for legacy single-session callers that don't pass one. +func TestNotifyTurnStart_EmptySessionIDResolvesToInitial(t *testing.T) { + in := make(chan acp.SessionNotification, 1) + tr := NewTranslator("run-1", in, "", slog.Default()) + tr.SetSessionID("initial-sid") + ch, _, _ := tr.Subscribe() + + tr.NotifyTurnStart("") + + ev := drainEvent(t, ch) + assert.Equal(t, "initial-sid", ev.SessionID, + "empty sessionID must resolve to the initial session") +} + +// TestRun_StampsSessionIDFromNotification verifies that content events +// produced by translating ACP SessionNotifications carry the notification's +// session id (not just the Translator's initial-session cache). Without +// this, agent_message events from a non-initial session would all be +// labelled with the initial id, defeating per-session watch filtering. +func TestRun_StampsSessionIDFromNotification(t *testing.T) { + in := make(chan acp.SessionNotification, 1) + tr := NewTranslator("run-1", in, "", slog.Default()) + tr.SetSessionID("initial-sid") + ch, _, _ := tr.Subscribe() + tr.Start() + defer tr.Stop() + + in <- acp.SessionNotification{ + SessionId: acp.SessionId("session-xyz"), + Update: acp.SessionUpdate{ + AgentMessageChunk: &acp.SessionUpdateAgentMessageChunk{ + Content: acp.TextBlock("hi"), + }, + }, + } + + ev := drainEvent(t, ch) + assert.Equal(t, "session-xyz", ev.SessionID, + "content event should carry the notification's sessionID, not the initial") +} From c96eca8364234848f70b6f11e1d859800dd2a33c Mon Sep 17 00:00:00 2001 From: wjueyao Date: Tue, 19 May 2026 13:15:37 +0800 Subject: [PATCH 15/16] fix(runtime/acp): guard activePrompts decrement against clearSessions race MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When Kill() races with an in-flight PromptSession, clearSessions resets m.activePrompts to 0 before the prompt's deferred cleanup runs. The deferred decrement then drove the counter to -1 — harmless for phaseFromActivePromptsLocked's `> 0` check (still reads Idle) but a broken invariant that the next contributor reading the code would rightly question. Extract the deferred cleanup into a named helper decrementPromptInflight(sessionID) and guard the activePrompts decrement with `if > 0`. sess.inflight already had a map-lookup guard for the same race; this matches the symmetry. Tests: - TestDecrementPromptInflight_GuardsAgainstNegativeAfterClear pins the post-clearSessions race contract. - TestDecrementPromptInflight_NormalPath confirms the happy path still decrements both fields. Co-Authored-By: Claude Opus 4.7 (1M context) --- pkg/agentrun/runtime/acp/runtime.go | 27 ++++++++---- .../runtime/acp/runtime_internal_test.go | 43 +++++++++++++++++++ 2 files changed, 61 insertions(+), 9 deletions(-) diff --git a/pkg/agentrun/runtime/acp/runtime.go b/pkg/agentrun/runtime/acp/runtime.go index 8cc60c1..eadb639 100644 --- a/pkg/agentrun/runtime/acp/runtime.go +++ b/pkg/agentrun/runtime/acp/runtime.go @@ -468,6 +468,23 @@ func (m *Manager) resolveSessionLocked(sessionID acp.SessionId) acp.SessionId { return sessionID } +// decrementPromptInflight is the deferred cleanup for PromptSession. The +// session may have been deleted by Kill/clearSessions in the window between +// PromptSession's entry and conn.Prompt returning; both decrements guard +// for that — sess.inflight only if the session still exists, and +// activePrompts only when positive so a racing clearSessions (which sets +// activePrompts to 0) doesn't drive the counter negative. +func (m *Manager) decrementPromptInflight(sessionID acp.SessionId) { + m.mu.Lock() + defer m.mu.Unlock() + if s, ok := m.sessions[sessionID]; ok { + s.inflight-- + } + if m.activePrompts > 0 { + m.activePrompts-- + } +} + // SessionIDs returns a snapshot of currently-active session IDs. Order is // not stable — caller should sort if a deterministic order is needed. func (m *Manager) SessionIDs() []acp.SessionId { @@ -516,15 +533,7 @@ func (m *Manager) PromptSession(ctx context.Context, sessionID acp.SessionId, pr m.activePrompts++ m.mu.Unlock() - defer func() { - m.mu.Lock() - // sess may have been deleted by Kill/Delete; guard the deref. - if s, ok := m.sessions[sessionID]; ok { - s.inflight-- - } - m.activePrompts-- - m.mu.Unlock() - }() + defer m.decrementPromptInflight(sessionID) m.logger.Debug("prompt started", "sessionID", sessionID, "blocks", len(prompt)) diff --git a/pkg/agentrun/runtime/acp/runtime_internal_test.go b/pkg/agentrun/runtime/acp/runtime_internal_test.go index b309977..e77be2a 100644 --- a/pkg/agentrun/runtime/acp/runtime_internal_test.go +++ b/pkg/agentrun/runtime/acp/runtime_internal_test.go @@ -287,3 +287,46 @@ func TestClearSessions(t *testing.T) { t.Errorf("conn not cleared: %v", m.conn) } } + +// TestDecrementPromptInflight_GuardsAgainstNegativeAfterClear pins the +// guard against a PromptSession defer racing with Kill: clearSessions +// resets activePrompts to 0, then the still-in-flight prompt's deferred +// cleanup runs. Without the > 0 check the counter would slip to -1 and +// stay there until next clearSessions; with it, the decrement is a +// no-op and the counter remains coherent. +func TestDecrementPromptInflight_GuardsAgainstNegativeAfterClear(t *testing.T) { + m := newManagerForSessionTest(t) + m.sessionID = "initial" + m.sessions["initial"] = &sessionState{id: "initial"} + + // Simulate the post-clearSessions state: counters reset, session map + // emptied (the in-flight prompt's session is already gone). + m.activePrompts = 0 + delete(m.sessions, "initial") + + // Run the deferred cleanup that PromptSession would have queued. + m.decrementPromptInflight("initial") + + if m.activePrompts != 0 { + t.Errorf("activePrompts should stay at 0 after racing clearSessions, got %d", m.activePrompts) + } +} + +// TestDecrementPromptInflight_NormalPath verifies that the happy-path +// decrement (no race) still works — sess.inflight and m.activePrompts +// both go down by one. +func TestDecrementPromptInflight_NormalPath(t *testing.T) { + m := newManagerForSessionTest(t) + m.sessionID = "initial" + m.sessions["initial"] = &sessionState{id: "initial", inflight: 1} + m.activePrompts = 1 + + m.decrementPromptInflight("initial") + + if got := m.sessions["initial"].inflight; got != 0 { + t.Errorf("session inflight should be 0, got %d", got) + } + if m.activePrompts != 0 { + t.Errorf("activePrompts should be 0, got %d", m.activePrompts) + } +} From d919a5b958222d4a6a83ad16e7694f68c12596ec Mon Sep 17 00:00:00 2001 From: wjueyao Date: Tue, 19 May 2026 13:18:10 +0800 Subject: [PATCH 16/16] fix(runtime/acp): don't clobber Stopped Phase from late prompt writeState MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Companion to c96eca8 (activePrompts negative-guard). Same race class: Kill() → clearSessions() nils m.conn and writes Phase=Stopped. Meanwhile a PromptSession that was mid-conn.Prompt finally observes the dead pipe and returns an error. Its outer code then runs: _ = m.writeState(m.phaseFromActivePromptsLocked, reason) phaseFromActivePromptsLocked reads m.activePrompts and writes Phase based on it (0 → Idle, >0 → Running) — silently clobbering Stopped. state.json then says Idle/Running even though the agent is dead, and operators / recovery code make the wrong call. Fix: phaseFromActivePromptsLocked early-returns when m.conn is nil. Kill / clearSessions owns lifecycle phase past that point; the apply callback respects it instead of fighting it. Tests: - TestPhaseFromActivePromptsLocked_SkipsAfterKill pins the new contract: Phase=Stopped with conn=nil + activePrompts=1 stays Stopped. - TestPhaseFromActivePrompts updated to set a non-nil conn so it exercises the live-agent branch as before. Co-Authored-By: Claude Opus 4.7 (1M context) --- pkg/agentrun/runtime/acp/runtime.go | 8 +++++++ .../runtime/acp/runtime_internal_test.go | 23 +++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/pkg/agentrun/runtime/acp/runtime.go b/pkg/agentrun/runtime/acp/runtime.go index eadb639..4e35284 100644 --- a/pkg/agentrun/runtime/acp/runtime.go +++ b/pkg/agentrun/runtime/acp/runtime.go @@ -563,7 +563,15 @@ func (m *Manager) PromptSession(ctx context.Context, sessionID acp.SessionId, pr // activePrompts counter. Used as the writeState apply callback for prompt // start/end so Phase reflects "any session running" rather than the last // caller's local view. Runs under m.mu (writeState locks before calling). +// +// Skips the write when m.conn is nil — Kill / clearSessions has taken +// ownership of lifecycle phase (Stopped), and a late-firing PromptSession +// writeState (from a prompt that was in flight when Kill ran) must not +// clobber that with Idle / Running. func (m *Manager) phaseFromActivePromptsLocked(s *apiruntime.State) { + if m.conn == nil { + return + } if m.activePrompts > 0 { s.Phase = apiruntime.PhaseRunning } else { diff --git a/pkg/agentrun/runtime/acp/runtime_internal_test.go b/pkg/agentrun/runtime/acp/runtime_internal_test.go index e77be2a..2b9187e 100644 --- a/pkg/agentrun/runtime/acp/runtime_internal_test.go +++ b/pkg/agentrun/runtime/acp/runtime_internal_test.go @@ -231,6 +231,10 @@ func TestPromptSession_EmptySessionIDResolvesToInitial(t *testing.T) { // could leave Phase=Idle while one was still running. func TestPhaseFromActivePrompts(t *testing.T) { m := newManagerForSessionTest(t) + // Non-nil conn means the agent is alive — phaseFromActivePromptsLocked + // only writes Phase when m.conn != nil (it cedes lifecycle to Kill + // otherwise; see TestPhaseFromActivePromptsLocked_SkipsAfterKill). + m.conn = &acp.ClientSideConnection{} cases := []struct { name string @@ -253,6 +257,25 @@ func TestPhaseFromActivePrompts(t *testing.T) { } } +// TestPhaseFromActivePromptsLocked_SkipsAfterKill verifies that the apply +// callback leaves state.Phase untouched when m.conn is nil. This is the +// post-Kill state — Kill / clearSessions owns lifecycle phase (Stopped), +// and a late-firing PromptSession writeState (from a prompt that was in +// flight when Kill ran and finally errored on the dead pipe) must NOT +// clobber Stopped with Idle/Running. +func TestPhaseFromActivePromptsLocked_SkipsAfterKill(t *testing.T) { + m := newManagerForSessionTest(t) + m.conn = nil // post-Kill: clearSessions nilled the conn + m.activePrompts = 1 // would normally cause apply to write Running + + s := apiruntime.State{Phase: apiruntime.PhaseStopped} + m.phaseFromActivePromptsLocked(&s) + + if s.Phase != apiruntime.PhaseStopped { + t.Fatalf("Phase should remain Stopped after Kill, got %q", s.Phase) + } +} + // TestClearSessions verifies that the bookkeeping reset on agent // teardown (Kill / process exit) drops every field that could leave a // stale view behind. Without this, Sessions() would return dead IDs and