feat(runtime/acp): multi-session-per-agentrun

[wjueyao]

Summary

Adds multi-session support to ACP agentruns. One agent process can now host multiple ACP sessions opened via agentrun new-session, each with its own cwd / model state / message history — multiplexed onto the same fork instead of paying fork+exec per task.

Motivation

A downstream caller (an Agent evaluation harness) spawns 3 ACP agentruns per case (executor / responder / reviewer). Running 41 cases means 123 fork+exec cycles of bunx claude-agent-acp (~200MB Node + ~300MB claude binary + MCP servers each). On memory-pressured dev hosts this thrashes swap and 2-min socket-wait deadlines start failing — the workflow becomes unusable.

With multi-session, the caller can keep 3 long-lived agentruns and open a fresh session per case via agentrun new-session. Total fork+exec drops from 3N → 3.

Design

ACP natively supports multiple session/new calls on a single agent process — each session has its own sessionId, cwd, and model state, multiplexed onto the same stdin/stdout. This PR builds the runtime + RPC + CLI layers on top of that primitive.

Per-session state. Manager.sessions map[acp.SessionId]*sessionState tracks each open session's (id, cwd, models). NewSession adds entries; EndSession removes them. The agent process owns its own session map; the runtime mirrors only what it needs to route RPCs.

Initial-session compat. Create() opens the first session via the normal handshake and populates both the new sessions map and the legacy m.sessionID / m.models fields. Existing callers of Prompt / Cancel / SetModel (which don't pass sessionId) keep working — these are now thin shims that delegate to the *Session variants with an empty sessionId.

Empty-sessionId convention. Client and wire layers pass sessionId: "" to mean "the initial session". Resolution happens in exactly one place: Manager's PromptSession / CancelSession / SetModelSession substitute m.sessionID when sessionId is empty. The Service layer is a thin pass-through — no per-layer "empty → initial" branching.

MCP scoping. Bundle-level mcpServers apply to every session; per-session NewSession(cwd, mcpServers) layers session-scoped servers on top (current behavior: append, no dedup — caller responsibility).

In-flight protection. EndSession refuses to release a session with active PromptSession work, tracked via a per-session refcount in Manager.inflight. Callers must CancelSession (or wait) before EndSession; otherwise the prompt completes against a session no longer in the map and bookkeeping drifts. The initial session cannot be EndSession'd at all — kill the whole agent to release it.

session/cancel wire compat. The method evolved from NullaryCommand (no params) to taking optional sessionId. To keep callers that send no params field working, the jsonrpc framework gains an OptionalUnaryCommand helper that delivers a zero-value Req instead of -32602 InvalidParams when params is absent.

Out of scope (follow-ups). state.json is still single-session — Phase and Session.Models reflect the initial session only. Per-session persistence is a follow-up; short-lived sessions don't need it, longer-lived multi-session deployments do. No max-sessions cap or idle timeout either; caller drives lifecycle today.

What changed

Runtime layer (pkg/agentrun/runtime/acp/runtime.go):

• Manager.sessions map[acp.SessionId]*sessionState — replaces single-session field; each entry tracks id / cwd / per-session model state.
• Manager.inflight map[acp.SessionId]int — per-session active-prompt counter; EndSession refuses to release a busy session.
• Backward-compat: keep m.sessionID / m.models pointing to the first session so existing Prompt / Cancel / SetModel shims keep working unchanged.
• New methods: NewSession(ctx, cwd, mcpServers), EndSession(sessionId), Sessions(), plus PromptSession / CancelSession / SetModelSession variants taking explicit sessionId (empty resolved to initial).

API layer (pkg/agentrun/api, pkg/ari/server):

• Method constants + RPC types for agentrun/new-session, agentrun/end-session, agentrun/list-sessions.
• Server handlers wire through to Manager. Existing agentrun/prompt / cancel / set-model gain optional sessionId — empty = initial session (back-compat).

Framework (pkg/jsonrpc):

• ErrNoParams sentinel returned by the unmarshal callback when the request has no params field.
• OptionalUnaryCommand helper — tolerates ErrNoParams, surfaces real JSON decode errors as InvalidParams. Used for session/cancel.

CLI (cmd/massctl/commands/agentrun):

• massctl agentrun new-session <name> -w <ws> --cwd <path> — returns sessionId on stdout.
• massctl agentrun end-session <name> -w <ws> --session-id <sid>.
• massctl agentrun list-sessions <name> -w <ws>.
• agentrun prompt gains --session-id flag.

ARI client (pkg/ari/client): typed wrappers for the new RPCs.

Tests:

• pkg/agentrun/runtime/acp/runtime_internal_test.go: bookkeeping tests for sessions map (add / remove / list / busy / state isolation).
• pkg/jsonrpc/server_test.go: wire-compat tests for OptionalUnaryCommand (missing params tolerated; malformed params surfaced).

Bundled fix (pkg/ari/api/domain.go): add explicit yaml: tags to AgentSpec fields. Without them, gopkg.in/yaml.v3 lowercases Go field names and silently drops camelCase keys like startupTimeoutSeconds on massctl agent apply -f round-trip. Tangentially related (camelCase fields matter for multi-session pool YAML), small enough to fold in.

Verification

Local go test ./... clean (one flake in TestAgentdRestartRecovery under full-suite contention; passes in isolation and on rerun — same behavior on main, not introduced by this PR).

Compat / migration

All existing single-session call sites unchanged. New methods are additive. State schema additive (sessions map) — no migration needed for daemons that don't open extra sessions. Pre-multi-session clients calling session/cancel with no params field continue to work via the OptionalUnaryCommand helper.

Next

A downstream pool adapter (wrapping new-session/end-session for the evaluation harness above) is implemented and unit-tested but blocked on this PR landing.

[wjueyao]

Reviewed end-to-end — commits, full diff, and downstream impact. Posting a structured set of findings so they're discoverable as follow-ups.

What's solid

• Back-compat layering: keeping m.sessionID / m.models pointing at the initial session and turning legacy Prompt/Cancel/SetModel into shims around the *Session variants is the right shape — existing callers stay unchanged, new callers opt into multi-session.
• Commit hygiene: 8 atomic commits, each compile-clean with a focused scope (scaffold → variants → wire types → server → client → ARI → CLI → tests). Easy to review / bisect.
• In-code documentation: sessionState rationale, why EndSession refuses initial, optional SessionID semantics — all spelled out in doc comments rather than hidden behind an external design doc.
• Type-boundary discipline: every wire boundary casts acp.SessionId ↔ string explicitly. No bare strings leaking into the runtime.

Suggestions worth addressing in this PR

M1 — Service.Prompt (and Cancel / SetModel) dispatch duplication

In pkg/agentrun/server/service.go:

if req.SessionID == "" {
    resp, err = s.mgr.Prompt(ctx, req.Prompt)
} else {
    resp, err = s.mgr.PromptSession(ctx, resolveSessionID(req.SessionID), req.Prompt)
}

Manager.Prompt is now itself a shim around PromptSession(initialSessionID), so this two-branch dispatch does the same thing twice. The "empty string ⇒ initial session" rule is duplicated across client, Service, and Manager layers — three places to update if/when the convention changes.

Suggest consolidating at the Manager boundary, e.g.:

func (m *Manager) PromptSession(ctx, sessionID acp.SessionId, prompt) (...) {
    m.mu.Lock()
    if sessionID == "" {
        sessionID = m.sessionID
    }
    // ... existing logic
}

Then Service / Manager.Prompt collapse to one-liners, and the eventual deprecation of the empty-string convention is a one-line change. Same pattern applies to Cancel / SetModel.

M2 — session/cancel wire shape change (NullaryCommand → UnaryCommand)

pkg/agentrun/server/register.go switches "cancel" from NullaryCommand to UnaryCommand. The Service handler nil-guards req, which is good — but the wire-level behavior depends on what jsonrpc.UnaryCommand does when a JSON-RPC request has no params field (vs params: null vs params: {}).

The PR description's compat claim ("Old clients calling Cancel with nil params: server handler tolerates nil req") covers the Go-side req *T == nil case, but not necessarily the JSON envelope path where the framework decides what reaches the handler. A quick look at pkg/jsonrpc/server.go's unmarshal callback suggests it returns "missing params" when req.Params == nil, which UnaryCommand then wraps as ErrInvalidParams — meaning callers omitting the params field never reach the Service nil-guard.

Suggest verifying with a wire-level test that pins the contract regardless of jsonrpc-framework changes:

• raw JSON-RPC client sends {"jsonrpc":"2.0","id":1,"method":"session/cancel"} (no params field) → expect 200, cancel routed to initial session.

If the framework does reject the request, a small OptionalUnaryCommand helper (or equivalent) in the jsonrpc package keeps the migration purely additive.

M3 — EndSession vs. in-flight prompt — contract is ambiguous

func (m *Manager) EndSession(sessionID acp.SessionId) error {
    // ... delete(m.sessions, sessionID) — no in-flight check
}

The doc comment says callers "must still call CancelSession (or wait for prompt completion) before EndSession" but EndSession itself doesn't enforce or detect it. If a caller ends a session mid-prompt:

1. The in-flight PromptSession already passed its _, known := m.sessions[sessionID] check at start, so it continues to completion.
2. The prompt's defer writes state.Phase = PhaseIdle after EndSession returned — bookkeeping diverges silently from the map.
3. Trailing log lines reference a sessionID no longer in the map.

Suggest one of:

• (a) Track an "in-flight prompts per session" counter; EndSession returns EBUSY if non-zero.
• (b) Keep current behavior but drop the misleading "(or wait for prompt completion)" hint from the docstring — make it explicit that the caller is solely responsible for prompt ordering and the runtime does no enforcement.

For callers that fully own session lifecycle, (b) is acceptable; (a) is a nicer interface for less disciplined consumers.

Smaller items (could be follow-ups)

• m1. MCP server merge has no dedup: merged := bundle; merged = append(merged, sessionLevel...) — if a session-scoped MCP server shares a name with a bundle-level one, both register. Consider dedup-by-Stdio.Name with session-scoped winning.
• m2. ListSessions returns IDs only: for debugging it'd be nicer to return []{id, cwd, currentModelId} so operators can see what each session is doing without further RPCs. Cheap upgrade — Manager already has the data.
• m3. Audit args for initial session omit sessionId: if sessionID != "" { auditArgs["sessionId"] = ... } — when replaying audit logs, initial-session cancels look like "untargeted cancel". Suggest always emitting the resolved session id (including initial).
• m4. AgentSpec yaml-tag bundling: clean bug-fix and the PR description is honest about it, but semantically unrelated to multi-session. Cleaner as its own commit/PR — easier to revert independently if needed. Non-blocking.

Test coverage gaps

• Unit tests for agentrun/server Service handlers: NewSession / EndSession / ListSessions validation paths (missing cwd, missing sessionId, empty req) have no direct coverage.
• Unit tests for the ari/server adapter handlers: same gap; NewSession / EndSession / ListSessions validation + agent not running path uncovered.
• Wire-compat test for M2 (above): paramless session/cancel request.
• The TestRuntimeSuite flakiness self-disclosed in the PR description is confirmed pre-existing — agree it shouldn't block this PR, but worth a tracking issue so it doesn't become normalized.

Known follow-ups already self-disclosed (not blocking)

• state.json schema has no per-session Phase / Models — multi-session sessions are not crash-recoverable. Fine for short-lived sessions; needs design before long-lived use.
• No max-sessions cap; no idle timeout. Pathological callers can leak sessions until process OOM. Not urgent.

Bottom line

Design is sound, motivation grounded in measured fork+exec cost, and the back-compat layering lets the migration happen one caller at a time. M1 is structural cleanup, M2 is the back-compat claim most worth pinning down with a test, M3 is contract clarity. None of the three should require a substantial rework — likely a half-day to address.