diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 23327ba1242..683a0d41dcb 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,9 @@ ## [Unreleased] +- Changed interactive, print, JSON, RPC, piped-stdin, and no-session clients to use the same daemon-owned runtime while preserving their existing commands, output protocols, and lifecycle behavior ([ENG-4685](https://linear.app/primeintellect/issue/ENG-4685)). +- Added RPC controls for schedules, heartbeats, agent messaging, and live session observation ([ENG-4685](https://linear.app/primeintellect/issue/ENG-4685)). +- Fixed daemon-backed headless startup, rollback routing, RPC wire compatibility, and duplicate client runtime preparation ([ENG-4685](https://linear.app/primeintellect/issue/ENG-4685)). - Fixed agents with heartbeat-owning subagents appearing completed and showing completion checkmarks in the Agents View. - Fixed daemon backpressure triggering redundant catch-up snapshots for events already queued by the socket. - Fixed incompatible daemon builds crashing startup or respawning after shutdown, with capability negotiation, verified provenance, and convergent force shutdown ([ENG-4687](https://linear.app/primeintellect/issue/ENG-4687/make-daemon-version-mismatches-self-healing)). diff --git a/packages/coding-agent/docs/daemon-implementation-summary.md b/packages/coding-agent/docs/daemon-implementation-summary.md index b889b0bb604..2581dab70c5 100644 --- a/packages/coding-agent/docs/daemon-implementation-summary.md +++ b/packages/coding-agent/docs/daemon-implementation-summary.md @@ -5,18 +5,14 @@ Prime Agent isolates each persistent root session tree in its own process. Proce ## Process Topology ```text -interactive clients - | - v -detached supervisor - |- catalog subprocess - |- resident worker: root A + RLM descendants + kernels - `- resident worker: root B + RLM descendants + kernels - -print / JSON / RPC / interactive --no-session - | - v -frontend process --private owner channel--> owned worker + root tree +interactive / print / JSON / RPC clients + | + v + detached supervisor + |- catalog subprocess + |- resident worker: root A + RLM descendants + kernels + |- resident worker: root B + RLM descendants + kernels + `- client-owned worker: hidden root + RLM descendants + kernels ``` The supervisor owns public sockets, client attachments, routing, global agent-message routing, worker health, and update coordination. Each resident worker owns scheduling and execution for its root tree. The supervisor does not execute providers, tools, compaction, bash, kernels, session schedules, or session-history scans. @@ -41,25 +37,27 @@ Normal interactive sessions use resident workers: There are no session, worker, client, or workload caps in this layer. -## Owned Workers +## Client-Owned Workers -Headless and ephemeral modes use the same process isolation without joining the resident daemon registry: +Headless and ephemeral clients use the same supervisor, worker process, `AgentSessionRuntime`, extension binding, scheduler, RLM host, and kernel lifecycle as interactive clients. Their workers have a client-owned lifecycle: -- `prime-agent -p`, piped stdin, and JSON mode proxy stdin/stdout/stderr through an owned worker and return its existing exit status. -- RPC keeps stdin/stdout in the frontend. It forwards strict LF-delimited JSONL to the worker and preserves every public command, response, event, and extension-UI shape. -- If a persisted RPC worker crashes, the frontend reports an ordinary failure for commands whose result is uncertain, reopens the exact persisted session, and continues accepting commands. It never replays an uncertain command. -- RPC EOF and frontend `SIGTERM`/`SIGHUP` stop only the owned worker and its subprocesses. -- Interactive `--no-session` uses an owned in-memory worker and creates no recovery descriptor. -- Owned workers never appear in the agents view and never require the global daemon. +- `prime-agent -p`, piped stdin, and JSON mode remain one-shot clients and return their existing output and exit status. +- RPC retains strict LF-delimited JSONL framing, accepts repeated prompts until EOF, and preserves every public command, response, event, and extension-UI shape. +- Interactive `--no-session` uses a client-owned in-memory session. +- Normal completion or RPC EOF explicitly removes the worker without archiving its session. `SIGTERM` and `SIGHUP` retain their existing exit codes. +- Unexpected client loss starts a bounded cleanup grace period. Reconnection with the same stable protocol client ID cancels cleanup and reattaches to the same worker. +- The supervisor can recover a crashed client-owned worker while the owner remains connected. After supervisor replacement, the owner resupplies its launch environment before a dead worker is restarted. +- Default session listings, the agents view, global peer routing, and global schedule aggregation omit client-owned workers. Owner-scoped heartbeat and cron operations still reach the hidden worker. +- The full launch environment is held only in supervisor memory and is never written to the worker descriptor. - Direct SDK calls to `runPrintMode(runtime, ...)` and `runRpcMode(runtime)` remain in-process and source-compatible. Non-serializable extension factories therefore continue to work for embedders. -An inherited, unreferenced IPC owner channel links the worker to its frontend. Frontend loss disconnects the channel and terminates the worker, which reaps its kernels and detached children; the unreferenced channel does not prevent clean RPC EOF or print completion. +The daemon is infrastructure, not a client mode. `interactive`, `print`, `json`, and `rpc` describe client behavior; the retained `--mode daemon` command is the compatible internal process entrypoint used to start the supervisor and session workers. ## Session Ownership -Persisted sessions use process-safe leases keyed by canonical JSONL path. Resident and owned workers both acquire leases. +Persisted sessions use process-safe leases keyed by canonical JSONL path. Resident and client-owned workers both acquire leases. -- Resident workers acquire a target lease before opening an existing session. Owned CLI workers acquire it before runtime writes, and all workers acquire the replacement lease before switching runtimes. +- Workers acquire a target lease before opening an existing session and acquire the replacement lease before switching runtimes. - The previous lease is released only after replacement succeeds. - Concurrent opens return `session_already_active` with the owning active-session ID. - Concurrent daemon creates for the same path share one worker launch. @@ -68,7 +66,7 @@ This prevents a script and resident daemon worker from writing the same session ## Scheduled Jobs -Each resident worker runs one scheduler for the root and RLM descendant sessions it owns. Schedule state is persisted per session at `session-artifacts//scheduled-jobs.json`; workers never scan or write a shared global cron file. +Each worker runs one scheduler for the root and RLM descendant sessions it owns. Schedule state is persisted per session at `session-artifacts//scheduled-jobs.json`; workers never scan or write a shared global cron file. - Creating or changing a heartbeat writes its session store and wakes the same worker's scheduler. - Due ticks are claimed and advanced durably before prompt delivery, so a crash never replays an uncertain prompt. @@ -80,9 +78,9 @@ Each resident worker runs one scheduler for the root and RLM descendant sessions The first worker-owned release migrates the legacy global `cron-jobs.json` into the corresponding session artifact directories before adopting workers. Archived, deleted, or descriptorless sessions are cancelled rather than revived. -## Public Protocol v2 +## Public Protocol v3 -The public daemon socket remains JSONL-framed. Version 2 adds: +The public daemon socket remains JSONL-framed. Versions 2 and 3 add: - stable command envelopes carrying protocol version, client ID, and command ID; - generation-aware event cursors `{ generation, sequence }`; @@ -91,6 +89,9 @@ The public daemon socket remains JSONL-framed. Version 2 adds: - attach acknowledgement followed by `session_snapshot_begin`, `session_snapshot_chunk`, and `session_snapshot_end`; - 512 KiB target transcript chunks; - file-backed transcript caches above 4 MiB. +- client-owned worker creation, completion, and promotion; +- owner-scoped recovery with in-memory launch environments; +- daemon-side headless completion, session-header, synchronous bash, and retry-setting commands. The one-release v1 update handoff uses raw v1 requests only to prepare and stop the old daemon. Busy older daemons that cannot produce a recovery manifest are left running. diff --git a/packages/coding-agent/docs/rpc.md b/packages/coding-agent/docs/rpc.md index 7a284453a7a..4b8a34c5451 100644 --- a/packages/coding-agent/docs/rpc.md +++ b/packages/coding-agent/docs/rpc.md @@ -211,6 +211,46 @@ Response: Messages are `AgentMessage` objects (see [Message Types](#message-types)). +### Daemon Coordination + +RPC clients can use the same daemon-owned coordination features as the interactive client. These commands are additive; existing commands and event shapes are unchanged. + +| Command | Fields | Successful response data | +|---------|--------|--------------------------| +| `send_message` | `targetActiveSessionId`, `message`, optional `deliveryMode` (`auto`, `steer`, `follow_up`) | Agent-message delivery receipt | +| `agent_messages_status` | none | Messaging safety status | +| `agent_messages_pause` | none | Updated messaging safety status | +| `agent_messages_resume` | none | Updated messaging safety status | +| `agent_messages_clear` | none | `{ "cleared": number }` | +| `list_schedules` | optional `includeInactive` | `{ "jobs": [...] }` | +| `add_schedule` | `schedule`, `prompt` | `{ "job": {...} }` | +| `cancel_schedule` | `jobId` | `{ "job": {...} }` | +| `list_heartbeats` | none | `{ "heartbeats": [...] }` | +| `get_heartbeat` | none | `{ "heartbeat": {...} }` or `null` | +| `set_heartbeat` | `schedule`, `prompt`, optional `deliveryMode` (`steer`, `follow_up`) | `{ "heartbeat": {...} }` | +| `update_heartbeat` | `action` (`pause`, `resume`, `clear`) | `{ "heartbeat": {...} }` or `null` | +| `manage_heartbeat` | `activeSessionId`, `jobId`, `action` (`pause`, `resume`, `stop`) | `{ "heartbeat": {...} }` | + +Adding a schedule or heartbeat promotes an invocation-local RPC session to a resident daemon session so the scheduled work remains available after RPC stdin closes. + +Use `observe` to subscribe to another active root or subagent session: + +```json +{"id":"watch-1","type":"observe","activeSessionId":"target-session-id"} +``` + +The response contains the target's current messages. Later target events are wrapped so they cannot be confused with events from the RPC client's own session: + +```json +{"type":"observed_session_event","activeSessionId":"target-session-id","event":{"type":"agent_start"}} +``` + +When the target closes, RPC emits `observed_session_closed`. Stop watching explicitly with: + +```json +{"type":"unobserve","activeSessionId":"target-session-id"} +``` + ### Model #### set_model diff --git a/packages/coding-agent/src/cli-main.ts b/packages/coding-agent/src/cli-main.ts index 752a239d8ff..34350b70ba2 100644 --- a/packages/coding-agent/src/cli-main.ts +++ b/packages/coding-agent/src/cli-main.ts @@ -1,11 +1,14 @@ import { enableCompileCache } from "node:module"; -import { maybeStartInteractiveDaemonEarly } from "./cli/daemon-launch.js"; +import { maybeStartDaemonEarly, shouldUseLegacyOwnedSessionWorkerFrontend } from "./cli/daemon-launch.js"; import { + classifyOwnedSessionWorkerInvocation, closeOwnedSessionWorkerOwnerWatch, installOwnedSessionWorkerOwnerWatch, + isOwnedSessionWorkerProcess, maybeRunOwnedSessionWorkerFrontend, } from "./cli/owned-session-worker.js"; import { APP_NAME } from "./config.js"; +import { defaultDaemonSocketPath } from "./modes/daemon/daemon-socket.js"; export async function runCli(): Promise { try { @@ -18,13 +21,20 @@ export async function runCli(): Promise { process.env.PI_CODING_AGENT = "true"; process.emitWarning = (() => {}) as typeof process.emitWarning; - // Boot a cold daemon concurrently with this process's heavy imports. - maybeStartInteractiveDaemonEarly(process.argv.slice(2)); - installOwnedSessionWorkerOwnerWatch(); - const handledByOwnedWorker = await maybeRunOwnedSessionWorkerFrontend(process.argv.slice(2)); + const args = process.argv.slice(2); + const ownedWorkerProfile = classifyOwnedSessionWorkerInvocation(args, process.stdin.isTTY); + const useLegacyOwnedWorker = + process.env.PRIME_AGENT_INTERNAL_LEGACY_OWNED_WORKER_FRONTEND !== "1" && + ownedWorkerProfile !== undefined && + (await shouldUseLegacyOwnedSessionWorkerFrontendFromArgs(args)); + const handledByOwnedWorker = await maybeRunOwnedSessionWorkerFrontend(args, useLegacyOwnedWorker); if (!handledByOwnedWorker) { + if (!isOwnedSessionWorkerProcess()) { + // Boot a cold daemon concurrently with this process's heavy imports. + maybeStartDaemonEarly(process.argv.slice(2)); + } const [{ EnvHttpProxyAgent, setGlobalDispatcher }, { main }] = await Promise.all([ import("undici"), import("./main.js"), @@ -41,3 +51,10 @@ export async function runCli(): Promise { } } } + +async function shouldUseLegacyOwnedSessionWorkerFrontendFromArgs(args: readonly string[]): Promise { + const socketIndex = args.indexOf("--daemon-socket"); + const socketPath = + socketIndex !== -1 && args[socketIndex + 1] ? (args[socketIndex + 1] as string) : defaultDaemonSocketPath(); + return shouldUseLegacyOwnedSessionWorkerFrontend(socketPath); +} diff --git a/packages/coding-agent/src/cli/daemon-launch.ts b/packages/coding-agent/src/cli/daemon-launch.ts index f50a6d9785c..96a824336a0 100644 --- a/packages/coding-agent/src/cli/daemon-launch.ts +++ b/packages/coding-agent/src/cli/daemon-launch.ts @@ -1,10 +1,8 @@ /** - * Interactive-daemon launch/readiness helpers. + * Daemon launch/readiness helpers. * - * This module is deliberately light on imports: cli.ts calls - * maybeStartInteractiveDaemonEarly() BEFORE the heavy main module graph loads, - * so a cold daemon boots concurrently with this process's own imports instead - * of serially after them. main.ts reuses the same memoized promise. + * This module stays light on imports so clients can start a cold daemon before + * the heavy main module graph loads. main.ts reuses the same memoized promise. */ import { spawn } from "node:child_process"; @@ -15,10 +13,12 @@ import { getProcessStartId } from "../core/session-lease.js"; import { DaemonClient, type DaemonHello } from "../modes/daemon/daemon-client.js"; import { DAEMON_PROTOCOL_VERSION, DAEMON_SCHEMA_ID } from "../modes/daemon/daemon-protocol.js"; import { getDaemonRuntimeIdentity } from "../modes/daemon/daemon-runtime-identity.js"; -import type { SessionSummary } from "../modes/daemon/daemon-session-list.js"; +import { isSessionSummaryBusy, type SessionSummary } from "../modes/daemon/daemon-session-list.js"; import { defaultDaemonSocketPath } from "../modes/daemon/daemon-socket.js"; import { isHelpCommandRequest, PUBLIC_COMMAND_NAMES, REMOVED_COMMAND_NAMES } from "./command-registry.js"; -import { formatCurrentCliCommand } from "./subprocess-launch.js"; +import { createCliSubprocessEnv, formatCurrentCliCommand } from "./subprocess-launch.js"; + +const DAEMON_STARTUP_TIMEOUT_MS = 30_000; export function isDaemonSessionSummary(value: unknown): value is SessionSummary { if (!value || typeof value !== "object") { @@ -102,12 +102,22 @@ export async function probeDaemonVersion(socketPath: string): Promise { +export async function listActiveDaemonSessionSummaries( + client: DaemonClient, + options: { includeClientOwned?: boolean } = {}, +): Promise { + return (await queryActiveDaemonSessions(client, options)).sessions; +} + +async function queryActiveDaemonSessions( + client: DaemonClient, + options: { includeClientOwned?: boolean } = {}, +): Promise<{ sessions: SessionSummary[]; busyClientOwnedSessionCount: number }> { const hello = await client.waitForHello(2000).catch(() => undefined); const response = hello && hello.protocol.version < DAEMON_PROTOCOL_VERSION - ? await client.requestLegacy({ type: "list" }) - : await client.request({ type: "list" }); + ? await client.requestLegacy({ type: "list", includeClientOwned: options.includeClientOwned }) + : await client.request({ type: "list", includeClientOwned: options.includeClientOwned }); if (!response.success) { throw new Error(response.error); } @@ -122,7 +132,16 @@ export async function listActiveDaemonSessionSummaries(client: DaemonClient): Pr if (!sessions.every(isDaemonSessionSummary)) { throw new Error("Daemon returned an invalid session list response"); } - return sessions; + const busyClientOwnedSessionCount = (data as { busyClientOwnedSessionCount?: unknown }).busyClientOwnedSessionCount; + if ( + busyClientOwnedSessionCount !== undefined && + (typeof busyClientOwnedSessionCount !== "number" || + !Number.isInteger(busyClientOwnedSessionCount) || + busyClientOwnedSessionCount < 0) + ) { + throw new Error("Daemon returned an invalid client-owned session count"); + } + return { sessions, busyClientOwnedSessionCount: busyClientOwnedSessionCount ?? 0 }; } /** Thrown when a stale-version daemon can't be replaced. The message is user-facing. */ @@ -250,18 +269,12 @@ export async function shutdownDaemonAndWait(socketPath: string, timeoutMs = 5000 // activeSessions is undefined when the daemon is reachable but its sessions couldn't // be listed — callers must treat that as "possibly busy", not idle. -export type RunningDaemonProbe = { reachable: false } | { reachable: true; activeSessions?: SessionSummary[] }; +export type RunningDaemonProbe = + | { reachable: false } + | { reachable: true; activeSessions?: SessionSummary[]; busyClientOwnedSessionCount?: number }; export function isSessionBusy(summary: SessionSummary): boolean { - // pendingMessageCount covers queued steering/follow-ups, which live only in - // memory and would be lost if the daemon were stopped. - return ( - summary.isStreaming || - summary.isCompacting || - summary.isBashRunning === true || - summary.hasRunningRlmChildren === true || - summary.pendingMessageCount > 0 - ); + return isSessionSummaryBusy(summary); } export async function probeRunningDaemonSessions(socketPath: string): Promise { @@ -273,8 +286,14 @@ export async function probeRunningDaemonSessions(socketPath: string): Promise summary.activeSessionId !== undefined) }; + const result = await queryActiveDaemonSessions(client, { includeClientOwned: true }); + return { + reachable: true, + activeSessions: result.sessions.filter((summary) => summary.activeSessionId !== undefined), + ...(result.busyClientOwnedSessionCount > 0 + ? { busyClientOwnedSessionCount: result.busyClientOwnedSessionCount } + : {}), + }; } catch { return { reachable: true }; } finally { @@ -282,6 +301,20 @@ export async function probeRunningDaemonSessions(socketPath: string): Promise { + const version = await probeDaemonVersion(socketPath); + if (version.status !== "stale") { + return false; + } + const running = await probeRunningDaemonSessions(socketPath); + return ( + running.reachable && + (running.activeSessions === undefined || + (running.busyClientOwnedSessionCount ?? 0) > 0 || + running.activeSessions.some((summary) => isSessionBusy(summary))) + ); +} + // Idle-but-loaded sessions reload from disk on the fresh daemon, so only a busy // session blocks replacing a stale daemon. async function shutdownStaleDaemonIfNotBusy(socketPath: string): Promise { @@ -293,9 +326,10 @@ async function shutdownStaleDaemonIfNotBusy(socketPath: string): Promise isSessionBusy(summary)); } catch { // Couldn't confirm idleness: treat as busy rather than risk interrupting work. hasBusySessions = true; @@ -346,7 +380,7 @@ async function ensureDaemonRunning(socketPath: string, spawnCwd?: string): Promi { cwd: spawnCwd ?? process.cwd(), detached: true, - env: process.env, + env: createCliSubprocessEnv(), // The daemon writes its own rotating log (see daemon-mode); nothing here // needs its stdout/stderr, so leave them detached. stdio: "ignore", @@ -354,7 +388,7 @@ async function ensureDaemonRunning(socketPath: string, spawnCwd?: string): Promi ); child.unref(); - const deadline = Date.now() + 10000; + const deadline = Date.now() + DAEMON_STARTUP_TIMEOUT_MS; while (Date.now() < deadline) { const started = await probeDaemonVersion(socketPath); if (started.status === "current") { @@ -389,66 +423,101 @@ export function ensureInteractiveDaemonRunning(socketPath: string, spawnCwd?: st return promise; } -const EARLY_LAUNCH_EXCLUDED_FLAGS = new Set([ +const EARLY_LAUNCH_EXCLUDED_FLAGS = new Set(["--help", "-h", "--version", "-v", "--list-models", "--export"]); +const EARLY_LAUNCH_VALUE_FLAGS = new Set([ "--mode", - "--print", - "-p", - "--help", - "-h", - "--version", - "-v", - "--no-session", - "--list-models", - "--export", + "--daemon-socket", + "--provider", + "--model", + "--api-key", + "--cwd", + "--system-prompt", + "--append-system-prompt", + "--fork", + "--session-dir", + "--models", + "--tools", + "-t", + "--thinking", + "--extension", + "-e", + "--skill", + "--prompt-template", + "--theme", + "--autonomous-gate", + "--autonomous-gate-retries", + "--autonomous-gate-timeout-ms", + "--autonomous-max-continuations", + "--autonomous-max-turns", + "--autonomous-max-tokens", + "--autonomous-timeout-ms", ]); -/** Conservative pre-parse of argv: true only when startup clearly heads into daemon-backed interactive mode. */ -export function shouldStartInteractiveDaemonEarly( - args: readonly string[], - stdinIsTTY: boolean | undefined, - startupBenchmark: boolean, -): boolean { - if (startupBenchmark || !stdinIsTTY) { +function findFirstEarlyLaunchPositional(args: readonly string[]): { index: number; value: string } | undefined { + for (let index = 0; index < args.length; index++) { + const arg = args[index]!; + if (arg === "--") { + return args[index + 1] === undefined ? undefined : { index: index + 1, value: args[index + 1]! }; + } + if (EARLY_LAUNCH_VALUE_FLAGS.has(arg)) { + index++; + continue; + } + if (arg === "--resume" || arg === "-r") { + if (args[index + 1] && !args[index + 1]!.startsWith("-")) { + index++; + } + continue; + } + if (!arg.startsWith("-")) { + return { index, value: arg }; + } + } + return undefined; +} + +export function shouldStartDaemonEarly(args: readonly string[], startupBenchmark: boolean): boolean { + if (startupBenchmark) { return false; } - const firstPositionalIndex = args.findIndex((arg) => arg.length > 0 && !arg.startsWith("-")); - const firstPositional = args[firstPositionalIndex]; - const isHelpCommand = firstPositional === "help" && isHelpCommandRequest(args.slice(firstPositionalIndex + 1)); + const modeIndex = args.indexOf("--mode"); + if (modeIndex !== -1 && args[modeIndex + 1] === "daemon") { + return false; + } + if (args.some((arg) => EARLY_LAUNCH_EXCLUDED_FLAGS.has(arg))) { + return false; + } + if (args.includes("--print") || args.includes("-p")) { + return true; + } + const firstPositional = findFirstEarlyLaunchPositional(args); + const isHelpCommand = + firstPositional?.value === "help" && isHelpCommandRequest(args.slice(firstPositional.index + 1)); if ( firstPositional && - (REMOVED_COMMAND_NAMES.has(firstPositional) || - (PUBLIC_COMMAND_NAMES.has(firstPositional) && - firstPositional !== "agents" && - (firstPositional !== "help" || isHelpCommand))) + (REMOVED_COMMAND_NAMES.has(firstPositional.value) || + (PUBLIC_COMMAND_NAMES.has(firstPositional.value) && + firstPositional.value !== "agents" && + (firstPositional.value !== "help" || isHelpCommand))) ) { return false; } - return !args.some((arg) => EARLY_LAUNCH_EXCLUDED_FLAGS.has(arg)); + return true; } -/** - * Fire-and-forget daemon launch for interactive startups, called from cli.ts - * before the heavy module graph is imported. A wrong "yes" merely starts the - * daemon that the next interactive run would need anyway; a wrong "no" only - * means main.ts starts it later, as before. Never throws — errors surface when - * main.ts awaits the memoized promise. - */ -export function maybeStartInteractiveDaemonEarly(args: readonly string[]): void { +export function maybeStartDaemonEarly(args: readonly string[]): void { const benchmarkFlag = (process.env.PI_STARTUP_BENCHMARK ?? "").toLowerCase(); const startupBenchmark = benchmarkFlag === "1" || benchmarkFlag === "true" || benchmarkFlag === "yes"; - if (!shouldStartInteractiveDaemonEarly(args, process.stdin.isTTY, startupBenchmark)) { + if (!shouldStartDaemonEarly(args, startupBenchmark)) { return; } const socketIndex = args.indexOf("--daemon-socket"); const socketPath = socketIndex !== -1 && args[socketIndex + 1] ? (args[socketIndex + 1] as string) : defaultDaemonSocketPath(); - // Honor --cwd: main() chdirs after parsing, but this runs before; spawn the - // daemon from the target directory so it matches the old post-chdir behavior. const cwdIndex = args.indexOf("--cwd"); const cwdArg = cwdIndex !== -1 ? args[cwdIndex + 1] : undefined; const spawnCwd = cwdArg ? resolve(expandTildePath(cwdArg)) : undefined; if (spawnCwd && !existsSync(spawnCwd)) { - // Invalid --cwd: skip the early spawn; main() reports the error. return; } void ensureInteractiveDaemonRunning(socketPath, spawnCwd); diff --git a/packages/coding-agent/src/cli/daemon-stop-confirm.ts b/packages/coding-agent/src/cli/daemon-stop-confirm.ts index 9b6d88f19a7..ab39b25b9db 100644 --- a/packages/coding-agent/src/cli/daemon-stop-confirm.ts +++ b/packages/coding-agent/src/cli/daemon-stop-confirm.ts @@ -61,11 +61,12 @@ export async function confirmDaemonSessionLoss( // Reachable but couldn't list sessions: assume work may be lost. detail = copy.unlistableDetail; } else { - const busySessions = probe.activeSessions.filter(isSessionBusy); - if (busySessions.length === 0) { + const busySessionCount = + probe.activeSessions.filter(isSessionBusy).length + (probe.busyClientOwnedSessionCount ?? 0); + if (busySessionCount === 0) { return true; } - detail = copy.busyDetail(busySessions.length); + detail = copy.busyDetail(busySessionCount); } if (!process.stdin.isTTY) { console.error(chalk.red(`${detail} ${copy.nonTtyHint}`)); diff --git a/packages/coding-agent/src/cli/owned-session-worker.ts b/packages/coding-agent/src/cli/owned-session-worker.ts index c7447862e52..f7ab3753550 100644 --- a/packages/coding-agent/src/cli/owned-session-worker.ts +++ b/packages/coding-agent/src/cli/owned-session-worker.ts @@ -24,6 +24,10 @@ let closeOwnerWatch: (() => void) | undefined; export type OwnedSessionWorkerProfile = "print" | "json" | "rpc" | "interactive-ephemeral"; +export function isOwnedSessionWorkerProcess(environment: NodeJS.ProcessEnv = process.env): boolean { + return environment[OWNED_WORKER_ENV] === "1"; +} + interface OwnedSessionRecoveryDescriptor { version: 1; profile: OwnedSessionWorkerProfile; @@ -63,7 +67,7 @@ export function classifyOwnedSessionWorkerInvocation( stdinIsTTY: boolean | undefined, environment: NodeJS.ProcessEnv = process.env, ): OwnedSessionWorkerProfile | undefined { - if (environment[OWNED_WORKER_ENV] === "1" || hasNonSessionOperation(args)) { + if (isOwnedSessionWorkerProcess(environment) || hasNonSessionOperation(args)) { return undefined; } @@ -472,7 +476,13 @@ export async function runOwnedSessionWorkerFrontend( } } -export async function maybeRunOwnedSessionWorkerFrontend(args: readonly string[]): Promise { +export async function maybeRunOwnedSessionWorkerFrontend( + args: readonly string[], + forceLegacyFrontend = false, +): Promise { + if (!forceLegacyFrontend && process.env.PRIME_AGENT_INTERNAL_LEGACY_OWNED_WORKER_FRONTEND !== "1") { + return false; + } const profile = classifyOwnedSessionWorkerInvocation(args, process.stdin.isTTY); if (!profile) { return false; @@ -482,7 +492,7 @@ export async function maybeRunOwnedSessionWorkerFrontend(args: readonly string[] } export function installOwnedSessionWorkerOwnerWatch(): void { - if (process.env[OWNED_WORKER_ENV] !== "1") { + if (!isOwnedSessionWorkerProcess()) { return; } if (!process.channel) { diff --git a/packages/coding-agent/src/cli/subprocess-launch.ts b/packages/coding-agent/src/cli/subprocess-launch.ts index 1a242148630..f5843b48d17 100644 --- a/packages/coding-agent/src/cli/subprocess-launch.ts +++ b/packages/coding-agent/src/cli/subprocess-launch.ts @@ -1,3 +1,5 @@ +import { existsSync } from "node:fs"; +import { dirname, isAbsolute, join, resolve } from "node:path"; import { isBunBinary } from "../config.js"; export interface CliSubprocessLaunchSpec { @@ -5,6 +7,30 @@ export interface CliSubprocessLaunchSpec { args: string[]; } +export function createCliSubprocessEnv( + source: NodeJS.ProcessEnv = process.env, + entrypoint = process.argv[1], + execArgs: readonly string[] = process.execArgv, +): NodeJS.ProcessEnv { + const environment = { ...source }; + if (environment.TSX_TSCONFIG_PATH !== undefined || !entrypoint || !execArgs.some((arg) => arg.includes("tsx"))) { + return environment; + } + let directory = dirname(resolve(entrypoint)); + while (true) { + const tsconfigPath = join(directory, "tsconfig.json"); + if (existsSync(tsconfigPath) && existsSync(join(directory, "node_modules", "tsx", "package.json"))) { + environment.TSX_TSCONFIG_PATH = tsconfigPath; + return environment; + } + const parent = dirname(directory); + if (parent === directory) { + return environment; + } + directory = parent; + } +} + function quoteCommandArgument(value: string): string { return /^[A-Za-z0-9_./:@%+=,-]+$/.test(value) ? value : `'${value.replaceAll("'", `'"'"'`)}'`; } @@ -30,5 +56,6 @@ export function createCliSubprocessLaunchSpec( if (!entrypoint) { throw new Error("Cannot determine current CLI entrypoint for subprocess launch"); } - return { command: executable, args: [...execArgs, entrypoint, ...args] }; + const resolvedEntrypoint = isAbsolute(entrypoint) ? entrypoint : resolve(entrypoint); + return { command: executable, args: [...execArgs, resolvedEntrypoint, ...args] }; } diff --git a/packages/coding-agent/src/main.ts b/packages/coding-agent/src/main.ts index 9e9e9b7a92d..ea7319f7c37 100644 --- a/packages/coding-agent/src/main.ts +++ b/packages/coding-agent/src/main.ts @@ -25,7 +25,7 @@ import { confirmDaemonSessionLoss, type DaemonSessionLossCopy, pluralizeSessions import { processFileArguments } from "./cli/file-processor.js"; import { buildInitialMessage } from "./cli/initial-message.js"; import { listModels } from "./cli/list-models.js"; -import { installOwnedSessionRecoveryTracking } from "./cli/owned-session-worker.js"; +import { installOwnedSessionRecoveryTracking, isOwnedSessionWorkerProcess } from "./cli/owned-session-worker.js"; import { handlePublicCommand } from "./cli/public-command.js"; import { selectSession } from "./cli/session-picker.js"; import { expandTildePath, getAgentDir, getSessionDirEnvOverride, VERSION } from "./config.js"; @@ -67,7 +67,8 @@ import { SettingsManager } from "./core/settings-manager.js"; import { printTimings, resetTimings, time } from "./core/timings.js"; import { runMigrations, showDeprecationWarnings } from "./migrations.js"; import { isDaemonCatalogProcess, runDaemonCatalogProcess } from "./modes/daemon/daemon-catalog-process.js"; -import { collectDaemonClientEnv } from "./modes/daemon/daemon-protocol.js"; +import { deserializeDaemonError } from "./modes/daemon/daemon-errors.js"; +import { collectDaemonClientEnv, collectDaemonLaunchEnv } from "./modes/daemon/daemon-protocol.js"; import { DAEMON_WORKER_ACTIVE_SESSION_ID_ENV, isDaemonWorkerProcess, @@ -80,6 +81,7 @@ import { createInteractiveModeLocalSessionHost, createInteractiveModeUiServicesFromServices, DaemonAgentConnection, + DaemonCapabilityUnavailableError, DaemonClient, defaultDaemonSocketPath, InProcessAgentConnection, @@ -89,7 +91,9 @@ import { runDaemonMode, runDaemonSupervisorMode, runPrintMode, + runPrintModeWithConnection, runRpcMode, + runRpcModeWithConnection, type SessionSummary, } from "./modes/index.js"; import { ExtensionSelectorComponent } from "./modes/interactive/components/extension-selector.js"; @@ -144,7 +148,9 @@ function isTruthyEnvFlag(value: string | undefined): boolean { return value === "1" || value.toLowerCase() === "true" || value.toLowerCase() === "yes"; } -export type AppMode = "interactive" | "print" | "json" | "rpc" | "daemon"; +export type ClientMode = "interactive" | "print" | "json" | "rpc"; +/** Compatibility view of the CLI's internal daemon process entrypoint. */ +export type AppMode = ClientMode | "daemon"; export function shouldRejectNonInteractiveAttach(attachAgent: string | undefined, appMode: AppMode): boolean { return attachAgent !== undefined && appMode !== "interactive"; @@ -178,14 +184,18 @@ export function parseAgentsViewCommand(args: string[]): { explicitAgentsView: bo return { explicitAgentsView: false, args }; } -export interface InteractiveDaemonStartupDecision { +export interface DaemonClientStartupDecision { appMode: AppMode; startupBenchmark: boolean; noSession?: boolean; + help?: boolean; listModels?: string | true; } -export function shouldUseDaemonInteractive(options: InteractiveDaemonStartupDecision): boolean { +export type InteractiveDaemonStartupDecision = DaemonClientStartupDecision; + +/** Retained for callers that only classify persistent interactive startup. */ +export function shouldUseDaemonInteractive(options: DaemonClientStartupDecision): boolean { return ( options.appMode === "interactive" && !options.startupBenchmark && @@ -194,6 +204,21 @@ export function shouldUseDaemonInteractive(options: InteractiveDaemonStartupDeci ); } +export function shouldUseDaemonClient(options: DaemonClientStartupDecision): boolean { + return ( + options.appMode !== "daemon" && !options.startupBenchmark && !options.help && options.listModels === undefined + ); +} + +export function shouldUseDaemonClientRuntime( + options: DaemonClientStartupDecision & { + ownedSessionWorker?: boolean; + hasProcessLocalExtensionFactories?: boolean; + }, +): boolean { + return shouldUseDaemonClient(options) && !options.ownedSessionWorker && !options.hasProcessLocalExtensionFactories; +} + export function shouldEnsureInteractiveDaemonForStartup( useDaemonInteractive: boolean, attachAgent: string | undefined, @@ -963,12 +988,15 @@ export function findActiveDaemonSessionSummaryForSessionFile( ); } -async function createDaemonInteractiveConnection(options: { +async function createDaemonClientConnection(options: { socketPath: string; config: AgentSessionRuntimeConfig; sessionPath?: string; continueRecent?: boolean; activeSessionId?: string; + clientOwned?: boolean; + noSession?: boolean; + supportsExtensionUi?: boolean; }): Promise<{ connection: DaemonAgentConnection; summary: SessionSummary }> { // Caller must have awaited ensureInteractiveDaemonRunning for this socket. const client = new DaemonClient(options.socketPath); @@ -979,6 +1007,8 @@ async function createDaemonInteractiveConnection(options: { const connection = await DaemonAgentConnection.attach(client, getDaemonSummaryActiveSessionId(summary), { closeClientOnDispose: true, sendClientEnv: true, + ownedSession: options.clientOwned, + supportsExtensionUi: options.supportsExtensionUi, recoverDaemon: () => ensureInteractiveDaemonRunning(options.socketPath), }); return { connection, summary }; @@ -989,7 +1019,7 @@ async function createDaemonInteractiveConnection(options: { return await attach(summary); } - if (options.sessionPath) { + if (options.sessionPath && !options.clientOwned) { const activeSummary = findActiveDaemonSessionSummaryForSessionFile( await listActiveDaemonSessionSummaries(client), options.sessionPath, @@ -998,16 +1028,25 @@ async function createDaemonInteractiveConnection(options: { return await attach(activeSummary); } } + if (options.clientOwned) { + await client.waitForHello(); + if (!client.supportsServerCapability("client_owned_sessions")) { + throw new DaemonCapabilityUnavailableError("create", "client_owned_sessions"); + } + } const response = await client.request({ type: "create", config: options.config, sessionPath: options.sessionPath, continueRecent: options.continueRecent, + noSession: options.noSession, env: collectDaemonClientEnv(), + lifecycle: options.clientOwned ? "client_owned" : "resident", + launchEnv: options.clientOwned ? collectDaemonLaunchEnv() : undefined, }); if (!response.success) { - throw new Error(response.error); + throw deserializeDaemonError(response); } if (!isDaemonSessionSummary(response.data)) { throw new Error("Daemon returned an invalid create response"); @@ -1143,12 +1182,17 @@ export async function main(args: string[], options?: MainOptions) { console.error(chalk.red("Error: PI_STARTUP_BENCHMARK only supports interactive mode")); process.exit(1); } - const useDaemonInteractive = shouldUseDaemonInteractive({ + // Programmatic factories are process-local functions and cannot be serialized to a daemon worker. + const hasProcessLocalExtensionFactories = (options?.extensionFactories?.length ?? 0) > 0; + const useDaemonClient = shouldUseDaemonClientRuntime({ appMode, startupBenchmark, noSession: parsed.noSession, listModels: parsed.listModels, + ownedSessionWorker: isOwnedSessionWorkerProcess(), + hasProcessLocalExtensionFactories, }); + const useDaemonInteractive = useDaemonClient && appMode === "interactive"; // Decide the final runtime cwd before creating cwd-bound runtime services. // --resume may select a session from another project, so project-local @@ -1162,7 +1206,7 @@ export async function main(args: string[], options?: MainOptions) { const daemonSocketPath = parsed.daemonSocket ?? defaultDaemonSocketPath(); // Kick off daemon spawn/readiness immediately so it overlaps session-manager // and runtime-services preparation; attach only connects to an existing daemon. - let daemonReady = shouldEnsureInteractiveDaemonForStartup(useDaemonInteractive, publicCommand.attachAgent) + let daemonReady = shouldEnsureInteractiveDaemonForStartup(useDaemonClient, publicCommand.attachAgent) ? ensureInteractiveDaemonRunning(daemonSocketPath) : undefined; // Errors are rethrown at the await sites below; this only avoids an unhandled @@ -1377,7 +1421,7 @@ export async function main(args: string[], options?: MainOptions) { }; if ( shouldOpenAgentsViewForDaemonInteractive({ - useDaemonInteractive, + useDaemonInteractive: useDaemonInteractive && !parsed.noSession, explicitAgentsView, needsOnboarding: shouldRunOnboarding({ settingsManager, @@ -1402,13 +1446,16 @@ export async function main(args: string[], options?: MainOptions) { // no DeferredAgentConnection is needed to avoid creating it up front. const isFreshDefaultSession = !activeDaemonSessionSummary && !getInteractiveDaemonSessionPath(parsed, sessionManager); - const { connection, summary } = await createDaemonInteractiveConnection({ + const { connection, summary } = await createDaemonClientConnection({ socketPath: daemonSocketPath, config: defaultSessionConfig, activeSessionId: activeDaemonSessionSummary ? getDaemonSummaryActiveSessionId(activeDaemonSessionSummary) : undefined, sessionPath: getInteractiveDaemonSessionPath(parsed, sessionManager), + clientOwned: parsed.noSession, + noSession: parsed.noSession, + supportsExtensionUi: true, }); const agentConnection: AgentConnection = connection; const attachModelFallbackMessage = isFreshDefaultSession @@ -1432,15 +1479,83 @@ export async function main(args: string[], options?: MainOptions) { // arrow takes them to the agents view like any other session. The agents // view was not rendered here, so we intentionally leave // agentsViewOwnsStartupNotices unset and let the in-session fallback run. - returnToAgentsView: true, + returnToAgentsView: !parsed.noSession, }); await preloadCodeHighlighter(); printTimings(); await interactiveMode.run(); + if (parsed.noSession) { + return; + } await launchAgentsView(false); return; } + if (useDaemonClient) { + const settingsManager = SettingsManager.create(sessionManager.getCwd(), agentDir); + let stdinContent: string | undefined; + if (appMode !== "rpc") { + stdinContent = await readPipedStdin(); + } + time("readPipedStdin"); + const { initialMessage, initialImages } = await prepareInitialMessage( + parsed, + settingsManager.getImageAutoResize(), + stdinContent, + ); + time("prepareInitialMessage"); + initTheme(settingsManager.getTheme(), false); + time("initTheme"); + + daemonReady = (await awaitDaemonReady(daemonReady)).ready; + let connection: DaemonAgentConnection; + let summary: SessionSummary; + try { + ({ connection, summary } = await createDaemonClientConnection({ + socketPath: daemonSocketPath, + config: defaultSessionConfig, + sessionPath: parsed.noSession ? undefined : sessionManager.getSessionFile(), + continueRecent: parsed.continue, + clientOwned: true, + noSession: parsed.noSession, + supportsExtensionUi: appMode === "rpc", + })); + } catch (error) { + if (error instanceof SessionAlreadyActiveError) { + console.error(chalk.red(`Error: ${error.message}`)); + process.exit(1); + } + throw error; + } + const diagnostics = summary.diagnostics ?? []; + reportDiagnostics(diagnostics); + if (diagnostics.some((diagnostic) => diagnostic.type === "error")) { + await connection.dispose(); + process.exit(1); + } + if (!summary.model) { + console.error(chalk.red(summary.modelFallbackMessage ?? formatNoModelsAvailableMessage())); + await connection.dispose(); + process.exit(1); + } + + printTimings(); + if (appMode === "rpc") { + return await runRpcModeWithConnection(connection); + } + const exitCode = await runPrintModeWithConnection(connection, { + mode: toPrintOutputMode(appMode), + messages: parsed.messages, + initialMessage, + initialImages, + }); + stopThemeWatcher(); + restoreStdout(); + if (exitCode !== 0) { + process.exitCode = exitCode; + } + return; + } let runtime: AgentSessionRuntime; try { diff --git a/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts b/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts index 5a4519992f0..00e31212b61 100644 --- a/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts +++ b/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts @@ -2,7 +2,14 @@ import { randomUUID } from "node:crypto"; import type { AgentMessage, ThinkingLevel } from "@earendil-works/pi-agent-core"; import type { ImageContent, ServiceTier, Transport } from "@earendil-works/pi-ai"; import { appendRotatingLog, getAgentLogPath, getDaemonLogPath } from "../../config.js"; +import type { + AgentSessionMessageDeliveryMode, + AgentSessionMessageReceipt, + AgentSessionMessageSafetyStatus, +} from "../../core/agent-messages.js"; import type { AgentSessionEvent } from "../../core/agent-session.js"; +import type { AgentAutonomousStatus } from "../../core/autonomous.js"; +import type { BashResult } from "../../core/bash-executor.js"; import type { CompactionResult } from "../../core/compaction/index.js"; import type { ContextTreeNode } from "../../core/context-tree.js"; import type { @@ -19,6 +26,7 @@ import { type DaemonClient, getDaemonSocketCloseReason } from "../daemon/daemon- import { deserializeDaemonError } from "../daemon/daemon-errors.js"; import { collectDaemonClientEnv, + collectDaemonLaunchEnv, type DaemonAttachResult, type DaemonCommand, type DaemonEventCursor, @@ -56,6 +64,7 @@ import type { AgentConnectionSavedSessionScope, AgentConnectionScopedModel, AgentConnectionSessionContext, + AgentConnectionSessionHeader, AgentConnectionSessionListCallbacks, AgentConnectionSessionTreeNode, AgentConnectionSessionWatcher, @@ -82,12 +91,14 @@ interface DaemonSnapshotAssembly { } export const DAEMON_REFINE_REQUEST_TIMEOUT_MS = 10 * 60 * 1000; +const DAEMON_LONG_RUNNING_REQUEST_TIMEOUT_MS = 24 * 60 * 60 * 1000; export const DAEMON_RECONNECT_TIMEOUT_MS = 60_000; export const DAEMON_SNAPSHOT_TIMEOUT_MS = 30_000; const MAX_IGNORED_SNAPSHOT_IDS = 128; const UPDATE_RECONNECT_TIMEOUT_MS = 120000; const UPDATE_RECONNECT_RETRY_MS = 100; const MAX_COMPLETED_SNAPSHOTS = 128; +const OWNED_SESSION_DISPOSE_RECONNECT_WAIT_MS = 10_000; const updateTransportReconnects = new WeakMap>(); function delay(ms: number): Promise { @@ -147,6 +158,10 @@ export interface DaemonAgentConnectionOptions { * rebinds, so watchers must not send env at all. */ sendClientEnv?: boolean; + /** Advertise support for interactive extension dialogs. */ + supportsExtensionUi?: boolean; + /** Dispose the connection by stopping its hidden worker instead of detaching. */ + ownedSession?: boolean; } /** @@ -160,6 +175,7 @@ export class DaemonAgentConnection implements AgentConnection { private readonly unsubscribeDaemonMessages: () => void; private readonly unsubscribeDaemonClose: () => void; private readonly clientId = `daemon-agent-connection:${randomUUID()}`; + private ownedSessionPromotionTail = Promise.resolve(); private lastEventCursor: DaemonEventCursor | undefined; private readonly retiredEventGenerations = new Set(); private lastEventSequence: number | undefined; @@ -179,6 +195,7 @@ export class DaemonAgentConnection implements AgentConnection { private readonly snapshotRecoveryPromises = new Map>(); private readonly ignoredSnapshotIds = new Set(); private reconnectPromise?: Promise; + private disposing = false; private disposed = false; constructor( @@ -243,13 +260,22 @@ export class DaemonAgentConnection implements AgentConnection { } async attach(): Promise { + const supportsExtensionUi = this.options.supportsExtensionUi !== false; const result = await this.requestData({ type: "attach", activeSessionId: this.activeSessionId, - supportsExtensionUi: true, + supportsExtensionUi, clientId: this.clientId, - capabilities: ["attach_snapshot", "event_sequence", "extension_ui", "slim_attach", "chunked_snapshot"], + capabilities: [ + "attach_snapshot", + "event_sequence", + ...(supportsExtensionUi ? (["extension_ui"] as const) : []), + "slim_attach", + "chunked_snapshot", + ...(this.options.ownedSession ? (["client_owned_sessions"] as const) : []), + ], env: this.options.sendClientEnv ? collectDaemonClientEnv() : undefined, + launchEnv: this.options.ownedSession ? collectDaemonLaunchEnv() : undefined, resumeCursor: this.lastEventCursor === undefined ? undefined @@ -370,6 +396,14 @@ export class DaemonAgentConnection implements AgentConnection { return data.messages; } + async getSessionHeader(): Promise { + const data = await this.requestData<{ header?: AgentConnectionSessionHeader | null }>({ + type: "get_session_header", + activeSessionId: this.activeSessionId, + }); + return data.header ?? undefined; + } + async getCommands(): Promise { const data = await this.requestData<{ commands: AgentConnectionSlashCommand[] }>({ type: "get_commands", @@ -478,7 +512,10 @@ export class DaemonAgentConnection implements AgentConnection { return []; } try { - const command = { type: "heartbeats_list" } as const; + const command = { + type: "heartbeats_list", + ...(this.options.ownedSession ? { activeSessionId: this.activeSessionId } : {}), + } as const; const data = hasCapability ? await this.requestData<{ heartbeats: AgentConnectionHeartbeat[] }>(command) : await this.requestLegacyData<{ heartbeats: AgentConnectionHeartbeat[] }>(command); @@ -520,17 +557,24 @@ export class DaemonAgentConnection implements AgentConnection { } async addCronJob(schedule: string, prompt: string): Promise { - const data = await this.requestData<{ job: AgentCronJob }>({ - type: "cron_add", - activeSessionId: this.activeSessionId, - schedule, - prompt, + return this.withOwnedSessionPromotion(async (promoteOwnedSession) => { + const data = await this.requestData<{ job: AgentCronJob }>({ + type: "cron_add", + activeSessionId: this.activeSessionId, + schedule, + prompt, + promoteOwnedSession, + }); + return data.job; }); - return data.job; } async cancelCronJob(jobId: string): Promise { - const data = await this.requestData<{ job: AgentCronJob }>({ type: "cron_cancel", jobId }); + const data = await this.requestData<{ job: AgentCronJob }>({ + type: "cron_cancel", + activeSessionId: this.activeSessionId, + jobId, + }); return data.job; } @@ -547,14 +591,17 @@ export class DaemonAgentConnection implements AgentConnection { instruction: string, deliveryMode?: AgentHeartbeatDeliveryMode, ): Promise { - const data = await this.requestData<{ heartbeat: AgentCronJob }>({ - type: "heartbeat_set", - activeSessionId: this.activeSessionId, - schedule, - prompt: instruction, - ...(deliveryMode ? { deliveryMode } : {}), + return this.withOwnedSessionPromotion(async (promoteOwnedSession) => { + const data = await this.requestData<{ heartbeat: AgentCronJob }>({ + type: "heartbeat_set", + activeSessionId: this.activeSessionId, + schedule, + prompt: instruction, + ...(deliveryMode ? { deliveryMode } : {}), + promoteOwnedSession, + }); + return data.heartbeat; }); - return data.heartbeat; } async updateHeartbeat(action: AgentHeartbeatUpdateAction): Promise { @@ -566,6 +613,48 @@ export class DaemonAgentConnection implements AgentConnection { return data.heartbeat ?? undefined; } + async sendAgentMessage( + targetActiveSessionId: string, + message: string, + deliveryMode?: AgentSessionMessageDeliveryMode, + ): Promise { + return this.requestData({ + type: "send_message", + targetActiveSessionId, + message, + fromActiveSessionId: this.activeSessionId, + deliveryMode, + }); + } + + async getAgentMessageStatus(): Promise { + return this.requestData({ + type: "agent_messages_status", + activeSessionId: this.activeSessionId, + }); + } + + async pauseAgentMessages(): Promise { + return this.requestData({ + type: "agent_messages_pause", + activeSessionId: this.activeSessionId, + }); + } + + async resumeAgentMessages(): Promise { + return this.requestData({ + type: "agent_messages_resume", + activeSessionId: this.activeSessionId, + }); + } + + async clearAgentMessages(): Promise { + return this.requestData({ + type: "agent_messages_clear", + activeSessionId: this.activeSessionId, + }); + } + async getUserMessagesForForking(): Promise { const data = await this.requestData<{ messages: AgentConnectionUserMessage[] }>({ type: "get_user_messages_for_forking", @@ -624,9 +713,24 @@ export class DaemonAgentConnection implements AgentConnection { message, images: options?.images, streamingBehavior: options?.streamingBehavior, + source: options?.source, }); } + async promptAndWait(message: string, options?: AgentConnectionPromptOptions): Promise { + await this.requestData( + { + type: "prompt_and_wait", + activeSessionId: this.activeSessionId, + message, + images: options?.images, + streamingBehavior: options?.streamingBehavior, + source: options?.source, + }, + DAEMON_LONG_RUNNING_REQUEST_TIMEOUT_MS, + ); + } + async startSideQuestion(id: string, question: string): Promise { this.activeSideQuestionIds.add(id); try { @@ -684,7 +788,20 @@ export class DaemonAgentConnection implements AgentConnection { } async waitForIdle(): Promise { - await this.requestOk({ type: "wait_for_idle", activeSessionId: this.activeSessionId }); + await this.requestData( + { type: "wait_for_idle", activeSessionId: this.activeSessionId }, + DAEMON_LONG_RUNNING_REQUEST_TIMEOUT_MS, + ); + } + + async waitForHeadlessCompletion(): Promise { + return this.requestData( + { + type: "wait_for_headless_completion", + activeSessionId: this.activeSessionId, + }, + DAEMON_LONG_RUNNING_REQUEST_TIMEOUT_MS, + ); } async executeBash(command: string, options?: AgentConnectionExecuteBashOptions): Promise { @@ -703,6 +820,17 @@ export class DaemonAgentConnection implements AgentConnection { } } + async executeBashAndWait(command: string): Promise { + return this.requestData( + { + type: "execute_bash_and_wait", + activeSessionId: this.activeSessionId, + command, + }, + DAEMON_LONG_RUNNING_REQUEST_TIMEOUT_MS, + ); + } + async abortBash(): Promise { try { await this.requestOk({ type: "abort_bash", activeSessionId: this.activeSessionId }); @@ -772,6 +900,10 @@ export class DaemonAgentConnection implements AgentConnection { await this.requestOk({ type: "set_auto_compaction", activeSessionId: this.activeSessionId, enabled }); } + async setAutoRetryEnabled(enabled: boolean): Promise { + await this.requestOk({ type: "set_auto_retry", activeSessionId: this.activeSessionId, enabled }); + } + async compact(customInstructions?: string): Promise { return this.requestData({ type: "compact", @@ -841,6 +973,9 @@ export class DaemonAgentConnection implements AgentConnection { if (!(error instanceof SessionAlreadyActiveError) || !error.activeSessionId) { throw error; } + if (this.options.ownedSession) { + throw error; + } if (error.activeSessionId === sourceActiveSessionId) { return { cancelled: false }; } @@ -868,14 +1003,23 @@ export class DaemonAgentConnection implements AgentConnection { this.pendingReattachActiveSessionIds.add(targetActiveSessionId); let reattached = false; try { + const supportsExtensionUi = this.options.supportsExtensionUi !== false; const result = await this.requestData({ type: "reattach", activeSessionId: sourceActiveSessionId, targetActiveSessionId, - supportsExtensionUi: true, + supportsExtensionUi, clientId: this.clientId, - capabilities: ["attach_snapshot", "event_sequence", "extension_ui", "slim_attach", "chunked_snapshot"], + capabilities: [ + "attach_snapshot", + "event_sequence", + ...(supportsExtensionUi ? (["extension_ui"] as const) : []), + "slim_attach", + "chunked_snapshot", + ...(this.options.ownedSession ? (["client_owned_sessions"] as const) : []), + ], env: this.options.sendClientEnv ? collectDaemonClientEnv() : undefined, + launchEnv: this.options.ownedSession ? collectDaemonLaunchEnv() : undefined, }); reattached = true; this.activeSessionId = result.activeSessionId; @@ -1000,18 +1144,56 @@ export class DaemonAgentConnection implements AgentConnection { } async dispose(): Promise { + if (this.disposed || this.disposing) { + return; + } + this.disposing = true; + if (this.options.ownedSession && !this.client.isConnected && this.reconnectPromise) { + await Promise.race([this.reconnectPromise, delay(OWNED_SESSION_DISPOSE_RECONNECT_WAIT_MS)]).catch( + () => undefined, + ); + } this.disposed = true; this.updateRestartPending = false; await Promise.allSettled([...this.activeSideQuestionIds].map((id) => this.abortSideQuestion(id))); this.unsubscribeDaemonMessages(); this.unsubscribeDaemonClose(); - await this.requestOk({ type: "detach", activeSessionId: this.activeSessionId }).catch(() => undefined); + if (this.options.ownedSession) { + await this.requestOk({ type: "complete_owned_session", activeSessionId: this.activeSessionId }).catch( + () => undefined, + ); + } else { + await this.requestOk({ type: "detach", activeSessionId: this.activeSessionId }).catch(() => undefined); + } if (this.options.closeClientOnDispose) { this.client.close(); } this.rejectSnapshotAssemblies(new Error("Daemon connection disposed during snapshot transfer")); } + async promoteToResident(): Promise { + await this.withOwnedSessionPromotion(async (promoteOwnedSession) => { + if (!promoteOwnedSession) return; + await this.requestOk({ type: "promote_owned_session", activeSessionId: this.activeSessionId }); + }); + } + + private withOwnedSessionPromotion(operation: (promoteOwnedSession: boolean) => Promise): Promise { + const run = this.ownedSessionPromotionTail.then(async () => { + const promoteOwnedSession = this.options.ownedSession === true; + const result = await operation(promoteOwnedSession); + if (promoteOwnedSession) { + this.options.ownedSession = false; + } + return result; + }); + this.ownedSessionPromotionTail = run.then( + () => undefined, + () => undefined, + ); + return run; + } + private async reconnect(cause: Error): Promise { if (this.reconnectPromise) { return this.reconnectPromise; @@ -1207,6 +1389,15 @@ export class DaemonAgentConnection implements AgentConnection { }); return; } + if (message.type === "extension_error") { + await this.emit({ + type: "extension_error", + extensionPath: message.extensionPath, + event: message.event, + error: message.error, + }); + return; + } if (message.type === "session_closed") { if (message.reason === "update") { this.captureDaemonLogPath(); diff --git a/packages/coding-agent/src/modes/agent-connection/in-process-agent-connection.ts b/packages/coding-agent/src/modes/agent-connection/in-process-agent-connection.ts index 596ac593bfa..3157bc1966b 100644 --- a/packages/coding-agent/src/modes/agent-connection/in-process-agent-connection.ts +++ b/packages/coding-agent/src/modes/agent-connection/in-process-agent-connection.ts @@ -1,7 +1,14 @@ import { resolve } from "node:path"; import type { AgentMessage, ThinkingLevel } from "@earendil-works/pi-agent-core"; import type { ImageContent, ServiceTier, Transport } from "@earendil-works/pi-ai"; +import type { + AgentSessionMessageDeliveryMode, + AgentSessionMessageReceipt, + AgentSessionMessageSafetyStatus, +} from "../../core/agent-messages.js"; import type { AgentSessionRuntime } from "../../core/agent-session-runtime.js"; +import type { AgentAutonomousStatus } from "../../core/autonomous.js"; +import type { BashResult } from "../../core/bash-executor.js"; import type { CompactionResult } from "../../core/compaction/index.js"; import type { ContextTreeNode } from "../../core/context-tree.js"; import type { @@ -10,11 +17,13 @@ import type { AgentHeartbeatManagementAction, AgentHeartbeatUpdateAction, } from "../../core/cron-jobs.js"; +import type { ExtensionUIContext } from "../../core/extensions/types.js"; import type { RefinementResult } from "../../core/refinement/index.js"; import { type DeleteSessionFileResult, deleteSessionFile } from "../../core/session-file-actions.js"; import { SessionManager } from "../../core/session-manager.js"; import type { SessionStats } from "../../core/session-stats.js"; import { type SideQuestionRun, startSideQuestion } from "../../core/side-question.js"; +import { waitForHeadlessCompletion } from "../headless-completion.js"; import { createAgentConnectionCommands, createAgentConnectionResourceSnapshot, @@ -44,6 +53,7 @@ import type { AgentConnectionSavedSessionScope, AgentConnectionScopedModel, AgentConnectionSessionContext, + AgentConnectionSessionHeader, AgentConnectionSessionListCallbacks, AgentConnectionSessionTreeNode, AgentConnectionSessionWatcher, @@ -55,22 +65,33 @@ import type { AgentConnectionUserMessage, } from "./types.js"; +export interface InProcessHeadlessExtensionOptions { + uiContext?: ExtensionUIContext; + shutdownHandler?: () => void; +} + export class InProcessAgentConnection implements AgentConnection { private readonly listeners = new Set(); private readonly beforeSessionInvalidateListeners = new Set(); private readonly sideQuestionRuns = new Map(); + private headlessExtensionOptions: InProcessHeadlessExtensionOptions | undefined; private unsubscribeSessionEvents: (() => void) | undefined; constructor(private readonly runtimeHost: AgentSessionRuntime) { this.bindCurrentSessionEvents(); - this.runtimeHost.setBeforeSessionInvalidate(() => { - this.abortAllSideQuestions(); - for (const listener of [...this.beforeSessionInvalidateListeners]) { - listener(); - } - }); + if (typeof this.runtimeHost.setBeforeSessionInvalidate === "function") { + this.runtimeHost.setBeforeSessionInvalidate(() => { + this.abortAllSideQuestions(); + for (const listener of [...this.beforeSessionInvalidateListeners]) { + listener(); + } + }); + } this.runtimeHost.setRebindSession(async () => { this.bindCurrentSessionEvents(); + if (this.headlessExtensionOptions) { + await this.bindCurrentSessionExtensions(); + } await this.emit({ type: "session_replaced", state: createAgentConnectionState(this.runtimeHost), @@ -79,6 +100,11 @@ export class InProcessAgentConnection implements AgentConnection { }); } + async bindHeadlessExtensions(options: InProcessHeadlessExtensionOptions = {}): Promise { + this.headlessExtensionOptions = options; + await this.bindCurrentSessionExtensions(); + } + subscribe(listener: AgentConnectionEventListener): () => void { this.listeners.add(listener); return () => { @@ -102,7 +128,11 @@ export class InProcessAgentConnection implements AgentConnection { } async getMessages(): Promise { - return this.session.messages; + return this.session.state.messages; + } + + async getSessionHeader(): Promise { + return this.session.sessionManager.getHeader() ?? undefined; } async getCommands(): Promise { @@ -207,6 +237,30 @@ export class InProcessAgentConnection implements AgentConnection { throw new Error("Heartbeats require daemon mode"); } + async sendAgentMessage( + _targetActiveSessionId: string, + _message: string, + _deliveryMode?: AgentSessionMessageDeliveryMode, + ): Promise { + throw new Error("Agent messaging requires daemon mode"); + } + + async getAgentMessageStatus(): Promise { + throw new Error("Agent messaging requires daemon mode"); + } + + async pauseAgentMessages(): Promise { + throw new Error("Agent messaging requires daemon mode"); + } + + async resumeAgentMessages(): Promise { + throw new Error("Agent messaging requires daemon mode"); + } + + async clearAgentMessages(): Promise { + throw new Error("Agent messaging requires daemon mode"); + } + async getUserMessagesForForking(): Promise { return this.session.getUserMessagesForForking(); } @@ -232,9 +286,39 @@ export class InProcessAgentConnection implements AgentConnection { } async prompt(message: string, options?: AgentConnectionPromptOptions): Promise { + await new Promise((resolve, reject) => { + let settled = false; + const resolveOnce = () => { + if (settled) return; + settled = true; + resolve(); + }; + const rejectOnce = (error: unknown) => { + if (settled) return; + settled = true; + reject(error); + }; + const prompt = this.session.prompt(message, { + images: options?.images, + streamingBehavior: options?.streamingBehavior, + source: options?.source, + preflightResult: (success) => { + if (success) resolveOnce(); + }, + }); + void prompt.then(resolveOnce, rejectOnce); + }); + } + + async promptAndWait(message: string, options?: AgentConnectionPromptOptions): Promise { + if (!options) { + await this.session.prompt(message); + return; + } await this.session.prompt(message, { - images: options?.images, - streamingBehavior: options?.streamingBehavior, + images: options.images, + streamingBehavior: options.streamingBehavior, + source: options.source, }); } @@ -281,10 +365,18 @@ export class InProcessAgentConnection implements AgentConnection { await this.session.agent.waitForIdle(); } + async waitForHeadlessCompletion(): Promise { + return waitForHeadlessCompletion(this.session); + } + async executeBash(command: string, options?: AgentConnectionExecuteBashOptions): Promise { await this.session.runUserBash(command, options); } + async executeBashAndWait(command: string): Promise { + return this.session.executeBash(command); + } + async abortBash(): Promise { this.session.abortBash(); } @@ -340,6 +432,10 @@ export class InProcessAgentConnection implements AgentConnection { this.session.setAutoCompactionEnabled(enabled); } + async setAutoRetryEnabled(enabled: boolean): Promise { + this.session.setAutoRetryEnabled(enabled); + } + async compact(customInstructions?: string): Promise { return this.session.compact(customInstructions); } @@ -458,7 +554,9 @@ export class InProcessAgentConnection implements AgentConnection { this.abortAllSideQuestions(); this.unsubscribeSessionEvents?.(); this.unsubscribeSessionEvents = undefined; - this.runtimeHost.setBeforeSessionInvalidate(undefined); + if (typeof this.runtimeHost.setBeforeSessionInvalidate === "function") { + this.runtimeHost.setBeforeSessionInvalidate(undefined); + } this.runtimeHost.setRebindSession(undefined); await this.runtimeHost.dispose(); } @@ -474,6 +572,36 @@ export class InProcessAgentConnection implements AgentConnection { }); } + private async bindCurrentSessionExtensions(): Promise { + const session = this.session; + await session.bindExtensions({ + uiContext: this.headlessExtensionOptions?.uiContext, + commandContextActions: { + waitForIdle: () => session.agent.waitForIdle(), + newSession: (options) => this.runtimeHost.newSession(options), + fork: async (entryId, options) => { + const result = await this.runtimeHost.fork(entryId, options); + return { cancelled: result.cancelled }; + }, + navigateTree: async (targetId, options) => { + const result = await session.navigateTree(targetId, options); + return { cancelled: result.cancelled }; + }, + switchSession: (sessionPath, options) => this.runtimeHost.switchSession(sessionPath, options), + reload: () => session.reload(), + }, + shutdownHandler: this.headlessExtensionOptions?.shutdownHandler, + onError: (error) => { + void this.emit({ + type: "extension_error", + extensionPath: error.extensionPath, + event: error.event, + error: error.error, + }); + }, + }); + } + private abortAllSideQuestions(): void { for (const run of this.sideQuestionRuns.values()) { run.abort(); diff --git a/packages/coding-agent/src/modes/agent-connection/types.ts b/packages/coding-agent/src/modes/agent-connection/types.ts index e1a9ec3e8ab..8bbd56a3e51 100644 --- a/packages/coding-agent/src/modes/agent-connection/types.ts +++ b/packages/coding-agent/src/modes/agent-connection/types.ts @@ -1,6 +1,13 @@ import type { AgentEvent, AgentMessage, ThinkingLevel } from "@earendil-works/pi-agent-core"; import type { Api, ImageContent, Model, ServiceTier, TextContent, Transport, Usage } from "@earendil-works/pi-ai"; +import type { + AgentSessionMessageDeliveryMode, + AgentSessionMessageReceipt, + AgentSessionMessageSafetyStatus, +} from "../../core/agent-messages.js"; import type { AuthSourceToken } from "../../core/auth-storage.js"; +import type { AgentAutonomousStatus } from "../../core/autonomous.js"; +import type { BashResult } from "../../core/bash-executor.js"; import type { CompactionResult } from "../../core/compaction/index.js"; import type { ContextTreeNode } from "../../core/context-tree.js"; import type { @@ -10,6 +17,7 @@ import type { AgentHeartbeatUpdateAction, } from "../../core/cron-jobs.js"; import type { ReplayBuiltInToolName } from "../../core/extensions/index.js"; +import type { InputSource } from "../../core/extensions/types.js"; import type { GoalState } from "../../core/goals.js"; import type { KernelSentAgentMessage } from "../../core/kernel/index.js"; import type { RefinementResult } from "../../core/refinement/index.js"; @@ -40,6 +48,20 @@ export type AgentConnectionQueueMode = "all" | "one-at-a-time"; export type AgentConnectionModel = Model; export type AgentConnectionSavedSessionScope = "current" | "all"; +export interface AgentConnectionSessionHeader { + type: "session"; + version?: number; + id: string; + timestamp: string; + cwd: string; + parentSession?: string; + git?: { + repoUrl?: string; + commit?: string; + branch?: string; + }; +} + export type AgentConnectionSavedSessionStateStatus = "active" | "archived" | "crash"; export type AgentConnectionSourceScope = "user" | "project" | "temporary"; @@ -402,6 +424,7 @@ export interface AgentConnectionToolDefinition { export interface AgentConnectionPromptOptions { images?: ImageContent[]; streamingBehavior?: "steer" | "followUp"; + source?: InputSource; } export interface AgentConnectionSideQuestionEvent { @@ -549,6 +572,7 @@ export type AgentConnectionEvent = | { type: "session_resynced"; snapshot: AgentConnectionSnapshot } | { type: "session_status"; recap?: string } | { type: "extension_ui_request"; request: AgentConnectionExtensionUiRequest } + | { type: "extension_error"; extensionPath: string; event: string; error: string } | { type: "connection_status"; status: "reconnecting" | "connected"; error?: string } | { type: "heartbeats_changed" } | { type: "closed"; error?: string }; @@ -563,6 +587,7 @@ export interface AgentConnection { getState(): Promise; getInitialSnapshot(): Promise; getMessages(): Promise; + getSessionHeader(): Promise; getCommands(): Promise; getResourceSnapshot(): Promise; getAvailableModels(): Promise; @@ -593,6 +618,15 @@ export interface AgentConnection { deliveryMode?: AgentHeartbeatDeliveryMode, ): Promise; updateHeartbeat(action: AgentHeartbeatUpdateAction): Promise; + sendAgentMessage( + targetActiveSessionId: string, + message: string, + deliveryMode?: AgentSessionMessageDeliveryMode, + ): Promise; + getAgentMessageStatus(): Promise; + pauseAgentMessages(): Promise; + resumeAgentMessages(): Promise; + clearAgentMessages(): Promise; getUserMessagesForForking(): Promise; getLastAssistantText(): Promise; /** The system prompt currently in effect for the model (with any per-turn extension changes). */ @@ -602,6 +636,7 @@ export interface AgentConnection { respondToExtensionUiRequest(requestId: string, response: AgentConnectionExtensionUiResponse): Promise; prompt(message: string, options?: AgentConnectionPromptOptions): Promise; + promptAndWait(message: string, options?: AgentConnectionPromptOptions): Promise; startSideQuestion(id: string, question: string): Promise; abortSideQuestion(id: string): Promise; steer(message: string, images?: ImageContent[]): Promise; @@ -610,6 +645,7 @@ export interface AgentConnection { abort(): Promise; cancelRlmChild(childId: string): Promise; waitForIdle(): Promise; + waitForHeadlessCompletion(): Promise; /** * Run a user-initiated bash command (! / !! prefix). Resolution timing is @@ -617,6 +653,7 @@ export interface AgentConnection { * bash_end session events, which reach every attached client. */ executeBash(command: string, options?: AgentConnectionExecuteBashOptions): Promise; + executeBashAndWait(command: string): Promise; abortBash(): Promise; setModel(provider: string, modelId: string): Promise; @@ -629,6 +666,7 @@ export interface AgentConnection { setSteeringMode(mode: AgentConnectionQueueMode): Promise; setFollowUpMode(mode: AgentConnectionQueueMode): Promise; setAutoCompactionEnabled(enabled: boolean): Promise; + setAutoRetryEnabled(enabled: boolean): Promise; compact(customInstructions?: string): Promise; refine(options?: { instructions?: string; rollbackId?: string; global?: boolean }): Promise; diff --git a/packages/coding-agent/src/modes/daemon/daemon-catalog-process.ts b/packages/coding-agent/src/modes/daemon/daemon-catalog-process.ts index 9b5362d29b4..27fc7e0e6b2 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-catalog-process.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-catalog-process.ts @@ -1,6 +1,6 @@ import { type ChildProcess, spawn } from "node:child_process"; import { randomUUID } from "node:crypto"; -import { createCliSubprocessLaunchSpec } from "../../cli/subprocess-launch.js"; +import { createCliSubprocessEnv, createCliSubprocessLaunchSpec } from "../../cli/subprocess-launch.js"; import type { DeleteSessionFileResult } from "../../core/session-file-actions.js"; import { deleteSessionFile } from "../../core/session-file-actions.js"; import { readSessionInfo, type SessionInfo, SessionManager } from "../../core/session-manager.js"; @@ -315,7 +315,7 @@ export class DaemonCatalogClient { const launch = createCliSubprocessLaunchSpec(["--version"]); const child = spawn(launch.command, launch.args, { cwd: process.cwd(), - env: { ...process.env, [DAEMON_CATALOG_ROLE_ENV]: "1" }, + env: createCliSubprocessEnv({ ...process.env, [DAEMON_CATALOG_ROLE_ENV]: "1" }), stdio: ["ignore", "ignore", "ignore", "ipc"], }); this.child = child; diff --git a/packages/coding-agent/src/modes/daemon/daemon-mode.ts b/packages/coding-agent/src/modes/daemon/daemon-mode.ts index d4d7b860a19..8314b1bbb94 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-mode.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-mode.ts @@ -23,7 +23,7 @@ import { import { createConnection, createServer, type Server, type Socket } from "node:net"; import { dirname, join, resolve } from "node:path"; import { getLogger } from "@earendil-works/pi-ai"; -import { createCliSubprocessLaunchSpec } from "../../cli/subprocess-launch.js"; +import { createCliSubprocessEnv, createCliSubprocessLaunchSpec } from "../../cli/subprocess-launch.js"; import { appendRotatingLog, getCronJobsPath, @@ -108,6 +108,7 @@ import { } from "../agent-connection/snapshot.js"; import { createAgentConnectionToolDefinition } from "../agent-connection/tool-definition.js"; import type { AgentConnectionHeartbeat } from "../agent-connection/types.js"; +import { waitForHeadlessCompletion } from "../headless-completion.js"; import { attachJsonlLineReader, serializeJsonLine } from "../rpc/jsonl.js"; import { encodePrivateFrame, PrivateFrameDecoder } from "../session-worker/private-framing.js"; import { @@ -208,6 +209,7 @@ const DAEMON_COMMAND_TYPES: ReadonlySet = new Set([ "kill", "rename", "prompt", + "prompt_and_wait", "steer", "follow_up", "restore_next_turn", @@ -222,9 +224,12 @@ const DAEMON_COMMAND_TYPES: ReadonlySet = new Set([ "start_side_question", "abort_side_question", "execute_bash", + "execute_bash_and_wait", "abort_bash", "cancel_rlm_child", "wait_for_idle", + "wait_for_headless_completion", + "get_session_header", "get_state", "get_connection_state", "get_messages", @@ -254,6 +259,7 @@ const DAEMON_COMMAND_TYPES: ReadonlySet = new Set([ "set_steering_mode", "set_follow_up_mode", "set_auto_compaction", + "set_auto_retry", "compact", "refine", "abort_compaction", @@ -459,6 +465,11 @@ export class AgentDaemon { } async start(): Promise { + if (this.options.worker) { + process.stderr.on("error", () => { + // A detached worker must survive the supervisor closing its diagnostic pipe. + }); + } this.installCrashHandlers(); await prepareDaemonSocketPath(this.socketPath); @@ -674,7 +685,7 @@ export class AgentDaemon { return; } const launch = createCliSubprocessLaunchSpec(["--mode", "daemon", "--daemon-socket", supervisorSocketPath]); - const environment = { ...process.env }; + const environment = createCliSubprocessEnv(); delete environment[DAEMON_WORKER_ROLE_ENV]; delete environment[DAEMON_WORKER_TOKEN_ENV]; delete environment[DAEMON_WORKER_ACTIVE_SESSION_ID_ENV]; @@ -1010,9 +1021,11 @@ export class AgentDaemon { sessionLease = acquireSessionLease(sessionPath, agentDir); sessionManager = sessionPath ? await SessionManager.openAsync(sessionPath, config.sessionDir, cwdOverride) - : command.continueRecent - ? SessionManager.continueRecent(cwd, config.sessionDir) - : SessionManager.create(cwd, config.sessionDir); + : command.noSession + ? SessionManager.inMemory(cwd) + : command.continueRecent + ? SessionManager.continueRecent(cwd, config.sessionDir) + : SessionManager.create(cwd, config.sessionDir); } catch (error) { sessionLease?.release(); throw error; @@ -2813,7 +2826,7 @@ export class AgentDaemon { agentMessageId: command.agentMessageId, customMessage: command.customMessage, skipInputHandlers: command.expandPromptTemplates === false ? true : undefined, - source: "rpc", + source: command.source, preflightResult: (didSucceed) => { if (didSucceed) { this.recordWorkerRecoveryState(state, "prompt_accepted", true); @@ -2834,6 +2847,24 @@ export class AgentDaemon { return undefined; } + case "prompt_and_wait": { + const state = this.getBoundSessionState(command.activeSessionId); + await this.promptWithAgentMessagePreparingGuard(state, command.message, { + content: command.content, + images: command.images, + streamingBehavior: command.streamingBehavior, + expandPromptTemplates: command.expandPromptTemplates, + skipInputHandlers: command.expandPromptTemplates === false ? true : undefined, + source: command.source, + preflightResult: (didSucceed) => { + if (didSucceed) { + this.recordWorkerRecoveryState(state, "prompt_accepted", true); + } + }, + }); + return success(command.id, command.type); + } + case "steer": { const state = this.getBoundSessionState(command.activeSessionId); if (command.expandPromptTemplates === false) { @@ -3026,6 +3057,15 @@ export class AgentDaemon { return success(command.id, "execute_bash"); } + case "execute_bash_and_wait": { + const state = this.getSessionState(command.activeSessionId); + return success( + command.id, + "execute_bash_and_wait", + await state.runtime.session.executeBash(command.command), + ); + } + case "abort_bash": { const state = this.getSessionState(command.activeSessionId); state.runtime.session.abortBash(); @@ -3044,6 +3084,22 @@ export class AgentDaemon { return success(command.id, "wait_for_idle"); } + case "wait_for_headless_completion": { + const state = this.getSessionState(command.activeSessionId); + return success( + command.id, + "wait_for_headless_completion", + await waitForHeadlessCompletion(state.runtime.session), + ); + } + + case "get_session_header": { + const state = this.getSessionState(command.activeSessionId); + return success(command.id, "get_session_header", { + header: state.runtime.session.sessionManager.getHeader(), + }); + } + case "get_state": { const state = this.getSessionState(command.activeSessionId); return success(command.id, "get_state", summaryForActiveSession(state)); @@ -3247,6 +3303,12 @@ export class AgentDaemon { return success(command.id, "set_auto_compaction"); } + case "set_auto_retry": { + const state = this.getSessionState(command.activeSessionId); + state.runtime.session.setAutoRetryEnabled(command.enabled); + return success(command.id, "set_auto_retry"); + } + case "compact": { const state = this.getSessionState(command.activeSessionId); const result = await state.runtime.session.compact(command.customInstructions); diff --git a/packages/coding-agent/src/modes/daemon/daemon-protocol.ts b/packages/coding-agent/src/modes/daemon/daemon-protocol.ts index e363db2e9a8..5243e7614d4 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-protocol.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-protocol.ts @@ -7,12 +7,15 @@ import type { } from "../../core/agent-messages.js"; import type { AgentSessionRuntimeConfig } from "../../core/agent-session-config.js"; import type { AgentSessionRuntimeMetadata } from "../../core/agent-session-runtime.js"; +import type { AgentAutonomousStatus } from "../../core/autonomous.js"; +import type { BashResult } from "../../core/bash-executor.js"; import type { AgentCronJob, AgentHeartbeatDeliveryMode, AgentHeartbeatManagementAction, AgentHeartbeatUpdateAction, } from "../../core/cron-jobs.js"; +import type { InputSource } from "../../core/extensions/types.js"; import type { CustomMessage } from "../../core/messages.js"; import type { SessionCwdIssue } from "../../core/session-cwd.js"; import type { DeleteSessionFileResult } from "../../core/session-file-actions.js"; @@ -27,6 +30,7 @@ import type { AgentConnectionScopedModel, AgentConnectionSessionContext, AgentConnectionSessionEvent, + AgentConnectionSessionHeader, AgentConnectionSessionTreeNode, AgentConnectionSideQuestionEvent, AgentConnectionState, @@ -46,7 +50,7 @@ export const DAEMON_PROTOCOL_NAME = "prime-agent.daemon"; export const DAEMON_PROTOCOL_VERSION = 4; export const DAEMON_COMMAND_ENVELOPE_MIN_PROTOCOL_VERSION = 2; export const DAEMON_SCHEMA_REVISION = 1; -export const DAEMON_SCHEMA_ID = "protocol-4-schema-1-07ce10c7f10a"; +export const DAEMON_SCHEMA_ID = "protocol-4-schema-1-f55df8742274"; export type DaemonProtocolName = typeof DAEMON_PROTOCOL_NAME; export type DaemonProtocolVersion = number; @@ -63,7 +67,8 @@ export type DaemonClientCapability = | "event_sequence" | "extension_ui" | "slim_attach" - | "chunked_snapshot"; + | "chunked_snapshot" + | "client_owned_sessions"; export type DaemonServerCapability = DaemonClientCapability | "heartbeat_catalog" | "heartbeat_management"; export type DaemonReplayStatus = "complete" | "partial" | "unavailable"; @@ -88,6 +93,7 @@ export const DAEMON_SUPPORTED_CLIENT_CAPABILITIES: readonly DaemonClientCapabili "extension_ui", "slim_attach", "chunked_snapshot", + "client_owned_sessions", ]; export const DAEMON_DEFAULT_SERVER_CAPABILITIES: readonly DaemonServerCapability[] = [ @@ -125,6 +131,12 @@ export interface DaemonClientEnv { env?: Record; } +export type DaemonSessionLifecycle = "resident" | "client_owned"; + +export interface DaemonLaunchEnv { + launchEnv?: Record; +} + /** * The allowlist of env vars a client may forward. One shared list because it * is the wire contract: clients filter before sending and the daemon @@ -150,6 +162,16 @@ export function collectDaemonClientEnv(source: NodeJS.ProcessEnv = process.env): return Object.keys(env).length > 0 ? env : undefined; } +export function collectDaemonLaunchEnv(source: NodeJS.ProcessEnv = process.env): Record { + const env: Record = {}; + for (const [key, value] of Object.entries(source)) { + if (value !== undefined && !key.startsWith("PRIME_AGENT_INTERNAL_")) { + env[key] = value; + } + } + return env; +} + export interface DaemonReplayInfo { status: DaemonReplayStatus; fromSequence?: DaemonEventSequence; @@ -296,17 +318,27 @@ export type DaemonSavedSessionListCommand = }; export type DaemonCommand = - | { id?: string; type: "list"; all?: boolean; cwd?: string; sessionDir?: string } + | { + id?: string; + type: "list"; + all?: boolean; + cwd?: string; + sessionDir?: string; + includeClientOwned?: boolean; + } | DaemonSavedSessionListCommand | ({ id?: string; type: "create"; sessionPath?: string; continueRecent?: boolean; + noSession?: boolean; name?: string; config?: AgentSessionRuntimeConfig; runtimeMetadata?: AgentSessionRuntimeMetadata; - } & DaemonClientEnv) + lifecycle?: DaemonSessionLifecycle; + } & DaemonClientEnv & + DaemonLaunchEnv) // Attach env is adopt-if-absent only: it fills identity for env-less // sessions (e.g. cron-created) but never rebinds one, since watchers // (agents view, subagent viewers) also attach. @@ -316,7 +348,8 @@ export type DaemonCommand = activeSessionId: string; supportsExtensionUi?: boolean; } & DaemonAttachClientMetadata & - DaemonClientEnv) + DaemonClientEnv & + DaemonLaunchEnv) | ({ id?: string; type: "reattach"; @@ -324,8 +357,11 @@ export type DaemonCommand = targetActiveSessionId: string; supportsExtensionUi?: boolean; } & DaemonAttachClientMetadata & - DaemonClientEnv) + DaemonClientEnv & + DaemonLaunchEnv) | { id?: string; type: "detach"; activeSessionId?: string } + | { id?: string; type: "complete_owned_session"; activeSessionId: string } + | { id?: string; type: "promote_owned_session"; activeSessionId: string } | { id?: string; type: "kill"; activeSessionId: string } | { id?: string; type: "rename"; activeSessionId: string; name: string } | { @@ -337,9 +373,21 @@ export type DaemonCommand = images?: ImageContent[]; streamingBehavior?: "steer" | "followUp"; expandPromptTemplates?: boolean; + source?: InputSource; agentMessageId?: string; customMessage?: CustomMessage; } + | { + id?: string; + type: "prompt_and_wait"; + activeSessionId: string; + message: string; + content?: (TextContent | ImageContent)[]; + images?: ImageContent[]; + streamingBehavior?: "steer" | "followUp"; + expandPromptTemplates?: boolean; + source?: InputSource; + } | { id?: string; type: "steer"; @@ -382,9 +430,9 @@ export type DaemonCommand = fromActiveSessionId?: string; deliveryMode?: AgentSessionMessageDeliveryMode; } - | { id?: string; type: "agent_messages_status" } - | { id?: string; type: "agent_messages_pause" } - | { id?: string; type: "agent_messages_resume" } + | { id?: string; type: "agent_messages_status"; activeSessionId?: string } + | { id?: string; type: "agent_messages_pause"; activeSessionId?: string } + | { id?: string; type: "agent_messages_resume"; activeSessionId?: string } | { id?: string; type: "agent_messages_clear"; activeSessionId: string } | { id?: string; type: "abort"; activeSessionId: string } | { @@ -405,6 +453,8 @@ export type DaemonCommand = | { id?: string; type: "abort_bash"; activeSessionId: string } | { id?: string; type: "cancel_rlm_child"; activeSessionId: string; childId: string } | { id?: string; type: "wait_for_idle"; activeSessionId: string } + | { id?: string; type: "wait_for_headless_completion"; activeSessionId: string } + | { id?: string; type: "get_session_header"; activeSessionId: string } | { id?: string; type: "get_state"; activeSessionId: string } | { id?: string; type: "get_connection_state"; activeSessionId: string } | { id?: string; type: "get_messages"; activeSessionId: string } @@ -417,7 +467,7 @@ export type DaemonCommand = | { id?: string; type: "clear_queue"; activeSessionId: string } | { id?: string; type: "abort_and_clear_queue"; activeSessionId: string } | { id?: string; type: "cron_list"; activeSessionId?: string; includeInactive?: boolean } - | { id?: string; type: "heartbeats_list" } + | { id?: string; type: "heartbeats_list"; activeSessionId?: string } | { id?: string; type: "heartbeat_manage"; @@ -425,8 +475,15 @@ export type DaemonCommand = jobId: string; action: AgentHeartbeatManagementAction; } - | { id?: string; type: "cron_add"; activeSessionId: string; schedule: string; prompt: string } - | { id?: string; type: "cron_cancel"; jobId: string } + | { + id?: string; + type: "cron_add"; + activeSessionId: string; + schedule: string; + prompt: string; + promoteOwnedSession?: boolean; + } + | { id?: string; type: "cron_cancel"; activeSessionId?: string; jobId: string } | { id?: string; type: "heartbeat_get"; activeSessionId: string } | { id?: string; @@ -435,6 +492,7 @@ export type DaemonCommand = schedule: string; prompt: string; deliveryMode?: AgentHeartbeatDeliveryMode; + promoteOwnedSession?: boolean; } | { id?: string; type: "heartbeat_update"; activeSessionId: string; action: AgentHeartbeatUpdateAction } | { id?: string; type: "set_model"; activeSessionId: string; provider: string; modelId: string } @@ -447,6 +505,7 @@ export type DaemonCommand = | { id?: string; type: "set_steering_mode"; activeSessionId: string; mode: AgentConnectionQueueMode } | { id?: string; type: "set_follow_up_mode"; activeSessionId: string; mode: AgentConnectionQueueMode } | { id?: string; type: "set_auto_compaction"; activeSessionId: string; enabled: boolean } + | { id?: string; type: "set_auto_retry"; activeSessionId: string; enabled: boolean } | { id?: string; type: "compact"; activeSessionId: string; customInstructions?: string } | { id?: string; @@ -459,6 +518,7 @@ export type DaemonCommand = | { id?: string; type: "abort_compaction"; activeSessionId: string } | { id?: string; type: "abort_branch_summary"; activeSessionId: string } | { id?: string; type: "abort_retry"; activeSessionId: string } + | { id?: string; type: "execute_bash_and_wait"; activeSessionId: string; command: string } | { id?: string; type: "reload"; activeSessionId: string } | { id?: string; type: "new_session"; activeSessionId: string; parentSession?: string } | { id?: string; type: "switch_session"; activeSessionId: string; sessionPath: string; cwdOverride?: string } @@ -507,6 +567,11 @@ export interface DaemonCommandCompatibility { } const LEGACY_DAEMON_COMMAND = { minProtocol: 1 } as const; +const CURRENT_DAEMON_COMMAND = { minProtocol: DAEMON_PROTOCOL_VERSION } as const; +const CLIENT_OWNED_DAEMON_COMMAND = { + minProtocol: DAEMON_PROTOCOL_VERSION, + capability: "client_owned_sessions", +} as const; export const DAEMON_COMMAND_COMPATIBILITY = { ack_result: LEGACY_DAEMON_COMMAND, @@ -516,9 +581,12 @@ export const DAEMON_COMMAND_COMPATIBILITY = { attach: LEGACY_DAEMON_COMMAND, reattach: LEGACY_DAEMON_COMMAND, detach: LEGACY_DAEMON_COMMAND, + complete_owned_session: CLIENT_OWNED_DAEMON_COMMAND, + promote_owned_session: CLIENT_OWNED_DAEMON_COMMAND, kill: LEGACY_DAEMON_COMMAND, rename: LEGACY_DAEMON_COMMAND, prompt: LEGACY_DAEMON_COMMAND, + prompt_and_wait: CURRENT_DAEMON_COMMAND, steer: LEGACY_DAEMON_COMMAND, follow_up: LEGACY_DAEMON_COMMAND, restore_next_turn: LEGACY_DAEMON_COMMAND, @@ -536,6 +604,8 @@ export const DAEMON_COMMAND_COMPATIBILITY = { abort_bash: LEGACY_DAEMON_COMMAND, cancel_rlm_child: LEGACY_DAEMON_COMMAND, wait_for_idle: LEGACY_DAEMON_COMMAND, + wait_for_headless_completion: CURRENT_DAEMON_COMMAND, + get_session_header: CURRENT_DAEMON_COMMAND, get_state: LEGACY_DAEMON_COMMAND, get_connection_state: LEGACY_DAEMON_COMMAND, get_messages: LEGACY_DAEMON_COMMAND, @@ -565,11 +635,13 @@ export const DAEMON_COMMAND_COMPATIBILITY = { set_steering_mode: LEGACY_DAEMON_COMMAND, set_follow_up_mode: LEGACY_DAEMON_COMMAND, set_auto_compaction: LEGACY_DAEMON_COMMAND, + set_auto_retry: CURRENT_DAEMON_COMMAND, compact: LEGACY_DAEMON_COMMAND, refine: LEGACY_DAEMON_COMMAND, abort_compaction: LEGACY_DAEMON_COMMAND, abort_branch_summary: LEGACY_DAEMON_COMMAND, abort_retry: LEGACY_DAEMON_COMMAND, + execute_bash_and_wait: CURRENT_DAEMON_COMMAND, reload: LEGACY_DAEMON_COMMAND, new_session: LEGACY_DAEMON_COMMAND, switch_session: LEGACY_DAEMON_COMMAND, @@ -662,6 +734,9 @@ export interface DaemonSavedSessionInfo { } export type DaemonDeleteSavedSessionResult = DeleteSessionFileResult; +export type DaemonAutonomousStatus = AgentAutonomousStatus; +export type DaemonBashResult = BashResult; +export type DaemonSessionHeader = AgentConnectionSessionHeader; export type DaemonResourceSnapshot = AgentConnectionResourceSnapshot; @@ -841,6 +916,7 @@ const READ_ONLY_DAEMON_COMMANDS: ReadonlySet = new Set([ "reattach", "agent_messages_status", "wait_for_idle", + "get_session_header", "get_state", "get_connection_state", "get_messages", diff --git a/packages/coding-agent/src/modes/daemon/daemon-session-list.ts b/packages/coding-agent/src/modes/daemon/daemon-session-list.ts index 71eec1d9e6d..cf4ae066fad 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-session-list.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-session-list.ts @@ -87,6 +87,16 @@ export function resolveAttachModelFallbackMessage( return summary.model ? undefined : startupModelFallbackMessage; } +export function isSessionSummaryBusy(summary: SessionSummary): boolean { + return ( + summary.isStreaming || + summary.isCompacting || + summary.isBashRunning === true || + summary.hasRunningRlmChildren === true || + summary.pendingMessageCount > 0 + ); +} + export function buildSessionList( activeSessions: readonly ActiveSessionState[], savedSessions: readonly SessionInfo[], diff --git a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts index f6f8d749de4..e5b2403be44 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts @@ -5,7 +5,7 @@ import { createServer, type Server, type Socket } from "node:net"; import { dirname, isAbsolute, join, resolve } from "node:path"; import { Writable } from "node:stream"; import { getLogger } from "@earendil-works/pi-ai"; -import { createCliSubprocessLaunchSpec } from "../../cli/subprocess-launch.js"; +import { createCliSubprocessEnv, createCliSubprocessLaunchSpec } from "../../cli/subprocess-launch.js"; import { appendRotatingLog, getCronJobsPath, @@ -27,7 +27,7 @@ import { ORPHAN_PROCESS_JOURNAL_ENV, readActiveOrphanProcesses, } from "../../core/orphan-process-journal.js"; -import { canonicalSessionPath, getProcessStartId } from "../../core/session-lease.js"; +import { canonicalSessionPath, getProcessStartId, SessionAlreadyActiveError } from "../../core/session-lease.js"; import type { SessionInfo } from "../../core/session-manager.js"; import { signalProcessGroupOrProcess } from "../../utils/child-process.js"; import type { AgentConnectionHeartbeat } from "../agent-connection/types.js"; @@ -60,7 +60,7 @@ import { } from "./daemon-protocol.js"; import { getDaemonRuntimeIdentity } from "./daemon-runtime-identity.js"; import { matchesSessionIdSuffix } from "./daemon-session-id.js"; -import { type SessionSummary, summaryForInactiveSession } from "./daemon-session-list.js"; +import { isSessionSummaryBusy, type SessionSummary, summaryForInactiveSession } from "./daemon-session-list.js"; import { acquireDaemonSocketPathLease, cleanupDaemonSocketPath, @@ -103,6 +103,7 @@ const WORKER_CONNECT_TIMEOUT_MS = 30_000; const WORKER_REQUEST_TIMEOUT_MS = 24 * 60 * 60 * 1000; const WORKER_RETRY_DELAYS_MS = [250, 1000, 5000] as const; const DEFERRED_RECOVERY_RECHECK_MS = 5000; +const OWNED_WORKER_DISCONNECT_GRACE_MS = 30_000; const SUPERVISOR_CONFIG_FILE_NAME = "supervisor-config"; const WORKER_STARTUP_GATE_FD = 3; @@ -114,9 +115,12 @@ const DAEMON_COMMAND_TYPES: ReadonlySet = new Set([ "attach", "reattach", "detach", + "complete_owned_session", + "promote_owned_session", "kill", "rename", "prompt", + "prompt_and_wait", "steer", "follow_up", "restore_next_turn", @@ -131,9 +135,12 @@ const DAEMON_COMMAND_TYPES: ReadonlySet = new Set([ "start_side_question", "abort_side_question", "execute_bash", + "execute_bash_and_wait", "abort_bash", "cancel_rlm_child", "wait_for_idle", + "wait_for_headless_completion", + "get_session_header", "get_state", "get_connection_state", "get_messages", @@ -163,6 +170,7 @@ const DAEMON_COMMAND_TYPES: ReadonlySet = new Set([ "set_steering_mode", "set_follow_up_mode", "set_auto_compaction", + "set_auto_retry", "compact", "refine", "abort_compaction", @@ -208,6 +216,9 @@ interface ResidentWorker { deferredRecovery?: Promise; intentionalStop: boolean; stopRevision: number; + launchEnv?: Record; + ownerCleanupTimer?: ReturnType; + promotedOwnerClientId?: string; } interface SnapshotDuplicateValidation { @@ -313,6 +324,11 @@ function withoutCommandId(command: DaemonCommand): DaemonCommandBody { return body as DaemonCommandBody; } +function withoutSupervisorCreateFields(command: DaemonCreateCommand): DaemonCreateCommand { + const { launchEnv: _launchEnv, lifecycle: _lifecycle, ...workerCommand } = command; + return workerCommand; +} + function responseWithId(response: DaemonResponse, id: string | undefined): DaemonResponse { return { ...response, id }; } @@ -339,6 +355,7 @@ function isDaemonWorkerDescriptor(value: unknown, socketPath: string): value is Number.isInteger(descriptor.pid) && (descriptor.pid ?? 0) > 0 && (descriptor.processStartId === undefined || typeof descriptor.processStartId === "string") && + (descriptor.ownerClientId === undefined || typeof descriptor.ownerClientId === "string") && typeof descriptor.socketPath === "string" && typeof descriptor.authenticationToken === "string" && typeof descriptor.rootActiveSessionId === "string" && @@ -520,6 +537,7 @@ export class DaemonSupervisor { private cleanupPromise?: Promise; private shuttingDown = false; private readonly clients = new Set(); + private readonly protocolClientIds = new WeakMap(); private readonly workers = new Map(); private readonly openingWorkers = new Map>(); private readonly signalCleanupHandlers: Array<() => void> = []; @@ -621,6 +639,9 @@ export class DaemonSupervisor { throw adoptionFailure; } await this.syncAgentPeers().catch((error) => this.log(`Could not synchronize agent peers: ${String(error)}`)); + for (const worker of this.workers.values()) { + this.scheduleOwnedWorkerCleanup(worker); + } await this.ownership.updatePhase("owner"); this.log(`Prime Agent daemon supervisor ${this.generation} listening on ${this.socketPath}`); this.markReady(); @@ -828,6 +849,7 @@ export class DaemonSupervisor { client.attachedActiveSessionIds.delete(activeSessionId); void this.syncWorkerExtensionUi(activeSessionId); } + this.scheduleOwnedWorkerCleanupForClient(this.protocolClientId(client)); }; socket.on("close", cleanup); socket.on("error", cleanup); @@ -841,6 +863,56 @@ export class DaemonSupervisor { }); } + private cancelOwnedWorkerCleanup(clientId: string): void { + for (const worker of this.workers.values()) { + if (worker.descriptor.ownerClientId !== clientId || !worker.ownerCleanupTimer) { + continue; + } + clearTimeout(worker.ownerCleanupTimer); + worker.ownerCleanupTimer = undefined; + } + } + + private protocolClientId(client: DaemonSocketClient): string { + return this.protocolClientIds.get(client) ?? client.id; + } + + private scheduleOwnedWorkerCleanupForClient(clientId: string): void { + if ([...this.clients].some((client) => this.protocolClientId(client) === clientId)) { + return; + } + for (const worker of this.workers.values()) { + if (worker.descriptor.ownerClientId === clientId) { + this.scheduleOwnedWorkerCleanup(worker); + } + } + } + + private scheduleOwnedWorkerCleanup(worker: ResidentWorker): void { + const ownerClientId = worker.descriptor.ownerClientId; + if ( + !ownerClientId || + worker.ownerCleanupTimer || + [...this.clients].some((client) => this.protocolClientId(client) === ownerClientId) + ) { + return; + } + worker.ownerCleanupTimer = setTimeout(() => { + worker.ownerCleanupTimer = undefined; + if ( + worker.descriptor.ownerClientId !== ownerClientId || + [...this.clients].some((client) => this.protocolClientId(client) === ownerClientId) || + this.workers.get(worker.descriptor.workerId) !== worker + ) { + return; + } + void this.stopWorker(worker, true).catch((error) => + this.log(`Could not clean up client-owned worker ${worker.descriptor.workerId}: ${String(error)}`), + ); + }, OWNED_WORKER_DISCONNECT_GRACE_MS); + worker.ownerCleanupTimer.unref(); + } + private async handleLine(client: DaemonSocketClient, line: string): Promise { try { await this.ready; @@ -856,8 +928,10 @@ export class DaemonSupervisor { : (parsed as { id?: unknown; type?: unknown }); if (isDaemonCommandEnvelope(parsed)) { envelopeClientId = parsed.clientId ?? client.id; + this.protocolClientIds.set(client, envelopeClientId); client.id = envelopeClientId; } + this.cancelOwnedWorkerCleanup(client.id); if (typeof candidate.type !== "string" || !DAEMON_COMMAND_TYPES.has(candidate.type)) { const commandName = typeof candidate.type === "string" ? candidate.type : "unknown"; this.write( @@ -938,11 +1012,11 @@ export class DaemonSupervisor { this.commandJournal.acknowledge(client.id, command.commandId); return undefined; case "list": - return this.handleList(command); + return this.handleList(client, command); case "list_saved_sessions": return this.handleSavedSessionList(client, command); case "create": { - const worker = await this.createOrReuseWorker(client.id, command); + const worker = await this.createOrReuseWorker(this.protocolClientId(client), command); const summary = worker.summaries.get(worker.descriptor.rootActiveSessionId); if (!summary) { throw new Error("Session worker started without a root session"); @@ -980,7 +1054,7 @@ export class DaemonSupervisor { return success(command.id, "attach", attached.result); } case "reattach": { - const target = await this.findWorker(command.targetActiveSessionId); + const target = await this.findWorkerForClient(client, command.targetActiveSessionId); const targetActiveSessionId = target.summary.activeSessionId ?? target.summary.id; if (targetActiveSessionId === command.activeSessionId) { return success(command.id, command.type, { cancelled: false }); @@ -1033,13 +1107,31 @@ export class DaemonSupervisor { case "detach": this.detachClient(client, command.activeSessionId); return success(command.id, "detach"); + case "complete_owned_session": { + const match = await this.findWorkerForClient(client, command.activeSessionId); + if (match.worker.descriptor.ownerClientId !== this.protocolClientId(client)) { + throw new Error("Session is not owned by this client"); + } + if (match.worker.ownerCleanupTimer) { + clearTimeout(match.worker.ownerCleanupTimer); + match.worker.ownerCleanupTimer = undefined; + } + await this.stopWorker(match.worker, true); + return success(command.id, command.type); + } + case "promote_owned_session": { + const match = await this.findWorkerForClient(client, command.activeSessionId); + await this.promoteOwnedWorker(client, match.worker); + return success(command.id, command.type, this.publicSummary(match.worker, match.summary)); + } case "retry_worker": { const direct = [...this.workers.values()].find( (worker) => worker.descriptor.rootActiveSessionId === command.activeSessionId || worker.descriptor.rootSessionId === command.activeSessionId, ); - const worker = direct ?? (await this.findWorker(command.activeSessionId)).worker; + const worker = direct ?? (await this.findWorkerForClient(client, command.activeSessionId)).worker; + this.assertWorkerAccessibleToClient(client, worker, command.activeSessionId); worker.intentionalStop = false; worker.descriptor.lifecycle = "recovering"; worker.descriptor.consecutiveFailures = 0; @@ -1062,7 +1154,11 @@ export class DaemonSupervisor { return success(command.id, "prepare_update_restart", manifest); } case "agent_messages_status": { - const first = [...this.workers.values()].find((worker) => worker.client); + if (command.activeSessionId) { + const match = await this.findWorkerForClient(client, command.activeSessionId); + return this.forwardToWorker(match.worker, command); + } + const first = [...this.workers.values()].find((worker) => this.isVisibleWorker(worker) && worker.client); if (!first) { return success(command.id, command.type, { paused: false, limits: {} }); } @@ -1070,9 +1166,13 @@ export class DaemonSupervisor { } case "agent_messages_pause": case "agent_messages_resume": { + if (command.activeSessionId) { + const match = await this.findWorkerForClient(client, command.activeSessionId); + return this.forwardToWorker(match.worker, command); + } const responses = await Promise.all( [...this.workers.values()] - .filter((worker) => worker.client) + .filter((worker) => this.isVisibleWorker(worker) && worker.client) .map((worker) => this.forwardToWorker(worker, command)), ); const failed = responses.find((response) => !response.success); @@ -1080,13 +1180,16 @@ export class DaemonSupervisor { } case "cron_list": { if (command.activeSessionId) { - const match = await this.findWorker(command.activeSessionId); + const match = await this.findWorkerForClient(client, command.activeSessionId); return this.forwardToWorker(match.worker, command); } const jobs = new Map(); const responses = await Promise.all( [...this.workers.values()] - .filter((worker) => worker.client && worker.descriptor.lifecycle === "ready") + .filter( + (worker) => + this.isVisibleWorker(worker) && worker.client && worker.descriptor.lifecycle === "ready", + ) .map((worker) => this.forwardToWorker(worker, command, 5000).catch((error: unknown) => failure(command.id, command.type, error, serializeDaemonError(error)), @@ -1105,7 +1208,11 @@ export class DaemonSupervisor { return success(command.id, "cron_list", { jobs: sortCronJobs([...jobs.values()]) }); } case "heartbeats_list": { - const workers = [...this.workers.values()]; + if (command.activeSessionId) { + const match = await this.findWorkerForClient(client, command.activeSessionId); + return this.forwardToWorker(match.worker, command); + } + const workers = [...this.workers.values()].filter((worker) => this.isVisibleWorker(worker)); const heartbeats = new Map(); const snapshots: Array<{ heartbeats?: AgentConnectionHeartbeat[]; response?: DaemonResponse }> = await Promise.all( @@ -1152,7 +1259,8 @@ export class DaemonSupervisor { heartbeat.job.id === command.jobId && heartbeat.job.activeSessionId === command.activeSessionId, ), ); - const worker = cachedWorker ?? (await this.findWorker(command.activeSessionId)).worker; + const worker = cachedWorker ?? (await this.findWorkerForClient(client, command.activeSessionId)).worker; + this.assertWorkerAccessibleToClient(client, worker, command.activeSessionId); const response = await this.forwardToWorker(worker, command); if ( response.success && @@ -1173,13 +1281,24 @@ export class DaemonSupervisor { return response; } case "cron_add": { - const match = await this.findWorker(command.activeSessionId); - return this.forwardToWorker(match.worker, command); + const match = await this.findWorkerForClient(client, command.activeSessionId); + const response = await this.forwardToWorker(match.worker, command); + if (response.success && command.promoteOwnedSession) { + await this.promoteOwnedWorker(client, match.worker); + } + return response; } case "cron_cancel": { + if (command.activeSessionId) { + const match = await this.findWorkerForClient(client, command.activeSessionId); + return this.forwardToWorker(match.worker, command); + } const listed = await Promise.all( [...this.workers.values()] - .filter((worker) => worker.client && worker.descriptor.lifecycle === "ready") + .filter( + (worker) => + this.isVisibleWorker(worker) && worker.client && worker.descriptor.lifecycle === "ready", + ) .map(async (worker) => ({ worker, response: await this.forwardToWorker( @@ -1200,15 +1319,19 @@ export class DaemonSupervisor { throw new Error(`No cron job found: ${command.jobId}`); } case "heartbeat_get": { - const match = await this.findWorker(command.activeSessionId); + const match = await this.findWorkerForClient(client, command.activeSessionId); return this.forwardToWorker(match.worker, command); } case "heartbeat_set": { - const match = await this.findWorker(command.activeSessionId); - return this.forwardToWorker(match.worker, command); + const match = await this.findWorkerForClient(client, command.activeSessionId); + const response = await this.forwardToWorker(match.worker, command); + if (response.success && command.promoteOwnedSession) { + await this.promoteOwnedWorker(client, match.worker); + } + return response; } case "heartbeat_update": { - const match = await this.findWorker(command.activeSessionId); + const match = await this.findWorkerForClient(client, command.activeSessionId); return this.forwardToWorker(match.worker, command); } case "rename_saved_session": @@ -1230,8 +1353,10 @@ export class DaemonSupervisor { } if (command.type === "send_message") { - const target = await this.findWorker(command.targetActiveSessionId); - const source = command.fromActiveSessionId ? await this.findWorker(command.fromActiveSessionId) : undefined; + const target = await this.findWorkerForClient(client, command.targetActiveSessionId); + const source = command.fromActiveSessionId + ? await this.findWorkerForClient(client, command.fromActiveSessionId) + : undefined; if (source && source.worker !== target.worker) { if (!target.worker.client) { throw new Error("Target session worker is not connected"); @@ -1260,7 +1385,7 @@ export class DaemonSupervisor { if (!("activeSessionId" in command) || typeof command.activeSessionId !== "string") { throw new Error(`Supervisor cannot route daemon command: ${command.type}`); } - const match = await this.findWorker(command.activeSessionId); + const match = await this.findWorkerForClient(client, command.activeSessionId); const resolvedCommand = { ...command, activeSessionId: match.summary.activeSessionId ?? match.summary.id, @@ -1281,20 +1406,35 @@ export class DaemonSupervisor { return response; } - private async handleList(command: Extract): Promise { + private async handleList( + client: DaemonSocketClient, + command: Extract, + ): Promise { await Promise.all( [...this.workers.values()].map((worker) => this.refreshWorkerSummaries(worker).catch(() => undefined)), ); await this.syncAgentPeers().catch((error) => this.log(`Could not synchronize agent peers: ${String(error)}`)); - const active = [...this.workers.values()].flatMap((worker) => - [...worker.summaries.values()].map((summary) => this.publicSummary(worker, summary)), - ); + const clientOwnedWorkers = [...this.workers.values()].filter((worker) => !this.isVisibleWorker(worker)); + const active = [...this.workers.values()] + .filter( + (worker) => + this.isVisibleWorker(worker) || + (command.includeClientOwned === true && this.isWorkerAccessibleToClient(client, worker)), + ) + .flatMap((worker) => [...worker.summaries.values()].map((summary) => this.publicSummary(worker, summary))); + const busyClientOwnedSessionCount = clientOwnedWorkers + .flatMap((worker) => [...worker.summaries.values()]) + .filter(isSessionSummaryBusy).length; + const data = { + sessions: active, + ...(command.includeClientOwned ? { busyClientOwnedSessionCount } : {}), + }; if (!command.all) { - return success(command.id, "list", { sessions: active }); + return success(command.id, "list", data); } const sessionDir = command.sessionDir ?? this.defaultSessionConfig.sessionDir; const saved = await this.catalog.list(command.cwd ? resolve(command.cwd) : undefined, sessionDir); - return success(command.id, "list", { sessions: mergeSessionLists(active, saved) }); + return success(command.id, "list", { ...data, sessions: mergeSessionLists(active, saved) }); } private async handleSavedSessionList( @@ -1305,7 +1445,7 @@ export class DaemonSupervisor { let sessionDir: string | undefined; let activeSessionId: string | undefined; if ("activeSessionId" in command) { - const match = await this.findWorker(command.activeSessionId); + const match = await this.findWorkerForClient(client, command.activeSessionId); cwd = match.summary.cwd; sessionDir = this.defaultSessionConfig.sessionDir; activeSessionId = match.summary.activeSessionId ?? match.summary.id; @@ -1340,10 +1480,11 @@ export class DaemonSupervisor { private async createOrReuseWorker(clientId: string, command: DaemonCreateCommand): Promise { let createCommand = command; + const ownerClientId = command.lifecycle === "client_owned" ? clientId : undefined; if (command.sessionPath) { const activeMatches = this.matchWorkers(command.sessionPath); if (activeMatches.length === 1) { - return activeMatches[0]!.worker; + return this.reuseWorkerForCreate(activeMatches[0]!.worker, ownerClientId, command.sessionPath); } if (activeMatches.length > 1) { throw new Error(`Ambiguous active session "${command.sessionPath}"`); @@ -1355,7 +1496,7 @@ export class DaemonSupervisor { createCommand = { ...command, sessionPath }; const existing = this.findWorkerBySessionFile(sessionPath); if (existing) { - return existing.worker; + return this.reuseWorkerForCreate(existing.worker, ownerClientId, sessionPath); } } const key = createCommand.sessionPath @@ -1365,7 +1506,7 @@ export class DaemonSupervisor { if (pending) { return pending; } - const opening = this.launchWorker(createCommand); + const opening = this.launchWorker(createCommand, undefined, ownerClientId); this.openingWorkers.set(key, opening); try { return await opening; @@ -1376,14 +1517,56 @@ export class DaemonSupervisor { } } - private async launchWorker(command: DaemonCreateCommand, existing?: ResidentWorker): Promise { + private reuseWorkerForCreate( + worker: ResidentWorker, + ownerClientId: string | undefined, + sessionPath: string, + ): ResidentWorker { + if (worker.descriptor.ownerClientId === ownerClientId) { + return worker; + } + throw new SessionAlreadyActiveError(sessionPath, worker.descriptor.rootActiveSessionId); + } + + private async promoteOwnedWorker(client: DaemonSocketClient, worker: ResidentWorker): Promise { + const clientId = this.protocolClientId(client); + if (worker.descriptor.ownerClientId === undefined && worker.promotedOwnerClientId === clientId) { + return; + } + if (worker.descriptor.ownerClientId !== clientId) { + throw new Error("Session is not owned by this client"); + } + const previousDescriptor = worker.descriptor; + worker.descriptor = { ...previousDescriptor, ownerClientId: undefined }; + try { + this.persistWorker(worker); + } catch (error) { + worker.descriptor = previousDescriptor; + throw error; + } + worker.promotedOwnerClientId = clientId; + if (worker.ownerCleanupTimer) { + clearTimeout(worker.ownerCleanupTimer); + worker.ownerCleanupTimer = undefined; + } + worker.launchEnv = undefined; + await this.syncAgentPeers().catch((error) => this.log(`Could not synchronize agent peers: ${String(error)}`)); + } + + private async launchWorker( + command: DaemonCreateCommand, + existing?: ResidentWorker, + ownerClientId?: string, + ): Promise { await this.assertRecoveryAllowed(); if (existing && this.isWorkerRecoveryCancelled(existing)) { throw new Error(`Session worker ${existing.descriptor.workerId} recovery was cancelled`); } const recoveryStopRevision = existing?.stopRevision; + const launchEnv = + ownerClientId || existing?.descriptor.ownerClientId ? (command.launchEnv ?? existing?.launchEnv) : undefined; const createCommand: DaemonCreateCommand = { - ...command, + ...withoutSupervisorCreateFields(command), config: mergeAgentSessionRuntimeConfig(this.defaultSessionConfig, command.config), }; const workerId = existing?.descriptor.workerId ?? createActiveSessionId(); @@ -1401,8 +1584,9 @@ export class DaemonSupervisor { const child: ChildProcess = spawn(launch.command, launch.args, { cwd: createCommand.config?.cwd ?? process.cwd(), detached: true, - env: { + env: createCliSubprocessEnv({ ...process.env, + ...launchEnv, [DAEMON_WORKER_ROLE_ENV]: "1", [DAEMON_WORKER_TOKEN_ENV]: token, [DAEMON_WORKER_ACTIVE_SESSION_ID_ENV]: rootActiveSessionId, @@ -1412,9 +1596,16 @@ export class DaemonSupervisor { [ORPHAN_PROCESS_JOURNAL_ENV]: orphanProcessJournalPath, [SESSION_LEASES_ENABLED_ENV]: "1", [SESSION_LEASE_OWNER_ID_ENV]: rootActiveSessionId, - }, - stdio: ["ignore", "ignore", "ignore", "pipe"], + }), + stdio: ["ignore", "ignore", "pipe", "pipe"], }); + const detachWorkerStderr = child.stderr + ? attachJsonlLineReader(child.stderr, (line) => this.log(`Session worker ${workerId} stderr: ${line}`), { + maxLineLength: 64 * 1024, + onLineOverflow: (prefix) => this.log(`Session worker ${workerId} stderr: ${prefix} [truncated]`), + }) + : () => {}; + child.once("close", detachWorkerStderr); const childClosed = new Promise((resolveClose) => child.once("close", () => resolveClose())); child.on("error", (error) => { this.log( @@ -1450,6 +1641,7 @@ export class DaemonSupervisor { supervisorSocketPath: this.socketPath, authenticationToken: token, rootActiveSessionId, + ownerClientId: existing?.descriptor.ownerClientId ?? ownerClientId, createdAt: existing?.descriptor.createdAt ?? now, updatedAt: now, lifecycle: "starting", @@ -1466,9 +1658,11 @@ export class DaemonSupervisor { snapshotLoads: new Map(), intentionalStop: false, stopRevision: 0, + launchEnv, }; await this.assertRecoveryAllowed(); worker.descriptor = descriptor; + worker.launchEnv = launchEnv; descriptorAssigned = true; this.persistWorker(worker); worker.intentionalStop = false; @@ -1982,6 +2176,12 @@ export class DaemonSupervisor { if (this.isWorkerRecoveryCancelled(worker)) { return; } + if (worker.descriptor.ownerClientId && !worker.launchEnv && !isProcessAlive(worker.descriptor.pid)) { + worker.descriptor.lifecycle = "failed"; + worker.descriptor.lastError = "Waiting for the owning client to reconnect"; + this.persistWorker(worker); + return; + } if (worker.recovery) { return worker.recovery; } @@ -2208,7 +2408,9 @@ export class DaemonSupervisor { .then(async () => { const readyWorkers = [...this.workers.values()].filter( (worker): worker is ResidentWorker & { client: DaemonWorkerClient } => - worker.descriptor.lifecycle === "ready" && worker.client !== undefined, + this.isVisibleWorker(worker) && + worker.descriptor.lifecycle === "ready" && + worker.client !== undefined, ); await Promise.all( readyWorkers.map(async (worker) => { @@ -2228,6 +2430,10 @@ export class DaemonSupervisor { return sync; } + private isVisibleWorker(worker: ResidentWorker): boolean { + return worker.descriptor.ownerClientId === undefined; + } + private agentPeerSummary(summary: SessionSummary): AgentSessionMessageAgentSummary { return { activeSessionId: summary.activeSessionId ?? summary.id, @@ -2253,13 +2459,16 @@ export class DaemonSupervisor { }; } - private async findWorker(selector: string): Promise { - let matches = this.matchWorkers(selector); + private async findWorker( + selector: string, + includeWorker?: (worker: ResidentWorker) => boolean, + ): Promise { + let matches = this.matchWorkers(selector, includeWorker); if (matches.length === 0) { await Promise.all( [...this.workers.values()].map((worker) => this.refreshWorkerSummaries(worker).catch(() => undefined)), ); - matches = this.matchWorkers(selector); + matches = this.matchWorkers(selector, includeWorker); } if (matches.length === 1) { return matches[0]!; @@ -2270,10 +2479,30 @@ export class DaemonSupervisor { throw new Error(`Unknown active session: ${selector}`); } - private matchWorkers(selector: string): WorkerMatch[] { + private findWorkerForClient(client: DaemonSocketClient, selector: string): Promise { + return this.findWorker(selector, (worker) => this.isWorkerAccessibleToClient(client, worker)); + } + + private isWorkerAccessibleToClient(client: DaemonSocketClient, worker: ResidentWorker): boolean { + return ( + worker.descriptor.ownerClientId === undefined || + worker.descriptor.ownerClientId === this.protocolClientId(client) + ); + } + + private assertWorkerAccessibleToClient(client: DaemonSocketClient, worker: ResidentWorker, selector: string): void { + if (!this.isWorkerAccessibleToClient(client, worker)) { + throw new Error(`Unknown active session: ${selector}`); + } + } + + private matchWorkers(selector: string, includeWorker?: (worker: ResidentWorker) => boolean): WorkerMatch[] { const exact: WorkerMatch[] = []; const suffix: WorkerMatch[] = []; for (const worker of this.workers.values()) { + if (includeWorker && !includeWorker(worker)) { + continue; + } for (const summary of worker.summaries.values()) { const activeSessionId = summary.activeSessionId ?? summary.id; const match = { worker, summary }; @@ -2325,7 +2554,31 @@ export class DaemonSupervisor { client: DaemonSocketClient, command: Extract, ): Promise { - const match = await this.findWorker(command.activeSessionId); + const ownedWorker = [...this.workers.values()].find( + (worker) => + worker.descriptor.ownerClientId !== undefined && + (worker.descriptor.rootActiveSessionId === command.activeSessionId || + worker.descriptor.rootSessionId === command.activeSessionId), + ); + if (ownedWorker) { + if (ownedWorker.descriptor.ownerClientId !== this.protocolClientId(client)) { + throw new Error(`Unknown active session: ${command.activeSessionId}`); + } + ownedWorker.launchEnv = command.launchEnv ?? ownedWorker.launchEnv; + if (!ownedWorker.client || ownedWorker.descriptor.lifecycle !== "ready") { + if (!ownedWorker.launchEnv) { + throw new Error("Client-owned session recovery requires the owning client environment"); + } + ownedWorker.intentionalStop = false; + ownedWorker.descriptor.stopRequestedAt = undefined; + ownedWorker.descriptor.archiveOnStop = undefined; + ownedWorker.descriptor.lifecycle = "recovering"; + ownedWorker.descriptor.consecutiveFailures = 0; + this.persistWorker(ownedWorker); + await this.recoverWorker(ownedWorker); + } + } + const match = await this.findWorkerForClient(client, command.activeSessionId); const activeSessionId = match.summary.activeSessionId ?? match.summary.id; const duplicateValidation = this.currentSnapshotGeneration(match.worker, activeSessionId)?.validation; if (duplicateValidation) { @@ -3400,6 +3653,10 @@ export class DaemonSupervisor { recoveryCleanup = false, directChild?: { child: ChildProcess; closed: Promise }, ): Promise { + if (worker.ownerCleanupTimer) { + clearTimeout(worker.ownerCleanupTimer); + worker.ownerCleanupTimer = undefined; + } if (!recoveryCleanup) { worker.stopRevision++; } @@ -3622,6 +3879,10 @@ export class DaemonSupervisor { } this.clients.clear(); for (const worker of this.workers.values()) { + if (worker.ownerCleanupTimer) { + clearTimeout(worker.ownerCleanupTimer); + worker.ownerCleanupTimer = undefined; + } await this.runCleanupStep(`worker client ${worker.descriptor.workerId}`, () => worker.client?.close()); worker.client = undefined; const transcripts = new Set(worker.transcriptCaches.values()); @@ -3729,7 +3990,7 @@ export class DaemonSupervisor { await this.runCleanupStep("daemon ownership", async () => ownership?.release()); if (relaunch) { const launch = createCliSubprocessLaunchSpec(["--mode", "daemon", "--daemon-socket", this.socketPath]); - const environment = { ...process.env }; + const environment = createCliSubprocessEnv(); delete environment[DAEMON_CATALOG_ROLE_ENV]; delete environment[DAEMON_WORKER_ROLE_ENV]; delete environment[DAEMON_WORKER_TOKEN_ENV]; diff --git a/packages/coding-agent/src/modes/daemon/daemon-worker-protocol.ts b/packages/coding-agent/src/modes/daemon/daemon-worker-protocol.ts index 29b55d141ef..f24e3efec6a 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-worker-protocol.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-worker-protocol.ts @@ -86,6 +86,8 @@ export interface DaemonWorkerDescriptor { supervisorSocketPath: string; authenticationToken: string; rootActiveSessionId: string; + /** Stable protocol client that owns this worker. Omitted for resident sessions. */ + ownerClientId?: string; rootSessionId?: string; sessionFile?: string; createdAt: string; diff --git a/packages/coding-agent/src/modes/headless-completion.ts b/packages/coding-agent/src/modes/headless-completion.ts new file mode 100644 index 00000000000..803e9b656a3 --- /dev/null +++ b/packages/coding-agent/src/modes/headless-completion.ts @@ -0,0 +1,78 @@ +import type { AssistantMessage } from "@earendil-works/pi-ai"; +import type { AgentSession } from "../core/agent-session.js"; +import { + type AgentAutonomousStatus, + autonomousLimitReason, + buildAutonomousGateFailureContinuation, +} from "../core/autonomous.js"; + +export function latestAutonomousGateAttempt(status: AgentAutonomousStatus): number { + return Math.max(status.lastGateFailure?.attempt ?? 0, 0, ...Object.values(status.gateAttempts)); +} + +function shouldContinueAutonomousGates(status: AgentAutonomousStatus): boolean { + return ( + status.enabled && + status.gates.commands.length > 0 && + !!status.lastGateFailure && + latestAutonomousGateAttempt(status) <= status.gates.maxRetries && + !autonomousLimitReason(status) + ); +} + +function autonomousProgressKey(status: AgentAutonomousStatus): string { + return [ + latestAutonomousGateAttempt(status), + status.continuationsUsed, + status.turnsUsed, + status.tokensUsed, + status.lastGateFailure?.exitText ?? "", + ].join(":"); +} + +export async function waitForHeadlessCompletion(session: AgentSession): Promise { + let lastPromptedProgressKey: string | undefined; + let repeatedProgressPrompts = 0; + while (true) { + await session.agent.waitForIdle(); + const status = session.getAutonomousStatus(); + if (!shouldContinueAutonomousGates(status) || !status.lastGateFailure) { + return status; + } + const progressKey = autonomousProgressKey(status); + if (progressKey === lastPromptedProgressKey) { + repeatedProgressPrompts++; + } else { + repeatedProgressPrompts = 0; + lastPromptedProgressKey = progressKey; + } + if (repeatedProgressPrompts > 0) { + await new Promise((resolve) => setTimeout(resolve, Math.min(1000, repeatedProgressPrompts * 50))); + } + session.recordHostAutonomousContinuation(); + await session.prompt( + buildAutonomousGateFailureContinuation( + { ...status.lastGateFailure, attempt: latestAutonomousGateAttempt(status) }, + status.gates.maxRetries, + ), + { + streamingBehavior: "followUp", + internalPrompt: true, + suppressAutonomousContinuation: true, + }, + ); + await session.agent.waitForIdle(); + await session.refreshAutonomousGates(); + const lastMessage = session.state.messages.at(-1); + if (lastMessage?.role === "assistant") { + const assistantMessage = lastMessage as AssistantMessage; + if (assistantMessage.stopReason === "error" || assistantMessage.stopReason === "aborted") { + const postErrorStatus = session.getAutonomousStatus(); + if (shouldContinueAutonomousGates(postErrorStatus) && postErrorStatus.lastGateFailure) { + continue; + } + return postErrorStatus; + } + } + } +} diff --git a/packages/coding-agent/src/modes/index.ts b/packages/coding-agent/src/modes/index.ts index ea68df90648..dd3ef262964 100644 --- a/packages/coding-agent/src/modes/index.ts +++ b/packages/coding-agent/src/modes/index.ts @@ -31,7 +31,11 @@ export { sectionTitle, shouldShowAgentsViewSession, } from "./agents-view/agents-view-state.js"; -export { DaemonClient, type DaemonClientMessageListener } from "./daemon/daemon-client.js"; +export { + DaemonCapabilityUnavailableError, + DaemonClient, + type DaemonClientMessageListener, +} from "./daemon/daemon-client.js"; export { type DaemonModeOptions, runDaemonMode } from "./daemon/daemon-mode.js"; export type { DaemonArtifactReference, @@ -81,7 +85,7 @@ export { type PromptStash, type PromptStashState, } from "./interactive/prompt-stash-state.js"; -export { type PrintModeOptions, runPrintMode } from "./print-mode.js"; +export { type PrintModeOptions, runPrintMode, runPrintModeWithConnection } from "./print-mode.js"; export { type ModelInfo, RpcClient, type RpcClientOptions, type RpcEventListener } from "./rpc/rpc-client.js"; -export { runRpcMode } from "./rpc/rpc-mode.js"; +export { runRpcMode, runRpcModeWithConnection } from "./rpc/rpc-mode.js"; export type { RpcCommand, RpcResponse, RpcSessionState } from "./rpc/rpc-types.js"; diff --git a/packages/coding-agent/src/modes/print-mode.ts b/packages/coding-agent/src/modes/print-mode.ts index efa82140210..ffcb7da8dbf 100644 --- a/packages/coding-agent/src/modes/print-mode.ts +++ b/packages/coding-agent/src/modes/print-mode.ts @@ -8,14 +8,12 @@ import type { AssistantMessage, ImageContent } from "@earendil-works/pi-ai"; import type { AgentSessionRuntime } from "../core/agent-session-runtime.js"; -import { - type AgentAutonomousStatus, - type AutonomousLimitReason, - autonomousLimitReason, - buildAutonomousGateFailureContinuation, -} from "../core/autonomous.js"; +import { type AgentAutonomousStatus, type AutonomousLimitReason, autonomousLimitReason } from "../core/autonomous.js"; import { flushRawStdout, writeRawStdout } from "../core/output-guard.js"; import { killTrackedDetachedChildren } from "../utils/shell.js"; +import { InProcessAgentConnection } from "./agent-connection/in-process-agent-connection.js"; +import type { AgentConnection } from "./agent-connection/types.js"; +import { latestAutonomousGateAttempt } from "./headless-completion.js"; /** * Options for print mode. @@ -31,10 +29,6 @@ export interface PrintModeOptions { initialImages?: ImageContent[]; } -function latestGateAttempt(status: AgentAutonomousStatus): number { - return Math.max(status.lastGateFailure?.attempt ?? 0, 0, ...Object.values(status.gateAttempts)); -} - function describeAutonomousLimit(status: AgentAutonomousStatus, reason: AutonomousLimitReason): string { if (reason === "maxContinuations") { return `maxContinuations reached (${status.continuationsUsed}/${status.limits.maxContinuations})`; @@ -49,198 +43,86 @@ function describeAutonomousLimit(status: AgentAutonomousStatus, reason: Autonomo return `timeoutMs reached (${elapsed}/${status.limits.timeoutMs})`; } -function shouldContinuePrintModeAutonomousGates(status: AgentAutonomousStatus): boolean { - return ( - status.enabled && - status.gates.commands.length > 0 && - !!status.lastGateFailure && - latestGateAttempt(status) <= status.gates.maxRetries && - !autonomousLimitReason(status) - ); -} - -function printModeAutonomousProgressKey(status: AgentAutonomousStatus): string { - return [ - latestGateAttempt(status), - status.continuationsUsed, - status.turnsUsed, - status.tokensUsed, - status.lastGateFailure?.exitText ?? "", - ].join(":"); -} - -async function waitForPrintModeIdleWithAutonomousGates( - getSession: () => AgentSessionRuntime["session"], -): Promise { - let lastPromptedProgressKey: string | undefined; - let repeatedProgressPrompts = 0; - while (true) { - const session = getSession(); - await session.agent.waitForIdle(); - const status = session.getAutonomousStatus(); - if (!shouldContinuePrintModeAutonomousGates(status) || !status.lastGateFailure) { - return; - } - const progressKey = printModeAutonomousProgressKey(status); - if (progressKey === lastPromptedProgressKey) { - repeatedProgressPrompts++; - } else { - repeatedProgressPrompts = 0; - lastPromptedProgressKey = progressKey; - } - // Print mode is the host loop for autonomous verifier runs. A still-failing - // gate is not terminal while the configured autonomous budgets remain. Keep - // feeding the failure back even if the observable status key repeats; the - // normal autonomous limits above bound retries, turns, tokens, and wall time. - // The small yield prevents a tight loop if a test double or broken session - // returns immediately without advancing state. - if (repeatedProgressPrompts > 0) { - await new Promise((resolve) => setTimeout(resolve, Math.min(1000, repeatedProgressPrompts * 50))); - } - session.recordHostAutonomousContinuation(); - await session.prompt( - buildAutonomousGateFailureContinuation( - { ...status.lastGateFailure, attempt: latestGateAttempt(status) }, - status.gates.maxRetries, - ), - { - streamingBehavior: "followUp", - internalPrompt: true, - suppressAutonomousContinuation: true, - }, - ); - // followUp prompts can be accepted/queued before the actual retry turn has - // produced a new assistant message. Wait for the session to become idle again - // so an earlier assistant error does not get mistaken for this retry result. - await session.agent.waitForIdle(); - await session.refreshAutonomousGates(); - const lastMessage = session.state.messages[session.state.messages.length - 1]; - if (lastMessage?.role === "assistant") { - const assistantMessage = lastMessage as AssistantMessage; - if (assistantMessage.stopReason === "error" || assistantMessage.stopReason === "aborted") { - const postErrorStatus = session.getAutonomousStatus(); - if (shouldContinuePrintModeAutonomousGates(postErrorStatus) && postErrorStatus.lastGateFailure) { - continue; - } - return; - } - } - } -} - /** * Run in print (single-shot) mode. * Sends prompts to the agent and outputs the result. */ export async function runPrintMode(runtimeHost: AgentSessionRuntime, options: PrintModeOptions): Promise { + const connection = new InProcessAgentConnection(runtimeHost); + return runPrintModeWithConnectionInternal(connection, options, () => connection.bindHeadlessExtensions()); +} + +export async function runPrintModeWithConnection( + connection: AgentConnection, + options: PrintModeOptions, +): Promise { + return runPrintModeWithConnectionInternal(connection, options); +} + +async function runPrintModeWithConnectionInternal( + connection: AgentConnection, + options: PrintModeOptions, + bindHeadlessExtensions?: () => Promise, +): Promise { const { mode, messages = [], initialMessage, initialImages } = options; let exitCode = 0; - let session = runtimeHost.session; - let unsubscribe: (() => void) | undefined; let disposed = false; + let unsubscribe: (() => void) | undefined; const signalCleanupHandlers: Array<() => void> = []; - const disposeRuntime = async (): Promise => { + const disposeConnection = async (): Promise => { if (disposed) return; disposed = true; unsubscribe?.(); - await runtimeHost.dispose(); + await connection.dispose(); }; - const registerSignalHandlers = (): void => { - const signals: NodeJS.Signals[] = ["SIGTERM"]; - if (process.platform !== "win32") { - signals.push("SIGHUP"); - } - - for (const signal of signals) { - const handler = () => { - killTrackedDetachedChildren(); - void disposeRuntime().finally(() => { - process.exit(signal === "SIGHUP" ? 129 : 143); - }); - }; - process.on(signal, handler); - signalCleanupHandlers.push(() => process.off(signal, handler)); - } - }; - - registerSignalHandlers(); - - runtimeHost.setRebindSession(async () => { - await rebindSession(); - }); - - const rebindSession = async (): Promise => { - session = runtimeHost.session; - await session.bindExtensions({ - commandContextActions: { - waitForIdle: () => session.agent.waitForIdle(), - newSession: async (newSessionOptions) => runtimeHost.newSession(newSessionOptions), - fork: async (entryId, forkOptions) => { - const result = await runtimeHost.fork(entryId, forkOptions); - return { cancelled: result.cancelled }; - }, - navigateTree: async (targetId, navigateOptions) => { - const result = await session.navigateTree(targetId, { - summarize: navigateOptions?.summarize, - customInstructions: navigateOptions?.customInstructions, - replaceInstructions: navigateOptions?.replaceInstructions, - label: navigateOptions?.label, - }); - return { cancelled: result.cancelled }; - }, - switchSession: async (sessionPath, switchOptions) => { - return runtimeHost.switchSession(sessionPath, switchOptions); - }, - reload: async () => { - await session.reload(); - }, - }, - onError: (err) => { - console.error(`Extension error (${err.extensionPath}): ${err.error}`); - }, - }); - - unsubscribe?.(); - unsubscribe = session.subscribe((event) => { - if (mode === "json") { - writeRawStdout(`${JSON.stringify(event)}\n`); - } - }); - }; + for (const signal of ["SIGTERM", ...(process.platform === "win32" ? [] : ["SIGHUP"])] as NodeJS.Signals[]) { + const handler = () => { + killTrackedDetachedChildren(); + void disposeConnection().finally(() => { + process.exit(signal === "SIGHUP" ? 129 : 143); + }); + }; + process.on(signal, handler); + signalCleanupHandlers.push(() => process.off(signal, handler)); + } try { if (mode === "json") { - const header = session.sessionManager.getHeader(); + const header = await connection.getSessionHeader(); if (header) { writeRawStdout(`${JSON.stringify(header)}\n`); } } - await rebindSession(); + unsubscribe = connection.subscribe((event) => { + if (mode === "json" && event.type === "session_event") { + writeRawStdout(`${JSON.stringify(event.event)}\n`); + } + if (event.type === "extension_error") { + console.error(`Extension error (${event.extensionPath}): ${event.error}`); + } + }); + await bindHeadlessExtensions?.(); if (initialMessage) { - await session.prompt(initialMessage, { images: initialImages }); + await connection.promptAndWait(initialMessage, { images: initialImages }); } - for (const message of messages) { - await session.prompt(message); + await connection.promptAndWait(message); } - await waitForPrintModeIdleWithAutonomousGates(() => session); - + const autonomousStatus = await connection.waitForHeadlessCompletion(); if (mode === "text") { - const state = session.state; - const lastMessage = state.messages[state.messages.length - 1]; - + const lastMessage = (await connection.getMessages()).at(-1); if (lastMessage?.role === "assistant") { - const assistantMsg = lastMessage as AssistantMessage; - if (assistantMsg.stopReason === "error" || assistantMsg.stopReason === "aborted") { - console.error(assistantMsg.errorMessage || `Request ${assistantMsg.stopReason}`); + const assistantMessage = lastMessage as AssistantMessage; + if (assistantMessage.stopReason === "error" || assistantMessage.stopReason === "aborted") { + console.error(assistantMessage.errorMessage || `Request ${assistantMessage.stopReason}`); exitCode = 1; } else { - for (const content of assistantMsg.content) { + for (const content of assistantMessage.content) { if (content.type === "text") { writeRawStdout(`${content.text}\n`); } @@ -249,14 +131,13 @@ export async function runPrintMode(runtimeHost: AgentSessionRuntime, options: Pr } } - const autonomousStatus = session.getAutonomousStatus(); const autonomousLimit = autonomousLimitReason(autonomousStatus); if (autonomousStatus.enabled && autonomousStatus.gates.commands.length > 0 && autonomousStatus.lastGateFailure) { const limitText = autonomousLimit ? `; autonomous limit reached: ${describeAutonomousLimit(autonomousStatus, autonomousLimit)}` : ""; console.error( - `Autonomous quality gate still failing after attempt ${latestGateAttempt(autonomousStatus)}/${autonomousStatus.gates.maxRetries}: ${autonomousStatus.lastGateFailure.exitText}${limitText}`, + `Autonomous quality gate still failing after attempt ${latestAutonomousGateAttempt(autonomousStatus)}/${autonomousStatus.gates.maxRetries}: ${autonomousStatus.lastGateFailure.exitText}${limitText}`, ); exitCode = 1; } else if (autonomousStatus.enabled && autonomousStatus.gates.commands.length === 0 && autonomousLimit) { @@ -274,7 +155,7 @@ export async function runPrintMode(runtimeHost: AgentSessionRuntime, options: Pr for (const cleanup of signalCleanupHandlers) { cleanup(); } - await disposeRuntime(); + await disposeConnection(); await flushRawStdout(); } } diff --git a/packages/coding-agent/src/modes/rpc/jsonl.ts b/packages/coding-agent/src/modes/rpc/jsonl.ts index b1fe98f0919..6ecb62f0a9c 100644 --- a/packages/coding-agent/src/modes/rpc/jsonl.ts +++ b/packages/coding-agent/src/modes/rpc/jsonl.ts @@ -11,6 +11,11 @@ export function serializeJsonLine(value: unknown): string { return `${JSON.stringify(value)}\n`; } +export interface JsonlLineReaderOptions { + maxLineLength?: number; + onLineOverflow?: (prefix: string) => void; +} + /** * Attach an LF-only JSONL reader to a stream. * @@ -18,7 +23,11 @@ export function serializeJsonLine(value: unknown): string { * Unicode separators that are valid inside JSON strings and therefore does not * implement strict JSONL framing. */ -export function attachJsonlLineReader(stream: Readable, onLine: (line: string) => void): () => void { +export function attachJsonlLineReader( + stream: Readable, + onLine: (line: string) => void, + options: JsonlLineReaderOptions = {}, +): () => void { const decoder = new StringDecoder("utf8"); // Segments of the current, not-yet-terminated line. We never concatenate into a // single growing buffer: appending to one accumulator and rescanning it with @@ -27,19 +36,48 @@ export function attachJsonlLineReader(stream: Readable, onLine: (line: string) = // Instead each chunk is scanned once (offset-advancing indexOf) and the segments // are joined exactly once, when the terminating newline arrives. let pending: string[] = []; + let pendingLength = 0; + let discardingOverflow = false; const emitLine = (line: string) => { onLine(line.endsWith("\r") ? line.slice(0, -1) : line); }; + const resetPending = () => { + pending = []; + pendingLength = 0; + }; + + const appendPending = (segment: string) => { + if (discardingOverflow || segment.length === 0) return; + const maxLineLength = options.maxLineLength; + if (maxLineLength !== undefined && pendingLength + segment.length > maxLineLength) { + const remaining = Math.max(0, maxLineLength - pendingLength); + if (remaining > 0) { + pending.push(segment.slice(0, remaining)); + } + options.onLineOverflow?.(pending.join("")); + resetPending(); + discardingOverflow = true; + return; + } + pending.push(segment); + pendingLength += segment.length; + }; + const emitFrom = (segment: string) => { - if (pending.length === 0) { - emitLine(segment); - } else { - pending.push(segment); - emitLine(pending.join("")); - pending = []; + if (discardingOverflow) { + discardingOverflow = false; + resetPending(); + return; + } + appendPending(segment); + if (discardingOverflow) { + discardingOverflow = false; + return; } + emitLine(pending.join("")); + resetPending(); }; const onData = (chunk: string | Buffer) => { @@ -53,19 +91,20 @@ export function attachJsonlLineReader(stream: Readable, onLine: (line: string) = newlineIndex = text.indexOf("\n", start); } if (start < text.length) { - pending.push(text.slice(start)); + appendPending(text.slice(start)); } }; const onEnd = () => { const tail = decoder.end(); if (tail.length > 0) { - pending.push(tail); + appendPending(tail); } - if (pending.length > 0) { + if (!discardingOverflow && pending.length > 0) { emitLine(pending.join("")); - pending = []; } + resetPending(); + discardingOverflow = false; }; stream.on("data", onData); diff --git a/packages/coding-agent/src/modes/rpc/rpc-client.ts b/packages/coding-agent/src/modes/rpc/rpc-client.ts index ffbd62ff584..fc29b2722d3 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-client.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-client.ts @@ -7,12 +7,30 @@ import { type ChildProcess, spawn } from "node:child_process"; import type { AgentEvent, AgentMessage, ThinkingLevel } from "@earendil-works/pi-agent-core"; import type { ImageContent } from "@earendil-works/pi-ai"; +import type { + AgentSessionMessageDeliveryMode, + AgentSessionMessageReceipt, + AgentSessionMessageSafetyStatus, +} from "../../core/agent-messages.js"; import type { BashResult } from "../../core/bash-executor.js"; import type { CompactionResult } from "../../core/compaction/index.js"; +import type { + AgentCronJob, + AgentHeartbeatDeliveryMode, + AgentHeartbeatManagementAction, + AgentHeartbeatUpdateAction, +} from "../../core/cron-jobs.js"; import type { RefinementResult } from "../../core/refinement/index.js"; import type { SessionStats } from "../../core/session-stats.js"; +import type { AgentConnectionHeartbeat } from "../agent-connection/types.js"; import { attachJsonlLineReader, serializeJsonLine } from "./jsonl.js"; -import type { RpcCommand, RpcResponse, RpcSessionState, RpcSlashCommand } from "./rpc-types.js"; +import type { + RpcCommand, + RpcObservedSessionEvent, + RpcResponse, + RpcSessionState, + RpcSlashCommand, +} from "./rpc-types.js"; // ============================================================================ // Types @@ -50,6 +68,7 @@ export interface ModelInfo { } export type RpcEventListener = (event: AgentEvent) => void; +export type RpcObservedSessionListener = (event: RpcObservedSessionEvent) => void; // ============================================================================ // RPC Client @@ -59,6 +78,7 @@ export class RpcClient { private process: ChildProcess | null = null; private stopReadingStdout: (() => void) | null = null; private eventListeners: RpcEventListener[] = []; + private observedSessionListeners: RpcObservedSessionListener[] = []; private pendingRequests: Map void; reject: (error: Error) => void }> = new Map(); private requestId = 0; @@ -152,6 +172,16 @@ export class RpcClient { }; } + onObservedSessionEvent(listener: RpcObservedSessionListener): () => void { + this.observedSessionListeners.push(listener); + return () => { + const index = this.observedSessionListeners.indexOf(listener); + if (index !== -1) { + this.observedSessionListeners.splice(index, 1); + } + }; + } + /** * Get collected stderr output (useful for debugging). */ @@ -407,6 +437,102 @@ export class RpcClient { return this.getData<{ messages: AgentMessage[] }>(response).messages; } + async sendAgentMessage( + targetActiveSessionId: string, + message: string, + deliveryMode?: AgentSessionMessageDeliveryMode, + ): Promise { + const response = await this.send({ + type: "send_message", + targetActiveSessionId, + message, + deliveryMode, + }); + return this.getData(response); + } + + async getAgentMessageStatus(): Promise { + return this.getData(await this.send({ type: "agent_messages_status" })); + } + + async pauseAgentMessages(): Promise { + return this.getData(await this.send({ type: "agent_messages_pause" })); + } + + async resumeAgentMessages(): Promise { + return this.getData(await this.send({ type: "agent_messages_resume" })); + } + + async clearAgentMessages(): Promise { + const response = await this.send({ type: "agent_messages_clear" }); + return this.getData<{ cleared: number }>(response).cleared; + } + + async listSchedules(includeInactive?: boolean): Promise { + const response = await this.send({ type: "list_schedules", includeInactive }); + return this.getData<{ jobs: AgentCronJob[] }>(response).jobs; + } + + async addSchedule(schedule: string, prompt: string): Promise { + const response = await this.send({ type: "add_schedule", schedule, prompt }); + return this.getData<{ job: AgentCronJob }>(response).job; + } + + async cancelSchedule(jobId: string): Promise { + const response = await this.send({ type: "cancel_schedule", jobId }); + return this.getData<{ job: AgentCronJob }>(response).job; + } + + async listHeartbeats(): Promise { + const response = await this.send({ type: "list_heartbeats" }); + return this.getData<{ heartbeats: AgentConnectionHeartbeat[] }>(response).heartbeats; + } + + async getHeartbeat(): Promise { + const response = await this.send({ type: "get_heartbeat" }); + return this.getData<{ heartbeat: AgentCronJob | null }>(response).heartbeat; + } + + async setHeartbeat( + schedule: string, + prompt: string, + deliveryMode?: AgentHeartbeatDeliveryMode, + ): Promise { + const response = await this.send({ type: "set_heartbeat", schedule, prompt, deliveryMode }); + const heartbeat = this.getData<{ heartbeat: AgentCronJob | null }>(response).heartbeat; + if (!heartbeat) { + throw new Error("Daemon did not return the created heartbeat"); + } + return heartbeat; + } + + async updateHeartbeat(action: AgentHeartbeatUpdateAction): Promise { + const response = await this.send({ type: "update_heartbeat", action }); + return this.getData<{ heartbeat: AgentCronJob | null }>(response).heartbeat; + } + + async manageHeartbeat( + activeSessionId: string, + jobId: string, + action: AgentHeartbeatManagementAction, + ): Promise { + const response = await this.send({ type: "manage_heartbeat", activeSessionId, jobId, action }); + const heartbeat = this.getData<{ heartbeat: AgentCronJob | null }>(response).heartbeat; + if (!heartbeat) { + throw new Error("Daemon did not return the managed heartbeat"); + } + return heartbeat; + } + + async observe(activeSessionId: string): Promise { + const response = await this.send({ type: "observe", activeSessionId }); + return this.getData<{ messages: AgentMessage[] }>(response).messages; + } + + async unobserve(activeSessionId: string): Promise { + await this.send({ type: "unobserve", activeSessionId }); + } + /** * Get available commands (extension commands, prompt templates, skills). */ @@ -486,10 +612,24 @@ export class RpcClient { pending.resolve(data as RpcResponse); return; } + if (data.type === "observed_session_event" || data.type === "observed_session_closed") { + for (const listener of [...this.observedSessionListeners]) { + try { + listener(data as RpcObservedSessionEvent); + } catch { + // Listener failures must not block other RPC subscribers. + } + } + return; + } // Otherwise it's an event - for (const listener of this.eventListeners) { - listener(data as AgentEvent); + for (const listener of [...this.eventListeners]) { + try { + listener(data as AgentEvent); + } catch { + // Listener failures must not block other RPC subscribers. + } } } catch { // Ignore non-JSON lines diff --git a/packages/coding-agent/src/modes/rpc/rpc-extension-ui-context.ts b/packages/coding-agent/src/modes/rpc/rpc-extension-ui-context.ts new file mode 100644 index 00000000000..d66e874a8e4 --- /dev/null +++ b/packages/coding-agent/src/modes/rpc/rpc-extension-ui-context.ts @@ -0,0 +1,157 @@ +import { randomUUID } from "node:crypto"; +import type { + ExtensionUIContext, + ExtensionUIDialogOptions, + ExtensionWidgetOptions, + WorkingIndicatorOptions, +} from "../../core/extensions/index.js"; +import { type Theme, theme } from "../interactive/theme/theme.js"; +import type { RpcExtensionUIRequest, RpcExtensionUIResponse } from "./rpc-types.js"; + +export interface RpcExtensionUiBridge { + uiContext: ExtensionUIContext; + handleResponse(response: RpcExtensionUIResponse): boolean; + close(): void; +} + +export function createRpcExtensionUiBridge(output: (request: RpcExtensionUIRequest) => void): RpcExtensionUiBridge { + const pending = new Map void>(); + let closed = false; + + const createDialogPromise = ( + opts: ExtensionUIDialogOptions | undefined, + defaultValue: T, + request: Record, + parseResponse: (response: RpcExtensionUIResponse) => T, + ): Promise => { + if (closed || opts?.signal?.aborted) { + return Promise.resolve(defaultValue); + } + const id = randomUUID(); + return new Promise((resolve) => { + let timeout: ReturnType | undefined; + const cleanup = () => { + if (timeout) clearTimeout(timeout); + opts?.signal?.removeEventListener("abort", onAbort); + pending.delete(id); + }; + const onAbort = () => { + cleanup(); + resolve(defaultValue); + }; + opts?.signal?.addEventListener("abort", onAbort, { once: true }); + if (opts?.timeout) { + timeout = setTimeout(onAbort, opts.timeout); + } + pending.set(id, (response) => { + cleanup(); + resolve(parseResponse(response)); + }); + output({ type: "extension_ui_request", id, ...request } as RpcExtensionUIRequest); + }); + }; + + const fireAndForget = (request: Record) => { + output({ type: "extension_ui_request", id: randomUUID(), ...request } as RpcExtensionUIRequest); + }; + + const uiContext: ExtensionUIContext = { + select: (title, options, opts) => + createDialogPromise( + opts, + undefined, + { method: "select", title, options, timeout: opts?.timeout }, + (response) => + "cancelled" in response && response.cancelled + ? undefined + : "value" in response + ? response.value + : undefined, + ), + confirm: (title, message, opts) => + createDialogPromise(opts, false, { method: "confirm", title, message, timeout: opts?.timeout }, (response) => + "cancelled" in response && response.cancelled + ? false + : "confirmed" in response + ? response.confirmed + : false, + ), + input: (title, placeholder, opts) => + createDialogPromise( + opts, + undefined, + { method: "input", title, placeholder, timeout: opts?.timeout }, + (response) => + "cancelled" in response && response.cancelled + ? undefined + : "value" in response + ? response.value + : undefined, + ), + notify: (message, notifyType) => fireAndForget({ method: "notify", message, notifyType }), + onTerminalInput: () => () => {}, + setStatus: (statusKey, statusText) => fireAndForget({ method: "setStatus", statusKey, statusText }), + setWorkingMessage: (_message?: string) => {}, + setWorkingVisible: (_visible: boolean) => {}, + setWorkingIndicator: (_options?: WorkingIndicatorOptions) => {}, + setHiddenThinkingLabel: (_label?: string) => {}, + setWidget: (widgetKey: string, content: unknown, options?: ExtensionWidgetOptions) => { + if (content === undefined || Array.isArray(content)) { + fireAndForget({ + method: "setWidget", + widgetKey, + widgetLines: content as string[] | undefined, + widgetPlacement: options?.placement, + }); + } + }, + setFooter: (_factory: unknown) => {}, + setHeader: (_factory: unknown) => {}, + setTitle: (title) => fireAndForget({ method: "setTitle", title }), + custom: async () => undefined as never, + pasteToEditor(text) { + this.setEditorText(text); + }, + setEditorText: (text) => fireAndForget({ method: "set_editor_text", text }), + getEditorText: () => "", + editor: (title, prefill) => + createDialogPromise(undefined, undefined, { method: "editor", title, prefill }, (response) => + "cancelled" in response && response.cancelled + ? undefined + : "value" in response + ? response.value + : undefined, + ), + addAutocompleteProvider: () => {}, + setEditorComponent: () => {}, + getEditorComponent: () => undefined, + get theme() { + return theme; + }, + getAllThemes: () => [], + getTheme: (_name: string) => undefined, + setTheme: (_theme: string | Theme) => ({ success: false, error: "Theme switching not supported in RPC mode" }), + getToolsExpanded: () => false, + setToolsExpanded: (_expanded: boolean) => {}, + }; + + return { + uiContext, + handleResponse(response) { + const resolve = pending.get(response.id); + if (!resolve) { + return false; + } + pending.delete(response.id); + resolve(response); + return true; + }, + close() { + closed = true; + for (const [id, resolve] of pending) { + resolve({ type: "extension_ui_response", id, cancelled: true }); + } + pending.clear(); + }, + }; +} diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index 2e64d14262d..637f72948f3 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -1,28 +1,18 @@ /** - * RPC mode: Headless operation with JSON stdin/stdout protocol. - * - * Used for embedding the agent in other applications. - * Receives commands as JSON on stdin, outputs events and responses as JSON on stdout. - * - * Protocol: - * - Commands: JSON objects with `type` field, optional `id` for correlation - * - Responses: JSON objects with `type: "response"`, `command`, `success`, and optional `data`/`error` - * - Events: AgentSessionEvent objects streamed as they occur - * - Extension UI: Extension UI requests are emitted, client responds with extension_ui_response + * RPC mode: headless operation with JSON commands on stdin and JSON responses/events on stdout. */ -import * as crypto from "node:crypto"; import type { AgentSessionRuntime } from "../../core/agent-session-runtime.js"; -import type { - ExtensionUIContext, - ExtensionUIDialogOptions, - ExtensionWidgetOptions, - WorkingIndicatorOptions, -} from "../../core/extensions/index.js"; import { takeOverStdout, writeRawStdout } from "../../core/output-guard.js"; import { killTrackedDetachedChildren } from "../../utils/shell.js"; -import { type Theme, theme } from "../interactive/theme/theme.js"; +import { InProcessAgentConnection } from "../agent-connection/in-process-agent-connection.js"; +import type { + AgentConnection, + AgentConnectionExtensionUiResponse, + AgentConnectionSessionWatcher, +} from "../agent-connection/types.js"; import { attachJsonlLineReader, serializeJsonLine } from "./jsonl.js"; +import { createRpcExtensionUiBridge } from "./rpc-extension-ui-context.js"; import type { RpcCommand, RpcExtensionUIRequest, @@ -32,636 +22,416 @@ import type { RpcSlashCommand, } from "./rpc-types.js"; -// Re-export types for consumers export type { RpcCommand, RpcExtensionUIRequest, RpcExtensionUIResponse, + RpcObservedSessionEvent, RpcResponse, RpcSessionState, } from "./rpc-types.js"; -/** - * Run in RPC mode. - * Listens for JSON commands on stdin, outputs events and responses on stdout. - */ +interface RpcModeConnectionOptions { + bindHeadlessExtensions?: (options: { + uiContext: ReturnType["uiContext"]; + shutdownHandler: () => void; + }) => Promise; +} + export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise { - takeOverStdout(); - let session = runtimeHost.session; - let unsubscribe: (() => void) | undefined; + const connection = new InProcessAgentConnection(runtimeHost); + return runRpcModeWithConnectionInternal(connection, { + bindHeadlessExtensions: (options) => connection.bindHeadlessExtensions(options), + }); +} - const output = (obj: RpcResponse | RpcExtensionUIRequest | object) => { - writeRawStdout(serializeJsonLine(obj)); - }; +export async function runRpcModeWithConnection(connection: AgentConnection): Promise { + return runRpcModeWithConnectionInternal(connection); +} +async function runRpcModeWithConnectionInternal( + connection: AgentConnection, + options: RpcModeConnectionOptions = {}, +): Promise { + takeOverStdout(); + const output = (value: RpcResponse | RpcExtensionUIRequest | object) => { + writeRawStdout(serializeJsonLine(value)); + }; const success = ( id: string | undefined, command: T, data?: object | null, - ): RpcResponse => { - if (data === undefined) { - return { id, type: "response", command, success: true } as RpcResponse; - } - return { id, type: "response", command, success: true, data } as RpcResponse; - }; - - const error = (id: string | undefined, command: string, message: string): RpcResponse => { - return { id, type: "response", command, success: false, error: message }; - }; - - // Pending extension UI requests waiting for response - const pendingExtensionRequests = new Map< - string, - { resolve: (value: any) => void; reject: (error: Error) => void } - >(); + ): RpcResponse => + (data === undefined + ? { id, type: "response", command, success: true } + : { id, type: "response", command, success: true, data }) as RpcResponse; + const error = (id: string | undefined, command: string, message: string): RpcResponse => ({ + id, + type: "response", + command, + success: false, + error: message, + }); - // Shutdown request flag let shutdownRequested = false; let shuttingDown = false; + let detachInput = () => {}; + let inputEnded = false; + let promptResponsePending = false; + let promptCommandTail = Promise.resolve(); + const bufferedConnectionOutputs: object[] = []; + const pendingConnectionUiRequests = new Set(); const signalCleanupHandlers: Array<() => void> = []; - - /** Helper for dialog methods with signal/timeout support */ - function createDialogPromise( - opts: ExtensionUIDialogOptions | undefined, - defaultValue: T, - request: Record, - parseResponse: (response: RpcExtensionUIResponse) => T, - ): Promise { - if (opts?.signal?.aborted) return Promise.resolve(defaultValue); - - const id = crypto.randomUUID(); - return new Promise((resolve, reject) => { - let timeoutId: ReturnType | undefined; - - const cleanup = () => { - if (timeoutId) clearTimeout(timeoutId); - opts?.signal?.removeEventListener("abort", onAbort); - pendingExtensionRequests.delete(id); - }; - - const onAbort = () => { - cleanup(); - resolve(defaultValue); - }; - opts?.signal?.addEventListener("abort", onAbort, { once: true }); - - if (opts?.timeout) { - timeoutId = setTimeout(() => { - cleanup(); - resolve(defaultValue); - }, opts.timeout); - } - - pendingExtensionRequests.set(id, { - resolve: (response: RpcExtensionUIResponse) => { - cleanup(); - resolve(parseResponse(response)); - }, - reject, - }); - output({ type: "extension_ui_request", id, ...request } as RpcExtensionUIRequest); - }); + const extensionUi = createRpcExtensionUiBridge(output); + const observationCommandTails = new Map>(); + + interface ActiveObservation { + watcher: AgentConnectionSessionWatcher; + unsubscribe: () => void; + ready: boolean; + closed: boolean; + pendingEvents: object[]; } + const observations = new Map(); + const isDialogMethod = (method: string) => + method === "select" || method === "confirm" || method === "input" || method === "editor"; - /** - * Create an extension UI context that uses the RPC protocol. - */ - const createExtensionUIContext = (): ExtensionUIContext => ({ - select: (title, options, opts) => - createDialogPromise(opts, undefined, { method: "select", title, options, timeout: opts?.timeout }, (r) => - "cancelled" in r && r.cancelled ? undefined : "value" in r ? r.value : undefined, - ), - - confirm: (title, message, opts) => - createDialogPromise(opts, false, { method: "confirm", title, message, timeout: opts?.timeout }, (r) => - "cancelled" in r && r.cancelled ? false : "confirmed" in r ? r.confirmed : false, - ), - - input: (title, placeholder, opts) => - createDialogPromise(opts, undefined, { method: "input", title, placeholder, timeout: opts?.timeout }, (r) => - "cancelled" in r && r.cancelled ? undefined : "value" in r ? r.value : undefined, - ), - - notify(message: string, type?: "info" | "warning" | "error"): void { - // Fire and forget - no response needed - output({ - type: "extension_ui_request", - id: crypto.randomUUID(), - method: "notify", - message, - notifyType: type, - } as RpcExtensionUIRequest); - }, - - onTerminalInput(): () => void { - // Raw terminal input not supported in RPC mode - return () => {}; - }, - - setStatus(key: string, text: string | undefined): void { - // Fire and forget - no response needed - output({ - type: "extension_ui_request", - id: crypto.randomUUID(), - method: "setStatus", - statusKey: key, - statusText: text, - } as RpcExtensionUIRequest); - }, - - setWorkingMessage(_message?: string): void { - // Working message not supported in RPC mode - requires TUI loader access - }, - - setWorkingVisible(_visible: boolean): void { - // Working visibility not supported in RPC mode - requires TUI loader access - }, - - setWorkingIndicator(_options?: WorkingIndicatorOptions): void { - // Working indicator customization not supported in RPC mode - requires TUI loader access - }, - - setHiddenThinkingLabel(_label?: string): void { - // Hidden thinking label not supported in RPC mode - requires TUI message rendering access - }, + const outputConnectionEvent = (event: object) => { + if (promptResponsePending) { + bufferedConnectionOutputs.push(event); + return; + } + output(event); + }; + const flushConnectionEvents = () => { + for (const event of bufferedConnectionOutputs.splice(0)) { + output(event); + } + }; + const stopObservation = async (activeSessionId: string) => { + const observation = observations.get(activeSessionId); + if (!observation) return; + observations.delete(activeSessionId); + observation.unsubscribe(); + await observation.watcher.close(); + }; + const activateObservation = (activeSessionId: string) => { + const observation = observations.get(activeSessionId); + if (!observation || observation.ready) return; + observation.ready = true; + for (const event of observation.pendingEvents.splice(0)) { + outputConnectionEvent(event); + } + if (observation.closed) { + void stopObservation(activeSessionId); + } + }; - setWidget(key: string, content: unknown, options?: ExtensionWidgetOptions): void { - // Only support string arrays in RPC mode - factory functions are ignored - if (content === undefined || Array.isArray(content)) { + const unsubscribe = connection.subscribe((event) => { + if (event.type === "session_event") { + outputConnectionEvent(event.event); + return; + } + if (event.type === "extension_error") { + outputConnectionEvent({ + type: "extension_error", + extensionPath: event.extensionPath, + event: event.event, + error: event.error, + }); + return; + } + if (event.type === "extension_ui_request") { + const method = event.request.method === "setEditorText" ? "set_editor_text" : event.request.method; + if (inputEnded && isDialogMethod(method)) { + void connection.respondToExtensionUiRequest(event.request.id, { cancelled: true }).catch(() => undefined); + return; + } + if ( + method === "select" || + method === "confirm" || + method === "input" || + method === "editor" || + method === "notify" || + method === "setStatus" || + method === "setWidget" || + method === "setTitle" || + method === "set_editor_text" + ) { + if (isDialogMethod(method)) { + pendingConnectionUiRequests.add(event.request.id); + } output({ type: "extension_ui_request", - id: crypto.randomUUID(), - method: "setWidget", - widgetKey: key, - widgetLines: content as string[] | undefined, - widgetPlacement: options?.placement, - } as RpcExtensionUIRequest); - } - // Component factories are not supported in RPC mode - would need TUI access - }, - - setFooter(_factory: unknown): void { - // Custom footer not supported in RPC mode - requires TUI access - }, - - setHeader(_factory: unknown): void { - // Custom header not supported in RPC mode - requires TUI access - }, - - setTitle(title: string): void { - // Fire and forget - host can implement terminal title control - output({ - type: "extension_ui_request", - id: crypto.randomUUID(), - method: "setTitle", - title, - } as RpcExtensionUIRequest); - }, - - async custom() { - // Custom UI not supported in RPC mode - return undefined as never; - }, - - pasteToEditor(text: string): void { - // Paste handling not supported in RPC mode - falls back to setEditorText - this.setEditorText(text); - }, - - setEditorText(text: string): void { - // Fire and forget - host can implement editor control - output({ - type: "extension_ui_request", - id: crypto.randomUUID(), - method: "set_editor_text", - text, - } as RpcExtensionUIRequest); - }, - - getEditorText(): string { - // Synchronous method can't wait for RPC response - // Host should track editor state locally if needed - return ""; - }, - - async editor(title: string, prefill?: string): Promise { - const id = crypto.randomUUID(); - return new Promise((resolve, reject) => { - pendingExtensionRequests.set(id, { - resolve: (response: RpcExtensionUIResponse) => { - if ("cancelled" in response && response.cancelled) { - resolve(undefined); - } else if ("value" in response) { - resolve(response.value); - } else { - resolve(undefined); - } - }, - reject, + id: event.request.id, + method, + ...event.request.payload, }); - output({ type: "extension_ui_request", id, method: "editor", title, prefill } as RpcExtensionUIRequest); - }); - }, - - addAutocompleteProvider(): void { - // Autocomplete provider composition is not supported in RPC mode - }, - - setEditorComponent(): void { - // Custom editor components not supported in RPC mode - }, - - getEditorComponent() { - // Custom editor components not supported in RPC mode - return undefined; - }, - - get theme() { - return theme; - }, - - getAllThemes() { - return []; - }, - - getTheme(_name: string) { - return undefined; - }, - - setTheme(_theme: string | Theme) { - // Theme switching not supported in RPC mode - return { success: false, error: "Theme switching not supported in RPC mode" }; - }, + } + return; + } + if (event.type === "closed") { + void shutdown(event.error ? 1 : 0); + } + }); - getToolsExpanded() { - // Tool expansion not supported in RPC mode - no TUI - return false; - }, + const cancelPendingExtensionUi = async () => { + extensionUi.close(); + const requestIds = [...pendingConnectionUiRequests]; + pendingConnectionUiRequests.clear(); + await Promise.allSettled(requestIds.map((id) => connection.respondToExtensionUiRequest(id, { cancelled: true }))); + }; - setToolsExpanded(_expanded: boolean) { - // Tool expansion not supported in RPC mode - no TUI - }, - }); + async function shutdown(exitCode = 0): Promise { + if (shuttingDown) { + process.exit(exitCode); + } + shuttingDown = true; + await cancelPendingExtensionUi(); + for (const cleanup of signalCleanupHandlers) cleanup(); + unsubscribe(); + detachInput(); + process.stdin.pause(); + await Promise.allSettled([...observations.keys()].map((activeSessionId) => stopObservation(activeSessionId))); + await connection.dispose(); + process.exit(exitCode); + } - runtimeHost.setRebindSession(async () => { - await rebindSession(); - }); + for (const signal of ["SIGTERM", ...(process.platform === "win32" ? [] : ["SIGHUP"])] as NodeJS.Signals[]) { + const handler = () => { + killTrackedDetachedChildren(); + void shutdown(signal === "SIGHUP" ? 129 : 143); + }; + process.on(signal, handler); + signalCleanupHandlers.push(() => process.off(signal, handler)); + } - const rebindSession = async (): Promise => { - session = runtimeHost.session; - await session.bindExtensions({ - uiContext: createExtensionUIContext(), - commandContextActions: { - waitForIdle: () => session.agent.waitForIdle(), - newSession: async (options) => runtimeHost.newSession(options), - fork: async (entryId, forkOptions) => { - const result = await runtimeHost.fork(entryId, forkOptions); - return { cancelled: result.cancelled }; - }, - navigateTree: async (targetId, options) => { - const result = await session.navigateTree(targetId, { - summarize: options?.summarize, - customInstructions: options?.customInstructions, - replaceInstructions: options?.replaceInstructions, - label: options?.label, - }); - return { cancelled: result.cancelled }; - }, - switchSession: async (sessionPath, options) => { - return runtimeHost.switchSession(sessionPath, options); - }, - reload: async () => { - await session.reload(); - }, - }, + if (options.bindHeadlessExtensions) { + await options.bindHeadlessExtensions({ + uiContext: extensionUi.uiContext, shutdownHandler: () => { shutdownRequested = true; }, - onError: (err) => { - output({ type: "extension_error", extensionPath: err.extensionPath, event: err.event, error: err.error }); - }, - }); - - unsubscribe?.(); - unsubscribe = session.subscribe((event) => { - output(event); }); - }; - - const registerSignalHandlers = (): void => { - const signals: NodeJS.Signals[] = ["SIGTERM"]; - if (process.platform !== "win32") { - signals.push("SIGHUP"); - } - - for (const signal of signals) { - const handler = () => { - killTrackedDetachedChildren(); - void shutdown(signal === "SIGHUP" ? 129 : 143); - }; - process.on(signal, handler); - signalCleanupHandlers.push(() => process.off(signal, handler)); - } - }; - - await rebindSession(); - registerSignalHandlers(); + } - // Handle a single command - const handleCommand = async (command: RpcCommand): Promise => { + const handleCommand = async (command: RpcCommand): Promise => { const id = command.id; - switch (command.type) { - // ================================================================= - // Prompting - // ================================================================= - - case "prompt": { - // Start prompt handling immediately, but emit the authoritative response only after - // prompt preflight succeeds. Queued and immediately handled prompts also count as success. - let preflightSucceeded = false; - void session - .prompt(command.message, { - images: command.images, - streamingBehavior: command.streamingBehavior, - source: "rpc", - preflightResult: (didSucceed) => { - if (didSucceed) { - preflightSucceeded = true; - output(success(id, "prompt")); - } - }, - }) - .catch((e) => { - if (!preflightSucceeded) { - output(error(id, "prompt", e.message)); - } - }); - return undefined; - } - - case "steer": { - await session.steer(command.message, command.images, { resumeIfIdle: true }); - return success(id, "steer"); - } - - case "follow_up": { - await session.followUp(command.message, command.images, { resumeIfIdle: true }); - return success(id, "follow_up"); - } - - case "abort": { - await session.abort(); - return success(id, "abort"); - } - - case "new_session": { - const options = command.parentSession ? { parentSession: command.parentSession } : undefined; - const result = await runtimeHost.newSession(options); - if (!result.cancelled) { - await rebindSession(); - } - return success(id, "new_session", result); - } - - // ================================================================= - // State - // ================================================================= - + case "prompt": + await connection.prompt(command.message, { + images: command.images, + streamingBehavior: command.streamingBehavior, + source: "rpc", + }); + return success(id, command.type); + case "steer": + await connection.steer(command.message, command.images); + return success(id, command.type); + case "follow_up": + await connection.followUp(command.message, command.images); + return success(id, command.type); + case "abort": + await connection.abort(); + return success(id, command.type); + case "new_session": + return success( + id, + command.type, + await connection.newSession( + command.parentSession ? { parentSession: command.parentSession } : undefined, + ), + ); case "get_state": { - const state: RpcSessionState = { - model: session.model, - thinkingLevel: session.thinkingLevel, - isStreaming: session.isStreaming, - isCompacting: session.isCompacting, - steeringMode: session.steeringMode, - followUpMode: session.followUpMode, - sessionFile: session.sessionFile, - sessionId: session.sessionId, - sessionName: session.sessionName, - autoCompactionEnabled: session.autoCompactionEnabled, - messageCount: session.messages.length, - pendingMessageCount: session.pendingMessageCount, - goal: session.goalState, + const state = await connection.getState(); + const rpcState: RpcSessionState = { + model: state.model, + thinkingLevel: state.thinkingLevel, + isStreaming: state.isStreaming, + isCompacting: state.isCompacting, + steeringMode: state.steeringMode, + followUpMode: state.followUpMode, + sessionFile: state.sessionFile, + sessionId: state.sessionId, + sessionName: state.sessionName, + autoCompactionEnabled: state.autoCompactionEnabled, + messageCount: state.messageCount, + pendingMessageCount: state.pendingMessageCount, + goal: state.goal, }; - return success(id, "get_state", state); - } - - // ================================================================= - // Model - // ================================================================= - - case "set_model": { - const models = await session.modelRegistry.refreshAvailableModels(); - const model = models.find((m) => m.provider === command.provider && m.id === command.modelId); - if (!model) { - return error(id, "set_model", `Model not found: ${command.provider}/${command.modelId}`); - } - await session.setModel(model); - return success(id, "set_model", model); - } - - case "cycle_model": { - const result = await session.cycleModel(); - if (!result) { - return success(id, "cycle_model", null); - } - return success(id, "cycle_model", result); - } - - case "get_available_models": { - const models = await session.modelRegistry.refreshAvailableModels(); - return success(id, "get_available_models", { models }); - } - - // ================================================================= - // Thinking - // ================================================================= - - case "set_thinking_level": { - session.setThinkingLevel(command.level); - return success(id, "set_thinking_level"); - } - + return success(id, command.type, rpcState); + } + case "set_model": + return success(id, command.type, await connection.setModel(command.provider, command.modelId)); + case "cycle_model": + return success(id, command.type, (await connection.cycleModel()) ?? null); + case "get_available_models": + return success(id, command.type, { models: await connection.getAvailableModels() }); + case "set_thinking_level": + await connection.setThinkingLevel(command.level); + return success(id, command.type); case "cycle_thinking_level": { - const level = session.cycleThinkingLevel(); - if (!level) { - return success(id, "cycle_thinking_level", null); - } - return success(id, "cycle_thinking_level", { level }); - } - - // ================================================================= - // Queue Modes - // ================================================================= - - case "set_steering_mode": { - session.setSteeringMode(command.mode); - return success(id, "set_steering_mode"); - } - - case "set_follow_up_mode": { - session.setFollowUpMode(command.mode); - return success(id, "set_follow_up_mode"); - } - - // ================================================================= - // Compaction - // ================================================================= - - case "compact": { - const result = await session.compact(command.customInstructions); - return success(id, "compact", result); - } - - case "refine": { - const result = await session.refine({ - instructions: command.instructions, - rollbackId: command.rollbackId, - global: command.global, - }); - return success(id, "refine", result); - } - - case "set_auto_compaction": { - session.setAutoCompactionEnabled(command.enabled); - return success(id, "set_auto_compaction"); - } - - // ================================================================= - // Retry - // ================================================================= - - case "set_auto_retry": { - session.setAutoRetryEnabled(command.enabled); - return success(id, "set_auto_retry"); - } - - case "abort_retry": { - session.abortRetry(); - return success(id, "abort_retry"); - } - - // ================================================================= - // Bash - // ================================================================= - - case "bash": { - const result = await session.executeBash(command.command); - return success(id, "bash", result); - } - - case "abort_bash": { - session.abortBash(); - return success(id, "abort_bash"); - } - - // ================================================================= - // Session - // ================================================================= - - case "get_session_stats": { - const stats = session.getSessionStats(); - return success(id, "get_session_stats", stats); - } - - case "export_html": { - const path = await session.exportToHtml(command.outputPath); - return success(id, "export_html", { path }); - } - - case "switch_session": { - const result = await runtimeHost.switchSession(command.sessionPath); - if (!result.cancelled) { - await rebindSession(); - } - return success(id, "switch_session", result); - } - + const level = await connection.cycleThinkingLevel(); + return success(id, command.type, level ? { level } : null); + } + case "set_steering_mode": + await connection.setSteeringMode(command.mode); + return success(id, command.type); + case "set_follow_up_mode": + await connection.setFollowUpMode(command.mode); + return success(id, command.type); + case "compact": + return success(id, command.type, await connection.compact(command.customInstructions)); + case "refine": + return success( + id, + command.type, + await connection.refine({ + instructions: command.instructions, + rollbackId: command.rollbackId, + global: command.global, + }), + ); + case "set_auto_compaction": + await connection.setAutoCompactionEnabled(command.enabled); + return success(id, command.type); + case "set_auto_retry": + await connection.setAutoRetryEnabled(command.enabled); + return success(id, command.type); + case "abort_retry": + await connection.abortRetry(); + return success(id, command.type); + case "bash": + return success(id, command.type, await connection.executeBashAndWait(command.command)); + case "abort_bash": + await connection.abortBash(); + return success(id, command.type); + case "get_session_stats": + return success(id, command.type, await connection.getSessionStats()); + case "export_html": + return success(id, command.type, { path: await connection.exportToHtml(command.outputPath) }); + case "switch_session": + return success(id, command.type, await connection.switchSession(command.sessionPath)); case "fork": { - const result = await runtimeHost.fork(command.entryId); - if (!result.cancelled) { - await rebindSession(); - } - return success(id, "fork", { text: result.selectedText, cancelled: result.cancelled }); + const result = await connection.fork(command.entryId); + return success(id, command.type, { text: result.selectedText, cancelled: result.cancelled }); } - case "clone": { - const leafId = session.sessionManager.getLeafId(); - if (!leafId) { - return error(id, "clone", "Cannot clone session: no current entry selected"); - } - const result = await runtimeHost.fork(leafId, { position: "at" }); - if (!result.cancelled) { - await rebindSession(); - } - return success(id, "clone", { cancelled: result.cancelled }); - } - - case "get_fork_messages": { - const messages = session.getUserMessagesForForking(); - return success(id, "get_fork_messages", { messages }); - } - - case "get_last_assistant_text": { - const text = session.getLastAssistantText(); - return success(id, "get_last_assistant_text", { text }); - } - + const { leafId } = await connection.getSessionTree(); + if (!leafId) return error(id, command.type, "Cannot clone session: no current entry selected"); + const result = await connection.fork(leafId, { position: "at" }); + return success(id, command.type, { cancelled: result.cancelled }); + } + case "get_fork_messages": + return success(id, command.type, { messages: await connection.getUserMessagesForForking() }); + case "get_last_assistant_text": + return success(id, command.type, { text: await connection.getLastAssistantText() }); case "set_session_name": { const name = command.name.trim(); - if (!name) { - return error(id, "set_session_name", "Session name cannot be empty"); + if (!name) return error(id, command.type, "Session name cannot be empty"); + await connection.setSessionName(name); + return success(id, command.type); + } + case "get_messages": + return success(id, command.type, { messages: await connection.getMessages() }); + case "send_message": + return success( + id, + command.type, + await connection.sendAgentMessage(command.targetActiveSessionId, command.message, command.deliveryMode), + ); + case "agent_messages_status": + return success(id, command.type, await connection.getAgentMessageStatus()); + case "agent_messages_pause": + return success(id, command.type, await connection.pauseAgentMessages()); + case "agent_messages_resume": + return success(id, command.type, await connection.resumeAgentMessages()); + case "agent_messages_clear": + return success(id, command.type, { cleared: await connection.clearAgentMessages() }); + case "list_schedules": + return success(id, command.type, { + jobs: await connection.listCronJobs({ includeInactive: command.includeInactive }), + }); + case "add_schedule": + return success(id, command.type, { + job: await connection.addCronJob(command.schedule, command.prompt), + }); + case "cancel_schedule": + return success(id, command.type, { job: await connection.cancelCronJob(command.jobId) }); + case "list_heartbeats": + return success(id, command.type, { heartbeats: await connection.listHeartbeats() }); + case "get_heartbeat": + return success(id, command.type, { heartbeat: (await connection.getHeartbeat()) ?? null }); + case "set_heartbeat": + return success(id, command.type, { + heartbeat: await connection.setHeartbeat(command.schedule, command.prompt, command.deliveryMode), + }); + case "update_heartbeat": + return success(id, command.type, { + heartbeat: (await connection.updateHeartbeat(command.action)) ?? null, + }); + case "manage_heartbeat": + return success(id, command.type, { + heartbeat: await connection.manageHeartbeat(command.activeSessionId, command.jobId, command.action), + }); + case "observe": { + const existing = observations.get(command.activeSessionId); + if (existing) { + return success(id, command.type, { messages: await existing.watcher.getMessages() }); + } + const watcher = await connection.watchSession(command.activeSessionId); + if (!watcher) throw new Error(`Unknown active session: ${command.activeSessionId}`); + const observation: ActiveObservation = { + watcher, + unsubscribe: () => {}, + ready: false, + closed: false, + pendingEvents: [], + }; + observations.set(command.activeSessionId, observation); + observation.unsubscribe = watcher.subscribe((event) => { + let observed: object | undefined; + if (event.type === "session_event") { + observed = { + type: "observed_session_event", + activeSessionId: command.activeSessionId, + event: event.event, + }; + } else if (event.type === "closed") { + observed = { + type: "observed_session_closed", + activeSessionId: command.activeSessionId, + error: event.error, + }; + observation.closed = true; + } + if (!observed) return; + if (observation.ready) { + outputConnectionEvent(observed); + if (observation.closed) void stopObservation(command.activeSessionId); + } else { + observation.pendingEvents.push(observed); + } + }); + try { + return success(id, command.type, { messages: await watcher.getMessages() }); + } catch (observationError) { + await stopObservation(command.activeSessionId); + throw observationError; } - session.setSessionName(name); - return success(id, "set_session_name"); - } - - // ================================================================= - // Messages - // ================================================================= - - case "get_messages": { - return success(id, "get_messages", { messages: session.messages }); } - - // ================================================================= - // Commands (available for invocation via prompt) - // ================================================================= - + case "unobserve": + await stopObservation(command.activeSessionId); + return success(id, command.type); case "get_commands": { - const commands: RpcSlashCommand[] = []; - - for (const command of session.extensionRunner.getRegisteredCommands()) { - commands.push({ - name: command.invocationName, - description: command.description, - source: "extension", - sourceInfo: command.sourceInfo, - }); - } - - for (const template of session.promptTemplates) { - commands.push({ - name: template.name, - description: template.description, - source: "prompt", - sourceInfo: template.sourceInfo, - }); - } - - for (const skill of session.resourceLoader.getSkills().skills) { - commands.push({ - name: `skill:${skill.name}`, - description: skill.description, - source: "skill", - sourceInfo: skill.sourceInfo, - }); - } - - return success(id, "get_commands", { commands }); + const commands: RpcSlashCommand[] = (await connection.getCommands()).map((available) => ({ + name: available.name, + description: available.description, + source: available.source, + sourceInfo: available.sourceInfo, + })); + return success(id, command.type, { commands }); } - default: { const unknownCommand = command as { type: string }; return error(undefined, unknownCommand.type, `Unknown command: ${unknownCommand.type}`); @@ -669,32 +439,6 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise {}; - - async function shutdown(exitCode = 0): Promise { - if (shuttingDown) { - process.exit(exitCode); - } - shuttingDown = true; - for (const cleanup of signalCleanupHandlers) { - cleanup(); - } - unsubscribe?.(); - await runtimeHost.dispose(); - detachInput(); - process.stdin.pause(); - process.exit(exitCode); - } - - async function checkShutdownRequested(): Promise { - if (!shutdownRequested) return; - await shutdown(); - } - const handleInputLine = async (line: string) => { let parsed: unknown; try { @@ -709,56 +453,101 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise undefined); return; } const command = parsed as RpcCommand; - try { - const response = await handleCommand(command); - if (response) { - output(response); + const executeCommand = async () => { + const isPrompt = command.type === "prompt"; + if (isPrompt) promptResponsePending = true; + try { + output(await handleCommand(command)); + if (command.type === "observe") activateObservation(command.activeSessionId); + } catch (commandError: unknown) { + output( + error( + command.id, + command.type, + commandError instanceof Error ? commandError.message : String(commandError), + ), + ); + } finally { + if (isPrompt) { + promptResponsePending = false; + flushConnectionEvents(); + } } - await checkShutdownRequested(); - } catch (commandError: unknown) { - output( - error( - command.id, - command.type, - commandError instanceof Error ? commandError.message : String(commandError), - ), + if (shutdownRequested) await shutdown(); + }; + if (command.type === "observe" || command.type === "unobserve") { + const previous = observationCommandTails.get(command.activeSessionId) ?? Promise.resolve(); + const ordered = previous.then(executeCommand, executeCommand); + const tail = ordered.then( + () => undefined, + () => undefined, + ); + observationCommandTails.set(command.activeSessionId, tail); + await ordered; + if (observationCommandTails.get(command.activeSessionId) === tail) { + observationCommandTails.delete(command.activeSessionId); + } + return; + } + if (command.type === "prompt") { + const ordered = promptCommandTail.then(executeCommand, executeCommand); + promptCommandTail = ordered.then( + () => undefined, + () => undefined, ); + await ordered; + return; } + await executeCommand(); }; + const pendingInputHandlers = new Set>(); + const dispatchInputLine = (line: string) => { + const pending = handleInputLine(line).finally(() => pendingInputHandlers.delete(pending)); + pendingInputHandlers.add(pending); + }; const onInputEnd = () => { - void shutdown(); + inputEnded = true; + detachInput(); + process.stdin.pause(); + queueMicrotask(() => { + void cancelPendingExtensionUi() + .then(() => Promise.allSettled([...pendingInputHandlers])) + .then(() => connection.waitForIdle()) + .then( + () => shutdown(), + () => shutdown(1), + ); + }); }; process.stdin.on("end", onInputEnd); - detachInput = (() => { - const detachJsonl = attachJsonlLineReader(process.stdin, (line) => { - void handleInputLine(line); - }); + const detachJsonl = attachJsonlLineReader(process.stdin, dispatchInputLine); return () => { detachJsonl(); process.stdin.off("end", onInputEnd); }; })(); - // Keep process alive forever return new Promise(() => {}); } diff --git a/packages/coding-agent/src/modes/rpc/rpc-types.ts b/packages/coding-agent/src/modes/rpc/rpc-types.ts index e768ff9e547..c1ebe266081 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-types.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-types.ts @@ -5,14 +5,25 @@ * Responses and events are emitted as JSON lines on stdout. */ -import type { AgentMessage, ThinkingLevel } from "@earendil-works/pi-agent-core"; +import type { AgentEvent, AgentMessage, ThinkingLevel } from "@earendil-works/pi-agent-core"; import type { ImageContent, Model } from "@earendil-works/pi-ai"; +import type { + AgentSessionMessageDeliveryMode, + AgentSessionMessageReceipt, + AgentSessionMessageSafetyStatus, +} from "../../core/agent-messages.js"; import type { BashResult } from "../../core/bash-executor.js"; import type { CompactionResult } from "../../core/compaction/index.js"; +import type { + AgentCronJob, + AgentHeartbeatDeliveryMode, + AgentHeartbeatManagementAction, + AgentHeartbeatUpdateAction, +} from "../../core/cron-jobs.js"; import type { GoalState } from "../../core/goals.js"; import type { RefinementResult } from "../../core/refinement/index.js"; import type { SessionStats } from "../../core/session-stats.js"; -import type { AgentConnectionSourceInfo } from "../agent-connection/types.js"; +import type { AgentConnectionHeartbeat, AgentConnectionSourceInfo } from "../agent-connection/types.js"; // ============================================================================ // RPC Commands (stdin) @@ -67,6 +78,43 @@ export type RpcCommand = // Messages | { id?: string; type: "get_messages" } + | { + id?: string; + type: "send_message"; + targetActiveSessionId: string; + message: string; + deliveryMode?: AgentSessionMessageDeliveryMode; + } + | { id?: string; type: "agent_messages_status" } + | { id?: string; type: "agent_messages_pause" } + | { id?: string; type: "agent_messages_resume" } + | { id?: string; type: "agent_messages_clear" } + + // Scheduling + | { id?: string; type: "list_schedules"; includeInactive?: boolean } + | { id?: string; type: "add_schedule"; schedule: string; prompt: string } + | { id?: string; type: "cancel_schedule"; jobId: string } + | { id?: string; type: "list_heartbeats" } + | { id?: string; type: "get_heartbeat" } + | { + id?: string; + type: "set_heartbeat"; + schedule: string; + prompt: string; + deliveryMode?: AgentHeartbeatDeliveryMode; + } + | { id?: string; type: "update_heartbeat"; action: AgentHeartbeatUpdateAction } + | { + id?: string; + type: "manage_heartbeat"; + activeSessionId: string; + jobId: string; + action: AgentHeartbeatManagementAction; + } + + // Active session and subagent observation + | { id?: string; type: "observe"; activeSessionId: string } + | { id?: string; type: "unobserve"; activeSessionId: string } // Commands (available for invocation via prompt) | { id?: string; type: "get_commands" }; @@ -197,6 +245,45 @@ export type RpcResponse = // Messages | { id?: string; type: "response"; command: "get_messages"; success: true; data: { messages: AgentMessage[] } } + | { id?: string; type: "response"; command: "send_message"; success: true; data: AgentSessionMessageReceipt } + | { + id?: string; + type: "response"; + command: "agent_messages_status" | "agent_messages_pause" | "agent_messages_resume"; + success: true; + data: AgentSessionMessageSafetyStatus; + } + | { id?: string; type: "response"; command: "agent_messages_clear"; success: true; data: { cleared: number } } + + // Scheduling + | { id?: string; type: "response"; command: "list_schedules"; success: true; data: { jobs: AgentCronJob[] } } + | { id?: string; type: "response"; command: "add_schedule"; success: true; data: { job: AgentCronJob } } + | { id?: string; type: "response"; command: "cancel_schedule"; success: true; data: { job: AgentCronJob } } + | { + id?: string; + type: "response"; + command: "list_heartbeats"; + success: true; + data: { heartbeats: AgentConnectionHeartbeat[] }; + } + | { + id?: string; + type: "response"; + command: "get_heartbeat"; + success: true; + data: { heartbeat: AgentCronJob | null }; + } + | { + id?: string; + type: "response"; + command: "set_heartbeat" | "update_heartbeat" | "manage_heartbeat"; + success: true; + data: { heartbeat: AgentCronJob | null }; + } + + // Active session and subagent observation + | { id?: string; type: "response"; command: "observe"; success: true; data: { messages: AgentMessage[] } } + | { id?: string; type: "response"; command: "unobserve"; success: true } // Commands | { @@ -267,3 +354,7 @@ export type RpcExtensionUIResponse = // ============================================================================ export type RpcCommandType = RpcCommand["type"]; + +export type RpcObservedSessionEvent = + | { type: "observed_session_event"; activeSessionId: string; event: AgentEvent } + | { type: "observed_session_closed"; activeSessionId: string; error?: string }; diff --git a/packages/coding-agent/src/package-manager-cli.ts b/packages/coding-agent/src/package-manager-cli.ts index e3ced746e84..6a64905515a 100644 --- a/packages/coding-agent/src/package-manager-cli.ts +++ b/packages/coding-agent/src/package-manager-cli.ts @@ -500,7 +500,12 @@ function confirmDaemonSessionLossBeforeUpdate(probe: RunningDaemonProbe, force: } function daemonProbeMayHaveBusySessions(probe: RunningDaemonProbe): boolean { - return probe.reachable && (probe.activeSessions === undefined || probe.activeSessions.some(isSessionBusy)); + return ( + probe.reachable && + (probe.activeSessions === undefined || + (probe.busyClientOwnedSessionCount ?? 0) > 0 || + probe.activeSessions.some(isSessionBusy)) + ); } function isRecord(value: unknown): value is Record { diff --git a/packages/coding-agent/test/agent-connection-daemon.test.ts b/packages/coding-agent/test/agent-connection-daemon.test.ts index 83e8c7f8a09..93133fe078a 100644 --- a/packages/coding-agent/test/agent-connection-daemon.test.ts +++ b/packages/coding-agent/test/agent-connection-daemon.test.ts @@ -46,6 +46,7 @@ class FakeDaemonClient { connectionStateFactory: ((activeSessionId: string) => AgentConnectionState) | undefined; abortBashUnknownCommand = false; abortAndClearQueueUnknownCommand = false; + cronAddGate: Promise | undefined; legacyHeartbeatCommandsSupported = false; serverCapabilities = new Set(); updateRestartSessions: Array> = []; @@ -223,7 +224,7 @@ class FakeDaemonClient { data: { steering: ["aborted"], followUp: ["cleared"] }, }; case "heartbeats_list": - return this.legacyHeartbeatCommandsSupported + return this.legacyHeartbeatCommandsSupported || this.serverCapabilities.has("heartbeat_catalog") ? { type: "response", command: command.type, success: true, data: { heartbeats: [] } } : { type: "response", @@ -232,7 +233,7 @@ class FakeDaemonClient { error: "Unknown daemon command: heartbeats_list", }; case "heartbeat_manage": - return this.legacyHeartbeatCommandsSupported + return this.legacyHeartbeatCommandsSupported || this.serverCapabilities.has("heartbeat_management") ? { type: "response", command: command.type, @@ -363,6 +364,29 @@ class FakeDaemonClient { harnessStatePath: "/tmp/harness_state.json", }, }; + case "cron_add": + await this.cronAddGate; + return { + type: "response", + command: command.type, + success: true, + data: { + job: { + id: `cron-${this.requests.filter((request) => request.type === "cron_add").length}`, + status: "active", + source: "cron", + activeSessionId: command.activeSessionId, + sessionId: "session-current", + sessionFile: "/tmp/session-current.jsonl", + cwd: "/tmp/project", + prompt: command.prompt, + schedule: { kind: "cron", expression: command.schedule }, + createdAt: "2026-01-01T00:00:00.000Z", + updatedAt: "2026-01-01T00:00:00.000Z", + runCount: 0, + }, + }, + }; case "switch_session": return { type: "response", @@ -630,6 +654,48 @@ function emitSequencedQueueUpdate(client: FakeDaemonClient, activeSessionId: str } describe("DaemonAgentConnection", () => { + it("uses fleet heartbeat scope for residents and session scope for owned workers", async () => { + const residentClient = new FakeDaemonClient(); + residentClient.serverCapabilities.add("heartbeat_catalog"); + const resident = new DaemonAgentConnection(asDaemonClient(residentClient), "resident-1"); + await resident.listHeartbeats(); + + const ownedClient = new FakeDaemonClient(); + ownedClient.serverCapabilities.add("heartbeat_catalog"); + const owned = new DaemonAgentConnection(asDaemonClient(ownedClient), "owned-1", { ownedSession: true }); + await owned.listHeartbeats(); + + expect(residentClient.requests.at(-1)).toEqual(expect.objectContaining({ type: "heartbeats_list" })); + expect(residentClient.requests.at(-1)).not.toHaveProperty("activeSessionId"); + expect(ownedClient.requests.at(-1)).toEqual( + expect.objectContaining({ type: "heartbeats_list", activeSessionId: "owned-1" }), + ); + }); + + it("serializes concurrent owned-session promotion commands", async () => { + const fakeClient = new FakeDaemonClient(); + let releaseFirst = () => {}; + fakeClient.cronAddGate = new Promise((resolve) => { + releaseFirst = resolve; + }); + const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { + ownedSession: true, + }); + + const first = connection.addCronJob("0 * * * *", "first"); + const second = connection.addCronJob("30 * * * *", "second"); + await vi.waitFor(() => { + expect(fakeClient.requests.filter((request) => request.type === "cron_add")).toHaveLength(1); + }); + releaseFirst(); + await Promise.all([first, second]); + + const commands = fakeClient.requests.filter( + (command): command is Extract => command.type === "cron_add", + ); + expect(commands.map((command) => command.promoteOwnedSession)).toEqual([true, false]); + }); + it("degrades an unavailable heartbeat catalog without sending an unsupported command", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); diff --git a/packages/coding-agent/test/daemon-launch.test.ts b/packages/coding-agent/test/daemon-launch.test.ts index 2afac0fff68..7572499d90e 100644 --- a/packages/coding-agent/test/daemon-launch.test.ts +++ b/packages/coding-agent/test/daemon-launch.test.ts @@ -8,7 +8,8 @@ import { ensureInteractiveDaemonRunning, probeDaemonVersion, probeRunningDaemonSessions, - shouldStartInteractiveDaemonEarly, + shouldStartDaemonEarly, + shouldUseLegacyOwnedSessionWorkerFrontend, shutdownDaemonAndWait, } from "../src/cli/daemon-launch.js"; import { VERSION } from "../src/config.js"; @@ -17,6 +18,7 @@ import { DAEMON_PROTOCOL_VERSION, DAEMON_SCHEMA_ID } from "../src/modes/daemon/d interface FakeDaemonOptions { /** Sessions returned for a `list` command. */ sessions?: Array>; + busyClientOwnedSessionCount?: number; /** When true, the `list` command responds with a failure. */ failList?: boolean; /** When false, the server ignores `shutdown` and stays up. */ @@ -24,6 +26,7 @@ interface FakeDaemonOptions { protocolVersion?: number; appVersion?: string; schemaId?: string; + serverCapabilities?: string[]; onCommand?: (command: { type: string }) => void; } @@ -48,7 +51,7 @@ async function startFakeDaemon(options: FakeDaemonOptions = {}): Promise { @@ -75,7 +78,13 @@ async function startFakeDaemon(options: FakeDaemonOptions = {}): Promise { }); }); -describe("shouldStartInteractiveDaemonEarly", () => { - it("starts early for default interactive use and the agents view", () => { - expect(shouldStartInteractiveDaemonEarly([], true, false)).toBe(true); - expect(shouldStartInteractiveDaemonEarly(["agents"], true, false)).toBe(true); - expect(shouldStartInteractiveDaemonEarly(["review this change"], true, false)).toBe(true); - expect(shouldStartInteractiveDaemonEarly(["help", "me", "fix", "this"], true, false)).toBe(true); +describe("shouldStartDaemonEarly", () => { + it.each([ + ["interactive", []], + ["print", ["--print", "hello"]], + ["json", ["--mode", "json", "hello"]], + ["rpc", ["--mode", "rpc"]], + ["no-session", ["--no-session"]], + ])("starts early for the %s client", (_label, args) => { + expect(shouldStartDaemonEarly(args, false)).toBe(true); }); - it("does not start a service for standalone management commands", () => { - for (const command of [ - "list", - "attach", - "status", - "doctor", - "shutdown", - "package", - "update", - "model", - "session", - ]) { - expect(shouldStartInteractiveDaemonEarly([command], true, false)).toBe(false); - } - }); - - it("does not start a service for help or removed command forms", () => { - expect(shouldStartInteractiveDaemonEarly(["help"], true, false)).toBe(false); - expect(shouldStartInteractiveDaemonEarly(["daemon", "list"], true, false)).toBe(false); + it.each([ + ["daemon process", ["--mode", "daemon"]], + ["help", ["--help"]], + ["version", ["--version"]], + ["model listing", ["--list-models"]], + ["management command after global flags", ["--daemon-socket", "/tmp/prime.sock", "status"]], + ["startup benchmark", []], + ])("does not start early for %s", (label, args) => { + expect(shouldStartDaemonEarly(args, label === "startup benchmark")).toBe(false); }); }); @@ -236,6 +238,34 @@ describe("ensureInteractiveDaemonRunning", () => { expect(commands).not.toContain("shutdown"); }); + it("does not replace a stale daemon with busy private client sessions", async () => { + const commands: string[] = []; + const daemon = await startFakeDaemon({ + protocolVersion: DAEMON_PROTOCOL_VERSION, + appVersion: VERSION, + schemaId: "stale-schema", + busyClientOwnedSessionCount: 1, + onCommand: (command) => commands.push(command.type), + }); + cleanups.push(daemon.close); + + await expect(shouldUseLegacyOwnedSessionWorkerFrontend(daemon.socketPath)).resolves.toBe(true); + await expect(ensureInteractiveDaemonRunning(daemon.socketPath)).rejects.toThrow("stale"); + expect(commands).toContain("list"); + expect(commands).not.toContain("shutdown"); + }); + + it("routes owned clients through the legacy frontend while a compatible legacy daemon is busy", async () => { + const daemon = await startFakeDaemon({ + protocolVersion: 3, + appVersion: VERSION, + sessions: [{ id: "active-1", activeSessionId: "active-1", isStreaming: true }], + }); + cleanups.push(daemon.close); + + await expect(shouldUseLegacyOwnedSessionWorkerFrontend(daemon.socketPath)).resolves.toBe(true); + }); + it("does not treat a live daemon as absent when cold startup delays the first connection", async () => { const daemon = await startFakeDaemon({ protocolVersion: DAEMON_PROTOCOL_VERSION, diff --git a/packages/coding-agent/test/daemon-protocol.test.ts b/packages/coding-agent/test/daemon-protocol.test.ts index fc99b528916..56a43d85f91 100644 --- a/packages/coding-agent/test/daemon-protocol.test.ts +++ b/packages/coding-agent/test/daemon-protocol.test.ts @@ -45,6 +45,10 @@ describe("daemon protocol helpers", () => { minProtocol: 3, capability: "heartbeat_management", }); + expect(DAEMON_COMMAND_COMPATIBILITY.complete_owned_session).toEqual({ + minProtocol: DAEMON_PROTOCOL_VERSION, + capability: "client_owned_sessions", + }); expect(DAEMON_OUTBOUND_COMPATIBILITY.heartbeats_changed).toEqual({ minProtocol: 3, capability: "heartbeat_catalog", @@ -95,6 +99,7 @@ describe("daemon protocol helpers", () => { it("keeps attachment routing out of the durable mutation journal", () => { expect(isDaemonMutatingCommand({ type: "attach" })).toBe(false); expect(isDaemonMutatingCommand({ type: "reattach" })).toBe(false); + expect(isDaemonMutatingCommand({ type: "wait_for_headless_completion" })).toBe(true); expect(isDaemonMutatingCommand({ type: "switch_session" })).toBe(true); }); diff --git a/packages/coding-agent/test/daemon-stop-confirm.test.ts b/packages/coding-agent/test/daemon-stop-confirm.test.ts index 37ca05bc16a..1ad47c7a418 100644 --- a/packages/coding-agent/test/daemon-stop-confirm.test.ts +++ b/packages/coding-agent/test/daemon-stop-confirm.test.ts @@ -68,6 +68,17 @@ describe("confirmDaemonSessionLoss", () => { expect(console.error).toHaveBeenCalledWith(expect.stringContaining("busy:1")); }); + it("counts busy client-owned sessions without exposing their identities", async () => { + setTTY(false); + const probe: RunningDaemonProbe = { + reachable: true, + activeSessions: [], + busyClientOwnedSessionCount: 2, + }; + expect(await confirmDaemonSessionLoss(probe, { force: false, copy: COPY })).toBe(false); + expect(console.error).toHaveBeenCalledWith(expect.stringContaining("busy:2")); + }); + it("aborts without prompting when not at a TTY and sessions cannot be listed", async () => { setTTY(false); const probe: RunningDaemonProbe = { reachable: true }; diff --git a/packages/coding-agent/test/daemon-supervisor-process.test.ts b/packages/coding-agent/test/daemon-supervisor-process.test.ts index fab69f118c7..c10ec52e196 100644 --- a/packages/coding-agent/test/daemon-supervisor-process.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-process.test.ts @@ -286,6 +286,153 @@ async function startBlockingBash(client: DaemonClient, activeSessionId: string, } describe("daemon supervisor resident workers", () => { + it("keeps client-owned workers hidden and removes them without archiving", async () => { + const root = tempDir(); + const agentDir = join(root, "agent"); + const projectDir = join(root, "project"); + const sessionDir = join(agentDir, "sessions"); + const socketPath = join(tmpdir(), `prime-supervisor-owned-${process.pid}-${randomUUID().slice(0, 8)}.sock`); + mkdirSync(projectDir, { recursive: true }); + const sessionManager = SessionManager.create(projectDir, sessionDir); + sessionManager.appendMessage({ role: "user", content: "owned worker fixture", timestamp: 1 }); + const sessionFile = sessionManager.getSessionFile(); + if (!sessionFile) { + throw new Error("Fixture session did not persist"); + } + + const supervisor = spawnSupervisor(agentDir, socketPath, projectDir); + const client = await connectEventually(socketPath, supervisor); + const launchEnvSentinel = `owned-env-${randomUUID()}`; + const created = await client.request({ + type: "create", + sessionPath: sessionFile, + lifecycle: "client_owned", + launchEnv: { PRIME_AGENT_OWNED_TEST: launchEnvSentinel }, + config: { cwd: projectDir, agentDir, sessionDir, noTools: true, noExtensions: true }, + }); + expect(created.success).toBe(true); + const summary = requireSummary(created.success ? created.data : undefined); + if (!summary.workerPid || !summary.activeSessionId) { + throw new Error("Client-owned worker did not expose its process identity"); + } + workerPids.add(summary.workerPid); + + const publicList = await client.request({ type: "list" }); + expect(publicList.success).toBe(true); + expect(requireSessionList(publicList.success ? publicList.data : undefined)).toEqual([]); + const internalList = await client.request({ type: "list", includeClientOwned: true }); + expect(internalList.success).toBe(true); + expect(requireSessionList(internalList.success ? internalList.data : undefined)).toHaveLength(1); + const otherClient = await connectEventually(socketPath); + const deniedList = await otherClient.request({ type: "list", includeClientOwned: true }); + expect(deniedList).toMatchObject({ + success: true, + data: { sessions: [], busyClientOwnedSessionCount: 0 }, + }); + const privateSelector = summary.sessionId.slice(-8); + const deniedAttach = await otherClient.request({ type: "attach", activeSessionId: privateSelector }); + expect(deniedAttach).toMatchObject({ + success: false, + error: `Unknown active session: ${privateSelector}`, + }); + const deniedCron = await otherClient.request({ + type: "cron_add", + activeSessionId: summary.activeSessionId, + schedule: "every 1h", + prompt: "unauthorized", + }); + expect(deniedCron).toMatchObject({ + success: false, + error: `Unknown active session: ${summary.activeSessionId}`, + }); + otherClient.close(); + const connection = await DaemonAgentConnection.attach(client, summary.activeSessionId, { + ownedSession: true, + supportsExtensionUi: false, + }); + await expect(connection.listHeartbeats()).resolves.toEqual([]); + await expect(connection.listCronJobs()).resolves.toEqual([]); + const cronResponse = await client.request({ + type: "cron_add", + activeSessionId: summary.activeSessionId, + schedule: "every 1h", + prompt: "check status", + }); + if (!cronResponse.success || !cronResponse.data || typeof cronResponse.data !== "object") { + throw new Error(cronResponse.success ? "Cron response was missing its job" : cronResponse.error); + } + const cronJob = (cronResponse.data as { job: { id: string } }).job; + await expect(connection.listCronJobs()).resolves.toEqual([expect.objectContaining({ id: cronJob.id })]); + await expect(connection.cancelCronJob(cronJob.id)).resolves.toMatchObject({ id: cronJob.id }); + + const descriptor = readWorkerDescriptor(agentDir); + expect(descriptor.ownerClientId).toEqual(expect.any(String)); + expect(JSON.stringify(descriptor)).not.toContain(launchEnvSentinel); + expect(descriptor.createCommand).not.toHaveProperty("launchEnv"); + expect(descriptor.createCommand).not.toHaveProperty("lifecycle"); + + await connection.dispose(); + await waitForProcessGone(summary.workerPid); + workerPids.delete(summary.workerPid); + await waitForCondition( + () => countWorkerDescriptors(agentDir) === 0, + "Client-owned worker descriptor was not removed", + ); + expect((await readSessionInfo(sessionFile))?.state?.status).not.toBe("archived"); + + await client.request({ type: "shutdown" }); + client.close(); + await waitForSocketGone(socketPath); + }, 60_000); + + it("releases an adopted client-owned worker when disposal races supervisor replacement", async () => { + const root = tempDir(); + const agentDir = join(root, "agent"); + const projectDir = join(root, "project"); + const socketPath = join(tmpdir(), `prime-supervisor-owned-adopt-${process.pid}-${randomUUID().slice(0, 8)}.sock`); + mkdirSync(projectDir, { recursive: true }); + + const supervisor = spawnSupervisor(agentDir, socketPath, projectDir); + const client = await connectEventually(socketPath, supervisor); + const created = await client.request({ + type: "create", + lifecycle: "client_owned", + noSession: true, + launchEnv: { TSX_TSCONFIG_PATH: resolve(__dirname, "../../../tsconfig.json") }, + config: { cwd: projectDir, agentDir, noTools: true, noExtensions: true }, + }); + if (!created.success) { + throw new Error(created.error); + } + const summary = requireSummary(created.data); + if (!summary.workerPid || !summary.activeSessionId) { + throw new Error("Client-owned worker did not expose its process identity"); + } + workerPids.add(summary.workerPid); + const connection = await DaemonAgentConnection.attach(client, summary.activeSessionId, { + closeClientOnDispose: true, + ownedSession: true, + recoverDaemon: async () => {}, + supportsExtensionUi: false, + }); + + supervisor.kill("SIGTERM"); + await waitForExit(supervisor); + children.delete(supervisor); + await connection.dispose(); + + await waitForProcessGone(summary.workerPid); + workerPids.delete(summary.workerPid); + await waitForCondition( + () => countWorkerDescriptors(agentDir) === 0, + "Adopted client-owned worker descriptor was not removed", + ); + const replacementClient = await connectEventually(socketPath); + await replacementClient.request({ type: "shutdown" }); + replacementClient.close(); + await waitForSocketGone(socketPath); + }, 60_000); + it("cancels an archived session heartbeat without spawning a worker", { tags: ["process-stress"] }, async () => { const root = tempDir(); const agentDir = join(root, "agent"); diff --git a/packages/coding-agent/test/daemon-supervisor-side-question.test.ts b/packages/coding-agent/test/daemon-supervisor-side-question.test.ts index 7c68f0e1b3d..5005183dfca 100644 --- a/packages/coding-agent/test/daemon-supervisor-side-question.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-side-question.test.ts @@ -21,6 +21,9 @@ describe("daemon supervisor side-question routing", () => { const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { ready: Promise.resolve(), ownership: { assertCurrent: vi.fn(async () => undefined) }, + workers: new Map(), + clients: new Set(), + protocolClientIds: new WeakMap(), handleCommand, write, log: vi.fn(), diff --git a/packages/coding-agent/test/fixtures/rpc-connection-mode-fixture.ts b/packages/coding-agent/test/fixtures/rpc-connection-mode-fixture.ts new file mode 100644 index 00000000000..7e76e901b58 --- /dev/null +++ b/packages/coding-agent/test/fixtures/rpc-connection-mode-fixture.ts @@ -0,0 +1,123 @@ +import type { AgentConnection, AgentConnectionEventListener } from "../../src/modes/agent-connection/types.js"; +import { runRpcModeWithConnection } from "../../src/modes/rpc/rpc-mode.js"; + +let listener: AgentConnectionEventListener = () => {}; +let resolveExtensionUi: (() => void) | undefined; +let slowWatcherCount = 0; +let activePrompt: Promise | undefined; + +const heartbeat = { + id: "heartbeat-1", + status: "active" as const, + source: "heartbeat" as const, + activeSessionId: "root-session", + sessionId: "root", + sessionFile: "/tmp/root.jsonl", + cwd: "/tmp", + prompt: "check status", + schedule: { kind: "interval" as const, expression: "every 1h", intervalMs: 3_600_000 }, + createdAt: "2026-01-01T00:00:00.000Z", + updatedAt: "2026-01-01T00:00:00.000Z", + runCount: 0, +}; + +const connection = { + subscribe(next: AgentConnectionEventListener) { + listener = next; + return () => {}; + }, + async prompt(message: string) { + if (message === "async-eof") { + activePrompt = (async () => { + await listener({ type: "session_event", event: { type: "agent_start" } }); + await new Promise((resolve) => setTimeout(resolve, 50)); + await listener({ type: "session_event", event: { type: "agent_end", messages: [] } }); + })(); + return; + } + if (message === "extension-ui") { + await new Promise((resolve) => { + resolveExtensionUi = resolve; + void listener({ + type: "extension_ui_request", + request: { + id: "extension-ui-1", + method: "confirm", + payload: { title: "Confirm", message: "Continue?" }, + }, + }); + }); + } + await listener({ type: "session_event", event: { type: "agent_start" } }); + await new Promise((resolve) => setTimeout(resolve, 10)); + }, + async respondToExtensionUiRequest(id: string) { + if (id === "extension-ui-1") { + resolveExtensionUi?.(); + resolveExtensionUi = undefined; + } + }, + async getAvailableModels() { + await new Promise((resolve) => setTimeout(resolve, 25)); + return []; + }, + async waitForIdle() { + await activePrompt; + }, + async getLastAssistantText() { + return undefined; + }, + async setSessionName(name: string) { + if (name !== name.trim()) { + throw new Error("Session name was not trimmed"); + } + }, + async listCronJobs() { + return [heartbeat]; + }, + async listHeartbeats() { + return [{ job: heartbeat }]; + }, + async getAgentMessageStatus() { + return { + paused: false, + maxMessageChars: 16384, + maxPendingPerSession: 20, + rateLimitCapacity: 3, + rateLimitRefillMs: 1000, + }; + }, + async sendAgentMessage(targetActiveSessionId: string, message: string) { + return { + id: "message-1", + source: "agent_message" as const, + target: { activeSessionId: targetActiveSessionId, sessionId: "target" }, + message, + deliveryStatus: "delivered" as const, + deliveredAt: "2026-01-01T00:00:00.000Z", + deliveryMode: "auto" as const, + }; + }, + async watchSession(activeSessionId: string) { + const watcherIndex = activeSessionId === "slow-child" ? slowWatcherCount++ : -1; + return { + async getMessages() { + if (watcherIndex === 0) { + await new Promise((resolve) => setTimeout(resolve, 30)); + } + return [{ role: "user" as const, content: "child message", timestamp: 1 }]; + }, + subscribe(next: AgentConnectionEventListener) { + void next({ type: "session_event", event: { type: "agent_start" } }); + return () => {}; + }, + async getToolDefinition() { + return undefined; + }, + async close() {}, + }; + }, + async dispose() {}, +} as unknown as AgentConnection; + +await runRpcModeWithConnection(connection); diff --git a/packages/coding-agent/test/fixtures/rpc-eof-faux-extension.ts b/packages/coding-agent/test/fixtures/rpc-eof-faux-extension.ts new file mode 100644 index 00000000000..65d153d368f --- /dev/null +++ b/packages/coding-agent/test/fixtures/rpc-eof-faux-extension.ts @@ -0,0 +1,26 @@ +import { fauxAssistantMessage, getApiProvider, registerFauxProvider } from "../../../ai/src/index.js"; +import type { ExtensionAPI } from "../../src/index.js"; + +export default function registerRpcEofFauxProvider(pi: ExtensionAPI): void { + const faux = registerFauxProvider({ + provider: "faux", + models: [{ id: "faux", reasoning: false }], + }); + faux.setResponses([ + async () => { + await new Promise((resolve) => setTimeout(resolve, 250)); + return fauxAssistantMessage("rpc eof response"); + }, + ]); + const apiProvider = getApiProvider(faux.api); + if (!apiProvider) { + throw new Error("Faux API provider was not registered"); + } + pi.registerProvider(faux.getModel().provider, { + api: faux.api, + apiKey: "faux-key", + baseUrl: faux.getModel().baseUrl, + streamSimple: apiProvider.streamSimple, + models: faux.models, + }); +} diff --git a/packages/coding-agent/test/main-interactive-routing.test.ts b/packages/coding-agent/test/main-interactive-routing.test.ts index c5db711992a..577c750ba8c 100644 --- a/packages/coding-agent/test/main-interactive-routing.test.ts +++ b/packages/coding-agent/test/main-interactive-routing.test.ts @@ -16,12 +16,65 @@ import { shouldEnsureInteractiveDaemonForStartup, shouldOpenAgentsViewForDaemonInteractive, shouldRejectNonInteractiveAttach, + shouldUseDaemonClient, + shouldUseDaemonClientRuntime, shouldUseDaemonInteractive, shouldUseEphemeralSessionManagerForDaemonInteractive, } from "../src/main.js"; import type { SessionSummary } from "../src/modes/index.js"; describe("interactive startup routing", () => { + test.each(["interactive", "print", "json", "rpc"] as const)( + "uses the daemon runtime for the %s client", + (appMode) => { + expect( + shouldUseDaemonClient({ + appMode, + startupBenchmark: false, + }), + ).toBe(true); + }, + ); + + test("uses a client-owned daemon session for --no-session", () => { + expect( + shouldUseDaemonClient({ + appMode: "interactive", + startupBenchmark: false, + noSession: true, + }), + ).toBe(true); + }); + + test("keeps process-local extension factories and rollback workers in process", () => { + expect( + shouldUseDaemonClientRuntime({ + appMode: "print", + startupBenchmark: false, + hasProcessLocalExtensionFactories: true, + }), + ).toBe(false); + expect( + shouldUseDaemonClientRuntime({ + appMode: "rpc", + startupBenchmark: false, + ownedSessionWorker: true, + }), + ).toBe(false); + }); + + test.each([ + ["daemon process", { appMode: "daemon", startupBenchmark: false }], + ["startup benchmark", { appMode: "interactive", startupBenchmark: true }], + ["help", { appMode: "interactive", startupBenchmark: false, help: true }], + ["model listing", { appMode: "interactive", startupBenchmark: false, listModels: true }], + ] satisfies Array<[string, InteractiveDaemonStartupDecision]>)( + "keeps %s out of daemon client routing", + (_label, decision) => { + expect(shouldUseDaemonClient(decision)).toBe(false); + }, + ); + test("uses daemon-backed interactive mode for normal interactive startup", () => { expect( shouldUseDaemonInteractive({ diff --git a/packages/coding-agent/test/owned-session-worker-process.test.ts b/packages/coding-agent/test/owned-session-worker-process.test.ts index caa81149fbc..ab73026ab27 100644 --- a/packages/coding-agent/test/owned-session-worker-process.test.ts +++ b/packages/coding-agent/test/owned-session-worker-process.test.ts @@ -49,6 +49,7 @@ function spawnFrontend( env: { ...process.env, ...environment, + PRIME_AGENT_INTERNAL_LEGACY_OWNED_WORKER_FRONTEND: "1", PRIME_AGENT_TEST_OWNED_PID_PATH: pidPath, ...(keepAlive ? { PRIME_AGENT_TEST_KEEP_ALIVE: "1" } : {}), TSX_TSCONFIG_PATH: resolve(__dirname, "../../../tsconfig.json"), diff --git a/packages/coding-agent/test/owned-session-worker.test.ts b/packages/coding-agent/test/owned-session-worker.test.ts index 11d8777c8e9..6b0eeb756f4 100644 --- a/packages/coding-agent/test/owned-session-worker.test.ts +++ b/packages/coding-agent/test/owned-session-worker.test.ts @@ -4,6 +4,7 @@ import { classifyOwnedSessionWorkerInvocation, createOwnedWorkerLaunchSpec, createRpcRecoveryArgs, + isOwnedSessionWorkerProcess, } from "../src/cli/owned-session-worker.js"; describe("owned session worker CLI routing", () => { @@ -31,6 +32,8 @@ describe("owned session worker CLI routing", () => { }); it("does not recursively route an owned worker", () => { + expect(isOwnedSessionWorkerProcess({ PRIME_AGENT_INTERNAL_OWNED_WORKER: "1" })).toBe(true); + expect(isOwnedSessionWorkerProcess({})).toBe(false); expect( classifyOwnedSessionWorkerInvocation(["--mode", "rpc"], true, { PRIME_AGENT_INTERNAL_OWNED_WORKER: "1", diff --git a/packages/coding-agent/test/rpc-jsonl.test.ts b/packages/coding-agent/test/rpc-jsonl.test.ts index 2b902d21bf1..a16d392eb2f 100644 --- a/packages/coding-agent/test/rpc-jsonl.test.ts +++ b/packages/coding-agent/test/rpc-jsonl.test.ts @@ -134,4 +134,22 @@ describe("RPC JSONL framing", () => { expect(JSON.parse(lines[0]).blob.length).toBe(8 * 1024 * 1024); expect(elapsed).toBeLessThan(1000); }); + + test("bounds oversized lines and resumes at the next record", () => { + const lines: string[] = []; + const overflows: string[] = []; + const emitter = new EventEmitter(); + attachJsonlLineReader(emitter as unknown as Readable, (line) => lines.push(line), { + maxLineLength: 5, + onLineOverflow: (prefix) => overflows.push(prefix), + }); + + emitter.emit("data", "abc"); + emitter.emit("data", "defghi"); + emitter.emit("data", "jkl\nok\n"); + emitter.emit("end"); + + expect(overflows).toEqual(["abcde"]); + expect(lines).toEqual(["ok"]); + }); }); diff --git a/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts b/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts index ed129a7e5d5..4fa4364425a 100644 --- a/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts +++ b/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts @@ -231,7 +231,7 @@ describe("RPC prompt response semantics", () => { }); it("emits one success response when prompt preflight succeeds", async () => { - const { lineHandler, cleanup } = await startRpcMode({ withAuth: true, responseDelayMs: 0 }); + const { lineHandler, session, cleanup } = await startRpcMode({ withAuth: true, responseDelayMs: 250 }); try { lineHandler(JSON.stringify({ id: "b2", type: "prompt", message: "Hello" })); @@ -246,6 +246,10 @@ describe("RPC prompt response semantics", () => { success: true, }); }); + await vi.waitFor(() => { + expect(parseOutputLines(rpcIo.outputLines).some((record) => record.type === "agent_start")).toBe(true); + }); + expect(session.isStreaming).toBe(true); } finally { await cleanup(); } diff --git a/packages/coding-agent/test/suite/regressions/4620-fast-mode-child-agents.test.ts b/packages/coding-agent/test/suite/regressions/4620-fast-mode-child-agents.test.ts index 7739805c050..8f28fc748cb 100644 --- a/packages/coding-agent/test/suite/regressions/4620-fast-mode-child-agents.test.ts +++ b/packages/coding-agent/test/suite/regressions/4620-fast-mode-child-agents.test.ts @@ -31,6 +31,9 @@ describe("ENG-4620 fast mode child agents", () => { const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { ready: Promise.resolve(), ownership: { assertCurrent: vi.fn(async () => undefined) }, + workers: new Map(), + clients: new Set(), + protocolClientIds: new WeakMap(), handleCommand, write, log: vi.fn(), diff --git a/packages/coding-agent/test/suite/regressions/4685-daemon-client-modes.test.ts b/packages/coding-agent/test/suite/regressions/4685-daemon-client-modes.test.ts new file mode 100644 index 00000000000..65f6c1361f7 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/4685-daemon-client-modes.test.ts @@ -0,0 +1,566 @@ +import { type ChildProcess, spawn } from "node:child_process"; +import { existsSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join, resolve } from "node:path"; +import { fauxAssistantMessage } from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { createCliSubprocessEnv, createCliSubprocessLaunchSpec } from "../../../src/cli/subprocess-launch.js"; +import { ENV_AGENT_DIR } from "../../../src/config.js"; +import type { AutonomousRuntimeState } from "../../../src/core/autonomous.js"; +import type { DaemonSocketClient } from "../../../src/modes/daemon/active-session-state.js"; +import { DaemonClient } from "../../../src/modes/daemon/daemon-client.js"; +import { DaemonSupervisor } from "../../../src/modes/daemon/daemon-supervisor.js"; +import { waitForHeadlessCompletion } from "../../../src/modes/headless-completion.js"; +import { RpcClient } from "../../../src/modes/rpc/rpc-client.js"; +import { createRpcExtensionUiBridge } from "../../../src/modes/rpc/rpc-extension-ui-context.js"; +import { createHarness, getAssistantTexts, getUserTexts, type Harness } from "../harness.js"; + +const fixturePath = resolve(__dirname, "../../fixtures/rpc-connection-mode-fixture.ts"); +const fauxExtensionPath = resolve(__dirname, "../../fixtures/eng-4600-faux-extension.ts"); +const rpcEofFauxExtensionPath = resolve(__dirname, "../../fixtures/rpc-eof-faux-extension.ts"); +const cliPath = resolve(__dirname, "../../../src/cli.ts"); +const tsxPath = resolve(__dirname, "../../../../../node_modules/tsx/dist/cli.mjs"); +const repoTsconfigPath = resolve(__dirname, "../../../../../tsconfig.json"); +const children = new Set(); +const harnesses: Harness[] = []; +const daemonSockets = new Set(); +const tempRoots = new Set(); + +afterEach(async () => { + for (const child of children) { + if (child.exitCode === null && child.signalCode === null) { + child.kill("SIGKILL"); + } + } + children.clear(); + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + for (const socketPath of daemonSockets) { + const client = new DaemonClient(socketPath); + try { + await client.connect(500); + await client.request({ type: "shutdown" }, 5000); + } catch { + // The process may have exited before publishing its socket. + } finally { + client.close(); + } + for (let attempt = 0; attempt < 50 && existsSync(socketPath); attempt++) { + await new Promise((resolveDelay) => setTimeout(resolveDelay, 20)); + } + } + daemonSockets.clear(); + for (const root of tempRoots) { + rmSync(root, { recursive: true, force: true, maxRetries: 10, retryDelay: 50 }); + } + tempRoots.clear(); +}); + +interface CliResult { + code: number | null; + signal: NodeJS.Signals | null; + stdout: string; + stderr: string; +} + +async function runCli( + args: string[], + options: { agentDir: string; stdin?: string; environment?: NodeJS.ProcessEnv }, +): Promise { + const child = spawn(process.execPath, [tsxPath, cliPath, ...args], { + env: { + ...process.env, + TSX_TSCONFIG_PATH: repoTsconfigPath, + [ENV_AGENT_DIR]: options.agentDir, + PI_SKIP_VERSION_CHECK: "1", + PRIME_AGENT_INTERNAL_LEGACY_OWNED_WORKER_FRONTEND: "0", + ...options.environment, + }, + stdio: ["pipe", "pipe", "pipe"], + }); + children.add(child); + let stdout = ""; + let stderr = ""; + child.stdout?.on("data", (chunk: Buffer) => { + stdout += chunk.toString("utf8"); + }); + child.stderr?.on("data", (chunk: Buffer) => { + stderr += chunk.toString("utf8"); + }); + child.stdin?.end(options.stdin ?? ""); + const exit = await new Promise<{ code: number | null; signal: NodeJS.Signals | null }>((resolveExit, reject) => { + const timeout = setTimeout(() => { + child.kill("SIGKILL"); + reject(new Error(`CLI timed out\n${stderr}`)); + }, 20_000); + child.once("exit", (code, signal) => { + clearTimeout(timeout); + resolveExit({ code, signal: signal as NodeJS.Signals | null }); + }); + }); + children.delete(child); + return { ...exit, stdout, stderr }; +} + +async function runRpc( + commands: unknown[], + options: { trailingNewline?: boolean } = {}, +): Promise<{ stdout: object[]; stderr: string }> { + const child = spawn(process.execPath, [tsxPath, fixturePath], { + env: { ...process.env, TSX_TSCONFIG_PATH: repoTsconfigPath }, + stdio: ["pipe", "pipe", "pipe"], + }); + children.add(child); + let stdout = ""; + let stderr = ""; + child.stdout?.on("data", (chunk: Buffer) => { + stdout += chunk.toString("utf8"); + }); + child.stderr?.on("data", (chunk: Buffer) => { + stderr += chunk.toString("utf8"); + }); + const input = commands.map((command) => JSON.stringify(command)).join("\n"); + child.stdin?.end(options.trailingNewline === false ? input : `${input}\n`); + const exit = await new Promise<{ code: number | null; signal: NodeJS.Signals | null }>((resolveExit, reject) => { + const timeout = setTimeout(() => reject(new Error(`RPC fixture timed out\n${stderr}`)), 10_000); + child.once("exit", (code, signal) => { + clearTimeout(timeout); + resolveExit({ code, signal: signal as NodeJS.Signals | null }); + }); + }); + children.delete(child); + expect(exit).toEqual({ code: 0, signal: null }); + return { + stdout: stdout + .trim() + .split("\n") + .filter(Boolean) + .map((line) => JSON.parse(line) as object), + stderr, + }; +} + +describe("ENG-4685 daemon-backed client modes", () => { + it("commits owned-worker promotion before best-effort peer synchronization", async () => { + const client = { id: "client-1" } as DaemonSocketClient; + const worker = { + descriptor: { ownerClientId: "protocol-client" }, + launchEnv: { TEST: "value" }, + }; + const persistWorker = vi.fn(); + const syncAgentPeers = vi.fn(async () => { + throw new Error("peer unavailable"); + }); + const log = vi.fn(); + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + protocolClientId: () => "protocol-client", + persistWorker, + syncAgentPeers, + log, + }) as { + promoteOwnedWorker(client: DaemonSocketClient, resident: typeof worker): Promise; + }; + + await supervisor.promoteOwnedWorker(client, worker); + await supervisor.promoteOwnedWorker(client, worker); + + expect(worker.descriptor.ownerClientId).toBeUndefined(); + expect(worker.launchEnv).toBeUndefined(); + expect(persistWorker).toHaveBeenCalledOnce(); + expect(syncAgentPeers).toHaveBeenCalledOnce(); + expect(log).toHaveBeenCalledWith(expect.stringContaining("peer unavailable")); + }); + + it("rolls back owned-worker promotion when persistence fails", async () => { + const client = { id: "client-1" } as DaemonSocketClient; + const descriptor = { ownerClientId: "protocol-client" }; + const timer = setTimeout(() => {}, 60_000); + const worker = { + descriptor, + launchEnv: { TEST: "value" }, + ownerCleanupTimer: timer, + }; + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + protocolClientId: () => "protocol-client", + persistWorker: () => { + throw new Error("disk full"); + }, + }) as { + promoteOwnedWorker(client: DaemonSocketClient, resident: typeof worker): Promise; + }; + + await expect(supervisor.promoteOwnedWorker(client, worker)).rejects.toThrow("disk full"); + + expect(worker.descriptor).toBe(descriptor); + expect(worker.launchEnv).toEqual({ TEST: "value" }); + expect(worker.ownerCleanupTimer).toBe(timer); + clearTimeout(timer); + }); + + it("runs host-owned autonomous gate retries through the shared completion loop", async () => { + const gate = `${process.execPath} -e "process.exit(0)"`; + const harness = await createHarness({ + autonomous: { + enabled: true, + maxContinuations: 2, + gates: { commands: [gate], maxRetries: 2 }, + }, + }); + harnesses.push(harness); + const state = ( + harness.session as unknown as { + _autonomousState: AutonomousRuntimeState; + } + )._autonomousState; + state.gateAttempts[gate] = 1; + state.lastGateFailure = { + command: gate, + attempt: 1, + exitText: "exited 1", + output: "gate failed", + }; + harness.setResponses([fauxAssistantMessage("I fixed the gate failure.")]); + + const status = await waitForHeadlessCompletion(harness.session); + + expect(getUserTexts(harness)).toHaveLength(1); + expect(getUserTexts(harness)[0]).toContain("Autonomous quality gate failed"); + expect(getAssistantTexts(harness)).toEqual(["I fixed the gate failure."]); + expect(status).toMatchObject({ + continuationsUsed: 1, + lastGateFailure: undefined, + }); + }); + + it("resolves source subprocesses independently of a spaced runtime cwd", () => { + const entrypoint = resolve(__dirname, "../../../src/cli.ts"); + const launch = createCliSubprocessLaunchSpec([], process.execPath, [], "packages/coding-agent/src/cli.ts"); + const environment = createCliSubprocessEnv({}, entrypoint, ["--import", "tsx"]); + + expect(launch.args[0]).toBe(resolve("packages/coding-agent/src/cli.ts")); + expect(environment.TSX_TSCONFIG_PATH).toBe(repoTsconfigPath); + }); + + it("launches real daemon workers for every migrated client surface", async () => { + const root = mkdtempSync(join(tmpdir(), "prime-agent-4685-clients-")); + tempRoots.add(root); + const agentDir = join(root, "agent dir"); + const socketPath = join(root, "daemon.sock"); + daemonSockets.add(socketPath); + const baseArgs = [ + "--daemon-socket", + socketPath, + "--model", + "faux/faux", + "--extension", + fauxExtensionPath, + "--no-tools", + "--no-skills", + "--no-prompt-templates", + "--no-themes", + "--no-context-files", + ]; + const cases = [ + { name: "print", args: ["--print"], stdin: "" }, + { name: "json", args: ["--mode", "json"], stdin: "" }, + { name: "rpc", args: ["--mode", "rpc"], stdin: '{"id":"state","type":"get_state"}\n' }, + { name: "piped stdin", args: [], stdin: " \n" }, + { name: "no-session", args: ["--print", "--no-session"], stdin: "" }, + ]; + + for (const testCase of cases) { + const result = await runCli([...baseArgs, ...testCase.args], { agentDir, stdin: testCase.stdin }); + expect(result, testCase.name).toMatchObject({ code: 0, signal: null }); + expect(result.stderr, testCase.name).not.toContain("Timed out waiting for daemon worker"); + if (testCase.name === "rpc") { + expect(result.stdout).toContain('"command":"get_state","success":true'); + } + } + expect(existsSync(socketPath)).toBe(true); + }, 90_000); + + it("keeps the rollback frontend fully off the daemon path", async () => { + const root = mkdtempSync(join(tmpdir(), "prime-agent-4685-rollback-")); + tempRoots.add(root); + const agentDir = join(root, "agent"); + const socketPath = join(root, "must-not-exist.sock"); + const result = await runCli( + [ + "--print", + "--daemon-socket", + socketPath, + "--model", + "faux/faux", + "--extension", + fauxExtensionPath, + "--no-tools", + "--no-skills", + "--no-prompt-templates", + "--no-themes", + "--no-context-files", + ], + { + agentDir, + environment: { PRIME_AGENT_INTERNAL_LEGACY_OWNED_WORKER_FRONTEND: "1" }, + }, + ); + + expect(result).toMatchObject({ code: 0, signal: null }); + expect(existsSync(socketPath)).toBe(false); + }, 30_000); + + it("loads headless runtime services only in the worker", async () => { + const root = mkdtempSync(join(tmpdir(), "prime-agent-4685-services-")); + tempRoots.add(root); + const agentDir = join(root, "agent"); + const socketPath = join(root, "daemon.sock"); + const markerPath = join(root, "extension-loads.txt"); + const extensionPath = join(root, "load-marker.ts"); + daemonSockets.add(socketPath); + writeFileSync( + extensionPath, + 'import { appendFileSync } from "node:fs";\nexport default function() { appendFileSync(process.env.PRIME_AGENT_TEST_EXTENSION_LOAD_MARKER, "loaded\\n"); }\n', + ); + + const result = await runCli( + [ + "--print", + "--daemon-socket", + socketPath, + "--model", + "faux/faux", + "--extension", + fauxExtensionPath, + "--extension", + extensionPath, + "--no-tools", + "--no-skills", + "--no-prompt-templates", + "--no-themes", + "--no-context-files", + ], + { + agentDir, + environment: { PRIME_AGENT_TEST_EXTENSION_LOAD_MARKER: markerPath }, + }, + ); + + expect(result).toMatchObject({ code: 0, signal: null }); + expect(readFileSync(markerPath, "utf8").trim().split("\n")).toEqual(["loaded"]); + }, 30_000); + + it("drains accepted RPC commands before EOF releases the connection", async () => { + const result = await runRpc([{ id: "models", type: "get_available_models" }]); + expect(result.stderr).toBe(""); + expect(result.stdout).toEqual([ + { + id: "models", + type: "response", + command: "get_available_models", + success: true, + data: { models: [] }, + }, + ]); + }); + + it("drains accepted RPC prompt work before EOF releases the connection", async () => { + const result = await runRpc([{ id: "prompt", type: "prompt", message: "async-eof" }]); + expect(result.stderr).toBe(""); + expect(result.stdout).toEqual([ + { id: "prompt", type: "response", command: "prompt", success: true }, + { type: "agent_start" }, + { type: "agent_end", messages: [] }, + ]); + }); + + it("drains accepted daemon RPC prompt work before EOF", async () => { + const root = mkdtempSync(join(tmpdir(), "prime-agent-4685-rpc-eof-")); + tempRoots.add(root); + const socketPath = join(root, "daemon.sock"); + daemonSockets.add(socketPath); + const result = await runCli( + [ + "--mode", + "rpc", + "--no-session", + "--daemon-socket", + socketPath, + "--model", + "faux/faux", + "--extension", + rpcEofFauxExtensionPath, + "--no-tools", + "--no-skills", + "--no-prompt-templates", + "--no-themes", + "--no-context-files", + ], + { + agentDir: join(root, "agent"), + stdin: '{"id":"prompt","type":"prompt","message":"finish before EOF"}\n', + }, + ); + + expect(result).toMatchObject({ code: 0, signal: null, stderr: "" }); + expect(result.stdout).toContain('{"id":"prompt","type":"response","command":"prompt","success":true}'); + expect(result.stdout).toContain("rpc eof response"); + expect(result.stdout).toContain('"type":"agent_end"'); + }, 30_000); + + it("drains an unterminated final RPC command before EOF", async () => { + const result = await runRpc([{ id: "models", type: "get_available_models" }], { + trailingNewline: false, + }); + expect(result.stderr).toBe(""); + expect(result.stdout).toEqual([ + { + id: "models", + type: "response", + command: "get_available_models", + success: true, + data: { models: [] }, + }, + ]); + }); + + it("returns a structured error for scalar JSON input", async () => { + const result = await runRpc([null]); + expect(result.stderr).toBe(""); + expect(result.stdout).toEqual([ + { + type: "response", + command: "parse", + success: false, + error: "Invalid command: expected an object with a string type", + }, + ]); + }); + + it("cancels pending extension dialogs on RPC EOF", async () => { + const result = await runRpc([{ id: "prompt", type: "prompt", message: "extension-ui" }]); + expect(result.stderr).toBe(""); + expect(result.stdout).toEqual([ + { + type: "extension_ui_request", + id: "extension-ui-1", + method: "confirm", + title: "Confirm", + message: "Continue?", + }, + { id: "prompt", type: "response", command: "prompt", success: true }, + { type: "agent_start" }, + ]); + }); + + it("closes local RPC extension dialogs and rejects future waits", async () => { + const bridge = createRpcExtensionUiBridge(() => {}); + const pending = bridge.uiContext.editor("Edit", "draft"); + + bridge.close(); + + await expect(pending).resolves.toBeUndefined(); + await expect(bridge.uiContext.editor("Edit again", "draft")).resolves.toBeUndefined(); + }); + + it("isolates throwing RPC event listeners", () => { + const client = new RpcClient(); + const observed: string[] = []; + client.onObservedSessionEvent(() => { + throw new Error("listener failed"); + }); + client.onObservedSessionEvent((event) => observed.push(event.type)); + + (client as unknown as { handleLine(line: string): void }).handleLine( + JSON.stringify({ type: "observed_session_closed", activeSessionId: "child" }), + ); + + expect(observed).toEqual(["observed_session_closed"]); + }); + + it("orders concurrent observation baselines before their live events", async () => { + const result = await runRpc([ + { id: "observe-1", type: "observe", activeSessionId: "slow-child" }, + { id: "observe-2", type: "observe", activeSessionId: "slow-child" }, + ]); + const firstResponse = result.stdout.findIndex((record) => "id" in record && record.id === "observe-1"); + const firstEvent = result.stdout.findIndex( + (record) => "type" in record && record.type === "observed_session_event", + ); + + expect(firstResponse).toBeGreaterThanOrEqual(0); + expect(firstEvent).toBeGreaterThan(firstResponse); + }); + + it("preserves prompt acknowledgements before events and repeated prompts", async () => { + const result = await runRpc([ + { id: "prompt-1", type: "prompt", message: "one" }, + { id: "prompt-2", type: "prompt", message: "two" }, + ]); + expect(result.stderr).toBe(""); + expect(result.stdout).toEqual([ + { id: "prompt-1", type: "response", command: "prompt", success: true }, + { type: "agent_start" }, + { id: "prompt-2", type: "response", command: "prompt", success: true }, + { type: "agent_start" }, + ]); + }); + + it("preserves RPC trimming and omitted-value encoding", async () => { + const result = await runRpc([ + { id: "last", type: "get_last_assistant_text" }, + { id: "name", type: "set_session_name", name: " exact name " }, + ]); + expect(result.stderr).toBe(""); + expect(result.stdout).toEqual([ + { id: "last", type: "response", command: "get_last_assistant_text", success: true, data: {} }, + { id: "name", type: "response", command: "set_session_name", success: true }, + ]); + }); + + it("exposes daemon schedules, heartbeats, messaging, and observation", async () => { + const result = await runRpc([ + { id: "schedules", type: "list_schedules" }, + { id: "heartbeats", type: "list_heartbeats" }, + { id: "status", type: "agent_messages_status" }, + { id: "message", type: "send_message", targetActiveSessionId: "child", message: "check in" }, + { id: "observe", type: "observe", activeSessionId: "child" }, + ]); + + expect(result.stderr).toBe(""); + expect(result.stdout).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + id: "schedules", + success: true, + data: expect.objectContaining({ jobs: expect.arrayContaining([expect.any(Object)]) }), + }), + expect.objectContaining({ + id: "heartbeats", + success: true, + data: expect.objectContaining({ heartbeats: expect.arrayContaining([expect.any(Object)]) }), + }), + expect.objectContaining({ + id: "status", + success: true, + data: expect.objectContaining({ paused: false }), + }), + expect.objectContaining({ + id: "message", + success: true, + data: expect.objectContaining({ deliveryStatus: "delivered" }), + }), + expect.objectContaining({ + id: "observe", + success: true, + data: expect.objectContaining({ messages: expect.arrayContaining([expect.any(Object)]) }), + }), + { + type: "observed_session_event", + activeSessionId: "child", + event: { type: "agent_start" }, + }, + ]), + ); + }); +});