diff --git a/.pylon/features.yaml b/.pylon/features.yaml index 5684e1935e..3d1a21fc36 100644 --- a/.pylon/features.yaml +++ b/.pylon/features.yaml @@ -170,3 +170,36 @@ 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. + + 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 28e5e8c30a..c061af6330 100644 --- a/.pylon/upstream-review.md +++ b/.pylon/upstream-review.md @@ -45,3 +45,24 @@ 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. + +## 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: 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/.changes/11-snapshot-recovery-integrity.md b/packages/coding-agent/.changes/11-snapshot-recovery-integrity.md new file mode 100644 index 0000000000..a4e207529e --- /dev/null +++ b/packages/coding-agent/.changes/11-snapshot-recovery-integrity.md @@ -0,0 +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..f7c85ab581 100644 --- a/packages/coding-agent/docs/agent-connection.md +++ b/packages/coding-agent/docs/agent-connection.md @@ -102,12 +102,25 @@ 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. + +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 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/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index b26fdfbef5..441325541a 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,13 @@ interface RlmChildRun { unsubscribe?: () => void; } +interface RetainedRlmChild { + session: AgentSession; + run?: RlmChildRun; + /** Durable terminal status for a rehydrated child that has no in-memory run. */ + terminalStatus?: "done" | "error"; +} + interface RlmSubagentModelSelection { model: Model; } @@ -1195,7 +1206,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 +4344,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 +4407,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 +10054,7 @@ export class AgentSession { }); recorded.add(run.id); } - for (const [childId, childSession] of this._rlmChildSessions) { + for (const [childId, { session: childSession, terminalStatus }] of this._rlmChildSessions) { if ( this._deletingRlmChildren.has(childId) || recorded.has(childId) || @@ -10063,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); } @@ -10131,7 +10142,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 +10419,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) { @@ -10431,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)) { @@ -10444,7 +10460,11 @@ 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), + terminalStatus, + }); if (unsubscribe) { this._rlmChildUnsubscribes.set(childId, unsubscribe); } @@ -10455,15 +10475,81 @@ 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, + terminalStatus: "done" | "error" = "done", + ): 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: 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. 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, + }; } /** Live recursive child roster from lifecycle state, including nested work under retained parents. */ @@ -10476,16 +10562,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 +10570,16 @@ export class AgentSession { snapshots.push(...child.getRlmChildSnapshots()); } } - for (const [childId, child] 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, terminalStatus); 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 +10598,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 +10608,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 +10679,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 +10689,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 +10719,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 +10736,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 +10856,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 +10865,7 @@ export class AgentSession { sessionDir: childSessionDir, model: modelSelection.model, status: "queued", + toolUseCount: 0, settled: false, abort: noopRlmChildAbort, publication: createAgentMessageDeferred(), @@ -10809,27 +10878,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 +10961,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 +10995,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 +11047,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 +11081,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/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/agent-connection/daemon-agent-connection.ts b/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts index 3b6f4be8d4..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 @@ -42,6 +42,7 @@ import { collectDaemonClientEnv, collectDaemonLaunchEnv, type DaemonAttachResult, + type DaemonClientCapability, type DaemonCommand, type DaemonEventCursor, type DaemonOutbound, @@ -113,12 +114,46 @@ type DaemonSessionReplaced = Extract; + settled: boolean; + attachmentEpoch: number; + requestAttemptId?: number; promise: Promise; resolve: (snapshot: DaemonSessionSnapshot) => void; reject: (error: Error) => void; 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; + 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 { activeSessionId: string; sessionId: string; @@ -314,9 +349,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 runtimeSnapshotAttempt: DaemonRuntimeSnapshotAttempt | undefined; 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; @@ -382,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; } @@ -437,44 +476,74 @@ export class DaemonAgentConnection implements AgentConnection { requestedActiveSessionId: string, resumeCursor: DaemonEventCursor | undefined, resetEventProgress: boolean, + maxAttempts = 2, + logicalPurpose: DaemonSnapshotPurpose = "attach", ): 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, + logicalPurpose, + "attach", + ); + try { + result = await this.requestData({ + type: "attach", + activeSessionId: requestedActiveSessionId, + snapshotGenerationNonce: randomUUID(), + 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) { + 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; + } + } 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); @@ -532,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; } @@ -1543,32 +1614,76 @@ 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, + "replacement", + attemptIndex === 0 ? "replacement" : "attach", + ); + 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, + snapshotGenerationNonce: randomUUID(), + 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) { + 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) { + requestAttempt.state = "failed"; + if (attemptIndex === 1 || !reattached || this.disposed || this.attachmentEpoch !== attachmentEpoch) { + throw error; + } + } finally { + this.finishSnapshotRequestAttempt(requestAttempt); + } } - const snapshot = result.snapshotStream - ? await this.waitForSnapshot(result.snapshotStream.id) - : result.snapshot; + if (this.disposed || this.attachmentEpoch !== attachmentEpoch) { + throw new Error("Daemon connection attachment changed during reattach"); + } + validateDaemonSnapshotIdentity(result.snapshot, result.activeSessionId); mapDaemonSessionSnapshot(result.snapshot, undefined, this.supportsCorrelatedPromptLifecycle()); const staged = this.stageSnapshotCommit(snapshot, { @@ -1712,6 +1827,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,60 +1954,142 @@ 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); } 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"); + if (!pendingReplacement && !requestAttempt && !acceptsUnsolicitedRuntimeSnapshot) return; + this.ignoredSnapshotIds.delete(message.snapshotId); if (pendingReplacement) { - if ( - message.purpose !== "replacement" || - (pendingReplacement.snapshotId !== undefined && pendingReplacement.snapshotId !== message.snapshotId) - ) { - 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) { + 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); + 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; + } + // 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; + if (!requestAttempt && (message.purpose === "replacement" || message.purpose === "resync")) { + this.startRuntimeSnapshotAttempt(message.purpose, message.activeSessionId, message.snapshotId); + } 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 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") - ? this.recoverFailedSnapshot(purpose, snapshotError) - : undefined; - if (recoveryPromise) { - this.snapshotRecoveryPromises.set(message.snapshotId, recoveryPromise); - } - this.rejectSnapshotAssembly(message.snapshotId, assembly, snapshotError); - this.ignoreSnapshotId(message.snapshotId); - if (quarantinesReplacement) this.failClosedReplacementReconciliation(); - if (recoveryPromise) { - try { - await recoveryPromise; - } finally { - this.snapshotRecoveryPromises.delete(message.snapshotId); - } - if (!this.terminalCloseEmitted) this.releaseChunkedReplacementFence(); - } + const existingAssembly = this.snapshotAssemblies.get(message.snapshotId); + const explicitPurpose = message.purpose; + const assembly = + existingAssembly ?? + (this.pendingChunkedReplacement || explicitPurpose + ? this.getSnapshotAssembly(message.snapshotId) + : undefined); + if (!assembly || assembly.settled || assembly.attachmentEpoch !== this.attachmentEpoch) return; + 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)) { @@ -2213,11 +2411,136 @@ 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 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; + } + + 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,40 +2550,62 @@ export class DaemonAgentConnection implements AgentConnection { void promise.catch(() => undefined); const timeout = setTimeout(() => { const current = this.snapshotAssemblies.get(snapshotId); - if (current) { - 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(); - } + if (current && !current.settled) { + 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(); 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); 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.snapshotRecoveryPromises.clear(); + this.completedSnapshotAttemptIds.clear(); + this.snapshotRequestAttempts.clear(); this.ignoredSnapshotIds.clear(); + if (!recoverRuntime) this.runtimeSnapshotAttempt = undefined; } private ignoreSnapshotId(snapshotId: string): void { @@ -2287,60 +2632,167 @@ 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 { - const snapshot = await this.getInitialSnapshot(); - if (this.disposed) { - return; - } - this.attachedSessionId = snapshot.state.sessionId; - this.attachedSessionFile = snapshot.state.sessionFile; - 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 }); } } - 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; + if (assembly.begin) this.validateSnapshotAttemptBegin(requestAttempt, assembly.begin); + } 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 +2916,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 +2974,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 +2984,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; @@ -2548,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 { @@ -2780,8 +3257,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/active-session-state.ts b/packages/coding-agent/src/modes/daemon/active-session-state.ts index fe062253ce..360fb6347b 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>; + /** Settles only after all queued private-frame reservations and their predecessors retire. */ + 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 0d58993757..fe4c9ed0fb 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,29 @@ 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; +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", @@ -345,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 = @@ -486,6 +510,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 +2729,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; } @@ -2983,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(); } @@ -3244,7 +3282,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(); @@ -3843,16 +3880,40 @@ export class AgentDaemon { const snapshotSignal = streamsSnapshot ? markClientSnapshotStreaming(client, state.activeSessionId) : undefined; - let result: DaemonAttachResult; + const snapshotAttemptSignal = snapshotSignal ? createSnapshotAttemptSignal(snapshotSignal) : undefined; + 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) { @@ -3864,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, @@ -3874,15 +3933,15 @@ 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)}:${command.snapshotGenerationNonce ?? "shared"}`, messages: result.snapshot.messages, - targetChunkBytes: SNAPSHOT_TARGET_CHUNK_BYTES, - signal: snapshotSignal, + signal: snapshotAttemptSignal, }); } catch (error) { state.clients.delete(client); @@ -3907,7 +3966,7 @@ export class AgentDaemon { streamedResult, transcript, "attach", - snapshotSignal, + snapshotAttemptSignal, true, ).catch((error) => this.log(`could not stream attach snapshot: ${String(error)}`)); }); @@ -5076,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 ? { @@ -5084,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"); @@ -5101,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], @@ -5127,37 +5184,126 @@ 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); - return { - activeSessionId: state.activeSessionId, - summary: summaryForActiveSession(state), - state: connectionState, - messages: session.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: { + 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 this.createWorkerSnapshotTransfer(options.activeSessionId, options.snapshotId, existing, payload); + } + 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: 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 this.createWorkerSnapshotTransfer(options.activeSessionId, options.snapshotId, generation, payload); + } catch (error) { + if (this.snapshotPayloadGenerations.get(options.activeSessionId) === generation) { + this.snapshotPayloadGenerations.delete(options.activeSessionId); + } + throw error; + } + } + } + + 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 [ + result.snapshot.state.sessionId, + cursor?.generation ?? "legacy", + cursor?.sequence ?? result.lastEventSequence, + result.snapshot.messages.length, + ].join(":"); } private async streamWorkerSnapshot( @@ -5179,7 +5325,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?.(); @@ -5193,7 +5340,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 = { @@ -5233,6 +5390,7 @@ export class AgentDaemon { activeSessionId: result.activeSessionId, snapshotId: stream.id, error: streamError.message, + purpose: purpose === "catchup" ? "resync" : purpose, }, purpose, undefined, @@ -5257,11 +5415,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, @@ -5269,10 +5433,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++; @@ -5281,7 +5443,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", @@ -5294,6 +5456,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`)); @@ -5336,33 +5499,108 @@ export class AgentDaemon { private async writeWorkerSnapshotBuffer( client: DaemonSocketClient, - buffer: Buffer, + buffer: SnapshotTranscriptWireChunk, message: DaemonOutbound, purpose: "attach" | "replacement" | "catchup", 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; + 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; + }); + 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(); } - 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); @@ -5376,9 +5614,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); }); } @@ -6560,6 +6796,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 +6924,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) => { @@ -6704,6 +6948,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, @@ -6717,12 +6962,12 @@ 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, + signal: snapshotAttemptSignal, }); this.write(client, { ...message, messages: [], snapshotFollows: true }); void this.streamWorkerSnapshot( @@ -6739,7 +6984,7 @@ export class AgentDaemon { }, transcript, "replacement", - snapshotSignal, + snapshotAttemptSignal, true, ).catch((error) => { this.log(`could not stream replacement snapshot: ${String(error)}`); @@ -6787,7 +7032,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; @@ -6797,33 +7041,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 && @@ -6831,13 +7048,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"; } @@ -6865,35 +7080,60 @@ 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 = `${activeSessionId}-${state.eventGeneration}-${state.lastEventSequence}`; + const snapshotId = createSnapshotTransferId(); const snapshotSignal = markClientSnapshotStreaming(client, activeSessionId); + const snapshotAttemptSignal = createSnapshotAttemptSignal(snapshotSignal); 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, + 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, @@ -6909,7 +7149,7 @@ export class AgentDaemon { }, transcript, purpose === "replacement" ? "replacement" : "catchup", - snapshotSignal, + snapshotAttemptSignal, true, ); continue; @@ -6934,15 +7174,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"; @@ -7007,37 +7258,40 @@ 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( - { - 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 } : {}), - }, - typeof line === "string" ? Buffer.from(line) : line, - ) - : line; - const accepted = client.socket.write(wireData); - if (!accepted) { - client.backpressured = true; - } - return accepted; + ? this.encodeWorkerPrivateFrameParts(line, message, payloadEncoding, snapshotPurpose) + : payloadParts; + 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 4331171d9c..d5a420947a 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-protocol.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-protocol.ts @@ -75,8 +75,11 @@ 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. +// Revision 26 correlates snapshot failures that occur before a begin frame can be emitted. +// 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; @@ -94,12 +97,17 @@ 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"; export interface DaemonPromptAdmissionCancellationResult { status: DaemonPromptAdmissionCancellationStatus; } +export type DaemonOwnedSessionCleanupStatus = "active" | "stopping" | "settled"; +export interface DaemonOwnedSessionCleanupResult { + status: DaemonOwnedSessionCleanupStatus; +} export type DaemonServerCapability = | DaemonClientCapability | "delete_rlm_subagent" @@ -122,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"; @@ -147,6 +156,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", ]; @@ -170,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; @@ -410,6 +426,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) @@ -424,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 } @@ -717,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", @@ -760,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, @@ -1061,6 +1086,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 } @@ -1210,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-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..4d53497da1 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, @@ -138,7 +139,13 @@ 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, + type SnapshotTranscriptWireChunk, +} from "./snapshot-transcript-cache.js"; import { WorkerRecoveryJournal } from "./worker-recovery-journal.js"; type DistributiveOmit = T extends unknown ? Omit : never; @@ -147,7 +154,22 @@ 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; + +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 @@ -182,6 +204,7 @@ const DAEMON_COMMAND_TYPES: ReadonlySet = new Set([ "reattach", "detach", "complete_owned_session", + "get_owned_session_cleanup", "promote_owned_session", "kill", "rename", @@ -292,13 +315,24 @@ interface ResidentWorker { transcriptCaches: Map; snapshotGenerations: Map>; snapshotLoads: Map>; + snapshotLoadSnapshotIds: Map; + /** 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; 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; @@ -313,13 +347,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 { @@ -346,6 +384,15 @@ interface WorkerAttachData { releaseTranscript?: () => void; } +interface SupervisorTranscriptPreparation { + result: DaemonAttachResult; + promise: Promise; + controller: AbortController; + workerClient: DaemonWorkerClient | undefined; + stopRevision: number; + selectionRevision: number; +} + interface SupervisorPromptAdmission { client: DaemonSocketClient; activeSessionId: string; @@ -668,11 +715,15 @@ 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(); 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 +740,13 @@ 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>; + private snapshotSelectionRevisions?: WeakMap>; private readonly pendingSessionNames = new Set(); private readonly catalog: DaemonCatalogClient; private readonly settingsManager: SettingsManager; @@ -720,7 +774,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 +800,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,11 +1086,19 @@ export class DaemonSupervisor { transcriptCaches: new Map(), snapshotGenerations: new Map(), snapshotLoads: new Map(), + snapshotLoadSnapshotIds: new Map(), + authorizedActiveSessionIds: new Set([durableDescriptor.rootActiveSessionId]), 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)}`); } @@ -1089,16 +1153,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 { @@ -1135,7 +1254,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, }); } }, @@ -1158,7 +1277,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)); @@ -1481,7 +1603,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}`)); @@ -1745,6 +1869,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; @@ -1767,34 +1895,58 @@ 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); + let 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")) { + 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"); + } + 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 +1959,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 +1991,7 @@ export class DaemonSupervisor { "replacement", releaseTranscript, releaseSnapshotReservation, + { selector: targetActiveSessionId, epoch: attachmentEpoch }, ); releaseTranscript = undefined; void streaming.catch((error) => @@ -1954,6 +2111,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)) { @@ -1982,6 +2141,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; @@ -2352,7 +2512,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 { @@ -2471,17 +2636,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 +2688,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`); + } + 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); } - 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) + ); } /** @@ -2550,11 +2776,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 @@ -2570,7 +2793,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 @@ -2615,10 +2838,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), @@ -2650,6 +2876,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, @@ -2686,6 +2913,7 @@ export class DaemonSupervisor { childPid = child.pid; childProcessStartId = getProcessStartId(childPid); await this.assertRecoveryAllowed(); + assertLaunchCurrent(); const descriptor: DaemonWorkerDescriptor = { version: 2, @@ -2715,12 +2943,15 @@ export class DaemonSupervisor { transcriptCaches: new Map(), snapshotGenerations: new Map(), snapshotLoads: new Map(), + snapshotLoadSnapshotIds: new Map(), + authorizedActiveSessionIds: new Set([descriptor.rootActiveSessionId]), intentionalStop: false, stopRevision: 0, launchEnv, transientCreateCommand: ownerClientId ? createCommand : undefined, }; await this.assertRecoveryAllowed(); + assertLaunchCurrent(); worker.descriptor = descriptor; worker.launchEnv = launchEnv; worker.transientCreateCommand = descriptor.ownerClientId ? createCommand : undefined; @@ -2751,7 +2982,9 @@ export class DaemonSupervisor { try { try { + assertLaunchCurrent(); await commitWorkerStartupGate(startupGate); + assertLaunchCurrent(); } catch (error) { startupGate.destroy(); await childClosed; @@ -2759,8 +2992,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); } @@ -2775,14 +3010,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; @@ -2855,11 +3091,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()))); @@ -2870,10 +3112,13 @@ export class DaemonSupervisor { 1000, ); await this.assertRecoveryAllowed(); - client.onFrame((frame) => this.handleWorkerFrame(worker, frame)); + assertCurrent?.(); + 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; @@ -2894,25 +3139,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) { @@ -2922,7 +3163,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 @@ -2933,9 +3190,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 { @@ -2944,34 +3203,49 @@ 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); - 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; } + 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)}`); @@ -2983,7 +3257,10 @@ export class DaemonSupervisor { if (worker.client !== client) { return; } + 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 ?? []) { @@ -3087,6 +3364,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, @@ -3094,6 +3383,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); @@ -3101,14 +3399,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 +3422,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 +3432,133 @@ 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 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, + _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; + 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( + 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; + } + if (snapshotPurpose === "replacement" || snapshotPurpose === "catchup" || snapshotPurpose === "resync") { + this.deliverPurposeBearingSnapshotFailure(activeSessionId, snapshotPurpose, error); + return; + } + this.log(`Ignored stale snapshot ${snapshotId} failure for ${activeSessionId}: ${error.message}`); + } + private retireWorkerSnapshotCache( worker: ResidentWorker, activeSessionId: string, @@ -3153,8 +3578,12 @@ export class DaemonSupervisor { generation.retired = true; 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(); } @@ -3226,12 +3655,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(); @@ -3272,23 +3727,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 = @@ -3296,16 +3753,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); @@ -3316,6 +3772,7 @@ export class DaemonSupervisor { throw error; } await this.assertRecoveryAllowed(); + assertCurrent(); worker.client?.close(); worker.client = undefined; if (retryIndex < WORKER_RETRY_DELAYS_MS.length - 1) { @@ -3334,6 +3791,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); @@ -3345,15 +3803,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)) { @@ -3361,6 +3818,7 @@ export class DaemonSupervisor { } try { await this.assertRecoveryAllowed(); + assertCurrent(); } catch { return; } @@ -3374,6 +3832,7 @@ export class DaemonSupervisor { } try { await this.assertRecoveryAllowed(); + assertCurrent(); } catch { return; } @@ -3395,6 +3854,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, @@ -3507,28 +3974,105 @@ export class DaemonSupervisor { return promptLifecycles; } - private async refreshWorkerSummaries(worker: ResidentWorker, recovery = false): Promise { - if (this.isWorkerStopping(worker)) { - throw new Error("Session worker is stopping"); + private workerAuthorizedActiveSessionIds(worker: ResidentWorker): Set { + if (!worker.authorizedActiveSessionIds) { + const rootActiveSessionId = worker.descriptor?.rootActiveSessionId; + worker.authorizedActiveSessionIds = new Set(rootActiveSessionId ? [rootActiveSessionId] : []); } - if (!worker.client) { - throw new Error("Session worker is not connected"); + return worker.authorizedActiveSessionIds; + } + + 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 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"); + 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 response = await worker.client.request({ type: "list" }, 5000); - const summaries = sessionSummariesFromResponse(response); - worker.summaries = new Map(summaries.map((summary) => [summary.activeSessionId ?? summary.id, summary])); - for (const summary of summaries) { - const activeSessionId = summary.activeSessionId ?? summary.id; - if (summary.streamingMessage?.role === "assistant") { - this.streamReconstructor.seed(activeSessionId, summary.streamingMessage); + 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, sourceClient?: DaemonWorkerClient): void { + const workerClient = sourceClient ?? worker.client; + if (!workerClient) return; + if (worker.client === workerClient) { + void this.handleWorkerClose(worker, workerClient, error); + } + workerClient.close(); + } + + 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; + if (!workerClient) { + throw new Error("Session worker is not connected"); + } + 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"); + } + 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); + if (recovery && !root) { + throw new Error(`Session worker omitted its root session during recovery`); + } + const nextAuthority = this.validateWorkerActiveSessionIds(worker, nextSummaries.keys()); + worker.authorizedActiveSessionIds = nextAuthority; + worker.summaries = nextSummaries; + for (const summary of summaries) { + const activeSessionId = summary.activeSessionId ?? summary.id; + if (summary.streamingMessage?.role === "assistant") { + this.streamReconstructor.seed(activeSessionId, summary.streamingMessage); } else if (!summary.isStreaming) { this.streamReconstructor.clear(activeSessionId); } } - const root = worker.summaries.get(worker.descriptor.rootActiveSessionId); if (root) { if (recovery) { await this.assertRecoveryAllowed(); + this.assertWorkerRecoveryCurrent(worker, expectedRecoveryRevision); } worker.descriptor.rootSessionId = root.sessionId; worker.descriptor.sessionFile = root.sessionFile; @@ -3901,6 +4445,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 +4488,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,35 +4523,51 @@ 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 = 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 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", + ]; + 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") + ) { + workerCapabilities.push("chunked_snapshot", "immutable_snapshot_transfer_v1"); + } + const response = await workerClient.request( + { + type: "attach", + activeSessionId, + snapshotGenerationNonce: randomUUID(), + capabilities: workerCapabilities, + supportsExtensionUi: false, + env: command.env ?? collectDaemonClientEnv(), + }, + 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); })(); match.worker.snapshotLoads.set(snapshotLoadKey, loading); + match.worker.snapshotLoadSnapshotIds.set(snapshotLoadKey, observedSnapshotId); void loading.then( async (loaded) => { try { @@ -4012,18 +4586,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 +4623,26 @@ 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"); + } + 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"); + } break; } } const releaseTranscript = transcript?.retain(); + assertAttachmentCurrent(); client.attachedActiveSessionIds.add(activeSessionId); try { const publicSummary = this.publicSummary(match.worker, result.snapshot.summary); @@ -4066,7 +4658,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, @@ -4165,60 +4759,166 @@ 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 { - const activeSessionId = result.activeSessionId; - const existing = worker.transcriptCaches.get(activeSessionId); - if (existing && (!result.snapshotStream || existing.snapshotId === result.snapshotStream.id)) { - return existing; + 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); } - if (result.snapshot.messages.length < result.snapshot.summary.messageCount) { - throw new Error("Session snapshot generation changed before its transcript could be selected"); + 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); } - if (existing) { - this.retireWorkerSnapshotCache(worker, activeSessionId, existing); + } + + 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, + ): Promise { + const activeSessionId = result.activeSessionId; + while (true) { + const existing = worker.transcriptCaches.get(activeSessionId); + 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 + ? 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 controller = new AbortController(); + const workerClient = worker.client; + const stopRevision = worker.stopRevision; + const selectionRevision = this.snapshotSelectionRevision(worker, activeSessionId); + const promise = prepareSnapshotTranscriptCache({ + activeSessionId, + snapshotId, + messages: result.snapshot.messages, + cacheRoot: this.snapshotCacheRoot, + targetChunkBytes: SNAPSHOT_TARGET_CHUNK_BYTES, + signal: AbortSignal.any([controller.signal, AbortSignal.timeout(SNAPSHOT_ATTEMPT_TIMEOUT_MS)]), + }); + const preparation: SupervisorTranscriptPreparation = { + result, + promise, + 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 || + (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 = { + ...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( 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, }, @@ -4233,10 +4933,23 @@ 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, + ); + const deadlineSignal = AbortSignal.timeout(SNAPSHOT_ATTEMPT_TIMEOUT_MS); + const attemptSignal = attachmentSignal ? AbortSignal.any([attachmentSignal, deadlineSignal]) : deadlineSignal; + if (!stream || client.socket.destroyed || !isAttachmentCurrent()) { releaseSnapshotReservation(); releaseTranscript(); return; @@ -4244,56 +4957,81 @@ 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) { - let chunk: Buffer | undefined; + if (!isAttachmentCurrent() || attemptSignal.aborted) return; + let chunk: SnapshotTranscriptWireChunk | undefined; try { - chunk = await transcript.waitForChunk(chunkCount); + chunk = await this.waitForSnapshotTransferChunk(transcript, stream.id, chunkCount, attemptSignal); } catch (error) { const streamError = error instanceof Error ? error : new Error(String(error)); - this.failWorkerSnapshotCache(worker, result.activeSessionId, streamError, false, stream.id); + const cacheGeneration = this.currentSnapshotGeneration(worker, result.activeSessionId); + if (cacheGeneration?.transcript === transcript) { + this.failCorrelatedSnapshotAttempt( + worker, + result.activeSessionId, + transcript.snapshotId, + streamError, + purpose, + false, + ); + } throw streamError; } - if (!chunk) { - break; - } - if (!(await this.writeSnapshotBuffer(client, chunk))) { - return; - } + if (!isAttachmentCurrent() || attachmentSignal?.aborted) return; + if (attemptSignal.aborted) throw new Error(`Snapshot ${stream.id} timed out`); + if (!chunk) break; + if (!(await this.writeSnapshotBuffer(client, chunk, attemptSignal))) return; chunkCount++; } - 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) { + 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) { + this.deliveredSnapshotFailures.add(streamError); } } catch (deliveryError) { client.socket.destroy(deliveryError instanceof Error ? deliveryError : new Error(String(deliveryError))); @@ -4306,6 +5044,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(); @@ -4340,47 +5090,105 @@ 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); }); } + 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); } @@ -4396,7 +5204,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; } @@ -4409,24 +5232,70 @@ export class DaemonSupervisor { snapshotPurpose, } = frame.header; if (outboundType === "heartbeats_changed") { + if (activeSessionId) { + 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, sourceClient)) { + this.rejectWorkerFrame( + worker, + new Error(`Worker emitted an unauthorized active session ${activeSessionId}`), + sourceClient, + ); + return; + } + if ( + (outboundType === "session_snapshot_begin" || + outboundType === "session_snapshot_chunk" || + outboundType === "session_snapshot_end" || + outboundType === "session_snapshot_failed") && + !activeSessionId + ) { + this.rejectWorkerFrame(worker, new Error("Worker returned an uncorrelatable snapshot frame"), sourceClient); + 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 +5329,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 +5351,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 +5375,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); } @@ -4544,12 +5415,15 @@ 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; 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 +5431,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(frame.payload); + generation.receivedMessageCount = receivedMessageCount; + } else { + 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; } + } 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 +5552,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 = 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 +5592,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 +5616,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 +5645,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" || @@ -4737,7 +5681,7 @@ export class DaemonSupervisor { return; } let publicPayload = frame.payload; - let decodedOutbound: DaemonOutbound | undefined; + let decodedOutbound: DaemonOutbound; if (payloadEncoding === "assistant-delta") { let compactValue: unknown; try { @@ -4746,10 +5690,18 @@ 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"), + sourceClient, + ); + return; + } const reconstructed = this.streamReconstructor.reconstruct(compactValue); if (!reconstructed) { this.scheduleCompactCatchup(worker, activeSessionId); @@ -4757,27 +5709,44 @@ 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"), sourceClient); + 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"), + sourceClient, + ); + return; + } + this.streamReconstructor.observe(decodedOutbound); + if ( + decodedOutbound.type === "session_event" && + decodedOutbound.event.type === "rlm_child_update" && + decodedOutbound.event.child.activeSessionId + ) { + try { + this.provisionWorkerActiveSessionIds(worker, [decodedOutbound.event.child.activeSessionId]); + } catch (error) { + this.rejectWorkerFrame(worker, error instanceof Error ? error : new Error(String(error)), sourceClient); + return; } } + 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 +5765,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; } @@ -4831,14 +5809,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); } } } @@ -4849,10 +5829,14 @@ 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) { - 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 +5905,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 +5914,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 +5978,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 +5987,30 @@ 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: createPublicSnapshotTransferId(), + error: snapshotError.message, + purpose, + }, + AbortSignal.timeout(1_000), + ); + } this.log(`Failed to catch up client ${client.id} for ${activeSessionId}: ${String(error)}`); } } @@ -5250,10 +6260,55 @@ 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 { 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(); } @@ -5271,6 +6326,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++; } @@ -5286,7 +6342,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`); @@ -5306,6 +6364,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())]) { @@ -5403,10 +6468,17 @@ 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); } + const preserveRecoveryRegistration = + !removeDescriptor && + recoveryCleanup && + 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(); } @@ -5418,26 +6490,60 @@ 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; - const stopRevision = worker.stopRevision; + 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.pid === pid && + worker.descriptor.processStartId === processStartId; + // 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; + } catch (error) { + this.reportCleanupFailure(`worker stop tombstone ${worker.descriptor.workerId}`, error); + await unrefDelay(STOP_FINALIZATION_RETRY_MS); + } + } + 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. @@ -5470,7 +6576,7 @@ export class DaemonSupervisor { }; const sigkillDeadline = Date.now() + STOP_FINALIZATION_SIGKILL_GRACE_MS; let killed = false; - while (!this.shuttingDown) { + while (true) { if (!isStopGenerationCurrent()) { return; } @@ -5484,7 +6590,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; } @@ -5492,14 +6598,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; - 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}`); @@ -5565,14 +6667,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; } @@ -5639,6 +6744,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; @@ -5699,16 +6805,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); @@ -5726,13 +6842,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/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..5dc2092883 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,524 @@ 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); + const reading = chunk.readPromise; + try { + return await reading; + } finally { + if (chunk.readPromise === reading) chunk.readPromise = undefined; + } + } + + 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 +686,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 +695,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 +741,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 }); @@ -190,47 +804,68 @@ export class SnapshotTranscriptCache { }); } - retain(): () => void { - if (this.disposed) { - throw new Error(`Snapshot transcript ${this.snapshotId} was disposed`); + 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(); 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 +875,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 +889,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 +913,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/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/src/modes/session-worker/private-framing.ts b/packages/coding-agent/src/modes/session-worker/private-framing.ts index 4caf143f6a..993ed61031 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), @@ -51,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, @@ -59,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"); + 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 frameLength = FRAME_PREFIX_BYTES + headerLength + payloadLength; - if (this.buffered.length - offset < frameLength) { - break; - } - - 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)}`, @@ -97,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 2476257dfd..aabbea4395 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, purpose: "attach" | "replacement") => { + const { messages: _messages, ...snapshot } = target.snapshot; + setImmediate(() => { + fakeClient.emitMessage({ + type: "session_snapshot_begin", + activeSessionId: "active-target", + snapshotId, + snapshot, + messageCount: 1, + targetChunkBytes: 512 * 1024, + purpose, + }); + 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, "replacement"); + fakeClient.attachResultFactory = () => emitTransfer("reattach-g2", false, "attach"); + 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..1cad16ee85 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,59 @@ 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); + 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 +717,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 +811,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 +916,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 +2273,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-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-mode.test.ts b/packages/coding-agent/test/daemon-mode.test.ts index f92101ba4f..2c4d713109 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(); }); @@ -3234,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; @@ -3242,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 { @@ -3351,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); @@ -3400,6 +3410,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 { @@ -4096,6 +4149,7 @@ describe("daemon mode helpers", () => { buildRlmChildSnapshotsWithPassiveRlmSubagents( state: ActiveSessionState, ): Promise>; + getOrHydrateBoundSessionState(selector: string): Promise; handleCommand(client: DaemonSocketClient, command: DaemonCommand): Promise; }; @@ -4217,6 +4271,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 }); } @@ -4464,13 +4528,113 @@ 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 }); + } + }); + + 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("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 { + 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("bounds snapshot stabilization when every child build replaces the runtime session", async () => { + 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); @@ -4496,11 +4660,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 }); @@ -5670,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, { @@ -5728,7 +5901,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 +5925,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 +5937,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 +5967,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 +9284,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..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", @@ -274,6 +284,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-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/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..e2a8225785 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,50 +1029,64 @@ 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; }; 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, ); - 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(); + // 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); 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 +1389,7 @@ describe("daemon worker supervisor monitoring", () => { archiveOnStop?: boolean; }; intentionalStop: boolean; + stopRevision: number; summaries: Map; }; type RetryHarness = { @@ -1346,9 +1412,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(); @@ -1613,7 +1681,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 +1690,177 @@ 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; + }; + + 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, + stopRevision: 0, + }; + 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; }; - expect(() => supervisor.reuseWorkerForCreate(worker, undefined, "/tmp/failed.jsonl")).toThrow(/failed worker/); + 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", () => { @@ -1664,20 +1899,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); }); @@ -1879,6 +2115,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(), @@ -1903,10 +2140,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"); @@ -1937,11 +2181,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"); @@ -1973,12 +2224,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(), @@ -1999,7 +2253,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(); @@ -2025,12 +2279,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(), @@ -2051,7 +2308,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); @@ -2072,7 +2329,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), { @@ -2101,7 +2358,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"); @@ -2126,11 +2383,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(), @@ -2151,7 +2411,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. @@ -2165,14 +2425,60 @@ describe("daemon worker supervisor monitoring", () => { } }); - it("aborts stale stop cleanup when the worker was relaunched during an await", async () => { + 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: "worker-relaunched-during-stop", - pid: 111_115, - rootActiveSessionId: "active-1", - stopRequestedAt: new Date().toISOString() as string | undefined, - archiveOnStop: true, + 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(), @@ -2180,55 +2486,463 @@ describe("daemon worker supervisor monitoring", () => { transcriptCaches: new Map(), snapshotGenerations: new Map(), snapshotLoads: new Map(), - intentionalStop: true, + intentionalStop: false, stopRevision: 0, + recovery: undefined as Promise | undefined, }; const workers = new Map([[worker.descriptor.workerId, worker]]); - const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + 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, - persistWorker: vi.fn(), - persistWorkerStopTombstone: vi.fn(), - reclaimStoppedWorkerCronLock: vi.fn(), - // Archival yields, and a retry relaunches the worker meanwhile. - finalizeArchivedWorkerStop: vi.fn(async () => { - worker.descriptor.pid = 222_222; - worker.descriptor.stopRequestedAt = undefined; + persistWorkerStopTombstone: vi.fn(() => { + worker.intentionalStop = true; + worker.descriptor.stopRequestedAt = new Date().toISOString(); }), - deleteWorkerDescriptor: vi.fn(), + abortTranscriptPreparations: vi.fn(), + deleteWorkerDescriptor, + invalidateWorkerSessionInputPauses: vi.fn(), broadcastHeartbeatsChanged: vi.fn(), log: vi.fn(), - reportCleanupFailure: vi.fn(), - }) as unknown as { - stopWorker( - target: object, - removeDescriptor: boolean, - force?: boolean, - archiveSession?: boolean, - ): Promise; - deleteWorkerDescriptor: ReturnType; - }; - const childProcessModule = await import("../src/utils/child-process.js"); + }); const existsSpy = vi.spyOn(childProcessModule, "processIdExists").mockReturnValue(false); - const aliveSpy = vi.spyOn(childProcessModule, "isProcessAlive").mockReturnValue(false); try { - await expect(supervisor.stopWorker(worker, true, true, true)).rejects.toThrow("was relaunched during stop"); - - // The relaunched worker's registration and descriptor must survive. - expect(workers.has(worker.descriptor.workerId)).toBe(true); - expect(supervisor.deleteWorkerDescriptor).not.toHaveBeenCalled(); + 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(); - aliveSpy.mockRestore(); } }); - it("aborts stale stop cleanup when the stop is rescinded before the relaunch lands", 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: { - workerId: "worker-rescinded-during-stop", - pid: 111_120, - rootActiveSessionId: "active-1", + 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, + 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(); + })(); + 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; + + 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 { + existsSpy.mockRestore(); + aliveSpy.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 { + 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 TestStopFinalization | 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?.promise; + + 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: { + workerId: "worker-relaunched-during-stop", + pid: 111_115, + rootActiveSessionId: "active-1", + stopRequestedAt: new Date().toISOString() as string | undefined, + archiveOnStop: true, + }, + 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 supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + workers, + shuttingDown: false, + persistWorker: vi.fn(), + persistWorkerStopTombstone: vi.fn(), + reclaimStoppedWorkerCronLock: vi.fn(), + // Archival yields, and a retry relaunches the worker meanwhile. + finalizeArchivedWorkerStop: vi.fn(async () => { + worker.descriptor.pid = 222_222; + worker.descriptor.stopRequestedAt = undefined; + }), + deleteWorkerDescriptor: vi.fn(), + broadcastHeartbeatsChanged: vi.fn(), + log: vi.fn(), + reportCleanupFailure: vi.fn(), + }) as unknown as { + stopWorker( + target: object, + removeDescriptor: boolean, + force?: boolean, + archiveSession?: boolean, + ): Promise; + deleteWorkerDescriptor: ReturnType; + }; + const childProcessModule = await import("../src/utils/child-process.js"); + const existsSpy = vi.spyOn(childProcessModule, "processIdExists").mockReturnValue(false); + const aliveSpy = vi.spyOn(childProcessModule, "isProcessAlive").mockReturnValue(false); + try { + await expect(supervisor.stopWorker(worker, true, true, true)).rejects.toThrow("was relaunched during stop"); + + // The relaunched worker's registration and descriptor must survive. + expect(workers.has(worker.descriptor.workerId)).toBe(true); + expect(supervisor.deleteWorkerDescriptor).not.toHaveBeenCalled(); + } finally { + existsSpy.mockRestore(); + aliveSpy.mockRestore(); + } + }); + + it("aborts stale stop cleanup when the stop is rescinded before the relaunch lands", async () => { + const worker = { + descriptor: { + workerId: "worker-rescinded-during-stop", + pid: 111_120, + rootActiveSessionId: "active-1", stopRequestedAt: new Date().toISOString() as string | undefined, archiveOnStop: true, }, @@ -2344,7 +3058,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 () => { @@ -2396,7 +3110,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 () => { @@ -2434,7 +3148,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 { @@ -2457,7 +3171,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), { @@ -2505,7 +3219,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), { @@ -2551,7 +3265,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 @@ -2578,7 +3292,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. @@ -2600,7 +3314,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 () => {}); @@ -2625,7 +3339,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 { @@ -2653,6 +3367,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 { @@ -2689,6 +3404,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 { @@ -2720,20 +3436,26 @@ 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, stopWorker, + persistWorkerStopTombstone: vi.fn(), log: vi.fn(), reportCleanupFailure: vi.fn(), }) as { @@ -2767,7 +3489,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]]), @@ -2809,7 +3531,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 () => { @@ -3401,6 +4123,297 @@ 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("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[] = []; @@ -3591,10 +4604,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 +4631,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/daemon-supervisor-process.test.ts b/packages/coding-agent/test/daemon-supervisor-process.test.ts index 5d973738df..346d83d126 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"; @@ -15,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"; @@ -198,7 +200,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 +558,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"); @@ -508,7 +691,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, @@ -534,7 +721,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( @@ -548,6 +747,163 @@ 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("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"); @@ -591,6 +947,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 new file mode 100644 index 0000000000..51940da6fb --- /dev/null +++ b/packages/coding-agent/test/daemon-version-compatibility.test.ts @@ -0,0 +1,323 @@ +import { type ChildProcess, spawn, spawnSync } from "node:child_process"; +import { randomUUID } from "node:crypto"; +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"; +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 workerPids = 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 waitForSupervisorExit(child: ChildProcess, timeoutMs: number): Promise { + if (child.exitCode !== null || child.signalCode !== null) return; + await new Promise((resolveExit, reject) => { + 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); +} + +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[]; +} + +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 }); +}); + +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"); + 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( + (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 waitForSupervisorExit(replacement, 10_000); + children.delete(replacement); + await waitForProcessExit(createdSummary.workerPid, 10_000); + expect(workerDescriptorFiles(agentDir)).toEqual([]); + workerPids.delete(createdSummary.workerPid); + }, + 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 ebb6116f3a..32f9a61973 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,136 @@ 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("]}")]; + 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..4b0ceda6ed 100644 --- a/packages/coding-agent/test/snapshot-transcript-cache.test.ts +++ b/packages/coding-agent/test/snapshot-transcript-cache.test.ts @@ -1,9 +1,20 @@ -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, + SNAPSHOT_MEMORY_CACHE_BYTES, + SnapshotTranscriptCache, + type SnapshotTranscriptCacheIo, + sweepAbandonedSnapshotCacheRoots, +} from "../src/modes/daemon/snapshot-transcript-cache.js"; const tempDirs: string[] = []; @@ -116,6 +127,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 +225,187 @@ 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]); + 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/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 38423943a7..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 @@ -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"; @@ -12,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"; @@ -113,6 +119,98 @@ 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("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?.(); + }, 60_000); + it.each([1, 2])( "keeps %i same-ID worker stream(s) bounded by socket drain without shared cleanup", async (count) => { @@ -222,16 +320,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), @@ -356,6 +453,7 @@ describe("ENG-4601 worker snapshot cache", () => { activeSessionId, snapshotId, error: `Snapshot ${snapshotId} was aborted`, + purpose: "attach", }); expect(written).toContainEqual( expect.objectContaining({ @@ -378,6 +476,143 @@ 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 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 deb503d15c..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 @@ -16,14 +16,18 @@ 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"; 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; @@ -63,7 +67,7 @@ function summary(): SessionSummary { }; } -function streamedResult(messages: AgentMessage[]): DaemonAttachResult { +function streamedResult(messages: AgentMessage[], resultSnapshotId = snapshotId): DaemonAttachResult { return { protocol: DAEMON_PROTOCOL_INFO, activeSessionId, @@ -82,7 +86,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"] }, }; } @@ -115,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(), @@ -128,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, @@ -183,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 = { @@ -231,6 +245,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 }, session: { sessionId: "session-4602" } }, + } 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,12 +392,132 @@ 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(); + registerWorker(supervisor, worker); + 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(); + registerWorker(supervisor, worker); + 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(); + registerWorker(supervisor, worker); + 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" }, 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", @@ -333,12 +553,49 @@ 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(); + registerWorker(supervisor, worker); const client = socketClient("public", new PassThrough()); const streamSnapshot = vi.fn(async () => {}); const internals = supervisor as unknown as { @@ -346,6 +603,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 +656,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 +683,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 () => { @@ -423,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 { @@ -473,11 +759,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 () => { @@ -486,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(); @@ -547,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; @@ -574,7 +864,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,36 +889,190 @@ 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(); + registerWorker(supervisor, completed.worker); 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(); + registerWorker(supervisor, mismatchedEnd.worker); 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(); + 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(); 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(); + registerWorker(supervisor, worker); + 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[0]).toMatchObject({ type: "session_snapshot_begin", purpose: "replacement" }); + expect(records[1]).toMatchObject({ + type: "session_snapshot_failed", + activeSessionId, + 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(); + }); + + 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(); + registerWorker(supervisor, worker); + 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(); + registerWorker(supervisor, worker); + 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", @@ -653,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/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) }; 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..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 @@ -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"; @@ -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,7 +151,178 @@ 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(); + 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); + registerWorker(supervisor, worker); + 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"), { + 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, + }), + ); + registerWorker(supervisor, worker); + 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; + }; + + 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); + 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("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 }), + ); + registerWorker(supervisor, worker); + 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"), { @@ -164,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"); @@ -281,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) => { @@ -377,11 +554,12 @@ 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); const worker = workerHarness(firstResult, firstTranscript); + registerWorker(supervisor, worker); const internals = supervisor as unknown as { handleWorkerFrame(worker: WorkerHarness, frame: PrivateFrame): void; }; @@ -422,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"); @@ -599,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; @@ -700,14 +880,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(); });