From 4dbbe925a9ebdedc33fa39caf71c6368e7577f00 Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Sat, 29 Aug 2026 14:17:56 -0600 Subject: [PATCH 01/12] fix: isolate snapshot catch-up failures Fixes #11 --- .pylon/features.yaml | 18 + .pylon/upstream-review.md | 11 + .../11-snapshot-recovery-integrity.md | 3 + .../coding-agent/src/core/agent-session.ts | 190 ++-- .../daemon-agent-connection.ts | 447 ++++++-- .../src/modes/daemon/daemon-mode.ts | 143 ++- .../src/modes/daemon/daemon-protocol.ts | 7 +- .../src/modes/daemon/daemon-session-list.ts | 91 +- .../src/modes/daemon/daemon-supervisor.ts | 1006 +++++++++++++---- .../src/modes/daemon/daemon-worker-client.ts | 6 +- .../modes/daemon/snapshot-transcript-cache.ts | 818 ++++++++++++-- .../modes/session-worker/private-framing.ts | 17 + .../test/agent-connection-daemon.test.ts | 294 ++++- .../test/agent-session-recursion.test.ts | 89 +- .../coding-agent/test/daemon-mode.test.ts | 69 +- .../coding-agent/test/daemon-protocol.test.ts | 6 + .../test/daemon-session-list.test.ts | 161 +-- .../daemon-supervisor-lazy-subagents.test.ts | 112 +- .../test/daemon-supervisor-monitor.test.ts | 214 +++- .../session-worker-private-framing.test.ts | 116 ++ .../test/snapshot-transcript-cache.test.ts | 255 ++++- .../4601-worker-snapshot-cache.test.ts | 58 +- ...4602-snapshot-transfer-idempotency.test.ts | 442 +++++++- .../4677-snapshot-catchup-replacement.test.ts | 56 +- 24 files changed, 3801 insertions(+), 828 deletions(-) create mode 100644 packages/coding-agent/.changes/11-snapshot-recovery-integrity.md diff --git a/.pylon/features.yaml b/.pylon/features.yaml index 5684e1935e..e72b2946e9 100644 --- a/.pylon/features.yaml +++ b/.pylon/features.yaml @@ -170,3 +170,21 @@ decisions: upstream_support: "Prime PR #1886 repairs protocol fixtures regressed by the #1839 merge commit, but is not yet on main." revisit_when: - "Prime merges PR #1886 and the exact fix reaches the mirrored main branch." + + snapshot-recovery-integrity: + area: runtime-reliability + state: candidate + owner: shared + decision: hybridize + pylon_refs: + - https://github.com/pylon-code/prime-agent/issues/11 + upstream_refs: + - https://github.com/PrimeIntellect-ai/prime-agent/issues/1229 + - https://github.com/PrimeIntellect-ai/prime-agent/pull/1845 + - https://github.com/PrimeIntellect-ai/prime-agent/pull/1756 + - https://github.com/PrimeIntellect-ai/prime-agent/pull/1864 + fork_change: immutable-snapshot-transfer-and-isolated-retry + upstream_support: Prime now provides one RLM child projection, recovery-aware session reuse, and ownership checks for in-flight opens, but still reuses the event cursor as snapshot transfer identity, serializes selected message objects lazily, and closes a healthy worker on an isolated duplicate-transfer mismatch. + revisit_when: + - Prime gives each transfer immutable bytes and a unique identity while preserving the event cursor as ordering metadata. + - Prime retries an isolated bad snapshot generation without closing the resident worker control channel. diff --git a/.pylon/upstream-review.md b/.pylon/upstream-review.md index 28e5e8c30a..e11872de3e 100644 --- a/.pylon/upstream-review.md +++ b/.pylon/upstream-review.md @@ -45,3 +45,14 @@ This ledger records Prime upstream evidence and the decision taken for each over - Fork candidate: [pylon-code/prime-agent#9](https://github.com/pylon-code/prime-agent/pull/9). - Validation: the post-rebase lifecycle, connection, supervisor, runtime, recovery, and ledger checkpoint passed 564 tests across 17 files. The final changed-boundary checkpoint passed 202 connection and supervisor-monitor tests; the clean supervisor process checkpoint passed 10 tests with 8 fixture-gated skips. `npm run check` passed Biome over 940 files, TypeScript, installer rendering, and browser smoke; `git diff --check` passed. Independent final and follow-up audits found no remaining P0 or P1 findings. - Revisit when Prime upstream ships an explicitly negotiated equivalent with the same delivery, provenance, scoped cancellation, recovery, generation, retry-integrity, privacy, and per-prompt usage guarantees. + +## 2026-08-29 — snapshot recovery integrity candidate + +- Audited baseline remains `PrimeIntellect-ai/prime-agent@d60fab8a76d9c169f945341f0ee3bde21903bb55`; latest audited release remains `v0.8.1`. The selected evidence commits are later than that baseline, but the intervening range was not fully audited and therefore does not advance `reviewed_upstream_commit`. +- Reviewed closed Prime issue #1229, merged Prime PRs #1845, #1756, and #1864, and their commits `c0334a176fb8f78b8c37327a6c76865792b03b52`, `ee8fd699617e8b66cb4bf36c97872c8fe52192e3`, and `5b6c0e94e11a97fcfdd7a9fc9dc4f7acbda9c853`. The pull requests solve child-snapshot projection fan-in, recovery-aware worker reuse, and in-flight open ownership. They do not fix the snapshot identity and failure-containment invariant recorded in #1229. +- `snapshot-recovery-integrity`: **hybridize**. Adopt and adapt those three upstream changes, while retaining Pylon's correlated prompt lifecycle and adding the missing transfer boundary: each attach, replacement, or catch-up gets an opaque unique snapshot ID; selected messages are encoded into immutable memory/file-backed chunks before deferred streaming; and a transfer-local mismatch retires only that generation and permits one bounded fresh-generation retry without closing the healthy resident worker. +- Protocol classification: **backward-compatible, capability-gated**. Schema revision 25 adds optional `immutable_snapshot_transfer_v1` negotiation without changing protocol version 7 or any existing command/event shape. A new supervisor requests worker chunk streams only when the worker advertises immutable identities. With an older worker it requests the full snapshot and re-encodes it under a supervisor-owned UUID, while old supervisors and stock clients retain their existing behavior. +- Resource boundary: transcript preparation checks cancellation before allocation, yields during large encodes, serializes per-session preparation, cleans partial spill files on every constructor/factory failure, uses private pid/start-identity/UUID cache roots, and sweeps abandoned roots only after ownership can be disproved. A single very large JSON message still incurs one synchronous `JSON.stringify` interval; realistic transcript size is not capped. +- The local `refs/review/current-candidate` at `a37efe92f88ae7234dcea1d065c0d3d76ac048cc` predates these corrections and includes the unaudited intervening range. Do not merge it. The upstream-review workflow must discard or regenerate it after this issue lands; this change does not mutate the protected review ref. +- Validation: the touched snapshot, protocol, connection, supervisor, recursion, daemon-mode, session-list, and lazy-subagent suites passed 616 tests across 11 files. The real-process supervisor control-channel regression for root isolation, chunked snapshot streaming, and worker adoption passed. `npm run check` passed Biome over 940 files, TypeScript, installer rendering, and browser smoke; `git diff --check` passed. +- Revisit when Prime gives every snapshot attempt immutable bytes and unique transfer identity, isolates bad generations from worker recovery, and preserves the same recovery and ownership guarantees without weakening Pylon's correlated lifecycle contract. diff --git a/packages/coding-agent/.changes/11-snapshot-recovery-integrity.md b/packages/coding-agent/.changes/11-snapshot-recovery-integrity.md new file mode 100644 index 0000000000..668a905f51 --- /dev/null +++ b/packages/coding-agent/.changes/11-snapshot-recovery-integrity.md @@ -0,0 +1,3 @@ +- Fixed daemon snapshot catch-up and recovery so transfer mismatches no longer interrupt healthy workers and recovered sessions preserve their state and owner ([#11](https://github.com/pylon-code/prime-agent/issues/11)). +- Fixed recovered child rosters and activity so parent sessions show current delegated work after restart ([#11](https://github.com/pylon-code/prime-agent/issues/11)). +- Fixed concurrent session opens so one owner and recovery attempt is reused instead of starting competing workers ([#11](https://github.com/pylon-code/prime-agent/issues/11)). diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index b26fdfbef5..2de6a12ab9 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -902,6 +902,10 @@ interface RlmChildRun { sessionDir: string; model: Model; status: RlmChildAgentStatus; + durationMs?: number; + answerPreview?: string; + toolUseCount: number; + activity?: RlmChildAgentActivity; error?: string; abort: () => void; publication: AgentMessageDeferred; @@ -931,6 +935,11 @@ interface RlmChildRun { unsubscribe?: () => void; } +interface RetainedRlmChild { + session: AgentSession; + run?: RlmChildRun; +} + interface RlmSubagentModelSelection { model: Model; } @@ -1195,7 +1204,7 @@ export class AgentSession { private _pendingRlmSubagentSessionNames = new Set(); // Inline mode keeps finished child sessions so the inspector can still read them; // the daemon does the same by leaving the child session resident in its registry. - private _rlmChildSessions = new Map(); + private _rlmChildSessions = new Map(); private _deletedRlmChildIds = new Set(); // Failed explicit deletes stay hidden from listings but retain their original // selector so a later delete can retry cleanup without orphaning the runtime. @@ -4333,7 +4342,7 @@ export class AgentSession { unsubscribe(); } this._rlmChildUnsubscribes.clear(); - for (const session of this._rlmChildSessions.values()) { + for (const { session } of this._rlmChildSessions.values()) { await session.disposeAsync().catch(() => undefined); } this._rlmChildSessions.clear(); @@ -4396,7 +4405,7 @@ export class AgentSession { unsubscribe(); } this._rlmChildUnsubscribes.clear(); - for (const session of this._rlmChildSessions.values()) { + for (const { session } of this._rlmChildSessions.values()) { session.dispose(); } this._rlmChildSessions.clear(); @@ -10043,7 +10052,7 @@ export class AgentSession { }); recorded.add(run.id); } - for (const [childId, childSession] of this._rlmChildSessions) { + for (const [childId, { session: childSession }] of this._rlmChildSessions) { if ( this._deletingRlmChildren.has(childId) || recorded.has(childId) || @@ -10131,7 +10140,7 @@ export class AgentSession { return result; } } - for (const retained of this._rlmChildSessions.values()) { + for (const { session: retained } of this._rlmChildSessions.values()) { const result = await retained.deleteInactiveRlmSubagent(childId, isExternallyRunning); if (result !== "not_found") { return result; @@ -10408,7 +10417,7 @@ export class AgentSession { } this._emitRlmSubagentRemoval(subagent); - const retained = this._rlmChildSessions.get(childId); + const retained = this._rlmChildSessions.get(childId)?.session; try { await this._deleteRlmSubagentSession(childId, retained); } catch (error) { @@ -10444,7 +10453,7 @@ export class AgentSession { void session.disposeAsync().catch(() => undefined); return false; } - this._rlmChildSessions.set(childId, session); + this._rlmChildSessions.set(childId, { session, run: this._activeRlmChildRuns.get(childId) }); if (unsubscribe) { this._rlmChildUnsubscribes.set(childId, unsubscribe); } @@ -10455,15 +10464,77 @@ export class AgentSession { const run = this._activeRlmChildRuns.get(childId); if (run?.session === session && run.status === "done") { const unsubscribe = run.unsubscribe ?? noopRlmChildEventUnsubscribe; - run.unsubscribe = undefined; - this._activeRlmChildRuns.delete(childId); - return unsubscribe; + return () => { + run.unsubscribe = undefined; + this._activeRlmChildRuns.delete(childId); + unsubscribe(); + }; } - if (this._rlmChildSessions.get(childId) !== session) return false; + if (this._rlmChildSessions.get(childId)?.session !== session) return false; const unsubscribe = this._rlmChildUnsubscribes.get(childId) ?? noopRlmChildEventUnsubscribe; - this._rlmChildUnsubscribes.delete(childId); - this._rlmChildSessions.delete(childId); - return unsubscribe; + return () => { + this._rlmChildUnsubscribes.delete(childId); + this._rlmChildSessions.delete(childId); + unsubscribe(); + }; + } + + private _rlmChildSnapshotForRun( + run: RlmChildRun, + child = run.session ?? this._rlmChildSessions.get(run.id)?.session, + ): RlmChildAgentSnapshot { + const model = child?.model ?? run.model; + return { + id: run.id, + parentId: this._rlmParentNodeId, + sessionName: child?.sessionName ?? run.sessionName, + model: `${model.provider}/${model.id}`, + label: rlmChildLabel(run.prompt), + status: run.status, + durationMs: run.durationMs, + answerPreview: run.answerPreview, + toolUseCount: run.toolUseCount > 0 ? run.toolUseCount : undefined, + tokenCount: child?._contextTokensForCurrentMessages(), + recap: child?.getCurrentRecap(), + sessionDir: run.sessionDir, + activity: run.activity, + repliedSinceTask: child?._repliedToParentSinceTask, + error: run.error, + }; + } + + private _rlmChildSnapshotForSession(childId: string, child: AgentSession): RlmChildAgentSnapshot { + let answerPreview: string | undefined; + let toolUseCount = 0; + const messages = + child.state.streamingMessage?.role === "assistant" + ? [...child.messages, child.state.streamingMessage] + : child.messages; + for (const message of messages) { + if (message.role !== "assistant") continue; + const text = compactRlmText(readAssistantText(message)); + if (text) answerPreview = text; + toolUseCount += message.content.filter((block) => block.type === "toolCall").length; + } + return { + id: childId, + parentId: this._rlmParentNodeId, + sessionName: child.sessionName, + model: child.model ? `${child.model.provider}/${child.model.id}` : undefined, + label: child.sessionName ?? "child agent", + status: "done", + answerPreview, + toolUseCount: toolUseCount > 0 ? toolUseCount : undefined, + tokenCount: child._contextTokensForCurrentMessages(), + recap: child.getCurrentRecap(), + sessionDir: child._rlmSessionDir ?? child.sessionManager.getSessionDir(), + // No run exists (e.g. a child rehydrated after daemon recovery), so live + // session state is the only source for in-flight follow-up work. Mirror + // the run projection's convention: status stays "done" (the recorded task + // finished) and current work surfaces through activity. + activity: child.isSessionActive ? { kind: child.isStreaming ? "writing" : "waiting" } : undefined, + repliedSinceTask: child._repliedToParentSinceTask, + }; } /** Live recursive child roster from lifecycle state, including nested work under retained parents. */ @@ -10476,16 +10547,7 @@ export class AgentSession { run.detachedDeletion || this._deletingRlmChildren.has(run.id) || this._deletedRlmChildIds.has(run.id); const child = run.session; if (!hidden) { - const model = child?.model ?? run.model; - snapshots.push({ - id: run.id, - parentId: this._rlmParentNodeId, - sessionName: child?.sessionName ?? run.sessionName, - model: `${model.provider}/${model.id}`, - label: rlmChildLabel(run.prompt), - status: run.status, - sessionDir: run.sessionDir, - }); + snapshots.push(this._rlmChildSnapshotForRun(run)); recorded.add(run.id); } if (child) { @@ -10493,20 +10555,16 @@ export class AgentSession { snapshots.push(...child.getRlmChildSnapshots()); } } - for (const [childId, child] of this._rlmChildSessions) { + for (const [childId, { session: child, run }] of this._rlmChildSessions) { if (recorded.has(childId) || traversed.has(childId)) continue; const hidden = this._deletingRlmChildren.has(childId) || this._deletedRlmChildIds.has(childId); if (!hidden) { + const snapshot = run + ? this._rlmChildSnapshotForRun(run, child) + : this._rlmChildSnapshotForSession(childId, child); snapshots.push({ - id: childId, - parentId: this._rlmParentNodeId, - sessionName: child.sessionName, - model: child.model ? `${child.model.provider}/${child.model.id}` : undefined, - label: child.sessionName ?? "child agent", - // A failed delete retains the session solely for cleanup retry. Preserve - // its cancellation truth in snapshots rather than reviving it as done. - status: this._rlmChildCleanupFailures.has(childId) ? "cancelled" : "done", - sessionDir: child._rlmSessionDir ?? child.sessionManager.getSessionDir(), + ...snapshot, + status: this._rlmChildCleanupFailures.has(childId) ? "cancelled" : snapshot.status, }); } snapshots.push(...child.getRlmChildSnapshots()); @@ -10525,7 +10583,7 @@ export class AgentSession { } } // A finished direct child can still have a running nested subagent. - for (const session of this._rlmChildSessions.values()) { + for (const { session } of this._rlmChildSessions.values()) { if (session.hasRunningRlmChildren()) { return true; } @@ -10535,7 +10593,7 @@ export class AgentSession { private _rlmChildSessionSnapshot(): AgentSession[] { const sessions = new Set(); - for (const [childId, session] of this._rlmChildSessions) { + for (const [childId, { session }] of this._rlmChildSessions) { if (!this._abandonedRlmQuiescenceChildIds.has(childId)) sessions.add(session); } for (const run of this._activeRlmChildRuns.values()) { @@ -10606,7 +10664,7 @@ export class AgentSession { // Inline (non-daemon) mode only; daemon clients attach to the child session directly. getRlmChildSession(childId: string): AgentSession | undefined { - const direct = this._activeRlmChildRuns.get(childId)?.session ?? this._rlmChildSessions.get(childId); + const direct = this._activeRlmChildRuns.get(childId)?.session ?? this._rlmChildSessions.get(childId)?.session; if (direct) { return direct; } @@ -10616,7 +10674,7 @@ export class AgentSession { return nested; } } - for (const retained of this._rlmChildSessions.values()) { + for (const { session: retained } of this._rlmChildSessions.values()) { const nested = retained.getRlmChildSession(childId); if (nested) { return nested; @@ -10646,7 +10704,7 @@ export class AgentSession { return true; } } - for (const retained of this._rlmChildSessions.values()) { + for (const { session: retained } of this._rlmChildSessions.values()) { if (retained.cancelRlmChildRun(childId, reason)) { return true; } @@ -10663,7 +10721,7 @@ export class AgentSession { [...this._activeRlmChildRuns.values()].some( (run) => run.session?.sessionName === name || (!run.session && run.sessionName === name), ) || - [...this._rlmChildSessions.values()].some((session) => session.sessionName === name) || + [...this._rlmChildSessions.values()].some(({ session }) => session.sessionName === name) || [...this._rlmChildCleanupFailures.values()].some((entry) => entry.session_name === name); if (localConflict) { throw new Error(formatAgentSessionNameUnavailable(name, depth)); @@ -10783,12 +10841,7 @@ export class AgentSession { if (!requestedSessionName) await this._assertRlmSubagentSessionNameAvailable(sessionName); const startedAt = Date.now(); const parentAssistantForUsage = this._findLastAssistantMessage(); - const label = rlmChildLabel(prompt); - let answerPreview: string | undefined; - let durationMs: number | undefined; - let toolUseCount = 0; let runningToolCount = 0; - let activity: RlmChildAgentActivity | undefined; let childSession: AgentSession | undefined; const run: RlmChildRun = { id: childNodeId, @@ -10797,6 +10850,7 @@ export class AgentSession { sessionDir: childSessionDir, model: modelSelection.model, status: "queued", + toolUseCount: 0, settled: false, abort: noopRlmChildAbort, publication: createAgentMessageDeferred(), @@ -10809,27 +10863,7 @@ export class AgentSession { this._activeRlmChildRuns.set(run.id, run); this._unsettledRlmChildRuns.add(run); const emitChildUpdate = () => { - const childModel = childSession?.model ?? modelSelection.model; - this._emit({ - type: "rlm_child_update", - child: { - id: childNodeId, - parentId: this._rlmParentNodeId, - sessionName: childSession?.sessionName ?? sessionName, - model: `${childModel.provider}/${childModel.id}`, - label, - status: run.status, - durationMs, - answerPreview, - toolUseCount: toolUseCount > 0 ? toolUseCount : undefined, - tokenCount: childSession?._contextTokensForCurrentMessages(), - recap: childSession?.getCurrentRecap(), - sessionDir: childSessionDir, - activity, - repliedSinceTask: childSession?._repliedToParentSinceTask, - error: run.error, - }, - }); + this._emit({ type: "rlm_child_update", child: this._rlmChildSnapshotForRun(run) }); }; run.emitUpdate = emitChildUpdate; emitChildUpdate(); @@ -10912,10 +10946,10 @@ export class AgentSession { return; } if (event.type === "agent_start") { - activity = { kind: "waiting" }; + run.activity = { kind: "waiting" }; emitChildUpdate(); } else if (event.type === "agent_end") { - activity = undefined; + run.activity = undefined; emitChildUpdate(); } else if (event.type === "message_end" && event.message.role === "assistant") { const assistant = event.message as AssistantMessage; @@ -10946,24 +10980,24 @@ export class AgentSession { } } const text = compactRlmText(readAssistantText(assistant)); - if (text) answerPreview = text; + if (text) run.answerPreview = text; void flushAgentTraceUpload(child.sessionManager).catch(() => undefined); emitChildUpdate(); } else if (event.type === "message_start" || event.type === "message_update") { if (event.message.role === "assistant") { const text = compactRlmText(readAssistantText(event.message as AssistantMessage)); - if (text) answerPreview = text; - activity = { kind: "writing" }; + if (text) run.answerPreview = text; + run.activity = { kind: "writing" }; emitChildUpdate(); } } else if (event.type === "tool_execution_start") { - toolUseCount += 1; + run.toolUseCount += 1; runningToolCount += 1; - activity = { kind: "executing", toolName: event.toolName }; + run.activity = { kind: "executing", toolName: event.toolName }; emitChildUpdate(); } else if (event.type === "tool_execution_end") { runningToolCount = Math.max(0, runningToolCount - 1); - if (runningToolCount === 0) activity = { kind: "waiting" }; + if (runningToolCount === 0) run.activity = { kind: "waiting" }; emitChildUpdate(); } else if (event.type === "session_info_changed" || event.type === "recap_update") { emitChildUpdate(); @@ -10998,8 +11032,8 @@ export class AgentSession { await child.waitForRlmQuiescence(); if (run.error) throw new Error(run.error); run.status = "done"; - durationMs = Date.now() - startedAt; - activity = undefined; + run.durationMs = Date.now() - startedAt; + run.activity = undefined; emitChildUpdate(); if ( !run.detachedDeletion && @@ -11032,8 +11066,8 @@ export class AgentSession { run.status = "error"; run.error = runError.message; } - durationMs = Date.now() - startedAt; - activity = undefined; + run.durationMs = Date.now() - startedAt; + run.activity = undefined; emitChildUpdate(); if (!run.detachedDeletion && !run.suppressTerminalNotice) { if (run.status === "error") { 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 3b6f4be8d4..d76ab50fef 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 @@ -42,6 +42,7 @@ import { collectDaemonClientEnv, collectDaemonLaunchEnv, type DaemonAttachResult, + type DaemonClientCapability, type DaemonCommand, type DaemonEventCursor, type DaemonOutbound, @@ -113,12 +114,23 @@ type DaemonSessionReplaced = Extract; + settled: boolean; + attachmentEpoch: number; + requestAttemptId?: number; promise: Promise; resolve: (snapshot: DaemonSessionSnapshot) => void; reject: (error: Error) => void; timeout: ReturnType; } +interface DaemonSnapshotRequestAttempt { + id: number; + activeSessionId: string; + attachmentEpoch: number; + candidateSnapshotIds: Set; + expectedSnapshotId?: string; +} + interface CorrelatedPromptRoute { activeSessionId: string; sessionId: string; @@ -314,9 +326,13 @@ export class DaemonAgentConnection implements AgentConnection { private readonly activeSideQuestionIds = new Set(); private readonly snapshotAssemblies = new Map(); private readonly completedSnapshots = new Map(); + private readonly completedSnapshotAttemptIds = new Map(); private readonly pendingReattachActiveSessionIds = new Set(); private readonly snapshotRecoveryPromises = new Map>(); private readonly ignoredSnapshotIds = new Set(); + private readonly snapshotRequestAttempts = new Map(); + private snapshotRequestAttemptSequence = 0; + private attachmentEpoch = 0; private reconnectPromise?: Promise; private readonly definitiveRequestErrors = new WeakSet(); private disposing = false; @@ -437,44 +453,64 @@ export class DaemonAgentConnection implements AgentConnection { requestedActiveSessionId: string, resumeCursor: DaemonEventCursor | undefined, resetEventProgress: boolean, + maxAttempts = 2, ): Promise { const supportsExtensionUi = this.options.supportsExtensionUi !== false; - const result = await this.requestData({ - type: "attach", - activeSessionId: requestedActiveSessionId, - supportsExtensionUi, - clientId: this.clientId, - capabilities: [ - "attach_snapshot", - "event_sequence", - ...(supportsExtensionUi ? (["extension_ui"] as const) : []), - "slim_attach", - "chunked_snapshot", - ...(this.supportsCorrelatedPromptLifecycle() ? (["correlated_prompt_lifecycle_v1"] as const) : []), - ...(this.options.ownedSession ? (["client_owned_sessions"] as const) : []), - ], - env: this.options.sendClientEnv ? collectDaemonClientEnv() : undefined, - launchEnv: this.options.ownedSession ? collectDaemonLaunchEnv() : undefined, - ...(this.options.ownedSession && - this.options.ownedSessionRecoveryConfig && - this.client.supportsServerCapability("owned_session_recovery_context") - ? { recoveryConfig: this.options.ownedSessionRecoveryConfig } - : {}), - telemetryDisabled: this.options.telemetryDisabled, - resumeCursor: - resumeCursor === undefined - ? undefined - : { - activeSessionId: requestedActiveSessionId, - ...resumeCursor, - }, - }); - if (this.disposed) throw new Error("Daemon connection disposed during attach"); + const attachmentEpoch = this.advanceAttachmentEpoch(); + let result!: SessionSummary | DaemonAttachResult; + let streamedSnapshot: DaemonSessionSnapshot | undefined; + for (let attemptIndex = 0; attemptIndex < maxAttempts; attemptIndex++) { + const requestAttempt = this.startSnapshotRequestAttempt(requestedActiveSessionId, attachmentEpoch); + try { + result = await this.requestData({ + type: "attach", + activeSessionId: requestedActiveSessionId, + supportsExtensionUi, + clientId: this.clientId, + capabilities: [ + "attach_snapshot", + "event_sequence", + ...(supportsExtensionUi ? (["extension_ui"] as const) : []), + "slim_attach", + "chunked_snapshot", + ...(this.supportsCorrelatedPromptLifecycle() ? (["correlated_prompt_lifecycle_v1"] as const) : []), + ...(this.options.ownedSession ? (["client_owned_sessions"] as const) : []), + ], + env: this.options.sendClientEnv ? collectDaemonClientEnv() : undefined, + launchEnv: this.options.ownedSession ? collectDaemonLaunchEnv() : undefined, + ...(this.options.ownedSession && + this.options.ownedSessionRecoveryConfig && + this.client.supportsServerCapability("owned_session_recovery_context") + ? { recoveryConfig: this.options.ownedSessionRecoveryConfig } + : {}), + telemetryDisabled: this.options.telemetryDisabled, + resumeCursor: + resumeCursor === undefined + ? undefined + : { activeSessionId: requestedActiveSessionId, ...resumeCursor }, + }); + if (this.disposed || this.attachmentEpoch !== attachmentEpoch) { + throw new Error("Daemon connection attachment changed during attach"); + } + if ("snapshot" in result && result.snapshotStream) { + requestAttempt.expectedSnapshotId = result.snapshotStream.id; + streamedSnapshot = await this.waitForSnapshot(result.snapshotStream.id, requestAttempt); + } else { + streamedSnapshot = undefined; + } + break; + } catch (error) { + if (attemptIndex === maxAttempts - 1 || this.disposed || this.attachmentEpoch !== attachmentEpoch) { + throw error; + } + } finally { + this.finishSnapshotRequestAttempt(requestAttempt); + } + } const nextActiveSessionId = getAttachActiveSessionId(result); + if ("snapshot" in result) { - const snapshot = result.snapshotStream - ? await this.waitForSnapshot(result.snapshotStream.id) - : result.snapshot; + const snapshot = streamedSnapshot ?? result.snapshot; let staged: StagedSnapshotCommit; try { validateDaemonSnapshotIdentity(result.snapshot, nextActiveSessionId); @@ -1543,32 +1579,65 @@ export class DaemonAgentConnection implements AgentConnection { let reattached = false; try { const supportsExtensionUi = this.options.supportsExtensionUi !== false; - const result = await this.requestData({ - type: "reattach", - activeSessionId: sourceActiveSessionId, - targetActiveSessionId, - supportsExtensionUi, - clientId: this.clientId, - capabilities: [ - "attach_snapshot", - "event_sequence", - ...(supportsExtensionUi ? (["extension_ui"] as const) : []), - "slim_attach", - "chunked_snapshot", - ...(this.supportsCorrelatedPromptLifecycle() ? (["correlated_prompt_lifecycle_v1"] as const) : []), - ...(this.options.ownedSession ? (["client_owned_sessions"] as const) : []), - ], - env: this.options.sendClientEnv ? collectDaemonClientEnv() : undefined, - launchEnv: this.options.ownedSession ? collectDaemonLaunchEnv() : undefined, - telemetryDisabled: this.options.telemetryDisabled, - }); - reattached = true; - if (result.activeSessionId !== targetActiveSessionId) { - throw new Error("Daemon returned an invalid session snapshot"); + const attachmentEpoch = this.advanceAttachmentEpoch(); + const capabilities: DaemonClientCapability[] = [ + "attach_snapshot", + "event_sequence", + ...(supportsExtensionUi ? (["extension_ui"] as const) : []), + "slim_attach", + "chunked_snapshot", + ...(this.supportsCorrelatedPromptLifecycle() ? (["correlated_prompt_lifecycle_v1"] as const) : []), + ...(this.options.ownedSession ? (["client_owned_sessions"] as const) : []), + ]; + let result!: DaemonAttachResult; + let snapshot!: DaemonSessionSnapshot; + for (let attemptIndex = 0; attemptIndex < 2; attemptIndex++) { + const requestAttempt = this.startSnapshotRequestAttempt(targetActiveSessionId, attachmentEpoch); + try { + result = await this.requestData( + attemptIndex === 0 + ? { + type: "reattach", + activeSessionId: sourceActiveSessionId, + targetActiveSessionId, + supportsExtensionUi, + clientId: this.clientId, + capabilities, + env: this.options.sendClientEnv ? collectDaemonClientEnv() : undefined, + launchEnv: this.options.ownedSession ? collectDaemonLaunchEnv() : undefined, + telemetryDisabled: this.options.telemetryDisabled, + } + : { + type: "attach", + activeSessionId: targetActiveSessionId, + supportsExtensionUi, + clientId: this.clientId, + capabilities, + env: this.options.sendClientEnv ? collectDaemonClientEnv() : undefined, + telemetryDisabled: this.options.telemetryDisabled, + }, + ); + reattached = true; + if (result.activeSessionId !== targetActiveSessionId) { + throw new Error("Daemon returned an invalid session snapshot"); + } + if (result.snapshotStream) { + requestAttempt.expectedSnapshotId = result.snapshotStream.id; + snapshot = await this.waitForSnapshot(result.snapshotStream.id, requestAttempt); + } else { + snapshot = result.snapshot; + } + break; + } catch (error) { + if (attemptIndex === 1 || this.disposed || this.attachmentEpoch !== attachmentEpoch) throw error; + } finally { + this.finishSnapshotRequestAttempt(requestAttempt); + } + } + if (this.disposed || this.attachmentEpoch !== attachmentEpoch) { + throw new Error("Daemon connection attachment changed during reattach"); } - const snapshot = result.snapshotStream - ? await this.waitForSnapshot(result.snapshotStream.id) - : result.snapshot; + validateDaemonSnapshotIdentity(result.snapshot, result.activeSessionId); mapDaemonSessionSnapshot(result.snapshot, undefined, this.supportsCorrelatedPromptLifecycle()); const staged = this.stageSnapshotCommit(snapshot, { @@ -1712,6 +1781,7 @@ export class DaemonAgentConnection implements AgentConnection { return; } this.disposing = true; + this.advanceAttachmentEpoch(); if (this.options.ownedSession && !this.client.isConnected && this.reconnectPromise) { await Promise.race([this.reconnectPromise, delay(OWNED_SESSION_DISPOSE_RECONNECT_WAIT_MS)]).catch( () => undefined, @@ -1838,7 +1908,11 @@ export class DaemonAgentConnection implements AgentConnection { if (!this.isMessageForActiveSession(message)) { return; } - if ("snapshotId" in message && this.ignoredSnapshotIds.has(message.snapshotId)) { + if ( + "snapshotId" in message && + message.type !== "session_snapshot_begin" && + this.ignoredSnapshotIds.has(message.snapshotId) + ) { if (message.type === "session_snapshot_end" || message.type === "session_snapshot_failed") { this.ignoredSnapshotIds.delete(message.snapshotId); } @@ -1846,36 +1920,102 @@ export class DaemonAgentConnection implements AgentConnection { } if (message.type === "session_snapshot_begin") { const pendingReplacement = this.pendingChunkedReplacement; + const requestAttempt = this.latestSnapshotRequestAttempt(message.activeSessionId); + const acceptsUnsolicitedRuntimeSnapshot = + this.attachedSessionId !== undefined && + (message.purpose === "resync" || + (message.purpose === "replacement" && !this.supportsCorrelatedPromptLifecycle())); + if (!pendingReplacement && !requestAttempt && !acceptsUnsolicitedRuntimeSnapshot) return; + this.ignoredSnapshotIds.delete(message.snapshotId); if (pendingReplacement) { - if ( - message.purpose !== "replacement" || - (pendingReplacement.snapshotId !== undefined && pendingReplacement.snapshotId !== message.snapshotId) - ) { + if (message.purpose !== "replacement") { this.failClosedReplacementReconciliation(); return; } + if (pendingReplacement.snapshotId && pendingReplacement.snapshotId !== message.snapshotId) { + const failedSnapshotId = pendingReplacement.snapshotId; + const failedAssembly = this.snapshotAssemblies.get(failedSnapshotId); + if (failedAssembly) { + this.rejectSnapshotAssembly( + failedSnapshotId, + failedAssembly, + new Error(`Snapshot ${failedSnapshotId} was replaced by retry ${message.snapshotId}`), + ); + this.snapshotAssemblies.delete(failedSnapshotId); + } + this.ignoreSnapshotId(failedSnapshotId); + } pendingReplacement.snapshotId = message.snapshotId; } - const assembly = this.getSnapshotAssembly(message.snapshotId); + let assembly = this.snapshotAssemblies.get(message.snapshotId); + if (assembly && (assembly.settled || assembly.attachmentEpoch !== this.attachmentEpoch)) { + clearTimeout(assembly.timeout); + this.snapshotAssemblies.delete(message.snapshotId); + assembly = undefined; + } + assembly ??= this.getSnapshotAssembly(message.snapshotId, requestAttempt); + if (requestAttempt) { + requestAttempt.candidateSnapshotIds.add(message.snapshotId); + assembly.requestAttemptId = requestAttempt.id; + } + try { + validateDaemonSnapshotBegin(message); + } catch (error) { + await this.rejectInvalidSnapshotAssembly( + message.snapshotId, + assembly, + error instanceof Error ? error : new Error(String(error)), + message.activeSessionId, + ); + return; + } + // A begin is the immutable boundary. Never retain chunks from an older transfer that reused an id. + assembly.chunks.clear(); + this.completedSnapshots.delete(message.snapshotId); + this.completedSnapshotAttemptIds.delete(message.snapshotId); assembly.begin = message; return; } if (message.type === "session_snapshot_chunk") { - this.getSnapshotAssembly(message.snapshotId).chunks.set(message.index, message.messages); + const assembly = this.snapshotAssemblies.get(message.snapshotId); + if (!assembly || assembly.settled || !assembly.begin || assembly.attachmentEpoch !== this.attachmentEpoch) { + return; + } + if ( + message.activeSessionId !== assembly.begin.activeSessionId || + !Number.isSafeInteger(message.index) || + message.index < 0 || + !Array.isArray(message.messages) || + assembly.chunks.has(message.index) + ) { + await this.rejectInvalidSnapshotAssembly( + message.snapshotId, + assembly, + new Error(`Snapshot ${message.snapshotId} returned an invalid chunk`), + message.activeSessionId, + ); + return; + } + assembly.chunks.set(message.index, message.messages); return; } if (message.type === "session_snapshot_end") { + const assembly = this.snapshotAssemblies.get(message.snapshotId); + if (!assembly || assembly.settled || assembly.attachmentEpoch !== this.attachmentEpoch) return; await this.completeSnapshotAssembly(message); return; } if (message.type === "session_snapshot_failed") { - const assembly = this.getSnapshotAssembly(message.snapshotId); - const purpose = assembly.begin?.purpose ?? "attach"; + const existingAssembly = this.snapshotAssemblies.get(message.snapshotId); + const assembly = + existingAssembly ?? + (this.pendingChunkedReplacement ? this.getSnapshotAssembly(message.snapshotId) : undefined); + if (!assembly || assembly.settled || assembly.attachmentEpoch !== this.attachmentEpoch) return; + const purpose = assembly.begin?.purpose ?? (this.pendingChunkedReplacement ? "replacement" : "attach"); const snapshotError = new Error(message.error); const pendingReattach = this.pendingReattachActiveSessionIds.has(message.activeSessionId); - const quarantinesReplacement = purpose === "replacement" && this.supportsCorrelatedPromptLifecycle(); const recoveryPromise = - !pendingReattach && !quarantinesReplacement && (purpose === "replacement" || purpose === "resync") + !pendingReattach && (purpose === "replacement" || purpose === "resync") ? this.recoverFailedSnapshot(purpose, snapshotError) : undefined; if (recoveryPromise) { @@ -1883,7 +2023,6 @@ export class DaemonAgentConnection implements AgentConnection { } this.rejectSnapshotAssembly(message.snapshotId, assembly, snapshotError); this.ignoreSnapshotId(message.snapshotId); - if (quarantinesReplacement) this.failClosedReplacementReconciliation(); if (recoveryPromise) { try { await recoveryPromise; @@ -2213,11 +2352,84 @@ export class DaemonAgentConnection implements AgentConnection { throw lastError ?? new Error("the restored session did not become available"); } - private getSnapshotAssembly(snapshotId: string): DaemonSnapshotAssembly { + private advanceAttachmentEpoch(): number { + const nextEpoch = ++this.attachmentEpoch; + for (const [snapshotId, assembly] of this.snapshotAssemblies) { + if (assembly.attachmentEpoch === nextEpoch) continue; + this.rejectSnapshotAssembly( + snapshotId, + assembly, + new Error(`Attachment changed while waiting for snapshot ${snapshotId}`), + ); + this.snapshotAssemblies.delete(snapshotId); + this.completedSnapshots.delete(snapshotId); + this.completedSnapshotAttemptIds.delete(snapshotId); + this.ignoreSnapshotId(snapshotId); + } + return nextEpoch; + } + + private startSnapshotRequestAttempt(activeSessionId: string, attachmentEpoch: number): DaemonSnapshotRequestAttempt { + const attempt: DaemonSnapshotRequestAttempt = { + id: ++this.snapshotRequestAttemptSequence, + activeSessionId, + attachmentEpoch, + candidateSnapshotIds: new Set(), + }; + this.snapshotRequestAttempts.set(attempt.id, attempt); + return attempt; + } + + private latestSnapshotRequestAttempt(activeSessionId: string): DaemonSnapshotRequestAttempt | undefined { + return [...this.snapshotRequestAttempts.values()] + .reverse() + .find( + (attempt) => + attempt.activeSessionId === activeSessionId && attempt.attachmentEpoch === this.attachmentEpoch, + ); + } + + private finishSnapshotRequestAttempt(attempt: DaemonSnapshotRequestAttempt): void { + if (this.snapshotRequestAttempts.get(attempt.id) !== attempt) return; + this.snapshotRequestAttempts.delete(attempt.id); + for (const snapshotId of attempt.candidateSnapshotIds) { + const assembly = this.snapshotAssemblies.get(snapshotId); + if (assembly?.requestAttemptId === attempt.id) { + this.rejectSnapshotAssembly( + snapshotId, + assembly, + new Error(`Snapshot ${snapshotId} arrived after its request attempt settled`), + ); + this.snapshotAssemblies.delete(snapshotId); + this.ignoreSnapshotId(snapshotId); + } + if (this.completedSnapshotAttemptIds.get(snapshotId) === attempt.id) { + this.completedSnapshots.delete(snapshotId); + this.completedSnapshotAttemptIds.delete(snapshotId); + } + } + } + + private getSnapshotAssembly( + snapshotId: string, + requestAttempt?: DaemonSnapshotRequestAttempt, + ): DaemonSnapshotAssembly { const existing = this.snapshotAssemblies.get(snapshotId); - if (existing) { + if ( + existing && + existing.attachmentEpoch === this.attachmentEpoch && + (!requestAttempt || existing.requestAttemptId === undefined || existing.requestAttemptId === requestAttempt.id) + ) { return existing; } + if (existing) { + this.rejectSnapshotAssembly( + snapshotId, + existing, + new Error(`Snapshot ${snapshotId} was superseded by a fresh transfer`), + ); + this.snapshotAssemblies.delete(snapshotId); + } let resolveSnapshot!: (snapshot: DaemonSessionSnapshot) => void; let rejectSnapshot!: (error: Error) => void; const promise = new Promise((resolve, reject) => { @@ -2227,7 +2439,7 @@ export class DaemonAgentConnection implements AgentConnection { void promise.catch(() => undefined); const timeout = setTimeout(() => { const current = this.snapshotAssemblies.get(snapshotId); - if (current) { + if (current && !current.settled) { current.reject(new Error(`Timed out waiting for snapshot ${snapshotId}`)); this.snapshotAssemblies.delete(snapshotId); this.ignoreSnapshotId(snapshotId); @@ -2243,9 +2455,20 @@ export class DaemonAgentConnection implements AgentConnection { timeout.unref(); const assembly: DaemonSnapshotAssembly = { chunks: new Map(), + settled: false, + attachmentEpoch: this.attachmentEpoch, + ...(requestAttempt ? { requestAttemptId: requestAttempt.id } : {}), promise, - resolve: resolveSnapshot, - reject: rejectSnapshot, + resolve(snapshot) { + if (assembly.settled) return; + assembly.settled = true; + resolveSnapshot(snapshot); + }, + reject(error) { + if (assembly.settled) return; + assembly.settled = true; + rejectSnapshot(error); + }, timeout, }; this.snapshotAssemblies.set(snapshotId, assembly); @@ -2259,6 +2482,8 @@ export class DaemonAgentConnection implements AgentConnection { } this.snapshotAssemblies.clear(); this.completedSnapshots.clear(); + this.completedSnapshotAttemptIds.clear(); + this.snapshotRequestAttempts.clear(); this.snapshotRecoveryPromises.clear(); this.ignoredSnapshotIds.clear(); } @@ -2305,12 +2530,9 @@ export class DaemonAgentConnection implements AgentConnection { private async recoverFailedSnapshot(purpose: "replacement" | "resync", snapshotError: Error): Promise { this.latestSnapshotIsFresh = false; try { + await this.attachSession(this.activeSessionId, this.lastEventCursor, false, 1); const snapshot = await this.getInitialSnapshot(); - if (this.disposed) { - return; - } - this.attachedSessionId = snapshot.state.sessionId; - this.attachedSessionFile = snapshot.state.sessionFile; + if (this.disposed) return; if (purpose === "replacement") { await this.emit({ type: "session_replaced", state: snapshot.state, messages: snapshot.messages }); } else { @@ -2328,19 +2550,31 @@ export class DaemonAgentConnection implements AgentConnection { } } - private async waitForSnapshot(snapshotId: string): Promise { + private async waitForSnapshot( + snapshotId: string, + requestAttempt?: DaemonSnapshotRequestAttempt, + ): Promise { const completed = this.completedSnapshots.get(snapshotId); - if (completed) { + const completedAttemptId = this.completedSnapshotAttemptIds.get(snapshotId); + if (completed && (!requestAttempt || completedAttemptId === requestAttempt.id)) { this.completedSnapshots.delete(snapshotId); + this.completedSnapshotAttemptIds.delete(snapshotId); return completed; } - const assembly = this.getSnapshotAssembly(snapshotId); + const assembly = this.getSnapshotAssembly(snapshotId, requestAttempt); + if (requestAttempt) { + requestAttempt.candidateSnapshotIds.add(snapshotId); + assembly.requestAttemptId = requestAttempt.id; + } try { return await assembly.promise; } finally { clearTimeout(assembly.timeout); - this.snapshotAssemblies.delete(snapshotId); + if (this.snapshotAssemblies.get(snapshotId) === assembly) { + this.snapshotAssemblies.delete(snapshotId); + } this.completedSnapshots.delete(snapshotId); + this.completedSnapshotAttemptIds.delete(snapshotId); } } @@ -2464,6 +2698,24 @@ export class DaemonAgentConnection implements AgentConnection { ); return; } + const beginCursor = assembly.begin.snapshot.lastEventCursor; + if ( + message.activeSessionId !== assembly.begin.activeSessionId || + !Number.isSafeInteger(message.chunkCount) || + message.chunkCount < 0 || + message.lastEventSequence !== assembly.begin.snapshot.lastEventSequence || + (message.lastEventCursor !== undefined && + (message.lastEventCursor.generation !== beginCursor?.generation || + message.lastEventCursor.sequence !== beginCursor?.sequence)) + ) { + await this.rejectInvalidSnapshotAssembly( + message.snapshotId, + assembly, + new Error(`Snapshot ${message.snapshotId} ended with invalid metadata`), + message.activeSessionId, + ); + return; + } if (assembly.chunks.size !== message.chunkCount) { await this.rejectInvalidSnapshotAssembly( message.snapshotId, @@ -2504,7 +2756,7 @@ export class DaemonAgentConnection implements AgentConnection { ...assembly.begin.snapshot, messages, lastEventSequence: message.lastEventSequence, - lastEventCursor: message.lastEventCursor, + lastEventCursor: message.lastEventCursor ?? assembly.begin.snapshot.lastEventCursor, }; const purpose = assembly.begin.purpose ?? "attach"; const pendingReattach = this.pendingReattachActiveSessionIds.has(message.activeSessionId); @@ -2514,10 +2766,12 @@ export class DaemonAgentConnection implements AgentConnection { if (pendingReattach) { this.snapshotAssemblies.delete(message.snapshotId); this.completedSnapshots.set(message.snapshotId, snapshot); + this.completedSnapshotAttemptIds.set(message.snapshotId, assembly.requestAttemptId); while (this.completedSnapshots.size > MAX_COMPLETED_SNAPSHOTS) { const oldest = this.completedSnapshots.keys().next().value; if (oldest === undefined) break; this.completedSnapshots.delete(oldest); + this.completedSnapshotAttemptIds.delete(oldest); } } return; @@ -2780,8 +3034,27 @@ function validateConnectionStateIdentity(state: AgentConnectionState, activeSess } } +function validateDaemonSnapshotBegin(begin: DaemonSnapshotBegin): void { + validateDaemonSnapshotIdentity(begin.snapshot, begin.activeSessionId); + if ( + !Number.isSafeInteger(begin.messageCount) || + begin.messageCount < 0 || + !Number.isSafeInteger(begin.targetChunkBytes) || + begin.targetChunkBytes <= 0 || + begin.snapshot.summary.messageCount !== begin.messageCount || + !Number.isSafeInteger(begin.snapshot.lastEventSequence) || + begin.snapshot.lastEventSequence < 0 || + (begin.snapshot.lastEventCursor !== undefined && + (begin.snapshot.lastEventCursor.sequence !== begin.snapshot.lastEventSequence || + typeof begin.snapshot.lastEventCursor.generation !== "string" || + begin.snapshot.lastEventCursor.generation.length === 0)) + ) { + throw new Error("Daemon returned an invalid snapshot begin frame"); + } +} + function validateDaemonSnapshotIdentity( - snapshot: DaemonSessionSnapshot, + snapshot: Omit, activeSessionId: string, expectedSessionId?: string, expectedState?: AgentConnectionState, diff --git a/packages/coding-agent/src/modes/daemon/daemon-mode.ts b/packages/coding-agent/src/modes/daemon/daemon-mode.ts index 0d58993757..ffed26bc70 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-mode.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-mode.ts @@ -12,6 +12,8 @@ import { existsSync, mkdirSync, readFileSync, renameSync, rmSync, writeFileSync import { readFile, stat } from "node:fs/promises"; import { createConnection, createServer, type Server, type Socket } from "node:net"; import { dirname, isAbsolute, join, resolve } from "node:path"; +import { isDeepStrictEqual } from "node:util"; +import type { AgentMessage } from "@earendil-works/pi-agent-core"; import { type Api, getLogger, type Model } from "@earendil-works/pi-ai"; import { createCliSubprocessEnv, createCliSubprocessLaunchSpec } from "../../cli/subprocess-launch.js"; import { @@ -117,7 +119,7 @@ import { createAgentConnectionToolDefinition } from "../agent-connection/tool-de import type { AgentConnectionHeartbeat, AgentConnectionRlmChildAgentSnapshot } 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 { encodePrivateFrameParts, PrivateFrameDecoder } from "../session-worker/private-framing.js"; import { type ActiveSessionState, AmbiguousActiveSessionError, @@ -206,9 +208,12 @@ import { } from "./rlm-subagent-display.js"; import { serializeSavedSessionInfo } from "./saved-session-info.js"; import { - createSnapshotTranscriptChunks, + createImmutableSnapshotMessages, + prepareSnapshotTranscriptPayload, SNAPSHOT_TARGET_CHUNK_BYTES, type SnapshotTranscriptChunkSource, + type SnapshotTranscriptPayloadCache, + type SnapshotTranscriptWireChunk, } from "./snapshot-transcript-cache.js"; import { WorkerRecoveryJournal } from "./worker-recovery-journal.js"; @@ -236,9 +241,14 @@ export { defaultDaemonSocketPath } from "./daemon-socket.js"; const structuredLog = getLogger("coding-agent.daemon"); const WORKER_SNAPSHOT_TERMINAL_DRAIN_TIMEOUT_MS = 1_000; +const WORKER_SNAPSHOT_PREPARATION_TIMEOUT_MS = 30_000; const UPDATE_RESTART_PREPARE_TIMEOUT_MS = 90_000; const MAX_SESSION_SNAPSHOT_STABILIZATION_RETRIES = 3; +function createSnapshotTransferId(): string { + return `snapshot-${randomUUID()}`; +} + const DAEMON_COMMAND_TYPES: ReadonlySet = new Set([ "ack_result", "list", @@ -486,6 +496,14 @@ export class AgentDaemon { private socketIdentity?: DaemonSocketIdentity; private readonly clients = new Set(); private readonly sessions = new Map(); + private readonly snapshotPayloadGenerations = new Map< + string, + { + key: string; + messages: readonly AgentMessage[]; + promise: Promise; + } + >(); private readonly openingSessions = new Map>(); /** Covers path resolution through publication in openingSessions, before the runtime promise exists. */ private readonly reservingSessionOpens = new Map>(); @@ -2697,12 +2715,9 @@ export class AgentDaemon { try { await this.closeSession(state, "shutdown", true, false); } catch (error) { - // A pre-removal close failure must not strand a resident child outside its - // parent's ownership map or disconnect its event forwarder. if ( this.sessions.get(state.activeSessionId) === state && - this.sessions.get(parentActiveSessionId) === parentState && - parentState.runtime.session.registerRlmChildSession(childId, state.runtime.session, unsubscribeChild) + this.sessions.get(parentActiveSessionId) === parentState ) { throw error; } @@ -3874,14 +3889,14 @@ export class AgentDaemon { }); } if (streamsSnapshot) { - const snapshotId = `${state.activeSessionId}-${state.eventGeneration}-${state.lastEventSequence}`; + const snapshotId = createSnapshotTransferId(); let transcript: SnapshotTranscriptChunkSource; try { - transcript = createSnapshotTranscriptChunks({ + transcript = await this.prepareWorkerSnapshotTranscript({ activeSessionId: state.activeSessionId, snapshotId, + generationKey: this.snapshotPayloadGenerationKey(result), messages: result.snapshot.messages, - targetChunkBytes: SNAPSHOT_TARGET_CHUNK_BYTES, signal: snapshotSignal, }); } catch (error) { @@ -5139,11 +5154,12 @@ export class AgentDaemon { } session = state.runtime.session; const connectionState = this.createConnectionState(state); + const messages = createImmutableSnapshotMessages(session.messages); return { activeSessionId: state.activeSessionId, summary: summaryForActiveSession(state), state: connectionState, - messages: session.messages, + messages, // Omit duplicate heavy payloads from attach. The client can derive render // context from messages + state, and fetch the full session tree lazily // when the tree/branch selector opens. @@ -5160,6 +5176,66 @@ export class AgentDaemon { }; } + private async prepareWorkerSnapshotTranscript(options: { + activeSessionId: string; + snapshotId: string; + generationKey: string; + messages: readonly AgentMessage[]; + signal?: AbortSignal; + }): Promise { + options.signal?.throwIfAborted(); + while (true) { + const existing = this.snapshotPayloadGenerations.get(options.activeSessionId); + if ( + existing && + existing.key === options.generationKey && + isDeepStrictEqual(existing.messages, options.messages) + ) { + const payload = await existing.promise; + options.signal?.throwIfAborted(); + return payload.createTransfer(options.activeSessionId, options.snapshotId); + } + if (existing) { + const stalePayload = await existing.promise.catch(() => undefined); + if (this.snapshotPayloadGenerations.get(options.activeSessionId) === existing) { + this.snapshotPayloadGenerations.delete(options.activeSessionId); + stalePayload?.dispose(); + } + continue; + } + const generation = { + key: options.generationKey, + messages: options.messages, + promise: prepareSnapshotTranscriptPayload({ + messages: options.messages, + targetChunkBytes: SNAPSHOT_TARGET_CHUNK_BYTES, + signal: AbortSignal.timeout(WORKER_SNAPSHOT_PREPARATION_TIMEOUT_MS), + }), + }; + this.snapshotPayloadGenerations.set(options.activeSessionId, generation); + try { + const payload = await generation.promise; + options.signal?.throwIfAborted(); + return payload.createTransfer(options.activeSessionId, options.snapshotId); + } catch (error) { + if (this.snapshotPayloadGenerations.get(options.activeSessionId) === generation) { + this.snapshotPayloadGenerations.delete(options.activeSessionId); + } + throw error; + } + } + } + + private snapshotPayloadGenerationKey(result: DaemonAttachResult): string { + const cursor = result.lastEventCursor; + return [ + result.snapshot.state.sessionId, + cursor?.generation ?? "legacy", + cursor?.sequence ?? result.lastEventSequence, + result.snapshot.messages.length, + ].join(":"); + } + private async streamWorkerSnapshot( client: DaemonSocketClient, result: DaemonAttachResult, @@ -5336,7 +5412,7 @@ export class AgentDaemon { private async writeWorkerSnapshotBuffer( client: DaemonSocketClient, - buffer: Buffer, + buffer: SnapshotTranscriptWireChunk, message: DaemonOutbound, purpose: "attach" | "replacement" | "catchup", signal?: AbortSignal, @@ -6560,6 +6636,14 @@ export class AgentDaemon { } state.clients.clear(); this.acpMcpOwners.delete(state.activeSessionId); + const snapshotPayload = this.snapshotPayloadGenerations.get(state.activeSessionId); + this.snapshotPayloadGenerations.delete(state.activeSessionId); + if (snapshotPayload) { + void snapshotPayload.promise.then( + (payload) => payload.dispose(), + () => undefined, + ); + } this.sessions.delete(state.activeSessionId); if (isEmptyDraftSession) { const sessionFile = state.runtime.session.sessionFile; @@ -6680,7 +6764,7 @@ export class AgentDaemon { state: ActiveSessionState, message: Extract, ): void { - const snapshotId = `${state.activeSessionId}-${state.eventGeneration}-${state.lastEventSequence}`; + const snapshotId = createSnapshotTransferId(); // Mark before the registry read so later events queue behind this snapshot. const snapshotSignal = markClientSnapshotStreaming(client, state.activeSessionId); void this.prepareReplacementSnapshot(client, state, message, snapshotId, snapshotSignal).catch((error) => { @@ -6717,11 +6801,11 @@ export class AgentDaemon { } return; } - const transcript = createSnapshotTranscriptChunks({ + const transcript = await this.prepareWorkerSnapshotTranscript({ activeSessionId: state.activeSessionId, snapshotId, + generationKey: this.snapshotPayloadGenerationKey(result), messages: result.snapshot.messages, - targetChunkBytes: SNAPSHOT_TARGET_CHUNK_BYTES, signal: snapshotSignal, }); this.write(client, { ...message, messages: [], snapshotFollows: true }); @@ -6880,15 +6964,15 @@ export class AgentDaemon { ), }); } - const snapshotId = `${activeSessionId}-${state.eventGeneration}-${state.lastEventSequence}`; + const snapshotId = createSnapshotTransferId(); const snapshotSignal = markClientSnapshotStreaming(client, activeSessionId); let transcript: SnapshotTranscriptChunkSource; try { - transcript = createSnapshotTranscriptChunks({ + transcript = await this.prepareWorkerSnapshotTranscript({ activeSessionId, snapshotId, + generationKey: this.snapshotPayloadGenerationKey(result), messages: result.snapshot.messages, - targetChunkBytes: SNAPSHOT_TARGET_CHUNK_BYTES, signal: snapshotSignal, }); } catch (error) { @@ -7007,17 +7091,16 @@ export class AgentDaemon { private writeSerialized( client: DaemonSocketClient, - line: string | Buffer, + line: string | Buffer | readonly Buffer[], message: DaemonOutbound, payloadEncoding: "jsonl" | "assistant-delta" = "jsonl", snapshotPurpose?: "attach" | "replacement" | "catchup", ): boolean { - if (client.socket.destroyed) { - return false; - } - const wireData = + if (client.socket.destroyed) return false; + const payloadParts = typeof line === "string" ? [Buffer.from(line)] : Buffer.isBuffer(line) ? [line] : [...line]; + const wireParts = client.transport === "private-framed" - ? encodePrivateFrame( + ? encodePrivateFrameParts( { kind: "outbound", outboundType: message.type, @@ -7030,13 +7113,17 @@ export class AgentDaemon { payloadEncoding, ...(snapshotPurpose ? { snapshotPurpose } : {}), }, - typeof line === "string" ? Buffer.from(line) : line, + payloadParts, ) - : line; - const accepted = client.socket.write(wireData); - if (!accepted) { - client.backpressured = true; + : payloadParts; + let accepted = true; + if (wireParts.length > 1) client.socket.cork?.(); + try { + for (const part of wireParts) accepted = client.socket.write(part) && accepted; + } finally { + if (wireParts.length > 1) client.socket.uncork?.(); } + if (!accepted) client.backpressured = true; return accepted; } diff --git a/packages/coding-agent/src/modes/daemon/daemon-protocol.ts b/packages/coding-agent/src/modes/daemon/daemon-protocol.ts index 4331171d9c..27dfb04aac 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-protocol.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-protocol.ts @@ -75,8 +75,9 @@ export const DAEMON_COMMAND_ENVELOPE_MIN_PROTOCOL_VERSION = 7; // Revision 22 scopes ACP MCP replacement and cleanup to a connection owner. // Revision 23 lets workers query the supervisor agent roster on demand. // Revision 24 adds capability-gated correlated prompt lifecycle and event provenance. -export const DAEMON_SCHEMA_REVISION = 24; -export const DAEMON_SCHEMA_ID = "protocol-7-schema-24-2c8fb17d3895"; +// Revision 25 negotiates immutable snapshot transfer identities with session workers. +export const DAEMON_SCHEMA_REVISION = 25; +export const DAEMON_SCHEMA_ID = "protocol-7-schema-25-2c8fb17d3895"; export type DaemonProtocolName = typeof DAEMON_PROTOCOL_NAME; export type DaemonProtocolVersion = number; @@ -94,6 +95,7 @@ export type DaemonClientCapability = | "extension_ui" | "slim_attach" | "chunked_snapshot" + | "immutable_snapshot_transfer_v1" | "client_owned_sessions" | "correlated_prompt_lifecycle_v1"; export type DaemonPromptAdmissionCancellationStatus = "cancelled" | "owned" | "unknown"; @@ -147,6 +149,7 @@ export const DAEMON_SUPPORTED_CLIENT_CAPABILITIES: readonly DaemonClientCapabili "extension_ui", "slim_attach", "chunked_snapshot", + "immutable_snapshot_transfer_v1", "client_owned_sessions", "correlated_prompt_lifecycle_v1", ]; 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 49b5a16dc7..d41e83a67a 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-session-list.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-session-list.ts @@ -2,8 +2,7 @@ import { statSync } from "node:fs"; import { resolve } from "node:path"; import type { AgentMessage, ThinkingLevel } from "@earendil-works/pi-agent-core"; import type { Api, Model } from "@earendil-works/pi-ai"; -import { compactRlmText, rlmChildLabel } from "../../core/agent-session.js"; -import type { AgentSessionRuntimeMetadata } from "../../core/agent-session-runtime.js"; +import { compactRlmText } from "../../core/agent-session.js"; import type { AgentSessionRuntimeDiagnostic } from "../../core/agent-session-services.js"; import { type AgentCronJob, isHeartbeatCronJob } from "../../core/cron-jobs.js"; import type { SessionActionSnapshot } from "../../core/session-action-store.js"; @@ -329,87 +328,23 @@ export function summaryForInactiveSession( }; } -/** - * Build snapshots for all RLM child sessions hosted by the daemon under the - * given session, including grandchildren. Mirrors the shape of live - * rlm_child_update events so attach clients can seed their subagent state - * from daemon memory instead of replaying the event stream. - */ +/** Build the root AgentSession projection with daemon-only active session ids. */ export function buildRlmChildSnapshots( rootActiveSessionId: string, activeSessions: readonly ActiveSessionState[], ): AgentConnectionRlmChildAgentSnapshot[] { - const childrenByParent = new Map(); - for (const candidate of activeSessions) { - const metadata = candidate.runtime.metadata; - if (metadata.kind !== "subagent" || !metadata.parentActiveSessionId) { - continue; - } - const siblings = childrenByParent.get(metadata.parentActiveSessionId) ?? []; - siblings.push(candidate); - childrenByParent.set(metadata.parentActiveSessionId, siblings); - } - - const snapshots: AgentConnectionRlmChildAgentSnapshot[] = []; - const visit = (parent: ActiveSessionState | undefined, parentActiveSessionId: string): void => { - const parentNodeId = parent?.runtime.metadata.rlmChildId; - for (const child of childrenByParent.get(parentActiveSessionId) ?? []) { - snapshots.push(rlmChildSnapshotForActiveSession(child, child.runtime.metadata, parentNodeId, parent)); - // A child passes its own node id to its children as their parent id. - visit(child, child.activeSessionId); - } - }; const root = activeSessions.find((candidate) => candidate.activeSessionId === rootActiveSessionId); - visit(root, rootActiveSessionId); - return snapshots; -} - -function rlmChildSnapshotForActiveSession( - activeSession: ActiveSessionState, - metadata: AgentSessionRuntimeMetadata, - parentNodeId: string | undefined, - parent: ActiveSessionState | undefined, -): AgentConnectionRlmChildAgentSnapshot { - const session = activeSession.runtime.session; - let answerPreview: string | undefined; - let toolUseCount = 0; - const messages = - session.state.streamingMessage?.role === "assistant" - ? [...session.messages, session.state.streamingMessage] - : session.messages; - for (const message of messages) { - if (message.role === "assistant") { - const text = compactRlmText(readMessageText(message.content)); - if (text) { - answerPreview = text; - } - toolUseCount += message.content.filter((block) => block.type === "toolCall").length; - } - } - // The parent session's run tracker is the source of truth for child status; - // a daemon-hosted child whose agent is momentarily idle is still part of an - // active run. The streaming heuristic only covers parents the daemon does - // not host (e.g. children attributed to a session created by an older build). - const runStatus = metadata.rlmChildId - ? parent?.runtime.session.getRlmChildRunStatus(metadata.rlmChildId) - : undefined; - const status = runStatus ?? (session.isSessionActive ? "running" : "done"); - const isActive = status === "running" || session.isSessionActive; - return { - id: metadata.rlmChildId ?? activeSession.activeSessionId, - parentId: parentNodeId, - activeSessionId: activeSession.activeSessionId, - sessionName: session.sessionName, - model: session.model ? `${session.model.provider}/${session.model.id}` : undefined, - label: rlmChildLabel(metadata.prompt ?? ""), - status, - answerPreview, - toolUseCount: toolUseCount > 0 ? toolUseCount : undefined, - tokenCount: session._contextTokensForCurrentMessages(), - recap: session.getCurrentRecap(), - sessionDir: metadata.sessionDir ?? session.sessionManager.getSessionDir(), - activity: isActive ? { kind: session.isStreaming ? "writing" : "waiting" } : undefined, - }; + if (!root) return []; + const activeSessionIds = new Map( + activeSessions.flatMap((candidate) => { + const childId = candidate.runtime.metadata.rlmChildId; + return childId ? [[childId, candidate.activeSessionId] as const] : []; + }), + ); + return root.runtime.session.getRlmChildSnapshots().map((snapshot) => ({ + ...snapshot, + activeSessionId: activeSessionIds.get(snapshot.id), + })); } function firstUserMessageText(session: ActiveSessionState["runtime"]["session"]): string | undefined { diff --git a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts index 2896357a91..54bdbc62b2 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts @@ -138,7 +138,12 @@ import { import { MutationDrainLatch } from "./mutation-drain-latch.js"; import { createRlmLedgerRegistrySeedSource, RlmSpawnLedger } from "./rlm-ledger.js"; import { serializeSavedSessionInfo } from "./saved-session-info.js"; -import { SNAPSHOT_TARGET_CHUNK_BYTES, SnapshotTranscriptCache } from "./snapshot-transcript-cache.js"; +import { + createSnapshotCacheProcessRoot, + prepareSnapshotTranscriptCache, + SNAPSHOT_TARGET_CHUNK_BYTES, + SnapshotTranscriptCache, +} from "./snapshot-transcript-cache.js"; import { WorkerRecoveryJournal } from "./worker-recovery-journal.js"; type DistributiveOmit = T extends unknown ? Omit : never; @@ -147,6 +152,7 @@ type DaemonCommandBody = DistributiveOmit; const structuredLog = getLogger("coding-agent.daemon-supervisor"); const WORKER_CONNECT_TIMEOUT_MS = 30_000; const WORKER_REQUEST_TIMEOUT_MS = 24 * 60 * 60 * 1000; +const SNAPSHOT_ATTEMPT_TIMEOUT_MS = 30_000; const INPUT_PAUSE_CLEANUP_TIMEOUT_MS = 5_000; const UPDATE_RESTART_MUTATION_DRAIN_TIMEOUT_MS = 80_000; const UPDATE_RESTART_WORKER_REQUEST_TIMEOUT_MS = 90_000; @@ -292,6 +298,7 @@ interface ResidentWorker { transcriptCaches: Map; snapshotGenerations: Map>; snapshotLoads: Map>; + snapshotLoadSnapshotIds: Map; recovery?: Promise; deferredRecovery?: Promise; intentionalStop: boolean; @@ -313,13 +320,17 @@ interface SnapshotDuplicateValidation { interface SnapshotTranscriptGeneration { transcript: SnapshotTranscriptCache; result: DaemonAttachResult; + materializedFrom?: DaemonAttachResult; begin?: Buffer; end?: Buffer; incoming: boolean; retired: boolean; + receivedMessageCount?: number; duplicateChunkIndex?: number; duplicateResult?: DaemonAttachResult; validation?: SnapshotDuplicateValidation; + deadline?: ReturnType; + purpose?: "attach" | "replacement" | "catchup" | "resync"; } interface DaemonSupervisorOptions { @@ -672,7 +683,10 @@ export class DaemonSupervisor { private updateRestartPhase?: "draining" | "fencing" | "prepared"; private readonly mutationDrain = new MutationDrainLatch(); private readonly clients = new Set(); + private readonly deliveredSnapshotFailures = new WeakSet(); private readonly connectionIds = new WeakMap(); + private attachmentEpochs = new WeakMap>(); + private attachmentAbortControllers = new WeakMap>(); private readonly sessionInputPauseEpochs = new WeakMap(); private readonly detachingInputPauseSessions = new WeakMap>(); private readonly protocolClientIds = new WeakMap(); @@ -689,10 +703,15 @@ export class DaemonSupervisor { private readonly generation = randomUUID(); private readonly supervisorConfigPath: string; private readonly defaultSessionConfig: AgentSessionRuntimeConfig; - private readonly snapshotCacheRoot: string; + private readonly snapshotCacheParent: string; + private snapshotCacheRoot: string; private commandJournal!: CommandRecoveryJournal; private readonly streamReconstructor = new CompactAssistantStreamReconstructor(); private readonly compactCatchupInProgress = new Set(); + private transcriptPreparations?: WeakMap< + ResidentWorker, + Map }> + >; private readonly pendingSessionNames = new Set(); private readonly catalog: DaemonCatalogClient; private readonly settingsManager: SettingsManager; @@ -720,7 +739,8 @@ export class DaemonSupervisor { options.defaultSessionConfig, this.loadPersistedSupervisorConfig(), ); - this.snapshotCacheRoot = join(this.descriptorDir, "snapshot-cache", this.generation); + this.snapshotCacheParent = join(this.descriptorDir, "snapshot-cache"); + this.snapshotCacheRoot = join(this.snapshotCacheParent, `inactive-${this.generation}`); this.catalog = new DaemonCatalogClient((message) => this.log(message)); this.settingsManager = SettingsManager.create(process.cwd(), this.defaultSessionConfig.agentDir ?? agentDir); } @@ -745,8 +765,9 @@ export class DaemonSupervisor { mkdirSync(this.descriptorDir, { recursive: true, mode: 0o700 }); chmodSync(this.descriptorDir, 0o700); this.persistSupervisorConfig(); - rmSync(this.snapshotCacheRoot, { recursive: true, force: true }); - mkdirSync(this.snapshotCacheRoot, { recursive: true, mode: 0o700 }); + // Ownership is exclusive here, so legacy unowned cache roots cannot still have a live writer. + rmSync(this.snapshotCacheParent, { recursive: true, force: true }); + this.snapshotCacheRoot = createSnapshotCacheProcessRoot(this.snapshotCacheParent, this.generation); this.commandJournal = new CommandRecoveryJournal(join(this.descriptorDir, "command-journal.jsonl")); this.loadWorkerDescriptors(); const workersToAdopt = [...this.workers.values()]; @@ -1030,6 +1051,7 @@ export class DaemonSupervisor { transcriptCaches: new Map(), snapshotGenerations: new Map(), snapshotLoads: new Map(), + snapshotLoadSnapshotIds: new Map(), intentionalStop: durableDescriptor.stopRequestedAt !== undefined, stopRevision: 0, }; @@ -1158,7 +1180,10 @@ export class DaemonSupervisor { this.clients.delete(client); this.cancelWaitingPromptAdmissionsForClient(client); for (const activeSessionId of [...client.attachedActiveSessionIds]) { + this.advanceAttachmentEpoch(client, activeSessionId); client.attachedActiveSessionIds.delete(activeSessionId); + client.catchupActiveSessionIds?.delete(activeSessionId); + client.catchupPurposes?.delete(activeSessionId); void this.syncWorkerExtensionUi(activeSessionId); } this.scheduleOwnedWorkerCleanupForClient(this.protocolClientId(client)); @@ -1767,34 +1792,50 @@ export class DaemonSupervisor { return success(command.id, "create", this.publicSummary(worker, summary)); } case "attach": { - const attached = await this.attachClient(client, command); - if (client.capabilities.has("chunked_snapshot")) { - const transcript = attached.transcript; - if (!transcript) { - throw new Error("Session worker did not provide a snapshot transcript"); - } - const streamedResult = this.createStreamedAttachResult(attached.result, transcript); - try { - this.write(client, success(command.id, "attach", streamedResult)); - void this.streamSnapshot( - client, - attached.worker, - streamedResult, - transcript, - "attach", - attached.releaseTranscript, - ).catch((error) => - this.log( - `Failed to stream attach snapshot for ${streamedResult.activeSessionId}: ${String(error)}`, - ), - ); - } catch (error) { - attached.releaseTranscript?.(); - throw error; + const attachmentEpoch = this.advanceAttachmentEpoch(client, command.activeSessionId); + const requestedCapabilities = normalizeCapabilities(command.capabilities, command.supportsExtensionUi); + const releaseSnapshotReservation = requestedCapabilities.has("chunked_snapshot") + ? this.reserveSnapshotStream(client, command.activeSessionId) + : undefined; + try { + const attached = await this.attachClient(client, command, { + selector: command.activeSessionId, + epoch: attachmentEpoch, + }); + if (client.capabilities.has("chunked_snapshot")) { + const transcript = attached.transcript; + if (!transcript) { + throw new Error("Session worker did not provide a snapshot transcript"); + } + const streamedResult = this.createStreamedAttachResult(attached.result, transcript); + try { + this.write(client, success(command.id, "attach", streamedResult)); + void this.streamSnapshot( + client, + attached.worker, + streamedResult, + transcript, + "attach", + attached.releaseTranscript, + releaseSnapshotReservation, + { selector: command.activeSessionId, epoch: attachmentEpoch }, + ).catch((error) => + this.log( + `Failed to stream attach snapshot for ${streamedResult.activeSessionId}: ${String(error)}`, + ), + ); + } catch (error) { + attached.releaseTranscript?.(); + throw error; + } + return undefined; } - return undefined; + releaseSnapshotReservation?.(); + return success(command.id, "attach", attached.result); + } catch (error) { + releaseSnapshotReservation?.(); + throw error; } - return success(command.id, "attach", attached.result); } case "reattach": { const target = await this.findWorkerForClient(client, command.targetActiveSessionId); @@ -1807,22 +1848,26 @@ export class DaemonSupervisor { return success(command.id, command.type, { cancelled: false }); } const targetWasAttached = client.attachedActiveSessionIds.has(targetActiveSessionId); + const attachmentEpoch = this.advanceAttachmentEpoch(client, targetActiveSessionId); const releaseSnapshotReservation = this.reserveSnapshotStream(client, targetActiveSessionId); let releaseTranscript: (() => void) | undefined; - client.attachedActiveSessionIds.add(targetActiveSessionId); try { - const attached = await this.attachClient(client, { - ...command, - type: "attach", - activeSessionId: targetActiveSessionId, - }); + const attached = await this.attachClient( + client, + { + ...command, + type: "attach", + activeSessionId: targetActiveSessionId, + }, + { selector: targetActiveSessionId, epoch: attachmentEpoch }, + ); const detachingSessions = this.detachingInputPauseSessions?.get(client); detachingSessions?.delete(command.activeSessionId); detachingSessions?.delete(command.targetActiveSessionId); detachingSessions?.delete(targetActiveSessionId); if (client.capabilities.has("chunked_snapshot")) { const transcript = - attached.transcript ?? this.getOrCreateTranscriptCache(attached.worker, attached.result); + attached.transcript ?? (await this.getOrCreateTranscriptCache(attached.worker, attached.result)); releaseTranscript = attached.releaseTranscript; const streamedResult = this.createStreamedAttachResult(attached.result, transcript); this.write(client, success(command.id, command.type, streamedResult)); @@ -1835,6 +1880,7 @@ export class DaemonSupervisor { "replacement", releaseTranscript, releaseSnapshotReservation, + { selector: targetActiveSessionId, epoch: attachmentEpoch }, ); releaseTranscript = undefined; void streaming.catch((error) => @@ -2471,17 +2517,30 @@ export class DaemonSupervisor { ? resolve(command.sessionPath) : await this.catalog.resolve(command.sessionPath, config.cwd ?? process.cwd(), config.sessionDir); createCommand = { ...createCommand, sessionPath }; - const existing = this.findWorkerBySessionFile(sessionPath); - if (existing && !(await this.reclaimStaleWorkerRegistration(existing, command.launchEnv !== undefined))) { - return this.reuseWorkerForCreate(existing, ownerClientId, sessionPath); - } } const key = createCommand.sessionPath ? canonicalSessionPath(createCommand.sessionPath) : `new:${command.id ? createCommandIdempotencyKey(clientId, command.id) : createActiveSessionId()}`; const pending = this.openingWorkers.get(key); if (pending) { - return pending; + return this.joinOpeningWorker(pending, ownerClientId, createCommand.sessionPath ?? key); + } + if (createCommand.sessionPath) { + const existing = this.findWorkerBySessionFile(createCommand.sessionPath); + if (existing && !(await this.reclaimStaleWorkerRegistration(existing, command.launchEnv !== undefined))) { + return this.reuseWorkerForCreate(existing, ownerClientId, createCommand.sessionPath); + } + // The reclaim await may have let a concurrent opener register; join it instead of double-launching. + const opened = this.openingWorkers.get(key); + if (opened) { + return this.joinOpeningWorker(opened, ownerClientId, createCommand.sessionPath); + } + } + if (createCommand.sessionPath) { + const existing = this.findWorkerBySessionFile(createCommand.sessionPath); + if (existing && !(await this.reclaimStaleWorkerRegistration(existing, command.launchEnv !== undefined))) { + return this.reuseWorkerForCreate(existing, ownerClientId, createCommand.sessionPath); + } } const opening = (async () => { if (!createCommand.name) return this.launchWorker(createCommand, undefined, ownerClientId); @@ -2510,20 +2569,68 @@ export class DaemonSupervisor { } } - private reuseWorkerForCreate( + private async reuseWorkerForCreate( worker: ResidentWorker, ownerClientId: string | undefined, sessionPath: string, - ): ResidentWorker { + ): Promise { if (worker.descriptor.lifecycle === "failed") { throw new Error( `Session "${sessionPath}" is registered to a failed worker that could not be safely reclaimed`, ); } - if (worker.descriptor.ownerClientId === ownerClientId) { - return worker; + this.assertWorkerCreateOwner(worker, ownerClientId, sessionPath); + if (!this.isWorkerReadyForCreate(worker)) { + if (worker.recovery) { + await worker.recovery; + } else if (this.isWorkerRecoveryEligible(worker)) { + await this.recoverWorker(worker); + } + } + const current = this.workers.get(worker.descriptor.workerId); + if (!current) { + throw new Error(`Session "${sessionPath}" worker recovery was interrupted; retry opening the session`); } - throw new SessionAlreadyActiveError(sessionPath, worker.descriptor.rootActiveSessionId); + this.assertWorkerCreateOwner(current, ownerClientId, sessionPath); + if (!this.isWorkerReadyForCreate(current)) { + if (!current.summaries.has(current.descriptor.rootActiveSessionId)) { + throw new Error( + `Session "${sessionPath}" worker is unavailable for reuse: assigned root session is missing`, + ); + } + const detail = current.descriptor.lastError ? `: ${current.descriptor.lastError}` : ""; + throw new Error(`Session "${sessionPath}" worker is ${this.effectiveWorkerState(current)}${detail}`); + } + return current; + } + + private async joinOpeningWorker( + pending: Promise, + ownerClientId: string | undefined, + sessionPath: string, + ): Promise { + const worker = await pending; + this.assertWorkerCreateOwner(worker, ownerClientId, sessionPath); + return worker; + } + + private assertWorkerCreateOwner( + worker: ResidentWorker, + ownerClientId: string | undefined, + sessionPath: string, + ): void { + if (worker.descriptor.ownerClientId !== ownerClientId) { + throw new SessionAlreadyActiveError(sessionPath, worker.descriptor.rootActiveSessionId); + } + } + + private isWorkerReadyForCreate(worker: ResidentWorker): boolean { + return ( + worker.descriptor.lifecycle === "ready" && + worker.client !== undefined && + worker.summaries.has(worker.descriptor.rootActiveSessionId) && + !this.isWorkerStopping(worker) + ); } /** @@ -2715,6 +2822,7 @@ export class DaemonSupervisor { transcriptCaches: new Map(), snapshotGenerations: new Map(), snapshotLoads: new Map(), + snapshotLoadSnapshotIds: new Map(), intentionalStop: false, stopRevision: 0, launchEnv, @@ -2894,25 +3002,21 @@ export class DaemonSupervisor { const supportsExtensionUi = [...this.clients].some( (client) => client.attachedActiveSessionIds.has(activeSessionId) && client.supportsExtensionUi, ); + const immutableSnapshotTransfer = worker.client.supportsServerCapability?.("immutable_snapshot_transfer_v1"); + const capabilities: DaemonClientCapability[] = [ + "attach_snapshot", + "event_sequence", + "slim_attach", + "correlated_prompt_lifecycle_v1", + ]; + if (supportsExtensionUi) capabilities.push("extension_ui"); + if (immutableSnapshotTransfer) { + capabilities.push("chunked_snapshot", "immutable_snapshot_transfer_v1"); + } const response = await worker.client.requestWorker({ type: "worker_subscribe", activeSessionId, - capabilities: supportsExtensionUi - ? [ - "attach_snapshot", - "event_sequence", - "extension_ui", - "slim_attach", - "chunked_snapshot", - "correlated_prompt_lifecycle_v1", - ] - : [ - "attach_snapshot", - "event_sequence", - "slim_attach", - "chunked_snapshot", - "correlated_prompt_lifecycle_v1", - ], + capabilities, supportsExtensionUi, }); if (!response.success) { @@ -3087,6 +3191,18 @@ export class DaemonSupervisor { } } + private invalidateSnapshotLoads(worker: ResidentWorker, activeSessionId: string, expectedSnapshotId?: string): void { + worker.snapshotLoadSnapshotIds ??= new Map(); + for (const mode of ["chunked", "full"] as const) { + const key = `${activeSessionId}:${mode}`; + if (expectedSnapshotId !== undefined && worker.snapshotLoadSnapshotIds.get(key) !== expectedSnapshotId) { + continue; + } + worker.snapshotLoads.delete(key); + worker.snapshotLoadSnapshotIds.delete(key); + } + } + private failWorkerSnapshotCache( worker: ResidentWorker, activeSessionId: string, @@ -3101,14 +3217,13 @@ export class DaemonSupervisor { this.failSnapshotGeneration(worker, activeSessionId, generation, error); } else { const transcript = worker.transcriptCaches.get(activeSessionId); - if (transcript?.snapshotId !== expectedSnapshotId) { - return; - } - transcript.markFailed(error); - transcript.dispose(); - worker.transcriptCaches.delete(activeSessionId); - if (worker.snapshotCache.get(activeSessionId)?.snapshotStream?.id === expectedSnapshotId) { - worker.snapshotCache.delete(activeSessionId); + if (transcript?.snapshotId === expectedSnapshotId) { + transcript.markFailed(error); + transcript.dispose(); + worker.transcriptCaches.delete(activeSessionId); + if (worker.snapshotCache.get(activeSessionId)?.snapshotStream?.id === expectedSnapshotId) { + worker.snapshotCache.delete(activeSessionId); + } } } } else { @@ -3125,6 +3240,7 @@ export class DaemonSupervisor { worker.transcriptCaches.delete(activeSessionId); worker.snapshotCache.delete(activeSessionId); } + this.invalidateSnapshotLoads(worker, activeSessionId, expectedSnapshotId); if (closeWorkerChannel) { const client = worker.client; if (client) { @@ -3134,6 +3250,95 @@ export class DaemonSupervisor { } } + private failWorkerSnapshotGeneration( + worker: ResidentWorker, + activeSessionId: string, + snapshotId: string, + error: Error, + snapshotPurpose?: "attach" | "replacement" | "catchup" | "resync", + notifyClients = true, + ): void { + const generation = this.snapshotGeneration(worker, activeSessionId, snapshotId); + if (!generation) return; + const result = generation.result; + const published = worker.transcriptCaches.get(activeSessionId) === generation.transcript; + this.failWorkerSnapshotCache(worker, activeSessionId, error, false, snapshotId); + if ( + notifyClients && + published && + (snapshotPurpose === "replacement" || snapshotPurpose === "catchup" || snapshotPurpose === "resync") + ) { + for (const client of this.clients) { + if (!client.attachedActiveSessionIds.has(activeSessionId) || !client.capabilities.has("chunked_snapshot")) { + continue; + } + void this.deliverWorkerSnapshotFailure(client, result, snapshotId, snapshotPurpose, error); + } + } + } + + private async deliverWorkerSnapshotFailure( + client: DaemonSocketClient, + result: DaemonAttachResult, + snapshotId: string, + purpose: "replacement" | "catchup" | "resync", + error: Error, + ): Promise { + if (client.socket.destroyed || !client.attachedActiveSessionIds.has(result.activeSessionId)) return; + const { messages: _messages, ...snapshot } = result.snapshot; + await this.writeSnapshotRecord(client, { + type: "session_snapshot_begin", + activeSessionId: result.activeSessionId, + snapshotId, + snapshot, + messageCount: result.snapshotStream?.messageCount ?? result.snapshot.summary.messageCount, + targetChunkBytes: result.snapshotStream?.targetChunkBytes ?? SNAPSHOT_TARGET_CHUNK_BYTES, + purpose: purpose === "replacement" ? "replacement" : "resync", + }); + await this.writeSnapshotRecord(client, { + type: "session_snapshot_failed", + activeSessionId: result.activeSessionId, + snapshotId, + error: error.message, + }); + } + + private failCorrelatedSnapshotAttempt( + worker: ResidentWorker, + activeSessionId: string, + snapshotId: string, + error: Error, + snapshotPurpose?: "attach" | "replacement" | "catchup" | "resync", + notifyClients = true, + ): void { + if (this.snapshotGeneration(worker, activeSessionId, snapshotId)) { + this.failWorkerSnapshotGeneration(worker, activeSessionId, snapshotId, error, snapshotPurpose, notifyClients); + return; + } + const transcript = worker.transcriptCaches.get(activeSessionId); + const result = worker.snapshotCache.get(activeSessionId); + if (transcript?.snapshotId === snapshotId || result?.snapshotStream?.id === snapshotId) { + this.failWorkerSnapshotCache(worker, activeSessionId, error, false, snapshotId); + if ( + notifyClients && + result && + (snapshotPurpose === "replacement" || snapshotPurpose === "catchup" || snapshotPurpose === "resync") + ) { + for (const client of this.clients) { + if ( + !client.attachedActiveSessionIds.has(activeSessionId) || + !client.capabilities.has("chunked_snapshot") + ) { + continue; + } + void this.deliverWorkerSnapshotFailure(client, result, snapshotId, snapshotPurpose, error); + } + } + return; + } + this.log(`Ignored stale snapshot ${snapshotId} failure for ${activeSessionId}: ${error.message}`); + } + private retireWorkerSnapshotCache( worker: ResidentWorker, activeSessionId: string, @@ -3151,6 +3356,7 @@ export class DaemonSupervisor { return; } generation.retired = true; + this.clearSnapshotGenerationDeadline(generation); this.settleSnapshotDuplicateValidation(generation); if (generation.incoming) { return; @@ -3226,12 +3432,38 @@ export class DaemonSupervisor { } } + private armSnapshotGenerationDeadline( + worker: ResidentWorker, + activeSessionId: string, + generation: SnapshotTranscriptGeneration, + ): void { + if (generation.deadline) clearTimeout(generation.deadline); + generation.deadline = setTimeout(() => { + if (!generation.incoming && generation.transcript.complete) return; + this.failWorkerSnapshotGeneration( + worker, + activeSessionId, + generation.transcript.snapshotId, + new Error(`Snapshot ${generation.transcript.snapshotId} timed out before completion`), + generation.purpose, + ); + }, SNAPSHOT_ATTEMPT_TIMEOUT_MS); + generation.deadline.unref(); + } + + private clearSnapshotGenerationDeadline(generation: SnapshotTranscriptGeneration): void { + if (!generation.deadline) return; + clearTimeout(generation.deadline); + generation.deadline = undefined; + } + private failSnapshotGeneration( worker: ResidentWorker, activeSessionId: string, generation: SnapshotTranscriptGeneration, error: Error, ): void { + this.clearSnapshotGenerationDeadline(generation); this.settleSnapshotDuplicateValidation(generation, error); generation.transcript.markFailed(error); generation.transcript.dispose(); @@ -3516,7 +3748,12 @@ export class DaemonSupervisor { } const response = await worker.client.request({ type: "list" }, 5000); const summaries = sessionSummariesFromResponse(response); - worker.summaries = new Map(summaries.map((summary) => [summary.activeSessionId ?? summary.id, summary])); + const nextSummaries = new Map(summaries.map((summary) => [summary.activeSessionId ?? summary.id, summary])); + const root = nextSummaries.get(worker.descriptor.rootActiveSessionId); + if (recovery && !root) { + throw new Error(`Session worker omitted its root session during recovery`); + } + worker.summaries = nextSummaries; for (const summary of summaries) { const activeSessionId = summary.activeSessionId ?? summary.id; if (summary.streamingMessage?.role === "assistant") { @@ -3525,7 +3762,6 @@ export class DaemonSupervisor { this.streamReconstructor.clear(activeSessionId); } } - const root = worker.summaries.get(worker.descriptor.rootActiveSessionId); if (root) { if (recovery) { await this.assertRecoveryAllowed(); @@ -3901,6 +4137,7 @@ export class DaemonSupervisor { private async attachClient( client: DaemonSocketClient, command: Extract, + attachmentFence?: { selector: string; epoch: number; requireAttached?: boolean }, ): Promise { const ownedWorker = [...this.workers.values()].find( (worker) => @@ -3943,9 +4180,22 @@ export class DaemonSupervisor { this.assertTelemetryAttachAllowed(match.worker, command.telemetryDisabled); this.requireAvailableWorkerClient(match.worker); const activeSessionId = match.summary.activeSessionId ?? match.summary.id; + const resolvedAttachmentEpoch = this.attachmentEpoch(client, activeSessionId); + const assertAttachmentCurrent = () => { + if ( + attachmentFence && + (this.attachmentEpoch(client, attachmentFence.selector) !== attachmentFence.epoch || + this.attachmentEpoch(client, activeSessionId) !== resolvedAttachmentEpoch || + (attachmentFence.requireAttached && !client.attachedActiveSessionIds.has(activeSessionId))) + ) { + throw new Error(`Attachment changed while preparing snapshot for ${activeSessionId}`); + } + }; + assertAttachmentCurrent(); const duplicateValidation = this.currentSnapshotGeneration(match.worker, activeSessionId)?.validation; if (duplicateValidation) { await duplicateValidation.promise; + assertAttachmentCurrent(); } if (command.clientId) { client.id = command.clientId; @@ -3965,6 +4215,7 @@ export class DaemonSupervisor { const snapshotLoadKey = `${activeSessionId}:${client.capabilities.has("chunked_snapshot") ? "chunked" : "full"}`; let retryInvalidatedLoad = true; while (!result) { + match.worker.snapshotLoadSnapshotIds ??= new Map(); let loading = match.worker.snapshotLoads.get(snapshotLoadKey); if (!loading) { const observedSnapshotId = @@ -3972,21 +4223,28 @@ export class DaemonSupervisor { match.worker.snapshotCache.get(activeSessionId)?.snapshotStream?.id; loading = (async () => { const workerClient = this.requireAvailableWorkerClient(match.worker); - const response = await workerClient.request({ - type: "attach", - activeSessionId, - capabilities: client.capabilities.has("chunked_snapshot") - ? [ - "attach_snapshot", - "event_sequence", - "slim_attach", - "chunked_snapshot", - "correlated_prompt_lifecycle_v1", - ] - : ["attach_snapshot", "event_sequence", "slim_attach", "correlated_prompt_lifecycle_v1"], - supportsExtensionUi: false, - env: command.env ?? collectDaemonClientEnv(), - }); + const workerCapabilities: DaemonClientCapability[] = [ + "attach_snapshot", + "event_sequence", + "slim_attach", + "correlated_prompt_lifecycle_v1", + ]; + if ( + client.capabilities.has("chunked_snapshot") && + workerClient.supportsServerCapability?.("immutable_snapshot_transfer_v1") + ) { + workerCapabilities.push("chunked_snapshot", "immutable_snapshot_transfer_v1"); + } + const response = await workerClient.request( + { + type: "attach", + activeSessionId, + capabilities: workerCapabilities, + supportsExtensionUi: false, + env: command.env ?? collectDaemonClientEnv(), + }, + SNAPSHOT_ATTEMPT_TIMEOUT_MS, + ); const loaded = attachResultFromResponse(response); if (match.worker.snapshotLoads.get(snapshotLoadKey) !== loading) { throw new SnapshotLoadInvalidatedError("Session snapshot changed during attach"); @@ -3994,6 +4252,7 @@ export class DaemonSupervisor { return this.cacheLoadedSnapshot(match.worker, activeSessionId, loaded, observedSnapshotId); })(); match.worker.snapshotLoads.set(snapshotLoadKey, loading); + match.worker.snapshotLoadSnapshotIds.set(snapshotLoadKey, observedSnapshotId); void loading.then( async (loaded) => { try { @@ -4012,18 +4271,21 @@ export class DaemonSupervisor { } finally { if (match.worker.snapshotLoads.get(snapshotLoadKey) === loading) { match.worker.snapshotLoads.delete(snapshotLoadKey); + match.worker.snapshotLoadSnapshotIds.delete(snapshotLoadKey); } } }, () => { if (match.worker.snapshotLoads.get(snapshotLoadKey) === loading) { match.worker.snapshotLoads.delete(snapshotLoadKey); + match.worker.snapshotLoadSnapshotIds.delete(snapshotLoadKey); } }, ); } try { result = await loading; + assertAttachmentCurrent(); } catch (error) { if (!(error instanceof SnapshotLoadInvalidatedError)) { throw error; @@ -4046,11 +4308,19 @@ export class DaemonSupervisor { continue; } result = match.worker.snapshotCache.get(activeSessionId) ?? result; - transcript = this.getOrCreateTranscriptCache(match.worker, result); + if (client.socket.destroyed) { + throw new Error("Daemon client disconnected during snapshot preparation"); + } + transcript = await this.getOrCreateTranscriptCache(match.worker, result); + assertAttachmentCurrent(); + if (client.socket.destroyed) { + throw new Error("Daemon client disconnected during snapshot preparation"); + } break; } } const releaseTranscript = transcript?.retain(); + assertAttachmentCurrent(); client.attachedActiveSessionIds.add(activeSessionId); try { const publicSummary = this.publicSummary(match.worker, result.snapshot.summary); @@ -4165,48 +4435,79 @@ export class DaemonSupervisor { } worker.transcriptCaches.set(activeSessionId, transcript); worker.snapshotCache.set(activeSessionId, loaded); + if (!transcript.complete) this.armSnapshotGenerationDeadline(worker, activeSessionId, generation); return loaded; } - private getOrCreateTranscriptCache(worker: ResidentWorker, result: DaemonAttachResult): SnapshotTranscriptCache { + private async getOrCreateTranscriptCache( + worker: ResidentWorker, + result: DaemonAttachResult, + ): Promise { const activeSessionId = result.activeSessionId; - const existing = worker.transcriptCaches.get(activeSessionId); - if (existing && (!result.snapshotStream || existing.snapshotId === result.snapshotStream.id)) { - return existing; - } - if (result.snapshot.messages.length < result.snapshot.summary.messageCount) { - throw new Error("Session snapshot generation changed before its transcript could be selected"); - } - if (existing) { - this.retireWorkerSnapshotCache(worker, activeSessionId, existing); + while (true) { + const existing = worker.transcriptCaches.get(activeSessionId); + const existingGeneration = existing + ? this.snapshotGeneration(worker, activeSessionId, existing.snapshotId) + : undefined; + if ( + existing && + (result.snapshotStream + ? existing.snapshotId === result.snapshotStream.id + : existingGeneration?.result === result || existingGeneration?.materializedFrom === result) + ) { + return existing; + } + if (result.snapshot.messages.length < result.snapshot.summary.messageCount) { + throw new Error("Session snapshot generation changed before its transcript could be selected"); + } + this.transcriptPreparations ??= new WeakMap(); + let preparations = this.transcriptPreparations.get(worker); + if (!preparations) { + preparations = new Map(); + this.transcriptPreparations.set(worker, preparations); + } + const pending = preparations.get(activeSessionId); + if (pending) { + await pending.promise; + continue; + } + if (existing) { + this.retireWorkerSnapshotCache(worker, activeSessionId, existing); + } + const snapshotId = `snapshot-${randomUUID()}`; + const promise = prepareSnapshotTranscriptCache({ + activeSessionId, + snapshotId, + messages: result.snapshot.messages, + cacheRoot: this.snapshotCacheRoot, + targetChunkBytes: SNAPSHOT_TARGET_CHUNK_BYTES, + signal: AbortSignal.timeout(SNAPSHOT_ATTEMPT_TIMEOUT_MS), + }); + preparations.set(activeSessionId, { result, promise }); + let transcript: SnapshotTranscriptCache; + try { + transcript = await promise; + } finally { + if (preparations.get(activeSessionId)?.promise === promise) { + preparations.delete(activeSessionId); + } + } + const cachedResult = { + ...result, + messages: result.messages ? [] : undefined, + snapshot: { ...result.snapshot, messages: [] }, + }; + worker.transcriptCaches.set(activeSessionId, transcript); + worker.snapshotCache.set(activeSessionId, cachedResult); + this.snapshotGenerationsFor(worker, activeSessionId).set(transcript.snapshotId, { + transcript, + result: cachedResult, + materializedFrom: result, + incoming: false, + retired: false, + }); + return transcript; } - const revision = createHash("sha256") - .update( - `${activeSessionId}:${result.snapshot.summary.sessionId}:${result.lastEventSequence}:${result.snapshot.messages.length}`, - ) - .digest("hex") - .slice(0, 16); - const transcript = new SnapshotTranscriptCache({ - activeSessionId, - snapshotId: `${activeSessionId}-${revision}`, - messages: result.snapshot.messages, - cacheRoot: this.snapshotCacheRoot, - targetChunkBytes: SNAPSHOT_TARGET_CHUNK_BYTES, - }); - worker.transcriptCaches.set(activeSessionId, transcript); - const cachedResult = { - ...result, - messages: result.messages ? [] : undefined, - snapshot: { ...result.snapshot, messages: [] }, - }; - worker.snapshotCache.set(activeSessionId, cachedResult); - this.snapshotGenerationsFor(worker, activeSessionId).set(transcript.snapshotId, { - transcript, - result: cachedResult, - incoming: false, - retired: false, - }); - return transcript; } private createStreamedAttachResult( @@ -4233,10 +4534,21 @@ export class DaemonSupervisor { purpose: "attach" | "replacement" | "resync" = "attach", retainedTranscriptRelease?: () => void, releaseSnapshotReservation = this.reserveSnapshotStream(client, result.activeSessionId), + attachmentFence?: { selector: string; epoch: number }, ): Promise { const stream = result.snapshotStream; const releaseTranscript = retainedTranscriptRelease ?? transcript.retain(); - if (!stream || client.socket.destroyed) { + const attachmentEpoch = this.attachmentEpoch(client, result.activeSessionId); + const isAttachmentCurrent = () => + this.attachmentEpoch(client, result.activeSessionId) === attachmentEpoch && + (!attachmentFence || this.attachmentEpoch(client, attachmentFence.selector) === attachmentFence.epoch) && + client.attachedActiveSessionIds.has(result.activeSessionId); + const attachmentSignal = this.attachmentAbortSignal( + client, + attachmentFence?.selector ?? result.activeSessionId, + attachmentFence?.epoch ?? attachmentEpoch, + ); + if (!stream || client.socket.destroyed || !isAttachmentCurrent()) { releaseSnapshotReservation(); releaseTranscript(); return; @@ -4258,14 +4570,36 @@ export class DaemonSupervisor { } let chunkCount = 0; while (true) { + if (!isAttachmentCurrent() || attachmentSignal?.aborted) return; let chunk: Buffer | undefined; try { - chunk = await transcript.waitForChunk(chunkCount); + if (!attachmentSignal) { + chunk = await transcript.waitForChunk(chunkCount); + } else { + chunk = await new Promise((resolveChunk, rejectChunk) => { + const onAbort = () => resolveChunk(undefined); + attachmentSignal.addEventListener("abort", onAbort, { once: true }); + void transcript + .waitForChunk(chunkCount) + .then(resolveChunk, rejectChunk) + .finally(() => { + attachmentSignal.removeEventListener("abort", onAbort); + }); + }); + } } catch (error) { const streamError = error instanceof Error ? error : new Error(String(error)); - this.failWorkerSnapshotCache(worker, result.activeSessionId, streamError, false, stream.id); + this.failCorrelatedSnapshotAttempt( + worker, + result.activeSessionId, + stream.id, + streamError, + purpose, + false, + ); throw streamError; } + if (!isAttachmentCurrent() || attachmentSignal?.aborted) return; if (!chunk) { break; } @@ -4274,7 +4608,7 @@ export class DaemonSupervisor { } chunkCount++; } - await this.writeSnapshotRecord(client, { + const completed = await this.writeSnapshotRecord(client, { type: "session_snapshot_end", activeSessionId: result.activeSessionId, snapshotId: stream.id, @@ -4282,9 +4616,10 @@ export class DaemonSupervisor { lastEventSequence: result.lastEventSequence, lastEventCursor: result.lastEventCursor, }); + if (!completed) return; } catch (error) { const streamError = error instanceof Error ? error : new Error(String(error)); - if (!client.socket.destroyed) { + if (!client.socket.destroyed && isAttachmentCurrent()) { try { const delivered = await this.writeSnapshotRecord(client, { type: "session_snapshot_failed", @@ -4294,6 +4629,8 @@ export class DaemonSupervisor { }); if (!delivered && !client.socket.destroyed) { client.socket.destroy(streamError); + } else if (delivered) { + this.deliveredSnapshotFailures.add(streamError); } } catch (deliveryError) { client.socket.destroy(deliveryError instanceof Error ? deliveryError : new Error(String(deliveryError))); @@ -4371,16 +4708,51 @@ export class DaemonSupervisor { }); } + private attachmentEpoch(client: DaemonSocketClient, activeSessionId: string): number { + return this.attachmentEpochs?.get(client)?.get(activeSessionId) ?? 0; + } + + private advanceAttachmentEpoch(client: DaemonSocketClient, activeSessionId: string): number { + this.attachmentEpochs ??= new WeakMap(); + let epochs = this.attachmentEpochs.get(client); + if (!epochs) { + epochs = new Map(); + this.attachmentEpochs.set(client, epochs); + } + const next = (epochs.get(activeSessionId) ?? 0) + 1; + epochs.set(activeSessionId, next); + this.attachmentAbortControllers ??= new WeakMap(); + let controllers = this.attachmentAbortControllers.get(client); + if (!controllers) { + controllers = new Map(); + this.attachmentAbortControllers.set(client, controllers); + } + controllers.get(activeSessionId)?.abort(); + controllers.set(activeSessionId, new AbortController()); + return next; + } + + private attachmentAbortSignal( + client: DaemonSocketClient, + activeSessionId: string, + epoch: number, + ): AbortSignal | undefined { + if (this.attachmentEpoch(client, activeSessionId) !== epoch) return AbortSignal.abort(); + return this.attachmentAbortControllers?.get(client)?.get(activeSessionId)?.signal; + } + private detachClient(client: DaemonSocketClient, activeSessionId?: string): void { const targets = activeSessionId ? [activeSessionId] : [...client.attachedActiveSessionIds]; for (const selector of targets) { const match = this.matchWorkers(selector)[0]; const resolvedId = match ? (match.summary.activeSessionId ?? match.summary.id) : selector; - if (!client.attachedActiveSessionIds.delete(resolvedId)) { - continue; - } + this.advanceAttachmentEpoch(client, selector); + if (resolvedId !== selector) this.advanceAttachmentEpoch(client, resolvedId); + client.catchupActiveSessionIds?.delete(selector); client.catchupActiveSessionIds?.delete(resolvedId); + client.catchupPurposes?.delete(selector); client.catchupPurposes?.delete(resolvedId); + if (!client.attachedActiveSessionIds.delete(resolvedId)) continue; this.write(client, { type: "session_detached", activeSessionId: resolvedId }); void this.syncWorkerExtensionUi(resolvedId); } @@ -4413,20 +4785,55 @@ export class DaemonSupervisor { this.broadcastHeartbeatsChanged(); return; } + if ( + (outboundType === "session_snapshot_begin" || + outboundType === "session_snapshot_chunk" || + outboundType === "session_snapshot_end" || + outboundType === "session_snapshot_failed") && + !activeSessionId + ) { + const error = new Error("Worker returned an uncorrelatable snapshot frame"); + const workerClient = worker.client; + if (workerClient) { + this.handleWorkerClose(worker, workerClient, error); + workerClient.close(); + } + return; + } if (outboundType === "session_snapshot_begin" && activeSessionId) { + let correlatedSnapshotId = frameSnapshotId; try { const begin = JSON.parse(frame.payload.toString("utf8")) as Extract< DaemonOutbound, { type: "session_snapshot_begin" } >; + if (typeof begin.snapshotId === "string") { + correlatedSnapshotId ??= begin.snapshotId; + } if ( begin.type !== "session_snapshot_begin" || begin.activeSessionId !== activeSessionId || typeof begin.snapshotId !== "string" || (frameSnapshotId !== undefined && frameSnapshotId !== begin.snapshotId) || - typeof begin.targetChunkBytes !== "number" || + !Number.isSafeInteger(begin.messageCount) || + begin.messageCount < 0 || + !Number.isSafeInteger(begin.targetChunkBytes) || + begin.targetChunkBytes <= 0 || !begin.snapshot || - !isSessionSummary(begin.snapshot.summary) + begin.snapshot.activeSessionId !== activeSessionId || + (begin.snapshot.state.activeSessionId !== undefined && + begin.snapshot.state.activeSessionId !== activeSessionId) || + !isSessionSummary(begin.snapshot.summary) || + (begin.snapshot.summary.activeSessionId !== undefined && + begin.snapshot.summary.activeSessionId !== activeSessionId) || + begin.snapshot.summary.sessionId !== begin.snapshot.state.sessionId || + begin.snapshot.summary.messageCount !== begin.messageCount || + !Number.isSafeInteger(begin.snapshot.lastEventSequence) || + begin.snapshot.lastEventSequence < 0 || + (begin.snapshot.lastEventCursor !== undefined && + (begin.snapshot.lastEventSequence !== begin.snapshot.lastEventCursor.sequence || + typeof begin.snapshot.lastEventCursor.generation !== "string" || + begin.snapshot.lastEventCursor.generation.length === 0)) ) { throw new Error("Worker returned an invalid snapshot begin frame"); } @@ -4460,12 +4867,12 @@ export class DaemonSupervisor { const generations = this.snapshotGenerationsFor(worker, activeSessionId); let generation = generations.get(begin.snapshotId); if (generation?.incoming) { - this.failWorkerSnapshotCache( + this.failWorkerSnapshotGeneration( worker, activeSessionId, - new Error(`Snapshot ${begin.snapshotId} restarted before completion`), - true, begin.snapshotId, + new Error(`Snapshot ${begin.snapshotId} restarted before completion`), + snapshotPurpose, ); return; } @@ -4482,12 +4889,12 @@ export class DaemonSupervisor { generation.result.snapshot.lastEventCursor?.generation === result.snapshot.lastEventCursor?.generation && generation.result.snapshot.lastEventCursor?.sequence === result.snapshot.lastEventCursor?.sequence; if (generation?.transcript.complete && !duplicate) { - this.failWorkerSnapshotCache( + this.failWorkerSnapshotGeneration( worker, activeSessionId, - new Error(`Snapshot ${begin.snapshotId} did not match the cached transfer`), - true, begin.snapshotId, + new Error(`Snapshot ${begin.snapshotId} did not match the cached transfer`), + snapshotPurpose, ); return; } @@ -4506,6 +4913,8 @@ export class DaemonSupervisor { generation.duplicateChunkIndex = 0; generation.duplicateResult = result; generation.validation = this.createSnapshotDuplicateValidation(); + generation.purpose = snapshotPurpose; + this.armSnapshotGenerationDeadline(worker, activeSessionId, generation); if (currentGeneration === generation) { worker.snapshotCache.delete(activeSessionId); } @@ -4547,9 +4956,12 @@ export class DaemonSupervisor { generation.begin = Buffer.from(frame.payload); generation.end = undefined; generation.incoming = true; + generation.receivedMessageCount = 0; generation.duplicateChunkIndex = undefined; generation.duplicateResult = undefined; generation.validation = undefined; + generation.purpose = snapshotPurpose; + this.armSnapshotGenerationDeadline(worker, activeSessionId, generation); if (!isOlderThanCurrent) { generation.retired = false; worker.transcriptCaches.set(activeSessionId, generation.transcript); @@ -4557,59 +4969,118 @@ export class DaemonSupervisor { } } catch (error) { this.log(`Invalid worker snapshot begin frame: ${String(error)}`); - this.failWorkerSnapshotCache( - worker, - activeSessionId, - error instanceof Error ? error : new Error(String(error)), - true, - ); + const frameError = error instanceof Error ? error : new Error(String(error)); + if (correlatedSnapshotId) { + this.failCorrelatedSnapshotAttempt( + worker, + activeSessionId, + correlatedSnapshotId, + frameError, + snapshotPurpose, + ); + } else { + this.failWorkerSnapshotCache(worker, activeSessionId, frameError, true); + } } return; } if (outboundType === "session_snapshot_chunk" && activeSessionId) { - const snapshotId = frameSnapshotId ?? worker.transcriptCaches.get(activeSessionId)?.snapshotId; + let chunk: Extract; + try { + chunk = JSON.parse(frame.payload.toString("utf8")) as typeof chunk; + } catch (error) { + const frameError = error instanceof Error ? error : new Error(String(error)); + if (frameSnapshotId) { + this.failCorrelatedSnapshotAttempt( + worker, + activeSessionId, + frameSnapshotId, + frameError, + snapshotPurpose, + ); + } else { + this.failWorkerSnapshotCache(worker, activeSessionId, frameError, true); + } + return; + } + const snapshotId = frameSnapshotId ?? (typeof chunk.snapshotId === "string" ? chunk.snapshotId : undefined); if (!snapshotId) { + this.failWorkerSnapshotCache( + worker, + activeSessionId, + new Error("Worker returned an uncorrelatable snapshot chunk frame"), + true, + ); return; } const generation = this.snapshotGeneration(worker, activeSessionId, snapshotId); - if (generation?.incoming) { - try { - const duplicateIndex = generation.duplicateChunkIndex; - if (duplicateIndex === undefined) { - generation.transcript.appendEncodedChunk(Buffer.from(frame.payload)); - } else { - const chunk = JSON.parse(frame.payload.toString("utf8")) as Extract< - DaemonOutbound, - { type: "session_snapshot_chunk" } - >; - if ( - chunk.type !== "session_snapshot_chunk" || - chunk.activeSessionId !== activeSessionId || - chunk.snapshotId !== generation.transcript.snapshotId || - chunk.index !== duplicateIndex || - !generation.transcript.readChunk(duplicateIndex).equals(Buffer.from(frame.payload)) - ) { - throw new Error( - `Duplicate snapshot ${generation.transcript.snapshotId} did not match cached bytes`, - ); - } - generation.duplicateChunkIndex = duplicateIndex + 1; + if (!generation?.incoming) { + return; + } + try { + const duplicateIndex = generation.duplicateChunkIndex; + const expectedIndex = duplicateIndex ?? generation.transcript.chunkCount; + if ( + chunk.type !== "session_snapshot_chunk" || + chunk.activeSessionId !== activeSessionId || + chunk.snapshotId !== generation.transcript.snapshotId || + chunk.index !== expectedIndex || + !Array.isArray(chunk.messages) + ) { + throw new Error(`Snapshot ${generation.transcript.snapshotId} returned an invalid chunk frame`); + } + if (duplicateIndex === undefined) { + const receivedMessageCount = (generation.receivedMessageCount ?? 0) + chunk.messages.length; + if (receivedMessageCount > (generation.result.snapshotStream?.messageCount ?? Number.MAX_SAFE_INTEGER)) { + throw new Error(`Snapshot ${generation.transcript.snapshotId} exceeded its declared message count`); } - } catch (error) { - this.failWorkerSnapshotCache( - worker, - activeSessionId, - error instanceof Error ? error : new Error(String(error)), - true, - generation.transcript.snapshotId, - ); + generation.transcript.appendEncodedChunk(Buffer.from(frame.payload)); + generation.receivedMessageCount = receivedMessageCount; + } else { + if (!generation.transcript.readChunk(duplicateIndex).equals(Buffer.from(frame.payload))) { + throw new Error(`Duplicate snapshot ${generation.transcript.snapshotId} did not match cached bytes`); + } + generation.duplicateChunkIndex = duplicateIndex + 1; } + } catch (error) { + this.failWorkerSnapshotGeneration( + worker, + activeSessionId, + generation.transcript.snapshotId, + error instanceof Error ? error : new Error(String(error)), + snapshotPurpose, + ); } return; } + if (outboundType === "session_snapshot_end" && activeSessionId) { - const snapshotId = frameSnapshotId ?? worker.transcriptCaches.get(activeSessionId)?.snapshotId; + let end: Extract; + try { + end = JSON.parse(frame.payload.toString("utf8")) as typeof end; + } catch (error) { + const frameError = error instanceof Error ? error : new Error(String(error)); + if (frameSnapshotId) { + this.failCorrelatedSnapshotAttempt( + worker, + activeSessionId, + frameSnapshotId, + frameError, + snapshotPurpose, + ); + } else { + this.failWorkerSnapshotCache(worker, activeSessionId, frameError, true); + } + return; + } + const snapshotId = frameSnapshotId ?? (typeof end.snapshotId === "string" ? end.snapshotId : undefined); if (!snapshotId) { + this.failWorkerSnapshotCache( + worker, + activeSessionId, + new Error("Worker returned an uncorrelatable snapshot end frame"), + true, + ); return; } const generation = this.snapshotGeneration(worker, activeSessionId, snapshotId); @@ -4619,25 +5090,35 @@ export class DaemonSupervisor { const transcript = generation.transcript; try { const duplicateChunkCount = generation.duplicateChunkIndex; + const expectedCursor = generation.result.lastEventCursor; + const endCursor = end.lastEventCursor; + if ( + end.type !== "session_snapshot_end" || + end.activeSessionId !== activeSessionId || + end.snapshotId !== transcript.snapshotId || + !Number.isSafeInteger(end.chunkCount) || + end.chunkCount < 0 || + end.chunkCount !== transcript.chunkCount || + end.lastEventSequence !== generation.result.lastEventSequence || + endCursor?.generation !== expectedCursor?.generation || + endCursor?.sequence !== expectedCursor?.sequence + ) { + throw new Error(`Snapshot ${transcript.snapshotId} ended with invalid metadata`); + } if (duplicateChunkCount === undefined) { - transcript.markComplete(); if (!generation.begin) { throw new Error(`Snapshot ${transcript.snapshotId} has no begin frame`); } - generation.end = Buffer.from(frame.payload); - } else { - const end = JSON.parse(frame.payload.toString("utf8")) as Extract< - DaemonOutbound, - { type: "session_snapshot_end" } - >; if ( - end.type !== "session_snapshot_end" || - end.activeSessionId !== activeSessionId || - end.snapshotId !== transcript.snapshotId || - end.chunkCount !== duplicateChunkCount || - end.chunkCount !== transcript.chunkCount || - !generation.end?.equals(frame.payload) + (generation.receivedMessageCount ?? 0) !== + (generation.result.snapshotStream?.messageCount ?? generation.receivedMessageCount ?? 0) ) { + throw new Error(`Snapshot ${transcript.snapshotId} ended with the wrong message count`); + } + generation.end = Buffer.from(frame.payload); + transcript.markComplete(); + } else { + if (end.chunkCount !== duplicateChunkCount || !generation.end?.equals(frame.payload)) { throw new Error(`Duplicate snapshot ${transcript.snapshotId} ended with different metadata`); } if (!generation.duplicateResult) { @@ -4649,16 +5130,18 @@ export class DaemonSupervisor { } this.settleSnapshotDuplicateValidation(generation); } + this.clearSnapshotGenerationDeadline(generation); generation.incoming = false; + generation.receivedMessageCount = undefined; generation.duplicateChunkIndex = undefined; generation.duplicateResult = undefined; } catch (error) { - this.failWorkerSnapshotCache( + this.failWorkerSnapshotGeneration( worker, activeSessionId, - error instanceof Error ? error : new Error(String(error)), - true, transcript.snapshotId, + error instanceof Error ? error : new Error(String(error)), + snapshotPurpose, ); return; } @@ -4671,19 +5154,26 @@ export class DaemonSupervisor { for (const client of this.clients) { if (!client.attachedActiveSessionIds.has(activeSessionId)) continue; this.queueCatchup(client, activeSessionId, snapshotPurpose === "replacement" ? "replacement" : "resync"); - void this.catchUpClient(client).catch((error) => - this.log(`Failed to catch up client ${client.id}: ${String(error)}`), - ); + if (!client.catchupPromise) { + void this.catchUpClient(client).catch((error) => + this.log(`Failed to catch up client ${client.id}: ${String(error)}`), + ); + } } } return; } + if (outboundType === "session_snapshot_failed" && activeSessionId) { + let correlatedSnapshotId = frameSnapshotId; try { const failed = JSON.parse(frame.payload.toString("utf8")) as Extract< DaemonOutbound, { type: "session_snapshot_failed" } >; + if (typeof failed.snapshotId === "string") { + correlatedSnapshotId ??= failed.snapshotId; + } if ( failed.type !== "session_snapshot_failed" || failed.activeSessionId !== activeSessionId || @@ -4693,38 +5183,30 @@ export class DaemonSupervisor { ) { throw new Error("Worker returned an invalid snapshot failure frame"); } - const currentGeneration = this.currentSnapshotGeneration(worker, activeSessionId); - const generation = - this.snapshotGeneration(worker, activeSessionId, failed.snapshotId) ?? - (currentGeneration?.transcript.snapshotId === failed.snapshotId ? currentGeneration : undefined); - if (!generation) { - return; - } - const published = worker.transcriptCaches.get(activeSessionId) === generation.transcript; - this.failWorkerSnapshotCache(worker, activeSessionId, new Error(failed.error), false, failed.snapshotId); - if (published && (snapshotPurpose === "replacement" || snapshotPurpose === "catchup")) { - for (const client of this.clients) { - if (!client.attachedActiveSessionIds.has(activeSessionId)) continue; - this.queueCatchup( - client, - activeSessionId, - snapshotPurpose === "replacement" ? "replacement" : "resync", - ); - void this.catchUpClient(client).catch((error) => - this.log(`Failed to catch up client ${client.id}: ${String(error)}`), - ); - } - } - } catch (error) { - this.failWorkerSnapshotCache( + this.failCorrelatedSnapshotAttempt( worker, activeSessionId, - error instanceof Error ? error : new Error(String(error)), - true, + failed.snapshotId, + new Error(failed.error), + snapshotPurpose, ); + } catch (error) { + const frameError = error instanceof Error ? error : new Error(String(error)); + if (correlatedSnapshotId) { + this.failCorrelatedSnapshotAttempt( + worker, + activeSessionId, + correlatedSnapshotId, + frameError, + snapshotPurpose, + ); + } else { + this.failWorkerSnapshotCache(worker, activeSessionId, frameError, true); + } } return; } + if ( outboundType === "daemon_hello" || outboundType === "response" || @@ -4778,6 +5260,9 @@ export class DaemonSupervisor { } const replacementSnapshotFollows = decodedOutbound?.type === "session_replaced" && decodedOutbound.snapshotFollows === true; + const reencodeLegacyWorkerReplacement = + decodedOutbound?.type === "session_replaced" && + worker.client?.supportsServerCapability?.("immutable_snapshot_transfer_v1") !== true; const transcriptChanged = outboundType === "session_replaced" || outboundType === "session_closed" || @@ -4796,6 +5281,15 @@ export class DaemonSupervisor { if (replacementSnapshotFollows && !client.capabilities.has("chunked_snapshot")) { continue; } + if (reencodeLegacyWorkerReplacement && client.capabilities.has("chunked_snapshot")) { + this.queueCatchup(client, activeSessionId, "replacement"); + if (!client.catchupPromise) { + void this.catchUpClient(client).catch((error) => + this.log(`Failed to re-encode legacy worker snapshot for ${client.id}: ${String(error)}`), + ); + } + continue; + } if (outboundType === "extension_ui_request" && !client.supportsExtensionUi) { continue; } @@ -4851,8 +5345,7 @@ export class DaemonSupervisor { ): void { worker.snapshotCache.delete(activeSessionId); if (invalidateLoad) { - worker.snapshotLoads.delete(`${activeSessionId}:chunked`); - worker.snapshotLoads.delete(`${activeSessionId}:full`); + this.invalidateSnapshotLoads(worker, activeSessionId); } if (!transcriptChanged) return; const transcript = worker.transcriptCaches.get(activeSessionId); @@ -4921,9 +5414,7 @@ export class DaemonSupervisor { } private async drainClientCatchups(client: DaemonSocketClient): Promise { - if (client.socket.destroyed) { - return; - } + if (client.socket.destroyed) return; const pending = [...(client.catchupActiveSessionIds ?? [])].map((activeSessionId) => ({ activeSessionId, purpose: client.catchupPurposes?.get(activeSessionId) ?? ("resync" as const), @@ -4932,20 +5423,29 @@ export class DaemonSupervisor { client.catchupPurposes?.clear(); for (let index = 0; index < pending.length; index++) { const { activeSessionId, purpose } = pending[index]!; + if (!client.attachedActiveSessionIds.has(activeSessionId)) continue; let releaseTranscript: (() => void) | undefined; + let attached: WorkerAttachData | undefined; + const attachmentEpoch = this.attachmentEpoch(client, activeSessionId); try { - const attached = await this.attachClient(client, { - type: "attach", - activeSessionId, - capabilities: [...client.capabilities], - supportsExtensionUi: client.supportsExtensionUi, - }); + attached = await this.attachClient( + client, + { + type: "attach", + activeSessionId, + capabilities: [...client.capabilities], + supportsExtensionUi: client.supportsExtensionUi, + }, + { selector: activeSessionId, epoch: attachmentEpoch, requireAttached: true }, + ); + if (!client.attachedActiveSessionIds.has(activeSessionId)) { + attached.releaseTranscript?.(); + continue; + } releaseTranscript = attached.releaseTranscript; if (client.capabilities.has("chunked_snapshot")) { const transcript = attached.transcript; - if (!transcript) { - throw new Error("Session worker did not provide a snapshot transcript"); - } + if (!transcript) throw new Error("Session worker did not provide a snapshot transcript"); if (purpose === "replacement") { this.write(client, { type: "session_replaced", @@ -4987,12 +5487,7 @@ export class DaemonSupervisor { messages: attached.result.snapshot.messages, meta, } - : { - type: "session_resynced", - activeSessionId, - snapshot: attached.result.snapshot, - meta, - }; + : { type: "session_resynced", activeSessionId, snapshot: attached.result.snapshot, meta }; if (!this.write(client, catchup)) { for (const remaining of pending.slice(index + 1)) { this.queueCatchup(client, remaining.activeSessionId, remaining.purpose); @@ -5001,6 +5496,25 @@ export class DaemonSupervisor { } } catch (error) { releaseTranscript?.(); + const snapshotError = error instanceof Error ? error : new Error(String(error)); + if (attached) { + if (!this.deliveredSnapshotFailures.has(snapshotError)) { + await this.deliverWorkerSnapshotFailure( + client, + attached.result, + attached.result.snapshotStream?.id ?? `snapshot-${randomUUID()}`, + purpose, + snapshotError, + ); + } + } else if (client.attachedActiveSessionIds.has(activeSessionId)) { + await this.writeSnapshotRecord(client, { + type: "session_snapshot_failed", + activeSessionId, + snapshotId: `snapshot-${randomUUID()}`, + error: snapshotError.message, + }); + } this.log(`Failed to catch up client ${client.id} for ${activeSessionId}: ${String(error)}`); } } diff --git a/packages/coding-agent/src/modes/daemon/daemon-worker-client.ts b/packages/coding-agent/src/modes/daemon/daemon-worker-client.ts index b3e01b95ee..9c435b2e72 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-worker-client.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-worker-client.ts @@ -1,7 +1,7 @@ import { createConnection, type Socket } from "node:net"; import { serializeJsonLine } from "../rpc/jsonl.js"; import { type PrivateFrame, PrivateFramedChannel } from "../session-worker/private-framing.js"; -import type { DaemonCommand, DaemonOutbound, DaemonResponse } from "./daemon-protocol.js"; +import type { DaemonCommand, DaemonOutbound, DaemonResponse, DaemonServerCapability } from "./daemon-protocol.js"; import { type DaemonWorkerCommand, type DaemonWorkerCommandBody, @@ -78,6 +78,10 @@ export class DaemonWorkerClient { socket.on("close", () => this.notifyClosed(socket, new Error("Daemon worker socket closed"))); } + supportsServerCapability(capability: DaemonServerCapability): boolean { + return this.hello?.serverCapabilities.includes(capability) === true; + } + waitForHello(timeoutMs = 3000): Promise { if (this.hello) { return Promise.resolve(this.hello); diff --git a/packages/coding-agent/src/modes/daemon/snapshot-transcript-cache.ts b/packages/coding-agent/src/modes/daemon/snapshot-transcript-cache.ts index 874c5632e7..f7f8e50921 100644 --- a/packages/coding-agent/src/modes/daemon/snapshot-transcript-cache.ts +++ b/packages/coding-agent/src/modes/daemon/snapshot-transcript-cache.ts @@ -1,15 +1,157 @@ -import { mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { createHash, randomUUID } from "node:crypto"; +import { + chmodSync, + lstatSync, + mkdirSync, + mkdtempSync, + readdirSync, + readFileSync, + rmSync, + writeFileSync, +} from "node:fs"; +import { appendFile, mkdir, readFile, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; import { join } from "node:path"; +import { isDeepStrictEqual } from "node:util"; import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import { getProcessStartId } from "../../core/session-lease.js"; +import { isProcessAlive } from "../../utils/child-process.js"; export const SNAPSHOT_TARGET_CHUNK_BYTES = 512 * 1024; export const SNAPSHOT_MEMORY_CACHE_BYTES = 4 * 1024 * 1024; +const SNAPSHOT_CACHE_UID = typeof process.getuid === "function" ? process.getuid() : "user"; +const SNAPSHOT_TRANSFER_CACHE_PARENT = join(tmpdir(), `prime-agent-snapshot-transfers-${SNAPSHOT_CACHE_UID}`); +const SNAPSHOT_CACHE_OWNER_PREFIX = "owner-"; +const SNAPSHOT_PREPARATION_YIELD_BYTES = 256 * 1024; +const SNAPSHOT_STRING_SLICE_CHARS = 64 * 1024; +let defaultSnapshotTransferCacheRoot: string | undefined; +let defaultSnapshotTransferCacheCleanupRegistered = false; + +export interface SnapshotTranscriptCacheIo { + mkdir(path: string, options: { recursive: boolean; mode: number }): Promise; + writeFile(path: string, data: Uint8Array, options: { mode: number; flag: "wx" }): Promise; + appendFile(path: string, data: Uint8Array): Promise; + readFile(path: string): Promise; + rm(path: string, options: { recursive: boolean; force: boolean }): Promise; +} + +const defaultSnapshotTranscriptCacheIo: SnapshotTranscriptCacheIo = { + mkdir, + writeFile, + appendFile, + readFile, + rm, +}; + +function processStartHash(processStartId: string | undefined): string { + return createHash("sha256") + .update(processStartId ?? "unknown") + .digest("hex") + .slice(0, 16); +} + +function assertSafeSnapshotCacheParent(parentRoot: string): void { + let metadata: ReturnType; + try { + metadata = lstatSync(parentRoot); + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== "ENOENT") throw error; + try { + mkdirSync(parentRoot, { recursive: false, mode: 0o700 }); + } catch (mkdirError) { + if ((mkdirError as NodeJS.ErrnoException).code !== "EEXIST") throw mkdirError; + } + metadata = lstatSync(parentRoot); + } + if (!metadata.isDirectory() || metadata.isSymbolicLink()) { + throw new Error(`Snapshot cache parent is not a private directory: ${parentRoot}`); + } + if (typeof process.getuid === "function" && metadata.uid !== process.getuid()) { + throw new Error(`Snapshot cache parent is not owned by the current user: ${parentRoot}`); + } + if ((metadata.mode & 0o077) !== 0) { + chmodSync(parentRoot, 0o700); + const secured = lstatSync(parentRoot); + if (!secured.isDirectory() || secured.isSymbolicLink() || (secured.mode & 0o077) !== 0) { + throw new Error(`Snapshot cache parent permissions are unsafe: ${parentRoot}`); + } + } +} + +/** Remove only process roots whose recorded pid/start identity can no longer own them. */ +export function sweepAbandonedSnapshotCacheRoots(parentRoot: string): void { + assertSafeSnapshotCacheParent(parentRoot); + for (const entry of readdirSync(parentRoot, { withFileTypes: true })) { + if (!entry.isDirectory() || !entry.name.startsWith(SNAPSHOT_CACHE_OWNER_PREFIX)) continue; + const match = /^owner-(\d+)-([a-f0-9]{16})-[a-zA-Z0-9_-]+$/.exec(entry.name); + if (!match) continue; + const ownerPid = Number(match[1]); + if (isProcessAlive(ownerPid)) { + const currentStartId = getProcessStartId(ownerPid); + if (currentStartId === undefined || processStartHash(currentStartId) === match[2]) continue; + } + try { + rmSync(join(parentRoot, entry.name), { recursive: true, force: true }); + } catch { + // One unreadable abandoned root must not block independent owners. + } + } +} + +export function createSnapshotCacheProcessRoot(parentRoot: string, ownerToken = randomUUID()): string { + sweepAbandonedSnapshotCacheRoots(parentRoot); + const prefix = join( + parentRoot, + `${SNAPSHOT_CACHE_OWNER_PREFIX}${process.pid}-${processStartHash(getProcessStartId(process.pid))}-${ownerToken}-`, + ); + const root = mkdtempSync(prefix); + const metadata = lstatSync(root); + if ( + !metadata.isDirectory() || + metadata.isSymbolicLink() || + (typeof process.getuid === "function" && metadata.uid !== process.getuid()) + ) { + try { + rmSync(root, { recursive: true, force: true }); + } catch { + // Preserve the ownership failure. + } + throw new Error(`Snapshot cache process root is unsafe: ${root}`); + } + if ((metadata.mode & 0o077) !== 0) chmodSync(root, 0o700); + return root; +} + +function getDefaultSnapshotTransferCacheRoot(): string { + if (!defaultSnapshotTransferCacheRoot) { + defaultSnapshotTransferCacheRoot = createSnapshotCacheProcessRoot(SNAPSHOT_TRANSFER_CACHE_PARENT); + } + if (!defaultSnapshotTransferCacheCleanupRegistered) { + defaultSnapshotTransferCacheCleanupRegistered = true; + process.once("exit", () => { + const root = defaultSnapshotTransferCacheRoot; + defaultSnapshotTransferCacheRoot = undefined; + if (!root) return; + try { + rmSync(root, { recursive: true, force: true }); + } catch { + // Startup sweep owns crash and permission cleanup on the next process. + } + }); + } + return defaultSnapshotTransferCacheRoot; +} interface SnapshotTranscriptChunk { buffer?: Buffer; path?: string; } +interface SnapshotPayloadChunk extends SnapshotTranscriptChunk { + messageCount: number; + readPromise?: Promise; +} + export interface SnapshotTranscriptCacheOptions { activeSessionId: string; snapshotId: string; @@ -17,60 +159,527 @@ export interface SnapshotTranscriptCacheOptions { cacheRoot: string; targetChunkBytes?: number; memoryCacheBytes?: number; + io?: SnapshotTranscriptCacheIo; } -export type SnapshotTranscriptChunkSource = (Iterable | AsyncIterable) & { +interface SnapshotTranscriptChunkLifecycle { markFailed?(error: Error): void; dispose?(): void; -}; +} + +export type SnapshotTranscriptWireChunk = Buffer | readonly Buffer[]; +export type SnapshotTranscriptChunkSource = ( + | Iterable + | AsyncIterable +) & + SnapshotTranscriptChunkLifecycle; + +type SnapshotTranscriptChunkIterable = Iterable & SnapshotTranscriptChunkLifecycle; + +function cloneSnapshotValue(value: T, seen = new Map()): T { + if (value === null || typeof value !== "object") return value; + const objectValue = value as object; + const existing = seen.get(objectValue); + if (existing !== undefined) return existing as T; + if (Buffer.isBuffer(value)) return Buffer.from(value) as T; + if (ArrayBuffer.isView(value)) return structuredClone(value); + if (value instanceof Date) return new Date(value.getTime()) as T; + if (Array.isArray(value)) { + const cloned: unknown[] = []; + seen.set(objectValue, cloned); + for (const item of value) cloned.push(cloneSnapshotValue(item, seen)); + return cloned as T; + } + const cloned: Record = {}; + seen.set(objectValue, cloned); + for (const key of Object.keys(value)) { + cloned[key] = cloneSnapshotValue((value as Record)[key], seen); + } + return cloned as T; +} + +/** Select a coherent transcript view without copying immutable string payload bytes. */ +export function createImmutableSnapshotMessages(messages: readonly AgentMessage[]): AgentMessage[] { + return cloneSnapshotValue([...messages]); +} + +function createGenerationGuard( + source: readonly AgentMessage[], + selected: readonly AgentMessage[], + validate?: () => void, +): () => void { + return () => { + validate?.(); + if (!isDeepStrictEqual(source, selected)) { + throw new Error("Session snapshot generation changed during transcript preparation"); + } + }; +} + +interface JsonParts { + parts: Buffer[]; + bytes: number; +} + +class IncrementalJsonEncoder { + private readonly parts: Buffer[] = []; + private bytes = 0; + private bytesSinceYield = 0; + private readonly ancestors = new Set(); + + constructor( + private readonly signal: AbortSignal | undefined, + private readonly assertGeneration: () => void, + ) {} + + async encode(value: unknown): Promise { + const encoded = await this.encodeValue(value, "", false); + if (!encoded) throw new Error("Snapshot message is not JSON serializable"); + await this.yieldIfNeeded(true); + return { parts: this.parts, bytes: this.bytes }; + } + + private emit(value: string | Buffer): void { + const buffer = typeof value === "string" ? Buffer.from(value) : value; + if (buffer.length === 0) return; + this.parts.push(buffer); + this.bytes += buffer.length; + this.bytesSinceYield += buffer.length; + } + + private async yieldIfNeeded(force = false): Promise { + if (!force && this.bytesSinceYield < SNAPSHOT_PREPARATION_YIELD_BYTES) return; + this.bytesSinceYield = 0; + await new Promise((resolveYield) => setImmediate(resolveYield)); + this.signal?.throwIfAborted(); + this.assertGeneration(); + } + + private resolveJsonValue(value: unknown, key: string): unknown { + if (value && typeof value === "object" && "toJSON" in value) { + const toJSON = (value as { toJSON?: unknown }).toJSON; + if (typeof toJSON === "function") return toJSON.call(value, key); + } + return value; + } + + private async encodeString(value: string): Promise { + this.emit('"'); + for (let start = 0; start < value.length; ) { + let end = Math.min(value.length, start + SNAPSHOT_STRING_SLICE_CHARS); + if ( + end < value.length && + value.charCodeAt(end - 1) >= 0xd800 && + value.charCodeAt(end - 1) <= 0xdbff && + value.charCodeAt(end) >= 0xdc00 && + value.charCodeAt(end) <= 0xdfff + ) { + end++; + } + const quoted = JSON.stringify(value.slice(start, end)); + this.emit(quoted.slice(1, -1)); + start = end; + await this.yieldIfNeeded(); + } + this.emit('"'); + } + + private async encodeValue(rawValue: unknown, key: string, arrayElement: boolean): Promise { + const value = this.resolveJsonValue(rawValue, key); + if (value === null) { + this.emit("null"); + return true; + } + switch (typeof value) { + case "string": + await this.encodeString(value); + return true; + case "number": + this.emit(Number.isFinite(value) ? String(value === 0 ? 0 : value) : "null"); + return true; + case "boolean": + this.emit(value ? "true" : "false"); + return true; + case "bigint": + throw new TypeError("Do not know how to serialize a BigInt"); + case "undefined": + case "function": + case "symbol": + if (arrayElement) { + this.emit("null"); + return true; + } + return false; + case "object": + break; + } + const objectValue = value as object; + if (this.ancestors.has(objectValue)) { + throw new TypeError("Converting circular structure to JSON"); + } + this.ancestors.add(objectValue); + try { + if (Array.isArray(value)) { + this.emit("["); + for (let index = 0; index < value.length; index++) { + if (index > 0) this.emit(","); + await this.encodeValue(value[index], String(index), true); + } + this.emit("]"); + return true; + } + this.emit("{"); + let emitted = 0; + for (const property of Object.keys(value)) { + const propertyValue = this.resolveJsonValue((value as Record)[property], property); + if (["undefined", "function", "symbol"].includes(typeof propertyValue)) continue; + if (emitted++ > 0) this.emit(","); + await this.encodeString(property); + this.emit(":"); + await this.encodeValue(propertyValue, property, false); + } + this.emit("}"); + return true; + } finally { + this.ancestors.delete(objectValue); + } + } +} + +export interface SnapshotTranscriptPayloadOptions { + messages: readonly AgentMessage[]; + cacheRoot?: string; + targetChunkBytes?: number; + memoryCacheBytes?: number; + signal?: AbortSignal; + validateGeneration?: () => void; + io?: SnapshotTranscriptCacheIo; +} + +export class SnapshotTranscriptPayloadCache { + private readonly chunks: SnapshotPayloadChunk[] = []; + private cacheDirectory?: string; + private totalBytes = 0; + private messageTotal = 0; + private readers = 0; + private disposeRequested = false; + private disposed = false; + private cleanupPromise?: Promise; + readonly targetChunkBytes: number; + + constructor( + private readonly options: { + cacheRoot: string; + targetChunkBytes?: number; + memoryCacheBytes?: number; + io?: SnapshotTranscriptCacheIo; + }, + ) { + this.targetChunkBytes = options.targetChunkBytes ?? SNAPSHOT_TARGET_CHUNK_BYTES; + } + + get chunkCount(): number { + return this.chunks.length; + } + + get messageCount(): number { + return this.messageTotal; + } + + get bytes(): number { + return this.totalBytes; + } + + get fileBacked(): boolean { + return this.cacheDirectory !== undefined; + } + + get activeReaders(): number { + return this.readers; + } + + get retainedPayloadBytes(): number { + return this.chunks.reduce((total, chunk) => total + (chunk.buffer?.length ?? 0), 0); + } + + private get io(): SnapshotTranscriptCacheIo { + return this.options.io ?? defaultSnapshotTranscriptCacheIo; + } + + async prepare( + sourceMessages: readonly AgentMessage[], + selectedMessages: readonly AgentMessage[], + signal?: AbortSignal, + validateGeneration?: () => void, + ): Promise { + signal?.throwIfAborted(); + const assertGeneration = createGenerationGuard(sourceMessages, selectedMessages, validateGeneration); + assertGeneration(); + let currentParts: Buffer[] = []; + let currentBytes = 0; + let currentMessages = 0; + const flush = async () => { + if (currentMessages === 0) return; + await this.storePayloadChunk(currentParts, currentBytes, currentMessages); + currentParts = []; + currentBytes = 0; + currentMessages = 0; + }; + for (const message of selectedMessages) { + signal?.throwIfAborted(); + const encoded = await new IncrementalJsonEncoder(signal, assertGeneration).encode(message); + const separatorBytes = currentMessages > 0 ? 1 : 0; + if (currentMessages > 0 && currentBytes + separatorBytes + encoded.bytes > this.targetChunkBytes) { + await flush(); + } + if (currentMessages > 0) { + currentParts.push(Buffer.from(",")); + currentBytes++; + } + currentParts.push(...encoded.parts); + currentBytes += encoded.bytes; + currentMessages++; + this.messageTotal++; + assertGeneration(); + } + await flush(); + assertGeneration(); + } + + async readPayloadChunk(index: number): Promise { + if (this.disposed) throw new Error(`Snapshot payload was disposed`); + const chunk = this.chunks[index]; + if (!chunk) throw new Error(`Unknown snapshot transcript chunk: ${index}`); + if (chunk.buffer) return chunk.buffer; + if (!chunk.path) throw new Error(`Snapshot transcript chunk ${index} has no backing storage`); + chunk.readPromise ??= this.io.readFile(chunk.path).then((buffer) => { + chunk.buffer = buffer; + return buffer; + }); + try { + return await chunk.readPromise; + } catch (error) { + chunk.readPromise = undefined; + throw error; + } + } + + readPayloadChunkSync(index: number): Buffer { + if (this.disposed) throw new Error(`Snapshot payload was disposed`); + const chunk = this.chunks[index]; + if (!chunk) throw new Error(`Unknown snapshot transcript chunk: ${index}`); + if (chunk.buffer) return chunk.buffer; + if (!chunk.path) throw new Error(`Snapshot transcript chunk ${index} has no backing storage`); + return readFileSync(chunk.path); + } + + chunkMessageCount(index: number): number { + const chunk = this.chunks[index]; + if (!chunk) throw new Error(`Unknown snapshot transcript chunk: ${index}`); + return chunk.messageCount; + } + + async transferChunk(activeSessionId: string, snapshotId: string, index: number): Promise { + return this.envelopeChunk(activeSessionId, snapshotId, index, await this.readPayloadChunk(index)); + } + + transferChunkSync(activeSessionId: string, snapshotId: string, index: number): Buffer { + return this.envelopeChunk(activeSessionId, snapshotId, index, this.readPayloadChunkSync(index)); + } + + createTransfer(activeSessionId: string, snapshotId: string): SnapshotTranscriptChunkSource { + const release = this.retain(); + return { + [Symbol.asyncIterator]: async function* ( + this: SnapshotTranscriptPayloadCache, + ): AsyncGenerator { + try { + for (let index = 0; index < this.chunkCount; index++) { + const payload = await this.readPayloadChunk(index); + const prefix = + `{"type":"session_snapshot_chunk","activeSessionId":${JSON.stringify(activeSessionId)},` + + `"snapshotId":${JSON.stringify(snapshotId)},"index":${index},"messages":[`; + yield [Buffer.from(prefix), payload, Buffer.from("]}\n")]; + } + } finally { + release(); + } + }.bind(this), + dispose: release, + }; + } + + retain(): () => void { + if (this.disposed || this.disposeRequested) throw new Error("Snapshot payload was disposed"); + this.readers++; + let released = false; + return () => { + if (released) return; + released = true; + this.readers--; + if (this.readers === 0) { + if (this.disposeRequested) this.dispose(); + else this.releaseLoadedFileBuffers(); + } + }; + } + + dispose(): void { + this.disposeRequested = true; + if (this.readers !== 0) return; + if (!this.options.io) { + const directory = this.cacheDirectory; + try { + if (directory) rmSync(directory, { recursive: true, force: true }); + } catch { + // Disposal is best effort. + } finally { + this.cacheDirectory = undefined; + this.chunks.length = 0; + this.disposed = true; + } + return; + } + void this.cleanup().catch(() => undefined); + } + + async disposeAsync(): Promise { + this.disposeRequested = true; + if (this.readers === 0) await this.cleanup(); + } + + private envelopeChunk(activeSessionId: string, snapshotId: string, index: number, payload: Buffer): Buffer { + const prefix = + `{"type":"session_snapshot_chunk","activeSessionId":${JSON.stringify(activeSessionId)},` + + `"snapshotId":${JSON.stringify(snapshotId)},"index":${index},"messages":[`; + return Buffer.concat([Buffer.from(prefix), payload, Buffer.from("]}\n")]); + } + + private async storePayloadChunk(parts: Buffer[], bytes: number, messageCount: number): Promise { + if (this.disposed || this.disposeRequested) throw new Error("Snapshot payload is not writable"); + this.totalBytes += bytes; + const memoryLimit = this.options.memoryCacheBytes ?? SNAPSHOT_MEMORY_CACHE_BYTES; + if (!this.cacheDirectory && this.totalBytes > memoryLimit) { + this.cacheDirectory = join(this.options.cacheRoot, `payload-${randomUUID()}`); + await this.io.mkdir(this.cacheDirectory, { recursive: false, mode: 0o700 }); + for (let index = 0; index < this.chunks.length; index++) { + const existing = this.chunks[index]!; + if (!existing.buffer) continue; + const path = join(this.cacheDirectory, `${index}.json`); + await this.io.writeFile(path, existing.buffer, { mode: 0o600, flag: "wx" }); + this.chunks[index] = { path, messageCount: existing.messageCount }; + } + } + if (this.cacheDirectory) { + const path = join(this.cacheDirectory, `${this.chunks.length}.json`); + const [first = Buffer.alloc(0), ...rest] = parts; + await this.io.writeFile(path, first, { mode: 0o600, flag: "wx" }); + for (const part of rest) await this.io.appendFile(path, part); + this.chunks.push({ path, messageCount }); + } else { + this.chunks.push({ buffer: Buffer.concat(parts, bytes), messageCount }); + } + } + + private releaseLoadedFileBuffers(): void { + for (const chunk of this.chunks) { + if (!chunk.path) continue; + chunk.buffer = undefined; + chunk.readPromise = undefined; + } + } + + private cleanup(): Promise { + if (this.cleanupPromise) return this.cleanupPromise; + this.cleanupPromise = (async () => { + const directory = this.cacheDirectory; + try { + if (directory) await this.io.rm(directory, { recursive: true, force: true }); + } finally { + this.cacheDirectory = undefined; + this.chunks.length = 0; + this.disposed = true; + } + })(); + return this.cleanupPromise; + } +} + +export async function prepareSnapshotTranscriptPayload( + options: SnapshotTranscriptPayloadOptions, +): Promise { + options.signal?.throwIfAborted(); + const selected = createImmutableSnapshotMessages(options.messages); + const payload = new SnapshotTranscriptPayloadCache({ + cacheRoot: options.cacheRoot ?? getDefaultSnapshotTransferCacheRoot(), + ...(options.targetChunkBytes !== undefined ? { targetChunkBytes: options.targetChunkBytes } : {}), + ...(options.memoryCacheBytes !== undefined ? { memoryCacheBytes: options.memoryCacheBytes } : {}), + ...(options.io ? { io: options.io } : {}), + }); + try { + await payload.prepare(options.messages, selected, options.signal, options.validateGeneration); + return payload; + } catch (error) { + try { + await payload.disposeAsync(); + } catch { + // Preserve the encoding or write failure that caused preparation to fail. + } + throw error; + } +} export function createSnapshotTranscriptChunks(options: { activeSessionId: string; snapshotId: string; messages: readonly AgentMessage[]; + cacheRoot?: string; targetChunkBytes?: number; + memoryCacheBytes?: number; signal?: AbortSignal; -}): Iterable { - const messages = [...options.messages]; - const targetChunkBytes = options.targetChunkBytes ?? SNAPSHOT_TARGET_CHUNK_BYTES; +}): SnapshotTranscriptChunkIterable { + options.signal?.throwIfAborted(); + const cache = new SnapshotTranscriptCache({ + activeSessionId: options.activeSessionId, + snapshotId: options.snapshotId, + messages: options.messages, + cacheRoot: options.cacheRoot ?? getDefaultSnapshotTransferCacheRoot(), + ...(options.targetChunkBytes !== undefined ? { targetChunkBytes: options.targetChunkBytes } : {}), + ...(options.memoryCacheBytes !== undefined ? { memoryCacheBytes: options.memoryCacheBytes } : {}), + }); return { *[Symbol.iterator](): Iterator { options.signal?.throwIfAborted(); - let serializedMessages: string[] = []; - let serializedBytes = 0; - let index = 0; - const flush = (): Buffer | undefined => { - if (serializedMessages.length === 0) { - return undefined; - } - const prefix = - `{"type":"session_snapshot_chunk","activeSessionId":${JSON.stringify(options.activeSessionId)},` + - `"snapshotId":${JSON.stringify(options.snapshotId)},"index":${index},"messages":[`; - const line = Buffer.from(`${prefix}${serializedMessages.join(",")}]}\n`); - serializedMessages = []; - serializedBytes = 0; - index++; - return line; - }; - - for (const message of messages) { + for (const chunk of cache) { options.signal?.throwIfAborted(); - const serialized = JSON.stringify(message); - const bytes = Buffer.byteLength(serialized) + (serializedMessages.length > 0 ? 1 : 0); - if (serializedMessages.length > 0 && serializedBytes + bytes > targetChunkBytes) { - const chunk = flush(); - if (chunk) yield chunk; - } - serializedMessages.push(serialized); - serializedBytes += bytes; + yield chunk; } options.signal?.throwIfAborted(); - const chunk = flush(); - if (chunk) yield chunk; + }, + markFailed(error) { + cache.markFailed(error); + }, + dispose() { + cache.dispose(); }, }; } +export async function prepareSnapshotTranscriptCache(options: { + activeSessionId: string; + snapshotId: string; + messages: readonly AgentMessage[]; + cacheRoot?: string; + targetChunkBytes?: number; + memoryCacheBytes?: number; + signal?: AbortSignal; + validateGeneration?: () => void; + io?: SnapshotTranscriptCacheIo; +}): Promise { + const payload = await prepareSnapshotTranscriptPayload(options); + return SnapshotTranscriptCache.fromPreparedPayload(options.activeSessionId, options.snapshotId, payload); +} + export class SnapshotTranscriptCache { private readonly chunks: SnapshotTranscriptChunk[] = []; private cacheDirectory?: string; @@ -80,6 +689,7 @@ export class SnapshotTranscriptCache { private disposeRequested = false; private disposed = false; private failure?: Error; + private preparedPayload?: SnapshotTranscriptPayloadCache; private readonly chunkWaiters = new Map< number, Array<{ resolve: (buffer: Buffer | undefined) => void; reject: (error: Error) => void }> @@ -88,18 +698,41 @@ export class SnapshotTranscriptCache { readonly snapshotId: string; readonly activeSessionId: string; + static fromPreparedPayload( + activeSessionId: string, + snapshotId: string, + payload: SnapshotTranscriptPayloadCache, + ): SnapshotTranscriptCache { + const cache = new SnapshotTranscriptCache({ + activeSessionId, + snapshotId, + cacheRoot: "", + targetChunkBytes: payload.targetChunkBytes, + }); + cache.preparedPayload = payload; + cache.totalBytes = payload.bytes; + cache.completed = true; + return cache; + } + constructor(private readonly options: SnapshotTranscriptCacheOptions) { this.targetChunkBytes = options.targetChunkBytes ?? SNAPSHOT_TARGET_CHUNK_BYTES; this.snapshotId = options.snapshotId; this.activeSessionId = options.activeSessionId; if (options.messages) { - this.encodeMessages(options.messages); - this.completed = true; + try { + this.encodeMessages(options.messages); + this.completed = true; + } catch (error) { + this.markFailed(error instanceof Error ? error : new Error(String(error))); + this.disposeNow(); + throw error; + } } } get chunkCount(): number { - return this.chunks.length; + return this.preparedPayload?.chunkCount ?? this.chunks.length; } get complete(): boolean { @@ -111,78 +744,62 @@ export class SnapshotTranscriptCache { } get fileBacked(): boolean { - return this.cacheDirectory !== undefined; + return this.preparedPayload?.fileBacked ?? this.cacheDirectory !== undefined; } readChunk(index: number): Buffer { - const chunk = this.chunks[index]; - if (!chunk) { - throw new Error(`Unknown snapshot transcript chunk: ${index}`); - } - if (chunk.buffer) { - return chunk.buffer; - } - if (!chunk.path) { - throw new Error(`Snapshot transcript chunk ${index} has no backing storage`); + if (this.preparedPayload) { + return this.preparedPayload.transferChunkSync(this.activeSessionId, this.snapshotId, index); } + const chunk = this.chunks[index]; + if (!chunk) throw new Error(`Unknown snapshot transcript chunk: ${index}`); + if (chunk.buffer) return chunk.buffer; + if (!chunk.path) throw new Error(`Snapshot transcript chunk ${index} has no backing storage`); return readFileSync(chunk.path); } *[Symbol.iterator](): Iterator { - for (let index = 0; index < this.chunkCount; index++) { - yield this.readChunk(index); - } + for (let index = 0; index < this.chunkCount; index++) yield this.readChunk(index); } appendEncodedChunk(buffer: Buffer): void { - if (this.completed || this.failure || this.disposed) { + if (this.completed || this.failure || this.disposed || this.preparedPayload) { throw new Error(`Snapshot transcript ${this.snapshotId} is not writable`); } this.storeChunk(buffer); } markComplete(): void { - if (this.completed) { - return; - } + if (this.completed) return; if (this.failure || this.disposed) { throw new Error(`Snapshot transcript ${this.snapshotId} cannot be completed`); } this.completed = true; for (const [index, waiters] of this.chunkWaiters) { - if (index < this.chunks.length) { - continue; - } - for (const waiter of waiters) { - waiter.resolve(undefined); - } + if (index < this.chunks.length) continue; + for (const waiter of waiters) waiter.resolve(undefined); this.chunkWaiters.delete(index); } } markFailed(error: Error): void { - if (this.failure) { - return; - } + if (this.failure) return; this.failure = error; for (const waiters of this.chunkWaiters.values()) { - for (const waiter of waiters) { - waiter.reject(error); - } + for (const waiter of waiters) waiter.reject(error); } this.chunkWaiters.clear(); } waitForChunk(index: number): Promise { - if (this.failure) { - return Promise.reject(this.failure); - } - if (index < this.chunks.length) { - return Promise.resolve(this.readChunk(index)); - } - if (this.completed) { - return Promise.resolve(undefined); + if (this.failure) return Promise.reject(this.failure); + if (this.preparedPayload) { + return index < this.preparedPayload.chunkCount + ? this.preparedPayload.transferChunk(this.activeSessionId, this.snapshotId, index) + : Promise.resolve(undefined); } + if (index < this.chunks.length) return Promise.resolve(this.readChunk(index)); + if (this.completed) return Promise.resolve(undefined); return new Promise((resolve, reject) => { const waiters = this.chunkWaiters.get(index) ?? []; waiters.push({ resolve, reject }); @@ -191,46 +808,46 @@ export class SnapshotTranscriptCache { } retain(): () => void { - if (this.disposed) { - throw new Error(`Snapshot transcript ${this.snapshotId} was disposed`); - } + if (this.disposed) throw new Error(`Snapshot transcript ${this.snapshotId} was disposed`); + const releasePayload = this.preparedPayload?.retain(); this.readers++; let released = false; return () => { if (released) return; released = true; + releasePayload?.(); this.readers--; - if (this.readers === 0 && this.disposeRequested) { - this.disposeNow(); - } + if (this.readers === 0 && this.disposeRequested) this.disposeNow(); }; } dispose(): void { if (this.disposed || this.disposeRequested) return; this.disposeRequested = true; - if (this.readers > 0) return; - this.disposeNow(); + this.preparedPayload?.dispose(); + if (this.readers === 0) this.disposeNow(); } private disposeNow(): void { if (this.disposed) return; this.disposed = true; this.markFailed(new Error(`Snapshot transcript ${this.snapshotId} was disposed`)); - if (this.cacheDirectory) { - rmSync(this.cacheDirectory, { recursive: true, force: true }); + const directory = this.cacheDirectory; + try { + if (directory) rmSync(directory, { recursive: true, force: true }); + } catch { + // Disposal is best effort; state is always cleared below. + } finally { this.cacheDirectory = undefined; + this.chunks.length = 0; } - this.chunks.length = 0; } private encodeMessages(messages: readonly AgentMessage[]): void { let serializedMessages: string[] = []; let serializedBytes = 0; const flush = () => { - if (serializedMessages.length === 0) { - return; - } + if (serializedMessages.length === 0) return; const index = this.chunks.length; const prefix = `{"type":"session_snapshot_chunk","activeSessionId":${JSON.stringify(this.options.activeSessionId)},` + @@ -240,13 +857,10 @@ export class SnapshotTranscriptCache { serializedMessages = []; serializedBytes = 0; }; - for (const message of messages) { const serialized = JSON.stringify(message); const bytes = Buffer.byteLength(serialized) + (serializedMessages.length > 0 ? 1 : 0); - if (serializedMessages.length > 0 && serializedBytes + bytes > this.targetChunkBytes) { - flush(); - } + if (serializedMessages.length > 0 && serializedBytes + bytes > this.targetChunkBytes) flush(); serializedMessages.push(serialized); serializedBytes += bytes; } @@ -257,22 +871,22 @@ export class SnapshotTranscriptCache { this.totalBytes += buffer.length; const memoryLimit = this.options.memoryCacheBytes ?? SNAPSHOT_MEMORY_CACHE_BYTES; if (!this.cacheDirectory && this.totalBytes > memoryLimit) { - this.cacheDirectory = join(this.options.cacheRoot, this.options.snapshotId.replaceAll(/[^a-zA-Z0-9_-]/g, "_")); - mkdirSync(this.cacheDirectory, { recursive: true, mode: 0o700 }); + this.cacheDirectory = join( + this.options.cacheRoot, + `${this.options.snapshotId.replaceAll(/[^a-zA-Z0-9_-]/g, "_")}-${randomUUID()}`, + ); + mkdirSync(this.cacheDirectory, { recursive: false, mode: 0o700 }); for (let index = 0; index < this.chunks.length; index++) { const existing = this.chunks[index]!; - if (!existing.buffer) { - continue; - } + if (!existing.buffer) continue; const path = join(this.cacheDirectory, `${index}.jsonl`); - writeFileSync(path, existing.buffer, { mode: 0o600 }); + writeFileSync(path, existing.buffer, { mode: 0o600, flag: "wx" }); this.chunks[index] = { path }; } } - if (this.cacheDirectory) { const path = join(this.cacheDirectory, `${this.chunks.length}.jsonl`); - writeFileSync(path, buffer, { mode: 0o600 }); + writeFileSync(path, buffer, { mode: 0o600, flag: "wx" }); this.chunks.push({ path }); } else { this.chunks.push({ buffer }); @@ -281,9 +895,7 @@ export class SnapshotTranscriptCache { const waiters = this.chunkWaiters.get(index); if (waiters) { const stored = this.readChunk(index); - for (const waiter of waiters) { - waiter.resolve(stored); - } + for (const waiter of waiters) waiter.resolve(stored); this.chunkWaiters.delete(index); } } diff --git a/packages/coding-agent/src/modes/session-worker/private-framing.ts b/packages/coding-agent/src/modes/session-worker/private-framing.ts index 4caf143f6a..3dca2672f1 100644 --- a/packages/coding-agent/src/modes/session-worker/private-framing.ts +++ b/packages/coding-agent/src/modes/session-worker/private-framing.ts @@ -29,6 +29,23 @@ function isObjectHeader(value: unknown): value is object { return typeof value === "object" && value !== null && !Array.isArray(value); } +export function encodePrivateFrameParts( + header: THeader, + payloadParts: readonly Uint8Array[], + limits: PrivateFrameLimits = DEFAULT_PRIVATE_FRAME_LIMITS, +): Buffer[] { + const headerBuffer = Buffer.from(JSON.stringify(header), "utf8"); + const payloadLength = payloadParts.reduce((total, part) => total + part.length, 0); + assertFrameLength("header length", headerBuffer.length, limits.maxHeaderBytes); + assertFrameLength("payload length", payloadLength, limits.maxPayloadBytes); + if (headerBuffer.length === 0) throw new Error("Private frame header cannot be empty"); + const prefix = Buffer.allocUnsafe(FRAME_PREFIX_BYTES + headerBuffer.length); + prefix.writeUInt32BE(headerBuffer.length, 0); + prefix.writeUInt32BE(payloadLength, 4); + headerBuffer.copy(prefix, FRAME_PREFIX_BYTES); + return [prefix, ...payloadParts.map((part) => (Buffer.isBuffer(part) ? part : Buffer.from(part)))]; +} + export function encodePrivateFrame( header: THeader, payload: Uint8Array = Buffer.alloc(0), diff --git a/packages/coding-agent/test/agent-connection-daemon.test.ts b/packages/coding-agent/test/agent-connection-daemon.test.ts index 2476257dfd..758bb3425f 100644 --- a/packages/coding-agent/test/agent-connection-daemon.test.ts +++ b/packages/coding-agent/test/agent-connection-daemon.test.ts @@ -3422,6 +3422,153 @@ describe("DaemonAgentConnection", () => { expect((connection as unknown as { snapshotAssemblies: Map }).snapshotAssemblies.size).toBe(0); }); + it("retries one failed attach request and clears attempt state after success", async () => { + const fakeClient = new FakeDaemonClient(); + fakeClient.attachFailures = 1; + const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); + + await expect(connection.attach()).resolves.toBeUndefined(); + expect(fakeClient.requests.filter((request) => request.type === "attach")).toHaveLength(2); + expect( + (connection as unknown as { snapshotRequestAttempts: Map }).snapshotRequestAttempts.size, + ).toBe(0); + }); + + it("reuses a settled same-id assembly with a fresh timer and resolves the original attach caller", async () => { + const fakeClient = new FakeDaemonClient(); + let attachAttempt = 0; + fakeClient.attachResultFactory = (command) => { + const attempt = ++attachAttempt; + const full = createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12, { + messages: [{ role: "user", content: `generation-${attempt}`, timestamp: attempt }], + }); + const { messages: _messages, ...snapshot } = full.snapshot; + queueMicrotask(() => { + fakeClient.emitMessage({ + type: "session_snapshot_begin", + activeSessionId: command.activeSessionId, + snapshotId: "same-id", + snapshot, + messageCount: 1, + targetChunkBytes: 512 * 1024, + }); + if (attempt === 1) { + fakeClient.emitMessage({ + type: "session_snapshot_failed", + activeSessionId: command.activeSessionId, + snapshotId: "same-id", + error: "first generation failed", + }); + return; + } + fakeClient.emitMessage({ + type: "session_snapshot_chunk", + activeSessionId: command.activeSessionId, + snapshotId: "same-id", + index: 0, + messages: full.snapshot.messages, + }); + fakeClient.emitMessage({ + type: "session_snapshot_end", + activeSessionId: command.activeSessionId, + snapshotId: "same-id", + chunkCount: 1, + lastEventSequence: full.lastEventSequence, + lastEventCursor: full.lastEventCursor, + }); + }); + return { + ...full, + snapshot: { ...full.snapshot, messages: [] }, + snapshotStream: { id: "same-id", messageCount: 1, targetChunkBytes: 512 * 1024 }, + }; + }; + const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); + + await expect(connection.attach()).resolves.toBeUndefined(); + await expect(connection.getInitialSnapshot()).resolves.toMatchObject({ + messages: [{ role: "user", content: "generation-2", timestamp: 2 }], + }); + expect(fakeClient.requests.filter((request) => request.type === "attach")).toHaveLength(2); + expect((connection as unknown as { snapshotAssemblies: Map }).snapshotAssemblies.size).toBe(0); + expect( + (connection as unknown as { snapshotRequestAttempts: Map }).snapshotRequestAttempts.size, + ).toBe(0); + }); + + it("aliases a failed command reattach to one fresh attach and resolves the switch caller", async () => { + const fakeClient = new FakeDaemonClient(); + const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-source", { + snapshotTimeoutMs: 50, + }); + await connection.attach(); + fakeClient.requests.length = 0; + fakeClient.switchSessionActiveSessionId = "active-target"; + const target = createAttachResult("active-target", "client-1", undefined, 1, { + state: createConnectionState("active-target", "session-target"), + messages: [{ role: "user", content: "target", timestamp: 1 }], + }); + const emitTransfer = (snapshotId: string, fail: boolean) => { + const { messages: _messages, ...snapshot } = target.snapshot; + setImmediate(() => { + fakeClient.emitMessage({ + type: "session_snapshot_begin", + activeSessionId: "active-target", + snapshotId, + snapshot, + messageCount: 1, + targetChunkBytes: 512 * 1024, + purpose: "replacement", + }); + if (fail) { + fakeClient.emitMessage({ + type: "session_snapshot_failed", + activeSessionId: "active-target", + snapshotId, + error: "reattach generation failed", + }); + return; + } + fakeClient.emitMessage({ + type: "session_snapshot_chunk", + activeSessionId: "active-target", + snapshotId, + index: 0, + messages: target.snapshot.messages, + }); + fakeClient.emitMessage({ + type: "session_snapshot_end", + activeSessionId: "active-target", + snapshotId, + chunkCount: 1, + lastEventSequence: target.lastEventSequence, + lastEventCursor: target.lastEventCursor, + }); + }); + return { + ...target, + snapshot: { ...target.snapshot, messages: [] }, + snapshotStream: { id: snapshotId, messageCount: 1, targetChunkBytes: 512 * 1024 }, + }; + }; + fakeClient.reattachResultFactory = () => emitTransfer("reattach-g1", true); + fakeClient.attachResultFactory = () => emitTransfer("reattach-g2", false); + const events: AgentConnectionEvent[] = []; + connection.subscribe((event) => { + events.push(event); + }); + + await expect(connection.switchSession("/tmp/target.jsonl")).resolves.toEqual({ cancelled: false }); + expect(fakeClient.requests.map((request) => request.type)).toEqual(["switch_session", "reattach", "attach"]); + expect(events).toContainEqual( + expect.objectContaining({ + type: "session_replaced", + state: expect.objectContaining({ sessionId: "session-target" }), + messages: [{ role: "user", content: "target", timestamp: 1 }], + }), + ); + }); + it("rejects one failed snapshot without interrupting another session on the shared client", async () => { const fakeClient = new FakeDaemonClient(); const sibling = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-2"); @@ -3461,6 +3608,13 @@ describe("DaemonAgentConnection", () => { const failed = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await expect(failed.attach()).rejects.toThrow("snapshot encoder failed"); + expect( + fakeClient.requests.filter((request) => request.type === "attach" && request.activeSessionId === "active-1"), + ).toHaveLength(2); + expect((failed as unknown as { snapshotAssemblies: Map }).snapshotAssemblies.size).toBe(0); + expect( + (failed as unknown as { snapshotRequestAttempts: Map }).snapshotRequestAttempts.size, + ).toBe(0); emitSequencedQueueUpdate(fakeClient, "active-2", 13); await vi.waitFor(() => expect(siblingEvents).toHaveLength(1)); @@ -3645,6 +3799,123 @@ describe("DaemonAgentConnection", () => { }); }); + it("keeps a correlated replacement fence across one client-owned attach retry", async () => { + const fakeClient = new FakeDaemonClient(); + fakeClient.serverCapabilities.add("correlated_prompt_lifecycle_v1"); + const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); + await connection.attach(); + const events: AgentConnectionEvent[] = []; + connection.subscribe((event) => { + events.push(event); + }); + const replacementState = createConnectionState("active-1", "session-retried"); + const full = createAttachResult("active-1", "client-1", ["correlated_prompt_lifecycle_v1"], 1, { + state: replacementState, + messages: [{ role: "user", content: "fresh", timestamp: 2 }], + }); + full.snapshot.lastEventCursor = { generation: "generation-retried", sequence: 1 }; + full.lastEventCursor = full.snapshot.lastEventCursor; + full.replay = { status: "complete", toSequence: 1, toCursor: full.lastEventCursor }; + const { messages: _messages, ...snapshot } = full.snapshot; + fakeClient.attachResultFactory = () => full; + fakeClient.requests.length = 0; + fakeClient.emitMessage({ + type: "session_replaced", + activeSessionId: "active-1", + state: replacementState, + messages: [], + snapshotFollows: true, + meta: { + id: "active-1:retried:1", + protocol: DAEMON_PROTOCOL_INFO, + activeSessionId: "active-1", + sequence: 1, + cursor: { generation: "generation-retried", sequence: 1 }, + emittedAt: "2026-01-01T00:00:00.000Z", + }, + }); + fakeClient.emitMessage({ + type: "session_snapshot_begin", + activeSessionId: "active-1", + snapshotId: "failed-generation", + snapshot, + messageCount: 1, + targetChunkBytes: 512 * 1024, + purpose: "replacement", + }); + fakeClient.emitMessage({ + type: "session_snapshot_failed", + activeSessionId: "active-1", + snapshotId: "failed-generation", + error: "first generation failed", + }); + + await vi.waitFor(() => expect(events).toHaveLength(1)); + expect(events).toEqual([ + expect.objectContaining({ + type: "session_replaced", + state: expect.objectContaining({ sessionId: "session-retried" }), + messages: [{ role: "user", content: "fresh", timestamp: 2 }], + }), + ]); + expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach"]); + expect(events.some((event) => event.type === "closed")).toBe(false); + await expect(connection.getInitialSnapshot()).resolves.toMatchObject({ + state: { sessionId: "session-retried" }, + messages: [{ role: "user", content: "fresh", timestamp: 2 }], + }); + }); + + it("fails a correlated replacement after its one client-owned attach retry fails", async () => { + const fakeClient = new FakeDaemonClient(); + fakeClient.serverCapabilities.add("correlated_prompt_lifecycle_v1"); + const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); + await connection.attach(); + const events: AgentConnectionEvent[] = []; + connection.subscribe((event) => { + events.push(event); + }); + const replacementState = createConnectionState("active-1", "session-terminal"); + const full = createAttachResult("active-1", "client-1", ["correlated_prompt_lifecycle_v1"], 1, { + state: replacementState, + }); + full.snapshot.lastEventCursor = { generation: "generation-terminal", sequence: 1 }; + const { messages: _messages, ...snapshot } = full.snapshot; + fakeClient.attachResultFactory = () => { + throw new Error("retry generation failed"); + }; + fakeClient.requests.length = 0; + fakeClient.emitMessage({ + type: "session_replaced", + activeSessionId: "active-1", + state: replacementState, + messages: [], + snapshotFollows: true, + }); + fakeClient.emitMessage({ + type: "session_snapshot_begin", + activeSessionId: "active-1", + snapshotId: "failed-once", + snapshot, + messageCount: full.snapshot.messages.length, + targetChunkBytes: 512 * 1024, + purpose: "replacement", + }); + fakeClient.emitMessage({ + type: "session_snapshot_failed", + activeSessionId: "active-1", + snapshotId: "failed-once", + error: "first generation failed", + }); + + await vi.waitFor(() => expect(events.some((event) => event.type === "closed")).toBe(true)); + expect(events).toContainEqual({ + type: "closed", + error: expect.stringContaining("retry generation failed"), + }); + expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach"]); + }); + it("distinguishes chunked catch-up snapshots from runtime replacements", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); @@ -3733,9 +4004,16 @@ describe("DaemonAgentConnection", () => { activeSessionId === "active-1" ? recoveredSessionId : "session-sibling", ); fakeClient.requests.length = 0; + const recoveredMessages: AgentMessage[] = [{ role: "user", content: "current prompt", timestamp: 4 }]; + fakeClient.attachResultFactory = (command) => + createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 13, { + state: createConnectionState(command.activeSessionId, recoveredSessionId), + messages: recoveredMessages, + }); const snapshotId = `snapshot-failed-${purpose}`; const full = createAttachResult("active-1", "client-1", undefined, 13, { state: createConnectionState("active-1", recoveredSessionId), + messages: recoveredMessages, }); const { messages: _messages, ...snapshot } = full.snapshot; if (purpose === "replacement") { @@ -3787,11 +4065,7 @@ describe("DaemonAgentConnection", () => { }, }); } - expect(fakeClient.requests.map((request) => request.type)).toEqual([ - "get_connection_state", - "get_messages", - "get_session_context", - ]); + expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach"]); emitSequencedQueueUpdate(fakeClient, "active-2", 13); await vi.waitFor(() => expect(siblingEvents).toHaveLength(1)); expect(siblingEvents[0]).toMatchObject({ type: "session_event", event: { type: "session_action_update" } }); @@ -3899,7 +4173,7 @@ describe("DaemonAgentConnection", () => { expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach", "attach"]); }); - it("resets a connected transport when reattach fails during supervisor recovery", async () => { + it("retries one reattach on the recovered transport before resetting it", async () => { const fakeClient = new FakeDaemonClient(); const recoverDaemon = vi.fn(async () => undefined); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { @@ -3919,9 +4193,9 @@ describe("DaemonAgentConnection", () => { fakeClient.emitClose(new Error("Daemon socket closed")); await vi.waitFor(() => expect(statuses).toEqual(["reconnecting", "connected"])); - expect(recoverDaemon).toHaveBeenCalledTimes(2); - expect(fakeClient.reconnectCount).toBe(2); - expect(fakeClient.resetTransportCount).toBe(1); + expect(recoverDaemon).toHaveBeenCalledOnce(); + expect(fakeClient.reconnectCount).toBe(1); + expect(fakeClient.resetTransportCount).toBe(0); expect(fakeClient.requests.filter((request) => request.type === "attach")).toHaveLength(3); }); @@ -4774,7 +5048,7 @@ describe("DaemonAgentConnection", () => { expect(fakeClient.getMessageListenerCount()).toBe(0); expect(fakeClient.getCloseListenerCount()).toBe(0); - expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach", "detach"]); + expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach", "attach", "detach"]); expect(fakeClient.closeCount).toBe(1); }); }); diff --git a/packages/coding-agent/test/agent-session-recursion.test.ts b/packages/coding-agent/test/agent-session-recursion.test.ts index 793801106e..33979fbe14 100644 --- a/packages/coding-agent/test/agent-session-recursion.test.ts +++ b/packages/coding-agent/test/agent-session-recursion.test.ts @@ -18,7 +18,7 @@ import { createAgentSessionMessage, isAgentSessionMessage, } from "../src/core/agent-messages.js"; -import { AgentSession } from "../src/core/agent-session.js"; +import { AgentSession, type RlmChildAgentSnapshot } from "../src/core/agent-session.js"; import { AuthStorage } from "../src/core/auth-storage.js"; import type { LoadExtensionsResult } from "../src/core/extensions/index.js"; import { type HostRequestHandlers, ReplKernelManager } from "../src/core/kernel/index.js"; @@ -120,7 +120,7 @@ interface InspectableRlmSession { } >; _rlmChildCleanupFailures: Map>["subagents"][number]>; - _rlmChildSessions: Map; + _rlmChildSessions: Map; _rlmChildUnsubscribes: Map void>; _deletedRlmChildIds: Set; _rlmQuiescenceWaitAborts: Set; @@ -523,6 +523,10 @@ describe("AgentSession rlm recursion", () => { mkdirSync(childDir, { recursive: true }); const child = createSession({ rlmSessionDir: childDir }); child.setSessionName("restored-worker"); + const restoredAnswer = assistantMessage("restored answer", usage(7, 3)); + restoredAnswer.content.push({ type: "toolCall", id: "tool-1", name: "ipython", arguments: {} }); + child.agent.state.messages.push(restoredAnswer); + child.setCurrentRecap("restored recap"); const disposeChild = vi.spyOn(child, "disposeAsync"); const root = createSession(); const childStatuses: string[] = []; @@ -533,6 +537,15 @@ describe("AgentSession rlm recursion", () => { }); expect(root.registerRlmChildSession(childId, child)).toBe(true); + expect(root.getRlmChildSnapshots()).toEqual([ + expect.objectContaining({ + id: childId, + answerPreview: "restored answer", + toolUseCount: 1, + tokenCount: 10, + recap: "restored recap", + }), + ]); expect((await root.listRlmSubagents()).subagents).toEqual([ expect.objectContaining({ rlm_child_id: childId, @@ -549,6 +562,40 @@ describe("AgentSession rlm recursion", () => { expect(childStatuses).toEqual(["cancelled"]); }); + it("projects live follow-up activity for a restored session-only child", async () => { + const childId = "restored-followup-child"; + const childDir = join(tempDir, childId); + mkdirSync(childDir, { recursive: true }); + const followUpGate = deferred(); + let followUpStarted = false; + const child = createSession({ + rlmSessionDir: childDir, + streamFn: () => { + const stream = createAssistantMessageEventStream(); + followUpStarted = true; + void followUpGate.promise.then(() => { + stream.push({ type: "done", reason: "stop", message: assistantMessage("follow-up answer") }); + }); + return stream; + }, + }); + const root = createSession(); + expect(root.registerRlmChildSession(childId, child)).toBe(true); + + const followUp = child.prompt("follow-up work"); + await waitFor(() => followUpStarted); + const busy = root.getRlmChildSnapshots(); + expect(busy).toEqual([expect.objectContaining({ id: childId, status: "done" })]); + expect(busy[0]?.activity).toBeDefined(); + + followUpGate.resolve(); + await followUp; + await child.waitForIdle(); + const idle = root.getRlmChildSnapshots(); + expect(idle).toEqual([expect.objectContaining({ id: childId, status: "done" })]); + expect(idle[0]?.activity).toBeUndefined(); + }); + it("retries and releases failed retained child cleanup on the next compaction", async () => { const childId = "retained-retry-child"; const childDir = join(tempDir, childId); @@ -651,9 +698,9 @@ describe("AgentSession rlm recursion", () => { }); // The same child can be visible in both lifecycle registries while // deletion settles; it must be traversed exactly once and remain hidden. - rootInternals._rlmChildSessions.set(id, parent); + rootInternals._rlmChildSessions.set(id, { session: parent }); } else { - rootInternals._rlmChildSessions.set(id, parent); + rootInternals._rlmChildSessions.set(id, { session: parent }); if (hiding === "deleted") { rootInternals._deletedRlmChildIds.add(id); } else { @@ -745,6 +792,14 @@ describe("AgentSession rlm recursion", () => { await waitFor(() => childUpdates.some((update) => update.status === "done")); const doneUpdate = [...childUpdates].reverse().find((update) => update.status === "done"); expect(doneUpdate?.answerPreview).toBe("child answer: summarize shard 1"); + expect(root.getRlmChildSnapshots()).toEqual([ + expect.objectContaining({ + id: result.rlm_child_id, + status: "done", + answerPreview: doneUpdate?.answerPreview, + durationMs: expect.any(Number), + }), + ]); const child = root.getRlmChildSession(result.rlm_child_id); expect(child?.messages[0]).toMatchObject({ role: "custom", @@ -842,6 +897,9 @@ describe("AgentSession rlm recursion", () => { }, }); const spawned = await root.runRlmChild("pending task", { name: "pending-child" }); + expect(root.getRlmChildSnapshots()).toEqual([ + expect.objectContaining({ id: spawned.rlm_child_id, status: "queued" }), + ]); const handlers = (root as unknown as InspectableRlmSession)._createKernelHostHandlers(); const send = handlers["agent_message.send"]; if (!send) throw new Error("Missing agent_message.send host handler"); @@ -2196,14 +2254,33 @@ describe("AgentSession rlm recursion", () => { if (!child) { throw new Error("Missing retained child session"); } + const rootInternals = root as unknown as InspectableRlmSession; + await waitFor(() => !rootInternals._activeRlmChildRuns.has(childId)); + const completeRelease = root.releaseRlmChildSession(childId, child); + if (!completeRelease) throw new Error("Failed to release retained child"); + child.setCurrentRecap("retained recap"); child.setSessionName("renamed-worker"); const childUpdates = events.filter( - (event): event is { type: "rlm_child_update"; child: { sessionName?: string } } => + (event): event is { type: "rlm_child_update"; child: RlmChildAgentSnapshot } => typeof event === "object" && event !== null && (event as { type?: string }).type === "rlm_child_update", ); - expect(childUpdates.at(-1)?.child.sessionName).toBe("renamed-worker"); + expect(childUpdates.at(-1)?.child).toMatchObject({ + sessionName: "renamed-worker", + durationMs: expect.any(Number), + tokenCount: 10, + recap: "retained recap", + repliedSinceTask: false, + }); + expect(root.getRlmChildSnapshots()).toEqual([ + expect.objectContaining({ + sessionName: "renamed-worker", + durationMs: expect.any(Number), + tokenCount: 10, + recap: "retained recap", + }), + ]); }); it("surfaces a child's recap on its snapshot once the summarizer sets it", async () => { diff --git a/packages/coding-agent/test/daemon-mode.test.ts b/packages/coding-agent/test/daemon-mode.test.ts index f92101ba4f..82290ff5d4 100644 --- a/packages/coding-agent/test/daemon-mode.test.ts +++ b/packages/coding-agent/test/daemon-mode.test.ts @@ -3400,6 +3400,49 @@ describe("daemon mode helpers", () => { expect(client.snapshotActiveSessionIds).not.toContain("active"); }); + it("shares one stable transcript payload across attach, replacement, and catch-up transfers", async () => { + const daemon = new AgentDaemon("/tmp/prime-agent-shared-snapshot-payload.sock", { + defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, + createRuntime: vi.fn(), + }); + const messages: DaemonAttachResult["snapshot"]["messages"] = [ + { role: "user", content: "shared payload", timestamp: 1 }, + ]; + type Transfer = AsyncIterable & { dispose?(): void }; + const internals = daemon as unknown as { + prepareWorkerSnapshotTranscript(options: { + activeSessionId: string; + snapshotId: string; + generationKey: string; + messages: DaemonAttachResult["snapshot"]["messages"]; + }): Promise; + snapshotPayloadGenerations: Map }>; + }; + const transfers = await Promise.all( + ["attach", "replacement", "catchup"].map((purpose) => + internals.prepareWorkerSnapshotTranscript({ + activeSessionId: "active-shared", + snapshotId: `snapshot-${purpose}`, + generationKey: "session:generation:1:1", + messages, + }), + ), + ); + const payload = await internals.snapshotPayloadGenerations.get("active-shared")!.promise; + expect(payload.activeReaders).toBe(3); + const iterators = transfers.map((transfer) => transfer[Symbol.asyncIterator]()); + const chunks = await Promise.all(iterators.map((iterator) => iterator.next())); + const payloadParts = chunks.map((chunk) => { + if (chunk.done || Buffer.isBuffer(chunk.value)) throw new Error("expected segmented snapshot transfer"); + return chunk.value[1]; + }); + expect(payloadParts[1]).toBe(payloadParts[0]); + expect(payloadParts[2]).toBe(payloadParts[0]); + await Promise.all(iterators.map((iterator) => iterator.return?.())); + expect(payload.activeReaders).toBe(0); + payload.dispose(); + }); + it("falls back to a full replacement when snapshot cache creation fails", async () => { const root = mkdtempSync(join(tmpdir(), "prime-agent-daemon-replacement-fallback-")); try { @@ -4464,7 +4507,7 @@ describe("daemon mode helpers", () => { expect(internals.buildRlmChildSnapshotsWithPassiveRlmSubagents).toHaveBeenCalledTimes(2); expect(snapshot.children).toEqual([expect.objectContaining({ id: "new-child" })]); - expect(snapshot.messages).toBe(replacementSession.messages); + expect(snapshot.messages).toStrictEqual(replacementSession.messages); } finally { rmSync(tempDir, { recursive: true, force: true }); } @@ -5728,7 +5771,7 @@ describe("daemon mode helpers", () => { } }); - it("re-adopts a resident child when its passivation close fails", async () => { + it("keeps ownership of a resident child until passivation succeeds", async () => { const tempDir = mkdtempSync(join(tmpdir(), "prime-agent-daemon-passivation-close-failure-")); let releaseAbort!: () => void; const abortGate = new Promise((resolve) => { @@ -5752,7 +5795,6 @@ describe("daemon mode helpers", () => { const childState = await internals.createRuntime({ type: "create", sessionPath: fixture.childSessionFile }); const parentSession = parentState.runtime.session as unknown as { releaseRlmChildSession: ReturnType; - registerRlmChildSession: ReturnType; }; let parentOwnsChild = true; let forwarderActive = true; @@ -5765,16 +5807,11 @@ describe("daemon mode helpers", () => { }); parentSession.releaseRlmChildSession = vi.fn(() => { if (!parentOwnsChild) return false; - parentOwnsChild = false; - return unsubscribeForwarder; - }); - parentSession.registerRlmChildSession = vi.fn( - (_childId: string, _childSession: unknown, unsubscribe: () => void) => { - parentOwnsChild = true; - forwarderActive = unsubscribe === unsubscribeForwarder; - return true; - }, - ); + return () => { + parentOwnsChild = false; + unsubscribeForwarder(); + }; + }); childState.unsubscribe = vi .fn() .mockImplementationOnce(() => { @@ -5800,11 +5837,6 @@ describe("daemon mode helpers", () => { await expect(delivery).resolves.toMatchObject({ deliveryStatus: "delivered" }); expect(internals.sessions.get(childState.activeSessionId)).toBe(childState); expect(parentOwnsChild).toBe(true); - expect(parentSession.registerRlmChildSession).toHaveBeenCalledWith( - fixture.childId, - childState.runtime.session, - unsubscribeForwarder, - ); expect(unsubscribeForwarder).not.toHaveBeenCalled(); emitChildUpdate("recap after failed close"); expect(parentUpdates).toEqual(["recap after failed close"]); @@ -9122,6 +9154,7 @@ function makeRuntimeSession( setSubagentRuntimeHost: vi.fn(), setSessionReplacementAdmissionGuard: vi.fn(), getRlmChildRunStatus: vi.fn(() => "running"), + getRlmChildSnapshots: vi.fn(() => []), registerRlmChildSession: vi.fn(() => true), releaseRlmChildSession: vi.fn(() => vi.fn()), subscribe: vi.fn(() => vi.fn()), diff --git a/packages/coding-agent/test/daemon-protocol.test.ts b/packages/coding-agent/test/daemon-protocol.test.ts index d516ec57fe..44b1a7863a 100644 --- a/packages/coding-agent/test/daemon-protocol.test.ts +++ b/packages/coding-agent/test/daemon-protocol.test.ts @@ -274,6 +274,12 @@ describe("daemon protocol helpers", () => { expect(isDaemonMutatingCommand({ type: "get_prompt_lifecycles" })).toBe(false); }); + it("negotiates immutable worker snapshots without changing legacy client defaults", () => { + expect(DAEMON_DEFAULT_CLIENT_CAPABILITIES).not.toContain("immutable_snapshot_transfer_v1"); + expect(DAEMON_SUPPORTED_CLIENT_CAPABILITIES).toContain("immutable_snapshot_transfer_v1"); + expect(DAEMON_DEFAULT_SERVER_CAPABILITIES).toContain("immutable_snapshot_transfer_v1"); + }); + it("removes correlated provenance from legacy daemon events", () => { const event: DaemonOutbound = { type: "session_event", diff --git a/packages/coding-agent/test/daemon-session-list.test.ts b/packages/coding-agent/test/daemon-session-list.test.ts index a332372bfb..6cb772e27f 100644 --- a/packages/coding-agent/test/daemon-session-list.test.ts +++ b/packages/coding-agent/test/daemon-session-list.test.ts @@ -1,6 +1,7 @@ import { resolve } from "node:path"; import type { AgentMessage } from "@earendil-works/pi-agent-core"; import { describe, expect, it } from "vitest"; +import type { RlmChildAgentSnapshot } from "../src/core/agent-session.js"; import type { AgentCronJob } from "../src/core/cron-jobs.js"; import type { SessionInfo } from "../src/core/session-manager.js"; import type { ActiveSessionState, DaemonSocketClient } from "../src/modes/daemon/active-session-state.js"; @@ -496,150 +497,42 @@ describe("summaryForActiveSession recap currency", () => { }); describe("buildRlmChildSnapshots", () => { - it("collects children and grandchildren with event-compatible parent ids", () => { - const parent = makeState({ activeSessionId: "parent", sessionFile: "/tmp/parent.jsonl" }); - const child = makeState({ - activeSessionId: "child", - model: { provider: "anthropic", id: "claude-opus-4-7" }, - isStreaming: true, - metadata: { - kind: "subagent", - createdAt: 1, - parentActiveSessionId: "parent", - rlmChildId: "sub-aaa", - rlmParentNodeId: "sub-aaa", - prompt: "Summarize the repo\nlayout", - sessionDir: "/tmp/artifacts/sub-aaa", - }, - messages: [ - { role: "user", content: "Summarize the repo layout" }, - { - role: "assistant", - content: [ - { type: "text", text: "The repo is an npm workspace." }, - { type: "toolCall", id: "tool-1", name: "ipython", arguments: {} }, - ], - }, - ] as AgentMessage[], - contextTokens: 41_000, - }); - const grandchild = makeState({ - activeSessionId: "grandchild", - metadata: { - kind: "subagent", - createdAt: 2, - parentActiveSessionId: "child", - rlmChildId: "sub-bbb", - rlmParentNodeId: "sub-bbb", - prompt: "Read the docs", - sessionDir: "/tmp/artifacts/sub-aaa/sub-bbb", - }, - }); - const unrelated = makeState({ - activeSessionId: "unrelated-child", - metadata: { - kind: "subagent", - createdAt: 3, - parentActiveSessionId: "someone-else", - rlmChildId: "sub-ccc", - }, - }); - - const snapshots = buildRlmChildSnapshots("parent", [parent, child, grandchild, unrelated]); - - expect(snapshots.map((snapshot) => [snapshot.id, snapshot.parentId, snapshot.status])).toEqual([ - ["sub-aaa", undefined, "running"], - ["sub-bbb", "sub-aaa", "done"], - ]); - expect(snapshots[0]).toMatchObject({ - model: "anthropic/claude-opus-4-7", - label: "Summarize the repo layout", - answerPreview: "The repo is an npm workspace.", - toolUseCount: 1, - tokenCount: 41_000, - sessionDir: "/tmp/artifacts/sub-aaa", - activeSessionId: "child", - }); - }); - - it("prefers the parent's run status over the streaming heuristic", () => { - // An idle child session is still part of an active run; only the parent's - // run tracker knows that. - const parent = makeState({ - activeSessionId: "parent", - sessionFile: "/tmp/parent.jsonl", - childRunStatuses: { "sub-aaa": "running" }, - }); - const idleChild = makeState({ - activeSessionId: "child", - isStreaming: false, - metadata: { - kind: "subagent", - createdAt: 1, - parentActiveSessionId: "parent", - rlmChildId: "sub-aaa", - rlmParentNodeId: "sub-aaa", - prompt: "Slow task", - sessionDir: "/tmp/artifacts/sub-aaa", - }, - }); - - const snapshots = buildRlmChildSnapshots("parent", [parent, idleChild]); - - expect(snapshots.map((snapshot) => [snapshot.id, snapshot.status])).toEqual([["sub-aaa", "running"]]); - }); - - it("keeps terminal run status while projecting a retained child's active follow-up", () => { + it("uses the AgentSession projection and adds resident active session ids", () => { + const queued = { + id: "sub-queued", + label: "Queued task", + status: "queued" as const, + sessionDir: "/tmp/artifacts/sub-queued", + }; + const executing = { + id: "sub-running", + label: "Running task", + status: "running" as const, + sessionDir: "/tmp/artifacts/sub-running", + activity: { kind: "executing" as const, toolName: "ipython" }, + }; const parent = makeState({ activeSessionId: "parent", - childRunStatuses: { "sub-aaa": "done" }, + childSnapshots: [queued, executing], }); - const activeRetainedChild = makeState({ - activeSessionId: "child", - isStreaming: true, + const residentChild = makeState({ + activeSessionId: "running-child", metadata: { kind: "subagent", createdAt: 1, parentActiveSessionId: "parent", - rlmChildId: "sub-aaa", + rlmChildId: "sub-running", }, }); - expect(buildRlmChildSnapshots("parent", [parent, activeRetainedChild])[0]).toMatchObject({ - status: "done", - activity: { kind: "writing" }, - }); - }); - - it("includes in-flight assistant output in child snapshots", () => { - const parent = makeState({ activeSessionId: "parent" }); - const child = makeState({ - activeSessionId: "child", - isStreaming: true, - metadata: { - kind: "subagent", - createdAt: 1, - parentActiveSessionId: "parent", - rlmChildId: "sub-aaa", - }, - streamingMessage: { - role: "assistant", - content: [ - { type: "text", text: "Still investigating" }, - { type: "toolCall", id: "tool-1", name: "search", arguments: {} }, - ], - } as AgentMessage, - }); - - expect(buildRlmChildSnapshots("parent", [parent, child])[0]).toMatchObject({ - answerPreview: "Still investigating", - toolUseCount: 1, - }); + expect(buildRlmChildSnapshots("parent", [parent, residentChild])).toEqual([ + { ...queued, activeSessionId: undefined }, + { ...executing, activeSessionId: "running-child" }, + ]); }); - it("returns no snapshots for sessions without children", () => { - const solo = makeState({ activeSessionId: "solo" }); - expect(buildRlmChildSnapshots("solo", [solo])).toEqual([]); + it("returns no snapshots when the root is not resident", () => { + expect(buildRlmChildSnapshots("missing", [])).toEqual([]); }); }); @@ -691,12 +584,12 @@ interface StateOptions { messages?: AgentMessage[]; hasUserContent?: boolean; summaryState?: ActiveSessionState["summaryState"]; - childRunStatuses?: Record; hasRunningRlmChildren?: boolean; hasAcceptedPromptInFlight?: boolean; unfinishedActionCount?: number; contextTokens?: number; streamingMessage?: AgentMessage; + childSnapshots?: RlmChildAgentSnapshot[]; rlmDepth?: number; metadata?: { kind: "top-level" | "subagent"; @@ -741,7 +634,7 @@ function makeState(options: StateOptions): ActiveSessionState { hasUserContent: () => options.hasUserContent ?? false, }, messages: options.messages ?? ([] as AgentMessage[]), - getRlmChildRunStatus: (childId: string) => options.childRunStatuses?.[childId], + getRlmChildSnapshots: () => options.childSnapshots ?? [], hasRunningRlmChildren: () => options.hasRunningRlmChildren ?? false, hasAcceptedPromptInFlight: options.hasAcceptedPromptInFlight ?? false, unfinishedActionCount: options.unfinishedActionCount ?? (options.hasAcceptedPromptInFlight ? 1 : 0), diff --git a/packages/coding-agent/test/daemon-supervisor-lazy-subagents.test.ts b/packages/coding-agent/test/daemon-supervisor-lazy-subagents.test.ts index 3233a551da..74332846c2 100644 --- a/packages/coding-agent/test/daemon-supervisor-lazy-subagents.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-lazy-subagents.test.ts @@ -22,7 +22,7 @@ interface SupervisorInternals { findSummaryInWorker(worker: WorkerFixture, selector: string): SessionSummary | undefined; createOrReuseWorker( clientId: string, - command: { type: "create"; name?: string; sessionPath?: string }, + command: { type: "create"; name?: string; sessionPath?: string; lifecycle?: "client_owned" }, ): Promise; assertSupervisorSavedSessionNameAvailable(sessionPath: string, name: string): Promise; assertSavedSiblingNameAvailable( @@ -37,13 +37,13 @@ interface SupervisorInternals { interface WorkerFixture { descriptor: { workerId: string; - lifecycle: "ready"; + lifecycle: "ready" | "starting"; rootActiveSessionId: string; rootSessionId: string; pid: number; authenticationToken: string; ownerClientId?: string; - createCommand: { config: { cwd: string } }; + createCommand: { config: { cwd: string }; sessionPath?: string }; }; client: { request: ReturnType; @@ -342,12 +342,118 @@ describe("daemon supervisor passive subagent topology", () => { }); const first = supervisor.createOrReuseWorker("client", { type: "create", name: "named", sessionPath }); + const starting = worker("starting"); + starting.descriptor.lifecycle = "starting"; + starting.descriptor.createCommand = { config: { cwd: "/tmp/project" }, sessionPath }; + supervisor.workers.set(starting.descriptor.workerId, starting); const second = supervisor.createOrReuseWorker("client", { type: "create", sessionPath }); releaseSiblings(); expect(await Promise.all([first, second])).toEqual([resident, resident]); expect(launchWorker).toHaveBeenCalledOnce(); }); + it("enforces session ownership when joining an in-flight open", async () => { + const directory = mkdtempSync(join(tmpdir(), "prime-supervisor-pending-owner-")); + tempDirs.push(directory); + const sessionPath = join(directory, "session.jsonl"); + let releaseLaunch!: () => void; + const launchGate = new Promise((resolve) => { + releaseLaunch = resolve; + }); + const supervisor = new DaemonSupervisor(join(directory, "daemon.sock"), { + defaultSessionConfig: { agentDir: directory, cwd: directory }, + descriptorDir: join(directory, "workers"), + }) as unknown as SupervisorInternals; + const resident = worker("opened"); + resident.descriptor.ownerClientId = "owner"; + const launchWorker = vi.fn(async () => { + await launchGate; + return resident; + }); + Object.assign(supervisor, { launchWorker }); + + const create = { type: "create" as const, sessionPath, lifecycle: "client_owned" as const }; + const first = supervisor.createOrReuseWorker("owner", create); + const sameOwner = supervisor.createOrReuseWorker("owner", create); + const otherClient = supervisor.createOrReuseWorker("intruder", create); + const expectations = Promise.all([ + expect(first).resolves.toBe(resident), + expect(sameOwner).resolves.toBe(resident), + expect(otherClient).rejects.toMatchObject({ code: "session_already_active" }), + ]); + releaseLaunch(); + await expectations; + expect(launchWorker).toHaveBeenCalledOnce(); + }); + + it("rejoins an open registered while reclaiming a stale worker registration", async () => { + const directory = mkdtempSync(join(tmpdir(), "prime-supervisor-reclaim-rejoin-")); + tempDirs.push(directory); + const sessionPath = join(directory, "session.jsonl"); + let releaseReclaim!: () => void; + const reclaimGate = new Promise((resolve) => { + releaseReclaim = resolve; + }); + const supervisor = new DaemonSupervisor(join(directory, "daemon.sock"), { + defaultSessionConfig: { agentDir: directory, cwd: directory }, + descriptorDir: join(directory, "workers"), + }) as unknown as SupervisorInternals; + const stale = worker("stale"); + stale.descriptor.createCommand = { config: { cwd: directory }, sessionPath }; + supervisor.workers.set(stale.descriptor.workerId, stale); + const resident = worker("opened"); + resident.descriptor.ownerClientId = "owner"; + const launchWorker = vi.fn(async () => resident); + const reclaimStaleWorkerRegistration = vi.fn(async () => { + await reclaimGate; + supervisor.workers.delete(stale.descriptor.workerId); + return true; + }); + Object.assign(supervisor, { launchWorker, reclaimStaleWorkerRegistration }); + + const create = { type: "create" as const, sessionPath, lifecycle: "client_owned" as const }; + const first = supervisor.createOrReuseWorker("owner", create); + const second = supervisor.createOrReuseWorker("owner", create); + const intruder = supervisor.createOrReuseWorker("intruder", create); + const expectations = Promise.all([ + expect(first).resolves.toBe(resident), + expect(second).resolves.toBe(resident), + expect(intruder).rejects.toMatchObject({ code: "session_already_active" }), + ]); + releaseReclaim(); + await expectations; + expect(launchWorker).toHaveBeenCalledOnce(); + }); + + it("propagates an in-flight open failure to joiners", async () => { + const directory = mkdtempSync(join(tmpdir(), "prime-supervisor-pending-failure-")); + tempDirs.push(directory); + const sessionPath = join(directory, "session.jsonl"); + let releaseLaunch!: () => void; + const launchGate = new Promise((resolve) => { + releaseLaunch = resolve; + }); + const supervisor = new DaemonSupervisor(join(directory, "daemon.sock"), { + defaultSessionConfig: { agentDir: directory, cwd: directory }, + descriptorDir: join(directory, "workers"), + }) as unknown as SupervisorInternals; + const launchWorker = vi.fn(async () => { + await launchGate; + throw new Error("launch exploded"); + }); + Object.assign(supervisor, { launchWorker }); + + const create = { type: "create" as const, sessionPath, lifecycle: "client_owned" as const }; + const first = supervisor.createOrReuseWorker("owner", create); + const joiner = supervisor.createOrReuseWorker("intruder", create); + const expectations = Promise.all([ + expect(first).rejects.toThrow("launch exploded"), + expect(joiner).rejects.toThrow("launch exploded"), + ]); + releaseLaunch(); + await expectations; + }); + it("uses injective structural session name reservation keys", () => { expect(sessionNameReservationKey({ name: "b:c", depth: 1, parentSessionPath: "/a" })).not.toBe( sessionNameReservationKey({ name: "c", depth: 1, parentSessionPath: "/a:b" }), diff --git a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts index 7a58e27411..6fa13ce0ed 100644 --- a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts @@ -1613,7 +1613,7 @@ describe("daemon worker supervisor monitoring", () => { expect(supervisor.persistWorker).toHaveBeenCalledWith(worker); }); - it("rejects create reuse when a failed worker cannot be safely reclaimed", () => { + it("rejects create reuse when a failed worker cannot be safely reclaimed", async () => { const worker = { descriptor: { workerId: "failed-unreclaimed", @@ -1622,10 +1622,173 @@ describe("daemon worker supervisor monitoring", () => { }, }; const supervisor = Object.create(DaemonSupervisor.prototype) as { - reuseWorkerForCreate(target: typeof worker, ownerClientId: undefined, sessionPath: string): typeof worker; + reuseWorkerForCreate( + target: typeof worker, + ownerClientId: undefined, + sessionPath: string, + ): Promise; }; - expect(() => supervisor.reuseWorkerForCreate(worker, undefined, "/tmp/failed.jsonl")).toThrow(/failed worker/); + await expect(supervisor.reuseWorkerForCreate(worker, undefined, "/tmp/failed.jsonl")).rejects.toThrow( + /failed worker/, + ); + }); + + it("waits for worker recovery before reusing a saved session", async () => { + const root = { id: "active-root", activeSessionId: "active-root", sessionId: "session-root", cwd: "/tmp" }; + const recovery = createDeferred(); + const worker = { + descriptor: { + workerId: "recovering-worker", + rootActiveSessionId: root.activeSessionId, + lifecycle: "recovering", + }, + client: undefined as object | undefined, + summaries: new Map(), + recovery: recovery.promise, + intentionalStop: false, + }; + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers: new Map([[worker.descriptor.workerId, worker]]), + }) as { + reuseWorkerForCreate( + target: typeof worker, + ownerClientId: undefined, + sessionPath: string, + ): Promise; + }; + + let settled = false; + const reused = supervisor.reuseWorkerForCreate(worker, undefined, "/tmp/session.jsonl").finally(() => { + settled = true; + }); + await Promise.resolve(); + expect(settled).toBe(false); + + worker.descriptor.lifecycle = "ready"; + worker.client = {}; + worker.summaries.set(root.activeSessionId, root as SessionSummary); + recovery.resolve(); + + await expect(reused).resolves.toBe(worker); + }); + + it("starts recovery before reusing a persisted recovering worker", async () => { + const root = { id: "active-root", activeSessionId: "active-root", sessionId: "session-root", cwd: "/tmp" }; + const worker = { + descriptor: { + workerId: "persisted-recovering-worker", + rootActiveSessionId: root.activeSessionId, + lifecycle: "recovering", + }, + client: undefined as object | undefined, + summaries: new Map(), + intentionalStop: false, + }; + const recoverWorker = vi.fn(async () => { + worker.descriptor.lifecycle = "ready"; + worker.client = {}; + worker.summaries.set(root.activeSessionId, root as SessionSummary); + }); + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers: new Map([[worker.descriptor.workerId, worker]]), + shuttingDown: false, + recoverWorker, + }) as { + reuseWorkerForCreate( + target: typeof worker, + ownerClientId: undefined, + sessionPath: string, + ): Promise; + }; + + await expect(supervisor.reuseWorkerForCreate(worker, undefined, "/tmp/session.jsonl")).resolves.toBe(worker); + expect(recoverWorker).toHaveBeenCalledOnce(); + }); + + it("starts recovery for a disconnected worker still marked ready", async () => { + const root = { id: "active-root", activeSessionId: "active-root", sessionId: "session-root", cwd: "/tmp" }; + const worker = { + descriptor: { + workerId: "disconnected-ready-worker", + rootActiveSessionId: root.activeSessionId, + lifecycle: "ready", + }, + client: undefined as object | undefined, + summaries: new Map([[root.activeSessionId, root as SessionSummary]]), + intentionalStop: false, + }; + const recoverWorker = vi.fn(async () => { + worker.client = {}; + }); + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers: new Map([[worker.descriptor.workerId, worker]]), + shuttingDown: false, + recoverWorker, + }) as { + reuseWorkerForCreate( + target: typeof worker, + ownerClientId: undefined, + sessionPath: string, + ): Promise; + }; + + await expect(supervisor.reuseWorkerForCreate(worker, undefined, "/tmp/session.jsonl")).resolves.toBe(worker); + expect(recoverWorker).toHaveBeenCalledOnce(); + }); + + it("rejects recovered workers whose assigned root is still missing", async () => { + const worker = { + descriptor: { + workerId: "rootless-worker", + rootActiveSessionId: "active-root", + lifecycle: "ready", + }, + client: {}, + summaries: new Map(), + recovery: Promise.resolve(), + intentionalStop: false, + }; + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers: new Map([[worker.descriptor.workerId, worker]]), + }) as { + reuseWorkerForCreate( + target: typeof worker, + ownerClientId: undefined, + sessionPath: string, + ): Promise; + }; + + await expect(supervisor.reuseWorkerForCreate(worker, undefined, "/tmp/session.jsonl")).rejects.toThrow( + "assigned root session is missing", + ); + }); + + it("preserves cached summaries when recovery omits the assigned root", async () => { + const root = { id: "active-root", activeSessionId: "active-root", sessionId: "session-root", cwd: "/tmp" }; + const worker = { + descriptor: { + workerId: "root-omitting-worker", + rootActiveSessionId: root.activeSessionId, + }, + client: { + request: vi.fn(async () => + success(undefined, "list", { + sessions: [{ id: "other", activeSessionId: "other", sessionId: "session-other", cwd: "/tmp" }], + }), + ), + }, + summaries: new Map([[root.activeSessionId, root as SessionSummary]]), + intentionalStop: false, + }; + const supervisor = Object.create(DaemonSupervisor.prototype) as { + refreshWorkerSummaries(target: typeof worker, recovery: boolean): Promise; + }; + + await expect(supervisor.refreshWorkerSummaries(worker, true)).rejects.toThrow( + "Session worker omitted its root session during recovery", + ); + expect(worker.summaries.get(root.activeSessionId)).toBe(root); }); it("ignores conflicting paths on workers unrelated to a session lookup", () => { @@ -3591,10 +3754,18 @@ describe("daemon worker supervisor monitoring", () => { it("subscribes to worker updates with chunked snapshots", async () => { type SubscriptionWorker = { - client: { requestWorker: (command: unknown) => Promise<{ success: boolean }> }; + client: { + requestWorker: (command: unknown) => Promise<{ success: boolean }>; + supportsServerCapability?: (capability: string) => boolean; + }; }; const requestWorker = vi.fn(async () => ({ success: true })); - const worker: SubscriptionWorker = { client: { requestWorker } }; + const worker = { + client: { + requestWorker, + supportsServerCapability: () => true, + }, + }; const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { clients: new Set(), }) as { @@ -3610,13 +3781,44 @@ describe("daemon worker supervisor monitoring", () => { "attach_snapshot", "event_sequence", "slim_attach", - "chunked_snapshot", "correlated_prompt_lifecycle_v1", + "chunked_snapshot", + "immutable_snapshot_transfer_v1", ], supportsExtensionUi: false, }); }); + it("uses a full snapshot subscription for a worker without the new transfer capability", async () => { + type SubscriptionWorker = { + client: { + requestWorker: (command: unknown) => Promise<{ success: boolean }>; + supportsServerCapability?: (capability: string) => boolean; + }; + }; + const requestWorker = vi.fn(async () => ({ success: true })); + const worker = { + client: { + requestWorker, + supportsServerCapability: () => false, + }, + }; + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + clients: new Set(), + }) as { + subscribeWorker(worker: SubscriptionWorker, activeSessionId: string): Promise; + }; + + await supervisor.subscribeWorker(worker, "active-1"); + + expect(requestWorker).toHaveBeenCalledWith({ + type: "worker_subscribe", + activeSessionId: "active-1", + capabilities: ["attach_snapshot", "event_sequence", "slim_attach", "correlated_prompt_lifecycle_v1"], + supportsExtensionUi: false, + }); + }); + it("does not retain an attachment when snapshot loading fails", async () => { type AttachClient = { id: string; diff --git a/packages/coding-agent/test/session-worker-private-framing.test.ts b/packages/coding-agent/test/session-worker-private-framing.test.ts index ebb6116f3a..e5c22e7486 100644 --- a/packages/coding-agent/test/session-worker-private-framing.test.ts +++ b/packages/coding-agent/test/session-worker-private-framing.test.ts @@ -1,7 +1,11 @@ +import { spawn } from "node:child_process"; +import { once } from "node:events"; +import { createServer, type Socket } from "node:net"; import { PassThrough } from "node:stream"; import { describe, expect, it } from "vitest"; import { encodePrivateFrame, + encodePrivateFrameParts, PrivateFrameDecoder, PrivateFramedChannel, type PrivateFrameHeaderValidator, @@ -42,6 +46,118 @@ describe("private worker framing", () => { ]); }); + it("keeps segmented new-worker payloads byte-compatible with the legacy decoder", () => { + const header = { type: "snapshot", requestId: "compat" }; + const payloadParts = [Buffer.from('{"messages":['), Buffer.from('{"content":"ok"}'), Buffer.from("]}")]; + const segmented = encodePrivateFrameParts(header, payloadParts); + const legacyWire = encodePrivateFrame(header, Buffer.concat(payloadParts)); + expect(Buffer.concat(segmented)).toEqual(legacyWire); + + const decoder = new PrivateFrameDecoder(isTestHeader); + const frames = segmented.flatMap((part) => decoder.push(part)); + decoder.finish(); + expect(frames).toEqual([{ header, payload: Buffer.concat(payloadParts) }]); + }); + + it("crosses a real process boundary with new segmented writes and legacy contiguous replies", async () => { + const server = createServer(); + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", () => resolve()); + }); + const address = server.address(); + if (!address || typeof address === "string") throw new Error("expected TCP test address"); + const connection = new Promise((resolve) => server.once("connection", resolve)); + const childScript = ` +const net = require("node:net"); +let buffered = Buffer.alloc(0); +const socket = net.createConnection({ host: "127.0.0.1", port: Number(process.argv[1]) }); +const encodeLegacy = (header, payload) => { + const encodedHeader = Buffer.from(JSON.stringify(header)); + const frame = Buffer.allocUnsafe(8 + encodedHeader.length + payload.length); + frame.writeUInt32BE(encodedHeader.length, 0); + frame.writeUInt32BE(payload.length, 4); + encodedHeader.copy(frame, 8); + payload.copy(frame, 8 + encodedHeader.length); + return frame; +}; +socket.on("data", (chunk) => { + buffered = buffered.length === 0 ? chunk : Buffer.concat([buffered, chunk]); + const replies = []; + while (buffered.length >= 8) { + const headerLength = buffered.readUInt32BE(0); + const payloadLength = buffered.readUInt32BE(4); + const frameLength = 8 + headerLength + payloadLength; + if (buffered.length < frameLength) break; + const header = JSON.parse(buffered.toString("utf8", 8, 8 + headerLength)); + const payload = Buffer.from(buffered.subarray(8 + headerLength, frameLength)); + replies.push(encodeLegacy({ type: "ack", requestId: header.requestId }, payload)); + buffered = Buffer.from(buffered.subarray(frameLength)); + } + if (replies.length > 0) socket.write(Buffer.concat(replies)); + if (replies.some((reply) => JSON.parse(reply.toString("utf8", 8, 8 + reply.readUInt32BE(0))).requestId === "three")) { + socket.end(); + } +});`; + const child = spawn(process.execPath, ["--eval", childScript, String(address.port)], { + stdio: ["ignore", "ignore", "pipe"], + }); + let socket: Socket | undefined; + try { + socket = await Promise.race([ + connection, + new Promise((_, reject) => + setTimeout(() => reject(new Error("legacy framing child did not connect")), 3000), + ), + ]); + const decoder = new PrivateFrameDecoder(isTestHeader); + const replies: Array<{ header: TestHeader; payload: Buffer }> = []; + const completed = new Promise((resolve, reject) => { + socket!.on("data", (chunk) => replies.push(...decoder.push(chunk))); + socket!.once("error", reject); + socket!.once("end", () => { + try { + decoder.finish(); + resolve(); + } catch (error) { + reject(error); + } + }); + }); + for (const [requestId, payloadParts] of [ + ["one", [Buffer.from("a"), Buffer.from("b")]], + ["two", [Buffer.from("c"), Buffer.from("d")]], + ["three", [Buffer.from("e"), Buffer.from("f")]], + ] as const) { + for (const part of encodePrivateFrameParts({ type: "snapshot", requestId }, payloadParts)) { + socket.write(part); + } + } + await Promise.race([ + completed, + new Promise((_, reject) => + setTimeout(() => reject(new Error("legacy framing child did not finish")), 3000), + ), + ]); + expect(replies).toEqual([ + { header: { type: "ack", requestId: "one" }, payload: Buffer.from("ab") }, + { header: { type: "ack", requestId: "two" }, payload: Buffer.from("cd") }, + { header: { type: "ack", requestId: "three" }, payload: Buffer.from("ef") }, + ]); + const [exitCode] = (await Promise.race([ + once(child, "close"), + new Promise((_, reject) => + setTimeout(() => reject(new Error("legacy framing child did not close")), 3000), + ), + ])) as [number | null]; + expect(exitCode).toBe(0); + } finally { + socket?.destroy(); + server.close(); + if (child.exitCode === null && child.signalCode === null) child.kill("SIGTERM"); + } + }); + it("rejects invalid lengths, JSON, and routing headers", () => { const oversized = Buffer.alloc(8); oversized.writeUInt32BE(1025, 0); diff --git a/packages/coding-agent/test/snapshot-transcript-cache.test.ts b/packages/coding-agent/test/snapshot-transcript-cache.test.ts index 287ee5301c..05a16a75bc 100644 --- a/packages/coding-agent/test/snapshot-transcript-cache.test.ts +++ b/packages/coding-agent/test/snapshot-transcript-cache.test.ts @@ -1,9 +1,19 @@ -import { mkdtempSync, rmSync } from "node:fs"; +import { spawnSync } from "node:child_process"; +import { chmodSync, existsSync, lstatSync, mkdirSync, mkdtempSync, readdirSync, rmSync, symlinkSync } from "node:fs"; +import { appendFile, mkdir, readFile, rm, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; +import { performance } from "node:perf_hooks"; import type { AgentMessage } from "@earendil-works/pi-agent-core"; import { afterEach, describe, expect, it } from "vitest"; -import { SnapshotTranscriptCache } from "../src/modes/daemon/snapshot-transcript-cache.js"; +import { + createSnapshotCacheProcessRoot, + prepareSnapshotTranscriptCache, + prepareSnapshotTranscriptPayload, + SnapshotTranscriptCache, + type SnapshotTranscriptCacheIo, + sweepAbandonedSnapshotCacheRoots, +} from "../src/modes/daemon/snapshot-transcript-cache.js"; const tempDirs: string[] = []; @@ -116,6 +126,89 @@ describe("snapshot transcript cache", () => { expect(cache.complete).toBe(false); }); + it("cleans a partial file-backed cache when message encoding throws", () => { + const cacheRoot = tempDir(); + const circular: { role: "user"; content: unknown; timestamp: number } = { + role: "user", + content: undefined, + timestamp: 2, + }; + circular.content = circular; + const invalidMessages = [ + { role: "user", content: "x".repeat(5 * 1024 * 1024), timestamp: 1 }, + circular, + ] as unknown as AgentMessage[]; + + expect( + () => + new SnapshotTranscriptCache({ + activeSessionId: "active-partial", + snapshotId: "snapshot-partial", + messages: invalidMessages, + cacheRoot, + memoryCacheBytes: 1, + }), + ).toThrow(); + expect(readdirSync(cacheRoot)).toEqual([]); + }); + + it("sweeps only abandoned process-owned cache roots", () => { + const parent = tempDir(); + const liveRoot = createSnapshotCacheProcessRoot(parent); + const liveName = liveRoot.slice(parent.length + 1); + const staleName = liveName.replace(/^owner-\d+-/, "owner-987654321-"); + const staleRoot = join(parent, staleName); + mkdirSync(staleRoot, { mode: 0o700 }); + + sweepAbandonedSnapshotCacheRoots(parent); + + expect(existsSync(liveRoot)).toBe(true); + expect(existsSync(staleRoot)).toBe(false); + }); + + it("does not allocate storage for an already-aborted preparation", async () => { + const cacheRoot = tempDir(); + const controller = new AbortController(); + controller.abort(); + + await expect( + prepareSnapshotTranscriptCache({ + activeSessionId: "active-aborted", + snapshotId: "snapshot-aborted", + messages: messages(2, 5 * 1024 * 1024), + cacheRoot, + memoryCacheBytes: 1, + signal: controller.signal, + }), + ).rejects.toThrow(); + expect(readdirSync(cacheRoot)).toEqual([]); + }); + + it("isolates concurrent file-backed preparations that use the same transfer id", async () => { + const cacheRoot = tempDir(); + const prepared = await Promise.all([ + prepareSnapshotTranscriptCache({ + activeSessionId: "active-concurrent", + snapshotId: "snapshot-concurrent", + messages: messages(3, 2 * 1024 * 1024), + cacheRoot, + memoryCacheBytes: 1, + }), + prepareSnapshotTranscriptCache({ + activeSessionId: "active-concurrent", + snapshotId: "snapshot-concurrent", + messages: messages(3, 2 * 1024 * 1024), + cacheRoot, + memoryCacheBytes: 1, + }), + ]); + + expect(readdirSync(cacheRoot)).toHaveLength(2); + expect(prepared.every((cache) => cache.fileBacked && cache.complete)).toBe(true); + for (const cache of prepared) cache.dispose(); + expect(readdirSync(cacheRoot)).toEqual([]); + }); + it("completes repeated endings idempotently", () => { const cache = new SnapshotTranscriptCache({ activeSessionId: "active-f", @@ -131,4 +224,162 @@ describe("snapshot transcript cache", () => { expect(() => cache.appendEncodedChunk(Buffer.from("stale"))).toThrow("is not writable"); cache.dispose(); }); + + it.each(["append", "nested mutation"] as const)( + "rejects a live transcript %s across preparation yields", + async (mutation) => { + const cacheRoot = tempDir(); + const source = messages(1, 3 * 1024 * 1024); + const preparing = prepareSnapshotTranscriptPayload({ messages: source, cacheRoot }); + setImmediate(() => { + if (mutation === "append") source.push({ role: "user", content: "late", timestamp: 2 }); + else (source[0] as { content: unknown }).content = "changed in place"; + }); + + await expect(preparing).rejects.toThrow("generation changed"); + expect(readdirSync(cacheRoot)).toEqual([]); + }, + ); + + it("cleans an asynchronously spilled partial payload after a later circular message", async () => { + const cacheRoot = tempDir(); + const circular: Record = { role: "user", timestamp: 3 }; + circular.content = circular; + const source = [ + { role: "user", content: "a".repeat(5 * 1024 * 1024), timestamp: 1 }, + { role: "user", content: "b".repeat(5 * 1024 * 1024), timestamp: 2 }, + circular, + ] as unknown as AgentMessage[]; + + await expect( + prepareSnapshotTranscriptPayload({ messages: source, cacheRoot, memoryCacheBytes: 1 }), + ).rejects.toThrow("circular"); + expect(readdirSync(cacheRoot)).toEqual([]); + }); + + it("preserves async write failures when cleanup also fails and clears state after remove failures", async () => { + const cacheRoot = tempDir(); + const writeFailure = new Error("injected snapshot write failure"); + const removeFailure = new Error("injected snapshot remove failure"); + let writeCount = 0; + const failingIo: SnapshotTranscriptCacheIo = { + mkdir, + async writeFile(path, data, options) { + if (++writeCount === 2) throw writeFailure; + await writeFile(path, data, options); + }, + appendFile, + readFile, + async rm() { + throw removeFailure; + }, + }; + await expect( + prepareSnapshotTranscriptPayload({ + messages: messages(3, 2 * 1024 * 1024), + cacheRoot, + memoryCacheBytes: 1, + io: failingIo, + }), + ).rejects.toBe(writeFailure); + + const removable = await prepareSnapshotTranscriptPayload({ + messages: messages(2, 3 * 1024 * 1024), + cacheRoot, + memoryCacheBytes: 1, + io: { ...failingIo, writeFile, rm: async () => Promise.reject(removeFailure) }, + }); + await expect(removable.disposeAsync()).rejects.toBe(removeFailure); + expect(removable.fileBacked).toBe(false); + expect(removable.retainedPayloadBytes).toBe(0); + }); + + it("surfaces injected async read failures without losing payload ownership", async () => { + const cacheRoot = tempDir(); + const readFailure = new Error("injected snapshot read failure"); + const payload = await prepareSnapshotTranscriptPayload({ + messages: messages(2, 3 * 1024 * 1024), + cacheRoot, + memoryCacheBytes: 1, + io: { mkdir, writeFile, appendFile, readFile: async () => Promise.reject(readFailure), rm }, + }); + const transfer = payload.createTransfer("active-read", "snapshot-read") as AsyncIterable; + await expect(transfer[Symbol.asyncIterator]().next()).rejects.toBe(readFailure); + payload.dispose(); + }); + + it("uses a private crash-recoverable process root and repairs parent permissions", () => { + const parent = tempDir(); + chmodSync(parent, 0o777); + const modulePath = join(process.cwd(), "src/modes/daemon/snapshot-transcript-cache.ts"); + const child = spawnSync( + "npx", + [ + "tsx", + "--eval", + `import { createSnapshotCacheProcessRoot } from ${JSON.stringify(modulePath)}; console.log(createSnapshotCacheProcessRoot(${JSON.stringify(parent)}));`, + ], + { cwd: process.cwd(), encoding: "utf8" }, + ); + expect(child.status, child.stderr).toBe(0); + const crashedRoot = child.stdout.trim().split("\n").at(-1)!; + expect(existsSync(crashedRoot)).toBe(true); + sweepAbandonedSnapshotCacheRoots(parent); + expect(existsSync(crashedRoot)).toBe(false); + expect(lstatSync(parent).mode & 0o077).toBe(0); + + const symlinkParent = join(tempDir(), "cache-parent-link"); + symlinkSync(parent, symlinkParent); + expect(() => createSnapshotCacheProcessRoot(symlinkParent)).toThrow("private directory"); + }); + + it("bounds 36 MiB single and multi-message preparation and shares one payload across three purposes", async () => { + const cacheRoot = tempDir(); + const run = async (source: AgentMessage[]) => { + let maxEventLoopDelayMs = 0; + let previousTick = performance.now(); + const monitor = setInterval(() => { + const current = performance.now(); + maxEventLoopDelayMs = Math.max(maxEventLoopDelayMs, current - previousTick - 5); + previousTick = current; + }, 5); + const rssBefore = process.memoryUsage().rss; + const started = performance.now(); + const payload = await prepareSnapshotTranscriptPayload({ messages: source, cacheRoot }); + await new Promise((resolve) => setTimeout(resolve, 10)); + clearInterval(monitor); + return { + payload, + elapsedMs: performance.now() - started, + maxEventLoopDelayMs, + rssDeltaBytes: process.memoryUsage().rss - rssBefore, + }; + }; + + const single = await run(messages(1, 36 * 1024 * 1024)); + expect(single.maxEventLoopDelayMs).toBeLessThan(50); + expect(single.rssDeltaBytes).toBeLessThan(160 * 1024 * 1024); + expect(single.payload.fileBacked).toBe(true); + const transfers = ["attach", "replacement", "resync"].map((purpose) => + single.payload.createTransfer("active-benchmark", `snapshot-${purpose}`), + ); + expect(single.payload.activeReaders).toBe(3); + const iterators = transfers.map((transfer) => + (transfer as AsyncIterable)[Symbol.asyncIterator](), + ); + const firstChunks = await Promise.all(iterators.map((iterator) => iterator.next())); + const payloadBuffers = firstChunks.map((entry) => entry.value![1]); + expect(payloadBuffers[1]).toBe(payloadBuffers[0]); + expect(payloadBuffers[2]).toBe(payloadBuffers[0]); + expect(single.payload.retainedPayloadBytes).toBeLessThanOrEqual(single.payload.bytes); + await Promise.all(iterators.map((iterator) => iterator.return?.())); + expect(single.payload.activeReaders).toBe(0); + single.payload.dispose(); + + const multiple = await run(messages(36, 1024 * 1024)); + expect(multiple.maxEventLoopDelayMs).toBeLessThan(50); + expect(multiple.rssDeltaBytes).toBeLessThan(192 * 1024 * 1024); + expect(multiple.payload.bytes).toBeGreaterThanOrEqual(36 * 1024 * 1024); + multiple.payload.dispose(); + }, 30_000); }); diff --git a/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts b/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts index 38423943a7..b517a58de3 100644 --- a/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts +++ b/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts @@ -1,9 +1,10 @@ -import { mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { existsSync, mkdirSync, mkdtempSync, readdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import type { Socket } from "node:net"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { PassThrough } from "node:stream"; import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import type { AssistantMessage } from "@earendil-works/pi-ai"; import { afterEach, describe, expect, it, vi } from "vitest"; import type { ActiveSessionState, DaemonSocketClient } from "../../../src/modes/daemon/active-session-state.js"; import { AgentDaemon, markClientSnapshotStreaming } from "../../../src/modes/daemon/daemon-mode.js"; @@ -113,6 +114,46 @@ describe("ENG-4601 worker snapshot cache", () => { expect([...createSnapshotTranscriptChunks({ activeSessionId, snapshotId, messages: [] })]).toEqual([]); }); + it("freezes nested transcript bytes when the selected message later mutates", () => { + const cacheRoot = tempDirectory(); + const message: AssistantMessage = { + role: "assistant", + content: [{ type: "text", text: "selected" }], + api: "openai-responses", + provider: "openai", + model: "gpt-4o-mini", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: 1, + }; + const chunks = createSnapshotTranscriptChunks({ + activeSessionId, + snapshotId, + messages: [message], + cacheRoot, + memoryCacheBytes: 1, + }); + const cacheDirectory = join(cacheRoot, readdirSync(cacheRoot)[0]!); + expect(existsSync(cacheDirectory)).toBe(true); + + const selected = message.content[0]; + if (selected?.type !== "text") throw new Error("expected text content"); + selected.text = "mutated"; + + const frozenBytes = Buffer.concat([...chunks]).toString("utf8"); + expect(frozenBytes).toContain('"text":"selected"'); + expect(frozenBytes).not.toContain('"text":"mutated"'); + chunks.dispose?.(); + expect(existsSync(cacheDirectory)).toBe(false); + }); + it.each([1, 2])( "keeps %i same-ID worker stream(s) bounded by socket drain without shared cleanup", async (count) => { @@ -222,16 +263,15 @@ describe("ENG-4601 worker snapshot cache", () => { return true; }; + const abortedTranscript = createSnapshotTranscriptChunks({ + activeSessionId, + snapshotId, + messages: messages("aborted"), + signal, + }); internals.detachClientFromSession(client, state); await Promise.all([ - internals.streamWorkerSnapshot( - client, - streamedResult(2), - createSnapshotTranscriptChunks({ activeSessionId, snapshotId, messages: messages("aborted"), signal }), - "attach", - signal, - true, - ), + internals.streamWorkerSnapshot(client, streamedResult(2), abortedTranscript, "attach", signal, true), internals.streamWorkerSnapshot( client, streamedResult(2, 1, siblingActiveSessionId, siblingSnapshotId), diff --git a/packages/coding-agent/test/suite/regressions/4602-snapshot-transfer-idempotency.test.ts b/packages/coding-agent/test/suite/regressions/4602-snapshot-transfer-idempotency.test.ts index deb503d15c..8fd0133616 100644 --- a/packages/coding-agent/test/suite/regressions/4602-snapshot-transfer-idempotency.test.ts +++ b/packages/coding-agent/test/suite/regressions/4602-snapshot-transfer-idempotency.test.ts @@ -16,7 +16,10 @@ import { type DaemonWorkerFrameHeader, isDaemonWorkerFrameHeader, } from "../../../src/modes/daemon/daemon-worker-protocol.js"; -import { SnapshotTranscriptCache } from "../../../src/modes/daemon/snapshot-transcript-cache.js"; +import { + SnapshotTranscriptCache, + type SnapshotTranscriptChunkSource, +} from "../../../src/modes/daemon/snapshot-transcript-cache.js"; import { type PrivateFrame, PrivateFrameDecoder } from "../../../src/modes/session-worker/private-framing.js"; const activeSessionId = "active-4602"; @@ -63,7 +66,7 @@ function summary(): SessionSummary { }; } -function streamedResult(messages: AgentMessage[]): DaemonAttachResult { +function streamedResult(messages: AgentMessage[], resultSnapshotId = snapshotId): DaemonAttachResult { return { protocol: DAEMON_PROTOCOL_INFO, activeSessionId, @@ -82,7 +85,7 @@ function streamedResult(messages: AgentMessage[]): DaemonAttachResult { }, lastEventSequence: 1, lastEventCursor: { generation: "generation-4602", sequence: 1 }, - snapshotStream: { id: snapshotId, messageCount: messages.length, targetChunkBytes: 512 * 1024 }, + snapshotStream: { id: resultSnapshotId, messageCount: messages.length, targetChunkBytes: 512 * 1024 }, client: { id: "worker", capabilities: ["chunked_snapshot"] }, }; } @@ -231,6 +234,92 @@ describe("ENG-4602 snapshot transfer containment", () => { expect(unhandled).not.toHaveBeenCalled(); }); + it("uses a fresh transfer identity when transcript bytes change at the same event cursor", async () => { + const daemon = new AgentDaemon("/tmp/eng-4602-transfer-identity.sock", { + defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, + createRuntime: async () => { + throw new Error("unexpected runtime creation"); + }, + }); + const state = { + activeSessionId, + clients: new Set(), + pendingAttaches: 0, + eventGeneration: "generation-4602", + lastEventSequence: 1, + runtime: { metadata: { kind: "top-level", createdAt: 1 } }, + } as unknown as ActiveSessionState; + const socket = new PassThrough(); + const client = { + id: "supervisor", + socket: socket as unknown as Socket, + transport: "private-framed", + attachedActiveSessionIds: new Set(), + detachInput: () => {}, + supportsExtensionUi: false, + capabilities: new Set(), + } as DaemonSocketClient; + const firstMessage = { role: "user", content: "before", timestamp: 1 } as const; + const secondMessage = { role: "user", content: "after", timestamp: 1 } as const; + const createAttachResult = vi + .fn<() => DaemonAttachResult>() + .mockReturnValueOnce(streamedResult([firstMessage])) + .mockReturnValueOnce(streamedResult([secondMessage])); + const streamWorkerSnapshot = vi.fn( + async ( + _client: DaemonSocketClient, + _result: DaemonAttachResult, + _transcript: SnapshotTranscriptChunkSource, + ) => {}, + ); + const internals = daemon as unknown as { + sessions: Map; + createAttachResult: typeof createAttachResult; + streamWorkerSnapshot: typeof streamWorkerSnapshot; + handleCommand(client: DaemonSocketClient, command: DaemonCommand): Promise; + }; + internals.sessions.set(activeSessionId, state); + internals.createAttachResult = createAttachResult; + internals.streamWorkerSnapshot = streamWorkerSnapshot; + const attach = { + type: "attach", + activeSessionId, + capabilities: ["attach_snapshot", "event_sequence", "slim_attach", "chunked_snapshot"], + } as const; + + const firstResponse = (await internals.handleCommand(client, attach)) as { + success: true; + data: DaemonAttachResult; + }; + const secondResponse = (await internals.handleCommand(client, attach)) as { + success: true; + data: DaemonAttachResult; + }; + await new Promise((resolve) => setImmediate(resolve)); + + expect(firstResponse.data.lastEventCursor).toEqual(secondResponse.data.lastEventCursor); + expect(firstResponse.data.snapshotStream?.id).not.toBe(secondResponse.data.snapshotStream?.id); + expect(streamWorkerSnapshot).toHaveBeenCalledTimes(2); + const transcriptContents = await Promise.all( + streamWorkerSnapshot.mock.calls.map(async ([, , transcript]) => { + const contents: unknown[] = []; + for await (const chunk of transcript) { + const wireChunk = Buffer.isBuffer(chunk) ? chunk : Buffer.concat(chunk); + const decoded = JSON.parse(wireChunk.toString("utf8")) as Extract< + DaemonOutbound, + { type: "session_snapshot_chunk" } + >; + contents.push( + ...decoded.messages.map((message) => (message.role === "user" ? message.content : undefined)), + ); + } + return contents; + }), + ); + expect(transcriptContents).toEqual([["before"], ["after"]]); + socket.destroy(); + }); + it("fails one worker snapshot without dropping another session on the supervisor channel", async () => { const daemon = new AgentDaemon("/tmp/eng-4602-stream.sock", { defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, @@ -292,6 +381,122 @@ describe("ENG-4602 snapshot transfer containment", () => { socket.destroy(); }); + it.each([ + { + name: "non-monotonic first chunk", + corrupt: (frames: ReturnType) => ({ ...frames.chunk, index: 1 }), + }, + { + name: "wrong fresh chunk snapshot id", + corrupt: (frames: ReturnType) => ({ + ...frames.chunk, + snapshotId: "payload-mismatch", + }), + }, + ] as const)("rejects a $name without closing the worker channel", ({ corrupt }) => { + const supervisor = new DaemonSupervisor("/tmp/eng-4602-supervisor-fresh-chunk.sock", { + defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, + descriptorDir: "/tmp/eng-4602-supervisor-fresh-chunk-state", + }); + const { close, worker } = workerHarness(); + const internals = supervisor as unknown as { + handleWorkerFrame(worker: WorkerHarness, frame: PrivateFrame): void; + }; + const frames = snapshotFrames([{ role: "user", content: "stable", timestamp: 1 }]); + internals.handleWorkerFrame(worker, frame(frames.begin)); + const corruptFrame = frame(corrupt(frames) as DaemonOutbound); + if (corruptFrame.header.kind !== "outbound") throw new Error("expected outbound frame"); + corruptFrame.header.snapshotId = snapshotId; + + internals.handleWorkerFrame(worker, corruptFrame); + + expect(close).not.toHaveBeenCalled(); + expect(worker.client).toBeDefined(); + expect(worker.snapshotCache.has(activeSessionId)).toBe(false); + expect(worker.transcriptCaches.has(activeSessionId)).toBe(false); + }); + + it.each([ + { + name: "chunk count", + corrupt: (end: ReturnType["end"]) => ({ ...end, chunkCount: 2 }), + }, + { + name: "event sequence", + corrupt: (end: ReturnType["end"]) => ({ ...end, lastEventSequence: 2 }), + }, + { + name: "event cursor", + corrupt: (end: ReturnType["end"]) => ({ + ...end, + lastEventCursor: { generation: "other", sequence: 1 }, + }), + }, + ] as const)("validates a fresh generation's $name before publishing completion", ({ corrupt }) => { + const supervisor = new DaemonSupervisor("/tmp/eng-4602-supervisor-fresh-end.sock", { + defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, + descriptorDir: "/tmp/eng-4602-supervisor-fresh-end-state", + }); + const { close, worker } = workerHarness(); + const internals = supervisor as unknown as { + handleWorkerFrame(worker: WorkerHarness, frame: PrivateFrame): void; + }; + const frames = snapshotFrames([{ role: "user", content: "stable", timestamp: 1 }]); + internals.handleWorkerFrame(worker, frame(frames.begin)); + internals.handleWorkerFrame(worker, frame(frames.chunk)); + const transcript = worker.transcriptCaches.get(activeSessionId); + if (!transcript) throw new Error("missing staged transcript"); + + internals.handleWorkerFrame(worker, frame(corrupt(frames.end) as DaemonOutbound)); + + expect(transcript.complete).toBe(false); + expect(close).not.toHaveBeenCalled(); + expect(worker.snapshotCache.has(activeSessionId)).toBe(false); + expect(worker.transcriptCaches.has(activeSessionId)).toBe(false); + }); + + it.each([ + { + name: "session identity", + corrupt: (begin: ReturnType["begin"]) => ({ + ...begin, + snapshot: { + ...begin.snapshot, + state: { ...begin.snapshot.state, sessionId: "wrong-session" }, + }, + }), + }, + { + name: "event progress", + corrupt: (begin: ReturnType["begin"]) => ({ + ...begin, + snapshot: { + ...begin.snapshot, + lastEventCursor: { generation: "generation-4602", sequence: 2 }, + }, + }), + }, + ] as const)("rejects invalid snapshot begin $name and ignores its late frames", ({ corrupt }) => { + const supervisor = new DaemonSupervisor("/tmp/eng-4602-supervisor-invalid-begin.sock", { + defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, + descriptorDir: "/tmp/eng-4602-supervisor-invalid-begin-state", + }); + const { close, worker } = workerHarness(); + const internals = supervisor as unknown as { + handleWorkerFrame(worker: WorkerHarness, frame: PrivateFrame): void; + }; + const frames = snapshotFrames([{ role: "user", content: "stable", timestamp: 1 }]); + + internals.handleWorkerFrame(worker, frame(corrupt(frames.begin) as DaemonOutbound)); + internals.handleWorkerFrame(worker, frame(frames.chunk)); + internals.handleWorkerFrame(worker, frame(frames.end)); + + expect(close).not.toHaveBeenCalled(); + expect(worker.snapshotCache.has(activeSessionId)).toBe(false); + expect(worker.transcriptCaches.has(activeSessionId)).toBe(false); + expect(worker.snapshotGenerations.has(activeSessionId)).toBe(false); + }); + it("keeps a multi-session worker connected after a scoped snapshot failure frame", () => { const supervisor = new DaemonSupervisor("/tmp/eng-4602-supervisor-failure.sock", { defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, @@ -333,12 +538,48 @@ describe("ENG-4602 snapshot transfer containment", () => { expect(worker.snapshotGenerations.has(activeSessionId)).toBe(false); }); + it("invalidates only the snapshot load associated with a stale generation", () => { + const supervisor = new DaemonSupervisor("/tmp/eng-4602-supervisor-overlap.sock", { + defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, + descriptorDir: "/tmp/eng-4602-supervisor-overlap-state", + }); + const staleKey = `${activeSessionId}:chunked`; + const freshKey = `${activeSessionId}:full`; + const worker = { + ...workerHarness().worker, + snapshotLoads: new Map([ + [staleKey, Promise.resolve(streamedResult([], "snapshot-stale"))], + [freshKey, Promise.resolve(streamedResult([], "snapshot-fresh"))], + ]), + snapshotLoadSnapshotIds: new Map([ + [staleKey, "snapshot-stale"], + [freshKey, "snapshot-fresh"], + ]), + }; + const internals = supervisor as unknown as { + failWorkerSnapshotCache( + worker: object, + activeSessionId: string, + error: Error, + closeWorkerChannel: boolean, + expectedSnapshotId: string, + ): void; + }; + + internals.failWorkerSnapshotCache(worker, activeSessionId, new Error("stale"), false, "snapshot-stale"); + + expect(worker.snapshotLoads.has(staleKey)).toBe(false); + expect(worker.snapshotLoads.has(freshKey)).toBe(true); + expect(worker.snapshotLoadSnapshotIds.has(staleKey)).toBe(false); + expect(worker.snapshotLoadSnapshotIds.get(freshKey)).toBe("snapshot-fresh"); + }); + it("quarantines a completed duplicate until transcript validation", async () => { const supervisor = new DaemonSupervisor("/tmp/eng-4602-supervisor.sock", { defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, descriptorDir: "/tmp/eng-4602-supervisor-state", }); - const { close, worker } = workerHarness(); + const { close, request, worker } = workerHarness(); const client = socketClient("public", new PassThrough()); const streamSnapshot = vi.fn(async () => {}); const internals = supervisor as unknown as { @@ -346,6 +587,7 @@ describe("ENG-4602 snapshot transfer containment", () => { workers: Map; syncWorkerExtensionUi: ReturnType; streamSnapshot: typeof streamSnapshot; + forwardToWorker(worker: WorkerHarness, command: DaemonCommand): Promise; handleWorkerFrame(worker: WorkerHarness, frame: PrivateFrame): void; }; internals.clients.add(client); @@ -398,6 +640,21 @@ describe("ENG-4602 snapshot transfer containment", () => { expect(streamSnapshot).toHaveBeenCalledOnce(); expect(close).not.toHaveBeenCalled(); streamSnapshot.mockClear(); + const replacementSnapshotId = "snapshot-4602-recovered"; + request.mockImplementation(async (command: { type: string }) => { + if (command.type === "attach") { + return { + type: "response", + command: "attach", + success: true, + data: streamedResult([], replacementSnapshotId), + }; + } + if (command.type === "prompt") { + return { type: "response", command: "prompt", success: true }; + } + throw new Error(`unexpected worker request: ${command.type}`); + }); internals.handleWorkerFrame(worker, frame(frames.begin, "replacement")); internals.handleWorkerFrame( @@ -410,11 +667,23 @@ describe("ENG-4602 snapshot transfer containment", () => { "replacement", ), ); + await new Promise((resolve) => setImmediate(resolve)); + expect(close).not.toHaveBeenCalled(); + expect(worker.client).toBeDefined(); + expect(worker.descriptor.lifecycle).toBe("ready"); expect(worker.transcriptCaches.has(activeSessionId)).toBe(false); expect(worker.snapshotCache.has(activeSessionId)).toBe(false); + expect(request.mock.calls.some(([command]) => command.type === "attach")).toBe(false); expect(streamSnapshot).not.toHaveBeenCalled(); - expect(close).toHaveBeenCalledOnce(); + await expect( + internals.forwardToWorker(worker, { + id: "prompt-after-snapshot-mismatch", + type: "prompt", + activeSessionId, + message: "still alive", + }), + ).resolves.toMatchObject({ id: "prompt-after-snapshot-mismatch", success: true }); }); it("holds catch-up behind duplicate validation and rejects it on mismatch", async () => { @@ -473,11 +742,13 @@ describe("ENG-4602 snapshot transfer containment", () => { ); await failedCatchup; - expect(request).not.toHaveBeenCalled(); + expect(request.mock.calls.some(([command]) => command.type === "attach")).toBe(false); expect(streamSnapshot).not.toHaveBeenCalled(); expect(worker.snapshotCache.has(activeSessionId)).toBe(false); expect(worker.transcriptCaches.has(activeSessionId)).toBe(false); - expect(close).toHaveBeenCalledOnce(); + expect(worker.client).toBeDefined(); + expect(worker.descriptor.lifecycle).toBe("ready"); + expect(close).not.toHaveBeenCalled(); }); it("rejects a quarantined catch-up before intentional worker stop", async () => { @@ -574,7 +845,7 @@ describe("ENG-4602 snapshot transfer containment", () => { expect(worker.snapshotGenerations.size).toBe(0); }); - it("rejects same-ID reentrant begins and mismatched duplicate metadata", async () => { + it("retires same-ID generation faults without recovering the worker", async () => { const supervisor = new DaemonSupervisor("/tmp/eng-4602-supervisor-invalid.sock", { defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, descriptorDir: "/tmp/eng-4602-supervisor-invalid-state", @@ -599,27 +870,29 @@ describe("ENG-4602 snapshot transfer containment", () => { internals.handleWorkerFrame(reentrant.worker, frame(frames.begin)); internals.handleWorkerFrame(reentrant.worker, frame(frames.begin)); await new Promise((resolve) => setImmediate(resolve)); - expect(reentrant.close).toHaveBeenCalledOnce(); + expect(reentrant.close).not.toHaveBeenCalled(); expect(reentrant.worker.transcriptCaches.has(activeSessionId)).toBe(false); - expect(reentrant.worker.client).toBeUndefined(); - expect(reentrant.worker.descriptor.lifecycle).toBe("recovering"); - expect(recoverWorker).toHaveBeenCalledWith(reentrant.worker); + expect(reentrant.worker.client).toBeDefined(); + expect(reentrant.worker.descriptor.lifecycle).toBe("ready"); + expect(recoverWorker).not.toHaveBeenCalled(); const completed = workerHarness(); for (const message of [frames.begin, frames.chunk, frames.end]) { internals.handleWorkerFrame(completed.worker, frame(message)); } internals.handleWorkerFrame(completed.worker, frame({ ...frames.begin, messageCount: 2 })); - expect(completed.close).toHaveBeenCalledOnce(); + expect(completed.close).not.toHaveBeenCalled(); expect(completed.worker.transcriptCaches.has(activeSessionId)).toBe(false); + expect(completed.worker.client).toBeDefined(); const mismatchedEnd = workerHarness(); for (const message of [frames.begin, frames.chunk, frames.end, frames.begin, frames.chunk]) { internals.handleWorkerFrame(mismatchedEnd.worker, frame(message)); } internals.handleWorkerFrame(mismatchedEnd.worker, frame({ ...frames.end, lastEventSequence: 2 })); - expect(mismatchedEnd.close).toHaveBeenCalledOnce(); + expect(mismatchedEnd.close).not.toHaveBeenCalled(); expect(mismatchedEnd.worker.transcriptCaches.has(activeSessionId)).toBe(false); + expect(mismatchedEnd.worker.client).toBeDefined(); const replaced = workerHarness(); internals.handleWorkerFrame(replaced.worker, frame(frames.begin)); @@ -628,7 +901,146 @@ describe("ENG-4602 snapshot transfer containment", () => { expect(replaced.worker.transcriptCaches.get(activeSessionId)?.snapshotId).toBe("snapshot-4602-new"); }); - it("fails one public snapshot without dropping another session on the shared client", async () => { + it("emits one terminal failure without a supervisor-side retry loop", async () => { + const supervisor = new DaemonSupervisor("/tmp/eng-4602-supervisor-bounded-retry.sock", { + defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, + descriptorDir: "/tmp/eng-4602-supervisor-bounded-retry-state", + }); + const firstSnapshotId = "snapshot-4602-failed-first"; + const retrySnapshotId = "snapshot-4602-failed-retry"; + const firstError = new Error("first generation failed"); + const retryError = new Error("retry generation failed"); + const transcripts = [ + new SnapshotTranscriptCache({ activeSessionId, snapshotId: firstSnapshotId, cacheRoot: "/tmp" }), + new SnapshotTranscriptCache({ activeSessionId, snapshotId: retrySnapshotId, cacheRoot: "/tmp" }), + ]; + transcripts[0]!.markFailed(firstError); + transcripts[1]!.markFailed(retryError); + const results = [streamedResult([], firstSnapshotId), streamedResult([], retrySnapshotId)]; + const socket = new PassThrough(); + socket.resume(); + const client = socketClient("bounded-retry", socket); + const { worker } = workerHarness(); + const records: DaemonOutbound[] = []; + let attachIndex = 0; + const attachClient = vi.fn(async () => { + const index = attachIndex++; + const transcript = transcripts[index]; + const result = results[index]; + if (!transcript || !result) throw new Error("snapshot retry was not bounded"); + return { result, worker, transcript, releaseTranscript: transcript.retain() }; + }); + const internals = supervisor as unknown as { + clients: Set; + attachClient: typeof attachClient; + writeSnapshotRecord(client: DaemonSocketClient, message: DaemonOutbound): Promise; + queueCatchup(client: DaemonSocketClient, activeSessionId: string, purpose: "replacement" | "resync"): void; + catchUpClient(client: DaemonSocketClient): Promise; + }; + internals.clients.add(client); + internals.attachClient = attachClient; + internals.writeSnapshotRecord = async (_client, message) => { + records.push(message); + return true; + }; + + internals.queueCatchup(client, activeSessionId, "replacement"); + await internals.catchUpClient(client); + + expect(attachClient).toHaveBeenCalledOnce(); + expect(records.map((record) => record.type)).toEqual(["session_snapshot_begin", "session_snapshot_failed"]); + expect(records[1]).toEqual({ + type: "session_snapshot_failed", + activeSessionId, + snapshotId: firstSnapshotId, + error: firstError.message, + }); + socket.destroy(); + for (const transcript of transcripts) transcript.dispose(); + }); + + it("expires an incomplete worker generation and settles cache and public waiters", async () => { + vi.useFakeTimers(); + try { + const supervisor = new DaemonSupervisor("/tmp/eng-4602-supervisor-deadline.sock", { + defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, + descriptorDir: "/tmp/eng-4602-supervisor-deadline-state", + }); + const { worker } = workerHarness(); + const client = socketClient("deadline", new PassThrough()); + const records: DaemonOutbound[] = []; + const internals = supervisor as unknown as { + clients: Set; + writeSnapshotRecord(client: DaemonSocketClient, message: DaemonOutbound): Promise; + handleWorkerFrame(worker: WorkerHarness, frame: PrivateFrame): void; + }; + internals.clients.add(client); + internals.writeSnapshotRecord = async (_client, message) => { + records.push(message); + return true; + }; + const frames = snapshotFrames([{ role: "user", content: "never completes", timestamp: 1 }]); + internals.handleWorkerFrame(worker, frame(frames.begin, "replacement")); + const transcript = worker.transcriptCaches.get(activeSessionId); + if (!transcript) throw new Error("snapshot deadline transcript was not created"); + const waiter = expect(transcript.waitForChunk(0)).rejects.toThrow("timed out before completion"); + + await vi.advanceTimersByTimeAsync(30_000); + await waiter; + await Promise.resolve(); + + expect(worker.snapshotCache.has(activeSessionId)).toBe(false); + expect(worker.transcriptCaches.has(activeSessionId)).toBe(false); + expect(worker.snapshotGenerations.has(activeSessionId)).toBe(false); + expect(records.filter((record) => record.type === "session_snapshot_failed")).toHaveLength(1); + client.socket.destroy(); + } finally { + vi.useRealTimers(); + } + }); + + it("settles a detached stream before a same-session reattach can revive zombie work", async () => { + const supervisor = new DaemonSupervisor("/tmp/eng-4602-supervisor-epoch.sock", { + defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, + descriptorDir: "/tmp/eng-4602-supervisor-epoch-state", + }); + const { worker } = workerHarness(); + const client = socketClient("epoch", new PassThrough()); + const transcript = new SnapshotTranscriptCache({ activeSessionId, snapshotId, cacheRoot: "/tmp" }); + const records: DaemonOutbound[] = []; + const internals = supervisor as unknown as { + advanceAttachmentEpoch(client: DaemonSocketClient, activeSessionId: string): number; + detachClient(client: DaemonSocketClient, activeSessionId: string): void; + writeSnapshotRecord(client: DaemonSocketClient, message: DaemonOutbound): Promise; + writeSnapshotBuffer(client: DaemonSocketClient, buffer: Uint8Array): Promise; + streamSnapshot( + client: DaemonSocketClient, + worker: WorkerHarness, + result: DaemonAttachResult, + transcript: SnapshotTranscriptCache, + ): Promise; + }; + internals.writeSnapshotRecord = async (_client, message) => { + records.push(message); + return true; + }; + internals.writeSnapshotBuffer = async () => true; + internals.advanceAttachmentEpoch(client, activeSessionId); + const streaming = internals.streamSnapshot(client, worker, streamedResult([]), transcript); + await Promise.resolve(); + expect(records.map((record) => record.type)).toEqual(["session_snapshot_begin"]); + + internals.detachClient(client, activeSessionId); + client.attachedActiveSessionIds.add(activeSessionId); + internals.advanceAttachmentEpoch(client, activeSessionId); + await streaming; + + expect(records.map((record) => record.type)).toEqual(["session_snapshot_begin"]); + transcript.dispose(); + client.socket.destroy(); + }); + + it("retries one public cache-read failure without dropping another session on the shared client", async () => { const supervisor = new DaemonSupervisor("/tmp/eng-4602-supervisor-public.sock", { defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, descriptorDir: "/tmp/eng-4602-supervisor-public-state", diff --git a/packages/coding-agent/test/suite/regressions/4677-snapshot-catchup-replacement.test.ts b/packages/coding-agent/test/suite/regressions/4677-snapshot-catchup-replacement.test.ts index 19258b6f41..15fd09d568 100644 --- a/packages/coding-agent/test/suite/regressions/4677-snapshot-catchup-replacement.test.ts +++ b/packages/coding-agent/test/suite/regressions/4677-snapshot-catchup-replacement.test.ts @@ -1,4 +1,4 @@ -import { existsSync, mkdtempSync, rmSync } from "node:fs"; +import { existsSync, mkdtempSync, readdirSync, rmSync } from "node:fs"; import type { Socket } from "node:net"; import { tmpdir } from "node:os"; import { join } from "node:path"; @@ -149,6 +149,58 @@ function workerHarness(result: DaemonAttachResult, transcript: SnapshotTranscrip } describe("ENG-4677 snapshot catch-up replacement", () => { + it("assigns supervisor-owned UUIDs to mixed-version full snapshots with the same cursor and count", async () => { + const root = tempDirectory(); + const supervisor = new DaemonSupervisor(join(root, "supervisor.sock"), { + defaultSessionConfig: { agentDir: root, cwd: root }, + descriptorDir: join(root, "state"), + }); + const large = "x".repeat(400 * 1024); + const firstMessages = [ + { role: "user", content: { nested: `before-${large}` }, timestamp: 1 }, + { role: "user", content: { nested: `before-2-${large}` }, timestamp: 2 }, + ] as unknown as AgentMessage[]; + const secondMessages = [ + { role: "user", content: { nested: "after" }, timestamp: 1 }, + { role: "user", content: { nested: "after-2" }, timestamp: 2 }, + ] as unknown as AgentMessage[]; + const withoutStream = (messages: AgentMessage[]): DaemonAttachResult => { + const result = streamedResult("legacy", messages.length, 7); + const { snapshotStream: _snapshotStream, ...full } = result; + full.snapshot = { ...full.snapshot, messages }; + return full; + }; + const firstResult = withoutStream(firstMessages); + const secondResult = withoutStream(secondMessages); + const worker = workerHarness( + firstResult, + new SnapshotTranscriptCache({ + activeSessionId, + snapshotId: "placeholder", + cacheRoot: root, + }), + ); + worker.transcriptCaches.clear(); + worker.snapshotCache.clear(); + worker.snapshotGenerations.clear(); + const internals = supervisor as unknown as { + getOrCreateTranscriptCache(worker: object, result: DaemonAttachResult): Promise; + }; + + const first = await internals.getOrCreateTranscriptCache(worker, firstResult); + const firstId = first.snapshotId; + const firstChunkCount = first.chunkCount; + const second = await internals.getOrCreateTranscriptCache(worker, secondResult); + + expect(firstId).not.toBe(second.snapshotId); + expect(firstId).toMatch(/^snapshot-[0-9a-f-]+$/); + expect(second.snapshotId).toMatch(/^snapshot-[0-9a-f-]+$/); + expect(firstChunkCount).toBe(2); + expect(second.chunkCount).toBe(1); + expect(Buffer.concat([...second]).toString("utf8")).toContain('"nested":"after"'); + second.dispose(); + }); + it("lets an incomplete retained snapshot finish after a newer generation begins", async () => { const root = tempDirectory(); const supervisor = new DaemonSupervisor(join(root, "supervisor.sock"), { @@ -377,7 +429,7 @@ describe("ENG-4677 snapshot catch-up replacement", () => { }); firstTranscript.appendEncodedChunk(Buffer.from("orphaned transcript chunk")); expect(firstTranscript.fileBacked).toBe(true); - const cacheDirectory = join(root, firstSnapshotId); + const cacheDirectory = join(root, readdirSync(root).find((entry) => entry.startsWith(firstSnapshotId))!); expect(existsSync(cacheDirectory)).toBe(true); const waiter = firstTranscript.waitForChunk(1); void waiter.catch(() => undefined); From 48eb15617803485976c072b74c0878059f62cbfb Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Sat, 29 Aug 2026 19:09:42 -0600 Subject: [PATCH 02/12] fix(daemon): harden immutable snapshot retries --- .../daemon-agent-connection.ts | 373 +++++++++++++---- .../src/modes/daemon/active-session-state.ts | 4 +- .../src/modes/daemon/daemon-mode.ts | 340 ++++++++++------ .../src/modes/daemon/daemon-protocol.ts | 9 +- .../src/modes/daemon/daemon-supervisor.ts | 382 ++++++++++++------ .../modes/daemon/snapshot-transcript-cache.ts | 34 +- .../modes/session-worker/private-framing.ts | 103 +++-- .../test/agent-connection-daemon.test.ts | 8 +- .../coding-agent/test/daemon-mode.test.ts | 31 +- .../test/daemon-supervisor-process.test.ts | 184 ++++++++- .../test/daemon-version-compatibility.test.ts | 230 +++++++++++ .../session-worker-private-framing.test.ts | 18 + .../test/snapshot-transcript-cache.test.ts | 30 +- .../4601-worker-snapshot-cache.test.ts | 53 +++ ...4602-snapshot-transfer-idempotency.test.ts | 11 +- .../4677-snapshot-catchup-replacement.test.ts | 76 +++- 16 files changed, 1504 insertions(+), 382 deletions(-) create mode 100644 packages/coding-agent/test/daemon-version-compatibility.test.ts 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 d76ab50fef..8574f386a1 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 @@ -123,12 +123,35 @@ interface DaemonSnapshotAssembly { timeout: ReturnType; } +interface DaemonSnapshotExpectation { + id: string; + activeSessionId: string; + sessionId: string; + messageCount: number; + targetChunkBytes: number; + purpose: DaemonSnapshotPurpose; + lastEventSequence: number; + lastEventCursor?: DaemonEventCursor; +} + interface DaemonSnapshotRequestAttempt { id: number; activeSessionId: string; attachmentEpoch: number; candidateSnapshotIds: Set; - expectedSnapshotId?: string; + logicalPurpose: DaemonSnapshotPurpose; + wirePurpose: DaemonSnapshotPurpose; + state: "requesting" | "receiving" | "succeeded" | "failed"; + expected?: DaemonSnapshotExpectation; +} + +interface DaemonRuntimeSnapshotAttempt { + activeSessionId: string; + attachmentEpoch: number; + purpose: "replacement" | "resync"; + state: "receiving" | "retrying" | "succeeded" | "terminal"; + snapshotIds: Set; + recovery?: Promise; } interface CorrelatedPromptRoute { @@ -328,7 +351,7 @@ export class DaemonAgentConnection implements AgentConnection { private readonly completedSnapshots = new Map(); private readonly completedSnapshotAttemptIds = new Map(); private readonly pendingReattachActiveSessionIds = new Set(); - private readonly snapshotRecoveryPromises = new Map>(); + private runtimeSnapshotAttempt: DaemonRuntimeSnapshotAttempt | undefined; private readonly ignoredSnapshotIds = new Set(); private readonly snapshotRequestAttempts = new Map(); private snapshotRequestAttemptSequence = 0; @@ -398,7 +421,7 @@ export class DaemonAgentConnection implements AgentConnection { const invalidatedInputPause = this.sessionInputPauses.size > 0; this.sessionInputPauses.clear(); this.sessionInputPauseGeneration++; - this.rejectSnapshotAssemblies(error); + this.rejectSnapshotAssemblies(error, true); if (this.disposed || this.terminalCloseEmitted) { return; } @@ -454,17 +477,24 @@ export class DaemonAgentConnection implements AgentConnection { resumeCursor: DaemonEventCursor | undefined, resetEventProgress: boolean, maxAttempts = 2, + logicalPurpose: DaemonSnapshotPurpose = "attach", ): Promise { const supportsExtensionUi = this.options.supportsExtensionUi !== false; const attachmentEpoch = this.advanceAttachmentEpoch(); let result!: SessionSummary | DaemonAttachResult; let streamedSnapshot: DaemonSessionSnapshot | undefined; for (let attemptIndex = 0; attemptIndex < maxAttempts; attemptIndex++) { - const requestAttempt = this.startSnapshotRequestAttempt(requestedActiveSessionId, attachmentEpoch); + const requestAttempt = this.startSnapshotRequestAttempt( + requestedActiveSessionId, + attachmentEpoch, + logicalPurpose, + "attach", + ); try { result = await this.requestData({ type: "attach", activeSessionId: requestedActiveSessionId, + snapshotGenerationNonce: randomUUID(), supportsExtensionUi, clientId: this.clientId, capabilities: [ @@ -493,13 +523,16 @@ export class DaemonAgentConnection implements AgentConnection { throw new Error("Daemon connection attachment changed during attach"); } if ("snapshot" in result && result.snapshotStream) { - requestAttempt.expectedSnapshotId = result.snapshotStream.id; + this.bindSnapshotExpectation(requestAttempt, result); + requestAttempt.state = "receiving"; streamedSnapshot = await this.waitForSnapshot(result.snapshotStream.id, requestAttempt); } else { streamedSnapshot = undefined; } + requestAttempt.state = "succeeded"; break; } catch (error) { + requestAttempt.state = "failed"; if (attemptIndex === maxAttempts - 1 || this.disposed || this.attachmentEpoch !== attachmentEpoch) { throw error; } @@ -568,6 +601,8 @@ export class DaemonAgentConnection implements AgentConnection { } async getInitialSnapshot(): Promise { + const snapshotRecovery = this.runtimeSnapshotAttempt?.recovery; + if (snapshotRecovery) await snapshotRecovery; if (this.latestSnapshotIsFresh && this.latestSnapshot) { return this.latestSnapshot; } @@ -1592,7 +1627,12 @@ export class DaemonAgentConnection implements AgentConnection { let result!: DaemonAttachResult; let snapshot!: DaemonSessionSnapshot; for (let attemptIndex = 0; attemptIndex < 2; attemptIndex++) { - const requestAttempt = this.startSnapshotRequestAttempt(targetActiveSessionId, attachmentEpoch); + const requestAttempt = this.startSnapshotRequestAttempt( + targetActiveSessionId, + attachmentEpoch, + "replacement", + attemptIndex === 0 ? "replacement" : "attach", + ); try { result = await this.requestData( attemptIndex === 0 @@ -1610,6 +1650,7 @@ export class DaemonAgentConnection implements AgentConnection { : { type: "attach", activeSessionId: targetActiveSessionId, + snapshotGenerationNonce: randomUUID(), supportsExtensionUi, clientId: this.clientId, capabilities, @@ -1622,14 +1663,19 @@ export class DaemonAgentConnection implements AgentConnection { throw new Error("Daemon returned an invalid session snapshot"); } if (result.snapshotStream) { - requestAttempt.expectedSnapshotId = result.snapshotStream.id; + this.bindSnapshotExpectation(requestAttempt, result); + requestAttempt.state = "receiving"; snapshot = await this.waitForSnapshot(result.snapshotStream.id, requestAttempt); } else { snapshot = result.snapshot; } + requestAttempt.state = "succeeded"; break; } catch (error) { - if (attemptIndex === 1 || this.disposed || this.attachmentEpoch !== attachmentEpoch) throw error; + requestAttempt.state = "failed"; + if (attemptIndex === 1 || !reattached || this.disposed || this.attachmentEpoch !== attachmentEpoch) { + throw error; + } } finally { this.finishSnapshotRequestAttempt(requestAttempt); } @@ -1919,17 +1965,32 @@ export class DaemonAgentConnection implements AgentConnection { return; } if (message.type === "session_snapshot_begin") { + const terminalRuntimeAttempt = this.runtimeSnapshotAttempt; + if ( + terminalRuntimeAttempt?.state === "terminal" && + terminalRuntimeAttempt.attachmentEpoch === this.attachmentEpoch && + terminalRuntimeAttempt.activeSessionId === message.activeSessionId + ) { + this.ignoreSnapshotId(message.snapshotId); + return; + } const pendingReplacement = this.pendingChunkedReplacement; const requestAttempt = this.latestSnapshotRequestAttempt(message.activeSessionId); + const retryingReplacement = + pendingReplacement !== undefined && + requestAttempt?.logicalPurpose === "replacement" && + requestAttempt.wirePurpose === "attach"; const acceptsUnsolicitedRuntimeSnapshot = - this.attachedSessionId !== undefined && - (message.purpose === "resync" || - (message.purpose === "replacement" && !this.supportsCorrelatedPromptLifecycle())); + this.attachedSessionId !== undefined && (message.purpose === "resync" || message.purpose === "replacement"); if (!pendingReplacement && !requestAttempt && !acceptsUnsolicitedRuntimeSnapshot) return; this.ignoredSnapshotIds.delete(message.snapshotId); if (pendingReplacement) { - if (message.purpose !== "replacement") { - this.failClosedReplacementReconciliation(); + if (message.purpose !== "replacement" && !retryingReplacement) { + void this.transitionRuntimeSnapshotFailure( + "replacement", + new Error("Daemon returned an invalid replacement snapshot purpose"), + message.snapshotId, + ); return; } if (pendingReplacement.snapshotId && pendingReplacement.snapshotId !== message.snapshotId) { @@ -1960,12 +2021,14 @@ export class DaemonAgentConnection implements AgentConnection { } try { validateDaemonSnapshotBegin(message); + if (requestAttempt) this.validateSnapshotAttemptBegin(requestAttempt, message); } catch (error) { await this.rejectInvalidSnapshotAssembly( message.snapshotId, assembly, error instanceof Error ? error : new Error(String(error)), message.activeSessionId, + message.purpose, ); return; } @@ -1974,6 +2037,9 @@ export class DaemonAgentConnection implements AgentConnection { this.completedSnapshots.delete(message.snapshotId); this.completedSnapshotAttemptIds.delete(message.snapshotId); assembly.begin = message; + if (!requestAttempt && (message.purpose === "replacement" || message.purpose === "resync")) { + this.startRuntimeSnapshotAttempt(message.purpose, message.activeSessionId, message.snapshotId); + } return; } if (message.type === "session_snapshot_chunk") { @@ -2007,30 +2073,23 @@ export class DaemonAgentConnection implements AgentConnection { } if (message.type === "session_snapshot_failed") { const existingAssembly = this.snapshotAssemblies.get(message.snapshotId); + const explicitPurpose = message.purpose; const assembly = existingAssembly ?? - (this.pendingChunkedReplacement ? this.getSnapshotAssembly(message.snapshotId) : undefined); + (this.pendingChunkedReplacement || explicitPurpose + ? this.getSnapshotAssembly(message.snapshotId) + : undefined); if (!assembly || assembly.settled || assembly.attachmentEpoch !== this.attachmentEpoch) return; - const purpose = assembly.begin?.purpose ?? (this.pendingChunkedReplacement ? "replacement" : "attach"); - const snapshotError = new Error(message.error); - const pendingReattach = this.pendingReattachActiveSessionIds.has(message.activeSessionId); - const recoveryPromise = - !pendingReattach && (purpose === "replacement" || purpose === "resync") - ? this.recoverFailedSnapshot(purpose, snapshotError) - : undefined; - if (recoveryPromise) { - this.snapshotRecoveryPromises.set(message.snapshotId, recoveryPromise); - } - this.rejectSnapshotAssembly(message.snapshotId, assembly, snapshotError); - this.ignoreSnapshotId(message.snapshotId); - if (recoveryPromise) { - try { - await recoveryPromise; - } finally { - this.snapshotRecoveryPromises.delete(message.snapshotId); - } - if (!this.terminalCloseEmitted) this.releaseChunkedReplacementFence(); - } + const purpose = + explicitPurpose ?? assembly.begin?.purpose ?? (this.pendingChunkedReplacement ? "replacement" : "attach"); + const recovery = this.transitionSnapshotFailure( + message.snapshotId, + assembly, + new Error(message.error), + message.activeSessionId, + purpose, + ); + if (recovery) await recovery; return; } if (this.isStaleSequencedMessage(message)) { @@ -2369,12 +2428,64 @@ export class DaemonAgentConnection implements AgentConnection { return nextEpoch; } - private startSnapshotRequestAttempt(activeSessionId: string, attachmentEpoch: number): DaemonSnapshotRequestAttempt { + private bindSnapshotExpectation(attempt: DaemonSnapshotRequestAttempt, result: DaemonAttachResult): void { + attempt.expected = this.snapshotExpectation(result, attempt.wirePurpose); + for (const snapshotId of attempt.candidateSnapshotIds) { + const begin = this.snapshotAssemblies.get(snapshotId)?.begin; + if (begin) this.validateSnapshotAttemptBegin(attempt, begin); + } + } + + private snapshotExpectation(result: DaemonAttachResult, purpose: DaemonSnapshotPurpose): DaemonSnapshotExpectation { + const stream = result.snapshotStream; + if (!stream) throw new Error("Daemon snapshot response omitted its stream descriptor"); + return { + id: stream.id, + activeSessionId: result.activeSessionId, + sessionId: result.snapshot.state.sessionId, + messageCount: stream.messageCount, + targetChunkBytes: stream.targetChunkBytes, + purpose, + lastEventSequence: result.snapshot.lastEventSequence, + lastEventCursor: result.snapshot.lastEventCursor, + }; + } + + private validateSnapshotAttemptBegin(attempt: DaemonSnapshotRequestAttempt, begin: DaemonSnapshotBegin): void { + const expected = attempt.expected; + if (!expected) return; + const cursor = begin.snapshot.lastEventCursor; + if ( + begin.snapshotId !== expected.id || + begin.activeSessionId !== expected.activeSessionId || + begin.snapshot.activeSessionId !== expected.activeSessionId || + begin.snapshot.state.sessionId !== expected.sessionId || + begin.snapshot.summary.sessionId !== expected.sessionId || + begin.messageCount !== expected.messageCount || + begin.targetChunkBytes !== expected.targetChunkBytes || + (begin.purpose ?? "attach") !== expected.purpose || + begin.snapshot.lastEventSequence !== expected.lastEventSequence || + cursor?.generation !== expected.lastEventCursor?.generation || + cursor?.sequence !== expected.lastEventCursor?.sequence + ) { + throw new Error(`Snapshot ${begin.snapshotId} did not match its response descriptor`); + } + } + + private startSnapshotRequestAttempt( + activeSessionId: string, + attachmentEpoch: number, + logicalPurpose: DaemonSnapshotPurpose, + wirePurpose: DaemonSnapshotPurpose, + ): DaemonSnapshotRequestAttempt { const attempt: DaemonSnapshotRequestAttempt = { id: ++this.snapshotRequestAttemptSequence, activeSessionId, attachmentEpoch, candidateSnapshotIds: new Set(), + logicalPurpose, + wirePurpose, + state: "requesting", }; this.snapshotRequestAttempts.set(attempt.id, attempt); return attempt; @@ -2440,16 +2551,15 @@ export class DaemonAgentConnection implements AgentConnection { const timeout = setTimeout(() => { const current = this.snapshotAssemblies.get(snapshotId); if (current && !current.settled) { - current.reject(new Error(`Timed out waiting for snapshot ${snapshotId}`)); - this.snapshotAssemblies.delete(snapshotId); - this.ignoreSnapshotId(snapshotId); - if ( - this.pendingChunkedReplacement || - (current.begin?.purpose === "replacement" && - !this.pendingReattachActiveSessionIds.has(current.begin.activeSessionId)) - ) { - this.failClosedReplacementReconciliation(); - } + const activeSessionId = + current.begin?.activeSessionId ?? requestAttempt?.activeSessionId ?? this.activeSessionId; + void this.transitionSnapshotFailure( + snapshotId, + current, + new Error(`Timed out waiting for snapshot ${snapshotId}`), + activeSessionId, + current.begin?.purpose, + ); } }, this.options.snapshotTimeoutMs ?? DAEMON_SNAPSHOT_TIMEOUT_MS); timeout.unref(); @@ -2475,17 +2585,27 @@ export class DaemonAgentConnection implements AgentConnection { return assembly; } - private rejectSnapshotAssemblies(error: Error): void { - for (const assembly of this.snapshotAssemblies.values()) { + private rejectSnapshotAssemblies(error: Error, recoverRuntime = false): void { + for (const [snapshotId, assembly] of [...this.snapshotAssemblies]) { clearTimeout(assembly.timeout); - assembly.reject(error); + if (recoverRuntime) { + void this.transitionSnapshotFailure( + snapshotId, + assembly, + error, + assembly.begin?.activeSessionId ?? this.activeSessionId, + assembly.begin?.purpose, + ); + } else { + assembly.reject(error); + } } this.snapshotAssemblies.clear(); this.completedSnapshots.clear(); this.completedSnapshotAttemptIds.clear(); this.snapshotRequestAttempts.clear(); - this.snapshotRecoveryPromises.clear(); this.ignoredSnapshotIds.clear(); + if (!recoverRuntime) this.runtimeSnapshotAttempt = undefined; } private ignoreSnapshotId(snapshotId: string): void { @@ -2512,41 +2632,138 @@ export class DaemonAgentConnection implements AgentConnection { assembly: DaemonSnapshotAssembly, error: Error, activeSessionId: string, + explicitPurpose?: DaemonSnapshotPurpose, ): Promise { + const recovery = this.transitionSnapshotFailure( + snapshotId, + assembly, + error, + activeSessionId, + explicitPurpose ?? assembly.begin?.purpose ?? (this.pendingChunkedReplacement ? "replacement" : undefined), + ); + if (this.supportsCorrelatedPromptLifecycle()) { + await this.emit({ type: "correlated_prompt_protocol_violation" }); + } + if (recovery) await recovery; + } + + private startRuntimeSnapshotAttempt( + purpose: "replacement" | "resync", + activeSessionId: string, + snapshotId: string, + ): DaemonRuntimeSnapshotAttempt { + const current = this.runtimeSnapshotAttempt; + if ( + current && + current.attachmentEpoch === this.attachmentEpoch && + current.activeSessionId === activeSessionId && + current.state !== "succeeded" + ) { + for (const previousId of current.snapshotIds) { + if (previousId === snapshotId) continue; + const previous = this.snapshotAssemblies.get(previousId); + if (previous && !previous.settled) { + this.rejectSnapshotAssembly( + previousId, + previous, + new Error(`Snapshot ${previousId} was superseded by ${snapshotId}`), + ); + this.snapshotAssemblies.delete(previousId); + this.ignoreSnapshotId(previousId); + } + } + current.snapshotIds.add(snapshotId); + return current; + } + const attempt: DaemonRuntimeSnapshotAttempt = { + activeSessionId, + attachmentEpoch: this.attachmentEpoch, + purpose, + state: "receiving", + snapshotIds: new Set([snapshotId]), + }; + this.runtimeSnapshotAttempt = attempt; + return attempt; + } + + private transitionSnapshotFailure( + snapshotId: string, + assembly: DaemonSnapshotAssembly, + error: Error, + activeSessionId: string, + explicitPurpose?: DaemonSnapshotPurpose, + ): Promise | undefined { this.rejectSnapshotAssembly(snapshotId, assembly, error); - this.snapshotAssemblies.delete(snapshotId); this.ignoreSnapshotId(snapshotId); + const requestAttempt = + assembly.requestAttemptId === undefined + ? undefined + : this.snapshotRequestAttempts.get(assembly.requestAttemptId); + if (requestAttempt) { + // Keep the settled assembly until the command response binds its descriptor. + // This covers begin/failed frames that beat the attach response continuation. + this.snapshotAssemblies.set(snapshotId, assembly); + requestAttempt.state = "failed"; + return undefined; + } + this.snapshotAssemblies.delete(snapshotId); + const purpose = explicitPurpose ?? assembly.begin?.purpose ?? "attach"; if ( - this.pendingChunkedReplacement || - (assembly.begin?.purpose === "replacement" && !this.pendingReattachActiveSessionIds.has(activeSessionId)) + (purpose === "replacement" || purpose === "resync") && + !this.pendingReattachActiveSessionIds.has(activeSessionId) ) { - this.failClosedReplacementReconciliation(); - } - if (this.supportsCorrelatedPromptLifecycle()) { - await this.emit({ type: "correlated_prompt_protocol_violation" }); + return this.transitionRuntimeSnapshotFailure(purpose, error, snapshotId, activeSessionId); } + return undefined; } - private async recoverFailedSnapshot(purpose: "replacement" | "resync", snapshotError: Error): Promise { + private transitionRuntimeSnapshotFailure( + purpose: "replacement" | "resync", + snapshotError: Error, + snapshotId: string, + activeSessionId = this.activeSessionId, + ): Promise { + const attempt = this.startRuntimeSnapshotAttempt(purpose, activeSessionId, snapshotId); + if (attempt.state === "terminal" || attempt.state === "succeeded") return Promise.resolve(); + if (attempt.recovery) return attempt.recovery; + attempt.state = "retrying"; this.latestSnapshotIsFresh = false; - try { - await this.attachSession(this.activeSessionId, this.lastEventCursor, false, 1); - const snapshot = await this.getInitialSnapshot(); - if (this.disposed) return; - if (purpose === "replacement") { - await this.emit({ type: "session_replaced", state: snapshot.state, messages: snapshot.messages }); - } else { - await this.emit({ type: "session_resynced", snapshot }); - } - } catch (recoveryError) { - if (this.disposed) { - return; - } - this.terminalCloseEmitted = true; - await this.emit({ - type: "closed", - error: `Failed to recover from a ${purpose} snapshot transfer. Snapshot error: ${formatErrorSentence(snapshotError)} Recovery error: ${formatErrorSentence(recoveryError)} ${this.formatDaemonDiagnosticContext()}`, - }); + const recovery = this.recoverFailedSnapshot(attempt.purpose).then( + async () => { + if (this.runtimeSnapshotAttempt !== attempt || attempt.state === "terminal") return; + attempt.state = "succeeded"; + if (attempt.purpose === "replacement") this.releaseChunkedReplacementFence(); + this.runtimeSnapshotAttempt = undefined; + }, + async (recoveryError: unknown) => { + if (this.disposed || this.runtimeSnapshotAttempt !== attempt || attempt.state === "terminal") return; + attempt.state = "terminal"; + for (const failedId of attempt.snapshotIds) this.ignoreSnapshotId(failedId); + if (attempt.purpose === "replacement") { + this.replacementReconciliationFailed = true; + this.pendingChunkedReplacement = undefined; + } + if (this.terminalCloseEmitted) return; + this.terminalCloseEmitted = true; + await this.emit({ + type: "closed", + error: `Failed to recover from a ${attempt.purpose} snapshot transfer. Snapshot error: ${formatErrorSentence(snapshotError)} Recovery error: ${formatErrorSentence(recoveryError)} ${this.formatDaemonDiagnosticContext()}`, + }); + }, + ); + attempt.recovery = recovery; + return recovery; + } + + private async recoverFailedSnapshot(purpose: "replacement" | "resync"): Promise { + await this.attachSession(this.activeSessionId, this.lastEventCursor, false, 1, purpose); + const snapshot = this.latestSnapshot; + if (!snapshot) throw new Error("Daemon snapshot retry returned no snapshot"); + if (this.disposed) return; + if (purpose === "replacement") { + await this.emit({ type: "session_replaced", state: snapshot.state, messages: snapshot.messages }); + } else { + await this.emit({ type: "session_resynced", snapshot }); } } @@ -2565,6 +2782,7 @@ export class DaemonAgentConnection implements AgentConnection { if (requestAttempt) { requestAttempt.candidateSnapshotIds.add(snapshotId); assembly.requestAttemptId = requestAttempt.id; + if (assembly.begin) this.validateSnapshotAttemptBegin(requestAttempt, assembly.begin); } try { return await assembly.promise; @@ -2802,6 +3020,11 @@ export class DaemonAgentConnection implements AgentConnection { assembly.resolve(snapshot); clearTimeout(assembly.timeout); this.snapshotAssemblies.delete(message.snapshotId); + const runtimeAttempt = this.runtimeSnapshotAttempt; + if (runtimeAttempt?.snapshotIds.has(message.snapshotId) && runtimeAttempt.state === "receiving") { + runtimeAttempt.state = "succeeded"; + this.runtimeSnapshotAttempt = undefined; + } if (purpose === "replacement") { await this.emit({ type: "session_replaced", state: snapshot.state, messages }); } else { diff --git a/packages/coding-agent/src/modes/daemon/active-session-state.ts b/packages/coding-agent/src/modes/daemon/active-session-state.ts index fe062253ce..9a925a8067 100644 --- a/packages/coding-agent/src/modes/daemon/active-session-state.ts +++ b/packages/coding-agent/src/modes/daemon/active-session-state.ts @@ -15,8 +15,6 @@ export interface DaemonSocketClient { catchupPurposes?: Map; /** The single catch-up drain currently serving this client. */ catchupPromise?: Promise; - /** Delayed retry after transient catch-up snapshot preparation failure. */ - catchupRetryTimer?: NodeJS.Timeout; backpressured?: boolean; authenticated?: boolean; transport?: "jsonl" | "private-framed"; @@ -25,6 +23,8 @@ export interface DaemonSocketClient { snapshotActiveSessionCounts?: Map; snapshotTransferAbortControllers?: Map; snapshotTransferTails?: Map>; + /** Serializes private frames when a large snapshot frame pauses between parts. */ + privateFrameWriteTail?: Promise; detachInput: () => void; supportsExtensionUi: boolean; capabilities: Set; diff --git a/packages/coding-agent/src/modes/daemon/daemon-mode.ts b/packages/coding-agent/src/modes/daemon/daemon-mode.ts index ffed26bc70..bda8b01c80 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-mode.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-mode.ts @@ -242,6 +242,7 @@ export { defaultDaemonSocketPath } from "./daemon-socket.js"; const structuredLog = getLogger("coding-agent.daemon"); const WORKER_SNAPSHOT_TERMINAL_DRAIN_TIMEOUT_MS = 1_000; const WORKER_SNAPSHOT_PREPARATION_TIMEOUT_MS = 30_000; +const WORKER_PRIVATE_FRAME_WRITE_BYTES = 64 * 1024; const UPDATE_RESTART_PREPARE_TIMEOUT_MS = 90_000; const MAX_SESSION_SNAPSHOT_STABILIZATION_RETRIES = 3; @@ -249,6 +250,20 @@ function createSnapshotTransferId(): string { return `snapshot-${randomUUID()}`; } +function createSnapshotAttemptSignal(attachmentSignal?: AbortSignal): AbortSignal { + const deadlineSignal = AbortSignal.timeout(WORKER_SNAPSHOT_PREPARATION_TIMEOUT_MS); + return attachmentSignal ? AbortSignal.any([attachmentSignal, deadlineSignal]) : deadlineSignal; +} + +function waitForSnapshotAttempt(promise: Promise, signal: AbortSignal, snapshotId: string): Promise { + if (signal.aborted) return Promise.reject(new Error(`Snapshot ${snapshotId} was aborted`)); + return new Promise((resolve, reject) => { + const onAbort = () => reject(new Error(`Snapshot ${snapshotId} was aborted`)); + signal.addEventListener("abort", onAbort, { once: true }); + void promise.then(resolve, reject).finally(() => signal.removeEventListener("abort", onAbort)); + }); +} + const DAEMON_COMMAND_TYPES: ReadonlySet = new Set([ "ack_result", "list", @@ -355,7 +370,6 @@ const DAEMON_COMMAND_TYPES: ReadonlySet = new Set([ ]); const DAEMON_CLIENT_CAPABILITY_SET: ReadonlySet = new Set(DAEMON_SUPPORTED_CLIENT_CAPABILITIES); -const CLIENT_CATCHUP_RETRY_MS = 250; const UPDATE_RESTART_ABORT_BASH_TIMEOUT_MS = 5000; const SUPERVISOR_FENCE_POLL_MS = 250; const UPDATE_RESTART_MARKER = @@ -3259,7 +3273,6 @@ export class AgentDaemon { cleanedUp = true; socket.off("close", cleanup); socket.off("error", cleanup); - this.clearClientCatchupRetry(client); for (const [pauseId, entry] of this.sessionInputPauses) { if (entry.owner !== client) continue; entry.pause.release(); @@ -3858,6 +3871,7 @@ export class AgentDaemon { const snapshotSignal = streamsSnapshot ? markClientSnapshotStreaming(client, state.activeSessionId) : undefined; + const snapshotAttemptSignal = snapshotSignal ? createSnapshotAttemptSignal(snapshotSignal) : undefined; let result: DaemonAttachResult; state.pendingAttaches++; try { @@ -3895,9 +3909,9 @@ export class AgentDaemon { transcript = await this.prepareWorkerSnapshotTranscript({ activeSessionId: state.activeSessionId, snapshotId, - generationKey: this.snapshotPayloadGenerationKey(result), + generationKey: `${this.snapshotPayloadGenerationKey(result)}:${command.snapshotGenerationNonce ?? "shared"}`, messages: result.snapshot.messages, - signal: snapshotSignal, + signal: snapshotAttemptSignal, }); } catch (error) { state.clients.delete(client); @@ -3922,7 +3936,7 @@ export class AgentDaemon { streamedResult, transcript, "attach", - snapshotSignal, + snapshotAttemptSignal, true, ).catch((error) => this.log(`could not stream attach snapshot: ${String(error)}`)); }); @@ -5193,7 +5207,7 @@ export class AgentDaemon { ) { const payload = await existing.promise; options.signal?.throwIfAborted(); - return payload.createTransfer(options.activeSessionId, options.snapshotId); + return this.createWorkerSnapshotTransfer(options.activeSessionId, options.snapshotId, existing, payload); } if (existing) { const stalePayload = await existing.promise.catch(() => undefined); @@ -5209,14 +5223,14 @@ export class AgentDaemon { promise: prepareSnapshotTranscriptPayload({ messages: options.messages, targetChunkBytes: SNAPSHOT_TARGET_CHUNK_BYTES, - signal: AbortSignal.timeout(WORKER_SNAPSHOT_PREPARATION_TIMEOUT_MS), + signal: options.signal ?? AbortSignal.timeout(WORKER_SNAPSHOT_PREPARATION_TIMEOUT_MS), }), }; this.snapshotPayloadGenerations.set(options.activeSessionId, generation); try { const payload = await generation.promise; options.signal?.throwIfAborted(); - return payload.createTransfer(options.activeSessionId, options.snapshotId); + return this.createWorkerSnapshotTransfer(options.activeSessionId, options.snapshotId, generation, payload); } catch (error) { if (this.snapshotPayloadGenerations.get(options.activeSessionId) === generation) { this.snapshotPayloadGenerations.delete(options.activeSessionId); @@ -5226,6 +5240,22 @@ export class AgentDaemon { } } + private createWorkerSnapshotTransfer( + activeSessionId: string, + snapshotId: string, + generation: { key: string; messages: readonly AgentMessage[]; promise: Promise }, + payload: SnapshotTranscriptPayloadCache, + ): SnapshotTranscriptChunkSource { + const transfer = payload.createTransfer(activeSessionId, snapshotId); + transfer.markFailed = () => { + if (this.snapshotPayloadGenerations.get(activeSessionId) === generation) { + this.snapshotPayloadGenerations.delete(activeSessionId); + payload.dispose(); + } + }; + return transfer; + } + private snapshotPayloadGenerationKey(result: DaemonAttachResult): string { const cursor = result.lastEventCursor; return [ @@ -5255,7 +5285,8 @@ export class AgentDaemon { if (snapshotAlreadyMarked && !signal) { throw new Error(`Snapshot ${stream.id} is missing its transfer signal`); } - const transferSignal = signal ?? markClientSnapshotStreaming(client, result.activeSessionId); + const attachmentSignal = signal ?? markClientSnapshotStreaming(client, result.activeSessionId); + const transferSignal = signal ?? createSnapshotAttemptSignal(attachmentSignal); if (client.socket.destroyed) { finishClientSnapshotStreaming(client, result.activeSessionId); transcript.dispose?.(); @@ -5269,7 +5300,17 @@ export class AgentDaemon { }); client.snapshotTransferTails.set(result.activeSessionId, transfer); if (previousTransfer) { - await previousTransfer; + try { + await waitForSnapshotAttempt(previousTransfer, transferSignal, stream.id); + } catch (error) { + finishTransfer(); + if (client.snapshotTransferTails.get(result.activeSessionId) === transfer) { + client.snapshotTransferTails.delete(result.activeSessionId); + } + finishClientSnapshotStreaming(client, result.activeSessionId); + transcript.dispose?.(); + throw error; + } } const { messages: _messages, ...snapshot } = result.snapshot; const snapshotBegin: DaemonOutbound = { @@ -5309,6 +5350,7 @@ export class AgentDaemon { activeSessionId: result.activeSessionId, snapshotId: stream.id, error: streamError.message, + purpose: purpose === "catchup" ? "resync" : purpose, }, purpose, undefined, @@ -5333,11 +5375,17 @@ export class AgentDaemon { return; } let chunkCount = 0; - for await (const chunk of transcript) { - if (transferSignal.aborted) { - await deliverSnapshotFailure(new Error(`Snapshot ${stream.id} was aborted`)); - return; - } + const asyncIteratorFactory = (transcript as Partial>)[ + Symbol.asyncIterator + ]; + const asyncIterator = asyncIteratorFactory?.call(transcript); + const syncIteratorFactory = (transcript as Partial>)[Symbol.iterator]; + const syncIterator = asyncIterator ? undefined : syncIteratorFactory?.call(transcript); + if (!asyncIterator && !syncIterator) throw new Error(`Snapshot ${stream.id} has no transcript iterator`); + while (true) { + const nextPromise = asyncIterator ? asyncIterator.next() : Promise.resolve(syncIterator!.next()); + const next = await waitForSnapshotAttempt(nextPromise, transferSignal, stream.id); + if (next.done) break; const headerMessage: DaemonOutbound = { type: "session_snapshot_chunk", activeSessionId: result.activeSessionId, @@ -5345,10 +5393,8 @@ export class AgentDaemon { index: chunkCount, messages: [], }; - if (!(await this.writeWorkerSnapshotBuffer(client, chunk, headerMessage, purpose, transferSignal))) { - if (transferSignal.aborted) { - await deliverSnapshotFailure(new Error(`Snapshot ${stream.id} was aborted`)); - } + if (!(await this.writeWorkerSnapshotBuffer(client, next.value, headerMessage, purpose, transferSignal))) { + await deliverSnapshotFailure(new Error(`Snapshot ${stream.id} was aborted`)); return; } chunkCount++; @@ -5357,7 +5403,7 @@ export class AgentDaemon { await deliverSnapshotFailure(new Error(`Snapshot ${stream.id} was aborted`)); return; } - await this.writeWorkerSnapshotRecord( + const ended = await this.writeWorkerSnapshotRecord( client, { type: "session_snapshot_end", @@ -5370,6 +5416,7 @@ export class AgentDaemon { purpose, transferSignal, ); + if (!ended) await deliverSnapshotFailure(new Error(`Snapshot ${stream.id} was aborted`)); } catch (error) { if (transferSignal.aborted) { await deliverSnapshotFailure(new Error(`Snapshot ${stream.id} was aborted`)); @@ -5418,27 +5465,85 @@ export class AgentDaemon { signal?: AbortSignal, drainTimeoutMs?: number, ): Promise { - if (signal?.aborted || client.socket.destroyed) { - return false; + if (signal?.aborted || client.socket.destroyed) return false; + if (client.transport === "private-framed" && !Object.hasOwn(this, "writeSerialized")) { + const parts = this.encodeWorkerPrivateFrameParts(buffer, message, "jsonl", purpose); + return this.writePrivateFrameParts(client, parts, signal, drainTimeoutMs); } - if (this.writeSerialized(client, buffer, message, "jsonl", purpose)) { + if (this.writeSerialized(client, buffer, message, "jsonl", purpose)) return true; + return this.waitForWorkerSnapshotDrain(client, signal, drainTimeoutMs); + } + + private encodeWorkerPrivateFrameParts( + line: string | Buffer | readonly Buffer[], + message: DaemonOutbound, + payloadEncoding: "jsonl" | "assistant-delta", + snapshotPurpose?: "attach" | "replacement" | "catchup", + ): Buffer[] { + const payloadParts = typeof line === "string" ? [Buffer.from(line)] : Buffer.isBuffer(line) ? [line] : [...line]; + return encodePrivateFrameParts( + { + kind: "outbound", + outboundType: message.type, + ...(hasDaemonOutboundActiveSessionId(message) ? { activeSessionId: message.activeSessionId } : {}), + ...("id" in message && typeof message.id === "string" ? { requestId: message.id } : {}), + ...("snapshotId" in message && typeof message.snapshotId === "string" + ? { snapshotId: message.snapshotId } + : {}), + ...(message.type === "session_event" ? { sessionEventType: message.event.type } : {}), + payloadEncoding, + ...(snapshotPurpose ? { snapshotPurpose } : {}), + }, + payloadParts, + ); + } + + private async writePrivateFrameParts( + client: DaemonSocketClient, + parts: readonly Buffer[], + signal?: AbortSignal, + drainTimeoutMs?: number, + ): Promise { + const previous = client.privateFrameWriteTail; + let resolveTail!: () => void; + const tail = new Promise((resolve) => { + resolveTail = resolve; + }); + client.privateFrameWriteTail = tail; + try { + if (previous) await waitForSnapshotAttempt(previous, signal ?? AbortSignal.timeout(30_000), "private-frame"); + for (const part of parts) { + for (let offset = 0; offset < part.length; offset += WORKER_PRIVATE_FRAME_WRITE_BYTES) { + if (signal?.aborted || client.socket.destroyed) return false; + const slice = part.subarray(offset, offset + WORKER_PRIVATE_FRAME_WRITE_BYTES); + if (client.socket.write(slice)) continue; + client.backpressured = true; + if (!(await this.waitForWorkerSnapshotDrain(client, signal, drainTimeoutMs))) return false; + } + } return true; + } finally { + resolveTail(); + if (client.privateFrameWriteTail === tail) client.privateFrameWriteTail = undefined; } - return new Promise((resolveDrain) => { + } + + private waitForWorkerSnapshotDrain( + client: DaemonSocketClient, + signal?: AbortSignal, + drainTimeoutMs?: number, + ): Promise { + return new Promise((resolveDrain) => { let settled = false; let drainTimeout: NodeJS.Timeout | undefined; const finish = (value: boolean) => { - if (settled) { - return; - } + if (settled) return; settled = true; client.socket.off("drain", onDrain); client.socket.off("close", onClose); client.socket.off("error", onClose); signal?.removeEventListener("abort", onAbort); - if (drainTimeout) { - clearTimeout(drainTimeout); - } + if (drainTimeout) clearTimeout(drainTimeout); resolveDrain(value); }; const onDrain = () => finish(true); @@ -5452,9 +5557,7 @@ export class AgentDaemon { drainTimeout = setTimeout(() => finish(false), drainTimeoutMs); drainTimeout.unref(); } - if (signal?.aborted || client.socket.destroyed) { - finish(false); - } + if (signal?.aborted || client.socket.destroyed) finish(false); }); } @@ -6788,6 +6891,7 @@ export class AgentDaemon { snapshotId: string, snapshotSignal: AbortSignal, ): Promise { + const snapshotAttemptSignal = createSnapshotAttemptSignal(snapshotSignal); const result = await this.createAttachResult(client, state, { type: "attach", activeSessionId: state.activeSessionId, @@ -6806,7 +6910,7 @@ export class AgentDaemon { snapshotId, generationKey: this.snapshotPayloadGenerationKey(result), messages: result.snapshot.messages, - signal: snapshotSignal, + signal: snapshotAttemptSignal, }); this.write(client, { ...message, messages: [], snapshotFollows: true }); void this.streamWorkerSnapshot( @@ -6823,7 +6927,7 @@ export class AgentDaemon { }, transcript, "replacement", - snapshotSignal, + snapshotAttemptSignal, true, ).catch((error) => { this.log(`could not stream replacement snapshot: ${String(error)}`); @@ -6871,7 +6975,6 @@ export class AgentDaemon { if (client.snapshotStreaming || client.backpressured) { return Promise.resolve(); } - this.clearClientCatchupRetry(client); const catchup = this.drainBackpressuredClientCatchupQueue(client).finally(() => { if (client.catchupPromise === catchup) { client.catchupPromise = undefined; @@ -6881,33 +6984,6 @@ export class AgentDaemon { return catchup; } - private clearClientCatchupRetry(client: DaemonSocketClient): void { - if (!client.catchupRetryTimer) { - return; - } - clearTimeout(client.catchupRetryTimer); - client.catchupRetryTimer = undefined; - } - - private scheduleClientCatchupRetry(client: DaemonSocketClient): void { - if (client.socket.destroyed || client.catchupRetryTimer) { - return; - } - client.catchupRetryTimer = setTimeout(() => { - client.catchupRetryTimer = undefined; - if (client.socket.destroyed || !client.catchupActiveSessionIds?.size) { - return; - } - if (client.snapshotStreaming || client.backpressured) { - this.scheduleClientCatchupRetry(client); - return; - } - void this.catchUpBackpressuredClient(client).catch((error) => - this.log(`could not retry catch-up for client ${client.id}: ${String(error)}`), - ); - }, CLIENT_CATCHUP_RETRY_MS); - } - private async drainBackpressuredClientCatchupQueue(client: DaemonSocketClient): Promise { while ( !client.socket.destroyed && @@ -6915,13 +6991,11 @@ export class AgentDaemon { !client.backpressured && client.catchupActiveSessionIds?.size ) { - if ((await this.drainBackpressuredClientCatchups(client)) === "retry-later") { - return; - } + await this.drainBackpressuredClientCatchups(client); } } - private async drainBackpressuredClientCatchups(client: DaemonSocketClient): Promise<"drained" | "retry-later"> { + private async drainBackpressuredClientCatchups(client: DaemonSocketClient): Promise<"drained"> { if (client.socket.destroyed) { return "drained"; } @@ -6949,23 +7023,9 @@ export class AgentDaemon { client.transport === "private-framed" && daemonClientCapabilitiesForSession(client, activeSessionId).has("chunked_snapshot") ) { - if (purpose === "replacement") { - this.write(client, { - type: "session_replaced", - activeSessionId, - state: result.snapshot.state, - messages: [], - snapshotFollows: true, - meta: createDaemonEventMeta( - activeSessionId, - state.lastEventSequence, - undefined, - state.eventGeneration, - ), - }); - } const snapshotId = createSnapshotTransferId(); const snapshotSignal = markClientSnapshotStreaming(client, activeSessionId); + const snapshotAttemptSignal = createSnapshotAttemptSignal(snapshotSignal); let transcript: SnapshotTranscriptChunkSource; try { transcript = await this.prepareWorkerSnapshotTranscript({ @@ -6973,11 +7033,50 @@ export class AgentDaemon { snapshotId, generationKey: this.snapshotPayloadGenerationKey(result), messages: result.snapshot.messages, - signal: snapshotSignal, + signal: snapshotAttemptSignal, }); } catch (error) { finishClientSnapshotStreaming(client, activeSessionId); - throw error; + const snapshotError = error instanceof Error ? error : new Error(String(error)); + await this.writeWorkerSnapshotRecord( + client, + { + type: "session_snapshot_failed", + activeSessionId, + snapshotId, + error: snapshotError.message, + purpose, + }, + purpose === "replacement" ? "replacement" : "catchup", + undefined, + WORKER_SNAPSHOT_TERMINAL_DRAIN_TIMEOUT_MS, + ); + continue; + } + if (this.sessions.get(activeSessionId) !== state || !state.clients.has(client)) { + finishClientSnapshotStreaming(client, activeSessionId); + transcript.dispose?.(); + continue; + } + if (purpose === "replacement") { + await this.writeWorkerSnapshotRecord( + client, + { + type: "session_replaced", + activeSessionId, + state: result.snapshot.state, + messages: [], + snapshotFollows: true, + meta: createDaemonEventMeta( + activeSessionId, + state.lastEventSequence, + undefined, + state.eventGeneration, + ), + }, + "replacement", + snapshotAttemptSignal, + ); } await this.streamWorkerSnapshot( client, @@ -6993,7 +7092,7 @@ export class AgentDaemon { }, transcript, purpose === "replacement" ? "replacement" : "catchup", - snapshotSignal, + snapshotAttemptSignal, true, ); continue; @@ -7018,15 +7117,26 @@ export class AgentDaemon { for (const remaining of pending.slice(index + 1)) { this.queueClientCatchup(client, remaining.activeSessionId, remaining.purpose); } - return "retry-later"; + return "drained"; } } catch (error) { - for (const remaining of pending.slice(index)) { - this.queueClientCatchup(client, remaining.activeSessionId, remaining.purpose); - } + const snapshotError = error instanceof Error ? error : new Error(String(error)); this.log(`could not catch up client ${client.id} for ${activeSessionId}: ${String(error)}`); - this.scheduleClientCatchupRetry(client); - return "retry-later"; + if (!client.socket.destroyed && client.attachedActiveSessionIds.has(activeSessionId)) { + await this.writeWorkerSnapshotRecord( + client, + { + type: "session_snapshot_failed", + activeSessionId, + snapshotId: createSnapshotTransferId(), + error: snapshotError.message, + purpose, + }, + purpose === "replacement" ? "replacement" : "catchup", + undefined, + WORKER_SNAPSHOT_TERMINAL_DRAIN_TIMEOUT_MS, + ); + } } } return "drained"; @@ -7100,31 +7210,31 @@ export class AgentDaemon { const payloadParts = typeof line === "string" ? [Buffer.from(line)] : Buffer.isBuffer(line) ? [line] : [...line]; const wireParts = client.transport === "private-framed" - ? encodePrivateFrameParts( - { - kind: "outbound", - outboundType: message.type, - ...("id" in message && typeof message.id === "string" ? { requestId: message.id } : {}), - ...(hasDaemonOutboundActiveSessionId(message) ? { activeSessionId: message.activeSessionId } : {}), - ...("snapshotId" in message && typeof message.snapshotId === "string" - ? { snapshotId: message.snapshotId } - : {}), - ...(message.type === "session_event" ? { sessionEventType: message.event.type } : {}), - payloadEncoding, - ...(snapshotPurpose ? { snapshotPurpose } : {}), - }, - payloadParts, - ) + ? this.encodeWorkerPrivateFrameParts(line, message, payloadEncoding, snapshotPurpose) : payloadParts; - let accepted = true; - if (wireParts.length > 1) client.socket.cork?.(); - try { - for (const part of wireParts) accepted = client.socket.write(part) && accepted; - } finally { - if (wireParts.length > 1) client.socket.uncork?.(); - } - if (!accepted) client.backpressured = true; - return accepted; + const writeNow = () => { + let accepted = true; + if (wireParts.length > 1) client.socket.cork?.(); + try { + for (const part of wireParts) accepted = client.socket.write(part) && accepted; + } finally { + if (wireParts.length > 1) client.socket.uncork?.(); + } + if (!accepted) client.backpressured = true; + return accepted; + }; + const previous = client.transport === "private-framed" ? client.privateFrameWriteTail : undefined; + if (!previous) return writeNow(); + const queued = previous + .catch(() => undefined) + .then(() => { + writeNow(); + }); + client.privateFrameWriteTail = queued; + void queued.finally(() => { + if (client.privateFrameWriteTail === queued) client.privateFrameWriteTail = undefined; + }); + return true; } private abortSideQuestionsFor(client: DaemonSocketClient, activeSessionId: string): void { diff --git a/packages/coding-agent/src/modes/daemon/daemon-protocol.ts b/packages/coding-agent/src/modes/daemon/daemon-protocol.ts index 27dfb04aac..2b5b200556 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-protocol.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-protocol.ts @@ -76,8 +76,9 @@ export const DAEMON_COMMAND_ENVELOPE_MIN_PROTOCOL_VERSION = 7; // Revision 23 lets workers query the supervisor agent roster on demand. // Revision 24 adds capability-gated correlated prompt lifecycle and event provenance. // Revision 25 negotiates immutable snapshot transfer identities with session workers. -export const DAEMON_SCHEMA_REVISION = 25; -export const DAEMON_SCHEMA_ID = "protocol-7-schema-25-2c8fb17d3895"; +// Revision 26 correlates snapshot failures that occur before a begin frame can be emitted. +export const DAEMON_SCHEMA_REVISION = 26; +export const DAEMON_SCHEMA_ID = "protocol-7-schema-26-1ac2610350df"; export type DaemonProtocolName = typeof DAEMON_PROTOCOL_NAME; export type DaemonProtocolVersion = number; @@ -413,6 +414,8 @@ export type DaemonCommand = type: "attach"; activeSessionId: string; supportsExtensionUi?: boolean; + /** A retry nonce forces a fresh transcript payload generation after a failed transfer. */ + snapshotGenerationNonce?: string; } & DaemonAttachClientMetadata & DaemonClientEnv & DaemonLaunchEnv) @@ -1064,6 +1067,8 @@ export type DaemonOutbound = activeSessionId: string; snapshotId: string; error: string; + /** Runtime failures may carry purpose when no begin frame could be prepared. */ + purpose?: "attach" | "replacement" | "resync"; } | { type: "session_detached"; activeSessionId: string } | { type: "session_closed"; activeSessionId: string; reason: DaemonSessionClosedReason; meta?: DaemonEventMeta } diff --git a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts index 54bdbc62b2..b3272b837a 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts @@ -143,6 +143,7 @@ import { prepareSnapshotTranscriptCache, SNAPSHOT_TARGET_CHUNK_BYTES, SnapshotTranscriptCache, + type SnapshotTranscriptWireChunk, } from "./snapshot-transcript-cache.js"; import { WorkerRecoveryJournal } from "./worker-recovery-journal.js"; @@ -154,6 +155,20 @@ const WORKER_CONNECT_TIMEOUT_MS = 30_000; const WORKER_REQUEST_TIMEOUT_MS = 24 * 60 * 60 * 1000; const SNAPSHOT_ATTEMPT_TIMEOUT_MS = 30_000; const INPUT_PAUSE_CLEANUP_TIMEOUT_MS = 5_000; + +function createPublicSnapshotTransferId(): string { + return `snapshot-${randomUUID()}`; +} + +function waitForAbortable(promise: Promise, signal: AbortSignal | undefined, message: string): Promise { + if (!signal) return promise; + if (signal.aborted) return Promise.reject(new Error(message)); + return new Promise((resolve, reject) => { + const onAbort = () => reject(new Error(message)); + signal.addEventListener("abort", onAbort, { once: true }); + void promise.then(resolve, reject).finally(() => signal.removeEventListener("abort", onAbort)); + }); +} const UPDATE_RESTART_MUTATION_DRAIN_TIMEOUT_MS = 80_000; const UPDATE_RESTART_WORKER_REQUEST_TIMEOUT_MS = 90_000; // The whole pre-commit prepare (drain + worker fencing) must finish inside the @@ -357,6 +372,14 @@ interface WorkerAttachData { releaseTranscript?: () => void; } +interface SupervisorTranscriptPreparation { + result: DaemonAttachResult; + promise: Promise; + controller: AbortController; + workerClient: DaemonWorkerClient | undefined; + stopRevision: number; +} + interface SupervisorPromptAdmission { client: DaemonSocketClient; activeSessionId: string; @@ -708,10 +731,7 @@ export class DaemonSupervisor { private commandJournal!: CommandRecoveryJournal; private readonly streamReconstructor = new CompactAssistantStreamReconstructor(); private readonly compactCatchupInProgress = new Set(); - private transcriptPreparations?: WeakMap< - ResidentWorker, - Map }> - >; + private transcriptPreparations?: WeakMap>; private readonly pendingSessionNames = new Set(); private readonly catalog: DaemonCatalogClient; private readonly settingsManager: SettingsManager; @@ -1794,7 +1814,7 @@ export class DaemonSupervisor { case "attach": { const attachmentEpoch = this.advanceAttachmentEpoch(client, command.activeSessionId); const requestedCapabilities = normalizeCapabilities(command.capabilities, command.supportsExtensionUi); - const releaseSnapshotReservation = requestedCapabilities.has("chunked_snapshot") + let releaseSnapshotReservation = requestedCapabilities.has("chunked_snapshot") ? this.reserveSnapshotStream(client, command.activeSessionId) : undefined; try { @@ -1803,6 +1823,14 @@ export class DaemonSupervisor { epoch: attachmentEpoch, }); if (client.capabilities.has("chunked_snapshot")) { + if (releaseSnapshotReservation) { + releaseSnapshotReservation = this.remapSnapshotStreamReservation( + client, + command.activeSessionId, + attached.result.activeSessionId, + releaseSnapshotReservation, + ); + } const transcript = attached.transcript; if (!transcript) { throw new Error("Session worker did not provide a snapshot transcript"); @@ -3087,6 +3115,7 @@ export class DaemonSupervisor { if (worker.client !== client) { return; } + this.abortTranscriptPreparations(worker, error); worker.client = undefined; this.invalidateWorkerSessionInputPauses(worker, "Session worker disconnected while input was paused"); const interrupted = new Map>(); @@ -3277,30 +3306,64 @@ export class DaemonSupervisor { } } + private deliverPurposeBearingSnapshotFailure( + activeSessionId: string, + purpose: "replacement" | "catchup" | "resync", + error: Error, + ): void { + for (const client of this.clients) { + if (!client.attachedActiveSessionIds.has(activeSessionId) || !client.capabilities.has("chunked_snapshot")) { + continue; + } + void this.writeSnapshotRecord( + client, + { + type: "session_snapshot_failed", + activeSessionId, + snapshotId: createPublicSnapshotTransferId(), + error: error.message, + purpose: purpose === "replacement" ? "replacement" : "resync", + }, + AbortSignal.timeout(1_000), + ); + } + } + private async deliverWorkerSnapshotFailure( client: DaemonSocketClient, result: DaemonAttachResult, - snapshotId: string, + _payloadSnapshotId: string, purpose: "replacement" | "catchup" | "resync", error: Error, ): Promise { if (client.socket.destroyed || !client.attachedActiveSessionIds.has(result.activeSessionId)) return; + const transferId = createPublicSnapshotTransferId(); const { messages: _messages, ...snapshot } = result.snapshot; - await this.writeSnapshotRecord(client, { - type: "session_snapshot_begin", - activeSessionId: result.activeSessionId, - snapshotId, - snapshot, - messageCount: result.snapshotStream?.messageCount ?? result.snapshot.summary.messageCount, - targetChunkBytes: result.snapshotStream?.targetChunkBytes ?? SNAPSHOT_TARGET_CHUNK_BYTES, - purpose: purpose === "replacement" ? "replacement" : "resync", - }); - await this.writeSnapshotRecord(client, { - type: "session_snapshot_failed", - activeSessionId: result.activeSessionId, - snapshotId, - error: error.message, - }); + const terminalSignal = AbortSignal.timeout(1_000); + await this.writeSnapshotRecord( + client, + { + type: "session_snapshot_begin", + activeSessionId: result.activeSessionId, + snapshotId: transferId, + snapshot, + messageCount: result.snapshotStream?.messageCount ?? result.snapshot.summary.messageCount, + targetChunkBytes: result.snapshotStream?.targetChunkBytes ?? SNAPSHOT_TARGET_CHUNK_BYTES, + purpose: purpose === "replacement" ? "replacement" : "resync", + }, + terminalSignal, + ); + await this.writeSnapshotRecord( + client, + { + type: "session_snapshot_failed", + activeSessionId: result.activeSessionId, + snapshotId: transferId, + error: error.message, + purpose: purpose === "replacement" ? "replacement" : "resync", + }, + terminalSignal, + ); } private failCorrelatedSnapshotAttempt( @@ -3336,6 +3399,10 @@ export class DaemonSupervisor { } return; } + if (snapshotPurpose === "replacement" || snapshotPurpose === "catchup" || snapshotPurpose === "resync") { + this.deliverPurposeBearingSnapshotFailure(activeSessionId, snapshotPurpose, error); + return; + } this.log(`Ignored stale snapshot ${snapshotId} failure for ${activeSessionId}: ${error.message}`); } @@ -3356,11 +3423,14 @@ export class DaemonSupervisor { return; } generation.retired = true; - this.clearSnapshotGenerationDeadline(generation); this.settleSnapshotDuplicateValidation(generation); if (generation.incoming) { + // A superseding generation cannot remove the only liveness deadline from + // an older begin that may never receive its end frame. + this.armSnapshotGenerationDeadline(worker, activeSessionId, generation); return; } + this.clearSnapshotGenerationDeadline(generation); this.deleteSnapshotGeneration(worker, activeSessionId, generation); expectedTranscript.dispose(); } @@ -4227,8 +4297,10 @@ export class DaemonSupervisor { "attach_snapshot", "event_sequence", "slim_attach", - "correlated_prompt_lifecycle_v1", ]; + if (workerClient.supportsServerCapability?.("correlated_prompt_lifecycle_v1")) { + workerCapabilities.push("correlated_prompt_lifecycle_v1"); + } if ( client.capabilities.has("chunked_snapshot") && workerClient.supportsServerCapability?.("immutable_snapshot_transfer_v1") @@ -4239,6 +4311,7 @@ export class DaemonSupervisor { { type: "attach", activeSessionId, + snapshotGenerationNonce: randomUUID(), capabilities: workerCapabilities, supportsExtensionUi: false, env: command.env ?? collectDaemonClientEnv(), @@ -4311,7 +4384,14 @@ export class DaemonSupervisor { if (client.socket.destroyed) { throw new Error("Daemon client disconnected during snapshot preparation"); } - transcript = await this.getOrCreateTranscriptCache(match.worker, result); + const preparationSignal = attachmentFence + ? this.attachmentAbortSignal(client, attachmentFence.selector, attachmentFence.epoch) + : undefined; + transcript = await waitForAbortable( + this.getOrCreateTranscriptCache(match.worker, result), + preparationSignal, + `Attachment changed while preparing snapshot for ${activeSessionId}`, + ); assertAttachmentCurrent(); if (client.socket.destroyed) { throw new Error("Daemon client disconnected during snapshot preparation"); @@ -4336,7 +4416,9 @@ export class DaemonSupervisor { } } const snapshot = client.capabilities.has("correlated_prompt_lifecycle_v1") - ? result.snapshot + ? result.snapshot.promptLifecycles + ? result.snapshot + : { ...result.snapshot, promptLifecycles: { records: [], expired: [] } } : withoutCorrelatedPromptLifecycleSnapshot(result.snapshot); const publicResult: DaemonAttachResult = { ...result, @@ -4439,6 +4521,13 @@ export class DaemonSupervisor { return loaded; } + private abortTranscriptPreparations(worker: ResidentWorker, reason: Error): void { + const preparations = this.transcriptPreparations?.get(worker); + if (!preparations) return; + for (const preparation of preparations.values()) preparation.controller.abort(reason); + preparations.clear(); + } + private async getOrCreateTranscriptCache( worker: ResidentWorker, result: DaemonAttachResult, @@ -4475,22 +4564,41 @@ export class DaemonSupervisor { this.retireWorkerSnapshotCache(worker, activeSessionId, existing); } const snapshotId = `snapshot-${randomUUID()}`; + const controller = new AbortController(); + const workerClient = worker.client; + const stopRevision = worker.stopRevision; const promise = prepareSnapshotTranscriptCache({ activeSessionId, snapshotId, messages: result.snapshot.messages, cacheRoot: this.snapshotCacheRoot, targetChunkBytes: SNAPSHOT_TARGET_CHUNK_BYTES, - signal: AbortSignal.timeout(SNAPSHOT_ATTEMPT_TIMEOUT_MS), + signal: AbortSignal.any([controller.signal, AbortSignal.timeout(SNAPSHOT_ATTEMPT_TIMEOUT_MS)]), }); - preparations.set(activeSessionId, { result, promise }); + const preparation: SupervisorTranscriptPreparation = { + result, + promise, + controller, + workerClient, + stopRevision, + }; + preparations.set(activeSessionId, preparation); let transcript: SnapshotTranscriptCache; try { transcript = await promise; } finally { - if (preparations.get(activeSessionId)?.promise === promise) { - preparations.delete(activeSessionId); - } + if (preparations.get(activeSessionId) === preparation) preparations.delete(activeSessionId); + } + const registeredWorker = this.workers?.get(worker.descriptor.workerId); + if ( + controller.signal.aborted || + this.shuttingDown || + (registeredWorker !== undefined && registeredWorker !== worker) || + worker.client !== workerClient || + worker.stopRevision !== stopRevision + ) { + transcript.dispose(); + throw new Error("Session worker changed during snapshot transcript preparation"); } const cachedResult = { ...result, @@ -4513,13 +4621,14 @@ export class DaemonSupervisor { private createStreamedAttachResult( result: DaemonAttachResult, transcript: SnapshotTranscriptCache, + transferId = createPublicSnapshotTransferId(), ): DaemonAttachResult { return { ...result, messages: result.messages ? [] : undefined, snapshot: { ...result.snapshot, messages: [] }, snapshotStream: { - id: transcript.snapshotId, + id: transferId, messageCount: result.snapshot.summary.messageCount, targetChunkBytes: transcript.targetChunkBytes, }, @@ -4548,6 +4657,8 @@ export class DaemonSupervisor { attachmentFence?.selector ?? result.activeSessionId, attachmentFence?.epoch ?? attachmentEpoch, ); + const deadlineSignal = AbortSignal.timeout(SNAPSHOT_ATTEMPT_TIMEOUT_MS); + const attemptSignal = attachmentSignal ? AbortSignal.any([attachmentSignal, deadlineSignal]) : deadlineSignal; if (!stream || client.socket.destroyed || !isAttachmentCurrent()) { releaseSnapshotReservation(); releaseTranscript(); @@ -4556,77 +4667,77 @@ export class DaemonSupervisor { const { messages: _messages, ...snapshotHeader } = result.snapshot; try { if ( - !(await this.writeSnapshotRecord(client, { - type: "session_snapshot_begin", - activeSessionId: result.activeSessionId, - snapshotId: stream.id, - snapshot: snapshotHeader, - messageCount: stream.messageCount, - targetChunkBytes: stream.targetChunkBytes, - purpose, - })) + !(await this.writeSnapshotRecord( + client, + { + type: "session_snapshot_begin", + activeSessionId: result.activeSessionId, + snapshotId: stream.id, + snapshot: snapshotHeader, + messageCount: stream.messageCount, + targetChunkBytes: stream.targetChunkBytes, + purpose, + }, + attemptSignal, + )) ) { return; } let chunkCount = 0; while (true) { - if (!isAttachmentCurrent() || attachmentSignal?.aborted) return; - let chunk: Buffer | undefined; + if (!isAttachmentCurrent() || attemptSignal.aborted) return; + let chunk: SnapshotTranscriptWireChunk | undefined; try { - if (!attachmentSignal) { - chunk = await transcript.waitForChunk(chunkCount); - } else { - chunk = await new Promise((resolveChunk, rejectChunk) => { - const onAbort = () => resolveChunk(undefined); - attachmentSignal.addEventListener("abort", onAbort, { once: true }); - void transcript - .waitForChunk(chunkCount) - .then(resolveChunk, rejectChunk) - .finally(() => { - attachmentSignal.removeEventListener("abort", onAbort); - }); - }); - } + chunk = await this.waitForSnapshotTransferChunk(transcript, stream.id, chunkCount, attemptSignal); } catch (error) { const streamError = error instanceof Error ? error : new Error(String(error)); - this.failCorrelatedSnapshotAttempt( - worker, - result.activeSessionId, - stream.id, - streamError, - purpose, - false, - ); + const cacheGeneration = this.currentSnapshotGeneration(worker, result.activeSessionId); + if (cacheGeneration?.transcript === transcript) { + this.failCorrelatedSnapshotAttempt( + worker, + result.activeSessionId, + transcript.snapshotId, + streamError, + purpose, + false, + ); + } throw streamError; } if (!isAttachmentCurrent() || attachmentSignal?.aborted) return; - if (!chunk) { - break; - } - if (!(await this.writeSnapshotBuffer(client, chunk))) { - return; - } + if (attemptSignal.aborted) throw new Error(`Snapshot ${stream.id} timed out`); + if (!chunk) break; + if (!(await this.writeSnapshotBuffer(client, chunk, attemptSignal))) return; chunkCount++; } - const completed = await this.writeSnapshotRecord(client, { - type: "session_snapshot_end", - activeSessionId: result.activeSessionId, - snapshotId: stream.id, - chunkCount, - lastEventSequence: result.lastEventSequence, - lastEventCursor: result.lastEventCursor, - }); + const completed = await this.writeSnapshotRecord( + client, + { + type: "session_snapshot_end", + activeSessionId: result.activeSessionId, + snapshotId: stream.id, + chunkCount, + lastEventSequence: result.lastEventSequence, + lastEventCursor: result.lastEventCursor, + }, + attemptSignal, + ); if (!completed) return; } catch (error) { const streamError = error instanceof Error ? error : new Error(String(error)); if (!client.socket.destroyed && isAttachmentCurrent()) { try { - const delivered = await this.writeSnapshotRecord(client, { - type: "session_snapshot_failed", - activeSessionId: result.activeSessionId, - snapshotId: stream.id, - error: streamError.message, - }); + const delivered = await this.writeSnapshotRecord( + client, + { + type: "session_snapshot_failed", + activeSessionId: result.activeSessionId, + snapshotId: stream.id, + error: streamError.message, + purpose, + }, + AbortSignal.timeout(1_000), + ); if (!delivered && !client.socket.destroyed) { client.socket.destroy(streamError); } else if (delivered) { @@ -4643,6 +4754,18 @@ export class DaemonSupervisor { } } + private remapSnapshotStreamReservation( + client: DaemonSocketClient, + selector: string, + activeSessionId: string, + releaseSelector: () => void, + ): () => void { + if (selector === activeSessionId) return releaseSelector; + const releaseCanonical = this.reserveSnapshotStream(client, activeSessionId); + releaseSelector(); + return releaseCanonical; + } + private reserveSnapshotStream(client: DaemonSocketClient, activeSessionId: string): () => void { client.snapshotStreaming = true; client.snapshotActiveSessionIds ??= new Set(); @@ -4677,34 +4800,57 @@ export class DaemonSupervisor { }; } - private writeSnapshotRecord(client: DaemonSocketClient, message: DaemonOutbound): Promise { - return this.writeSnapshotBuffer(client, Buffer.from(serializeJsonLine(message))); + private waitForSnapshotTransferChunk( + transcript: SnapshotTranscriptCache, + transferId: string, + index: number, + signal: AbortSignal, + ): Promise { + if (signal.aborted) return Promise.resolve(undefined); + return new Promise((resolve, reject) => { + const onAbort = () => resolve(undefined); + signal.addEventListener("abort", onAbort, { once: true }); + void transcript + .waitForTransferChunk(transferId, index) + .then(resolve, reject) + .finally(() => signal.removeEventListener("abort", onAbort)); + }); } - private async writeSnapshotBuffer(client: DaemonSocketClient, buffer: Uint8Array): Promise { - if (client.socket.destroyed) { - return false; - } - if (this.writeSerialized(client, buffer)) { - return true; - } + private writeSnapshotRecord( + client: DaemonSocketClient, + message: DaemonOutbound, + signal?: AbortSignal, + ): Promise { + return this.writeSnapshotBuffer(client, Buffer.from(serializeJsonLine(message)), signal); + } + + private async writeSnapshotBuffer( + client: DaemonSocketClient, + buffer: SnapshotTranscriptWireChunk, + signal?: AbortSignal, + ): Promise { + if (signal?.aborted || client.socket.destroyed) return false; + if (this.writeSerialized(client, buffer)) return true; return new Promise((resolveDrain) => { let settled = false; const finish = (value: boolean) => { - if (settled) { - return; - } + if (settled) return; settled = true; client.socket.off("drain", onDrain); client.socket.off("close", onClose); client.socket.off("error", onClose); + signal?.removeEventListener("abort", onAbort); resolveDrain(value); }; const onDrain = () => finish(true); const onClose = () => finish(false); + const onAbort = () => finish(false); client.socket.once("drain", onDrain); client.socket.once("close", onClose); client.socket.once("error", onClose); + signal?.addEventListener("abort", onAbort, { once: true }); + if (signal?.aborted || client.socket.destroyed) finish(false); }); } @@ -4953,7 +5099,7 @@ export class DaemonSupervisor { this.snapshotGenerationsFor(worker, activeSessionId).set(begin.snapshotId, generation); } generation.result = result; - generation.begin = Buffer.from(frame.payload); + generation.begin = frame.payload; generation.end = undefined; generation.incoming = true; generation.receivedMessageCount = 0; @@ -5034,10 +5180,10 @@ export class DaemonSupervisor { if (receivedMessageCount > (generation.result.snapshotStream?.messageCount ?? Number.MAX_SAFE_INTEGER)) { throw new Error(`Snapshot ${generation.transcript.snapshotId} exceeded its declared message count`); } - generation.transcript.appendEncodedChunk(Buffer.from(frame.payload)); + generation.transcript.appendEncodedChunk(frame.payload); generation.receivedMessageCount = receivedMessageCount; } else { - if (!generation.transcript.readChunk(duplicateIndex).equals(Buffer.from(frame.payload))) { + if (!generation.transcript.readChunk(duplicateIndex).equals(frame.payload)) { throw new Error(`Duplicate snapshot ${generation.transcript.snapshotId} did not match cached bytes`); } generation.duplicateChunkIndex = duplicateIndex + 1; @@ -5115,7 +5261,7 @@ export class DaemonSupervisor { ) { throw new Error(`Snapshot ${transcript.snapshotId} ended with the wrong message count`); } - generation.end = Buffer.from(frame.payload); + generation.end = frame.payload; transcript.markComplete(); } else { if (end.chunkCount !== duplicateChunkCount || !generation.end?.equals(frame.payload)) { @@ -5508,12 +5654,17 @@ export class DaemonSupervisor { ); } } else if (client.attachedActiveSessionIds.has(activeSessionId)) { - await this.writeSnapshotRecord(client, { - type: "session_snapshot_failed", - activeSessionId, - snapshotId: `snapshot-${randomUUID()}`, - error: snapshotError.message, - }); + await this.writeSnapshotRecord( + client, + { + type: "session_snapshot_failed", + activeSessionId, + snapshotId: createPublicSnapshotTransferId(), + error: snapshotError.message, + purpose, + }, + AbortSignal.timeout(1_000), + ); } this.log(`Failed to catch up client ${client.id} for ${activeSessionId}: ${String(error)}`); } @@ -5785,6 +5936,7 @@ export class DaemonSupervisor { clearTimeout(worker.ownerCleanupTimer); worker.ownerCleanupTimer = undefined; } + this.abortTranscriptPreparations(worker, new Error("Session worker stopped during snapshot preparation")); if (!recoveryCleanup) { worker.stopRevision++; } @@ -6079,14 +6231,17 @@ export class DaemonSupervisor { } } - private writeSerialized(client: DaemonSocketClient, line: string | Uint8Array): boolean { - if (client.socket.destroyed) { - return false; - } - const accepted = client.socket.write(line); - if (!accepted) { - client.backpressured = true; + private writeSerialized(client: DaemonSocketClient, line: string | SnapshotTranscriptWireChunk): boolean { + if (client.socket.destroyed) return false; + const parts = typeof line === "string" || Buffer.isBuffer(line) ? [line] : [...line]; + let accepted = true; + if (parts.length > 1) client.socket.cork?.(); + try { + for (const part of parts) accepted = client.socket.write(part) && accepted; + } finally { + if (parts.length > 1) client.socket.uncork?.(); } + if (!accepted) client.backpressured = true; return accepted; } @@ -6153,6 +6308,7 @@ export class DaemonSupervisor { } this.clients.clear(); for (const worker of this.workers.values()) { + this.abortTranscriptPreparations(worker, new Error("Daemon supervisor stopped during snapshot preparation")); if (worker.ownerCleanupTimer) { clearTimeout(worker.ownerCleanupTimer); worker.ownerCleanupTimer = undefined; diff --git a/packages/coding-agent/src/modes/daemon/snapshot-transcript-cache.ts b/packages/coding-agent/src/modes/daemon/snapshot-transcript-cache.ts index f7f8e50921..5dc2092883 100644 --- a/packages/coding-agent/src/modes/daemon/snapshot-transcript-cache.ts +++ b/packages/coding-agent/src/modes/daemon/snapshot-transcript-cache.ts @@ -452,15 +452,12 @@ export class SnapshotTranscriptPayloadCache { if (!chunk) throw new Error(`Unknown snapshot transcript chunk: ${index}`); if (chunk.buffer) return chunk.buffer; if (!chunk.path) throw new Error(`Snapshot transcript chunk ${index} has no backing storage`); - chunk.readPromise ??= this.io.readFile(chunk.path).then((buffer) => { - chunk.buffer = buffer; - return buffer; - }); + chunk.readPromise ??= this.io.readFile(chunk.path); + const reading = chunk.readPromise; try { - return await chunk.readPromise; - } catch (error) { - chunk.readPromise = undefined; - throw error; + return await reading; + } finally { + if (chunk.readPromise === reading) chunk.readPromise = undefined; } } @@ -807,6 +804,27 @@ export class SnapshotTranscriptCache { }); } + async waitForTransferChunk(snapshotId: string, index: number): Promise { + const chunk = await this.waitForChunk(index); + if (!chunk || snapshotId === this.snapshotId) return chunk; + const messagesMarker = Buffer.from('"messages":['); + const markerIndex = chunk.indexOf(messagesMarker); + const payloadStart = markerIndex + messagesMarker.length; + const hasNewline = chunk.at(-1) === 0x0a; + const payloadEnd = chunk.length - (hasNewline ? 3 : 2); + if ( + markerIndex < 0 || + payloadEnd < payloadStart || + chunk.subarray(payloadEnd, payloadEnd + 2).toString() !== "]}" + ) { + throw new Error(`Snapshot transcript ${this.snapshotId} has an invalid chunk envelope`); + } + const prefix = + `{"type":"session_snapshot_chunk","activeSessionId":${JSON.stringify(this.activeSessionId)},` + + `"snapshotId":${JSON.stringify(snapshotId)},"index":${index},"messages":[`; + return Buffer.concat([Buffer.from(prefix), chunk.subarray(payloadStart, payloadEnd), Buffer.from("]}\n")]); + } + retain(): () => void { if (this.disposed) throw new Error(`Snapshot transcript ${this.snapshotId} was disposed`); const releasePayload = this.preparedPayload?.retain(); diff --git a/packages/coding-agent/src/modes/session-worker/private-framing.ts b/packages/coding-agent/src/modes/session-worker/private-framing.ts index 3dca2672f1..993ed61031 100644 --- a/packages/coding-agent/src/modes/session-worker/private-framing.ts +++ b/packages/coding-agent/src/modes/session-worker/private-framing.ts @@ -68,7 +68,11 @@ export function encodePrivateFrame( } export class PrivateFrameDecoder { - private buffered: Buffer = Buffer.alloc(0); + private readonly buffers: Buffer[] = []; + private headOffset = 0; + private totalBufferedBytes = 0; + private totalCoalescedBytes = 0; + private pendingLengths?: { header: number; payload: number }; constructor( private readonly validateHeader: PrivateFrameHeaderValidator, @@ -76,36 +80,42 @@ export class PrivateFrameDecoder { ) {} get bufferedBytes(): number { - return this.buffered.length; + return this.totalBufferedBytes; + } + + get coalescedBytes(): number { + return this.totalCoalescedBytes; } push(chunk: Uint8Array): PrivateFrame[] { if (chunk.length > 0) { - const buffer = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk); - this.buffered = this.buffered.length === 0 ? buffer : Buffer.concat([this.buffered, buffer]); + this.buffers.push( + Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk.buffer, chunk.byteOffset, chunk.byteLength), + ); + this.totalBufferedBytes += chunk.length; } const frames: PrivateFrame[] = []; - let offset = 0; - while (this.buffered.length - offset >= FRAME_PREFIX_BYTES) { - const headerLength = this.buffered.readUInt32BE(offset); - const payloadLength = this.buffered.readUInt32BE(offset + 4); - assertFrameLength("header length", headerLength, this.limits.maxHeaderBytes); - assertFrameLength("payload length", payloadLength, this.limits.maxPayloadBytes); - if (headerLength === 0) { - throw new Error("Private frame header cannot be empty"); - } - - const frameLength = FRAME_PREFIX_BYTES + headerLength + payloadLength; - if (this.buffered.length - offset < frameLength) { - break; + while (true) { + if (!this.pendingLengths) { + if (this.totalBufferedBytes < FRAME_PREFIX_BYTES) break; + const prefix = this.consume(FRAME_PREFIX_BYTES); + const headerLength = prefix.readUInt32BE(0); + const payloadLength = prefix.readUInt32BE(4); + assertFrameLength("header length", headerLength, this.limits.maxHeaderBytes); + assertFrameLength("payload length", payloadLength, this.limits.maxPayloadBytes); + if (headerLength === 0) throw new Error("Private frame header cannot be empty"); + this.pendingLengths = { header: headerLength, payload: payloadLength }; } - const headerStart = offset + FRAME_PREFIX_BYTES; - const payloadStart = headerStart + headerLength; + const { header: headerLength, payload: payloadLength } = this.pendingLengths; + if (this.totalBufferedBytes < headerLength + payloadLength) break; + const header = this.consume(headerLength); + const payload = this.consume(payloadLength); + this.pendingLengths = undefined; let decoded: unknown; try { - decoded = JSON.parse(this.buffered.toString("utf8", headerStart, payloadStart)); + decoded = JSON.parse(header.toString("utf8")); } catch (error) { throw new Error( `Invalid private frame header JSON: ${error instanceof Error ? error.message : String(error)}`, @@ -114,27 +124,54 @@ export class PrivateFrameDecoder { if (!isObjectHeader(decoded) || !this.validateHeader(decoded)) { throw new Error("Invalid private frame routing header"); } - - frames.push({ - header: decoded, - payload: Buffer.from(this.buffered.subarray(payloadStart, payloadStart + payloadLength)), - }); - offset += frameLength; - } - - if (offset > 0) { - this.buffered = Buffer.from(this.buffered.subarray(offset)); + frames.push({ header: decoded, payload }); } return frames; } finish(): void { - if (this.buffered.length !== 0) { - throw new Error(`Private frame channel ended with ${this.buffered.length} incomplete bytes`); + if (this.totalBufferedBytes !== 0 || this.pendingLengths) { + throw new Error( + `Private frame channel ended with ${this.totalBufferedBytes + (this.pendingLengths ? FRAME_PREFIX_BYTES : 0)} incomplete bytes`, + ); } } -} + private consume(length: number): Buffer { + if (length === 0) return Buffer.alloc(0); + const first = this.buffers[0]; + if (!first) throw new Error("Private frame decoder buffer underflow"); + const firstAvailable = first.length - this.headOffset; + if (firstAvailable >= length) { + const value = first.subarray(this.headOffset, this.headOffset + length); + this.headOffset += length; + this.totalBufferedBytes -= length; + if (this.headOffset === first.length) { + this.buffers.shift(); + this.headOffset = 0; + } + return value; + } + const value = Buffer.allocUnsafe(length); + this.totalCoalescedBytes += length; + let written = 0; + while (written < length) { + const current = this.buffers[0]; + if (!current) throw new Error("Private frame decoder buffer underflow"); + const available = current.length - this.headOffset; + const selected = Math.min(available, length - written); + current.copy(value, written, this.headOffset, this.headOffset + selected); + written += selected; + this.headOffset += selected; + this.totalBufferedBytes -= selected; + if (this.headOffset === current.length) { + this.buffers.shift(); + this.headOffset = 0; + } + } + return value; + } +} export type PrivateFrameListener = (frame: PrivateFrame) => void; export class PrivateFramedChannel { diff --git a/packages/coding-agent/test/agent-connection-daemon.test.ts b/packages/coding-agent/test/agent-connection-daemon.test.ts index 758bb3425f..aabbea4395 100644 --- a/packages/coding-agent/test/agent-connection-daemon.test.ts +++ b/packages/coding-agent/test/agent-connection-daemon.test.ts @@ -3508,7 +3508,7 @@ describe("DaemonAgentConnection", () => { state: createConnectionState("active-target", "session-target"), messages: [{ role: "user", content: "target", timestamp: 1 }], }); - const emitTransfer = (snapshotId: string, fail: boolean) => { + const emitTransfer = (snapshotId: string, fail: boolean, purpose: "attach" | "replacement") => { const { messages: _messages, ...snapshot } = target.snapshot; setImmediate(() => { fakeClient.emitMessage({ @@ -3518,7 +3518,7 @@ describe("DaemonAgentConnection", () => { snapshot, messageCount: 1, targetChunkBytes: 512 * 1024, - purpose: "replacement", + purpose, }); if (fail) { fakeClient.emitMessage({ @@ -3551,8 +3551,8 @@ describe("DaemonAgentConnection", () => { snapshotStream: { id: snapshotId, messageCount: 1, targetChunkBytes: 512 * 1024 }, }; }; - fakeClient.reattachResultFactory = () => emitTransfer("reattach-g1", true); - fakeClient.attachResultFactory = () => emitTransfer("reattach-g2", false); + fakeClient.reattachResultFactory = () => emitTransfer("reattach-g1", true, "replacement"); + fakeClient.attachResultFactory = () => emitTransfer("reattach-g2", false, "attach"); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); diff --git a/packages/coding-agent/test/daemon-mode.test.ts b/packages/coding-agent/test/daemon-mode.test.ts index 82290ff5d4..9deb8527d8 100644 --- a/packages/coding-agent/test/daemon-mode.test.ts +++ b/packages/coding-agent/test/daemon-mode.test.ts @@ -3115,7 +3115,7 @@ describe("daemon mode helpers", () => { expect(client.catchupActiveSessionIds).toEqual(new Set()); }); - it("automatically retries every pending catch-up after snapshot creation rejects", async () => { + it("terminates a failed catch-up attempt without an unbounded self-retry", async () => { const daemon = new AgentDaemon("/tmp/prime-agent-test.sock", { defaultSessionConfig: { agentDir: "/tmp/prime-agent-test-agent", cwd: "/tmp" }, createRuntime: async () => { @@ -3164,23 +3164,20 @@ describe("daemon mode helpers", () => { await internals.catchUpBackpressuredClient(client); - expect(client.catchupActiveSessionIds).toEqual( - new Set([firstState.activeSessionId, secondState.activeSessionId]), - ); - expect(client.catchupPurposes).toEqual( - new Map([ - [firstState.activeSessionId, "replacement"], - [secondState.activeSessionId, "resync"], - ]), - ); - expect(createAttachResult).toHaveBeenCalledOnce(); - - await vi.waitFor(() => expect(createAttachResult).toHaveBeenCalledTimes(3)); - expect(client.catchupActiveSessionIds).toEqual(new Set()); expect(client.catchupPurposes).toEqual(new Map()); - const messages = write.mock.calls.map(([data]) => JSON.parse(String(data)) as { type: string }); - expect(messages.map((message) => message.type)).toEqual(["session_replaced", "session_resynced"]); + expect(createAttachResult).toHaveBeenCalledTimes(2); + const messages = write.mock.calls.map( + ([data]) => JSON.parse(String(data)) as { type: string; activeSessionId?: string; purpose?: string }, + ); + expect(messages).toEqual([ + expect.objectContaining({ + type: "session_snapshot_failed", + activeSessionId: firstState.activeSessionId, + purpose: "replacement", + }), + expect.objectContaining({ type: "session_resynced", activeSessionId: secondState.activeSessionId }), + ]); }); it("clears a scheduled catch-up retry when the client disconnects", async () => { @@ -3220,10 +3217,8 @@ describe("daemon mode helpers", () => { await internals.catchUpBackpressuredClient(client); - expect(client.catchupRetryTimer).toBeDefined(); socketState.destroyed = true; socket.emit("close"); - expect(client.catchupRetryTimer).toBeUndefined(); await new Promise((resolve) => setTimeout(resolve, 300)); expect(createAttachResult).toHaveBeenCalledOnce(); }); diff --git a/packages/coding-agent/test/daemon-supervisor-process.test.ts b/packages/coding-agent/test/daemon-supervisor-process.test.ts index 5d973738df..790db2091b 100644 --- a/packages/coding-agent/test/daemon-supervisor-process.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-process.test.ts @@ -1,6 +1,7 @@ import { type ChildProcess, spawn } from "node:child_process"; import { randomUUID } from "node:crypto"; import { existsSync, mkdirSync, mkdtempSync, readdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { createConnection, createServer, type Server, type Socket } from "node:net"; import { tmpdir } from "node:os"; import { dirname, join, resolve } from "node:path"; import { afterEach, describe, expect, it } from "vitest"; @@ -198,7 +199,114 @@ async function connectEventually(socketPath: string, child?: ChildProcess): Prom await new Promise((resolveDelay) => setTimeout(resolveDelay, 25)); } } - throw new Error(`Timed out waiting for supervisor: ${String(lastError)}`); + const diagnostics = child ? childDiagnostics.get(child) : undefined; + throw new Error( + `Timed out waiting for supervisor: ${String(lastError)}\nstdout:\n${diagnostics?.stdout ?? ""}\nstderr:\n${diagnostics?.stderr ?? ""}`, + ); +} + +async function createSnapshotRetryProxy( + proxyPath: string, + targetPath: string, +): Promise<{ + server: Server; + attachRequests: () => number; + transferIds: () => string[]; +}> { + let attachCount = 0; + const observedTransferIds: string[] = []; + const server = createServer((downstream: Socket) => { + const upstream = createConnection(targetPath); + let requestBuffer = Buffer.alloc(0); + let responseBuffer = Buffer.alloc(0); + let firstTransferId: string | undefined; + let firstChunk: Buffer | undefined; + let firstEnd: Buffer | undefined; + downstream.on("data", (chunk: Buffer) => { + requestBuffer = Buffer.concat([requestBuffer, chunk]); + while (true) { + const newline = requestBuffer.indexOf(0x0a); + if (newline < 0) break; + const line = requestBuffer.subarray(0, newline); + requestBuffer = requestBuffer.subarray(newline + 1); + try { + const parsed = JSON.parse(line.toString("utf8")) as { + type?: unknown; + command?: { type?: unknown }; + }; + if (parsed.type === "attach" || parsed.command?.type === "attach") attachCount++; + } catch { + // The real daemon owns command validation. + } + upstream.write(Buffer.concat([line, Buffer.from("\n")])); + } + }); + upstream.on("data", (chunk: Buffer) => { + responseBuffer = Buffer.concat([responseBuffer, chunk]); + while (true) { + const newline = responseBuffer.indexOf(0x0a); + if (newline < 0) break; + const original = Buffer.from(responseBuffer.subarray(0, newline + 1)); + responseBuffer = responseBuffer.subarray(newline + 1); + let parsed: { type?: unknown; snapshotId?: unknown } | undefined; + try { + parsed = JSON.parse(original.toString("utf8")) as { type?: unknown; snapshotId?: unknown }; + } catch { + downstream.write(original); + continue; + } + if (parsed.type === "session_snapshot_begin" && typeof parsed.snapshotId === "string") { + if (!observedTransferIds.includes(parsed.snapshotId)) observedTransferIds.push(parsed.snapshotId); + firstTransferId ??= parsed.snapshotId; + } + if ( + parsed.type === "session_snapshot_chunk" && + typeof parsed.snapshotId === "string" && + parsed.snapshotId === firstTransferId && + !firstChunk + ) { + firstChunk = original; + parsed.snapshotId = `${parsed.snapshotId}-corrupted`; + downstream.write(`${JSON.stringify(parsed)}\n`); + continue; + } + if ( + parsed.type === "session_snapshot_end" && + typeof parsed.snapshotId === "string" && + parsed.snapshotId === firstTransferId + ) { + firstEnd = original; + } + downstream.write(original); + if ( + parsed.type === "session_snapshot_end" && + typeof parsed.snapshotId === "string" && + parsed.snapshotId !== firstTransferId && + firstChunk && + firstEnd + ) { + // Old-generation frames arrive after the fresh attempt's terminal end. + downstream.write(firstChunk); + downstream.write(firstEnd); + firstChunk = undefined; + firstEnd = undefined; + } + } + }); + upstream.on("error", (error) => downstream.destroy(error)); + downstream.on("error", () => upstream.destroy()); + upstream.on("close", () => downstream.destroy()); + downstream.on("close", () => upstream.destroy()); + }); + await new Promise((resolveListen, reject) => { + server.once("error", reject); + server.listen(proxyPath, () => resolveListen()); + }); + return { + server, + attachRequests: () => attachCount, + transferIds: () => [...observedTransferIds], + }; } async function waitForSocketGone(socketPath: string): Promise { @@ -449,6 +557,80 @@ describe("daemon supervisor resident workers", () => { await waitForSocketGone(socketPath); }, 60_000); + it("retries one corrupted real worker snapshot with fresh transfer identity and ignores stale terminal frames", async () => { + const root = tempDir(); + const agentDir = join(root, "agent"); + const projectDir = join(root, "project"); + const sessionDir = join(agentDir, "sessions"); + const suffix = `${process.pid}-${randomUUID().slice(0, 8)}`; + const socketPath = join(tmpdir(), `prime-snapshot-retry-${suffix}.sock`); + const proxyPath = join(tmpdir(), `prime-snapshot-proxy-${suffix}.sock`); + mkdirSync(projectDir, { recursive: true }); + const sessionManager = SessionManager.create(projectDir, sessionDir); + sessionManager.appendMessage({ role: "user", content: "real snapshot retry fixture", timestamp: 1 }); + sessionManager.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "fixture complete" }], + api: "openai-responses", + provider: "faux", + model: "faux", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: 2, + }); + const sessionFile = sessionManager.getSessionFile(); + if (!sessionFile) throw new Error("Snapshot retry fixture did not persist"); + + const supervisor = spawnSupervisor(agentDir, socketPath, projectDir); + const creator = await connectEventually(socketPath, supervisor); + const created = await creator.request({ + type: "create", + sessionPath: sessionFile, + config: { cwd: projectDir, agentDir, sessionDir, noTools: true, noExtensions: true }, + }); + if (!created.success) throw new Error(created.error); + const summary = requireSummary(created.data); + if (!summary.activeSessionId || !summary.workerPid) throw new Error("Snapshot retry worker was incomplete"); + workerPids.add(summary.workerPid); + creator.close(); + + const proxy = await createSnapshotRetryProxy(proxyPath, socketPath); + const client = new DaemonClient(proxyPath); + try { + await client.connect(3_000); + await client.waitForHello(3_000); + const connection = await DaemonAgentConnection.attach(client, summary.activeSessionId, { + supportsExtensionUi: false, + snapshotTimeoutMs: 3_000, + }); + const snapshot = await connection.getInitialSnapshot(); + expect(snapshot.messages).toContainEqual( + expect.objectContaining({ role: "user", content: "real snapshot retry fixture" }), + ); + expect(proxy.attachRequests()).toBe(2); + expect(proxy.transferIds()).toHaveLength(2); + expect(new Set(proxy.transferIds()).size).toBe(2); + await connection.dispose(); + } finally { + client.close(); + await new Promise((resolveClose) => proxy.server.close(() => resolveClose())); + } + + const shutdownClient = await connectEventually(socketPath); + await shutdownClient.request({ type: "shutdown" }); + shutdownClient.close(); + await waitForProcessGone(summary.workerPid); + workerPids.delete(summary.workerPid); + await waitForSocketGone(socketPath); + }, 60_000); + it("keeps client-owned workers hidden and removes them without archiving", async () => { const root = tempDir(); const agentDir = join(root, "agent"); diff --git a/packages/coding-agent/test/daemon-version-compatibility.test.ts b/packages/coding-agent/test/daemon-version-compatibility.test.ts new file mode 100644 index 0000000000..a44ade8b2a --- /dev/null +++ b/packages/coding-agent/test/daemon-version-compatibility.test.ts @@ -0,0 +1,230 @@ +import { type ChildProcess, spawn, spawnSync } from "node:child_process"; +import { randomUUID } from "node:crypto"; +import { mkdirSync, mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join, resolve } from "node:path"; +import { afterEach, describe, expect, it } from "vitest"; +import { ENV_AGENT_DIR } from "../src/config.js"; +import { SessionManager } from "../src/core/session-manager.js"; +import { DaemonAgentConnection } from "../src/modes/agent-connection/daemon-agent-connection.js"; +import { DaemonClient } from "../src/modes/daemon/daemon-client.js"; +import type { SessionSummary } from "../src/modes/daemon/daemon-session-list.js"; + +/** + * Run with PRIME_AGENT_HISTORICAL_081_CLI=. + * The artifact must be built from the pinned git tag v0.8.1, not a protocol mock. + */ +const historicalCli = process.env.PRIME_AGENT_HISTORICAL_081_CLI; +const currentCli = resolve(__dirname, "../dist/bundle/cli.js"); +const compiledCli = process.env.PRIME_AGENT_COMPILED_CLI; +const processTests = historicalCli ? describe : describe.skip; +const children = new Set(); +const roots: string[] = []; + +function scrubbedEnvironment(agentDir: string): NodeJS.ProcessEnv { + const environment = { ...process.env, [ENV_AGENT_DIR]: agentDir, PI_OFFLINE: "1" }; + for (const name of [ + "PRIME_AGENT_INTERNAL_DAEMON_WORKER", + "PRIME_AGENT_INTERNAL_DAEMON_WORKER_TOKEN", + "PRIME_AGENT_INTERNAL_DAEMON_WORKER_ACTIVE_SESSION_ID", + "PRIME_AGENT_INTERNAL_DAEMON_SUPERVISOR_SOCKET", + "PRIME_AGENT_INTERNAL_DAEMON_WORKER_RECOVERY_JOURNAL", + ]) { + delete environment[name]; + } + return environment; +} + +function launch(cli: string, agentDir: string, socketPath: string, cwd: string): ChildProcess { + const child = spawn(process.execPath, [cli, "--mode", "daemon", "--daemon-socket", socketPath, "--offline"], { + cwd, + env: scrubbedEnvironment(agentDir), + stdio: ["ignore", "ignore", "pipe"], + }); + children.add(child); + return child; +} + +async function connectEventually(socketPath: string, process?: ChildProcess): Promise { + const deadline = Date.now() + 15_000; + let lastError: unknown; + while (Date.now() < deadline) { + if (process && (process.exitCode !== null || process.signalCode !== null)) { + throw new Error(`Compatibility supervisor exited early (${process.exitCode ?? process.signalCode})`); + } + const client = new DaemonClient(socketPath); + try { + await client.connect(250); + await client.waitForHello(1_000); + return client; + } catch (error) { + lastError = error; + client.close(); + await new Promise((resolveDelay) => setTimeout(resolveDelay, 25)); + } + } + throw new Error(`Compatibility supervisor did not become ready: ${String(lastError)}`); +} + +async function stopSupervisor(child: ChildProcess): Promise { + if (child.exitCode !== null || child.signalCode !== null) return; + child.kill("SIGTERM"); + await new Promise((resolveExit, reject) => { + const timeout = setTimeout(() => reject(new Error("Compatibility supervisor did not exit")), 10_000); + child.once("exit", () => { + clearTimeout(timeout); + resolveExit(); + }); + }); + children.delete(child); +} + +function summaries(data: unknown): SessionSummary[] { + if (!data || typeof data !== "object" || !("sessions" in data)) throw new Error("Missing session list"); + const sessions = (data as { sessions: unknown }).sessions; + if (!Array.isArray(sessions)) throw new Error("Invalid session list"); + return sessions as SessionSummary[]; +} + +afterEach(() => { + for (const child of children) { + if (child.exitCode === null && child.signalCode === null) child.kill("SIGTERM"); + } + children.clear(); + for (const root of roots.splice(0)) rmSync(root, { recursive: true, force: true }); +}); + +processTests("stock v0.8.1 daemon entrypoint compatibility", () => { + it.each([ + ["new supervisor adopts a stock 0.8.1 worker", historicalCli!, currentCli], + ["stock 0.8.1 supervisor adopts a new worker", currentCli, historicalCli!], + ] as const)( + "%s", + async (_name, firstCli, replacementCli) => { + const root = mkdtempSync(join(tmpdir(), "prime-081-compat-")); + roots.push(root); + const agentDir = join(root, "agent"); + const projectDir = join(root, "project"); + const socketPath = join(tmpdir(), `prime-081-${process.pid}-${randomUUID().slice(0, 8)}.sock`); + mkdirSync(projectDir, { recursive: true }); + + const firstSupervisor = launch(firstCli, agentDir, socketPath, projectDir); + const firstClient = await connectEventually(socketPath, firstSupervisor); + const created = await firstClient.request({ + type: "create", + noSession: true, + config: { cwd: projectDir, agentDir, noTools: true, noExtensions: true }, + }); + if (!created.success || !created.data || typeof created.data !== "object") { + throw new Error(created.success ? "Compatibility create omitted its summary" : created.error); + } + const createdSummary = created.data as SessionSummary; + if (!createdSummary.activeSessionId || !createdSummary.workerPid) + throw new Error("Compatibility worker was incomplete"); + firstClient.close(); + await stopSupervisor(firstSupervisor); + + const replacement = launch(replacementCli, agentDir, socketPath, projectDir); + const client = await connectEventually(socketPath, replacement); + const listed = await client.request({ type: "list" }); + if (!listed.success) throw new Error(listed.error); + const adopted = summaries(listed.data).find( + (summary) => summary.activeSessionId === createdSummary.activeSessionId, + ); + expect(adopted?.workerPid).toBe(createdSummary.workerPid); + const connection = await DaemonAgentConnection.attach(client, createdSummary.activeSessionId, { + supportsExtensionUi: false, + }); + await expect(connection.getInitialSnapshot()).resolves.toMatchObject({ + state: { activeSessionId: createdSummary.activeSessionId }, + }); + await connection.dispose(); + await client.request({ type: "shutdown" }); + client.close(); + await stopSupervisor(replacement); + }, + 60_000, + ); +}); + +const performanceTests = compiledCli ? describe : describe.skip; + +function residentSetBytes(pid: number): number { + const sampled = spawnSync("ps", ["-o", "rss=", "-p", String(pid)], { encoding: "utf8" }); + if (sampled.status !== 0) return 0; + return Number.parseInt(sampled.stdout.trim(), 10) * 1024; +} + +performanceTests("compiled snapshot process performance", () => { + it("bounds a 36 MiB worker-to-supervisor-to-public transfer", async () => { + const root = mkdtempSync(join(tmpdir(), "prime-compiled-snapshot-")); + roots.push(root); + const agentDir = join(root, "agent"); + const projectDir = join(root, "project"); + const sessionDir = join(agentDir, "sessions"); + const socketPath = join(tmpdir(), `prime-compiled-${process.pid}-${randomUUID().slice(0, 8)}.sock`); + mkdirSync(projectDir, { recursive: true }); + const sessionManager = SessionManager.create(projectDir, sessionDir); + const largePrompt = `compiled:${"x".repeat(36 * 1024 * 1024)}`; + sessionManager.appendMessage({ role: "user", content: largePrompt, timestamp: 1 }); + sessionManager.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "complete" }], + api: "openai-responses", + provider: "faux", + model: "faux", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: 2, + }); + const sessionFile = sessionManager.getSessionFile(); + if (!sessionFile) throw new Error("Compiled snapshot fixture did not persist"); + + const supervisor = launch(compiledCli!, agentDir, socketPath, projectDir); + const client = await connectEventually(socketPath, supervisor); + const created = await client.request({ + type: "create", + sessionPath: sessionFile, + config: { cwd: projectDir, agentDir, sessionDir, noTools: true, noExtensions: true }, + }); + if (!created.success || !created.data || typeof created.data !== "object") { + throw new Error(created.success ? "Compiled create omitted its summary" : created.error); + } + const summary = created.data as SessionSummary; + if (!summary.activeSessionId || !summary.workerPid || !supervisor.pid) { + throw new Error("Compiled snapshot processes were incomplete"); + } + const baselineRss = residentSetBytes(supervisor.pid) + residentSetBytes(summary.workerPid); + let peakRss = baselineRss; + let maxEventLoopDelayMs = 0; + let previousTick = performance.now(); + const monitor = setInterval(() => { + const current = performance.now(); + maxEventLoopDelayMs = Math.max(maxEventLoopDelayMs, current - previousTick - 25); + previousTick = current; + peakRss = Math.max(peakRss, residentSetBytes(supervisor.pid!) + residentSetBytes(summary.workerPid!)); + }, 25); + try { + const connection = await DaemonAgentConnection.attach(client, summary.activeSessionId, { + supportsExtensionUi: false, + }); + const snapshot = await connection.getInitialSnapshot(); + expect(snapshot.messages[0]).toMatchObject({ role: "user", content: largePrompt }); + await connection.dispose(); + } finally { + clearInterval(monitor); + } + expect(peakRss - baselineRss).toBeLessThan(192 * 1024 * 1024); + expect(maxEventLoopDelayMs).toBeLessThan(500); + await client.request({ type: "shutdown" }); + client.close(); + await stopSupervisor(supervisor); + }, 120_000); +}); diff --git a/packages/coding-agent/test/session-worker-private-framing.test.ts b/packages/coding-agent/test/session-worker-private-framing.test.ts index e5c22e7486..32f9a61973 100644 --- a/packages/coding-agent/test/session-worker-private-framing.test.ts +++ b/packages/coding-agent/test/session-worker-private-framing.test.ts @@ -46,6 +46,24 @@ describe("private worker framing", () => { ]); }); + it("coalesces a 36 MiB fragmented payload at most once", () => { + const payload = Buffer.alloc(36 * 1024 * 1024, 0x61); + const frame = encodePrivateFrame({ type: "snapshot", requestId: "large" }, payload); + const decoder = new PrivateFrameDecoder(isTestHeader); + const frames = []; + let peakBufferedBytes = 0; + for (let offset = 0; offset < frame.length; offset += 64 * 1024) { + frames.push(...decoder.push(frame.subarray(offset, offset + 64 * 1024))); + peakBufferedBytes = Math.max(peakBufferedBytes, decoder.bufferedBytes); + } + decoder.finish(); + + expect(frames).toHaveLength(1); + expect(frames[0]?.payload.equals(payload)).toBe(true); + expect(peakBufferedBytes).toBeLessThanOrEqual(payload.length); + expect(decoder.coalescedBytes).toBeLessThanOrEqual(frame.length); + }); + it("keeps segmented new-worker payloads byte-compatible with the legacy decoder", () => { const header = { type: "snapshot", requestId: "compat" }; const payloadParts = [Buffer.from('{"messages":['), Buffer.from('{"content":"ok"}'), Buffer.from("]}")]; diff --git a/packages/coding-agent/test/snapshot-transcript-cache.test.ts b/packages/coding-agent/test/snapshot-transcript-cache.test.ts index 05a16a75bc..4b0ceda6ed 100644 --- a/packages/coding-agent/test/snapshot-transcript-cache.test.ts +++ b/packages/coding-agent/test/snapshot-transcript-cache.test.ts @@ -10,6 +10,7 @@ import { createSnapshotCacheProcessRoot, prepareSnapshotTranscriptCache, prepareSnapshotTranscriptPayload, + SNAPSHOT_MEMORY_CACHE_BYTES, SnapshotTranscriptCache, type SnapshotTranscriptCacheIo, sweepAbandonedSnapshotCacheRoots, @@ -371,15 +372,40 @@ describe("snapshot transcript cache", () => { const payloadBuffers = firstChunks.map((entry) => entry.value![1]); expect(payloadBuffers[1]).toBe(payloadBuffers[0]); expect(payloadBuffers[2]).toBe(payloadBuffers[0]); - expect(single.payload.retainedPayloadBytes).toBeLessThanOrEqual(single.payload.bytes); - await Promise.all(iterators.map((iterator) => iterator.return?.())); + const consume = async (iterator: (typeof iterators)[number]) => { + while (!(await iterator.next()).done) { + // Deliberately drain one reader while its siblings remain at the first chunk. + } + }; + await consume(iterators[0]!); + expect(single.payload.activeReaders).toBe(2); + expect(single.payload.retainedPayloadBytes).toBeLessThanOrEqual(SNAPSHOT_MEMORY_CACHE_BYTES); + await consume(iterators[1]!); + expect(single.payload.activeReaders).toBe(1); + expect(single.payload.retainedPayloadBytes).toBeLessThanOrEqual(SNAPSHOT_MEMORY_CACHE_BYTES); + await consume(iterators[2]!); expect(single.payload.activeReaders).toBe(0); + expect(single.payload.retainedPayloadBytes).toBeLessThanOrEqual(SNAPSHOT_MEMORY_CACHE_BYTES); single.payload.dispose(); const multiple = await run(messages(36, 1024 * 1024)); expect(multiple.maxEventLoopDelayMs).toBeLessThan(50); expect(multiple.rssDeltaBytes).toBeLessThan(192 * 1024 * 1024); expect(multiple.payload.bytes).toBeGreaterThanOrEqual(36 * 1024 * 1024); + const skewed = ["fast", "middle", "slow"].map((purpose) => + multiple.payload.createTransfer("active-skewed", `snapshot-${purpose}`), + ); + const skewedIterators = skewed.map((transfer) => + (transfer as AsyncIterable)[Symbol.asyncIterator](), + ); + await Promise.all(skewedIterators.map((iterator) => iterator.next())); + for (const iterator of skewedIterators) { + while (!(await iterator.next()).done) { + // Each earlier reader fully drains while every later reader stays pinned at chunk zero. + } + expect(multiple.payload.retainedPayloadBytes).toBeLessThanOrEqual(SNAPSHOT_MEMORY_CACHE_BYTES); + } + expect(multiple.payload.activeReaders).toBe(0); multiple.payload.dispose(); }, 30_000); }); diff --git a/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts b/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts index b517a58de3..b0b6290e63 100644 --- a/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts +++ b/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts @@ -154,6 +154,58 @@ describe("ENG-4601 worker snapshot cache", () => { expect(existsSync(cacheDirectory)).toBe(false); }); + it("evicts an exact failed spill generation before a fresh nonce re-encodes it", async () => { + const root = tempDirectory(); + const daemon = new AgentDaemon(join(root, "worker.sock"), { + defaultSessionConfig: { agentDir: root, cwd: root }, + createRuntime: async () => { + throw new Error("unexpected runtime creation"); + }, + }); + const internals = daemon as unknown as { + snapshotPayloadGenerations: Map; + prepareWorkerSnapshotTranscript(options: { + activeSessionId: string; + snapshotId: string; + generationKey: string; + messages: readonly AgentMessage[]; + }): Promise; + }; + const source: AgentMessage[] = [{ role: "user", content: "x".repeat(5 * 1024 * 1024), timestamp: 1 }]; + const first = await internals.prepareWorkerSnapshotTranscript({ + activeSessionId, + snapshotId: "generation-one", + generationKey: "stable:nonce-one", + messages: source, + }); + const firstGeneration = internals.snapshotPayloadGenerations.get(activeSessionId); + const firstIterator = (first as AsyncIterable)[Symbol.asyncIterator](); + const firstChunk = await firstIterator.next(); + first.markFailed?.(new Error("corrupted generation one")); + await firstIterator.return?.(); + expect(internals.snapshotPayloadGenerations.has(activeSessionId)).toBe(false); + + const second = await internals.prepareWorkerSnapshotTranscript({ + activeSessionId, + snapshotId: "generation-two", + generationKey: "stable:nonce-two", + messages: source, + }); + const secondGeneration = internals.snapshotPayloadGenerations.get(activeSessionId); + const secondIterator = (second as AsyncIterable)[Symbol.asyncIterator](); + const secondChunk = await secondIterator.next(); + expect(secondGeneration).not.toBe(firstGeneration); + expect(secondChunk.value?.[1]).not.toBe(firstChunk.value?.[1]); + expect(JSON.parse(Buffer.concat(secondChunk.value ?? []).toString("utf8"))).toMatchObject({ + type: "session_snapshot_chunk", + snapshotId: "generation-two", + }); + while (!(await secondIterator.next()).done) { + // Fully consume the fresh generation to prove its new backing is readable. + } + second.dispose?.(); + }); + it.each([1, 2])( "keeps %i same-ID worker stream(s) bounded by socket drain without shared cleanup", async (count) => { @@ -396,6 +448,7 @@ describe("ENG-4601 worker snapshot cache", () => { activeSessionId, snapshotId, error: `Snapshot ${snapshotId} was aborted`, + purpose: "attach", }); expect(written).toContainEqual( expect.objectContaining({ diff --git a/packages/coding-agent/test/suite/regressions/4602-snapshot-transfer-idempotency.test.ts b/packages/coding-agent/test/suite/regressions/4602-snapshot-transfer-idempotency.test.ts index 8fd0133616..b777ed77b6 100644 --- a/packages/coding-agent/test/suite/regressions/4602-snapshot-transfer-idempotency.test.ts +++ b/packages/coding-agent/test/suite/regressions/4602-snapshot-transfer-idempotency.test.ts @@ -949,12 +949,19 @@ describe("ENG-4602 snapshot transfer containment", () => { expect(attachClient).toHaveBeenCalledOnce(); expect(records.map((record) => record.type)).toEqual(["session_snapshot_begin", "session_snapshot_failed"]); - expect(records[1]).toEqual({ + expect(records[0]).toMatchObject({ type: "session_snapshot_begin", purpose: "replacement" }); + expect(records[1]).toMatchObject({ type: "session_snapshot_failed", activeSessionId, - snapshotId: firstSnapshotId, error: firstError.message, + purpose: "replacement", }); + expect("snapshotId" in records[0]! && "snapshotId" in records[1]! && records[1].snapshotId).toBe( + (records[0] as Extract).snapshotId, + ); + expect((records[0] as Extract).snapshotId).not.toBe( + firstSnapshotId, + ); socket.destroy(); for (const transcript of transcripts) transcript.dispose(); }); diff --git a/packages/coding-agent/test/suite/regressions/4677-snapshot-catchup-replacement.test.ts b/packages/coding-agent/test/suite/regressions/4677-snapshot-catchup-replacement.test.ts index 15fd09d568..9a099dcbde 100644 --- a/packages/coding-agent/test/suite/regressions/4677-snapshot-catchup-replacement.test.ts +++ b/packages/coding-agent/test/suite/regressions/4677-snapshot-catchup-replacement.test.ts @@ -149,6 +149,69 @@ function workerHarness(result: DaemonAttachResult, transcript: SnapshotTranscrip } describe("ENG-4677 snapshot catch-up replacement", () => { + it("atomically remaps an alias snapshot fence before canonical events can pass", () => { + const root = tempDirectory(); + const supervisor = new DaemonSupervisor(join(root, "supervisor.sock"), { + defaultSessionConfig: { agentDir: root, cwd: root }, + descriptorDir: join(root, "state"), + }); + const client = socketClient("alias-client"); + client.capabilities.add("correlated_prompt_lifecycle_v1"); + const transcript = new SnapshotTranscriptCache({ + activeSessionId, + snapshotId: "alias-snapshot", + cacheRoot: root, + }); + const worker = workerHarness(streamedResult("alias-snapshot", 0, 1), transcript); + const written: Buffer[] = []; + const internals = supervisor as unknown as { + clients: Set; + writeSerialized(client: DaemonSocketClient, payload: Buffer): boolean; + reserveSnapshotStream(client: DaemonSocketClient, selector: string): () => void; + remapSnapshotStreamReservation( + client: DaemonSocketClient, + selector: string, + activeSessionId: string, + release: () => void, + ): () => void; + handleWorkerFrame(worker: WorkerHarness, frame: PrivateFrame): void; + }; + internals.clients.add(client); + internals.writeSerialized = (_client, payload) => { + written.push(payload); + return true; + }; + const alias = "session-4677"; + const releaseAlias = internals.reserveSnapshotStream(client, alias); + const releaseCanonical = internals.remapSnapshotStreamReservation(client, alias, activeSessionId, releaseAlias); + + expect(client.snapshotActiveSessionIds).toEqual(new Set([activeSessionId])); + internals.handleWorkerFrame(worker, { + header: { + kind: "outbound", + outboundType: "session_event", + activeSessionId, + sessionEventType: "agent_start", + payloadEncoding: "jsonl", + }, + payload: Buffer.from( + JSON.stringify({ + type: "session_event", + activeSessionId, + event: { type: "agent_start" }, + attribution: { scope: "session" }, + }), + ), + }); + expect(written).toEqual([]); + expect(client.catchupActiveSessionIds).toEqual(new Set([activeSessionId])); + client.catchupActiveSessionIds?.clear(); + releaseCanonical(); + expect(client.snapshotActiveSessionIds).toEqual(new Set()); + transcript.dispose(); + client.socket.destroy(); + }); + it("assigns supervisor-owned UUIDs to mixed-version full snapshots with the same cursor and count", async () => { const root = tempDirectory(); const supervisor = new DaemonSupervisor(join(root, "supervisor.sock"), { @@ -752,14 +815,13 @@ describe("ENG-4677 snapshot catch-up replacement", () => { internals.queueCatchup(client, activeSessionId, "replacement"); await internals.catchUpClient(client); + const publicReplacement = written.find( + (message) => message.type === "session_snapshot_begin" && message.purpose === "replacement", + ); + expect(publicReplacement).toBeDefined(); expect( - written.some( - (message) => - message.type === "session_snapshot_begin" && - message.snapshotId === replacementSnapshotId && - message.purpose === "replacement", - ), - ).toBe(true); + publicReplacement && "snapshotId" in publicReplacement ? publicReplacement.snapshotId : undefined, + ).not.toBe(replacementSnapshotId); expect(worker.transcriptCaches.get(activeSessionId)?.snapshotId).toBe(replacementSnapshotId); client.socket.destroy(); }); From a6d21697b26235f6d0d061188af270680bf398db Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Sat, 29 Aug 2026 19:35:44 -0600 Subject: [PATCH 03/12] test(daemon): prove compatibility worker cleanup --- .../test/daemon-version-compatibility.test.ts | 90 +++++++++++++++++-- 1 file changed, 82 insertions(+), 8 deletions(-) diff --git a/packages/coding-agent/test/daemon-version-compatibility.test.ts b/packages/coding-agent/test/daemon-version-compatibility.test.ts index a44ade8b2a..aad93dc3a4 100644 --- a/packages/coding-agent/test/daemon-version-compatibility.test.ts +++ b/packages/coding-agent/test/daemon-version-compatibility.test.ts @@ -1,6 +1,6 @@ import { type ChildProcess, spawn, spawnSync } from "node:child_process"; import { randomUUID } from "node:crypto"; -import { mkdirSync, mkdtempSync, rmSync } from "node:fs"; +import { existsSync, mkdirSync, mkdtempSync, readdirSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join, resolve } from "node:path"; import { afterEach, describe, expect, it } from "vitest"; @@ -19,6 +19,7 @@ const currentCli = resolve(__dirname, "../dist/bundle/cli.js"); const compiledCli = process.env.PRIME_AGENT_COMPILED_CLI; const processTests = historicalCli ? describe : describe.skip; const children = new Set(); +const workerPids = new Set(); const roots: string[] = []; function scrubbedEnvironment(agentDir: string): NodeJS.ProcessEnv { @@ -66,16 +67,20 @@ async function connectEventually(socketPath: string, process?: ChildProcess): Pr throw new Error(`Compatibility supervisor did not become ready: ${String(lastError)}`); } -async function stopSupervisor(child: ChildProcess): Promise { +async function waitForSupervisorExit(child: ChildProcess, timeoutMs: number): Promise { if (child.exitCode !== null || child.signalCode !== null) return; - child.kill("SIGTERM"); await new Promise((resolveExit, reject) => { - const timeout = setTimeout(() => reject(new Error("Compatibility supervisor did not exit")), 10_000); + const timeout = setTimeout(() => reject(new Error("Compatibility supervisor did not exit")), timeoutMs); child.once("exit", () => { clearTimeout(timeout); resolveExit(); }); }); +} + +async function stopSupervisor(child: ChildProcess): Promise { + if (child.exitCode === null && child.signalCode === null) child.kill("SIGTERM"); + await waitForSupervisorExit(child, 10_000); children.delete(child); } @@ -86,11 +91,75 @@ function summaries(data: unknown): SessionSummary[] { return sessions as SessionSummary[]; } -afterEach(() => { - for (const child of children) { - if (child.exitCode === null && child.signalCode === null) child.kill("SIGTERM"); +function workerDescriptorFiles(agentDir: string): string[] { + const workersRoot = join(agentDir, "daemon-workers"); + if (!existsSync(workersRoot)) return []; + return readdirSync(workersRoot, { recursive: true }) + .map(String) + .filter((path) => path.endsWith(".json")); +} + +function isProcessAlive(pid: number): boolean { + try { + process.kill(pid, 0); + return true; + } catch (error) { + if ((error as NodeJS.ErrnoException).code === "ESRCH") return false; + throw error; + } +} + +async function waitForProcessExit(pid: number, timeoutMs: number): Promise { + const deadline = Date.now() + timeoutMs; + while (isProcessAlive(pid)) { + if (Date.now() >= deadline) throw new Error(`Compatibility worker ${pid} did not exit`); + await new Promise((resolveDelay) => setTimeout(resolveDelay, 25)); + } +} + +async function stopWorkerProcess(pid: number): Promise { + try { + await waitForProcessExit(pid, 250); + return; + } catch { + // Test-only fallback after the natural shutdown assertion has already failed or been skipped. + } + try { + process.kill(-pid, "SIGTERM"); + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== "ESRCH") throw error; + try { + process.kill(pid, "SIGTERM"); + } catch (pidError) { + if ((pidError as NodeJS.ErrnoException).code === "ESRCH") return; + throw pidError; + } + } + try { + await waitForProcessExit(pid, 2_000); + return; + } catch { + try { + process.kill(-pid, "SIGKILL"); + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== "ESRCH") throw error; + process.kill(pid, "SIGKILL"); + } + await waitForProcessExit(pid, 2_000); + } +} + +afterEach(async () => { + for (const child of [...children]) { + try { + await stopSupervisor(child); + } catch { + if (child.exitCode === null && child.signalCode === null) child.kill("SIGKILL"); + } } children.clear(); + for (const pid of workerPids) await stopWorkerProcess(pid); + workerPids.clear(); for (const root of roots.splice(0)) rmSync(root, { recursive: true, force: true }); }); @@ -121,6 +190,7 @@ processTests("stock v0.8.1 daemon entrypoint compatibility", () => { const createdSummary = created.data as SessionSummary; if (!createdSummary.activeSessionId || !createdSummary.workerPid) throw new Error("Compatibility worker was incomplete"); + workerPids.add(createdSummary.workerPid); firstClient.close(); await stopSupervisor(firstSupervisor); @@ -141,7 +211,11 @@ processTests("stock v0.8.1 daemon entrypoint compatibility", () => { await connection.dispose(); await client.request({ type: "shutdown" }); client.close(); - await stopSupervisor(replacement); + await waitForSupervisorExit(replacement, 10_000); + children.delete(replacement); + await waitForProcessExit(createdSummary.workerPid, 10_000); + expect(workerDescriptorFiles(agentDir)).toEqual([]); + workerPids.delete(createdSummary.workerPid); }, 60_000, ); From 9a957dd10258d93d305c92af24746cedb20eeb36 Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Sat, 29 Aug 2026 19:37:15 -0600 Subject: [PATCH 04/12] fix(daemon): fence snapshots and worker frames --- .../coding-agent/src/core/agent-session.ts | 37 ++-- .../src/modes/daemon/active-session-state.ts | 2 +- .../src/modes/daemon/daemon-mode.ts | 99 ++++++---- .../src/modes/daemon/daemon-supervisor.ts | 169 ++++++++++++++++-- .../test/agent-session-recursion.test.ts | 19 ++ .../coding-agent/test/daemon-mode.test.ts | 77 +++++++- .../test/daemon-supervisor-monitor.test.ts | 156 ++++++++++++++++ .../4601-worker-snapshot-cache.test.ts | 81 ++++++++- .../4677-snapshot-catchup-replacement.test.ts | 49 +++++ 9 files changed, 615 insertions(+), 74 deletions(-) diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 2de6a12ab9..441325541a 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -938,6 +938,8 @@ interface RlmChildRun { interface RetainedRlmChild { session: AgentSession; run?: RlmChildRun; + /** Durable terminal status for a rehydrated child that has no in-memory run. */ + terminalStatus?: "done" | "error"; } interface RlmSubagentModelSelection { @@ -10052,7 +10054,7 @@ export class AgentSession { }); recorded.add(run.id); } - for (const [childId, { session: childSession }] of this._rlmChildSessions) { + for (const [childId, { session: childSession, terminalStatus }] of this._rlmChildSessions) { if ( this._deletingRlmChildren.has(childId) || recorded.has(childId) || @@ -10072,7 +10074,7 @@ export class AgentSession { session_name: daemonChild?.sessionName ?? childSession.sessionName ?? createDefaultRlmSubagentSessionName("", childId), session_dir: sessionDir, - status: "completed", + status: terminalStatus === "error" ? "error" : "completed", }); recorded.add(childId); } @@ -10440,7 +10442,12 @@ export class AgentSession { * the child) when the parent is already tearing down, so the caller can drop the * matching event forwarder too. */ - registerRlmChildSession(childId: string, session: AgentSession, unsubscribe?: () => void): boolean { + registerRlmChildSession( + childId: string, + session: AgentSession, + unsubscribe?: () => void, + terminalStatus: "done" | "error" = "done", + ): boolean { // A child can finish concurrently while the parent is (or has) torn down; don't // resurrect the map (it would never be disposed), just drop the child now. if (this._deletingRlmChildren.has(childId) || this._deletedRlmChildIds.has(childId)) { @@ -10453,7 +10460,11 @@ export class AgentSession { void session.disposeAsync().catch(() => undefined); return false; } - this._rlmChildSessions.set(childId, { session, run: this._activeRlmChildRuns.get(childId) }); + this._rlmChildSessions.set(childId, { + session, + run: this._activeRlmChildRuns.get(childId), + terminalStatus, + }); if (unsubscribe) { this._rlmChildUnsubscribes.set(childId, unsubscribe); } @@ -10503,7 +10514,11 @@ export class AgentSession { }; } - private _rlmChildSnapshotForSession(childId: string, child: AgentSession): RlmChildAgentSnapshot { + private _rlmChildSnapshotForSession( + childId: string, + child: AgentSession, + terminalStatus: "done" | "error" = "done", + ): RlmChildAgentSnapshot { let answerPreview: string | undefined; let toolUseCount = 0; const messages = @@ -10522,16 +10537,16 @@ export class AgentSession { sessionName: child.sessionName, model: child.model ? `${child.model.provider}/${child.model.id}` : undefined, label: child.sessionName ?? "child agent", - status: "done", + status: terminalStatus, answerPreview, toolUseCount: toolUseCount > 0 ? toolUseCount : undefined, tokenCount: child._contextTokensForCurrentMessages(), recap: child.getCurrentRecap(), sessionDir: child._rlmSessionDir ?? child.sessionManager.getSessionDir(), // No run exists (e.g. a child rehydrated after daemon recovery), so live - // session state is the only source for in-flight follow-up work. Mirror - // the run projection's convention: status stays "done" (the recorded task - // finished) and current work surfaces through activity. + // session state is the only source for in-flight follow-up work. The + // registry owns the recorded task's terminal status while current work + // surfaces separately through activity. activity: child.isSessionActive ? { kind: child.isStreaming ? "writing" : "waiting" } : undefined, repliedSinceTask: child._repliedToParentSinceTask, }; @@ -10555,13 +10570,13 @@ export class AgentSession { snapshots.push(...child.getRlmChildSnapshots()); } } - for (const [childId, { session: child, run }] of this._rlmChildSessions) { + for (const [childId, { session: child, run, terminalStatus }] of this._rlmChildSessions) { if (recorded.has(childId) || traversed.has(childId)) continue; const hidden = this._deletingRlmChildren.has(childId) || this._deletedRlmChildIds.has(childId); if (!hidden) { const snapshot = run ? this._rlmChildSnapshotForRun(run, child) - : this._rlmChildSnapshotForSession(childId, child); + : this._rlmChildSnapshotForSession(childId, child, terminalStatus); snapshots.push({ ...snapshot, status: this._rlmChildCleanupFailures.has(childId) ? "cancelled" : snapshot.status, diff --git a/packages/coding-agent/src/modes/daemon/active-session-state.ts b/packages/coding-agent/src/modes/daemon/active-session-state.ts index 9a925a8067..360fb6347b 100644 --- a/packages/coding-agent/src/modes/daemon/active-session-state.ts +++ b/packages/coding-agent/src/modes/daemon/active-session-state.ts @@ -23,7 +23,7 @@ export interface DaemonSocketClient { snapshotActiveSessionCounts?: Map; snapshotTransferAbortControllers?: Map; snapshotTransferTails?: Map>; - /** Serializes private frames when a large snapshot frame pauses between parts. */ + /** Settles only after all queued private-frame reservations and their predecessors retire. */ privateFrameWriteTail?: Promise; detachInput: () => void; supportsExtensionUi: boolean; diff --git a/packages/coding-agent/src/modes/daemon/daemon-mode.ts b/packages/coding-agent/src/modes/daemon/daemon-mode.ts index bda8b01c80..0eb002e7c2 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-mode.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-mode.ts @@ -3012,7 +3012,16 @@ export class AgentDaemon { ); // The session transcript is authoritative for mutable metadata such as a // later user-assigned name; the registry value is only the spawn snapshot. - if (!parentState.runtime.session.registerRlmChildSession(entry.childId, runtime.session)) { + const registered = + entry.status === "completed" + ? parentState.runtime.session.registerRlmChildSession(entry.childId, runtime.session) + : parentState.runtime.session.registerRlmChildSession( + entry.childId, + runtime.session, + undefined, + "error", + ); + if (!registered) { await this.closeSession(state, "replaced"); throw new RuntimeOpenCancelledError(); } @@ -5156,38 +5165,50 @@ export class AgentDaemon { ...(metadata.rlmChildId ? { childId: metadata.rlmChildId } : {}), } : undefined; - let session = state.runtime.session; - let children = await this.buildRlmChildSnapshotsWithPassiveRlmSubagents(state); - for ( - let retries = 0; - retries < MAX_SESSION_SNAPSHOT_STABILIZATION_RETRIES && state.runtime.session !== session; - retries++ - ) { - session = state.runtime.session; - children = await this.buildRlmChildSnapshotsWithPassiveRlmSubagents(state); + for (let attempt = 0; attempt <= MAX_SESSION_SNAPSHOT_STABILIZATION_RETRIES; attempt++) { + const session = state.runtime.session; + const eventGeneration = state.eventGeneration; + const lastEventSequence = state.lastEventSequence; + const children = await this.buildRlmChildSnapshotsWithPassiveRlmSubagents(state); + if ( + state.runtime.session !== session || + state.eventGeneration !== eventGeneration || + state.lastEventSequence !== lastEventSequence + ) { + continue; + } + const connectionState = this.createConnectionState(state); + const messages = createImmutableSnapshotMessages(session.messages); + const summary = summaryForActiveSession(state); + const promptLifecycles = includePromptLifecycles + ? (session.getPromptLifecycles?.() ?? { records: [], expired: [] }) + : undefined; + if ( + state.runtime.session !== session || + state.eventGeneration !== eventGeneration || + state.lastEventSequence !== lastEventSequence + ) { + continue; + } + return { + activeSessionId: state.activeSessionId, + summary, + state: connectionState, + messages, + // Omit duplicate heavy payloads from attach. The client can derive render + // context from messages + state, and fetch the full session tree lazily + // when the tree/branch selector opens. + lastEventSequence, + lastEventCursor: { + generation: eventGeneration, + sequence: lastEventSequence, + }, + ...(parent ? { parent } : {}), + children, + ...(promptLifecycles ? { promptLifecycles } : {}), + }; } - session = state.runtime.session; - const connectionState = this.createConnectionState(state); - const messages = createImmutableSnapshotMessages(session.messages); - return { - activeSessionId: state.activeSessionId, - summary: summaryForActiveSession(state), - state: connectionState, - messages, - // Omit duplicate heavy payloads from attach. The client can derive render - // context from messages + state, and fetch the full session tree lazily - // when the tree/branch selector opens. - lastEventSequence: state.lastEventSequence, - lastEventCursor: { - generation: state.eventGeneration, - sequence: state.lastEventSequence, - }, - ...(parent ? { parent } : {}), - children, - ...(includePromptLifecycles - ? { promptLifecycles: session.getPromptLifecycles?.() ?? { records: [], expired: [] } } - : {}), - }; + throw new Error(`Session ${state.activeSessionId} changed while its snapshot was being prepared`); } private async prepareWorkerSnapshotTranscript(options: { @@ -5505,11 +5526,16 @@ export class AgentDaemon { drainTimeoutMs?: number, ): Promise { const previous = client.privateFrameWriteTail; - let resolveTail!: () => void; - const tail = new Promise((resolve) => { - resolveTail = resolve; + const previousSettled = previous?.catch(() => undefined); + let releaseCurrent!: () => void; + const currentReleased = new Promise((resolve) => { + releaseCurrent = resolve; }); + const tail = previousSettled ? previousSettled.then(() => currentReleased) : currentReleased; client.privateFrameWriteTail = tail; + void tail.finally(() => { + if (client.privateFrameWriteTail === tail) client.privateFrameWriteTail = undefined; + }); try { if (previous) await waitForSnapshotAttempt(previous, signal ?? AbortSignal.timeout(30_000), "private-frame"); for (const part of parts) { @@ -5523,8 +5549,7 @@ export class AgentDaemon { } return true; } finally { - resolveTail(); - if (client.privateFrameWriteTail === tail) client.privateFrameWriteTail = undefined; + releaseCurrent(); } } diff --git a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts index b3272b837a..7befb30ab6 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts @@ -314,6 +314,8 @@ interface ResidentWorker { snapshotGenerations: Map>; snapshotLoads: Map>; snapshotLoadSnapshotIds: Map; + /** Active-session IDs this authenticated resident worker may emit. */ + authorizedActiveSessionIds: Set; recovery?: Promise; deferredRecovery?: Promise; intentionalStop: boolean; @@ -378,6 +380,7 @@ interface SupervisorTranscriptPreparation { controller: AbortController; workerClient: DaemonWorkerClient | undefined; stopRevision: number; + selectionRevision: number; } interface SupervisorPromptAdmission { @@ -732,6 +735,7 @@ export class DaemonSupervisor { private readonly streamReconstructor = new CompactAssistantStreamReconstructor(); private readonly compactCatchupInProgress = new Set(); private transcriptPreparations?: WeakMap>; + private snapshotSelectionRevisions?: WeakMap>; private readonly pendingSessionNames = new Set(); private readonly catalog: DaemonCatalogClient; private readonly settingsManager: SettingsManager; @@ -1072,6 +1076,7 @@ export class DaemonSupervisor { snapshotGenerations: new Map(), snapshotLoads: new Map(), snapshotLoadSnapshotIds: new Map(), + authorizedActiveSessionIds: new Set([durableDescriptor.rootActiveSessionId]), intentionalStop: durableDescriptor.stopRequestedAt !== undefined, stopRevision: 0, }; @@ -2851,6 +2856,7 @@ export class DaemonSupervisor { snapshotGenerations: new Map(), snapshotLoads: new Map(), snapshotLoadSnapshotIds: new Map(), + authorizedActiveSessionIds: new Set([descriptor.rootActiveSessionId]), intentionalStop: false, stopRevision: 0, launchEnv, @@ -3239,6 +3245,15 @@ export class DaemonSupervisor { closeWorkerChannel = false, expectedSnapshotId?: string, ): void { + const selectedSnapshotId = worker.snapshotCache.get(activeSessionId)?.snapshotStream?.id; + const selectedTranscriptId = worker.transcriptCaches.get(activeSessionId)?.snapshotId; + if ( + expectedSnapshotId === undefined || + selectedSnapshotId === expectedSnapshotId || + selectedTranscriptId === expectedSnapshotId + ) { + this.invalidateSnapshotSelection(worker, activeSessionId, error); + } const generations = worker.snapshotGenerations?.get(activeSessionId); if (expectedSnapshotId) { const generation = generations?.get(expectedSnapshotId); @@ -3809,6 +3824,44 @@ export class DaemonSupervisor { return promptLifecycles; } + private workerAuthorizedActiveSessionIds(worker: ResidentWorker): Set { + if (!worker.authorizedActiveSessionIds) { + const rootActiveSessionId = worker.descriptor?.rootActiveSessionId; + worker.authorizedActiveSessionIds = new Set(rootActiveSessionId ? [rootActiveSessionId] : []); + } + return worker.authorizedActiveSessionIds; + } + + private workerOwnsActiveSessionId(worker: ResidentWorker, activeSessionId: string): boolean { + // Prototype-only unit harnesses omit both identity fields. Every production + // worker has a descriptor and remains registered while its channel is valid. + if (!worker.descriptor) return !this.workers; + if (this.workers && this.workers.get(worker.descriptor.workerId) !== worker) return false; + return this.workerAuthorizedActiveSessionIds(worker).has(activeSessionId); + } + + private bindWorkerActiveSessionIds(worker: ResidentWorker, activeSessionIds: Iterable): void { + const candidates = new Set(activeSessionIds); + for (const activeSessionId of candidates) { + if (!activeSessionId) throw new Error("Session worker returned an empty active-session ID"); + for (const other of this.workers?.values() ?? []) { + if (other === worker) continue; + if (this.workerAuthorizedActiveSessionIds(other).has(activeSessionId)) { + throw new Error(`Active session ${activeSessionId} is already bound to another worker`); + } + } + } + const authorized = this.workerAuthorizedActiveSessionIds(worker); + for (const activeSessionId of candidates) authorized.add(activeSessionId); + } + + private rejectWorkerFrame(worker: ResidentWorker, error: Error): void { + const workerClient = worker.client; + if (!workerClient) return; + this.handleWorkerClose(worker, workerClient, error); + workerClient.close(); + } + private async refreshWorkerSummaries(worker: ResidentWorker, recovery = false): Promise { if (this.isWorkerStopping(worker)) { throw new Error("Session worker is stopping"); @@ -3823,6 +3876,7 @@ export class DaemonSupervisor { if (recovery && !root) { throw new Error(`Session worker omitted its root session during recovery`); } + this.bindWorkerActiveSessionIds(worker, nextSummaries.keys()); worker.summaries = nextSummaries; for (const summary of summaries) { const activeSessionId = summary.activeSessionId ?? summary.id; @@ -4291,6 +4345,7 @@ export class DaemonSupervisor { const observedSnapshotId = match.worker.transcriptCaches.get(activeSessionId)?.snapshotId ?? match.worker.snapshotCache.get(activeSessionId)?.snapshotStream?.id; + const selectionRevision = this.snapshotSelectionRevision(match.worker, activeSessionId); loading = (async () => { const workerClient = this.requireAvailableWorkerClient(match.worker); const workerCapabilities: DaemonClientCapability[] = [ @@ -4319,7 +4374,10 @@ export class DaemonSupervisor { SNAPSHOT_ATTEMPT_TIMEOUT_MS, ); const loaded = attachResultFromResponse(response); - if (match.worker.snapshotLoads.get(snapshotLoadKey) !== loading) { + if ( + match.worker.snapshotLoads.get(snapshotLoadKey) !== loading || + this.snapshotSelectionRevision(match.worker, activeSessionId) !== selectionRevision + ) { throw new SnapshotLoadInvalidatedError("Session snapshot changed during attach"); } return this.cacheLoadedSnapshot(match.worker, activeSessionId, loaded, observedSnapshotId); @@ -4521,6 +4579,26 @@ export class DaemonSupervisor { return loaded; } + private snapshotSelectionRevision(worker: ResidentWorker, activeSessionId: string): number { + return this.snapshotSelectionRevisions?.get(worker)?.get(activeSessionId) ?? 0; + } + + private invalidateSnapshotSelection(worker: ResidentWorker, activeSessionId: string, reason: Error): void { + this.snapshotSelectionRevisions ??= new WeakMap(); + let revisions = this.snapshotSelectionRevisions.get(worker); + if (!revisions) { + revisions = new Map(); + this.snapshotSelectionRevisions.set(worker, revisions); + } + revisions.set(activeSessionId, this.snapshotSelectionRevision(worker, activeSessionId) + 1); + const preparations = this.transcriptPreparations?.get(worker); + const preparation = preparations?.get(activeSessionId); + if (preparation) { + preparations?.delete(activeSessionId); + preparation.controller.abort(reason); + } + } + private abortTranscriptPreparations(worker: ResidentWorker, reason: Error): void { const preparations = this.transcriptPreparations?.get(worker); if (!preparations) return; @@ -4538,6 +4616,17 @@ export class DaemonSupervisor { const existingGeneration = existing ? this.snapshotGeneration(worker, activeSessionId, existing.snapshotId) : undefined; + const selectedResult = worker.snapshotCache.get(activeSessionId); + if ( + selectedResult !== result && + !( + existingGeneration && + existingGeneration.result === selectedResult && + existingGeneration.materializedFrom === result + ) + ) { + throw new SnapshotLoadInvalidatedError("Session snapshot changed before transcript preparation"); + } if ( existing && (result.snapshotStream @@ -4567,6 +4656,7 @@ export class DaemonSupervisor { const controller = new AbortController(); const workerClient = worker.client; const stopRevision = worker.stopRevision; + const selectionRevision = this.snapshotSelectionRevision(worker, activeSessionId); const promise = prepareSnapshotTranscriptCache({ activeSessionId, snapshotId, @@ -4581,23 +4671,39 @@ export class DaemonSupervisor { controller, workerClient, stopRevision, + selectionRevision, }; preparations.set(activeSessionId, preparation); let transcript: SnapshotTranscriptCache; try { transcript = await promise; + } catch (error) { + if ( + this.snapshotSelectionRevision(worker, activeSessionId) !== selectionRevision || + worker.snapshotCache.get(activeSessionId) !== result + ) { + throw new SnapshotLoadInvalidatedError("Session snapshot changed during transcript preparation"); + } + throw error; } finally { if (preparations.get(activeSessionId) === preparation) preparations.delete(activeSessionId); } const registeredWorker = this.workers?.get(worker.descriptor.workerId); + const selectionChanged = + this.snapshotSelectionRevision(worker, activeSessionId) !== selectionRevision || + worker.snapshotCache.get(activeSessionId) !== result; if ( + selectionChanged || controller.signal.aborted || this.shuttingDown || - (registeredWorker !== undefined && registeredWorker !== worker) || + (this.workers !== undefined && registeredWorker !== worker) || worker.client !== workerClient || worker.stopRevision !== stopRevision ) { transcript.dispose(); + if (selectionChanged) { + throw new SnapshotLoadInvalidatedError("Session snapshot changed during transcript preparation"); + } throw new Error("Session worker changed during snapshot transcript preparation"); } const cachedResult = { @@ -4927,10 +5033,18 @@ export class DaemonSupervisor { snapshotPurpose, } = frame.header; if (outboundType === "heartbeats_changed") { + if (activeSessionId) { + this.rejectWorkerFrame(worker, new Error("Worker heartbeat frame was unexpectedly session-scoped")); + return; + } worker.heartbeatSnapshotStale = true; this.broadcastHeartbeatsChanged(); return; } + if (activeSessionId && !this.workerOwnsActiveSessionId(worker, activeSessionId)) { + this.rejectWorkerFrame(worker, new Error(`Worker emitted an unauthorized active session ${activeSessionId}`)); + return; + } if ( (outboundType === "session_snapshot_begin" || outboundType === "session_snapshot_chunk" || @@ -5365,7 +5479,7 @@ export class DaemonSupervisor { return; } let publicPayload = frame.payload; - let decodedOutbound: DaemonOutbound | undefined; + let decodedOutbound: DaemonOutbound; if (payloadEncoding === "assistant-delta") { let compactValue: unknown; try { @@ -5374,10 +5488,14 @@ export class DaemonSupervisor { this.scheduleCompactCatchup(worker, activeSessionId); return; } - if (!isCompactAssistantDelta(compactValue) || compactValue.activeSessionId !== activeSessionId) { + if (!isCompactAssistantDelta(compactValue)) { this.scheduleCompactCatchup(worker, activeSessionId); return; } + if (compactValue.activeSessionId !== activeSessionId) { + this.rejectWorkerFrame(worker, new Error("Worker compact payload did not match its active-session header")); + return; + } const reconstructed = this.streamReconstructor.reconstruct(compactValue); if (!reconstructed) { this.scheduleCompactCatchup(worker, activeSessionId); @@ -5385,25 +5503,35 @@ export class DaemonSupervisor { } decodedOutbound = reconstructed; publicPayload = Buffer.from(serializeJsonLine(reconstructed)); - } else if ( - sessionEventType === "message_start" || - sessionEventType === "message_end" || - outboundType === "session_replaced" || - outboundType === "session_resynced" || - outboundType === "session_closed" || - [...this.clients].some( - (client) => - client.attachedActiveSessionIds.has(activeSessionId) && - !client.capabilities.has("correlated_prompt_lifecycle_v1"), - ) - ) { + } else { try { decodedOutbound = JSON.parse(frame.payload.toString("utf8")) as DaemonOutbound; - this.streamReconstructor.observe(decodedOutbound); } catch { - // A malformed worker event is still isolated to this worker connection. + this.rejectWorkerFrame(worker, new Error("Worker returned malformed session-scoped JSON")); + return; } } + if ( + decodedOutbound.type !== outboundType || + (decodedOutbound as { activeSessionId?: unknown }).activeSessionId !== activeSessionId + ) { + this.rejectWorkerFrame(worker, new Error("Worker payload did not match its routed frame header")); + return; + } + this.streamReconstructor.observe(decodedOutbound); + if ( + decodedOutbound.type === "session_event" && + decodedOutbound.event.type === "rlm_child_update" && + decodedOutbound.event.child.activeSessionId + ) { + try { + this.bindWorkerActiveSessionIds(worker, [decodedOutbound.event.child.activeSessionId]); + } catch (error) { + this.rejectWorkerFrame(worker, error instanceof Error ? error : new Error(String(error))); + return; + } + } + const replacementSnapshotFollows = decodedOutbound?.type === "session_replaced" && decodedOutbound.snapshotFollows === true; const reencodeLegacyWorkerReplacement = @@ -5489,6 +5617,11 @@ export class DaemonSupervisor { transcriptChanged = true, invalidateLoad = transcriptChanged, ): void { + this.invalidateSnapshotSelection( + worker, + activeSessionId, + new SnapshotLoadInvalidatedError("Session snapshot changed during selection"), + ); worker.snapshotCache.delete(activeSessionId); if (invalidateLoad) { this.invalidateSnapshotLoads(worker, activeSessionId); diff --git a/packages/coding-agent/test/agent-session-recursion.test.ts b/packages/coding-agent/test/agent-session-recursion.test.ts index 33979fbe14..1cad16ee85 100644 --- a/packages/coding-agent/test/agent-session-recursion.test.ts +++ b/packages/coding-agent/test/agent-session-recursion.test.ts @@ -562,6 +562,25 @@ describe("AgentSession rlm recursion", () => { expect(childStatuses).toEqual(["cancelled"]); }); + it("preserves an errored terminal outcome for a restored session-only child", async () => { + const childId = "restored-error-child"; + const childDir = join(tempDir, childId); + mkdirSync(childDir, { recursive: true }); + const child = createSession({ rlmSessionDir: childDir }); + child.setSessionName("restored-error-worker"); + const root = createSession(); + + expect(root.registerRlmChildSession(childId, child, undefined, "error")).toBe(true); + expect(root.getRlmChildSnapshots()).toEqual([expect.objectContaining({ id: childId, status: "error" })]); + expect((await root.listRlmSubagents()).subagents).toEqual([ + expect.objectContaining({ rlm_child_id: childId, status: "error" }), + ]); + const release = root.releaseRlmChildSession(childId, child); + expect(release).not.toBe(false); + expect(root.getRlmChildSnapshots()).toEqual([expect.objectContaining({ id: childId, status: "error" })]); + if (release) release(); + }); + it("projects live follow-up activity for a restored session-only child", async () => { const childId = "restored-followup-child"; const childDir = join(tempDir, childId); diff --git a/packages/coding-agent/test/daemon-mode.test.ts b/packages/coding-agent/test/daemon-mode.test.ts index 9deb8527d8..c843ddc278 100644 --- a/packages/coding-agent/test/daemon-mode.test.ts +++ b/packages/coding-agent/test/daemon-mode.test.ts @@ -4134,6 +4134,7 @@ describe("daemon mode helpers", () => { buildRlmChildSnapshotsWithPassiveRlmSubagents( state: ActiveSessionState, ): Promise>; + getOrHydrateBoundSessionState(selector: string): Promise; handleCommand(client: DaemonSocketClient, command: DaemonCommand): Promise; }; @@ -4255,6 +4256,16 @@ describe("daemon mode helpers", () => { }), ]), ); + + const hydratedChild = await internals.getOrHydrateBoundSessionState(fixture.childId); + expect(hydratedChild.runtime.session).toBe(fixture.runtimeSessions[1]); + expect(parentState.runtime.session.registerRlmChildSession).toHaveBeenCalledWith( + fixture.childId, + hydratedChild.runtime.session, + undefined, + "error", + ); + expect(fixture.createRuntime).toHaveBeenCalledTimes(2); } finally { rmSync(tempDir, { recursive: true, force: true }); } @@ -4508,7 +4519,63 @@ describe("daemon mode helpers", () => { } }); - it("bounds snapshot stabilization when every child build replaces the runtime session", async () => { + it("recomputes snapshot children when the event cursor advances during the passive walk", async () => { + const tempDir = mkdtempSync(join(tmpdir(), "prime-agent-daemon-snapshot-child-event-")); + try { + const fixture = makePersistedRlmDaemonFixture(tempDir); + const internals = fixture.daemon as unknown as { + createRuntime(command: Extract): Promise; + createSessionSnapshot(state: ActiveSessionState): Promise; + createConnectionState: ReturnType; + buildRlmChildSnapshotsWithPassiveRlmSubagents: ReturnType; + }; + const state = await internals.createRuntime({ type: "create", sessionPath: fixture.parentSessionFile }); + internals.createConnectionState = vi.fn(() => ({})); + let calls = 0; + internals.buildRlmChildSnapshotsWithPassiveRlmSubagents = vi.fn(async () => { + calls++; + if (calls === 1) state.lastEventSequence++; + return [{ id: calls === 1 ? "running-child" : "finished-child", status: "done", sessionDir: tempDir }]; + }); + + const snapshot = await internals.createSessionSnapshot(state); + + expect(internals.buildRlmChildSnapshotsWithPassiveRlmSubagents).toHaveBeenCalledTimes(2); + expect(snapshot.children).toEqual([expect.objectContaining({ id: "finished-child" })]); + expect(snapshot.lastEventSequence).toBe(1); + expect(snapshot.lastEventCursor).toMatchObject({ generation: state.eventGeneration, sequence: 1 }); + } finally { + rmSync(tempDir, { recursive: true, force: true }); + } + }); + + it("fails snapshot stabilization when the event cursor advances on every child build", async () => { + const tempDir = mkdtempSync(join(tmpdir(), "prime-agent-daemon-snapshot-event-bound-")); + try { + const fixture = makePersistedRlmDaemonFixture(tempDir); + const internals = fixture.daemon as unknown as { + createRuntime(command: Extract): Promise; + createSessionSnapshot(state: ActiveSessionState): Promise; + createConnectionState: ReturnType; + buildRlmChildSnapshotsWithPassiveRlmSubagents: ReturnType; + }; + const state = await internals.createRuntime({ type: "create", sessionPath: fixture.parentSessionFile }); + internals.createConnectionState = vi.fn(() => ({})); + internals.buildRlmChildSnapshotsWithPassiveRlmSubagents = vi.fn(async () => { + state.lastEventSequence++; + return [{ id: `child-${state.lastEventSequence}`, status: "done", sessionDir: tempDir }]; + }); + + await expect(internals.createSessionSnapshot(state)).rejects.toThrow( + "changed while its snapshot was being prepared", + ); + expect(internals.buildRlmChildSnapshotsWithPassiveRlmSubagents).toHaveBeenCalledTimes(4); + } finally { + rmSync(tempDir, { recursive: true, force: true }); + } + }); + + it("fails snapshot stabilization when every child build replaces the runtime session", async () => { const tempDir = mkdtempSync(join(tmpdir(), "prime-agent-daemon-snapshot-stabilization-bound-")); try { const fixture = makePersistedRlmDaemonFixture(tempDir); @@ -4534,11 +4601,9 @@ describe("daemon mode helpers", () => { return [{ id: `child-${calls}`, status: "done", sessionDir: tempDir }]; }); - await expect(internals.createSessionSnapshot(state)).resolves.toMatchObject({ - children: [expect.objectContaining({ id: "child-4" })], - messages: [{ content: "transcript 4" }], - summary: { sessionId: "session-4" }, - }); + await expect(internals.createSessionSnapshot(state)).rejects.toThrow( + "changed while its snapshot was being prepared", + ); expect(internals.buildRlmChildSnapshotsWithPassiveRlmSubagents).toHaveBeenCalledTimes(4); } finally { rmSync(tempDir, { recursive: true, force: true }); diff --git a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts index 6fa13ce0ed..a091f4cea8 100644 --- a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts @@ -3564,6 +3564,162 @@ describe("daemon worker supervisor monitoring", () => { expect(catchUpClient).not.toHaveBeenCalled(); }); + it("rejects worker frames that target another resident or mismatch the routed payload", () => { + const writes: string[] = []; + const publicClient = { + id: "public-b", + socket: { + destroyed: false, + write: vi.fn((value: unknown) => { + writes.push(String(value)); + return true; + }), + }, + attachedActiveSessionIds: new Set(["active-b"]), + catchupActiveSessionIds: new Set(), + backpressured: false, + supportsExtensionUi: false, + capabilities: new Set(), + } as unknown as DaemonSocketClient; + const makeWorker = (workerId: string, activeId: string) => ({ + descriptor: { workerId, rootActiveSessionId: activeId }, + client: { close: vi.fn() }, + authorizedActiveSessionIds: new Set([activeId]), + summaries: new Map(), + snapshotCache: new Map(), + transcriptCaches: new Map(), + snapshotGenerations: new Map(), + snapshotLoads: new Map(), + }); + const workerA = makeWorker("worker-a", "active-a"); + const workerB = makeWorker("worker-b", "active-b"); + const handleWorkerClose = vi.fn(); + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers: new Map([ + ["worker-a", workerA], + ["worker-b", workerB], + ]), + clients: new Set([publicClient]), + streamReconstructor: { observe: vi.fn() }, + handleWorkerClose, + invalidateWorkerSnapshot: vi.fn(), + refreshWorkerSummaries: vi.fn(async () => undefined), + }) as { + handleWorkerFrame(residentWorker: typeof workerA, frame: PrivateFrame): void; + }; + + supervisor.handleWorkerFrame(workerA, { + header: { + kind: "outbound", + outboundType: "session_snapshot_begin", + activeSessionId: "active-b", + snapshotId: "foreign-snapshot", + }, + payload: Buffer.from("{}"), + }); + expect(handleWorkerClose).toHaveBeenCalledOnce(); + expect(workerA.client.close).toHaveBeenCalledOnce(); + expect(workerA.snapshotCache.size).toBe(0); + expect(workerA.snapshotGenerations.size).toBe(0); + expect(writes).toEqual([]); + + handleWorkerClose.mockClear(); + workerA.client.close.mockClear(); + supervisor.handleWorkerFrame(workerA, { + header: { kind: "outbound", outboundType: "extension_error", activeSessionId: "active-a" }, + payload: Buffer.from( + JSON.stringify({ + type: "extension_error", + activeSessionId: "active-b", + extensionPath: "/tmp/mismatch.ts", + event: "load", + error: "mismatch", + }), + ), + }); + expect(handleWorkerClose).toHaveBeenCalledOnce(); + expect(workerA.client.close).toHaveBeenCalledOnce(); + expect(writes).toEqual([]); + }); + + it("pre-binds child session authority from an authorized parent update", () => { + const makeWorker = (workerId: string, activeId: string) => ({ + descriptor: { workerId, rootActiveSessionId: activeId }, + client: { close: vi.fn() }, + authorizedActiveSessionIds: new Set([activeId]), + summaries: new Map(), + snapshotCache: new Map(), + transcriptCaches: new Map(), + snapshotGenerations: new Map(), + snapshotLoads: new Map(), + }); + const workerA = makeWorker("worker-a", "active-a"); + const workerB = makeWorker("worker-b", "active-b"); + const handleWorkerClose = vi.fn(); + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers: new Map([ + ["worker-a", workerA], + ["worker-b", workerB], + ]), + clients: new Set(), + streamReconstructor: { observe: vi.fn() }, + handleWorkerClose, + invalidateWorkerSnapshot: vi.fn(), + refreshWorkerSummaries: vi.fn(async () => undefined), + }) as { + handleWorkerFrame(residentWorker: typeof workerA, frame: PrivateFrame): void; + }; + const childUpdate = (childActiveSessionId: string) => ({ + header: { + kind: "outbound" as const, + outboundType: "session_event" as const, + activeSessionId: "active-a", + sessionEventType: "rlm_child_update", + }, + payload: Buffer.from( + JSON.stringify({ + type: "session_event", + activeSessionId: "active-a", + event: { + type: "rlm_child_update", + child: { + id: "child-a", + activeSessionId: childActiveSessionId, + label: "child", + status: "running", + sessionDir: "/tmp/child", + }, + }, + attribution: { scope: "session" }, + }), + ), + }); + + supervisor.handleWorkerFrame(workerA, childUpdate("active-child")); + supervisor.handleWorkerFrame(workerA, { + header: { + kind: "outbound", + outboundType: "session_event", + activeSessionId: "active-child", + sessionEventType: "turn_start", + }, + payload: Buffer.from( + JSON.stringify({ + type: "session_event", + activeSessionId: "active-child", + event: { type: "turn_start" }, + attribution: { scope: "session" }, + }), + ), + }); + expect(workerA.authorizedActiveSessionIds).toContain("active-child"); + expect(handleWorkerClose).not.toHaveBeenCalled(); + + supervisor.handleWorkerFrame(workerA, childUpdate("active-b")); + expect(handleWorkerClose).toHaveBeenCalledOnce(); + expect(workerA.client.close).toHaveBeenCalledOnce(); + }); + it("shapes correlated worker events separately for capable and legacy clients", () => { const activeSessionId = "active-correlated"; const capableWrites: string[] = []; diff --git a/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts b/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts index b0b6290e63..1f24ee62b8 100644 --- a/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts +++ b/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts @@ -13,11 +13,16 @@ import { type DaemonAttachResult, type DaemonOutbound, } from "../../../src/modes/daemon/daemon-protocol.js"; +import { + type DaemonWorkerFrameHeader, + isDaemonWorkerFrameHeader, +} from "../../../src/modes/daemon/daemon-worker-protocol.js"; import { createSnapshotTranscriptChunks, SnapshotTranscriptCache, type SnapshotTranscriptChunkSource, } from "../../../src/modes/daemon/snapshot-transcript-cache.js"; +import { encodePrivateFrameParts, PrivateFrameDecoder } from "../../../src/modes/session-worker/private-framing.js"; const tempDirectories: string[] = []; const activeSessionId = "active-4601"; @@ -204,7 +209,7 @@ describe("ENG-4601 worker snapshot cache", () => { // Fully consume the fresh generation to prove its new backing is readable. } second.dispose?.(); - }); + }, 15_000); it.each([1, 2])( "keeps %i same-ID worker stream(s) bounded by socket drain without shared cleanup", @@ -471,6 +476,80 @@ describe("ENG-4601 worker snapshot cache", () => { socket.destroy(); }); + it("keeps an aborted queued private frame behind its still-writing predecessor", async () => { + const root = tempDirectory(); + const daemon = new AgentDaemon(join(root, "worker.sock"), { + defaultSessionConfig: { agentDir: root, cwd: root }, + createRuntime: async () => { + throw new Error("unexpected runtime creation"); + }, + }); + const { client, socket } = socketClient("worker-private-frame-reservation"); + const written: Buffer[] = []; + let writes = 0; + vi.spyOn(socket, "write").mockImplementation(((chunk: Uint8Array) => { + written.push(Buffer.from(chunk)); + writes++; + return writes !== 1; + }) as typeof socket.write); + const frame = (snapshot: string, payload: string) => + encodePrivateFrameParts( + { + kind: "outbound", + outboundType: "session_snapshot_chunk", + activeSessionId, + snapshotId: snapshot, + payloadEncoding: "jsonl", + }, + [Buffer.from(payload)], + ); + const internals = daemon as unknown as { + writePrivateFrameParts( + client: DaemonSocketClient, + parts: readonly Buffer[], + signal?: AbortSignal, + ): Promise; + write(client: DaemonSocketClient, message: DaemonOutbound): boolean; + }; + + const first = internals.writePrivateFrameParts(client, frame("frame-a", "payload-a")); + for (let attempt = 0; attempt < 10 && socket.listenerCount("drain") === 0; attempt++) { + await Promise.resolve(); + } + const abortQueued = new AbortController(); + const second = internals.writePrivateFrameParts(client, frame("frame-b", "payload-b"), abortQueued.signal); + internals.write(client, { + type: "session_closed", + activeSessionId, + reason: "killed", + }); + abortQueued.abort(); + + await expect(second).rejects.toThrow("was aborted"); + expect(written).toHaveLength(1); + expect(client.privateFrameWriteTail).toBeDefined(); + + socket.emit("drain"); + await expect(first).resolves.toBe(true); + await client.privateFrameWriteTail; + await Promise.resolve(); + + const decoder = new PrivateFrameDecoder(isDaemonWorkerFrameHeader); + const decoded = decoder.push(Buffer.concat(written)); + decoder.finish(); + expect( + decoded.map((entry) => + entry.header.kind === "outbound" ? (entry.header.snapshotId ?? entry.header.outboundType) : "command", + ), + ).toEqual(["frame-a", "session_closed"]); + expect(decoded.map((entry) => entry.payload.toString("utf8"))).not.toContain("payload-b"); + expect(client.privateFrameWriteTail).toBeUndefined(); + expect(socket.listenerCount("drain")).toBe(0); + expect(socket.listenerCount("close")).toBe(0); + expect(socket.listenerCount("error")).toBe(0); + socket.destroy(); + }); + it("closes a permanently backpressured worker channel after an aborted snapshot", async () => { vi.useFakeTimers(); try { diff --git a/packages/coding-agent/test/suite/regressions/4677-snapshot-catchup-replacement.test.ts b/packages/coding-agent/test/suite/regressions/4677-snapshot-catchup-replacement.test.ts index 9a099dcbde..ef3d009c76 100644 --- a/packages/coding-agent/test/suite/regressions/4677-snapshot-catchup-replacement.test.ts +++ b/packages/coding-agent/test/suite/regressions/4677-snapshot-catchup-replacement.test.ts @@ -245,6 +245,7 @@ describe("ENG-4677 snapshot catch-up replacement", () => { ); worker.transcriptCaches.clear(); worker.snapshotCache.clear(); + worker.snapshotCache.set(activeSessionId, firstResult); worker.snapshotGenerations.clear(); const internals = supervisor as unknown as { getOrCreateTranscriptCache(worker: object, result: DaemonAttachResult): Promise; @@ -253,6 +254,7 @@ describe("ENG-4677 snapshot catch-up replacement", () => { const first = await internals.getOrCreateTranscriptCache(worker, firstResult); const firstId = first.snapshotId; const firstChunkCount = first.chunkCount; + worker.snapshotCache.set(activeSessionId, secondResult); const second = await internals.getOrCreateTranscriptCache(worker, secondResult); expect(firstId).not.toBe(second.snapshotId); @@ -264,6 +266,53 @@ describe("ENG-4677 snapshot catch-up replacement", () => { second.dispose(); }); + it("does not publish a legacy transcript after its snapshot selection is invalidated", async () => { + const root = tempDirectory(); + const supervisor = new DaemonSupervisor(join(root, "supervisor.sock"), { + defaultSessionConfig: { agentDir: root, cwd: root }, + descriptorDir: join(root, "state"), + }); + const withoutStream = (content: string, sequence: number): DaemonAttachResult => { + const result = streamedResult("legacy", 1, sequence); + const { snapshotStream: _snapshotStream, ...full } = result; + full.snapshot = { + ...full.snapshot, + messages: [{ role: "user", content, timestamp: sequence }] as unknown as AgentMessage[], + }; + return full; + }; + const staleResult = withoutStream(`stale-${"x".repeat(2 * 1024 * 1024)}`, 7); + const freshResult = withoutStream("fresh", 8); + const worker = workerHarness( + staleResult, + new SnapshotTranscriptCache({ activeSessionId, snapshotId: "placeholder", cacheRoot: root }), + ); + worker.transcriptCaches.clear(); + worker.snapshotGenerations.clear(); + worker.snapshotCache.set(activeSessionId, staleResult); + const internals = supervisor as unknown as { + workers: Map; + getOrCreateTranscriptCache(worker: object, result: DaemonAttachResult): Promise; + invalidateWorkerSnapshot(worker: object, activeSessionId: string, transcriptChanged?: boolean): void; + transcriptPreparations?: WeakMap>; + }; + internals.workers.set(worker.descriptor.workerId, worker); + + const stalePreparation = internals.getOrCreateTranscriptCache(worker, staleResult); + expect(internals.transcriptPreparations?.get(worker)?.has(activeSessionId)).toBe(true); + internals.invalidateWorkerSnapshot(worker, activeSessionId, false); + + await expect(stalePreparation).rejects.toThrow("changed during transcript preparation"); + expect(worker.snapshotCache.has(activeSessionId)).toBe(false); + expect(worker.transcriptCaches.has(activeSessionId)).toBe(false); + expect(worker.snapshotGenerations.get(activeSessionId)?.size ?? 0).toBe(0); + + worker.snapshotCache.set(activeSessionId, freshResult); + const freshTranscript = await internals.getOrCreateTranscriptCache(worker, freshResult); + expect(Buffer.concat([...freshTranscript]).toString("utf8")).toContain("fresh"); + freshTranscript.dispose(); + }); + it("lets an incomplete retained snapshot finish after a newer generation begins", async () => { const root = tempDirectory(); const supervisor = new DaemonSupervisor(join(root, "supervisor.sock"), { From ba86fb349ba43bf4179b58bde19b34eacce4d8f2 Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Sat, 29 Aug 2026 19:41:13 -0600 Subject: [PATCH 05/12] test(daemon): model worker frame authority --- ...4602-snapshot-transfer-idempotency.test.ts | 30 +++++++++++++++++-- .../4677-snapshot-catchup-replacement.test.ts | 18 ++++++++++- 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/packages/coding-agent/test/suite/regressions/4602-snapshot-transfer-idempotency.test.ts b/packages/coding-agent/test/suite/regressions/4602-snapshot-transfer-idempotency.test.ts index b777ed77b6..572b8a2e81 100644 --- a/packages/coding-agent/test/suite/regressions/4602-snapshot-transfer-idempotency.test.ts +++ b/packages/coding-agent/test/suite/regressions/4602-snapshot-transfer-idempotency.test.ts @@ -26,7 +26,8 @@ const activeSessionId = "active-4602"; const snapshotId = "snapshot-4602"; interface WorkerHarness { - descriptor: { workerId: string; lifecycle: "ready" | "recovering"; pid: number }; + descriptor: { workerId: string; rootActiveSessionId: string; lifecycle: "ready" | "recovering"; pid: number }; + authorizedActiveSessionIds: Set; client?: { close: ReturnType; request: ReturnType }; summaries: Map; snapshotCache: Map; @@ -118,7 +119,13 @@ function workerHarness() { throw new Error("unexpected snapshot reload"); }); const worker: WorkerHarness = { - descriptor: { workerId: "worker-4602", lifecycle: "ready", pid: 987_654_321 }, + descriptor: { + workerId: "worker-4602", + rootActiveSessionId: activeSessionId, + lifecycle: "ready", + pid: 987_654_321, + }, + authorizedActiveSessionIds: new Set([activeSessionId]), client: { close, request }, summaries: new Map([[activeSessionId, summary()]]), snapshotCache: new Map(), @@ -131,6 +138,10 @@ function workerHarness() { return { close, request, worker }; } +function registerWorker(supervisor: DaemonSupervisor, worker: WorkerHarness): void { + (supervisor as unknown as { workers: Map }).workers.set(worker.descriptor.workerId, worker); +} + function socketClient(id: string, socket: PassThrough): DaemonSocketClient { return { id, @@ -399,6 +410,7 @@ describe("ENG-4602 snapshot transfer containment", () => { descriptorDir: "/tmp/eng-4602-supervisor-fresh-chunk-state", }); const { close, worker } = workerHarness(); + registerWorker(supervisor, worker); const internals = supervisor as unknown as { handleWorkerFrame(worker: WorkerHarness, frame: PrivateFrame): void; }; @@ -438,6 +450,7 @@ describe("ENG-4602 snapshot transfer containment", () => { descriptorDir: "/tmp/eng-4602-supervisor-fresh-end-state", }); const { close, worker } = workerHarness(); + registerWorker(supervisor, worker); const internals = supervisor as unknown as { handleWorkerFrame(worker: WorkerHarness, frame: PrivateFrame): void; }; @@ -482,6 +495,7 @@ describe("ENG-4602 snapshot transfer containment", () => { descriptorDir: "/tmp/eng-4602-supervisor-invalid-begin-state", }); const { close, worker } = workerHarness(); + registerWorker(supervisor, worker); const internals = supervisor as unknown as { handleWorkerFrame(worker: WorkerHarness, frame: PrivateFrame): void; }; @@ -503,6 +517,7 @@ describe("ENG-4602 snapshot transfer containment", () => { descriptorDir: "/tmp/eng-4602-supervisor-failure-state", }); const { close, worker } = workerHarness(); + registerWorker(supervisor, worker); worker.summaries.set("active-4602-sibling", { ...summary(), id: "active-4602-sibling", @@ -580,6 +595,7 @@ describe("ENG-4602 snapshot transfer containment", () => { descriptorDir: "/tmp/eng-4602-supervisor-state", }); const { close, request, worker } = workerHarness(); + registerWorker(supervisor, worker); const client = socketClient("public", new PassThrough()); const streamSnapshot = vi.fn(async () => {}); const internals = supervisor as unknown as { @@ -692,6 +708,7 @@ describe("ENG-4602 snapshot transfer containment", () => { descriptorDir: "/tmp/eng-4602-supervisor-gate-state", }); const { close, request, worker } = workerHarness(); + registerWorker(supervisor, worker); const client = socketClient("catchup", new PassThrough()); const streamSnapshot = vi.fn(async () => {}); const internals = supervisor as unknown as { @@ -757,6 +774,7 @@ describe("ENG-4602 snapshot transfer containment", () => { descriptorDir: "/tmp/eng-4602-supervisor-stop-state", }); const { close, request, worker } = workerHarness(); + registerWorker(supervisor, worker); const client = socketClient("catchup", new PassThrough()); const streamSnapshot = vi.fn(async () => {}); const persistWorker = vi.fn(); @@ -818,6 +836,7 @@ describe("ENG-4602 snapshot transfer containment", () => { descriptorDir: "/tmp/eng-4602-supervisor-reader-stop-state", }); const { worker } = workerHarness(); + registerWorker(supervisor, worker); const persistWorker = vi.fn(); const internals = supervisor as unknown as { persistWorker: typeof persistWorker; @@ -877,6 +896,7 @@ describe("ENG-4602 snapshot transfer containment", () => { expect(recoverWorker).not.toHaveBeenCalled(); const completed = workerHarness(); + registerWorker(supervisor, completed.worker); for (const message of [frames.begin, frames.chunk, frames.end]) { internals.handleWorkerFrame(completed.worker, frame(message)); } @@ -886,6 +906,7 @@ describe("ENG-4602 snapshot transfer containment", () => { expect(completed.worker.client).toBeDefined(); const mismatchedEnd = workerHarness(); + registerWorker(supervisor, mismatchedEnd.worker); for (const message of [frames.begin, frames.chunk, frames.end, frames.begin, frames.chunk]) { internals.handleWorkerFrame(mismatchedEnd.worker, frame(message)); } @@ -895,6 +916,7 @@ describe("ENG-4602 snapshot transfer containment", () => { expect(mismatchedEnd.worker.client).toBeDefined(); const replaced = workerHarness(); + registerWorker(supervisor, replaced.worker); internals.handleWorkerFrame(replaced.worker, frame(frames.begin)); internals.handleWorkerFrame(replaced.worker, frame({ ...frames.begin, snapshotId: "snapshot-4602-new" })); expect(replaced.close).not.toHaveBeenCalled(); @@ -921,6 +943,7 @@ describe("ENG-4602 snapshot transfer containment", () => { socket.resume(); const client = socketClient("bounded-retry", socket); const { worker } = workerHarness(); + registerWorker(supervisor, worker); const records: DaemonOutbound[] = []; let attachIndex = 0; const attachClient = vi.fn(async () => { @@ -974,6 +997,7 @@ describe("ENG-4602 snapshot transfer containment", () => { descriptorDir: "/tmp/eng-4602-supervisor-deadline-state", }); const { worker } = workerHarness(); + registerWorker(supervisor, worker); const client = socketClient("deadline", new PassThrough()); const records: DaemonOutbound[] = []; const internals = supervisor as unknown as { @@ -1012,6 +1036,7 @@ describe("ENG-4602 snapshot transfer containment", () => { descriptorDir: "/tmp/eng-4602-supervisor-epoch-state", }); const { worker } = workerHarness(); + registerWorker(supervisor, worker); const client = socketClient("epoch", new PassThrough()); const transcript = new SnapshotTranscriptCache({ activeSessionId, snapshotId, cacheRoot: "/tmp" }); const records: DaemonOutbound[] = []; @@ -1072,6 +1097,7 @@ describe("ENG-4602 snapshot transfer containment", () => { failedClient.attachedActiveSessionIds.add("active-4602-sibling"); const siblingClient = socketClient("sibling", siblingSocket); const { close, worker } = workerHarness(); + registerWorker(supervisor, worker); worker.snapshotCache.set(activeSessionId, result); worker.transcriptCaches.set(activeSessionId, transcript); const internals = supervisor as unknown as { diff --git a/packages/coding-agent/test/suite/regressions/4677-snapshot-catchup-replacement.test.ts b/packages/coding-agent/test/suite/regressions/4677-snapshot-catchup-replacement.test.ts index ef3d009c76..afbad7d52c 100644 --- a/packages/coding-agent/test/suite/regressions/4677-snapshot-catchup-replacement.test.ts +++ b/packages/coding-agent/test/suite/regressions/4677-snapshot-catchup-replacement.test.ts @@ -23,7 +23,8 @@ const directories: string[] = []; interface WorkerHarness { descriptor: { workerId: string; rootActiveSessionId: string; lifecycle: "ready"; pid: number }; - client?: { request: ReturnType }; + client?: { close: ReturnType; request: ReturnType }; + authorizedActiveSessionIds: Set; summaries: Map; snapshotCache: Map; transcriptCaches: Map; @@ -134,10 +135,12 @@ function workerHarness(result: DaemonAttachResult, transcript: SnapshotTranscrip pid: 4677, }, client: { + close: vi.fn(), request: vi.fn(async () => { throw new Error("unexpected snapshot reload"); }), }, + authorizedActiveSessionIds: new Set([activeSessionId]), summaries: new Map([[activeSessionId, result.snapshot.summary]]), snapshotCache: new Map([[activeSessionId, result]]), transcriptCaches: new Map([[activeSessionId, transcript]]), @@ -148,6 +151,10 @@ function workerHarness(result: DaemonAttachResult, transcript: SnapshotTranscrip }; } +function registerWorker(supervisor: DaemonSupervisor, worker: WorkerHarness): void { + (supervisor as unknown as { workers: Map }).workers.set(worker.descriptor.workerId, worker); +} + describe("ENG-4677 snapshot catch-up replacement", () => { it("atomically remaps an alias snapshot fence before canonical events can pass", () => { const root = tempDirectory(); @@ -163,6 +170,7 @@ describe("ENG-4677 snapshot catch-up replacement", () => { cacheRoot: root, }); const worker = workerHarness(streamedResult("alias-snapshot", 0, 1), transcript); + registerWorker(supervisor, worker); const written: Buffer[] = []; const internals = supervisor as unknown as { clients: Set; @@ -243,6 +251,7 @@ describe("ENG-4677 snapshot catch-up replacement", () => { cacheRoot: root, }), ); + registerWorker(supervisor, worker); worker.transcriptCaches.clear(); worker.snapshotCache.clear(); worker.snapshotCache.set(activeSessionId, firstResult); @@ -287,6 +296,7 @@ describe("ENG-4677 snapshot catch-up replacement", () => { staleResult, new SnapshotTranscriptCache({ activeSessionId, snapshotId: "placeholder", cacheRoot: root }), ); + registerWorker(supervisor, worker); worker.transcriptCaches.clear(); worker.snapshotGenerations.clear(); worker.snapshotCache.set(activeSessionId, staleResult); @@ -328,6 +338,7 @@ describe("ENG-4677 snapshot catch-up replacement", () => { targetChunkBytes: 1, }); const worker = workerHarness(streamedResult(firstSnapshotId, 1, 1), placeholder); + registerWorker(supervisor, worker); worker.snapshotCache.clear(); worker.transcriptCaches.clear(); const client = socketClient("incomplete-reader"); @@ -445,10 +456,12 @@ describe("ENG-4677 snapshot catch-up replacement", () => { targetChunkBytes: 1, }); const worker = workerHarness(firstResult, placeholder); + registerWorker(supervisor, worker); worker.snapshotCache.clear(); worker.transcriptCaches.clear(); let resolveAttach!: (response: { success: true; data: DaemonAttachResult }) => void; worker.client = { + close: vi.fn(), request: vi.fn( () => new Promise<{ success: true; data: DaemonAttachResult }>((resolve) => { @@ -546,6 +559,7 @@ describe("ENG-4677 snapshot catch-up replacement", () => { const waiter = firstTranscript.waitForChunk(1); void waiter.catch(() => undefined); const worker = workerHarness(firstResult, firstTranscript); + registerWorker(supervisor, worker); const internals = supervisor as unknown as { handleWorkerFrame(worker: WorkerHarness, frame: PrivateFrame): void; }; @@ -586,6 +600,7 @@ describe("ENG-4677 snapshot catch-up replacement", () => { targetChunkBytes: 1, }); const worker = workerHarness(firstResult, placeholder); + registerWorker(supervisor, worker); worker.snapshotCache.clear(); worker.transcriptCaches.clear(); const client = socketClient("validation-waiter"); @@ -763,6 +778,7 @@ describe("ENG-4677 snapshot catch-up replacement", () => { targetChunkBytes: 1, }); const worker = workerHarness(firstResult, firstTranscript); + registerWorker(supervisor, worker); const client = socketClient("slow-client"); const written: DaemonOutbound[] = []; let releaseFirstChunk!: (accepted: boolean) => void; From 5d1c9e224b04c351bb1cdb69b216bea0fdd82d3b Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Sat, 29 Aug 2026 20:06:30 -0600 Subject: [PATCH 06/12] fix(daemon): fence attach and worker authority --- .../src/modes/daemon/daemon-mode.ts | 74 +++++++--- .../src/modes/daemon/daemon-supervisor.ts | 106 ++++++++++---- .../coding-agent/test/daemon-mode.test.ts | 76 +++++++++- .../test/daemon-supervisor-monitor.test.ts | 135 ++++++++++++++++++ .../4601-worker-snapshot-cache.test.ts | 63 ++++++++ ...4602-snapshot-transfer-idempotency.test.ts | 4 +- 6 files changed, 407 insertions(+), 51 deletions(-) diff --git a/packages/coding-agent/src/modes/daemon/daemon-mode.ts b/packages/coding-agent/src/modes/daemon/daemon-mode.ts index 0eb002e7c2..fe4c9ed0fb 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-mode.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-mode.ts @@ -3881,16 +3881,39 @@ export class AgentDaemon { ? markClientSnapshotStreaming(client, state.activeSessionId) : undefined; const snapshotAttemptSignal = snapshotSignal ? createSnapshotAttemptSignal(snapshotSignal) : undefined; - let result: DaemonAttachResult; + let result: DaemonAttachResult | undefined; state.pendingAttaches++; try { - result = await this.createAttachResult(client, state, command); - if ( - this.sessions.get(state.activeSessionId) !== state || - this.closingSessions.has(state.activeSessionId) - ) { - throw new BoundSessionUnavailableError( - `Active session ${state.activeSessionId} closed during attach`, + for (let attempt = 0; attempt <= MAX_SESSION_SNAPSHOT_STABILIZATION_RETRIES; attempt++) { + const selectedSession = state.runtime.session; + const candidate = await this.createAttachResult(client, state, command); + if ( + this.sessions.get(state.activeSessionId) === state && + !this.closingSessions.has(state.activeSessionId) && + state.runtime.session === selectedSession && + state.runtime.session.sessionId === candidate.snapshot.state.sessionId && + state.eventGeneration === candidate.snapshot.lastEventCursor?.generation && + state.lastEventSequence === candidate.snapshot.lastEventSequence + ) { + result = candidate; + // This is the attach admission cut. There must be no await between + // validating the snapshot cursor and registering the client. + state.clients.add(client); + client.attachedActiveSessionIds.add(state.activeSessionId); + break; + } + } + if (!result) { + if ( + this.sessions.get(state.activeSessionId) !== state || + this.closingSessions.has(state.activeSessionId) + ) { + throw new BoundSessionUnavailableError( + `Active session ${state.activeSessionId} closed during attach`, + ); + } + throw new Error( + `Session ${state.activeSessionId} changed while its attach snapshot was being prepared`, ); } } catch (error) { @@ -3902,8 +3925,6 @@ export class AgentDaemon { } finally { state.pendingAttaches--; } - state.clients.add(client); - client.attachedActiveSessionIds.add(state.activeSessionId); if (deferClientEnv && clientEnv) { this.updateRestart?.deferredClientEnv.push({ client, @@ -5114,6 +5135,10 @@ export class AgentDaemon { ): Promise { const capabilities = daemonClientCapabilitiesForSession(client, state.activeSessionId); const snapshot = await this.createSessionSnapshot(state, capabilities.has("correlated_prompt_lifecycle_v1")); + const snapshotCursor = snapshot.lastEventCursor ?? { + generation: state.eventGeneration, + sequence: snapshot.lastEventSequence, + }; const replay = command.resumeCursor?.activeSessionId && command.resumeCursor.activeSessionId !== state.activeSessionId ? { @@ -5122,14 +5147,11 @@ export class AgentDaemon { "sequence" in command.resumeCursor ? command.resumeCursor.sequence : command.resumeCursor.eventSequence, - toSequence: state.lastEventSequence, - toCursor: { - generation: state.eventGeneration, - sequence: state.lastEventSequence, - }, + toSequence: snapshot.lastEventSequence, + toCursor: snapshotCursor, reason: "resume_cursor_session_mismatch", } - : createDaemonReplayInfo(command.resumeCursor, state.lastEventSequence, state.eventGeneration); + : createDaemonReplayInfo(command.resumeCursor, snapshot.lastEventSequence, snapshotCursor.generation); // Slim clients read summary/messages from the snapshot; duplicating them at // the top level would serialize the full history twice more per attach. const slim = capabilities.has("slim_attach"); @@ -5139,11 +5161,8 @@ export class AgentDaemon { ...(slim ? {} : { state: snapshot.summary, messages: snapshot.messages }), snapshot, replay, - lastEventSequence: state.lastEventSequence, - lastEventCursor: { - generation: state.eventGeneration, - sequence: state.lastEventSequence, - }, + lastEventSequence: snapshot.lastEventSequence, + lastEventCursor: snapshotCursor, client: { id: client.id, capabilities: [...capabilities], @@ -5536,19 +5555,32 @@ export class AgentDaemon { void tail.finally(() => { if (client.privateFrameWriteTail === tail) client.privateFrameWriteTail = undefined; }); + let committed = false; + let completed = false; try { if (previous) await waitForSnapshotAttempt(previous, signal ?? AbortSignal.timeout(30_000), "private-frame"); for (const part of parts) { for (let offset = 0; offset < part.length; offset += WORKER_PRIVATE_FRAME_WRITE_BYTES) { if (signal?.aborted || client.socket.destroyed) return false; const slice = part.subarray(offset, offset + WORKER_PRIVATE_FRAME_WRITE_BYTES); + if (slice.length === 0) continue; + // socket.write(false) still accepts the bytes. Once the length prefix + // is committed, an interrupted write must close the channel rather + // than let a later frame fill the announced payload. + committed = true; if (client.socket.write(slice)) continue; client.backpressured = true; if (!(await this.waitForWorkerSnapshotDrain(client, signal, drainTimeoutMs))) return false; } } + completed = true; return true; } finally { + if (committed && !completed && !client.socket.destroyed) { + client.socket.destroy( + new Error("Private frame write was interrupted after its length prefix was committed"), + ); + } releaseCurrent(); } } diff --git a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts index 7befb30ab6..f245317e20 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts @@ -314,8 +314,10 @@ interface ResidentWorker { snapshotGenerations: Map>; snapshotLoads: Map>; snapshotLoadSnapshotIds: Map; - /** Active-session IDs this authenticated resident worker may emit. */ + /** Active-session IDs the current authenticated worker channel may emit. */ authorizedActiveSessionIds: Set; + /** Invalidates stale roster refreshes and provisional child grants across channel changes. */ + authorityRevision?: number; recovery?: Promise; deferredRecovery?: Promise; intentionalStop: boolean; @@ -3012,10 +3014,12 @@ export class DaemonSupervisor { 1000, ); await this.assertRecoveryAllowed(); - client.onFrame((frame) => this.handleWorkerFrame(worker, frame)); + client.onFrame((frame) => this.handleWorkerFrame(worker, frame, client)); client.onClose((error) => void this.handleWorkerClose(worker, client, error)); worker.client?.close(); worker.client = client; + worker.authorityRevision = (worker.authorityRevision ?? 0) + 1; + worker.authorizedActiveSessionIds = new Set([worker.descriptor.rootActiveSessionId]); return client; } catch (error) { lastError = error; @@ -3123,6 +3127,8 @@ export class DaemonSupervisor { } this.abortTranscriptPreparations(worker, error); worker.client = undefined; + worker.authorityRevision = (worker.authorityRevision ?? 0) + 1; + worker.authorizedActiveSessionIds = new Set([worker.descriptor.rootActiveSessionId]); this.invalidateWorkerSessionInputPauses(worker, "Session worker disconnected while input was paused"); const interrupted = new Map>(); for (const [activeSessionId, generations] of worker.snapshotGenerations ?? []) { @@ -3832,15 +3838,20 @@ export class DaemonSupervisor { return worker.authorizedActiveSessionIds; } - private workerOwnsActiveSessionId(worker: ResidentWorker, activeSessionId: string): boolean { + private workerOwnsActiveSessionId( + worker: ResidentWorker, + activeSessionId: string, + sourceClient?: DaemonWorkerClient, + ): boolean { // Prototype-only unit harnesses omit both identity fields. Every production // worker has a descriptor and remains registered while its channel is valid. if (!worker.descriptor) return !this.workers; if (this.workers && this.workers.get(worker.descriptor.workerId) !== worker) return false; + if (sourceClient && worker.client !== sourceClient) return false; return this.workerAuthorizedActiveSessionIds(worker).has(activeSessionId); } - private bindWorkerActiveSessionIds(worker: ResidentWorker, activeSessionIds: Iterable): void { + private validateWorkerActiveSessionIds(worker: ResidentWorker, activeSessionIds: Iterable): Set { const candidates = new Set(activeSessionIds); for (const activeSessionId of candidates) { if (!activeSessionId) throw new Error("Session worker returned an empty active-session ID"); @@ -3851,14 +3862,22 @@ export class DaemonSupervisor { } } } + return candidates; + } + + private provisionWorkerActiveSessionIds(worker: ResidentWorker, activeSessionIds: Iterable): void { + const candidates = this.validateWorkerActiveSessionIds(worker, activeSessionIds); + worker.authorityRevision = (worker.authorityRevision ?? 0) + 1; const authorized = this.workerAuthorizedActiveSessionIds(worker); for (const activeSessionId of candidates) authorized.add(activeSessionId); } - private rejectWorkerFrame(worker: ResidentWorker, error: Error): void { - const workerClient = worker.client; + private rejectWorkerFrame(worker: ResidentWorker, error: Error, sourceClient?: DaemonWorkerClient): void { + const workerClient = sourceClient ?? worker.client; if (!workerClient) return; - this.handleWorkerClose(worker, workerClient, error); + if (worker.client === workerClient) { + void this.handleWorkerClose(worker, workerClient, error); + } workerClient.close(); } @@ -3866,17 +3885,28 @@ export class DaemonSupervisor { if (this.isWorkerStopping(worker)) { throw new Error("Session worker is stopping"); } - if (!worker.client) { + const workerClient = worker.client; + if (!workerClient) { throw new Error("Session worker is not connected"); } - const response = await worker.client.request({ type: "list" }, 5000); + const authorityRevision = (worker.authorityRevision ?? 0) + 1; + worker.authorityRevision = authorityRevision; + const response = await workerClient.request({ type: "list" }, 5000); + if ( + worker.client !== workerClient || + worker.authorityRevision !== authorityRevision || + (this.workers && this.workers.get(worker.descriptor.workerId) !== worker) + ) { + throw new Error("Session worker authority changed during roster refresh"); + } const summaries = sessionSummariesFromResponse(response); const nextSummaries = new Map(summaries.map((summary) => [summary.activeSessionId ?? summary.id, summary])); const root = nextSummaries.get(worker.descriptor.rootActiveSessionId); if (recovery && !root) { throw new Error(`Session worker omitted its root session during recovery`); } - this.bindWorkerActiveSessionIds(worker, nextSummaries.keys()); + const nextAuthority = this.validateWorkerActiveSessionIds(worker, nextSummaries.keys()); + worker.authorizedActiveSessionIds = nextAuthority; worker.summaries = nextSummaries; for (const summary of summaries) { const activeSessionId = summary.activeSessionId ?? summary.id; @@ -5020,7 +5050,22 @@ export class DaemonSupervisor { ); } - private handleWorkerFrame(worker: ResidentWorker, frame: PrivateFrame): void { + private handleWorkerFrame( + worker: ResidentWorker, + frame: PrivateFrame, + sourceClient?: DaemonWorkerClient, + ): void { + if ( + (worker.descriptor && this.workers && this.workers.get(worker.descriptor.workerId) !== worker) || + (sourceClient && worker.client !== sourceClient) + ) { + this.rejectWorkerFrame( + worker, + new Error("Worker frame came from a stale authenticated channel"), + sourceClient, + ); + return; + } if (frame.header.kind !== "outbound") { return; } @@ -5034,15 +5079,23 @@ export class DaemonSupervisor { } = frame.header; if (outboundType === "heartbeats_changed") { if (activeSessionId) { - this.rejectWorkerFrame(worker, new Error("Worker heartbeat frame was unexpectedly session-scoped")); + this.rejectWorkerFrame( + worker, + new Error("Worker heartbeat frame was unexpectedly session-scoped"), + sourceClient, + ); return; } worker.heartbeatSnapshotStale = true; this.broadcastHeartbeatsChanged(); return; } - if (activeSessionId && !this.workerOwnsActiveSessionId(worker, activeSessionId)) { - this.rejectWorkerFrame(worker, new Error(`Worker emitted an unauthorized active session ${activeSessionId}`)); + if (activeSessionId && !this.workerOwnsActiveSessionId(worker, activeSessionId, sourceClient)) { + this.rejectWorkerFrame( + worker, + new Error(`Worker emitted an unauthorized active session ${activeSessionId}`), + sourceClient, + ); return; } if ( @@ -5052,12 +5105,7 @@ export class DaemonSupervisor { outboundType === "session_snapshot_failed") && !activeSessionId ) { - const error = new Error("Worker returned an uncorrelatable snapshot frame"); - const workerClient = worker.client; - if (workerClient) { - this.handleWorkerClose(worker, workerClient, error); - workerClient.close(); - } + this.rejectWorkerFrame(worker, new Error("Worker returned an uncorrelatable snapshot frame"), sourceClient); return; } if (outboundType === "session_snapshot_begin" && activeSessionId) { @@ -5493,7 +5541,11 @@ export class DaemonSupervisor { return; } if (compactValue.activeSessionId !== activeSessionId) { - this.rejectWorkerFrame(worker, new Error("Worker compact payload did not match its active-session header")); + this.rejectWorkerFrame( + worker, + new Error("Worker compact payload did not match its active-session header"), + sourceClient, + ); return; } const reconstructed = this.streamReconstructor.reconstruct(compactValue); @@ -5507,7 +5559,7 @@ export class DaemonSupervisor { try { decodedOutbound = JSON.parse(frame.payload.toString("utf8")) as DaemonOutbound; } catch { - this.rejectWorkerFrame(worker, new Error("Worker returned malformed session-scoped JSON")); + this.rejectWorkerFrame(worker, new Error("Worker returned malformed session-scoped JSON"), sourceClient); return; } } @@ -5515,7 +5567,11 @@ export class DaemonSupervisor { decodedOutbound.type !== outboundType || (decodedOutbound as { activeSessionId?: unknown }).activeSessionId !== activeSessionId ) { - this.rejectWorkerFrame(worker, new Error("Worker payload did not match its routed frame header")); + this.rejectWorkerFrame( + worker, + new Error("Worker payload did not match its routed frame header"), + sourceClient, + ); return; } this.streamReconstructor.observe(decodedOutbound); @@ -5525,9 +5581,9 @@ export class DaemonSupervisor { decodedOutbound.event.child.activeSessionId ) { try { - this.bindWorkerActiveSessionIds(worker, [decodedOutbound.event.child.activeSessionId]); + this.provisionWorkerActiveSessionIds(worker, [decodedOutbound.event.child.activeSessionId]); } catch (error) { - this.rejectWorkerFrame(worker, error instanceof Error ? error : new Error(String(error))); + this.rejectWorkerFrame(worker, error instanceof Error ? error : new Error(String(error)), sourceClient); return; } } diff --git a/packages/coding-agent/test/daemon-mode.test.ts b/packages/coding-agent/test/daemon-mode.test.ts index c843ddc278..2c4d713109 100644 --- a/packages/coding-agent/test/daemon-mode.test.ts +++ b/packages/coding-agent/test/daemon-mode.test.ts @@ -3229,6 +3229,8 @@ describe("daemon mode helpers", () => { createRuntime: vi.fn(), }); const state = makeState("active"); + state.eventGeneration = "attach-generation"; + (state.runtime as unknown as { session: { sessionId: string } }).session = { sessionId: "session-active" }; const client = makeClient("client-1", state.activeSessionId); client.attachedActiveSessionIds.clear(); let releaseSnapshot!: () => void; @@ -3237,7 +3239,13 @@ describe("daemon mode helpers", () => { }); const result = { activeSessionId: state.activeSessionId, - snapshot: { summary: {}, state: {}, messages: [] }, + snapshot: { + summary: {}, + state: { sessionId: "session-active" }, + messages: [], + lastEventSequence: 0, + lastEventCursor: { generation: state.eventGeneration, sequence: 0 }, + }, lastEventSequence: 0, } as unknown as DaemonAttachResult; const internals = daemon as unknown as { @@ -3346,11 +3354,18 @@ describe("daemon mode helpers", () => { }); const state = makeState("active"); state.eventGeneration = "generation-1"; + (state.runtime as unknown as { session: { sessionId: string } }).session = { sessionId: "session-active" }; const client = makeClient("client-1", state.activeSessionId); client.transport = "private-framed"; const result = { activeSessionId: state.activeSessionId, - snapshot: { summary: {}, state: {}, messages: [] }, + snapshot: { + summary: {}, + state: { sessionId: "session-active" }, + messages: [], + lastEventSequence: 0, + lastEventCursor: { generation: state.eventGeneration, sequence: 0 }, + }, lastEventSequence: 0, } as unknown as DaemonAttachResult; const streamWorkerSnapshot = vi.fn(async () => undefined); @@ -4549,6 +4564,50 @@ describe("daemon mode helpers", () => { } }); + it("retries attach admission when the event cursor advances after snapshot creation", async () => { + const tempDir = mkdtempSync(join(tmpdir(), "prime-agent-daemon-attach-admission-cursor-")); + try { + const fixture = makePersistedRlmDaemonFixture(tempDir); + const internals = fixture.daemon as unknown as { + createRuntime(command: Extract): Promise; + createAttachResult( + client: DaemonSocketClient, + state: ActiveSessionState, + command: Extract, + ): Promise; + handleCommand(client: DaemonSocketClient, command: DaemonCommand): Promise; + }; + const state = await internals.createRuntime({ type: "create", sessionPath: fixture.parentSessionFile }); + const client = makeClient("attach-cursor-client", state.activeSessionId); + client.attachedActiveSessionIds.clear(); + const createAttachResult = internals.createAttachResult.bind(internals); + let calls = 0; + internals.createAttachResult = vi.fn(async (attachClient, attachState, attachCommand) => { + expect(state.clients.has(client)).toBe(false); + const candidate = await createAttachResult(attachClient, attachState, attachCommand); + calls++; + if (calls === 1) state.lastEventSequence++; + return candidate; + }); + + const response = (await internals.handleCommand(client, { + type: "attach", + activeSessionId: state.activeSessionId, + capabilities: ["attach_snapshot", "event_sequence", "slim_attach"], + })) as { success: true; data: DaemonAttachResult }; + + expect(internals.createAttachResult).toHaveBeenCalledTimes(2); + expect(state.clients.has(client)).toBe(true); + expect(client.attachedActiveSessionIds).toContain(state.activeSessionId); + expect(response.data.snapshot.lastEventSequence).toBe(1); + expect(response.data.lastEventSequence).toBe(response.data.snapshot.lastEventSequence); + expect(response.data.lastEventCursor).toEqual(response.data.snapshot.lastEventCursor); + expect(response.data.replay.toCursor).toEqual(response.data.snapshot.lastEventCursor); + } finally { + rmSync(tempDir, { recursive: true, force: true }); + } + }); + it("fails snapshot stabilization when the event cursor advances on every child build", async () => { const tempDir = mkdtempSync(join(tmpdir(), "prime-agent-daemon-snapshot-event-bound-")); try { @@ -5773,7 +5832,18 @@ describe("daemon mode helpers", () => { ).releaseRlmChildSession = vi.fn(() => true); internals.createAttachResult = vi.fn(async () => { await snapshotGate; - return { activeSessionId: childState.activeSessionId, snapshot: {}, lastEventSequence: 0 }; + return { + activeSessionId: childState.activeSessionId, + snapshot: { + state: { sessionId: childState.runtime.session.sessionId }, + lastEventSequence: childState.lastEventSequence, + lastEventCursor: { + generation: childState.eventGeneration, + sequence: childState.lastEventSequence, + }, + }, + lastEventSequence: childState.lastEventSequence, + }; }); const attach = internals.handleCommand(client, { diff --git a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts index a091f4cea8..92c48e27ae 100644 --- a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts @@ -3720,6 +3720,141 @@ describe("daemon worker supervisor monitoring", () => { expect(workerA.client.close).toHaveBeenCalledOnce(); }); + it("revokes removed roster authority and rejects stale worker channels", async () => { + const root = { id: "active-root", activeSessionId: "active-root", sessionId: "session-root", cwd: "/tmp" }; + const removedChild = "active-removed-child"; + const oldClient = { close: vi.fn() }; + const currentClient = { + close: vi.fn(), + request: vi.fn(async () => success(undefined, "list", { sessions: [root] })), + }; + const worker = { + descriptor: { + workerId: "worker-roster-revoke", + rootActiveSessionId: root.activeSessionId, + createCommand: { type: "create" as const }, + }, + client: currentClient, + authorizedActiveSessionIds: new Set([root.activeSessionId, removedChild]), + authorityRevision: 1, + summaries: new Map([ + [root.activeSessionId, root as SessionSummary], + [ + removedChild, + { + id: removedChild, + activeSessionId: removedChild, + sessionId: "session-child", + cwd: "/tmp", + } as SessionSummary, + ], + ]), + snapshotCache: new Map(), + transcriptCaches: new Map(), + snapshotGenerations: new Map(), + snapshotLoads: new Map(), + intentionalStop: false, + stopRevision: 0, + }; + const handleWorkerClose = vi.fn(); + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers: new Map([[worker.descriptor.workerId, worker]]), + clients: new Set(), + streamReconstructor: { seed: vi.fn(), clear: vi.fn(), observe: vi.fn() }, + handleWorkerClose, + invalidateWorkerSnapshot: vi.fn(), + persistWorker: vi.fn(), + }) as { + refreshWorkerSummaries(target: typeof worker): Promise; + handleWorkerFrame( + target: typeof worker, + frame: PrivateFrame, + sourceClient?: typeof currentClient, + ): void; + }; + const removedChildFrame: PrivateFrame = { + header: { + kind: "outbound", + outboundType: "session_event", + activeSessionId: removedChild, + sessionEventType: "turn_start", + }, + payload: Buffer.from( + JSON.stringify({ + type: "session_event", + activeSessionId: removedChild, + event: { type: "turn_start" }, + attribution: { scope: "session" }, + }), + ), + }; + + await supervisor.refreshWorkerSummaries(worker); + expect(worker.authorizedActiveSessionIds).toEqual(new Set([root.activeSessionId])); + supervisor.handleWorkerFrame(worker, removedChildFrame, currentClient); + expect(handleWorkerClose).toHaveBeenCalledOnce(); + expect(currentClient.close).toHaveBeenCalledOnce(); + expect(worker.snapshotCache.size).toBe(0); + + handleWorkerClose.mockClear(); + currentClient.close.mockClear(); + supervisor.handleWorkerFrame(worker, removedChildFrame, oldClient as typeof currentClient); + expect(oldClient.close).toHaveBeenCalledOnce(); + expect(currentClient.close).not.toHaveBeenCalled(); + expect(handleWorkerClose).not.toHaveBeenCalled(); + }); + + it("does not let an older roster refresh resurrect revoked authority", async () => { + const root = { id: "active-root", activeSessionId: "active-root", sessionId: "session-root", cwd: "/tmp" }; + const removedChild = { + id: "active-old-child", + activeSessionId: "active-old-child", + sessionId: "session-old-child", + cwd: "/tmp", + }; + let resolveOlder!: (value: unknown) => void; + let resolveNewer!: (value: unknown) => void; + const older = new Promise((resolve) => { + resolveOlder = resolve; + }); + const newer = new Promise((resolve) => { + resolveNewer = resolve; + }); + const client = { + request: vi.fn().mockReturnValueOnce(older).mockReturnValueOnce(newer), + }; + const worker = { + descriptor: { + workerId: "worker-roster-order", + rootActiveSessionId: root.activeSessionId, + createCommand: { type: "create" as const }, + }, + client, + authorizedActiveSessionIds: new Set([root.activeSessionId, removedChild.activeSessionId]), + authorityRevision: 0, + summaries: new Map(), + intentionalStop: false, + stopRevision: 0, + }; + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers: new Map([[worker.descriptor.workerId, worker]]), + streamReconstructor: { seed: vi.fn(), clear: vi.fn() }, + persistWorker: vi.fn(), + }) as { + refreshWorkerSummaries(target: typeof worker): Promise; + }; + + const olderRefresh = supervisor.refreshWorkerSummaries(worker); + const newerRefresh = supervisor.refreshWorkerSummaries(worker); + resolveNewer(success(undefined, "list", { sessions: [root] })); + await newerRefresh; + resolveOlder(success(undefined, "list", { sessions: [root, removedChild] })); + await expect(olderRefresh).rejects.toThrow("authority changed during roster refresh"); + + expect(worker.authorizedActiveSessionIds).toEqual(new Set([root.activeSessionId])); + expect(worker.summaries.has(removedChild.activeSessionId)).toBe(false); + }); + it("shapes correlated worker events separately for capable and legacy clients", () => { const activeSessionId = "active-correlated"; const capableWrites: string[] = []; diff --git a/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts b/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts index 1f24ee62b8..42de6719bc 100644 --- a/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts +++ b/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts @@ -550,6 +550,69 @@ describe("ENG-4601 worker snapshot cache", () => { socket.destroy(); }); + it("closes a private channel when an active frame aborts after committing its length", async () => { + const root = tempDirectory(); + const daemon = new AgentDaemon(join(root, "worker.sock"), { + defaultSessionConfig: { agentDir: root, cwd: root }, + createRuntime: async () => { + throw new Error("unexpected runtime creation"); + }, + }); + const { client, socket } = socketClient("worker-private-frame-active-abort"); + socket.on("error", () => {}); + const written: Buffer[] = []; + let writes = 0; + vi.spyOn(socket, "write").mockImplementation(((chunk: Uint8Array) => { + written.push(Buffer.from(chunk)); + writes++; + return writes !== 3; + }) as typeof socket.write); + const parts = encodePrivateFrameParts( + { + kind: "outbound", + outboundType: "session_snapshot_chunk", + activeSessionId, + snapshotId: "active-abort-frame", + payloadEncoding: "jsonl", + }, + [Buffer.alloc(128 * 1024, 0x61)], + ); + const internals = daemon as unknown as { + writePrivateFrameParts( + client: DaemonSocketClient, + parts: readonly Buffer[], + signal?: AbortSignal, + ): Promise; + write(client: DaemonSocketClient, message: DaemonOutbound): boolean; + }; + const controller = new AbortController(); + + const activeWrite = internals.writePrivateFrameParts(client, parts, controller.signal); + for (let attempt = 0; attempt < 10 && socket.listenerCount("drain") === 0; attempt++) { + await Promise.resolve(); + } + expect(writes).toBe(3); + expect(socket.listenerCount("drain")).toBe(1); + controller.abort(); + + await expect(activeWrite).resolves.toBe(false); + await client.privateFrameWriteTail; + await Promise.resolve(); + const writesAfterAbort = written.length; + expect(socket.destroyed).toBe(true); + expect( + internals.write(client, { + type: "session_closed", + activeSessionId, + reason: "killed", + }), + ).toBe(false); + expect(written).toHaveLength(writesAfterAbort); + expect(client.privateFrameWriteTail).toBeUndefined(); + expect(socket.listenerCount("drain")).toBe(0); + expect(socket.listenerCount("close")).toBe(0); + }); + it("closes a permanently backpressured worker channel after an aborted snapshot", async () => { vi.useFakeTimers(); try { diff --git a/packages/coding-agent/test/suite/regressions/4602-snapshot-transfer-idempotency.test.ts b/packages/coding-agent/test/suite/regressions/4602-snapshot-transfer-idempotency.test.ts index 572b8a2e81..e14dded443 100644 --- a/packages/coding-agent/test/suite/regressions/4602-snapshot-transfer-idempotency.test.ts +++ b/packages/coding-agent/test/suite/regressions/4602-snapshot-transfer-idempotency.test.ts @@ -197,7 +197,7 @@ describe("ENG-4602 snapshot transfer containment", () => { clients: new Set(), eventGeneration: "generation-4602", lastEventSequence: 1, - runtime: { metadata: { kind: "top-level", createdAt: 1 } }, + runtime: { metadata: { kind: "top-level", createdAt: 1 }, session: { sessionId: "session-4602" } }, } as unknown as ActiveSessionState; const socket = new PassThrough(); const client = { @@ -258,7 +258,7 @@ describe("ENG-4602 snapshot transfer containment", () => { pendingAttaches: 0, eventGeneration: "generation-4602", lastEventSequence: 1, - runtime: { metadata: { kind: "top-level", createdAt: 1 } }, + runtime: { metadata: { kind: "top-level", createdAt: 1 }, session: { sessionId: "session-4602" } }, } as unknown as ActiveSessionState; const socket = new PassThrough(); const client = { From 3ab9110d11643c8bc1da9472843a1b1c17b45d2c Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Sat, 29 Aug 2026 20:25:17 -0600 Subject: [PATCH 07/12] feat(daemon): prove owned session cleanup --- .pylon/features.yaml | 15 ++ .pylon/upstream-review.md | 10 + .../11-snapshot-recovery-integrity.md | 1 + .../coding-agent/docs/agent-connection.md | 13 +- packages/coding-agent/src/index.ts | 2 + .../src/modes/daemon/daemon-protocol.ts | 26 ++- .../src/modes/daemon/daemon-supervisor.ts | 148 ++++++++++--- packages/coding-agent/src/modes/index.ts | 2 + .../coding-agent/test/daemon-client.test.ts | 60 ++++++ .../coding-agent/test/daemon-protocol.test.ts | 10 + .../test/daemon-supervisor-monitor.test.ts | 202 +++++++++++++++++- .../test/daemon-supervisor-process.test.ts | 92 +++++++- .../test/daemon-version-compatibility.test.ts | 19 ++ 13 files changed, 566 insertions(+), 34 deletions(-) diff --git a/.pylon/features.yaml b/.pylon/features.yaml index e72b2946e9..3d1a21fc36 100644 --- a/.pylon/features.yaml +++ b/.pylon/features.yaml @@ -188,3 +188,18 @@ decisions: revisit_when: - Prime gives each transfer immutable bytes and a unique identity while preserving the event cursor as ordering metadata. - Prime retries an isolated bad snapshot generation without closing the resident worker control channel. + + authoritative-owned-session-cleanup: + area: runtime-reliability + state: candidate + owner: shared + decision: hybridize + pylon_refs: + - https://github.com/pylon-code/prime-agent/issues/11 + - https://github.com/pylon-code/pylon/issues/114 + - https://github.com/rynfar/comet/issues/5 + fork_change: negotiated-authoritative-owned-session-cleanup-v1 + upstream_support: Prime v0.8.1 supports client-owned sessions and best-effort owner cleanup, but another host process cannot prove exact worker-generation and durable-descriptor absence after the owner connection disappears. + revisit_when: + - Prime upstream exposes an equivalent privacy-safe, capability-negotiated cleanup proof with durable descriptor-last retries. + - Comet and Pylon can remove the fork capability without weakening crash cleanup or stock-Prime fallback behavior. diff --git a/.pylon/upstream-review.md b/.pylon/upstream-review.md index e11872de3e..d1a0cd5f9d 100644 --- a/.pylon/upstream-review.md +++ b/.pylon/upstream-review.md @@ -56,3 +56,13 @@ This ledger records Prime upstream evidence and the decision taken for each over - The local `refs/review/current-candidate` at `a37efe92f88ae7234dcea1d065c0d3d76ac048cc` predates these corrections and includes the unaudited intervening range. Do not merge it. The upstream-review workflow must discard or regenerate it after this issue lands; this change does not mutate the protected review ref. - Validation: the touched snapshot, protocol, connection, supervisor, recursion, daemon-mode, session-list, and lazy-subagent suites passed 616 tests across 11 files. The real-process supervisor control-channel regression for root isolation, chunked snapshot streaming, and worker adoption passed. `npm run check` passed Biome over 940 files, TypeScript, installer rendering, and browser smoke; `git diff --check` passed. - Revisit when Prime gives every snapshot attempt immutable bytes and unique transfer identity, isolates bad generations from worker recovery, and preserves the same recovery and ownership guarantees without weakening Pylon's correlated lifecycle contract. + +## 2026-08-30 — snapshot, worker-authority, and authoritative owned-session cleanup candidate + +- Pylon base: `pylon@f728316dabbaa85aa561e6d6b08550ed337574be`, which includes merged correlated-lifecycle PR [#12](https://github.com/pylon-code/prime-agent/pull/12). Upstream evidence remains `PrimeIntellect-ai/prime-agent@a903d4b6768f484bd6d459b7b0aa7dee38e461e2`; latest compatibility release remains stock `v0.8.1`. +- The recovered [#11](https://github.com/pylon-code/prime-agent/issues/11) candidate keeps protocol version 7. Schema revision 26 adds the optional pre-begin `session_snapshot_failed.purpose` correlation used by immutable snapshot transfer. Schema revision 27 adds the supervisor-only `authoritative_owned_session_cleanup_v1` offer and read-only `get_owned_session_cleanup` query. Neither capability is inferred from package version, schema alone, method presence, or the legacy `client_owned_sessions` offer. +- `snapshot-recovery-integrity`: **hybridize**. In addition to immutable transfer identities and isolated generation failure, the candidate preserves terminal child errors through passive hydration, stabilizes snapshot cursors across asynchronous roster selection and the exact attach-registration cut, invalidates stale supervisor selection/materialization, closes a private channel if an announced frame is interrupted, and binds worker frames to the current authenticated channel and latest authoritative roster. +- `authoritative-owned-session-cleanup`: **hybridize**. `complete_owned_session` cannot succeed until the exact process generation is absent, ancillary journals are absent, the durable descriptor is removed and verified last, and only then the in-memory registration is removed. Tombstone, archive, and descriptor failures retain a stopping registration and arm a single exact-generation retry finalizer. A different client that already knows the exact host-private active-session ID may poll only `active | stopping | settled`; no PID, path, owner identity, descriptor, token, or diagnostics cross the public boundary. +- Compatibility classification: **backward-compatible, additive, capability-gated**. Private session workers continue to advertise the shared default server capabilities and do not falsely offer the supervisor-only cleanup query. Current clients reject the query locally against stock `v0.8.1` or schema 26 without sending an unknown command. The proof applies only to registrations created or durably adopted by a capable supervisor; it cannot retroactively certify a descriptorless historical orphan. +- Validation: 688 focused protocol/client/connection/session/snapshot/supervisor tests passed across 13 files; 40 correlated-lifecycle/queue/continuation tests passed; 12 real-process supervisor tests passed with 8 fixture-gated skips, including direct completion, real owner-socket loss with visible `stopping`, replacement-client `settled` proof, supervisor replacement, exact worker exit, and zero descriptors. Both stock/current `v0.8.1` adoption directions passed with local cleanup-query rejection on the stock supervisor. `npm run check` passed Biome over 941 files, TypeScript, installer rendering, and browser smoke. The root build passed; the live generated model catalog was restored afterward. `git diff --check` passed. +- Revisit when Prime upstream provides equivalent immutable snapshot, attach admission, worker-channel authority, and descriptor-last crash cleanup contracts without weakening the stock fallback or exposing host-private process/session details. diff --git a/packages/coding-agent/.changes/11-snapshot-recovery-integrity.md b/packages/coding-agent/.changes/11-snapshot-recovery-integrity.md index 668a905f51..a4e207529e 100644 --- a/packages/coding-agent/.changes/11-snapshot-recovery-integrity.md +++ b/packages/coding-agent/.changes/11-snapshot-recovery-integrity.md @@ -1,3 +1,4 @@ - Fixed daemon snapshot catch-up and recovery so transfer mismatches no longer interrupt healthy workers and recovered sessions preserve their state and owner ([#11](https://github.com/pylon-code/prime-agent/issues/11)). - Fixed recovered child rosters and activity so parent sessions show current delegated work after restart ([#11](https://github.com/pylon-code/prime-agent/issues/11)). - Fixed concurrent session opens so one owner and recovery attempt is reused instead of starting competing workers ([#11](https://github.com/pylon-code/prime-agent/issues/11)). +- Added capability-gated authoritative cleanup proof for client-owned daemon sessions, including crash-safe descriptor-last retries and privacy-safe polling from a replacement host ([#11](https://github.com/pylon-code/prime-agent/issues/11)). diff --git a/packages/coding-agent/docs/agent-connection.md b/packages/coding-agent/docs/agent-connection.md index 1c76f3c31b..9b9eee318d 100644 --- a/packages/coding-agent/docs/agent-connection.md +++ b/packages/coding-agent/docs/agent-connection.md @@ -102,12 +102,23 @@ The protocol does not promise that every historical event remains replayable. Du ## Command Lifecycle and Idempotency -The public daemon protocol is JSONL-framed and currently at protocol v4. Commands may be sent in versioned envelopes containing protocol metadata, client ID, and command ID. +The public daemon protocol is JSONL-framed and currently at protocol v7. Commands may be sent in versioned envelopes containing protocol metadata, client ID, and command ID. Mutating commands are recorded before dispatch. A repeated completed command returns its recorded result. A command known to have been received but lacking a durable result is reported as uncertain instead of being replayed blindly. Clients acknowledge durable results so old journal entries can be compacted. The `AgentConnection` method promise is a client convenience. It should not be treated as a general accepted/running/completed remote workflow API. +## Client-Owned Session Cleanup + +`DaemonAgentConnection.dispose()` performs best-effort cleanup for a client-owned session, but it deliberately does not turn cleanup failure into a disposal failure. An authoritative local host must use raw `DaemonClient` commands instead: + +1. Require the supervisor-only `authoritative_owned_session_cleanup_v1` server capability. +2. Send `complete_owned_session` from the owning client and check its response. +3. If the owner connection disappears or completion is uncertain, poll `get_owned_session_cleanup` from a replacement client with the exact host-private `activeSessionId`. +4. Accept only `settled` as proof that the exact worker generation, cleanup journals, durable descriptor, and in-memory registration are absent. Treat `active`, `stopping`, a missing capability, or transport uncertainty as not cleaned up. + +The query returns only `{ status: "active" | "stopping" | "settled" }`. It does not expose process IDs, worker or owner IDs, filesystem paths, tokens, descriptor contents, or diagnostics. Stock v0.8.1 does not offer this capability, so `DaemonClient` rejects the query locally without sending an unknown command. The proof applies to registrations created or durably adopted by a capable supervisor; it cannot retroactively certify a descriptorless orphan created before that supervisor observed it. + ## Session Replacement New, switch, fork, import, and tree-navigation operations may replace the runtime behind an active connection. The adapter owns rebinding and emits a replacement snapshot. The UI applies the new state and transcript; it does not rewire `AgentSession` listeners directly. diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 4988d065ec..d878e30ff2 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -326,6 +326,8 @@ export { type DaemonEventSequence, type DaemonModeOptions, type DaemonOutbound, + type DaemonOwnedSessionCleanupResult, + type DaemonOwnedSessionCleanupStatus, type DaemonProtocolInfo, type DaemonProtocolName, type DaemonProtocolVersion, diff --git a/packages/coding-agent/src/modes/daemon/daemon-protocol.ts b/packages/coding-agent/src/modes/daemon/daemon-protocol.ts index 2b5b200556..d5a420947a 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-protocol.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-protocol.ts @@ -77,8 +77,9 @@ export const DAEMON_COMMAND_ENVELOPE_MIN_PROTOCOL_VERSION = 7; // Revision 24 adds capability-gated correlated prompt lifecycle and event provenance. // Revision 25 negotiates immutable snapshot transfer identities with session workers. // Revision 26 correlates snapshot failures that occur before a begin frame can be emitted. -export const DAEMON_SCHEMA_REVISION = 26; -export const DAEMON_SCHEMA_ID = "protocol-7-schema-26-1ac2610350df"; +// Revision 27 adds a capability-gated authoritative owned-session cleanup query. +export const DAEMON_SCHEMA_REVISION = 27; +export const DAEMON_SCHEMA_ID = "protocol-7-schema-27-7c0c21a689b5"; export type DaemonProtocolName = typeof DAEMON_PROTOCOL_NAME; export type DaemonProtocolVersion = number; @@ -103,6 +104,10 @@ export type DaemonPromptAdmissionCancellationStatus = "cancelled" | "owned" | "u export interface DaemonPromptAdmissionCancellationResult { status: DaemonPromptAdmissionCancellationStatus; } +export type DaemonOwnedSessionCleanupStatus = "active" | "stopping" | "settled"; +export interface DaemonOwnedSessionCleanupResult { + status: DaemonOwnedSessionCleanupStatus; +} export type DaemonServerCapability = | DaemonClientCapability | "delete_rlm_subagent" @@ -125,7 +130,8 @@ export type DaemonServerCapability = | "rlm_quiescence_barrier" | "session_input_pause" | "owned_prompt_cancellation" - | "acp_mcp_servers"; + | "acp_mcp_servers" + | "authoritative_owned_session_cleanup_v1"; export type DaemonReplayStatus = "complete" | "partial" | "unavailable"; @@ -174,6 +180,12 @@ export const DAEMON_DEFAULT_SERVER_CAPABILITIES: readonly DaemonServerCapability "acp_mcp_servers", ]; +/** Supervisor-only offers. Private session workers must not advertise these commands. */ +export const DAEMON_SUPERVISOR_SERVER_CAPABILITIES: readonly DaemonServerCapability[] = [ + ...DAEMON_DEFAULT_SERVER_CAPABILITIES, + "authoritative_owned_session_cleanup_v1", +]; + export interface DaemonRuntimeIdentity { buildId: string; executablePath: string; @@ -430,6 +442,7 @@ export type DaemonCommand = DaemonLaunchEnv) | { id?: string; type: "detach"; activeSessionId?: string } | { id?: string; type: "complete_owned_session"; activeSessionId: string } + | { id?: string; type: "get_owned_session_cleanup"; activeSessionId: string } | { id?: string; type: "promote_owned_session"; activeSessionId: string } | { id?: string; type: "kill"; activeSessionId: string } | { id?: string; type: "rename"; activeSessionId: string; name: string } @@ -723,6 +736,11 @@ const CLIENT_OWNED_DAEMON_COMMAND = { minProtocol: 7, capability: "client_owned_sessions", } as const; +const AUTHORITATIVE_OWNED_SESSION_CLEANUP_COMMAND = { + minProtocol: 7, + minSchemaRevision: 27, + capability: "authoritative_owned_session_cleanup_v1", +} as const; const DELETE_RLM_SUBAGENT_COMMAND = { minProtocol: 7, capability: "delete_rlm_subagent", @@ -766,6 +784,7 @@ export const DAEMON_COMMAND_COMPATIBILITY = { reattach: LEGACY_DAEMON_COMMAND, detach: LEGACY_DAEMON_COMMAND, complete_owned_session: CLIENT_OWNED_DAEMON_COMMAND, + get_owned_session_cleanup: AUTHORITATIVE_OWNED_SESSION_CLEANUP_COMMAND, promote_owned_session: CLIENT_OWNED_DAEMON_COMMAND, kill: LEGACY_DAEMON_COMMAND, rename: LEGACY_DAEMON_COMMAND, @@ -1218,6 +1237,7 @@ const READ_ONLY_DAEMON_COMMANDS: ReadonlySet = new Set([ "list_agent_peers", "attach", "reattach", + "get_owned_session_cleanup", "agent_messages_status", "wait_for_idle", "get_session_header", diff --git a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts index f245317e20..7ba29cc91b 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts @@ -1,6 +1,6 @@ import { type ChildProcess, spawn } from "node:child_process"; import { createHash, randomBytes, randomUUID } from "node:crypto"; -import { chmodSync, mkdirSync, readdirSync, readFileSync, renameSync, rmSync, writeFileSync } from "node:fs"; +import { chmodSync, lstatSync, mkdirSync, readdirSync, readFileSync, renameSync, rmSync, writeFileSync } from "node:fs"; import { createServer, type Server, type Socket } from "node:net"; import { dirname, join, resolve } from "node:path"; import { Writable } from "node:stream"; @@ -71,16 +71,17 @@ import { DAEMON_COMMAND_COMPATIBILITY, DAEMON_COMMAND_ENVELOPE_MIN_PROTOCOL_VERSION, DAEMON_DEFAULT_CLIENT_CAPABILITIES, - DAEMON_DEFAULT_SERVER_CAPABILITIES, DAEMON_PROTOCOL_INFO, DAEMON_SCHEMA_ID, DAEMON_SCHEMA_REVISION, + DAEMON_SUPERVISOR_SERVER_CAPABILITIES, DAEMON_UPDATE_RESTART_FORMAT_VERSION, type DaemonAttachResult, type DaemonClientCapability, type DaemonClosingReason, type DaemonCommand, type DaemonOutbound, + type DaemonOwnedSessionCleanupResult, type DaemonResponse, type DaemonUpdateRestartManifest, daemonOutboundForCorrelatedPromptCapability, @@ -203,6 +204,7 @@ const DAEMON_COMMAND_TYPES: ReadonlySet = new Set([ "reattach", "detach", "complete_owned_session", + "get_owned_session_cleanup", "promote_owned_session", "kill", "rename", @@ -1082,8 +1084,14 @@ export class DaemonSupervisor { intentionalStop: durableDescriptor.stopRequestedAt !== undefined, stopRevision: 0, }; - this.persistWorker(worker); this.workers.set(durableDescriptor.workerId, worker); + try { + this.persistWorker(worker); + } catch (error) { + // Keep a valid durable registration in memory even when its + // recovery-state rewrite fails. Cleanup queries must fail closed. + this.log(`Could not refresh worker descriptor ${path}: ${String(error)}`); + } } catch (error) { this.log(`Ignoring invalid worker descriptor ${path}: ${String(error)}`); } @@ -1138,16 +1146,71 @@ export class DaemonSupervisor { renameSync(tempPath, worker.descriptorPath); } + private removeWorkerCleanupFile(path: string): void { + rmSync(path, { force: true }); + try { + lstatSync(path); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === "ENOENT") return; + throw error; + } + throw new Error("Worker cleanup file still exists after removal"); + } + private deleteWorkerDescriptor(worker: ResidentWorker): void { try { - rmSync(worker.descriptorPath, { force: true }); - rmSync(worker.descriptor.recoveryJournalPath, { force: true }); + // Ancillary journals are removed first. The descriptor is the durable + // registration and must remain present until every other cleanup step + // has succeeded and been verified. + this.removeWorkerCleanupFile(worker.descriptor.recoveryJournalPath); if (worker.descriptor.orphanProcessJournalPath) { - rmSync(worker.descriptor.orphanProcessJournalPath, { force: true }); + this.removeWorkerCleanupFile(worker.descriptor.orphanProcessJournalPath); } + this.removeWorkerCleanupFile(worker.descriptorPath); } catch (error) { this.log(`Failed to remove worker descriptor ${worker.descriptorPath}: ${String(error)}`); + throw new Error("Worker cleanup could not verify durable registration removal"); + } + } + + private hasDurableOwnedSessionRegistration(activeSessionId: string): boolean { + try { + for (const name of readdirSync(this.descriptorDir)) { + if (name === SUPERVISOR_CONFIG_FILE_NAME || !name.endsWith(".json")) continue; + const path = join(this.descriptorDir, name); + let descriptor: unknown; + try { + descriptor = JSON.parse(readFileSync(path, "utf8")); + } catch (error) { + this.log(`Could not inspect worker descriptor ${path}: ${String(error)}`); + return true; + } + if (!isDaemonWorkerDescriptor(descriptor, this.socketPath)) { + this.log(`Could not validate worker descriptor ${path} while checking cleanup`); + return true; + } + if (descriptor.rootActiveSessionId === activeSessionId) return true; + } + return false; + } catch (error) { + this.log(`Could not inspect durable worker registrations: ${String(error)}`); + return true; + } + } + + private getOwnedSessionCleanup(activeSessionId: string): DaemonOwnedSessionCleanupResult { + const worker = [...this.workers.values()].find( + (candidate) => candidate.descriptor.rootActiveSessionId === activeSessionId, + ); + if (worker) { + const stopping = + worker.descriptor.stopRequestedAt !== undefined || + worker.intentionalStop || + (worker.stopFinalization !== undefined && worker.stopFinalization !== null) || + (this.workerStopCounts?.get(worker) ?? 0) > 0; + return { status: stopping ? "stopping" : "active" }; } + return { status: this.hasDurableOwnedSessionRegistration(activeSessionId) ? "stopping" : "settled" }; } private handleConnection(socket: Socket): void { @@ -1184,7 +1247,7 @@ export class DaemonSupervisor { supervisorProcessStartId: this.ownership?.record.processStartId, supervisorSocketPath: this.ownership?.record.socketPath, clientId: client.id, - serverCapabilities: DAEMON_DEFAULT_SERVER_CAPABILITIES, + serverCapabilities: DAEMON_SUPERVISOR_SERVER_CAPABILITIES, }); } }, @@ -2035,6 +2098,8 @@ export class DaemonSupervisor { await this.releaseClientSessionInputPauses(client, command.activeSessionId, true); return success(command.id, "detach"); } + case "get_owned_session_cleanup": + return success(command.id, command.type, this.getOwnedSessionCleanup(command.activeSessionId)); case "complete_owned_session": { const match = await this.findWorkerForClient(client, command.activeSessionId); if (match.worker.descriptor.ownerClientId !== this.protocolClientId(client)) { @@ -2433,7 +2498,12 @@ export class DaemonSupervisor { } return await forward(); } - this.persistWorkerStopTombstone(match.worker, true); + try { + this.persistWorkerStopTombstone(match.worker, true); + } catch (error) { + this.scheduleWorkerStopFinalization(match.worker); + throw error; + } const releaseStopOwnership = this.acquireWorkerStopOwnership(match.worker); let response: DaemonResponse; try { @@ -2692,11 +2762,8 @@ export class DaemonSupervisor { if (identity !== "gone" && identity !== "replaced") { return false; } - worker.intentionalStop = true; await this.recoverUncertainWorkerOperations(worker, false); - this.invalidateWorkerSessionInputPauses(worker, "Session worker stopped while input was paused"); - this.workers.delete(worker.descriptor.workerId); - this.deleteWorkerDescriptor(worker); + await this.stopWorker(worker, true, true); return true; } // Fail fast before waiting on anything: only a confirmed-dead process is @@ -3091,7 +3158,12 @@ export class DaemonSupervisor { } catch (error) { worker.descriptor.lifecycle = "failed"; worker.descriptor.lastError = error instanceof Error ? error.message : String(error); - this.persistWorker(worker); + try { + this.persistWorker(worker); + } catch (persistError) { + this.reportCleanupFailure(`worker adoption tombstone ${worker.descriptor.workerId}`, persistError); + } + this.scheduleWorkerStopFinalization(worker); this.log(`Could not complete intentional stop for worker ${worker.descriptor.workerId}: ${String(error)}`); } return; @@ -5655,14 +5727,16 @@ export class DaemonSupervisor { activeSessionId === worker.descriptor.rootActiveSessionId && !this.shuttingDown ) { - worker.intentionalStop = true; // An exact stop owns its registration and descriptor cleanup until its - // tuple assertions complete. A synchronous root shutdown event can arrive - // before its request resolves, so leave both intact while it is active. + // tuple assertions complete. Without one, retain a durable tombstone and + // let the exact-generation finalizer prove process and descriptor absence. if ((this.workerStopCounts?.get(worker) ?? 0) === 0) { - this.invalidateWorkerSessionInputPauses(worker, "Session worker stopped while input was paused"); - this.workers.delete(worker.descriptor.workerId); - this.deleteWorkerDescriptor(worker); + try { + this.persistWorkerStopTombstone(worker); + } catch (error) { + this.reportCleanupFailure(`worker shutdown tombstone ${worker.descriptor.workerId}`, error); + } + this.scheduleWorkerStopFinalization(worker); } } } @@ -6108,6 +6182,16 @@ export class DaemonSupervisor { const releaseStopOwnership = this.acquireWorkerStopOwnership(worker); try { await this.stopWorkerUntracked(worker, removeDescriptor, force, archiveSession, recoveryCleanup, directChild); + } catch (error) { + if ( + removeDescriptor && + !directChild && + this.workers.get(worker.descriptor.workerId) === worker && + worker.descriptor.stopRequestedAt !== undefined + ) { + this.scheduleWorkerStopFinalization(worker); + } + throw error; } finally { releaseStopOwnership(); } @@ -6141,7 +6225,9 @@ export class DaemonSupervisor { return; } if ( + this.workers.get(worker.descriptor.workerId) !== worker || worker.descriptor.pid !== entryPid || + worker.descriptor.processStartId !== entryStartId || (removeDescriptor && worker.descriptor.stopRequestedAt === undefined) ) { throw new Error(`Session worker ${worker.descriptor.workerId} was relaunched during stop`); @@ -6258,10 +6344,10 @@ export class DaemonSupervisor { assertStopStillApplies(); } this.invalidateWorkerSessionInputPauses(worker, "Session worker stopped while input was paused"); - this.workers.delete(worker.descriptor.workerId); if (removeDescriptor) { this.deleteWorkerDescriptor(worker); } + this.workers.delete(worker.descriptor.workerId); if (!this.shuttingDown) { this.broadcastHeartbeatsChanged(); } @@ -6287,12 +6373,23 @@ export class DaemonSupervisor { // pid. The finalizer must never follow either successor. const pid = worker.descriptor.pid; const processStartId = worker.descriptor.processStartId; - const stopRevision = worker.stopRevision; const isStopGenerationCurrent = () => this.workers.get(worker.descriptor.workerId) === worker && - worker.stopRevision === stopRevision && worker.descriptor.stopRequestedAt !== undefined && - worker.descriptor.pid === pid; + worker.descriptor.pid === pid && + worker.descriptor.processStartId === processStartId; + // Do not signal or remove anything until stop intent is durable. A failed + // first tombstone write is retried here without losing the registration. + while (!this.shuttingDown && isStopGenerationCurrent()) { + try { + this.persistWorkerStopTombstone(worker, worker.descriptor.archiveOnStop === true); + break; + } catch (error) { + this.reportCleanupFailure(`worker stop tombstone ${worker.descriptor.workerId}`, error); + await unrefDelay(STOP_FINALIZATION_RETRY_MS); + } + } + if (this.shuttingDown || !isStopGenerationCurrent()) return; // A replaced pid counts as gone (never SIGKILL a recycled pid); an // unobservable identity counts as alive (never clean up a possibly-live // worker). kill(0) probes every poll; ps-backed checks are throttled. @@ -6339,7 +6436,7 @@ export class DaemonSupervisor { // unobservable identity skips this attempt but keeps escalation // armed so a wedged worker is still killed on a later pass. const observedNow = processStartId === undefined ? undefined : getProcessStartId(pid); - if (processStartId === undefined || observedNow === processStartId) { + if (processStartId !== undefined && observedNow === processStartId) { signalProcessGroupOrProcess(pid, "SIGKILL"); killed = true; } @@ -6353,7 +6450,8 @@ export class DaemonSupervisor { const isCleanupStillWanted = () => this.workers.get(worker.descriptor.workerId) === worker && worker.descriptor.stopRequestedAt !== undefined && - worker.descriptor.pid === pid; + worker.descriptor.pid === pid && + worker.descriptor.processStartId === processStartId; while (!this.shuttingDown && isCleanupStillWanted()) { try { await this.stopWorker(worker, true, true, worker.descriptor.archiveOnStop === true); diff --git a/packages/coding-agent/src/modes/index.ts b/packages/coding-agent/src/modes/index.ts index 89ece8f474..349e06e340 100644 --- a/packages/coding-agent/src/modes/index.ts +++ b/packages/coding-agent/src/modes/index.ts @@ -81,6 +81,8 @@ export type { DaemonEventMeta, DaemonEventSequence, DaemonOutbound, + DaemonOwnedSessionCleanupResult, + DaemonOwnedSessionCleanupStatus, DaemonProtocolInfo, DaemonProtocolName, DaemonProtocolVersion, diff --git a/packages/coding-agent/test/daemon-client.test.ts b/packages/coding-agent/test/daemon-client.test.ts index 2a28c47da3..8e291ac05c 100644 --- a/packages/coding-agent/test/daemon-client.test.ts +++ b/packages/coding-agent/test/daemon-client.test.ts @@ -215,6 +215,66 @@ describe("DaemonClient", () => { client.close(); }); + it("does not send authoritative cleanup queries to stock-compatible capability offers", async () => { + const client = new DaemonClient("/tmp/prime-agent.sock"); + const connect = client.connect(); + const socket = netMock.sockets[0]!; + socket.emit("connect"); + await connect; + emitHello(socket, DAEMON_PROTOCOL_VERSION, ["client_owned_sessions"], DAEMON_SCHEMA_REVISION); + + await expect(client.request({ type: "get_owned_session_cleanup", activeSessionId: "active-1" })).rejects.toThrow( + "does not support authoritative_owned_session_cleanup_v1", + ); + expect(socket.writes).toEqual([]); + client.close(); + }); + + it("schema-gates authoritative cleanup queries before writing", async () => { + const client = new DaemonClient("/tmp/prime-agent.sock"); + const connect = client.connect(); + const socket = netMock.sockets[0]!; + socket.emit("connect"); + await connect; + emitHello(socket, DAEMON_PROTOCOL_VERSION, ["authoritative_owned_session_cleanup_v1"], 26); + + await expect(client.request({ type: "get_owned_session_cleanup", activeSessionId: "active-1" })).rejects.toThrow( + "does not support authoritative_owned_session_cleanup_v1", + ); + expect(socket.writes).toEqual([]); + client.close(); + }); + + it("sends authoritative cleanup queries only to a capable current supervisor", async () => { + const client = new DaemonClient("/tmp/prime-agent.sock"); + const connect = client.connect(); + const socket = netMock.sockets[0]!; + socket.emit("connect"); + await connect; + emitHello(socket, DAEMON_PROTOCOL_VERSION, ["authoritative_owned_session_cleanup_v1"], DAEMON_SCHEMA_REVISION); + + const response = client.request({ type: "get_owned_session_cleanup", activeSessionId: "active-1" }); + expect(socket.writes).toHaveLength(1); + const envelope = JSON.parse(socket.writes[0]!.trim()) as { + id: string; + command: { type: string; activeSessionId: string }; + }; + expect(envelope.command).toMatchObject({ type: "get_owned_session_cleanup", activeSessionId: "active-1" }); + socket.emit( + "data", + `${JSON.stringify({ + id: envelope.id, + type: "response", + command: "get_owned_session_cleanup", + success: true, + data: { status: "settled" }, + })} +`, + ); + await expect(response).resolves.toMatchObject({ success: true, data: { status: "settled" } }); + client.close(); + }); + it("does not send subagent deletion to an old daemon without the capability", async () => { const client = new DaemonClient("/tmp/prime-agent.sock"); const connect = client.connect(); diff --git a/packages/coding-agent/test/daemon-protocol.test.ts b/packages/coding-agent/test/daemon-protocol.test.ts index 44b1a7863a..455a6a67c2 100644 --- a/packages/coding-agent/test/daemon-protocol.test.ts +++ b/packages/coding-agent/test/daemon-protocol.test.ts @@ -15,6 +15,7 @@ import { DAEMON_PROTOCOL_VERSION, DAEMON_SCHEMA_ID, DAEMON_SCHEMA_REVISION, + DAEMON_SUPERVISOR_SERVER_CAPABILITIES, DAEMON_SUPPORTED_CLIENT_CAPABILITIES, type DaemonCommand, type DaemonOutbound, @@ -113,6 +114,15 @@ describe("daemon protocol helpers", () => { minProtocol: 7, capability: "client_owned_sessions", }); + expect(DAEMON_COMMAND_COMPATIBILITY.get_owned_session_cleanup).toEqual({ + minProtocol: 7, + minSchemaRevision: 27, + capability: "authoritative_owned_session_cleanup_v1", + }); + expect(DAEMON_DEFAULT_CLIENT_CAPABILITIES).not.toContain("authoritative_owned_session_cleanup_v1"); + expect(DAEMON_DEFAULT_SERVER_CAPABILITIES).not.toContain("authoritative_owned_session_cleanup_v1"); + expect(DAEMON_SUPERVISOR_SERVER_CAPABILITIES).toContain("authoritative_owned_session_cleanup_v1"); + expect(isDaemonMutatingCommand({ type: "get_owned_session_cleanup" })).toBe(false); expect(DAEMON_OUTBOUND_COMPATIBILITY.heartbeats_changed).toEqual({ minProtocol: 7, capability: "heartbeat_catalog", diff --git a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts index 92c48e27ae..23d5d4f729 100644 --- a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts @@ -1827,20 +1827,21 @@ describe("daemon worker supervisor monitoring", () => { }; const workers = new Map([[worker.descriptor.workerId, worker]]); const recoverUncertainWorkerOperations = vi.fn(async () => {}); - const deleteWorkerDescriptor = vi.fn(); + const stopWorker = vi.fn(async () => { + workers.delete(worker.descriptor.workerId); + }); const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { workers, - sessionInputPauses: new Map(), processIdentity: vi.fn(() => "gone"), recoverUncertainWorkerOperations, - deleteWorkerDescriptor, + stopWorker, }) as { reclaimStaleWorkerRegistration(target: typeof worker): Promise; }; await expect(supervisor.reclaimStaleWorkerRegistration(worker)).resolves.toBe(true); expect(recoverUncertainWorkerOperations).toHaveBeenCalledWith(worker, false); - expect(deleteWorkerDescriptor).toHaveBeenCalledWith(worker); + expect(stopWorker).toHaveBeenCalledWith(worker, true, true); expect(workers.has(worker.descriptor.workerId)).toBe(false); }); @@ -2328,6 +2329,196 @@ describe("daemon worker supervisor monitoring", () => { } }); + it("reports exact privacy-safe owned cleanup states", () => { + const descriptorDir = mkdtempSync(join(tmpdir(), "prime-supervisor-owned-cleanup-status-")); + try { + const worker = { + descriptor: { + workerId: "owned-cleanup-worker", + rootActiveSessionId: "active-owned-cleanup", + }, + intentionalStop: false, + stopFinalization: undefined, + }; + const workers = new Map([[worker.descriptor.workerId, worker]]); + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + descriptorDir, + socketPath: "/tmp/owned-cleanup-status.sock", + workers, + workerStopCounts: new Map(), + log: vi.fn(), + }) as { + getOwnedSessionCleanup(activeSessionId: string): { status: "active" | "stopping" | "settled" }; + }; + + expect(supervisor.getOwnedSessionCleanup(worker.descriptor.rootActiveSessionId)).toEqual({ status: "active" }); + expect(supervisor.getOwnedSessionCleanup("owned-cleanup")).toEqual({ status: "settled" }); + worker.intentionalStop = true; + expect(supervisor.getOwnedSessionCleanup(worker.descriptor.rootActiveSessionId)).toEqual({ + status: "stopping", + }); + workers.clear(); + writeFileSync(join(descriptorDir, "uncertain.json"), "{not-json"); + expect(supervisor.getOwnedSessionCleanup(worker.descriptor.rootActiveSessionId)).toEqual({ + status: "stopping", + }); + rmSync(join(descriptorDir, "uncertain.json"), { force: true }); + const settled = supervisor.getOwnedSessionCleanup(worker.descriptor.rootActiveSessionId); + expect(settled).toEqual({ status: "settled" }); + expect(Object.keys(settled)).toEqual(["status"]); + } finally { + rmSync(descriptorDir, { recursive: true, force: true }); + } + }); + + it("retains a tombstoned registration when descriptor removal fails and settles only after retry", async () => { + const descriptorDir = mkdtempSync(join(tmpdir(), "prime-supervisor-owned-cleanup-unlink-")); + const worker = { + descriptor: { + workerId: "owned-cleanup-unlink-worker", + pid: 111_116, + processStartId: "proc:owned-cleanup", + rootActiveSessionId: "active-owned-cleanup-unlink", + stopRequestedAt: new Date().toISOString(), + }, + client: undefined, + summaries: new Map(), + snapshotCache: new Map(), + transcriptCaches: new Map(), + snapshotGenerations: new Map(), + snapshotLoads: new Map(), + intentionalStop: true, + stopRevision: 0, + }; + const workers = new Map([[worker.descriptor.workerId, worker]]); + const deleteWorkerDescriptor = vi.fn<() => void>().mockImplementationOnce(() => { + throw new Error("descriptor unlink failed"); + }); + const scheduleWorkerStopFinalization = vi.fn(); + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + descriptorDir, + socketPath: "/tmp/owned-cleanup-unlink.sock", + workers, + shuttingDown: false, + persistWorkerStopTombstone: vi.fn(), + deleteWorkerDescriptor, + scheduleWorkerStopFinalization, + invalidateWorkerSessionInputPauses: vi.fn(), + broadcastHeartbeatsChanged: vi.fn(), + log: vi.fn(), + }) as unknown as { + stopWorker(target: typeof worker, removeDescriptor: boolean, force?: boolean): Promise; + getOwnedSessionCleanup(activeSessionId: string): { status: "active" | "stopping" | "settled" }; + }; + const childProcessModule = await import("../src/utils/child-process.js"); + const existsSpy = vi.spyOn(childProcessModule, "processIdExists").mockReturnValue(false); + try { + await expect(supervisor.stopWorker(worker, true, true)).rejects.toThrow("descriptor unlink failed"); + expect(workers.get(worker.descriptor.workerId)).toBe(worker); + expect(scheduleWorkerStopFinalization).toHaveBeenCalledOnce(); + expect(supervisor.getOwnedSessionCleanup(worker.descriptor.rootActiveSessionId)).toEqual({ + status: "stopping", + }); + + await supervisor.stopWorker(worker, true, true); + expect(workers.has(worker.descriptor.workerId)).toBe(false); + expect(deleteWorkerDescriptor).toHaveBeenCalledTimes(2); + expect(supervisor.getOwnedSessionCleanup(worker.descriptor.rootActiveSessionId)).toEqual({ + status: "settled", + }); + } finally { + existsSpy.mockRestore(); + rmSync(descriptorDir, { recursive: true, force: true }); + } + }); + + it("arms authoritative retry when the first stop tombstone write fails", async () => { + const worker = { + descriptor: { + workerId: "owned-cleanup-tombstone-worker", + pid: 111_117, + processStartId: "proc:owned-cleanup-tombstone", + rootActiveSessionId: "active-owned-cleanup-tombstone", + stopRequestedAt: undefined as string | undefined, + }, + client: undefined, + summaries: new Map(), + snapshotCache: new Map(), + transcriptCaches: new Map(), + snapshotGenerations: new Map(), + snapshotLoads: new Map(), + intentionalStop: false, + stopRevision: 0, + }; + const workers = new Map([[worker.descriptor.workerId, worker]]); + const scheduleWorkerStopFinalization = vi.fn(); + const persistWorkerStopTombstone = vi.fn(() => { + worker.intentionalStop = true; + worker.descriptor.stopRequestedAt = new Date().toISOString(); + throw new Error("tombstone write failed"); + }); + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers, + shuttingDown: false, + persistWorkerStopTombstone, + scheduleWorkerStopFinalization, + }) as unknown as { + stopWorker(target: typeof worker, removeDescriptor: boolean): Promise; + }; + + await expect(supervisor.stopWorker(worker, true)).rejects.toThrow("tombstone write failed"); + expect(workers.get(worker.descriptor.workerId)).toBe(worker); + expect(worker.descriptor.stopRequestedAt).toBeDefined(); + expect(scheduleWorkerStopFinalization).toHaveBeenCalledOnce(); + }); + + it("retries durable stop intent before finalizing a dead owned worker", async () => { + vi.useFakeTimers(); + const worker = { + descriptor: { + workerId: "owned-cleanup-durable-retry-worker", + pid: 111_118, + processStartId: "proc:owned-cleanup-durable-retry", + rootActiveSessionId: "active-owned-cleanup-durable-retry", + stopRequestedAt: new Date().toISOString(), + }, + intentionalStop: true, + stopRevision: 0, + stopFinalization: undefined as Promise | undefined, + }; + const workers = new Map([[worker.descriptor.workerId, worker]]); + const persistWorker = vi.fn<() => void>().mockImplementationOnce(() => { + throw new Error("durable tombstone unavailable"); + }); + const stopWorker = vi.fn(async () => { + workers.delete(worker.descriptor.workerId); + }); + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers, + shuttingDown: false, + persistWorker, + stopWorker, + log: vi.fn(), + reportCleanupFailure: vi.fn(), + }) as { + scheduleWorkerStopFinalization(target: typeof worker): void; + }; + const childProcessModule = await import("../src/utils/child-process.js"); + const existsSpy = vi.spyOn(childProcessModule, "processIdExists").mockReturnValue(false); + try { + supervisor.scheduleWorkerStopFinalization(worker); + expect(stopWorker).not.toHaveBeenCalled(); + await vi.advanceTimersByTimeAsync(20_000); + await worker.stopFinalization; + + expect(persistWorker).toHaveBeenCalledTimes(2); + expect(stopWorker).toHaveBeenCalledWith(worker, true, true, false); + expect(persistWorker.mock.invocationCallOrder.at(-1)).toBeLessThan(stopWorker.mock.invocationCallOrder[0]!); + } finally { + existsSpy.mockRestore(); + } + }); + it("aborts stale stop cleanup when the worker was relaunched during an await", async () => { const worker = { descriptor: { @@ -2816,6 +3007,7 @@ describe("daemon worker supervisor monitoring", () => { const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { workers, stopWorker, + persistWorkerStopTombstone: vi.fn(), log: vi.fn(), reportCleanupFailure: vi.fn(), }) as { @@ -2852,6 +3044,7 @@ describe("daemon worker supervisor monitoring", () => { const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { workers, stopWorker, + persistWorkerStopTombstone: vi.fn(), log: vi.fn(), reportCleanupFailure: vi.fn(), }) as { @@ -2897,6 +3090,7 @@ describe("daemon worker supervisor monitoring", () => { const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { workers, stopWorker, + persistWorkerStopTombstone: vi.fn(), log: vi.fn(), reportCleanupFailure: vi.fn(), }) as { diff --git a/packages/coding-agent/test/daemon-supervisor-process.test.ts b/packages/coding-agent/test/daemon-supervisor-process.test.ts index 790db2091b..dba48b9036 100644 --- a/packages/coding-agent/test/daemon-supervisor-process.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-process.test.ts @@ -690,7 +690,11 @@ describe("daemon supervisor resident workers", () => { success: false, error: `Unknown active session: ${summary.activeSessionId}`, }); - otherClient.close(); + const activeCleanup = await otherClient.request({ + type: "get_owned_session_cleanup", + activeSessionId: summary.activeSessionId, + }); + expect(activeCleanup).toMatchObject({ success: true, data: { status: "active" } }); const connection = await DaemonAgentConnection.attach(client, summary.activeSessionId, { ownedSession: true, supportsExtensionUi: false, @@ -716,7 +720,19 @@ describe("daemon supervisor resident workers", () => { expect(descriptor.createCommand).not.toHaveProperty("launchEnv"); expect(descriptor.createCommand).not.toHaveProperty("lifecycle"); + const completed = await client.request({ + type: "complete_owned_session", + activeSessionId: summary.activeSessionId, + }); + expect(completed.success).toBe(true); + const settledCleanup = await otherClient.request({ + type: "get_owned_session_cleanup", + activeSessionId: summary.activeSessionId, + }); + expect(settledCleanup).toEqual(expect.objectContaining({ success: true, data: { status: "settled" } })); + expect(settledCleanup.success ? Object.keys(settledCleanup.data as object) : []).toEqual(["status"]); await connection.dispose(); + otherClient.close(); await waitForProcessGone(summary.workerPid); workerPids.delete(summary.workerPid); await waitForCondition( @@ -730,6 +746,75 @@ describe("daemon supervisor resident workers", () => { await waitForSocketGone(socketPath); }, 60_000); + it("lets a different client prove cleanup after the owner socket disappears", async () => { + const root = tempDir(); + const agentDir = join(root, "agent"); + const projectDir = join(root, "project"); + const socketPath = join(tmpdir(), `prime-supervisor-owned-crash-${process.pid}-${randomUUID().slice(0, 8)}.sock`); + mkdirSync(projectDir, { recursive: true }); + + const supervisor = spawnSupervisor(agentDir, socketPath, projectDir); + const owner = await connectEventually(socketPath, supervisor); + const created = await owner.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 observer = await connectEventually(socketPath); + try { + await expect( + observer.request({ type: "get_owned_session_cleanup", activeSessionId: summary.activeSessionId }), + ).resolves.toMatchObject({ success: true, data: { status: "active" } }); + process.kill(summary.workerPid, "SIGSTOP"); + owner.close(); + + let sawStopping = false; + let settled = false; + const deadline = Date.now() + 50_000; + while (Date.now() < deadline) { + const response = await observer.request({ + type: "get_owned_session_cleanup", + activeSessionId: summary.activeSessionId, + }); + if (!response.success) throw new Error(response.error); + const status = (response.data as { status?: string } | undefined)?.status; + if (status === "stopping") sawStopping = true; + if (status === "settled") { + settled = true; + break; + } + await new Promise((resolveDelay) => setTimeout(resolveDelay, 50)); + } + expect(sawStopping).toBe(true); + expect(settled).toBe(true); + await waitForProcessGone(summary.workerPid); + workerPids.delete(summary.workerPid); + expect(countWorkerDescriptors(agentDir)).toBe(0); + await observer.request({ type: "shutdown" }); + await waitForSocketGone(socketPath); + } finally { + owner.close(); + observer.close(); + if (workerPids.has(summary.workerPid)) { + try { + process.kill(summary.workerPid, "SIGCONT"); + process.kill(summary.workerPid, "SIGKILL"); + } catch { + // The authoritative cleanup path already removed it. + } + workerPids.delete(summary.workerPid); + } + } + }, 75_000); + it("releases an adopted client-owned worker when disposal races supervisor replacement", async () => { const root = tempDir(); const agentDir = join(root, "agent"); @@ -773,6 +858,11 @@ describe("daemon supervisor resident workers", () => { "Adopted client-owned worker descriptor was not removed", ); const replacementClient = await connectEventually(socketPath); + const replacementCleanup = await replacementClient.request({ + type: "get_owned_session_cleanup", + activeSessionId: summary.activeSessionId, + }); + expect(replacementCleanup).toMatchObject({ success: true, data: { status: "settled" } }); await replacementClient.request({ type: "shutdown" }); replacementClient.close(); await waitForSocketGone(socketPath); diff --git a/packages/coding-agent/test/daemon-version-compatibility.test.ts b/packages/coding-agent/test/daemon-version-compatibility.test.ts index aad93dc3a4..51940da6fb 100644 --- a/packages/coding-agent/test/daemon-version-compatibility.test.ts +++ b/packages/coding-agent/test/daemon-version-compatibility.test.ts @@ -191,11 +191,30 @@ processTests("stock v0.8.1 daemon entrypoint compatibility", () => { if (!createdSummary.activeSessionId || !createdSummary.workerPid) throw new Error("Compatibility worker was incomplete"); workerPids.add(createdSummary.workerPid); + if (firstCli === historicalCli) { + expect(firstClient.supportsServerCapability("authoritative_owned_session_cleanup_v1")).toBe(false); + await expect( + firstClient.request({ + type: "get_owned_session_cleanup", + activeSessionId: createdSummary.activeSessionId, + }), + ).rejects.toThrow("authoritative_owned_session_cleanup_v1"); + await expect(firstClient.request({ type: "list" })).resolves.toMatchObject({ success: true }); + } firstClient.close(); await stopSupervisor(firstSupervisor); const replacement = launch(replacementCli, agentDir, socketPath, projectDir); const client = await connectEventually(socketPath, replacement); + if (replacementCli === historicalCli) { + expect(client.supportsServerCapability("authoritative_owned_session_cleanup_v1")).toBe(false); + await expect( + client.request({ + type: "get_owned_session_cleanup", + activeSessionId: createdSummary.activeSessionId, + }), + ).rejects.toThrow("authoritative_owned_session_cleanup_v1"); + } const listed = await client.request({ type: "list" }); if (!listed.success) throw new Error(listed.error); const adopted = summaries(listed.data).find( From f163bddc8a76d5ecd02b848cdbcc80f7b9fa0753 Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Sat, 29 Aug 2026 21:04:45 -0600 Subject: [PATCH 08/12] fix(daemon): fence owned cleanup races --- .pylon/upstream-review.md | 2 +- .../coding-agent/docs/agent-connection.md | 2 + .../src/modes/daemon/daemon-supervisor.ts | 275 ++++++++++--- .../test/daemon-supervisor-monitor.test.ts | 372 +++++++++++++++--- .../test/daemon-supervisor-process.test.ts | 89 +++++ 5 files changed, 632 insertions(+), 108 deletions(-) diff --git a/.pylon/upstream-review.md b/.pylon/upstream-review.md index d1a0cd5f9d..b16277a334 100644 --- a/.pylon/upstream-review.md +++ b/.pylon/upstream-review.md @@ -64,5 +64,5 @@ This ledger records Prime upstream evidence and the decision taken for each over - `snapshot-recovery-integrity`: **hybridize**. In addition to immutable transfer identities and isolated generation failure, the candidate preserves terminal child errors through passive hydration, stabilizes snapshot cursors across asynchronous roster selection and the exact attach-registration cut, invalidates stale supervisor selection/materialization, closes a private channel if an announced frame is interrupted, and binds worker frames to the current authenticated channel and latest authoritative roster. - `authoritative-owned-session-cleanup`: **hybridize**. `complete_owned_session` cannot succeed until the exact process generation is absent, ancillary journals are absent, the durable descriptor is removed and verified last, and only then the in-memory registration is removed. Tombstone, archive, and descriptor failures retain a stopping registration and arm a single exact-generation retry finalizer. A different client that already knows the exact host-private active-session ID may poll only `active | stopping | settled`; no PID, path, owner identity, descriptor, token, or diagnostics cross the public boundary. - Compatibility classification: **backward-compatible, additive, capability-gated**. Private session workers continue to advertise the shared default server capabilities and do not falsely offer the supervisor-only cleanup query. Current clients reject the query locally against stock `v0.8.1` or schema 26 without sending an unknown command. The proof applies only to registrations created or durably adopted by a capable supervisor; it cannot retroactively certify a descriptorless historical orphan. -- Validation: 688 focused protocol/client/connection/session/snapshot/supervisor tests passed across 13 files; 40 correlated-lifecycle/queue/continuation tests passed; 12 real-process supervisor tests passed with 8 fixture-gated skips, including direct completion, real owner-socket loss with visible `stopping`, replacement-client `settled` proof, supervisor replacement, exact worker exit, and zero descriptors. Both stock/current `v0.8.1` adoption directions passed with local cleanup-query rejection on the stock supervisor. `npm run check` passed Biome over 941 files, TypeScript, installer rendering, and browser smoke. The root build passed; the live generated model catalog was restored afterward. `git diff --check` passed. +- Validation: 693 focused protocol/client/connection/session/snapshot/supervisor tests passed across 13 files; 40 correlated-lifecycle/queue/continuation tests passed; 13 real-process supervisor tests passed with 8 fixture-gated skips, including direct completion, owner loss before and after worker registration, visible `stopping`, replacement-client `settled` proof, supervisor replacement, exact worker exit, and zero descriptors. Final exact-head review also forced recovery-join, shutdown-drain, concurrent-completion, and generation-keyed-finalizer repairs before renewed exact-head review. Both stock/current `v0.8.1` adoption directions passed with local cleanup-query rejection on the stock supervisor. `npm run check` passed Biome over 941 files, TypeScript, installer rendering, and browser smoke. The root build passed; the live generated model catalog was restored afterward. `git diff --check` passed. - Revisit when Prime upstream provides equivalent immutable snapshot, attach admission, worker-channel authority, and descriptor-last crash cleanup contracts without weakening the stock fallback or exposing host-private process/session details. diff --git a/packages/coding-agent/docs/agent-connection.md b/packages/coding-agent/docs/agent-connection.md index 9b9eee318d..c61913be63 100644 --- a/packages/coding-agent/docs/agent-connection.md +++ b/packages/coding-agent/docs/agent-connection.md @@ -119,6 +119,8 @@ The `AgentConnection` method promise is a client convenience. It should not be t The query returns only `{ status: "active" | "stopping" | "settled" }`. It does not expose process IDs, worker or owner IDs, filesystem paths, tokens, descriptor contents, or diagnostics. Stock v0.8.1 does not offer this capability, so `DaemonClient` rejects the query locally without sending an unknown command. The proof applies to registrations created or durably adopted by a capable supervisor; it cannot retroactively certify a descriptorless orphan created before that supervisor observed it. +Cleanup closes and joins any admitted recovery before it can publish a successor. Concurrent completion requests for the same resident generation share one stop operation. Finalizers are keyed by process identity and stop revision, and daemon shutdown keeps registry ownership and retry services alive until every requested worker cleanup settles. A client-owned create rechecks owner liveness after registration, so a disconnect while daemon readiness or process launch is still pending cannot miss its cleanup timer. + ## Session Replacement New, switch, fork, import, and tree-navigation operations may replace the runtime behind an active connection. The adapter owns rebinding and emits a replacement snapshot. The UI applies the new state and transcript; it does not rewire `AgentSession` listeners directly. diff --git a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts index 7ba29cc91b..df9835e68c 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts @@ -326,7 +326,13 @@ interface ResidentWorker { stopRevision: number; launchEnv?: Record; transientCreateCommand?: DaemonCreateCommand; - stopFinalization?: Promise; + stopOperation?: Promise; + stopFinalization?: { + pid: number; + processStartId?: string; + stopRevision: number; + promise: Promise; + }; ownerCleanupTimer?: ReturnType; promotedOwnerClientId?: string; updateRestartPrepareClient?: DaemonWorkerClient; @@ -1596,7 +1602,9 @@ export class DaemonSupervisor { this.protocolClientIds.set(client, envelopeClientId); client.id = envelopeClientId; } - this.cancelOwnedWorkerCleanup(client.id); + if (this.clients?.has(client)) { + this.cancelOwnedWorkerCleanup(client.id); + } if (!DAEMON_COMMAND_TYPES.has(command.type)) { correlatedOrder?.release(); this.write(client, failure(command.id, command.type, `Unknown daemon command: ${command.type}`)); @@ -1860,6 +1868,10 @@ export class DaemonSupervisor { return this.handleSavedSessionList(client, command); case "create": { const worker = await this.createOrReuseWorker(this.protocolClientId(client), command); + // The owner may disconnect while creation is waiting for daemon readiness, + // process launch, or session materialization. Re-evaluate ownership only + // after the worker is registered so the disconnect edge cannot be missed. + this.scheduleOwnedWorkerCleanup(worker); const requestedSummary = command.sessionPath ? this.findSummaryInWorker(worker, command.sessionPath) : undefined; @@ -2128,6 +2140,7 @@ export class DaemonSupervisor { if ((this.workerStopCounts?.get(worker) ?? 0) > 0) { throw new Error("Session worker is stopping; retry after it finishes"); } + worker.stopRevision++; worker.intentionalStop = false; worker.descriptor.stopRequestedAt = undefined; worker.descriptor.archiveOnStop = undefined; @@ -2779,7 +2792,7 @@ export class DaemonSupervisor { this.scheduleWorkerStopFinalization(worker); const finalization = worker.stopFinalization; if (finalization) { - await Promise.race([finalization.catch(() => undefined), unrefDelay(STALE_RECLAIM_WAIT_MS)]); + await Promise.race([finalization.promise.catch(() => undefined), unrefDelay(STALE_RECLAIM_WAIT_MS)]); } if (this.workers.get(worker.descriptor.workerId) === worker) { // The process is confirmed dead, so the registration must never be @@ -2824,10 +2837,13 @@ export class DaemonSupervisor { promptLifecycleRecovery?: ReadonlyMap, ): 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 assertLaunchCurrent = () => { + if (existing && recoveryStopRevision !== undefined) { + this.assertWorkerRecoveryCurrent(existing, recoveryStopRevision); + } + }; + assertLaunchCurrent(); const launchEnv = command.launchEnv ?? existing?.launchEnv; const createCommand: DaemonCreateCommand = { ...withoutSupervisorCreateFields(command), @@ -2859,6 +2875,7 @@ export class DaemonSupervisor { }); delete workerEnvironment.RLM_DEPTH; await this.assertRecoveryAllowed(); + assertLaunchCurrent(); const child: ChildProcess = spawn(launch.command, launch.args, { cwd: createCommand.config?.cwd ?? process.cwd(), detached: true, @@ -2895,6 +2912,7 @@ export class DaemonSupervisor { childPid = child.pid; childProcessStartId = getProcessStartId(childPid); await this.assertRecoveryAllowed(); + assertLaunchCurrent(); const descriptor: DaemonWorkerDescriptor = { version: 2, @@ -2932,6 +2950,7 @@ export class DaemonSupervisor { transientCreateCommand: ownerClientId ? createCommand : undefined, }; await this.assertRecoveryAllowed(); + assertLaunchCurrent(); worker.descriptor = descriptor; worker.launchEnv = launchEnv; worker.transientCreateCommand = descriptor.ownerClientId ? createCommand : undefined; @@ -2962,7 +2981,9 @@ export class DaemonSupervisor { try { try { + assertLaunchCurrent(); await commitWorkerStartupGate(startupGate); + assertLaunchCurrent(); } catch (error) { startupGate.destroy(); await childClosed; @@ -2970,8 +2991,10 @@ export class DaemonSupervisor { } finally { child.unref(); } - const client = await this.connectWorker(worker, WORKER_CONNECT_TIMEOUT_MS); + const client = await this.connectWorker(worker, WORKER_CONNECT_TIMEOUT_MS, assertLaunchCurrent); + assertLaunchCurrent(); const response = await client.request(withoutCommandId(createCommand), WORKER_REQUEST_TIMEOUT_MS); + assertLaunchCurrent(); if (!response.success) { throw deserializeDaemonError(response); } @@ -2986,14 +3009,15 @@ export class DaemonSupervisor { worker.descriptor.rootSessionId = summary.sessionId; worker.descriptor.sessionFile = summary.sessionFile; await this.subscribeWorker(worker, rootActiveSessionId); - await this.refreshWorkerSummaries(worker, true); + assertLaunchCurrent(); + await this.refreshWorkerSummaries(worker, true, recoveryStopRevision); + assertLaunchCurrent(); if (promptLifecycleRecovery !== undefined) { await this.restoreWorkerPromptLifecycles(worker, promptLifecycleRecovery); - } - if (existing && (this.isWorkerRecoveryCancelled(worker) || worker.stopRevision !== recoveryStopRevision)) { - throw new Error(`Session worker ${workerId} recovery was cancelled`); + assertLaunchCurrent(); } await this.assertRecoveryAllowed(); + assertLaunchCurrent(); worker.descriptor.lifecycle = "ready"; worker.descriptor.consecutiveFailures = 0; worker.descriptor.lastError = undefined; @@ -3066,11 +3090,17 @@ export class DaemonSupervisor { } } - private async connectWorker(worker: ResidentWorker, timeoutMs: number): Promise { + private async connectWorker( + worker: ResidentWorker, + timeoutMs: number, + assertCurrent?: () => void, + ): Promise { const deadline = Date.now() + timeoutMs; let lastError: unknown; while (Date.now() < deadline) { + assertCurrent?.(); await this.assertRecoveryAllowed(); + assertCurrent?.(); const client = new DaemonWorkerClient(worker.descriptor.socketPath); try { await client.connect(Math.min(500, Math.max(50, deadline - Date.now()))); @@ -3081,6 +3111,7 @@ export class DaemonSupervisor { 1000, ); await this.assertRecoveryAllowed(); + assertCurrent?.(); client.onFrame((frame) => this.handleWorkerFrame(worker, frame, client)); client.onClose((error) => void this.handleWorkerClose(worker, client, error)); worker.client?.close(); @@ -3131,7 +3162,23 @@ export class DaemonSupervisor { private async adoptOrRecoverWorker(worker: ResidentWorker): Promise { await this.assertRecoveryAllowed(); + const adoptionRevision = worker.stopRevision; if (worker.descriptor.stopRequestedAt) { + const stopPid = worker.descriptor.pid; + let stopProcessStartId = worker.descriptor.processStartId; + const assertStopAdoptionCurrent = () => { + if ( + this.workers.get(worker.descriptor.workerId) !== worker || + worker.stopRevision !== adoptionRevision || + worker.descriptor.stopRequestedAt === undefined || + worker.descriptor.pid !== stopPid || + worker.descriptor.processStartId !== stopProcessStartId + ) { + throw new SupervisorRecoveryCancelledError( + `Session worker ${worker.descriptor.workerId} stop adoption was cancelled`, + ); + } + }; try { // A descriptor persisted before identity tracking has no // processStartId, so stopWorker could neither signal the live @@ -3142,9 +3189,11 @@ export class DaemonSupervisor { if (worker.descriptor.processStartId === undefined && isProcessAlive(worker.descriptor.pid)) { const observedProcessStartId = getProcessStartId(worker.descriptor.pid); try { - await this.connectWorker(worker, 2000); + await this.connectWorker(worker, 2000, assertStopAdoptionCurrent); + assertStopAdoptionCurrent(); if (observedProcessStartId) { worker.descriptor.processStartId = observedProcessStartId; + stopProcessStartId = observedProcessStartId; this.persistWorker(worker); } } catch { @@ -3153,9 +3202,11 @@ export class DaemonSupervisor { // waiting rather than signalling a possibly-recycled pid. } } + assertStopAdoptionCurrent(); await this.stopWorker(worker, true, true, worker.descriptor.archiveOnStop === true); this.log(`Completed intentional stop for worker ${worker.descriptor.workerId} during supervisor adoption`); } catch (error) { + if (this.workers.get(worker.descriptor.workerId) !== worker) return; worker.descriptor.lifecycle = "failed"; worker.descriptor.lastError = error instanceof Error ? error.message : String(error); try { @@ -3168,24 +3219,32 @@ export class DaemonSupervisor { } return; } + const assertAdoptionCurrent = () => this.assertWorkerRecoveryCurrent(worker, adoptionRevision); try { + assertAdoptionCurrent(); if (!isProcessAlive(worker.descriptor.pid)) { throw new Error("Session worker process is no longer running"); } const observedProcessStartId = getProcessStartId(worker.descriptor.pid); - await this.connectWorker(worker, 2000); + await this.connectWorker(worker, 2000, assertAdoptionCurrent); + assertAdoptionCurrent(); await this.subscribeWorker(worker, worker.descriptor.rootActiveSessionId); - await this.refreshWorkerSummaries(worker, true); + assertAdoptionCurrent(); + await this.refreshWorkerSummaries(worker, true, adoptionRevision); + assertAdoptionCurrent(); if (worker.descriptor.processStartId === undefined && observedProcessStartId) { worker.descriptor.processStartId = observedProcessStartId; } await this.assertRecoveryAllowed(); + assertAdoptionCurrent(); worker.descriptor.lifecycle = "ready"; worker.descriptor.consecutiveFailures = 0; this.persistWorker(worker); this.broadcastHeartbeatsChanged(); } catch (error) { if (isSupervisorRecoveryCancelled(error)) { + worker.client?.close(); + worker.client = undefined; return; } this.log(`Could not adopt worker ${worker.descriptor.workerId}: ${String(error)}`); @@ -3667,23 +3726,25 @@ export class DaemonSupervisor { if (this.isWorkerRecoveryCancelled(worker)) { return; } + if (worker.recovery) { + return worker.recovery; + } + const recoveryRevision = worker.stopRevision; + const assertCurrent = () => this.assertWorkerRecoveryCurrent(worker, recoveryRevision); if (worker.descriptor.ownerClientId && !worker.launchEnv && !isProcessAlive(worker.descriptor.pid)) { + assertCurrent(); worker.descriptor.lifecycle = "failed"; worker.descriptor.lastError = "Waiting for the owning client to reconnect"; this.persistWorker(worker); return; } - if (worker.recovery) { - return worker.recovery; - } worker.recovery = (async () => { for (const [retryIndex, retryDelay] of WORKER_RETRY_DELAYS_MS.entries()) { await delay(retryDelay); - if (this.isWorkerRecoveryCancelled(worker)) { - return; - } try { + assertCurrent(); await this.assertRecoveryAllowed(); + assertCurrent(); const processAlive = isProcessAlive(worker.descriptor.pid); const observedProcessStartId = processAlive ? getProcessStartId(worker.descriptor.pid) : undefined; const processIdentityMatches = @@ -3691,16 +3752,15 @@ export class DaemonSupervisor { observedProcessStartId === worker.descriptor.processStartId; if (processAlive && processIdentityMatches) { try { - await this.connectWorker(worker, 1500); + await this.connectWorker(worker, 1500, assertCurrent); + assertCurrent(); await this.subscribeWorker(worker, worker.descriptor.rootActiveSessionId); - await this.refreshWorkerSummaries(worker, true); - if (this.isWorkerRecoveryCancelled(worker)) { - return; - } + assertCurrent(); + await this.refreshWorkerSummaries(worker, true, recoveryRevision); + assertCurrent(); if (worker.descriptor.processStartId === undefined && observedProcessStartId) { worker.descriptor.processStartId = observedProcessStartId; } - await this.assertRecoveryAllowed(); worker.descriptor.lifecycle = "ready"; worker.descriptor.consecutiveFailures = 0; this.persistWorker(worker); @@ -3711,6 +3771,7 @@ export class DaemonSupervisor { throw error; } await this.assertRecoveryAllowed(); + assertCurrent(); worker.client?.close(); worker.client = undefined; if (retryIndex < WORKER_RETRY_DELAYS_MS.length - 1) { @@ -3729,6 +3790,7 @@ export class DaemonSupervisor { const recoveryCommand = worker.descriptor.ownerClientId ? worker.transientCreateCommand : undefined; if (!recoveryCommand || !worker.launchEnv) { await this.recoverUncertainWorkerOperations(worker, false); + assertCurrent(); worker.descriptor.lifecycle = "failed"; worker.descriptor.lastError = "Waiting for a client with fresh runtime context"; this.persistWorker(worker); @@ -3740,15 +3802,14 @@ export class DaemonSupervisor { worker, safeToKillWorkerProcess, ); - if (this.isWorkerRecoveryCancelled(worker)) { - return; - } + assertCurrent(); await this.launchWorker( recoveryCommand, worker, worker.descriptor.ownerClientId, promptLifecycleRecovery, ); + assertCurrent(); return; } catch (error) { if (isSupervisorRecoveryCancelled(error) || this.isWorkerRecoveryCancelled(worker)) { @@ -3756,6 +3817,7 @@ export class DaemonSupervisor { } try { await this.assertRecoveryAllowed(); + assertCurrent(); } catch { return; } @@ -3769,6 +3831,7 @@ export class DaemonSupervisor { } try { await this.assertRecoveryAllowed(); + assertCurrent(); } catch { return; } @@ -3790,6 +3853,14 @@ export class DaemonSupervisor { ); } + private assertWorkerRecoveryCurrent(worker: ResidentWorker, stopRevision: number): void { + if (worker.stopRevision !== stopRevision || this.isWorkerRecoveryCancelled(worker)) { + throw new SupervisorRecoveryCancelledError( + `Session worker ${worker.descriptor.workerId} recovery was cancelled`, + ); + } + } + private async restoreWorkerPromptLifecycles( worker: ResidentWorker, recovery: ReadonlyMap, @@ -3953,8 +4024,14 @@ export class DaemonSupervisor { workerClient.close(); } - private async refreshWorkerSummaries(worker: ResidentWorker, recovery = false): Promise { - if (this.isWorkerStopping(worker)) { + private async refreshWorkerSummaries( + worker: ResidentWorker, + recovery = false, + expectedRecoveryRevision = worker.stopRevision, + ): Promise { + if (recovery) { + this.assertWorkerRecoveryCurrent(worker, expectedRecoveryRevision); + } else if (this.isWorkerStopping(worker)) { throw new Error("Session worker is stopping"); } const workerClient = worker.client; @@ -3971,6 +4048,9 @@ export class DaemonSupervisor { ) { throw new Error("Session worker authority changed during roster refresh"); } + if (recovery) { + this.assertWorkerRecoveryCurrent(worker, expectedRecoveryRevision); + } const summaries = sessionSummariesFromResponse(response); const nextSummaries = new Map(summaries.map((summary) => [summary.activeSessionId ?? summary.id, summary])); const root = nextSummaries.get(worker.descriptor.rootActiveSessionId); @@ -3991,6 +4071,7 @@ export class DaemonSupervisor { if (root) { if (recovery) { await this.assertRecoveryAllowed(); + this.assertWorkerRecoveryCurrent(worker, expectedRecoveryRevision); } worker.descriptor.rootSessionId = root.sessionId; worker.descriptor.sessionFile = root.sessionFile; @@ -6178,6 +6259,41 @@ export class DaemonSupervisor { archiveSession = false, recoveryCleanup = false, directChild?: { child: ChildProcess; closed: Promise }, + ): Promise { + if (removeDescriptor && !directChild) { + const pending = worker.stopOperation; + if (pending) { + if (archiveSession && !worker.descriptor.archiveOnStop) { + this.persistWorkerStopTombstone(worker, true); + } + return pending; + } + const operation = this.performWorkerStop( + worker, + removeDescriptor, + force, + archiveSession, + recoveryCleanup, + directChild, + ); + worker.stopOperation = operation; + try { + await operation; + } finally { + if (worker.stopOperation === operation) worker.stopOperation = undefined; + } + return; + } + await this.performWorkerStop(worker, removeDescriptor, force, archiveSession, recoveryCleanup, directChild); + } + + private async performWorkerStop( + worker: ResidentWorker, + removeDescriptor: boolean, + force: boolean, + archiveSession: boolean, + recoveryCleanup: boolean, + directChild?: { child: ChildProcess; closed: Promise }, ): Promise { const releaseStopOwnership = this.acquireWorkerStopOwnership(worker); try { @@ -6247,6 +6363,13 @@ export class DaemonSupervisor { } this.reportCleanupFailure(`worker rollback state ${worker.descriptor.workerId}`, error); } + if (!recoveryCleanup && !directChild) { + // The tombstone/revision above closes recovery admission. Join every + // recovery task that could already have passed an earlier async gate + // before proving the process and durable registration absent. + await Promise.all([worker.recovery, worker.deferredRecovery].flatMap((task) => (task ? [task] : []))); + assertStopStillApplies(); + } const transferError = new Error("Session worker stopped during snapshot transfer"); const generationTranscripts = new Set(); for (const [activeSessionId, generations] of [...(worker.snapshotGenerations ?? new Map())]) { @@ -6359,28 +6482,51 @@ export class DaemonSupervisor { * interrupted cleanup instead of leaving a dead worker registered forever. */ private scheduleWorkerStopFinalization(worker: ResidentWorker): void { - if (worker.stopFinalization) { + const pid = worker.descriptor.pid; + const processStartId = worker.descriptor.processStartId; + const stopRevision = worker.stopRevision; + const existing = worker.stopFinalization; + if ( + existing?.pid === pid && + existing.processStartId === processStartId && + existing.stopRevision === stopRevision + ) { return; } - worker.stopFinalization = this.finalizeTimedOutWorkerStop(worker).finally(() => { + let promise!: Promise; + promise = this.finalizeTimedOutWorkerStop(worker, pid, processStartId, stopRevision).finally(() => { + if (worker.stopFinalization?.promise !== promise) return; worker.stopFinalization = undefined; + if ( + this.workers.get(worker.descriptor.workerId) === worker && + worker.descriptor.stopRequestedAt !== undefined + ) { + this.scheduleWorkerStopFinalization(worker); + } }); + worker.stopFinalization = { pid, processStartId, stopRevision, promise }; + void promise.catch(() => undefined); } - private async finalizeTimedOutWorkerStop(worker: ResidentWorker): Promise { - // Bind to the exact process generation being stopped: a retry can rescind - // the stop and relaunch with a new pid, and the OS can recycle the old - // pid. The finalizer must never follow either successor. - const pid = worker.descriptor.pid; - const processStartId = worker.descriptor.processStartId; + private async finalizeTimedOutWorkerStop( + worker: ResidentWorker, + pid: number, + processStartId: string | undefined, + stopRevision: number, + ): Promise { + // Bind to the exact process generation and stop intent being finalized. + // A retry may rescind this stop without changing pid, and the OS may + // recycle the pid after exit; neither case may inherit this finalizer. const isStopGenerationCurrent = () => this.workers.get(worker.descriptor.workerId) === worker && + worker.stopRevision === stopRevision && worker.descriptor.stopRequestedAt !== undefined && worker.descriptor.pid === pid && worker.descriptor.processStartId === processStartId; - // Do not signal or remove anything until stop intent is durable. A failed - // first tombstone write is retried here without losing the registration. - while (!this.shuttingDown && isStopGenerationCurrent()) { + // Do not signal or remove anything until stop intent is durable. This + // remains active during supervisor shutdown so shutdown cannot outrun a + // failed tombstone write and leave an adoptable live registration. + while (isStopGenerationCurrent()) { try { this.persistWorkerStopTombstone(worker, worker.descriptor.archiveOnStop === true); break; @@ -6389,7 +6535,7 @@ export class DaemonSupervisor { await unrefDelay(STOP_FINALIZATION_RETRY_MS); } } - if (this.shuttingDown || !isStopGenerationCurrent()) return; + if (!isStopGenerationCurrent()) return; // A replaced pid counts as gone (never SIGKILL a recycled pid); an // unobservable identity counts as alive (never clean up a possibly-live // worker). kill(0) probes every poll; ps-backed checks are throttled. @@ -6422,7 +6568,7 @@ export class DaemonSupervisor { }; const sigkillDeadline = Date.now() + STOP_FINALIZATION_SIGKILL_GRACE_MS; let killed = false; - while (!this.shuttingDown) { + while (true) { if (!isStopGenerationCurrent()) { return; } @@ -6444,15 +6590,10 @@ export class DaemonSupervisor { await unrefDelay(STOP_FINALIZATION_RECHECK_MS); } // Retry transient cleanup failures (for example catalog archival) so a - // dead worker's registration is never stranded permanently. Each attempt - // bumps the worker's stopRevision, so rescission is detected through the - // registration and tombstone instead of the waiting-phase snapshot. - const isCleanupStillWanted = () => - this.workers.get(worker.descriptor.workerId) === worker && - worker.descriptor.stopRequestedAt !== undefined && - worker.descriptor.pid === pid && - worker.descriptor.processStartId === processStartId; - while (!this.shuttingDown && isCleanupStillWanted()) { + // dead worker's registration is never stranded permanently. A failed + // attempt advances stopRevision and installs a generation-keyed successor + // finalizer; this exact finalizer then retires without clearing it. + while (isStopGenerationCurrent()) { try { await this.stopWorker(worker, true, true, worker.descriptor.archiveOnStop === true); this.log(`Finalized timed-out stop for worker ${worker.descriptor.workerId}`); @@ -6683,13 +6824,31 @@ export class DaemonSupervisor { try { await this.stopWorker(worker, true, forceWorkers, true); } catch (error) { - if (!(error instanceof WorkerStopTimeoutError)) { - throw error; - } this.log( - `Worker ${worker.descriptor.workerId} remains tombstoned for recovery after shutdown: ${error.message}`, + `Worker ${worker.descriptor.workerId} remains tombstoned while shutdown drains cleanup: ${String(error)}`, ); } + // Shutdown may not disable the only retry path after its public + // response succeeds. Keep ownership, catalog, and the server alive + // until this supervisor has proved the exact registration absent. + while (this.workers.get(worker.descriptor.workerId) === worker) { + if (worker.descriptor.stopRequestedAt === undefined) { + try { + this.persistWorkerStopTombstone(worker, true); + } catch (error) { + this.reportCleanupFailure(`shutdown worker tombstone ${worker.descriptor.workerId}`, error); + } + } + this.scheduleWorkerStopFinalization(worker); + const finalization = worker.stopFinalization; + if (finalization) { + await finalization.promise.catch((error) => + this.reportCleanupFailure(`shutdown worker ${worker.descriptor.workerId}`, error), + ); + } else { + await unrefDelay(STOP_FINALIZATION_RECHECK_MS); + } + } }), ); if (!this.hasPersistedWorkerDescriptors()) { diff --git a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts index 23d5d4f729..73ad65a8d2 100644 --- a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts @@ -150,6 +150,13 @@ interface DeferredRecoveryWorker { stopRevision: number; } +interface TestStopFinalization { + pid: number; + processStartId?: string; + stopRevision: number; + promise: Promise; +} + interface DeferredRecoveryHarness { workers: Map; shuttingDown: boolean; @@ -580,6 +587,51 @@ describe("daemon worker supervisor monitoring", () => { expect(workers.size).toBe(0); }); + it("cancels an existing recovery at the final pre-spawn generation fence", async () => { + workerLaunchTestState.capture = true; + const root = mkdtempSync(join(tmpdir(), "prime-supervisor-recovery-pre-spawn-test-")); + const descriptorDir = join(root, "descriptors"); + mkdirSync(descriptorDir, { recursive: true }); + supervisorRegistryDirs.add(root); + const existing = createExistingLaunchWorker(root, descriptorDir); + const workers = new Map([[existing.descriptor.workerId, existing]]); + const gateReached = createDeferred(); + const releaseGate = createDeferred(); + let assertionCount = 0; + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + ...createSupervisorSnapshotState(), + defaultSessionConfig: { cwd: root, agentDir: root }, + descriptorDir, + socketPath: join(root, "supervisor.sock"), + workers, + shuttingDown: false, + assertRecoveryAllowed: vi.fn(async () => { + assertionCount++; + if (assertionCount === 2) { + gateReached.resolve(); + await releaseGate.promise; + } + }), + log: vi.fn(), + }) as { + launchWorker( + command: { type: "create"; config: { cwd: string; agentDir: string } }, + worker: object, + ): Promise; + }; + + const launch = supervisor.launchWorker({ type: "create", config: { cwd: root, agentDir: root } }, existing); + await gateReached.promise; + existing.stopRevision++; + existing.intentionalStop = true; + existing.descriptor.stopRequestedAt = new Date().toISOString(); + releaseGate.resolve(); + + await expect(launch).rejects.toMatchObject({ code: "supervisor_recovery_cancelled" }); + expect(workerLaunchTestState.spawned).toHaveLength(0); + expect(readdirSync(descriptorDir).filter((name) => name.endsWith(".json"))).toEqual([]); + }); + it("rolls back promptly when the child closes its startup gate before commit", async () => { const root = mkdtempSync(join(tmpdir(), "prime-supervisor-closed-gate-test-")); const descriptorDir = join(root, "descriptors"); @@ -958,8 +1010,7 @@ describe("daemon worker supervisor monitoring", () => { } }); - it("completes shutdown without awaiting an unsignalable worker finalizer", async () => { - vi.useFakeTimers(); + it("keeps shutdown cleanup retries active before releasing daemon ownership", async () => { const root = mkdtempSync(join(tmpdir(), "prime-supervisor-shutdown-finalization-test-")); supervisorRegistryDirs.add(root); const worker = { @@ -978,28 +1029,36 @@ describe("daemon worker supervisor monitoring", () => { snapshotLoads: new Map(), intentionalStop: true, stopRevision: 1, - stopFinalization: new Promise(() => {}), + stopFinalization: undefined as TestStopFinalization | undefined, }; const workers = new Map([[worker.descriptor.workerId, worker]]); + const cleanupStarted = createDeferred(); + const releaseCleanup = createDeferred(); + const stopWorker = vi + .fn(async () => { + cleanupStarted.resolve(); + await releaseCleanup.promise; + workers.delete(worker.descriptor.workerId); + }) + .mockRejectedValueOnce(new Error("descriptor unlink failed")); const exit = vi.spyOn(process, "exit").mockImplementation(((code?: string | number | null) => { throw new Error(`exit ${code}`); }) as typeof process.exit); - const killSpy = vi.spyOn(childProcessModule, "signalProcessGroupOrProcess").mockImplementation(() => {}); - const existsSpy = vi.spyOn(childProcessModule, "processIdExists").mockReturnValue(true); - const aliveSpy = vi.spyOn(childProcessModule, "isProcessAlive").mockReturnValue(true); + const existsSpy = vi.spyOn(childProcessModule, "processIdExists").mockReturnValue(false); const catalogStop = vi.fn(async () => undefined); - const log = vi.fn(); const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { shuttingDown: false, signalCleanupHandlers: [], workers, clients: new Set(), - persistWorkerStopTombstone: vi.fn(), - hasPersistedWorkerDescriptors: vi.fn(() => true), + stopWorker, + persistWorker: vi.fn(), + hasPersistedWorkerDescriptors: vi.fn(() => false), + supervisorConfigPath: join(root, "supervisor-config"), catalog: { stop: catalogStop }, cleanupSocket: vi.fn(), snapshotCacheRoot: join(root, "cache"), - log, + log: vi.fn(), }) as { shutdown(exitCode: number, stopWorkers: boolean, relaunch?: boolean, forceWorkers?: boolean): Promise; }; @@ -1009,19 +1068,18 @@ describe("daemon worker supervisor monitoring", () => { () => undefined, (error: unknown) => error, ); - await vi.advanceTimersByTimeAsync(2000); - await expect(shutdown).resolves.toEqual(new Error("exit 0")); + await cleanupStarted.promise; + expect(catalogStop).not.toHaveBeenCalled(); + expect(exit).not.toHaveBeenCalled(); - expect(workers.has(worker.descriptor.workerId)).toBe(true); - expect(killSpy).not.toHaveBeenCalled(); + releaseCleanup.resolve(); + await expect(shutdown).resolves.toEqual(new Error("exit 0")); + expect(stopWorker).toHaveBeenCalledTimes(2); expect(catalogStop).toHaveBeenCalledOnce(); - expect(log).toHaveBeenCalledWith(expect.stringContaining("remains tombstoned for recovery")); expect(exit).toHaveBeenCalledWith(0); } finally { exit.mockRestore(); - killSpy.mockRestore(); existsSpy.mockRestore(); - aliveSpy.mockRestore(); } }); @@ -1324,6 +1382,7 @@ describe("daemon worker supervisor monitoring", () => { archiveOnStop?: boolean; }; intentionalStop: boolean; + stopRevision: number; summaries: Map; }; type RetryHarness = { @@ -1346,9 +1405,11 @@ describe("daemon worker supervisor monitoring", () => { archiveOnStop: true, }, intentionalStop: true, + stopRevision: 0, summaries: new Map(), }; const persistWorker = vi.fn(() => { + expect(worker.stopRevision).toBe(1); expect(worker.intentionalStop).toBe(false); expect(worker.descriptor.stopRequestedAt).toBeUndefined(); expect(worker.descriptor.archiveOnStop).toBeUndefined(); @@ -1780,8 +1841,12 @@ describe("daemon worker supervisor monitoring", () => { }, summaries: new Map([[root.activeSessionId, root as SessionSummary]]), intentionalStop: false, + stopRevision: 0, }; - const supervisor = Object.create(DaemonSupervisor.prototype) as { + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers: new Map([[worker.descriptor.workerId, worker]]), + shuttingDown: false, + }) as { refreshWorkerSummaries(target: typeof worker, recovery: boolean): Promise; }; @@ -2043,6 +2108,7 @@ describe("daemon worker supervisor monitoring", () => { }; const stopWorker = vi.fn(async () => {}); const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers: new Map([[worker.descriptor.workerId, worker]]), assertRecoveryAllowed: vi.fn(async () => undefined), stopWorker, log: vi.fn(), @@ -2067,10 +2133,17 @@ describe("daemon worker supervisor monitoring", () => { pid: process.pid, stopRequestedAt: new Date().toISOString(), archiveOnStop: false, - } as { processStartId?: string }, + } as { + workerId: string; + pid: number; + stopRequestedAt: string; + archiveOnStop: boolean; + processStartId?: string; + }, }; const stopOrder: string[] = []; const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers: new Map([[worker.descriptor.workerId, worker]]), assertRecoveryAllowed: vi.fn(async () => undefined), connectWorker: vi.fn(async () => { stopOrder.push("connect"); @@ -2101,11 +2174,18 @@ describe("daemon worker supervisor monitoring", () => { pid: process.pid, stopRequestedAt: new Date().toISOString(), archiveOnStop: false, - } as { processStartId?: string }, + } as { + workerId: string; + pid: number; + stopRequestedAt: string; + archiveOnStop: boolean; + processStartId?: string; + }, }; const persistWorker = vi.fn(); const stopWorker = vi.fn(async () => {}); const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers: new Map([[worker.descriptor.workerId, worker]]), assertRecoveryAllowed: vi.fn(async () => undefined), connectWorker: vi.fn(async () => { throw new Error("connect refused"); @@ -2137,12 +2217,15 @@ describe("daemon worker supervisor monitoring", () => { }, intentionalStop: true, stopRevision: 0, - stopFinalization: undefined as Promise | undefined, + stopFinalization: undefined as TestStopFinalization | undefined, }; let alive = true; - const stopWorker = vi.fn(async () => {}); + const workers = new Map([[worker.descriptor.workerId, worker]]); + const stopWorker = vi.fn(async () => { + workers.delete(worker.descriptor.workerId); + }); const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { - workers: new Map([[worker.descriptor.workerId, worker]]), + workers, shuttingDown: false, stopWorker, persistWorker: vi.fn(), @@ -2163,7 +2246,7 @@ describe("daemon worker supervisor monitoring", () => { alive = false; await vi.advanceTimersByTimeAsync(500); - await finalization; + await finalization?.promise; expect(stopWorker).toHaveBeenCalledWith(worker, true, true, false); expect(worker.stopFinalization).toBeUndefined(); @@ -2189,12 +2272,15 @@ describe("daemon worker supervisor monitoring", () => { }, intentionalStop: true, stopRevision: 0, - stopFinalization: undefined as Promise | undefined, + stopFinalization: undefined as TestStopFinalization | undefined, }; let alive = true; - const stopWorker = vi.fn(async () => {}); + const workers = new Map([[worker.descriptor.workerId, worker]]); + const stopWorker = vi.fn(async () => { + workers.delete(worker.descriptor.workerId); + }); const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { - workers: new Map([[worker.descriptor.workerId, worker]]), + workers, shuttingDown: false, stopWorker, persistWorker: vi.fn(), @@ -2215,7 +2301,7 @@ describe("daemon worker supervisor monitoring", () => { const finalization = worker.stopFinalization; await vi.advanceTimersByTimeAsync(10_000); - await finalization; + await finalization?.promise; expect(killSpy).toHaveBeenCalledWith(worker.descriptor.pid, "SIGKILL"); expect(stopWorker).toHaveBeenCalledWith(worker, true, true, true); @@ -2236,7 +2322,7 @@ describe("daemon worker supervisor monitoring", () => { }, intentionalStop: true, stopRevision: 3, - stopFinalization: undefined as Promise | undefined, + stopFinalization: undefined as TestStopFinalization | undefined, }; const stopWorker = vi.fn(async () => {}); const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { @@ -2265,7 +2351,7 @@ describe("daemon worker supervisor monitoring", () => { worker.descriptor.pid = 222_222; await vi.advanceTimersByTimeAsync(20_000); - await finalization; + await finalization?.promise; // The healthy relaunched worker must never be signalled or stopped. expect(killSpy).not.toHaveBeenCalledWith(222_222, "SIGKILL"); @@ -2290,11 +2376,14 @@ describe("daemon worker supervisor monitoring", () => { }, intentionalStop: true, stopRevision: 0, - stopFinalization: undefined as Promise | undefined, + stopFinalization: undefined as TestStopFinalization | undefined, }; - const stopWorker = vi.fn(async () => {}); + const workers = new Map([[worker.descriptor.workerId, worker]]); + const stopWorker = vi.fn(async () => { + workers.delete(worker.descriptor.workerId); + }); const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { - workers: new Map([[worker.descriptor.workerId, worker]]), + workers, shuttingDown: false, stopWorker, persistWorker: vi.fn(), @@ -2315,7 +2404,7 @@ describe("daemon worker supervisor monitoring", () => { const finalization = worker.stopFinalization; await vi.advanceTimersByTimeAsync(20_000); - await finalization; + await finalization?.promise; // The original worker is gone, so the stop is finalized without ever // signalling the unrelated pid owner. @@ -2329,6 +2418,186 @@ describe("daemon worker supervisor monitoring", () => { } }); + it("arms owner cleanup after a client-owned create registers without its disconnected client", async () => { + const client = { id: "owner-create-race" } as DaemonSocketClient; + const worker = { + descriptor: { + workerId: "owner-create-race-worker", + ownerClientId: client.id, + rootActiveSessionId: "owner-create-race-active", + }, + summaries: new Map([ + [ + "owner-create-race-active", + { + id: "owner-create-race-active", + activeSessionId: "owner-create-race-active", + sessionId: "owner-create-race-session", + } as SessionSummary, + ], + ]), + ownerCleanupTimer: undefined as ReturnType | undefined, + }; + const opening = createDeferred(); + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + clients: new Set(), + workers: new Map([[worker.descriptor.workerId, worker]]), + protocolClientIds: new WeakMap(), + createOrReuseWorker: vi.fn(() => opening.promise), + publicSummary: vi.fn((_worker, summary) => summary), + stopWorker: vi.fn(async () => undefined), + log: vi.fn(), + }) as { + handleCommand( + requestClient: DaemonSocketClient, + command: { type: "create"; lifecycle: "client_owned" }, + ): Promise; + }; + + const creating = supervisor.handleCommand(client, { type: "create", lifecycle: "client_owned" }); + opening.resolve(worker); + await creating; + + expect(worker.ownerCleanupTimer).toBeDefined(); + if (worker.ownerCleanupTimer) clearTimeout(worker.ownerCleanupTimer); + }); + + it("joins an admitted recovery before proving exact cleanup settled", async () => { + const recoveryGate = createDeferred(); + let staleRecoveryRejected = false; + const worker = { + descriptor: { + workerId: "owned-cleanup-recovery-worker", + pid: 111_124, + processStartId: "proc:owned-cleanup-recovery", + rootActiveSessionId: "active-owned-cleanup-recovery", + stopRequestedAt: undefined as string | undefined, + }, + client: undefined, + summaries: new Map(), + snapshotCache: new Map(), + transcriptCaches: new Map(), + snapshotGenerations: new Map(), + snapshotLoads: new Map(), + intentionalStop: false, + stopRevision: 0, + recovery: undefined as Promise | undefined, + }; + const workers = new Map([[worker.descriptor.workerId, worker]]); + const deleteWorkerDescriptor = vi.fn(); + let supervisor!: { + stopWorker(target: typeof worker, removeDescriptor: boolean, force?: boolean): Promise; + assertWorkerRecoveryCurrent(target: typeof worker, stopRevision: number): void; + }; + worker.recovery = recoveryGate.promise.then(() => { + try { + supervisor.assertWorkerRecoveryCurrent(worker, 0); + } catch { + staleRecoveryRejected = true; + } + }); + supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers, + shuttingDown: false, + persistWorkerStopTombstone: vi.fn(() => { + worker.intentionalStop = true; + worker.descriptor.stopRequestedAt = new Date().toISOString(); + }), + abortTranscriptPreparations: vi.fn(), + deleteWorkerDescriptor, + invalidateWorkerSessionInputPauses: vi.fn(), + broadcastHeartbeatsChanged: vi.fn(), + log: vi.fn(), + }); + const existsSpy = vi.spyOn(childProcessModule, "processIdExists").mockReturnValue(false); + try { + const stopping = supervisor.stopWorker(worker, true, true); + await Promise.resolve(); + expect(worker.descriptor.stopRequestedAt).toBeDefined(); + expect(deleteWorkerDescriptor).not.toHaveBeenCalled(); + expect(workers.get(worker.descriptor.workerId)).toBe(worker); + + recoveryGate.resolve(); + await stopping; + expect(staleRecoveryRejected).toBe(true); + expect(deleteWorkerDescriptor).toHaveBeenCalledOnce(); + expect(workers.has(worker.descriptor.workerId)).toBe(false); + } finally { + existsSpy.mockRestore(); + } + }); + + it("single-flights concurrent exact-generation completion stops", async () => { + const worker = { + descriptor: { workerId: "owned-cleanup-concurrent-worker", rootActiveSessionId: "active-concurrent" }, + }; + const stopStarted = createDeferred(); + const releaseStop = createDeferred(); + const stopWorkerUntracked = vi.fn(async () => { + stopStarted.resolve(); + await releaseStop.promise; + }); + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers: new Map([[worker.descriptor.workerId, worker]]), + workerStopCounts: new Map(), + stopWorkerUntracked, + }) as { + stopWorker(target: typeof worker, removeDescriptor: boolean): Promise; + }; + + const first = supervisor.stopWorker(worker, true); + await stopStarted.promise; + const second = supervisor.stopWorker(worker, true); + expect(stopWorkerUntracked).toHaveBeenCalledOnce(); + releaseStop.resolve(); + await expect(Promise.all([first, second])).resolves.toEqual([undefined, undefined]); + expect(stopWorkerUntracked).toHaveBeenCalledOnce(); + }); + + it("keeps a newer stop generation finalizer when the old generation retires", async () => { + const worker = { + descriptor: { + workerId: "owned-cleanup-finalizer-generation-worker", + pid: 111_125, + processStartId: "proc:owned-cleanup-finalizer", + rootActiveSessionId: "active-owned-cleanup-finalizer", + stopRequestedAt: new Date().toISOString(), + }, + intentionalStop: true, + stopRevision: 1, + stopFinalization: undefined as TestStopFinalization | undefined, + }; + const workers = new Map([[worker.descriptor.workerId, worker]]); + const oldGeneration = createDeferred(); + const newGeneration = createDeferred(); + const finalizeTimedOutWorkerStop = vi.fn( + (_target: typeof worker, _pid: number, _startId: string | undefined, revision: number) => + revision === 1 ? oldGeneration.promise : newGeneration.promise, + ); + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers, + finalizeTimedOutWorkerStop, + }) as { + scheduleWorkerStopFinalization(target: typeof worker): void; + }; + + supervisor.scheduleWorkerStopFinalization(worker); + const oldFinalization = worker.stopFinalization; + worker.stopRevision = 2; + supervisor.scheduleWorkerStopFinalization(worker); + const currentFinalization = worker.stopFinalization; + expect(currentFinalization).not.toBe(oldFinalization); + + oldGeneration.resolve(); + await oldFinalization?.promise; + expect(worker.stopFinalization).toBe(currentFinalization); + + workers.delete(worker.descriptor.workerId); + newGeneration.resolve(); + await currentFinalization?.promise; + expect(worker.stopFinalization).toBeUndefined(); + }); + it("reports exact privacy-safe owned cleanup states", () => { const descriptorDir = mkdtempSync(join(tmpdir(), "prime-supervisor-owned-cleanup-status-")); try { @@ -2484,7 +2753,7 @@ describe("daemon worker supervisor monitoring", () => { }, intentionalStop: true, stopRevision: 0, - stopFinalization: undefined as Promise | undefined, + stopFinalization: undefined as TestStopFinalization | undefined, }; const workers = new Map([[worker.descriptor.workerId, worker]]); const persistWorker = vi.fn<() => void>().mockImplementationOnce(() => { @@ -2509,7 +2778,7 @@ describe("daemon worker supervisor monitoring", () => { supervisor.scheduleWorkerStopFinalization(worker); expect(stopWorker).not.toHaveBeenCalled(); await vi.advanceTimersByTimeAsync(20_000); - await worker.stopFinalization; + await worker.stopFinalization?.promise; expect(persistWorker).toHaveBeenCalledTimes(2); expect(stopWorker).toHaveBeenCalledWith(worker, true, true, false); @@ -2698,7 +2967,7 @@ describe("daemon worker supervisor monitoring", () => { }, intentionalStop: true, stopRevision: 0, - stopFinalization: undefined as Promise | undefined, + stopFinalization: undefined as TestStopFinalization | undefined, }; const workers = new Map([[worker.descriptor.workerId, worker]]); const stopWorker = vi.fn(async () => { @@ -2750,7 +3019,7 @@ describe("daemon worker supervisor monitoring", () => { }, intentionalStop: true, stopRevision: 0, - stopFinalization: undefined as Promise | undefined, + stopFinalization: undefined as TestStopFinalization | undefined, }; const workers = new Map([[worker.descriptor.workerId, worker]]); const stopWorker = vi.fn(async () => { @@ -2788,7 +3057,7 @@ describe("daemon worker supervisor monitoring", () => { // ...but once identity is observable again, escalation still fires. startIdSpy.mockReturnValue("proc:original"); await vi.advanceTimersByTimeAsync(5000); - await finalization; + await finalization?.promise; expect(killSpy).toHaveBeenCalledWith(worker.descriptor.pid, "SIGKILL"); expect(stopWorker).toHaveBeenCalled(); } finally { @@ -2811,7 +3080,7 @@ describe("daemon worker supervisor monitoring", () => { }, intentionalStop: true, stopRevision: 0, - stopFinalization: undefined as Promise | undefined, + stopFinalization: undefined as TestStopFinalization | undefined, }; const stopWorker = vi.fn(async () => {}); const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { @@ -2859,7 +3128,7 @@ describe("daemon worker supervisor monitoring", () => { }, intentionalStop: true, stopRevision: 0, - stopFinalization: undefined as Promise | undefined, + stopFinalization: undefined as TestStopFinalization | undefined, }; const stopWorker = vi.fn(async () => {}); const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { @@ -2905,7 +3174,7 @@ describe("daemon worker supervisor monitoring", () => { }, intentionalStop: true, stopRevision: 0, - stopFinalization: undefined as Promise | undefined, + stopFinalization: undefined as TestStopFinalization | undefined, }; const workers = new Map([[worker.descriptor.workerId, worker]]); const stopWorker = vi @@ -2932,7 +3201,7 @@ describe("daemon worker supervisor monitoring", () => { const finalization = worker.stopFinalization; await vi.advanceTimersByTimeAsync(20_000); - await finalization; + await finalization?.promise; // The first attempt failed transiently; the registration is still // cleaned up by a retry instead of being stranded forever. @@ -2954,7 +3223,7 @@ describe("daemon worker supervisor monitoring", () => { }, intentionalStop: true, stopRevision: 0, - stopFinalization: undefined as Promise | undefined, + stopFinalization: undefined as TestStopFinalization | undefined, }; let alive = true; const stopWorker = vi.fn(async () => {}); @@ -2979,7 +3248,7 @@ describe("daemon worker supervisor monitoring", () => { worker.intentionalStop = false; alive = false; await vi.advanceTimersByTimeAsync(500); - await finalization; + await finalization?.promise; expect(stopWorker).not.toHaveBeenCalled(); } finally { @@ -3076,16 +3345,21 @@ describe("daemon worker supervisor monitoring", () => { recovery: undefined, intentionalStop: true, stopRevision: 0, - stopFinalization: undefined as Promise | undefined, + stopFinalization: undefined as TestStopFinalization | undefined, }; const workers = new Map([[worker.descriptor.workerId, worker]]); let releaseFinalization!: () => void; - worker.stopFinalization = new Promise((resolveFinalization) => { + const finalizationPromise = new Promise((resolveFinalization) => { releaseFinalization = () => { workers.delete(worker.descriptor.workerId); resolveFinalization(); }; }); + worker.stopFinalization = { + pid: worker.descriptor.pid, + stopRevision: worker.stopRevision, + promise: finalizationPromise, + }; const stopWorker = vi.fn(async () => {}); const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { workers, @@ -3124,7 +3398,7 @@ describe("daemon worker supervisor monitoring", () => { recovery: undefined, intentionalStop: true, stopRevision: 0, - stopFinalization: undefined as Promise | undefined, + stopFinalization: undefined as TestStopFinalization | undefined, }; const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { workers: new Map([[worker.descriptor.workerId, worker]]), @@ -3166,7 +3440,7 @@ describe("daemon worker supervisor monitoring", () => { recovery: undefined, intentionalStop: true, stopRevision: 0, - stopFinalization: undefined as Promise | undefined, + stopFinalization: undefined as TestStopFinalization | undefined, }; const workers = new Map([[worker.descriptor.workerId, worker]]); const stopWorker = vi.fn(async () => { diff --git a/packages/coding-agent/test/daemon-supervisor-process.test.ts b/packages/coding-agent/test/daemon-supervisor-process.test.ts index dba48b9036..346d83d126 100644 --- a/packages/coding-agent/test/daemon-supervisor-process.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-process.test.ts @@ -16,6 +16,7 @@ import { import { readSessionInfo, SessionManager } from "../src/core/session-manager.js"; import { DaemonAgentConnection } from "../src/modes/agent-connection/daemon-agent-connection.js"; import { DaemonClient, getDaemonSocketCloseReason } from "../src/modes/daemon/daemon-client.js"; +import { createDaemonCommandEnvelope } from "../src/modes/daemon/daemon-protocol.js"; import type { SessionSummary } from "../src/modes/daemon/daemon-session-list.js"; import type { DaemonWorkerDescriptor } from "../src/modes/daemon/daemon-worker-protocol.js"; @@ -815,6 +816,94 @@ describe("daemon supervisor resident workers", () => { } }, 75_000); + it("cleans up a client-owned create whose owner disconnects before registration", async () => { + const root = tempDir(); + const agentDir = join(root, "agent"); + const projectDir = join(root, "project"); + const socketPath = join( + tmpdir(), + `prime-supervisor-owned-opening-${process.pid}-${randomUUID().slice(0, 8)}.sock`, + ); + mkdirSync(projectDir, { recursive: true }); + + const supervisor = spawnSupervisor(agentDir, socketPath, projectDir); + const readiness = await connectEventually(socketPath, supervisor); + readiness.close(); + const openingOwner = createConnection(socketPath); + openingOwner.on("error", () => undefined); + await new Promise((resolveConnect, rejectConnect) => { + openingOwner.once("connect", resolveConnect); + openingOwner.once("error", rejectConnect); + }); + const commandId = `opening-${randomUUID()}`; + const wire = `${JSON.stringify( + createDaemonCommandEnvelope( + { + type: "create", + lifecycle: "client_owned", + noSession: true, + launchEnv: { TSX_TSCONFIG_PATH: resolve(__dirname, "../../../tsconfig.json") }, + config: { cwd: projectDir, agentDir, noTools: true, noExtensions: true }, + }, + commandId, + "opening-owner", + ), + )}\n`; + const closed = new Promise((resolveClose) => openingOwner.once("close", resolveClose)); + await new Promise((resolveWrite, rejectWrite) => { + openingOwner.end(wire, (error?: Error | null) => (error ? rejectWrite(error) : resolveWrite())); + }); + await new Promise((resolveDelay) => setTimeout(resolveDelay, 5)); + expect(countWorkerDescriptors(agentDir)).toBe(0); + openingOwner.destroy(); + // The command bytes are committed to the socket, but the client is gone + // before launchWorker can publish the new resident registration. + await closed; + + let descriptor: DaemonWorkerDescriptor | undefined; + const descriptorDeadline = Date.now() + 15_000; + while (!descriptor && Date.now() < descriptorDeadline) { + try { + descriptor = readWorkerDescriptor(agentDir); + } catch { + await new Promise((resolveDelay) => setTimeout(resolveDelay, 10)); + } + } + if (!descriptor) throw new Error(`Opening client-owned create was not registered\n${readDaemonLogs(agentDir)}`); + expect(descriptor.ownerClientId).toBe("opening-owner"); + workerPids.add(descriptor.pid); + + const observer = await connectEventually(socketPath, supervisor); + try { + let sawStopping = false; + let settled = false; + const deadline = Date.now() + 50_000; + while (Date.now() < deadline) { + const response = await observer.request({ + type: "get_owned_session_cleanup", + activeSessionId: descriptor.rootActiveSessionId, + }); + if (!response.success) throw new Error(response.error); + const status = (response.data as { status?: string } | undefined)?.status; + if (status === "stopping") sawStopping = true; + if (status === "settled") { + settled = true; + break; + } + await new Promise((resolveDelay) => setTimeout(resolveDelay, 50)); + } + expect(sawStopping).toBe(true); + expect(settled).toBe(true); + await waitForProcessGone(descriptor.pid); + workerPids.delete(descriptor.pid); + expect(countWorkerDescriptors(agentDir)).toBe(0); + await observer.request({ type: "shutdown" }); + await waitForSocketGone(socketPath); + } finally { + observer.close(); + } + }, 75_000); + it("releases an adopted client-owned worker when disposal races supervisor replacement", async () => { const root = tempDir(); const agentDir = join(root, "agent"); From 90e0f092c0b79d5ef2c4538b1fa2be69d61da6e9 Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Sat, 29 Aug 2026 21:15:55 -0600 Subject: [PATCH 09/12] fix(daemon): preserve cleanup authority --- .pylon/upstream-review.md | 2 +- .../coding-agent/docs/agent-connection.md | 2 +- .../src/modes/daemon/daemon-supervisor.ts | 29 +++++-- .../test/daemon-supervisor-monitor.test.ts | 82 ++++++++++++++++++- 4 files changed, 107 insertions(+), 8 deletions(-) diff --git a/.pylon/upstream-review.md b/.pylon/upstream-review.md index b16277a334..3b9c65cb4c 100644 --- a/.pylon/upstream-review.md +++ b/.pylon/upstream-review.md @@ -64,5 +64,5 @@ This ledger records Prime upstream evidence and the decision taken for each over - `snapshot-recovery-integrity`: **hybridize**. In addition to immutable transfer identities and isolated generation failure, the candidate preserves terminal child errors through passive hydration, stabilizes snapshot cursors across asynchronous roster selection and the exact attach-registration cut, invalidates stale supervisor selection/materialization, closes a private channel if an announced frame is interrupted, and binds worker frames to the current authenticated channel and latest authoritative roster. - `authoritative-owned-session-cleanup`: **hybridize**. `complete_owned_session` cannot succeed until the exact process generation is absent, ancillary journals are absent, the durable descriptor is removed and verified last, and only then the in-memory registration is removed. Tombstone, archive, and descriptor failures retain a stopping registration and arm a single exact-generation retry finalizer. A different client that already knows the exact host-private active-session ID may poll only `active | stopping | settled`; no PID, path, owner identity, descriptor, token, or diagnostics cross the public boundary. - Compatibility classification: **backward-compatible, additive, capability-gated**. Private session workers continue to advertise the shared default server capabilities and do not falsely offer the supervisor-only cleanup query. Current clients reject the query locally against stock `v0.8.1` or schema 26 without sending an unknown command. The proof applies only to registrations created or durably adopted by a capable supervisor; it cannot retroactively certify a descriptorless historical orphan. -- Validation: 693 focused protocol/client/connection/session/snapshot/supervisor tests passed across 13 files; 40 correlated-lifecycle/queue/continuation tests passed; 13 real-process supervisor tests passed with 8 fixture-gated skips, including direct completion, owner loss before and after worker registration, visible `stopping`, replacement-client `settled` proof, supervisor replacement, exact worker exit, and zero descriptors. Final exact-head review also forced recovery-join, shutdown-drain, concurrent-completion, and generation-keyed-finalizer repairs before renewed exact-head review. Both stock/current `v0.8.1` adoption directions passed with local cleanup-query rejection on the stock supervisor. `npm run check` passed Biome over 941 files, TypeScript, installer rendering, and browser smoke. The root build passed; the live generated model catalog was restored afterward. `git diff --check` passed. +- Validation: 694 focused protocol/client/connection/session/snapshot/supervisor tests passed across 13 files; 40 correlated-lifecycle/queue/continuation tests passed; 13 real-process supervisor tests passed with 8 fixture-gated skips, including direct completion, owner loss before and after worker registration, visible `stopping`, replacement-client `settled` proof, supervisor replacement, exact worker exit, and zero descriptors. Final exact-head review also forced recovery-join, published-replacement rollback retention, reentrant shutdown single-flight, concurrent-completion, and generation-keyed-finalizer repairs before renewed exact-head review. Both stock/current `v0.8.1` adoption directions passed with local cleanup-query rejection on the stock supervisor. `npm run check` passed Biome over 941 files, TypeScript, installer rendering, and browser smoke. The root build passed; the live generated model catalog was restored afterward. `git diff --check` passed. - Revisit when Prime upstream provides equivalent immutable snapshot, attach admission, worker-channel authority, and descriptor-last crash cleanup contracts without weakening the stock fallback or exposing host-private process/session details. diff --git a/packages/coding-agent/docs/agent-connection.md b/packages/coding-agent/docs/agent-connection.md index c61913be63..f7c85ab581 100644 --- a/packages/coding-agent/docs/agent-connection.md +++ b/packages/coding-agent/docs/agent-connection.md @@ -119,7 +119,7 @@ The `AgentConnection` method promise is a client convenience. It should not be t The query returns only `{ status: "active" | "stopping" | "settled" }`. It does not expose process IDs, worker or owner IDs, filesystem paths, tokens, descriptor contents, or diagnostics. Stock v0.8.1 does not offer this capability, so `DaemonClient` rejects the query locally without sending an unknown command. The proof applies to registrations created or durably adopted by a capable supervisor; it cannot retroactively certify a descriptorless orphan created before that supervisor observed it. -Cleanup closes and joins any admitted recovery before it can publish a successor. Concurrent completion requests for the same resident generation share one stop operation. Finalizers are keyed by process identity and stop revision, and daemon shutdown keeps registry ownership and retry services alive until every requested worker cleanup settles. A client-owned create rechecks owner liveness after registration, so a disconnect while daemon readiness or process launch is still pending cannot miss its cleanup timer. +Cleanup closes and joins any admitted recovery before it can publish a successor. Concurrent completion requests for the same resident generation share one stop operation. Finalizers are keyed by process identity and stop revision, and daemon shutdown keeps registry ownership and retry services alive until every requested worker cleanup settles. Repeated shutdown commands and signals join that same drain rather than forcing an early process exit. A client-owned create rechecks owner liveness after registration, so a disconnect while daemon readiness or process launch is still pending cannot miss its cleanup timer. ## Session Replacement diff --git a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts index df9835e68c..362253e3e8 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts @@ -715,6 +715,7 @@ export class DaemonSupervisor { private socketLease?: DaemonSocketPathLease; private ownership?: Awaited>; private cleanupPromise?: Promise; + private shutdownTask?: Promise; private shuttingDown = false; private updateRestartPhase?: "draining" | "fencing" | "prepared"; private readonly mutationDrain = new MutationDrainLatch(); @@ -6470,7 +6471,15 @@ export class DaemonSupervisor { if (removeDescriptor) { this.deleteWorkerDescriptor(worker); } - this.workers.delete(worker.descriptor.workerId); + const preserveRecoveryRegistration = + !removeDescriptor && + recoveryCleanup && + directChild !== undefined && + this.workers.get(worker.descriptor.workerId) === worker && + (worker.descriptor.stopRequestedAt !== undefined || this.shuttingDown); + if (!preserveRecoveryRegistration) { + this.workers.delete(worker.descriptor.workerId); + } if (!this.shuttingDown) { this.broadcastHeartbeatsChanged(); } @@ -6797,16 +6806,26 @@ export class DaemonSupervisor { } } - private async shutdown( + private shutdown( exitCode: number, stopWorkers: boolean, relaunch = false, forceWorkers = false, closingReason?: DaemonClosingReason, ): Promise { - if (this.shuttingDown) { - process.exit(exitCode); - } + if (this.shutdownTask) return this.shutdownTask; + const task = this.runShutdown(exitCode, stopWorkers, relaunch, forceWorkers, closingReason); + this.shutdownTask = task; + return task; + } + + private async runShutdown( + exitCode: number, + stopWorkers: boolean, + relaunch: boolean, + forceWorkers: boolean, + closingReason?: DaemonClosingReason, + ): Promise { this.shuttingDown = true; this.clearIdleEvictionTimer(); await this.idleEvictionSweep?.catch(() => undefined); diff --git a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts index 73ad65a8d2..9c7b67f52a 100644 --- a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts @@ -1064,7 +1064,8 @@ describe("daemon worker supervisor monitoring", () => { }; try { - const shutdown = supervisor.shutdown(0, true, false, true).then( + const shutdownTask = supervisor.shutdown(0, true, false, true); + const shutdown = shutdownTask.then( () => undefined, (error: unknown) => error, ); @@ -1072,6 +1073,12 @@ describe("daemon worker supervisor monitoring", () => { expect(catalogStop).not.toHaveBeenCalled(); expect(exit).not.toHaveBeenCalled(); + // A second public command or signal joins the authoritative drain; + // it cannot force process exit while the descriptor is still present. + const reentrantShutdown = supervisor.shutdown(99, true, false, true); + expect(reentrantShutdown).toBe(shutdownTask); + expect(exit).not.toHaveBeenCalled(); + releaseCleanup.resolve(); await expect(shutdown).resolves.toEqual(new Error("exit 0")); expect(stopWorker).toHaveBeenCalledTimes(2); @@ -2527,6 +2534,79 @@ describe("daemon worker supervisor monitoring", () => { } }); + it("preserves a tombstoned registration when replacement-child rollback joins exact cleanup", async () => { + const rollbackGate = createDeferred(); + const worker = { + descriptor: { + workerId: "owned-cleanup-published-recovery-worker", + pid: 111_126, + processStartId: "proc:owned-cleanup-published-recovery", + rootActiveSessionId: "active-owned-cleanup-published-recovery", + lifecycle: "starting" as const, + stopRequestedAt: undefined as string | undefined, + }, + client: undefined, + summaries: new Map(), + snapshotCache: new Map(), + transcriptCaches: new Map(), + snapshotGenerations: new Map(), + snapshotLoads: new Map(), + intentionalStop: false, + stopRevision: 0, + recovery: undefined as Promise | undefined, + }; + const workers = new Map([[worker.descriptor.workerId, worker]]); + const child = { + exitCode: 0, + signalCode: null, + kill: vi.fn(() => true), + } as unknown as ChildProcess; + const deleteWorkerDescriptor = vi.fn(); + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers, + shuttingDown: false, + persistWorker: vi.fn(), + abortTranscriptPreparations: vi.fn(), + deleteWorkerDescriptor, + invalidateWorkerSessionInputPauses: vi.fn(), + broadcastHeartbeatsChanged: vi.fn(), + log: vi.fn(), + }) as { + stopWorker( + target: typeof worker, + removeDescriptor: boolean, + force?: boolean, + archiveSession?: boolean, + recoveryCleanup?: boolean, + directChild?: { child: ChildProcess; closed: Promise }, + ): Promise; + }; + worker.recovery = (async () => { + await rollbackGate.promise; + await supervisor.stopWorker(worker, false, true, false, true, { child, closed: Promise.resolve() }); + // The replacement child is gone, but its tombstoned descriptor and + // registration still belong to the outer authoritative stop. + expect(workers.get(worker.descriptor.workerId)).toBe(worker); + expect(worker.descriptor.stopRequestedAt).toBeDefined(); + })(); + const existsSpy = vi.spyOn(childProcessModule, "processIdExists").mockReturnValue(false); + const aliveSpy = vi.spyOn(childProcessModule, "isProcessAlive").mockReturnValue(false); + try { + const stopping = supervisor.stopWorker(worker, true, true); + await Promise.resolve(); + expect(worker.descriptor.stopRequestedAt).toBeDefined(); + rollbackGate.resolve(); + await stopping; + + expect(child.kill).toHaveBeenCalledWith("SIGTERM"); + expect(deleteWorkerDescriptor).toHaveBeenCalledOnce(); + expect(workers.has(worker.descriptor.workerId)).toBe(false); + } finally { + existsSpy.mockRestore(); + aliveSpy.mockRestore(); + } + }); + it("single-flights concurrent exact-generation completion stops", async () => { const worker = { descriptor: { workerId: "owned-cleanup-concurrent-worker", rootActiveSessionId: "active-concurrent" }, From 49c592ffc863263123020c7a8aaa97158bbbc7ec Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Sat, 29 Aug 2026 21:21:44 -0600 Subject: [PATCH 10/12] fix(daemon): retain generic recovery cleanup --- .pylon/upstream-review.md | 2 +- .../src/modes/daemon/daemon-supervisor.ts | 1 - .../test/daemon-supervisor-monitor.test.ts | 19 +++++++++++++++---- 3 files changed, 16 insertions(+), 6 deletions(-) diff --git a/.pylon/upstream-review.md b/.pylon/upstream-review.md index 3b9c65cb4c..c061af6330 100644 --- a/.pylon/upstream-review.md +++ b/.pylon/upstream-review.md @@ -64,5 +64,5 @@ This ledger records Prime upstream evidence and the decision taken for each over - `snapshot-recovery-integrity`: **hybridize**. In addition to immutable transfer identities and isolated generation failure, the candidate preserves terminal child errors through passive hydration, stabilizes snapshot cursors across asynchronous roster selection and the exact attach-registration cut, invalidates stale supervisor selection/materialization, closes a private channel if an announced frame is interrupted, and binds worker frames to the current authenticated channel and latest authoritative roster. - `authoritative-owned-session-cleanup`: **hybridize**. `complete_owned_session` cannot succeed until the exact process generation is absent, ancillary journals are absent, the durable descriptor is removed and verified last, and only then the in-memory registration is removed. Tombstone, archive, and descriptor failures retain a stopping registration and arm a single exact-generation retry finalizer. A different client that already knows the exact host-private active-session ID may poll only `active | stopping | settled`; no PID, path, owner identity, descriptor, token, or diagnostics cross the public boundary. - Compatibility classification: **backward-compatible, additive, capability-gated**. Private session workers continue to advertise the shared default server capabilities and do not falsely offer the supervisor-only cleanup query. Current clients reject the query locally against stock `v0.8.1` or schema 26 without sending an unknown command. The proof applies only to registrations created or durably adopted by a capable supervisor; it cannot retroactively certify a descriptorless historical orphan. -- Validation: 694 focused protocol/client/connection/session/snapshot/supervisor tests passed across 13 files; 40 correlated-lifecycle/queue/continuation tests passed; 13 real-process supervisor tests passed with 8 fixture-gated skips, including direct completion, owner loss before and after worker registration, visible `stopping`, replacement-client `settled` proof, supervisor replacement, exact worker exit, and zero descriptors. Final exact-head review also forced recovery-join, published-replacement rollback retention, reentrant shutdown single-flight, concurrent-completion, and generation-keyed-finalizer repairs before renewed exact-head review. Both stock/current `v0.8.1` adoption directions passed with local cleanup-query rejection on the stock supervisor. `npm run check` passed Biome over 941 files, TypeScript, installer rendering, and browser smoke. The root build passed; the live generated model catalog was restored afterward. `git diff --check` passed. +- Validation: 695 focused protocol/client/connection/session/snapshot/supervisor tests passed across 13 files; 40 correlated-lifecycle/queue/continuation tests passed; 13 real-process supervisor tests passed with 8 fixture-gated skips, including direct completion, owner loss before and after worker registration, visible `stopping`, replacement-client `settled` proof, supervisor replacement, exact worker exit, and zero descriptors. Final exact-head review also forced recovery-join, published-replacement rollback retention, reentrant shutdown single-flight, concurrent-completion, and generation-keyed-finalizer repairs before renewed exact-head review. Both stock/current `v0.8.1` adoption directions passed with local cleanup-query rejection on the stock supervisor. `npm run check` passed Biome over 941 files, TypeScript, installer rendering, and browser smoke. The root build passed; the live generated model catalog was restored afterward. `git diff --check` passed. - Revisit when Prime upstream provides equivalent immutable snapshot, attach admission, worker-channel authority, and descriptor-last crash cleanup contracts without weakening the stock fallback or exposing host-private process/session details. diff --git a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts index 362253e3e8..4d53497da1 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts @@ -6474,7 +6474,6 @@ export class DaemonSupervisor { const preserveRecoveryRegistration = !removeDescriptor && recoveryCleanup && - directChild !== undefined && this.workers.get(worker.descriptor.workerId) === worker && (worker.descriptor.stopRequestedAt !== undefined || this.shuttingDown); if (!preserveRecoveryRegistration) { diff --git a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts index 9c7b67f52a..e2a8225785 100644 --- a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts @@ -2534,7 +2534,10 @@ describe("daemon worker supervisor monitoring", () => { } }); - it("preserves a tombstoned registration when replacement-child rollback joins exact cleanup", async () => { + it.each([ + { name: "direct-child cancellation rollback", passDirectChild: true }, + { name: "ordinary failed-recovery rollback", passDirectChild: false }, + ])("preserves a tombstoned registration during $name", async ({ passDirectChild }) => { const rollbackGate = createDeferred(); const worker = { descriptor: { @@ -2583,8 +2586,15 @@ describe("daemon worker supervisor monitoring", () => { }; worker.recovery = (async () => { await rollbackGate.promise; - await supervisor.stopWorker(worker, false, true, false, true, { child, closed: Promise.resolve() }); - // The replacement child is gone, but its tombstoned descriptor and + await supervisor.stopWorker( + worker, + false, + true, + false, + true, + passDirectChild ? { child, closed: Promise.resolve() } : undefined, + ); + // The replacement process is gone, but its tombstoned descriptor and // registration still belong to the outer authoritative stop. expect(workers.get(worker.descriptor.workerId)).toBe(worker); expect(worker.descriptor.stopRequestedAt).toBeDefined(); @@ -2598,7 +2608,8 @@ describe("daemon worker supervisor monitoring", () => { rollbackGate.resolve(); await stopping; - expect(child.kill).toHaveBeenCalledWith("SIGTERM"); + if (passDirectChild) expect(child.kill).toHaveBeenCalledWith("SIGTERM"); + else expect(child.kill).not.toHaveBeenCalled(); expect(deleteWorkerDescriptor).toHaveBeenCalledOnce(); expect(workers.has(worker.descriptor.workerId)).toBe(false); } finally { From c37f4bddb360826b5878a9989f3779bea5ee2e33 Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Sat, 29 Aug 2026 21:33:51 -0600 Subject: [PATCH 11/12] test(coding-agent): align daemon fixtures with safety contracts --- .../coding-agent/test/daemon-supervisor-heartbeats.test.ts | 1 + .../test/suite/regressions/4257-update-restart-resume.test.ts | 4 ++-- .../test/suite/regressions/4601-worker-snapshot-cache.test.ts | 2 +- .../test/suite/regressions/4656-resume-active-session.test.ts | 2 +- 4 files changed, 5 insertions(+), 4 deletions(-) diff --git a/packages/coding-agent/test/daemon-supervisor-heartbeats.test.ts b/packages/coding-agent/test/daemon-supervisor-heartbeats.test.ts index 68f8f98e45..b8c7d6c84b 100644 --- a/packages/coding-agent/test/daemon-supervisor-heartbeats.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-heartbeats.test.ts @@ -98,6 +98,7 @@ describe("daemon supervisor heartbeat aggregation", () => { const supervisor = createSupervisorHarness(); const target = { ...worker("ready"), + descriptor: { workerId: "target", lifecycle: "ready" }, heartbeatSnapshot: [{ job: { id: "heartbeat-1" } }], heartbeatSnapshotStale: false, }; diff --git a/packages/coding-agent/test/suite/regressions/4257-update-restart-resume.test.ts b/packages/coding-agent/test/suite/regressions/4257-update-restart-resume.test.ts index ddfd58b603..bfd979480a 100644 --- a/packages/coding-agent/test/suite/regressions/4257-update-restart-resume.test.ts +++ b/packages/coding-agent/test/suite/regressions/4257-update-restart-resume.test.ts @@ -523,8 +523,8 @@ describe("issue #4257 update restart resume", () => { id, socket: { destroyed: false, - write: (chunk: string) => { - writes.push(chunk); + write: (chunk: string | Uint8Array) => { + writes.push(typeof chunk === "string" ? chunk : Buffer.from(chunk).toString("utf8")); return true; }, }, diff --git a/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts b/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts index 42de6719bc..68c2e6a8d0 100644 --- a/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts +++ b/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts @@ -209,7 +209,7 @@ describe("ENG-4601 worker snapshot cache", () => { // Fully consume the fresh generation to prove its new backing is readable. } second.dispose?.(); - }, 15_000); + }, 30_000); it.each([1, 2])( "keeps %i same-ID worker stream(s) bounded by socket drain without shared cleanup", diff --git a/packages/coding-agent/test/suite/regressions/4656-resume-active-session.test.ts b/packages/coding-agent/test/suite/regressions/4656-resume-active-session.test.ts index ac1b8fc757..e5b43ac75c 100644 --- a/packages/coding-agent/test/suite/regressions/4656-resume-active-session.test.ts +++ b/packages/coding-agent/test/suite/regressions/4656-resume-active-session.test.ts @@ -263,7 +263,7 @@ describe("ENG-4656 active session resume", () => { return releaseReservation; }), attachClient: vi.fn(async (client: DaemonSocketClient) => { - expect(client.attachedActiveSessionIds.has(targetActiveSessionId)).toBe(true); + expect(client.attachedActiveSessionIds.has(targetActiveSessionId)).toBe(false); operations.push("attach"); client.attachedActiveSessionIds.add(targetActiveSessionId); return { worker: targetWorker, result: createAttachResult(targetActiveSessionId) }; From c5a34dffcd44eaaa4d66ea43786d19260a10e7f2 Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Sat, 29 Aug 2026 21:35:09 -0600 Subject: [PATCH 12/12] test(coding-agent): bound spill regression for loaded CI --- .../test/suite/regressions/4601-worker-snapshot-cache.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts b/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts index 68c2e6a8d0..71913db6dc 100644 --- a/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts +++ b/packages/coding-agent/test/suite/regressions/4601-worker-snapshot-cache.test.ts @@ -209,7 +209,7 @@ describe("ENG-4601 worker snapshot cache", () => { // Fully consume the fresh generation to prove its new backing is readable. } second.dispose?.(); - }, 30_000); + }, 60_000); it.each([1, 2])( "keeps %i same-ID worker stream(s) bounded by socket drain without shared cleanup",