From 2e4de30bd61265fec33c11573af29b07b0d6d57c Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Tue, 22 Sep 2026 05:48:16 -0600 Subject: [PATCH 01/26] docs(pylon): record post-v0.9.4 upstream triage and dispositions --- .pylon/features.yaml | 189 ++++++++++++++++++++++++++++++++++++++ .pylon/upstream-review.md | 38 ++++++++ 2 files changed, 227 insertions(+) diff --git a/.pylon/features.yaml b/.pylon/features.yaml index 2ba7d3ef67d..c1f093a17cc 100644 --- a/.pylon/features.yaml +++ b/.pylon/features.yaml @@ -527,3 +527,192 @@ decisions: upstream_support: Current upstream allocates a second timestamp when persisting custom messages, so reloading or compacting changes live transcript identities. Preserve the original message timestamp at every AgentSession persistence boundary. revisit_when: - Upstream preserves exact custom-message identities across live emission, persistence and reconstruction. + + kernel-stderr-stream-errors: + area: kernel + state: candidate + owner: upstream + decision: adopt + pylon_refs: + - https://github.com/pylon-code/prime-agent/issues/1 + - https://github.com/pylon-code/pylon/issues/114 + upstream_refs: + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2383 + fork_change: absorb-kernel-child-stderr-stream-errors + upstream_support: Absorb child stderr stream errors to prevent worker crashes on broken pipe or process death without altering process ownership. + revisit_when: + - Upstream redesigns kernel stdio streams. + + repl-bounded-frames: + area: kernel + state: candidate + owner: upstream + decision: adopt + pylon_refs: + - https://github.com/pylon-code/prime-agent/issues/1 + upstream_refs: + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2423 + fork_change: bound-repl-protocol-frame-sizes + upstream_support: Bound REPL message sizes to prevent memory exhaustion and protocol desync. + revisit_when: + - Kernel REPL migrates to chunked framing. + + owner-only-kernel-stderr-log: + area: kernel-security + state: candidate + owner: upstream + decision: adopt + pylon_refs: + - https://github.com/pylon-code/prime-agent/issues/1 + upstream_refs: + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2425 + fork_change: owner-only-kernel-stderr-log-permissions + upstream_support: Set mode 0o600 on kernel stderr log files. + revisit_when: + - Kernel logging subsystem is replaced. + + heartbeats-list-bounded-live-state: + area: daemon-health + state: candidate + owner: upstream + decision: adopt + pylon_refs: + - https://github.com/pylon-code/prime-agent/issues/1 + upstream_refs: + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2378 + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2342 + fork_change: heartbeats-list-from-live-state-with-deadline + upstream_support: Read active heartbeats from in-memory state with 10s deadline without queueing behind busy worker turns; must not fabricate liveness for quarantined sessions. + revisit_when: + - Heartbeat storage moves out-of-process. + + session-open-restart-wait: + area: daemon-lifecycle + state: candidate + owner: upstream + decision: adopt + pylon_refs: + - https://github.com/pylon-code/prime-agent/issues/1 + upstream_refs: + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2391 + fork_change: session-open-wait-during-daemon-restart + upstream_support: Let session opens wait through daemon update restart instead of failing immediately; re-verifies daemon hello and supervisor generation on connect. + revisit_when: + - Daemon lifecycle management changes. + + spawn-name-reservation-admission: + area: session-lifecycle + state: candidate + owner: upstream + decision: adopt + pylon_refs: + - https://github.com/pylon-code/prime-agent/issues/1 + upstream_refs: + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2396 + fork_change: hold-spawn-name-reservations-until-admission + upstream_support: Hold spawn name reservations until admission is durable, respecting caller-owned admission proofs and reconnect handles. + revisit_when: + - Session admission protocol is redesigned. + + async-bash-notice-turn-boundary: + area: runtime-reliability + state: candidate + owner: shared + decision: hybridize + pylon_refs: + - https://github.com/pylon-code/prime-agent/issues/64 + - https://github.com/pylon-code/prime-agent/issues/1 + upstream_refs: + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2372 + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2386 + fork_change: async-bash-notice-re-parking-with-handle-withdrawal + upstream_support: Hybridize turn-boundary notice fixes with fork #72 per-handle withdrawal and single terminal settlement. + revisit_when: + - Upstream introduces handle-bound async command receipts. + + compaction-tool-identity-and-anchoring: + area: context-compaction + state: candidate + owner: upstream + decision: adopt + pylon_refs: + - https://github.com/pylon-code/prime-agent/issues/1 + upstream_refs: + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2424 + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2385 + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2463 + fork_change: compaction-tool-identity-anchoring-and-validation + upstream_support: Preserve tool identity in summarizer input, anchor summaries to kept tail, validate refinement writes, and preserve custom message timestamps without altering cursor order. + revisit_when: + - Compaction architecture or prompt format changes. + + codex-stale-chain-recovery: + area: provider-codex + state: candidate + owner: upstream + decision: adopt + pylon_refs: + - https://github.com/pylon-code/prime-agent/issues/1 + upstream_refs: + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2374 + fork_change: codex-stale-chain-recovery-after-metadata + upstream_support: Recover stale Codex chains after metadata while preserving connection binding. + revisit_when: + - Codex completion API changes. + + prime-inference-glm-no-thinking: + area: provider-glm + state: candidate + owner: upstream + decision: adopt + pylon_refs: + - https://github.com/pylon-code/prime-agent/issues/1 + upstream_refs: + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2459 + fork_change: prime-inference-glm-disable-thinking + upstream_support: Stop sending enable_thinking to Prime Inference GLM models. + revisit_when: + - GLM adds reasoning support. + + abort-and-send-queued: + area: session-input + state: candidate + owner: shared + decision: hybridize + pylon_refs: + - https://github.com/pylon-code/pylon/issues/114 + upstream_refs: + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2426 + fork_change: capability-gated-abort-and-send-queued-schema34 + upstream_support: Renumber schema revision above fork 33 to 34 and gate behind abort_and_send_queued_v1 capability. + revisit_when: + - Upstream negotiates queued aborts via capability. + + quota-park-and-auto-resume: + area: provider-quota + state: deferred + owner: shared + decision: redesign + pylon_refs: + - https://github.com/pylon-code/pylon/issues/114 + upstream_refs: + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2375 + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2280 + fork_change: none + upstream_support: Quota waiting and auto-resume daemon behavior. Pending maintainer decision. + revisit_when: + - Maintainer guides quota park and auto-resume projection. + + mcp-service-catalog-oauth: + area: mcp-ownership + state: deferred + owner: shared + decision: retain + pylon_refs: + - https://github.com/pylon-code/pylon/issues/114 + upstream_refs: + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2330 + fork_change: none + upstream_support: Daemon MCP service catalog and OAuth engine. Retaining Pylon MCP ownership pending maintainer decision. + revisit_when: + - Maintainer guidance on MCP service catalog. diff --git a/.pylon/upstream-review.md b/.pylon/upstream-review.md index 9357423e67f..a66c4814b45 100644 --- a/.pylon/upstream-review.md +++ b/.pylon/upstream-review.md @@ -389,3 +389,41 @@ The SDK token gates the new local API; daemon wire shapes and stock behavior do not change. Validation: focused malformed/private/path/registration cases and a real isolated supervisor/owned-worker test prove registered before cleanup and settled after exact owned cleanup and supervisor exit. Pylon delivery tracked in [pylon#557](https://github.com/pylon-code/pylon/issues/557). + +## 2026-09-22 — post-v0.9.4 upstream review (cfb2d3fe3..18a56bf35 and open Sept 14 triage) + +- Review range: upstream commits `cfb2d3fe35c7c23cf6678ec235412b0d3f286fae..18a56bf3557063c0cd0be12f99f0bcce2a46da32` (72 commits, including release v0.9.5 `a7d791bc1`) plus open Sept 14 triage items. +- Merge-base and complete integration checkpoint remain frozen at v0.9.4 `1eee2938b4eeb7a4d72e17035adda669a89b63de`. Selective adoption does not advance the full integration checkpoint. +- Fork product base: `72899e4e4fdaef52cfffca58152f769a6aef91b1`. Tracking: #1 and Pylon #114. + +### Dispositions and Revisit Triggers + +| PR / Commit | Subject / Area | Disposition | Notes & Revisit Triggers | +| --- | --- | --- | --- | +| #2383 (`b6b94c78a`) | Kernel child stderr stream error absorption | **Adopt** | Prevents worker crashes on unhandled child stderr errors (EPIPE). Does not touch process ownership, isolation, or cleanup receipts. Revisit if upstream redesigns kernel stdio streams. | +| #2423 (`c91e6e991`) | Bound REPL protocol frame sizes in kernel and host | **Adopt** | Guards against memory exhaustion and protocol desync by bounding message frames. Revisit if kernel REPL switches to streaming chunked wire format. | +| #2425 (`b6d955a9b`) | Keep kernel stderr log owner-only (mode 0o600) | **Adopt** | Security fix ensuring private process logs are not world-readable. Revisit if kernel logging subsystem is replaced. | +| #2378 (`3dfb0276b`) + #2342 (`ad426c672`) | Heartbeats list from live state with 10s deadline | **Adopt** | Reads active heartbeats from in-memory state without blocking behind serialized session worker queues. Must not fabricate liveness for terminating or quarantined sessions. Revisit if heartbeat storage moves out-of-process. | +| #2391 (`54da82e97`) | Session opens wait through daemon update restart | **Adopt** | Retries and waits for restart completion instead of immediately throwing connection errors. Wait must re-verify daemon hello, schema revision, and supervisor generation upon reconnect to prevent attaching across unauthorized generation boundaries. Revisit if daemon process manager lifecycle changes. | +| #2396 (`126fe7010`) | Hold spawn name reservations until admission is durable | **Adopt** | Prevents concurrent duplicate session creation collision in the registry. Must respect caller-owned admission proofs and not reject reconnects or recovery handles under the same session ID. Revisit if session creation protocol is redesigned. | +| #2372 (`66df06fc1`) + #2386 (`b3e04b58b`) | Async-bash completion notices and re-parking | **Hybridize** | Overlaps fork #72 (acknowledged per-handle background notice withdrawal). Preserves fork #72 random handle identities, host-acknowledged withdrawal, and single terminal turn settlement. Revisit if upstream introduces handle-bound async receipts. | +| #2424 (`5fb3d9acb`) | Compaction: keep tool identity in serialized summarizer input | **Adopt** | Retains tool name and call parameters in summarizer prompt. Preserves custom message timestamps (#76) and does not alter event cursor ordering. Revisit if compaction serialization schema changes. | +| #2385 (`27f32ddb7`) | Compaction: anchor summaries to kept-tail state and stop re-summarizing file lists | **Adopt** | Fixes repeated file listings and anchors summaries to recent state. Does not alter transcript disk format or break replay verification. Revisit if compaction prompt architecture is modified. | +| #2463 (`41e4e0e18`) | Validate refinement writes and skip malformed entries in digest | **Adopt** | Bounds digest errors on corrupted refinement records. Logs malformed entries without corrupting the digest hash or causing divergence between replicas. Revisit if digest format evolves. | +| #2374 (`e683fcb53`) | AI: recover stale Codex chains after metadata | **Adopt** | Follow-up to adopted #2314. Retains producing connection binding and one-shot retry before output. Revisit if Codex completion API changes. | +| #2459 (`63d88319b`) | AI: stop sending enable_thinking to Prime Inference GLM models | **Adopt** | Fixes provider rejection for GLM models. Revisit if Prime Inference GLM adds reasoning support. | +| Perf: #2414, #2416, #2433, #2389, #2393, #2399, #2379, #2387, #2405 | In-memory caching, frame decoding, kernel startup perf | **Adopt** | Pure performance improvements that do not mutate snapshot/catalog ownership boundaries or wire protocols. #2433 must not alter catalog disk schema, serialized format, or bypass immutable snapshot boundaries. Revisit if underlying subsystems are overhauled. | +| Excluded: #2398, #2296, #2273, #2260, #2242, #2246, #2276 | On-demand catalog, transcript download cache, persistent SupervisorLink, orphan GC, session append | **Retain (Deferred)** | These touch snapshot/catalog/supervisor/ownership boundaries. Fork preserves immutable bounded snapshots, attachment generations, per-event recovery cursors, authenticated ownership, and cleanup receipts. Revisit in separate review with real-daemon and managed-artifact integration gates. | +| #2426 (`e2fb7bfa1`) | Send queued messages after interrupt (`abort_and_send_queued`) | **Hybridize / Redesign** | Daemon wire change: renumber schema above fork 33 to schema 34 and gate behind `abort_and_send_queued_v1` capability. Pylon negotiates capability and falls back to interrupt-then-send. Revisit when upstream capability-gates queued aborts. | +| #2282 (`ebe530d28`) | Child progress notes (`rlm_progress_note`) and child snapshots (`rlm.collect`) | **Redesign (Pylon issue)** | Stock 0.9.5 events are safely ignored in Pylon via `Ignored/unknown-event`. Pylon issue tracks Agents surface integration. Revisit when Agents UI exposes child progress. | +| #2453 (`3b1aa5ff3`) | Image-model routing | **Redesign (Pylon issue)** | Dynamic model switching impacts model identity, usage attribution, and pricing. Pylon issue tracks provider routing design. Revisit with multimodal routing architecture. | +| #2455 (`b09e37d84`) | ACP model and effort pickers | **Adopt (Pylon issue)** | Exposes model/effort options for ACP fallback. Pylon issue tracks UI pickers. Revisit during ACP settings review. | +| #2334 (`dca77ecfe`) + #2465 (`b08f08efa`) | Human message priority and continuation holds while bash runs | **Redesign (Pylon issue)** | Evaluated against Pylon queue projection and occurrence ownership. Pylon issue tracks queue prioritization semantics. Revisit with turn scheduling audit. | +| Stock 0.9.5 compiled binary without Node SDK | Runtime detection and fallback | **Retain (Pylon issue)** | Stock 0.9.5 installer defaults to Bun-compiled binary omitting Node SDK. Pylon issue tracks fallback to ACP with actionable error guidance. | +| #2375 (`2883a7843`) + #2280 | Quota park and auto-resume (`providerBackupModel`) | **Maintainer Decision** | Daemon parks quota-blocked sessions until reset and auto-resumes. Options: disable for Pylon-owned sessions, or project unsolicited resumed turns as first-class runtime state. Pending maintainer decision. | +| #2330 (`b6ac5d014`) | MCP service catalog and OAuth engine | **Maintainer Decision** | Overlaps Pylon-owned MCP server management. Retain Pylon MCP ownership pending maintainer confirmation. | +| Context/Refinement: #2217, #2226, #2220, #2394, #2400, #2450, #2392 | Context summarization and harness ranking | **Adopt / Hybridize** | Improves compaction accuracy and prompt-cache stability. #2226 already hybridized in fork #71. Revisit if context architecture changes. | +| Models/CLI: #2505, #2462, #2381, #2410, #2409, #2216, #2380, #2515, #2336 | Catalog updates, CLI diagnostics, test repairs | **Adopt** | Additive model options (Grok 4.7), diagnostics (`PI_REQUEST_TIMING`), flag validations, and test flakiness fixes. Revisit when affected providers change. | +| TUI / CLI UI / CI: #2446, #2447, #2430, #2466, #2340, #2331, #2328, #2324, #2325, #2301, #2302, #2306, #2449, #2376, #2456, #2323 | Terminal UI interactions, upstream CI/benchmarks, install docs | **Skip** | Prime-specific TUI screens, terminal commands, and upstream internal CI are not consumed by Pylon web/desktop/mobile clients. No daemon protocol or SDK impact. Revisit if features move to daemon protocol. | +| #2458 (`2e9ab77b8`) | Reconnect attached windows when daemon restarts | **Skip / Retain** | TUI client window reconnect is skipped because Pylon manages its own reconnection; supervisor update-restart phase closing reason is retained/deferred with supervisor lifecycle audit. | +| Distribution: #2366, #2370, #2369, #2368, #2361, #2344, #2319 | Native binaries, updater manifests, musl packages | **Retain** | Retain Pylon recipe 2 deterministic SDK archives per fork #75 decision. Revisit if upstream provides importable Node SDK in binary archives. | From 19e3940f8c7b10e352ab47733b4a86f3fbdf2763 Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:44:26 -0400 Subject: [PATCH 02/26] [RSI, bug] fix(coding-agent): absorb kernel child stderr stream errors like stdin/stdout (#2383) * fix(coding-agent): absorb kernel child stderr stream errors like stdin/stdout wireChild guards stdin and stdout pipe streams with error listeners, but the stderr stream only had data/end/close handlers, so an error emitted on it reached the daemon uncaughtException handler, which process.exit(1)s the worker and every hosted session with it. Add the matching stderr error listener that records a kernel diagnostic, and extend the pipe-errors regression test to emit on stderr and assert the kernel keeps running. * test(coding-agent): trim the duplicated pipe-error comment to fit the test-line budget The added comment block restated the guard the test name and the new stderr assertions already carry; one line keeps the read-side rationale. Net test lines drop to +2 against 3 added source lines. (cherry picked from commit b6b94c78a7747997e0ec6c690c0e8f2a563dbe14) --- packages/coding-agent/.changes/kernel-stderr-pipe-error.md | 1 + packages/coding-agent/src/core/kernel/repl-manager.ts | 7 ++++++- packages/coding-agent/test/repl-kernel-pipe-errors.test.ts | 4 +++- 3 files changed, 10 insertions(+), 2 deletions(-) create mode 100644 packages/coding-agent/.changes/kernel-stderr-pipe-error.md diff --git a/packages/coding-agent/.changes/kernel-stderr-pipe-error.md b/packages/coding-agent/.changes/kernel-stderr-pipe-error.md new file mode 100644 index 00000000000..514dee4849b --- /dev/null +++ b/packages/coding-agent/.changes/kernel-stderr-pipe-error.md @@ -0,0 +1 @@ +- Fixed an unhandled error on the kernel child stderr pipe crashing the daemon worker: the stream now records a kernel diagnostic like stdin and stdout instead of reaching the daemon uncaughtException handler and exiting with every hosted session. diff --git a/packages/coding-agent/src/core/kernel/repl-manager.ts b/packages/coding-agent/src/core/kernel/repl-manager.ts index 31a1408f30b..5159524700a 100644 --- a/packages/coding-agent/src/core/kernel/repl-manager.ts +++ b/packages/coding-agent/src/core/kernel/repl-manager.ts @@ -436,7 +436,8 @@ export class ReplKernelManager { // stream (write EPIPE); without a listener Node rethrows it and takes // down the whole worker. The pending writeLine rejection and the child // 'exit' handler below own the fallout, so this only records the - // diagnosis. + // diagnosis. Read-side pipe errors on stdout/stderr land the same way + // and need the same guard. child.stdin?.on("error", (error) => { if (this.child !== child) return; this.appendKernelDiagnostic(`kernel stdin error: ${errorMessage(error)}`); @@ -445,6 +446,10 @@ export class ReplKernelManager { if (this.child !== child) return; this.appendKernelDiagnostic(`kernel stdout error: ${errorMessage(error)}`); }); + child.stderr?.on("error", (error) => { + if (this.child !== child) return; + this.appendKernelDiagnostic(`kernel stderr error: ${errorMessage(error)}`); + }); child.once("exit", () => { // One turn for the poll phase to deliver the bytes the kernel wrote // before dying (the pipe buffer bounds them), then destroy: EOF may diff --git a/packages/coding-agent/test/repl-kernel-pipe-errors.test.ts b/packages/coding-agent/test/repl-kernel-pipe-errors.test.ts index 46266cd42cb..d40026b3f74 100644 --- a/packages/coding-agent/test/repl-kernel-pipe-errors.test.ts +++ b/packages/coding-agent/test/repl-kernel-pipe-errors.test.ts @@ -47,11 +47,13 @@ describeIf("ReplKernelManager pipe errors (real runtime)", () => { const child = (manager as unknown as { child?: ChildProcess }).child; expect(child?.stdin).toBeDefined(); expect(child?.stdout).toBeDefined(); + expect(child?.stderr).toBeDefined(); // A write racing the kernel's death lands as an 'error' event on the pipe; // without a listener Node would crash the worker (observed in production as - // "uncaught exception: Error: write EPIPE"). + // "write EPIPE"). Read-side pipes need the same guard. expect(() => child?.stdin?.emit("error", new Error("write EPIPE"))).not.toThrow(); expect(() => child?.stdout?.emit("error", new Error("read ECONNRESET"))).not.toThrow(); + expect(() => child?.stderr?.emit("error", new Error("read ECONNRESET"))).not.toThrow(); // The kernel is still healthy: further cells keep executing and shutdown // (afterEach) still completes. From 8d2dcaa8314289591efa19cab7dfffd7f822dd7e Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Mon, 21 Sep 2026 09:43:32 -0400 Subject: [PATCH 03/26] [RSI, bug] fix(kernel): bound REPL protocol frame sizes in kernel and host (#2423) * [RSI, bug] fix(kernel): bound REPL protocol frame sizes in kernel and host The host buffers whole protocol lines before its per-execution truncation, so a print()/repr() of huge data from one kernel cell buffered the full payload twice in the host worker, and a newline-less flood from a corrupt runtime ballooned host memory forever because the protocol repair only ever saw complete lines. The kernel now chunks Python-level stdout/stderr writes into 64 Ki-char stream frames under a per-writer lock, caps result repr text at 1 Mi chars with a trailing truncation marker, and refuses display payloads whose JSON encoding exceeds 16 Mi characters (emit() raises ValueError in the cell). The host caps protocol line accumulation at 32 Mi chars, failing an oversized line through the existing protocol repair instead of buffering it, and stops draining a poisoned child's residue. Each regression test fails pre-fix: the chunked-write test saw one 300,000-char frame, the repr/display test saw the untruncated 3,000,000-char result, and the oversized unterminated protocol line hung the manager until vitest's test timeout. * docs(runtime): result text cap excludes the appended truncation marker * fix(kernel): mark output truncated when frames arrive after the buffer is exactly full A 64 KiB stream frame fills the 65536-char buffer exactly; later frames were dropped without setting stdoutTruncated/stderrTruncated, so the model saw no truncation marker. --------- Co-authored-by: prime-agent (cherry picked from commit c91e6e991fc006098ea21c3640eed3bdd268c222) --- .../.changes/bound-repl-protocol-frames.md | 1 + .../src/core/kernel/repl-manager.ts | 20 ++++++++++- .../repl-kernel-protocol-corruption.test.ts | 11 ++++++ prime-agent-runtime/src/rlm/repl.md | 12 +++++-- prime-agent-runtime/src/rlm/repl.py | 34 ++++++++++++++++--- prime-agent-runtime/test/test_repl.py | 15 ++++++++ 6 files changed, 85 insertions(+), 8 deletions(-) create mode 100644 packages/coding-agent/.changes/bound-repl-protocol-frames.md diff --git a/packages/coding-agent/.changes/bound-repl-protocol-frames.md b/packages/coding-agent/.changes/bound-repl-protocol-frames.md new file mode 100644 index 00000000000..f36de25766f --- /dev/null +++ b/packages/coding-agent/.changes/bound-repl-protocol-frames.md @@ -0,0 +1 @@ +- Fixed kernel REPL protocol output that could buffer unbounded memory in the host: Python-level stdout/stderr writes now ship as 64 Ki-char frames, oversized result reprs are capped at 1 Mi chars with a marker, oversized display payloads fail the cell, and the host repairs a kernel that streams an oversized protocol line. diff --git a/packages/coding-agent/src/core/kernel/repl-manager.ts b/packages/coding-agent/src/core/kernel/repl-manager.ts index 5159524700a..f7fb48ee80e 100644 --- a/packages/coding-agent/src/core/kernel/repl-manager.ts +++ b/packages/coding-agent/src/core/kernel/repl-manager.ts @@ -59,6 +59,10 @@ const REPAIR_STEP_TIMEOUT_MS = 30_000; const MAX_HANDLED_HOST_REQUEST_IDS = 1024; // Cap for unattributed background output buffered between and during cells. const MAX_BACKGROUND_OUTPUT_CHARS = 64 * 1024; +// Largest legit frame is an attachment display event, base64 capped at +// MAX_ATTACHMENT_DATA_CHARS; a line that cannot complete within this ceiling is +// corruption the protocol repair owns, not output worth buffering until OOM. +const MAX_PROTOCOL_LINE_CHARS = 32 * 1024 * 1024; const MAX_KERNEL_STDERR_CHARS = 8 * 1024; const MAX_KERNEL_STDERR_LOG_BYTES = 5 * 1024 * 1024; @@ -363,9 +367,18 @@ export class ReplKernelManager { private wireChild(child: ChildProcess): void { const decoder = new StringDecoder("utf8"); let buffered = ""; + // A poisoned child's residue must not grow the buffer again before the + // protocol repair kills it. + let poisoned = false; child.stdout?.on("data", (buf: Buffer) => { - if (this.child !== child) return; + if (this.child !== child || poisoned) return; buffered += decoder.write(buf); + if (buffered.length > MAX_PROTOCOL_LINE_CHARS) { + poisoned = true; + buffered = ""; + this.failProtocolFrame(child, `oversized protocol line: exceeds ${MAX_PROTOCOL_LINE_CHARS} chars`); + return; + } let newline = buffered.indexOf("\n"); while (newline !== -1) { if (this.child !== child) return; @@ -825,6 +838,9 @@ export class ReplKernelManager { execution.stdout = execution.stdout.slice(0, execution.maxChars); execution.stdoutTruncated = true; } + } else if (text.length > 0) { + // The buffer filled exactly on an earlier frame; the dropped remainder still counts as truncation. + execution.stdoutTruncated = true; } } else { if (execution.stderr.length < execution.maxChars) { @@ -833,6 +849,8 @@ export class ReplKernelManager { execution.stderr = execution.stderr.slice(0, execution.maxChars); execution.stderrTruncated = true; } + } else if (text.length > 0) { + execution.stderrTruncated = true; } } execution.opts.onStream?.(text, type); diff --git a/packages/coding-agent/test/repl-kernel-protocol-corruption.test.ts b/packages/coding-agent/test/repl-kernel-protocol-corruption.test.ts index 3e2ee779f31..89093ab2e28 100644 --- a/packages/coding-agent/test/repl-kernel-protocol-corruption.test.ts +++ b/packages/coding-agent/test/repl-kernel-protocol-corruption.test.ts @@ -63,6 +63,7 @@ input.on("line", (line) => { process.stdout.write("BROKEN-" + "x".repeat(400) + "\\n"); return; } + if (request.code === "corrupt-huge-line") return process.stdout.write("x".repeat(33 * 1024 * 1024)); if (request.code === "corrupt-idle") { process.stdout.write(JSON.stringify({ event: "done", id: request.id, status: "ok" }) + "\\n42\\n"); return; @@ -356,6 +357,16 @@ describe("ReplKernelManager corrupt protocol repair", () => { } }); + it("repairs after an oversized unterminated protocol line", async () => { + const { manager } = newManager(); + try { + await expect(manager.execute("corrupt-huge-line")).rejects.toThrow(/oversized protocol line/); + expect((await manager.execute("read")).status).toBe("ok"); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + it("repairs an object frame with an unknown event kind instead of hanging the request", async () => { const { manager, countPath } = newManager(); try { diff --git a/prime-agent-runtime/src/rlm/repl.md b/prime-agent-runtime/src/rlm/repl.md index c0cb0127754..122bda80fe0 100644 --- a/prime-agent-runtime/src/rlm/repl.md +++ b/prime-agent-runtime/src/rlm/repl.md @@ -49,12 +49,18 @@ runtime keeps serving. Closing stdin is equivalent to `shutdown`. tasks inherit the spawning cell's id (even after that cell finished). `null` for user threads, raw fd writes (`os.write`, C extensions, subprocesses), and anything else without provable ownership — bytes read from the fd pipes - are never attributed to a cell. + are never attributed to a cell. A Python-level write ships at most 64 Ki + characters per frame; a larger write arrives as multiple events in order. - `{"event":"result","id":str,"text":str}` — `repr` of the cell's trailing expression when the body ends in an expression whose value is not `None`. - The value is also bound to `_` in the namespace. + The value is also bound to `_` in the namespace. The `repr` content is capped + at 1,048,576 characters; a longer `repr` is truncated to the cap and a trailing + truncation marker is appended after it, so the total `text` can exceed the cap + by the marker's length. - `{"event":"display","id":str|null,"data":{mime:payload,...}}` — one dict of - MIME type to JSON payload, shipped verbatim from `emit()`. `id` rides task + MIME type to JSON payload, shipped verbatim from `emit()`. A payload whose + JSON encoding exceeds 16 Mi characters is refused: `emit()` raises + `ValueError` in the calling cell. `id` rides task context: an asyncio task spawned by a cell keeps that cell's id even after the cell finishes; user threads emit `null`. - `{"event":"host_request","id":str,"data":{...}}` — one typed request from diff --git a/prime-agent-runtime/src/rlm/repl.py b/prime-agent-runtime/src/rlm/repl.py index 413b2a83acf..0af5c834935 100644 --- a/prime-agent-runtime/src/rlm/repl.py +++ b/prime-agent-runtime/src/rlm/repl.py @@ -36,6 +36,17 @@ DEFAULT_SNAPSHOT_MAX_BYTES = 256 * 1024 * 1024 DEFAULT_SNAPSHOT_MAX_VARIABLE_BYTES = 16 * 1024 * 1024 +# Stream writes must fit one protocol frame: the host buffers whole lines +# before its per-execution truncation, and raw fd writes already arrive as +# 64 KiB pump chunks. +_STREAM_FRAME_TEXT_CAP = 64 * 1024 +# The host truncates results at a smaller per-execution maxChars, so this only +# bounds a pathological repr in transit. +_RESULT_TEXT_CAP = 1_048_576 +_RESULT_TRUNCATION_MARKER = f"\n[... result truncated at {_RESULT_TEXT_CAP} characters ...]" +# Oversized display payloads fail the cell instead of wedging host memory. +_DISPLAY_PAYLOAD_CAP = 16 * 1024 * 1024 + # Names the session bootstrap re-creates on every start; never snapshotted. _ALWAYS_SKIP = {"rlm", "mcp", "bash", "asyncio", "In", "Out", "get_ipython", "exit", "quit", "open"} # IPython-injected names that may appear in a snapshot payload; never restored. @@ -97,9 +108,11 @@ def emit(data: dict[str, Any]) -> None: # Strict-dumps validation: default allow_nan=True would let NaN/Infinity # serialize as non-JSON text and tear the host's protocol framing (a # non-serializable value already raises in _send before any bytes are - # written, so NaN is the only corruption vector). Payloads are small, so - # the throwaway serialization here is cheap; _send re-serializes. - json.dumps(data, allow_nan=False) + # written, so NaN is the only corruption vector). The encoded length + # enforces the display frame cap; _send re-serializes. + encoded = json.dumps(data, allow_nan=False) + if len(encoded) > _DISPLAY_PAYLOAD_CAP: + raise ValueError(f"display payload exceeds the {_DISPLAY_PAYLOAD_CAP}-character frame cap") _send({"event": "display", "id": _current_cell.get(), "data": data}) @@ -293,13 +306,24 @@ class _TaggedWriter(io.TextIOBase): def __init__(self, stream: str, fallback_fd: int) -> None: self._stream = stream self._fallback_fd = fallback_fd + # Keeps one write()'s frames contiguous under concurrent writers. + self._frame_lock = threading.Lock() self._buffer = _TaggedBuffer(fallback_fd) def write(self, text: str) -> int: if not isinstance(text, str): raise TypeError(f"write() argument must be str, not {type(text).__name__}") if text: - _send({"event": self._stream, "id": _current_cell.get(), "text": text}) + cell_id = _current_cell.get() + with self._frame_lock: + for start in range(0, len(text), _STREAM_FRAME_TEXT_CAP): + _send( + { + "event": self._stream, + "id": cell_id, + "text": text[start : start + _STREAM_FRAME_TEXT_CAP], + } + ) return len(text) def flush(self) -> None: @@ -584,6 +608,8 @@ async def _handle_execute(req: dict[str, Any], ns: dict[str, Any]) -> None: result_text = repr(value) except BaseException as exc: # noqa: BLE001 - a broken __repr__ is a cell error status, error = "error", _error_event(cell_id, exc) + if result_text is not None and len(result_text) > _RESULT_TEXT_CAP: + result_text = result_text[:_RESULT_TEXT_CAP] + _RESULT_TRUNCATION_MARKER _drain_output() finally: # Close the interrupt window before the protocol sends so a diff --git a/prime-agent-runtime/test/test_repl.py b/prime-agent-runtime/test/test_repl.py index 69ccf69fe09..7bc22a49beb 100644 --- a/prime-agent-runtime/test/test_repl.py +++ b/prime-agent-runtime/test/test_repl.py @@ -708,6 +708,21 @@ def test_emit_nan_payload_errors_in_cell_and_keeps_framing(self): self.assertEqual(one(follow, "result")["text"], "2") self.assertEqual(one(follow, "done")["status"], "ok") + def test_large_write_ships_bounded_stream_frames(self): + events = self.repl.execute("chunk", "import sys\nsys.stdout.write('x' * 300_000)") + frames = [e["text"] for e in events if e.get("event") == "stdout"] + self.assertEqual("".join(frames), "x" * 300_000) + self.assertTrue(all(len(frame) <= 65536 for frame in frames)) + + def test_oversized_repr_and_display_payloads_are_capped(self): + code = "class C:\n def __repr__(self):\n return 'x' * 3_000_000\nC()" + events = self.repl.execute("big-repr", code) + expected = "x" * 1_048_576 + "\n[... result truncated at 1048576 characters ...]" + self.assertEqual(one(events, "result")["text"], expected) + events = self.repl.execute("emit-big", "from rlm.repl import emit\nemit({'text/plain': 'x' * 17_000_000})") + self.assertEqual(one(events, "error")["ename"], "ValueError") + self.assertEqual(one(events, "done")["status"], "error") + def test_bash_integration(self): events = self.repl.execute( "sh1", "from rlm import bash\nresult = await bash('echo repl-bash')\nresult.output.strip()" From cd7b57789e0043e2dd17b8efeccbf1d14dac243c Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:44:04 -0400 Subject: [PATCH 04/26] [RSI, security] fix(kernel): keep the kernel stderr log owner-only (#2425) * [RSI, security] fix(kernel): keep the kernel stderr log owner-only The kernel stderr log (session-artifacts//kernel-stderr.log) was created world-readable (0644) while every sibling artifact in the same private directory is owner-only (0600 files, 0700 dirs). Kernel stderr can carry exception payloads and library warnings, so the log and its directory must match the other private session artifacts. Create both with owner-only bits: mode 0o700 for the directory and 0o600 for the file, with fchmodSync on the opened descriptor to enforce the exact bits despite the umask (and tighten a pre-existing loose log on next kernel start). * fix(kernel): tighten the rotated kernel stderr log to 0600 Rotating an oversized kernel stderr log renamed it to `.old` before any chmod, so a pre-existing world-readable (0644) log kept those bits and its historical exception payloads stayed readable by other users. Only the new current file was tightened. chmod the log to KERNEL_STDERR_LOG_MODE immediately before the rename. The path is known to exist (statSync just succeeded), a chmod failure lands in the existing rotation catch and keeps appending instead, and on Windows the chmod clears the read-only bit so the rename still proceeds. Cover it with a regression test that pre-creates an oversized 0644 log and asserts the rotated file is 0600. The failing-start tests now share fakeRuntime/withManager helpers to keep the added coverage line-neutral. (cherry picked from commit b6d955a9b0e61a100088ff13757a440470b38af3) --- .../.changes/kernel-stderr-log-owner-only.md | 1 + .../src/core/kernel/repl-manager.ts | 32 +++++++++++++++++-- .../test/repl-kernel-startup.test.ts | 23 ++++++++++++- 3 files changed, 52 insertions(+), 4 deletions(-) create mode 100644 packages/coding-agent/.changes/kernel-stderr-log-owner-only.md diff --git a/packages/coding-agent/.changes/kernel-stderr-log-owner-only.md b/packages/coding-agent/.changes/kernel-stderr-log-owner-only.md new file mode 100644 index 00000000000..442591b395e --- /dev/null +++ b/packages/coding-agent/.changes/kernel-stderr-log-owner-only.md @@ -0,0 +1 @@ +- Fixed the kernel stderr log (`kernel-stderr.log` in the session artifact directory) being created world-readable: the log is now owner-only (0600), and its directory is created owner-only (0700) when the kernel manager creates it, matching the kernel state snapshot and the other private session artifacts, because kernel stderr can carry exception payloads. diff --git a/packages/coding-agent/src/core/kernel/repl-manager.ts b/packages/coding-agent/src/core/kernel/repl-manager.ts index f7fb48ee80e..f1d008c4532 100644 --- a/packages/coding-agent/src/core/kernel/repl-manager.ts +++ b/packages/coding-agent/src/core/kernel/repl-manager.ts @@ -2,7 +2,18 @@ // (`python -m rlm.repl`) — requests on stdin, events on stdout, stderr kept as // a diagnostics tail. The protocol is documented in prime-agent-runtime/src/rlm/repl.md. import type { ChildProcess } from "node:child_process"; -import { closeSync, existsSync, mkdirSync, openSync, renameSync, rmSync, statSync, writeSync } from "node:fs"; +import { + chmodSync, + closeSync, + existsSync, + fchmodSync, + mkdirSync, + openSync, + renameSync, + rmSync, + statSync, + writeSync, +} from "node:fs"; import { dirname } from "node:path"; import { StringDecoder } from "node:string_decoder"; import { v4 as uuid } from "uuid"; @@ -67,6 +78,10 @@ const MAX_PROTOCOL_LINE_CHARS = 32 * 1024 * 1024; const MAX_KERNEL_STDERR_CHARS = 8 * 1024; const MAX_KERNEL_STDERR_LOG_BYTES = 5 * 1024 * 1024; const KERNEL_STDERR_LOG_BUDGET_MARKER = "[stderr log budget exhausted]\n"; +// Owner-only directory for the kernel stderr log, matching the other private session artifacts. +const KERNEL_STDERR_LOG_DIR_MODE = 0o700; +// Owner-only file bits; kernel stderr can carry exception payloads. +const KERNEL_STDERR_LOG_MODE = 0o600; /** fs.writeSync may write fewer bytes than asked (partial ENOSPC, signals); loop until done. */ function writeFullySync(fd: number, data: Buffer): void { @@ -251,10 +266,13 @@ export class ReplKernelManager { const path = this.options.stderrLogPath; if (!path) return undefined; try { - mkdirSync(dirname(path), { recursive: true }); + mkdirSync(dirname(path), { recursive: true, mode: KERNEL_STDERR_LOG_DIR_MODE }); let size = existsSync(path) ? statSync(path).size : 0; if (size > MAX_KERNEL_STDERR_LOG_BYTES) { try { + // Tighten before the move: a renamed log keeps its mode, and the + // rotated file holds the exception payloads worth protecting. + chmodSync(path, KERNEL_STDERR_LOG_MODE); // Drop any prior .old first: rename fails on Windows if it exists. rmSync(`${path}.old`, { force: true }); renameSync(path, `${path}.old`); @@ -264,7 +282,15 @@ export class ReplKernelManager { this.appendKernelDiagnostic(`cannot rotate kernel stderr log: ${errorMessage(error)}`); } } - return { fd: openSync(path, "a"), budget: Math.max(0, MAX_KERNEL_STDERR_LOG_BYTES - size) }; + const fd = openSync(path, "a", KERNEL_STDERR_LOG_MODE); + // Exact bits despite the umask; tightens a pre-existing loose log. + try { + fchmodSync(fd, KERNEL_STDERR_LOG_MODE); + } catch (error) { + closeSync(fd); + throw error; + } + return { fd, budget: Math.max(0, MAX_KERNEL_STDERR_LOG_BYTES - size) }; } catch (error) { this.appendKernelDiagnostic(`cannot open kernel stderr log: ${errorMessage(error)}`); return undefined; diff --git a/packages/coding-agent/test/repl-kernel-startup.test.ts b/packages/coding-agent/test/repl-kernel-startup.test.ts index 3c4707954c0..5a0dc1f69e5 100644 --- a/packages/coding-agent/test/repl-kernel-startup.test.ts +++ b/packages/coding-agent/test/repl-kernel-startup.test.ts @@ -127,7 +127,7 @@ describe("ReplKernelManager startup", () => { "", ].join("\n"), ); - const stderrLogPath = join(tempDir, "kernel-stderr.log"); + const stderrLogPath = join(tempDir, "artifacts", "kernel-stderr.log"); const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); const manager = new ReplKernelManager({ python, cwd: tempDir, stderrLogPath }); @@ -135,6 +135,27 @@ describe("ReplKernelManager startup", () => { // Well under the 30s ready timeout: teardown must not wait for the // grandchild (destroying the pipe kills it with SIGPIPE). await expect(manager.execute("print(1)")).rejects.toThrow(/Kernel exited before ready/); + expect(statSync(stderrLogPath).mode & 0o777).toBe(0o600); + expect(statSync(join(tempDir, "artifacts")).mode & 0o777).toBe(0o700); + } finally { + errorSpy.mockRestore(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }, 15000); + + it("tightens a rotated kernel stderr log that was world-readable", async () => { + const stderrLogPath = join(tempDir, "kernel-stderr.log"); + writeFileSync(stderrLogPath, Buffer.alloc(5 * 1024 * 1024 + 1), { mode: 0o644 }); + const python = join(tempDir, "python"); + writeExecutable(python, ["#!/bin/sh", "exit 42", ""].join("\n")); + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + const manager = new ReplKernelManager({ python, cwd: tempDir, stderrLogPath }); + + try { + await expect(manager.execute("print(1)")).rejects.toThrow(/Kernel exited before ready/); + // The rotated file holds the historical exception payloads. + expect(statSync(`${stderrLogPath}.old`).mode & 0o777).toBe(0o600); + expect(statSync(stderrLogPath).mode & 0o777).toBe(0o600); } finally { errorSpy.mockRestore(); await manager.shutdown({ snapshot: true, drainHostRequests: true }); From d990a30eb8eed366ae081713b5533cf7818cecb5 Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Tue, 15 Sep 2026 01:23:53 -0400 Subject: [PATCH 05/26] fix(coding-agent): bound heartbeat view catalog fetch with a 10s deadline (#2342) --- .../.changes/heartbeat-view-fetch-deadline.md | 1 + .../components/heartbeat-manager.ts | 11 +- .../src/modes/interactive/interactive-mode.ts | 162 ++++- .../interactive-heartbeat-management.test.ts | 556 +++++++++++++++++- 4 files changed, 712 insertions(+), 18 deletions(-) create mode 100644 packages/coding-agent/.changes/heartbeat-view-fetch-deadline.md diff --git a/packages/coding-agent/.changes/heartbeat-view-fetch-deadline.md b/packages/coding-agent/.changes/heartbeat-view-fetch-deadline.md new file mode 100644 index 00000000000..8fc2163e757 --- /dev/null +++ b/packages/coding-agent/.changes/heartbeat-view-fetch-deadline.md @@ -0,0 +1 @@ +- Fixed the TUI heartbeat view freezing on busy sessions: catalog fetches now keep the last snapshot after a 10-second deadline and retry on the next heartbeats_changed event, instead of waiting behind an active turn indefinitely. A timed-out fetch still applies its late answer and re-arms a short retry so the open view always converges, and a failed fetch now shows an error in the view instead of silently showing a stale catalog. diff --git a/packages/coding-agent/src/modes/interactive/components/heartbeat-manager.ts b/packages/coding-agent/src/modes/interactive/components/heartbeat-manager.ts index 28e9344ec7c..ab0dfbb2dbf 100644 --- a/packages/coding-agent/src/modes/interactive/components/heartbeat-manager.ts +++ b/packages/coding-agent/src/modes/interactive/components/heartbeat-manager.ts @@ -15,6 +15,8 @@ type HeartbeatManagerMode = { type: "list" } | { type: "actions"; heartbeatId: s export interface HeartbeatManagerOptions { getHeartbeats: () => readonly AgentConnectionHeartbeat[]; + /** Last catalog fetch failure to surface in the list view, if any. */ + getFetchError?: () => string | undefined; getRows: () => number; onAction: (heartbeat: AgentConnectionHeartbeat, action: AgentHeartbeatManagementAction) => Promise; onClose: () => void; @@ -111,6 +113,13 @@ export class HeartbeatManagerComponent implements Component, Focusable { const list = new MenuList({ compact: this.getListLayout().compact }); this.populateHeartbeatList(list); panel.addChild(list); + const fetchError = this.options.getFetchError?.(); + if (fetchError) { + panel.addChild(new Spacer(1)); + panel.addChild( + new TruncatedText(theme.fg("warning", `Heartbeat refresh failed: ${this.singleLine(fetchError)}`)), + ); + } if (this.error) { panel.addChild(new Spacer(1)); panel.addChild(new TruncatedText(theme.fg("error", `Error: ${this.error}`))); @@ -270,7 +279,7 @@ export class HeartbeatManagerComponent implements Component, Focusable { getRows: this.options.getRows, preferredVisibleItems: PREFERRED_VISIBLE_HEARTBEATS, totalItems: this.heartbeats.length, - reservedRows: HEARTBEAT_LIST_RESERVED_ROWS + (this.error ? 2 : 0), + reservedRows: HEARTBEAT_LIST_RESERVED_ROWS + (this.error || this.options.getFetchError?.() ? 2 : 0), comfortableItemRows: 3, compactItemRows: 2, scrollIndicatorRows: HEARTBEAT_SCROLL_INDICATOR_ROWS, diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index c9b69ac794c..9beb32362b4 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -589,6 +589,39 @@ const DEAD_TERMINAL_ERROR_CODES = new Set(["EIO", "EPIPE", "ENOTCONN"]); // evicted past the cap to keep a long session bounded. const MAX_PASTED_IMAGE_BYTES = 64 * 1024 * 1024; const INITIAL_TRANSCRIPT_RENDER_MESSAGE_LIMIT = 400; +// Coalesce at most this many heartbeats_changed refreshes into one refresh +// promise; sustained churn must not hold rebindCurrentSession in a drain loop. +const HEARTBEAT_REFRESH_DRAIN_LIMIT = 25; +// A busy daemon can serialize a catalog fetch behind session work or worker +// recovery for minutes; the TUI waits at most this long before keeping the +// last catalog and letting the next heartbeats_changed event retry. +export const HEARTBEAT_REFRESH_FETCH_TIMEOUT_MS = 10_000; +// After a deadline expiry, retry the open manager's refresh on this short +// cadence so the view always has a next scheduled retrieval even when no +// heartbeats_changed event arrives and the catalog has no nextRunAt to +// schedule from. +export const HEARTBEAT_REFRESH_RETRY_DELAY_MS = 5_000; + +/** Race a catalog fetch against the refresh deadline; resolves undefined on expiry. */ +async function fetchHeartbeatsWithinDeadline( + fetch: Promise, +): Promise { + // The deadline can settle first; a late failure must not surface as an + // unhandled rejection after this refresh has moved on. + void fetch.catch(() => undefined); + let timeout: ReturnType | undefined; + try { + return await Promise.race([ + fetch, + new Promise((resolve) => { + timeout = setTimeout(() => resolve(undefined), HEARTBEAT_REFRESH_FETCH_TIMEOUT_MS); + timeout.unref?.(); + }), + ]); + } finally { + clearTimeout(timeout); + } +} function initialRenderMessages(messages: AgentMessage[]): AgentMessage[] { if (messages.length <= INITIAL_TRANSCRIPT_RENDER_MESSAGE_LIMIT) { @@ -965,6 +998,7 @@ export class InteractiveMode { private readonly retainedSubmissionGenerations = new WeakMap(); private admitPendingStartupPrompts: (() => Promise) | undefined; private agentsViewRequest: InteractiveModeRunResult["type"] | undefined; + private isReturningToAgentsView = false; private loadingAnimation: Loader | undefined = undefined; private workingMessage: string | undefined = undefined; private workingVisible = true; @@ -1073,6 +1107,15 @@ export class InteractiveMode { private heartbeatManagerHandle: OverlayHandle | undefined; private heartbeatManagerRefreshTimer: ReturnType | undefined; private heartbeatManagerRefreshAt: number | undefined; + private heartbeatManagerFetchRetryTimer: ReturnType | undefined; + // Last catalog fetch failure surfaced in the open manager view; cleared by + // the next successful apply and on manager open. + private heartbeatCatalogFetchError: string | undefined; + // Monotonic issue-order sequence for catalog writes (fetches and heartbeat + // management actions). A fetch write sequenced behind an already-applied + // newer one is dropped, so a late/older snapshot cannot revert newer state. + private heartbeatCatalogSeq = 0; + private heartbeatCatalogAppliedSeq = 0; // Registry of images pasted this session, keyed by the `[image #N]` marker // shown to the user. Insertion-ordered; the bytes persist (bounded by @@ -2619,23 +2662,52 @@ export class InteractiveMode { } const connection = this.agentConnection; const refresh = (async () => { - do { + // Bound the drain loop: sustained heartbeats_changed churn must not + // starve an awaiting rebind (and its transcript render) indefinitely. + for (let drain = 0; drain <= HEARTBEAT_REFRESH_DRAIN_LIMIT; drain++) { this.heartbeatRefreshRequested = false; - const heartbeats = await connection.listHeartbeats(); - if (this.agentConnection !== connection) return; - this.applyHeartbeatCatalog(heartbeats); - } while (this.heartbeatRefreshRequested); + const fetchSeq = ++this.heartbeatCatalogSeq; + const fetch = connection.listHeartbeats(); + const heartbeats = await fetchHeartbeatsWithinDeadline(fetch); + if (this.isShuttingDown || this.isReturningToAgentsView || this.agentConnection !== connection) return; + if (heartbeats === undefined) { + // The fetch deadline expired: keep the last catalog, retain the + // in-flight result so the open view still converges when the + // late answer lands, and arm a bounded retry so the open view + // always has a next scheduled retrieval even if it never does. + // A change event that arrived mid-fetch still re-triggers a + // refresh through the follow-up scheduling below. + this.retainHeartbeatCatalogFetch(fetch, connection, fetchSeq); + this.armHeartbeatManagerFetchRetry(); + return; + } + this.applyHeartbeatCatalog(heartbeats, fetchSeq); + if (!this.heartbeatRefreshRequested) return; + } })().finally(() => { if (this.heartbeatRefreshPromise === refresh) { this.heartbeatRefreshPromise = undefined; + // The drain bound was hit with a newer heartbeats_changed + // pending: schedule the follow-up refresh so the catalog + // converges instead of staying stale until the next event. + if (this.heartbeatRefreshRequested) { + void this.refreshHeartbeatCatalog().catch(() => undefined); + } } }); this.heartbeatRefreshPromise = refresh; return refresh; } - private applyHeartbeatCatalog(heartbeats: AgentConnectionHeartbeat[]): void { + private applyHeartbeatCatalog(heartbeats: AgentConnectionHeartbeat[], seq?: number): void { + // A sequenced fetch write that lost the race to a newer catalog write + // (a later fetch or a management action) must not revert it. + if (seq !== undefined) { + if (seq < this.heartbeatCatalogAppliedSeq) return; + this.heartbeatCatalogAppliedSeq = seq; + } this.heartbeatCatalog = heartbeats; + this.heartbeatCatalogFetchError = undefined; this.scheduleHeartbeatManagerRefresh(); this.updateSubagentSummaryLine(); this.ui.requestRender(); @@ -4234,7 +4306,7 @@ export class InteractiveMode { this.defaultEditor.onAction("app.thinking.toggle", () => this.toggleThinkingBlockVisibility()); this.defaultEditor.onAction("app.subagents.focus", () => this.focusSubagentSummary()); this.defaultEditor.onAction("app.heartbeats.open", () => { - void this.showHeartbeatManager(); + this.showHeartbeatManager(); }); this.defaultEditor.onAction("app.editor.external", () => this.openExternalEditor()); this.defaultEditor.onAction("app.prompt.stash", () => this.handlePromptStash()); @@ -6931,7 +7003,8 @@ export class InteractiveMode { } private async returnToAgentsView(request: InteractiveModeRunResult["type"] = "agents_view"): Promise { - if (this.isShuttingDown || this.agentsViewRequest) return; + if (this.isShuttingDown || this.isReturningToAgentsView || this.agentsViewRequest) return; + this.isReturningToAgentsView = true; this.stashDraftForAgentsView(); this.agentsViewRequest = request; this.isShuttingDown = true; @@ -9711,19 +9784,22 @@ export class InteractiveMode { } } - private async showHeartbeatManager(): Promise { + private showHeartbeatManager(): void { if (this.heartbeatManagerHandle) { this.heartbeatManagerHandle.focus(); return; } - try { - await this.refreshHeartbeatCatalog(); - } catch (error) { - this.showError(error instanceof Error ? error.message : String(error)); - return; - } + // Stale-while-revalidate: open immediately with the cached catalog. A + // busy session can take minutes to answer the fetch, and the overlay + // must not wait for it; the background refresh updates the open view + // when fresh data lands (the component reads the catalog per render). + // A failed fetch surfaces inside the view instead of silently showing a + // possibly-stale catalog; each open starts with a clean notice state. + this.heartbeatCatalogFetchError = undefined; + void this.refreshHeartbeatCatalog().catch((error) => this.onHeartbeatCatalogFetchError(error)); const manager = new HeartbeatManagerComponent({ getHeartbeats: () => this.getScopedHeartbeats(), + getFetchError: () => this.heartbeatCatalogFetchError, getRows: () => this.ui.terminal.rows, onAction: (heartbeat, action) => this.manageHeartbeat(heartbeat, action), onClose: () => this.closeHeartbeatManager(), @@ -9739,6 +9815,7 @@ export class InteractiveMode { private closeHeartbeatManager(): void { this.clearHeartbeatManagerRefreshTimer(); + this.clearHeartbeatManagerFetchRetryTimer(); this.heartbeatManagerHandle?.hide(); this.heartbeatManagerHandle = undefined; this.heartbeatManager = undefined; @@ -9793,10 +9870,64 @@ export class InteractiveMode { this.heartbeatManagerRefreshAt = undefined; } + private retainHeartbeatCatalogFetch( + fetch: Promise, + connection: AgentConnection, + seq: number, + ): void { + // The deadline expired while this fetch was still in flight: apply the + // late result when it lands so the open view converges even without a + // further heartbeats_changed event. Bounded by the same context guards + // as the drain loop and by the write sequence, so a superseded result + // cannot revert newer catalog state; a late failure is swallowed (the + // retry chain owns recovery from it). + void fetch + .then((heartbeats) => { + if (this.isShuttingDown || this.isReturningToAgentsView || this.agentConnection !== connection) { + return; + } + this.applyHeartbeatCatalog(heartbeats, seq); + }) + .catch(() => undefined); + } + + private armHeartbeatManagerFetchRetry(): void { + // Guarantee the open manager always has a next scheduled retrieval + // after a deadline expiry; a still-armed retry keeps the earliest slot. + if (!this.heartbeatManager || this.heartbeatManagerFetchRetryTimer) return; + this.heartbeatManagerFetchRetryTimer = setTimeout(() => { + this.heartbeatManagerFetchRetryTimer = undefined; + if (!this.heartbeatManager) return; + void this.refreshHeartbeatCatalog().catch((error) => this.onHeartbeatCatalogFetchError(error)); + }, HEARTBEAT_REFRESH_RETRY_DELAY_MS); + this.heartbeatManagerFetchRetryTimer.unref?.(); + } + + private clearHeartbeatManagerFetchRetryTimer(): void { + if (this.heartbeatManagerFetchRetryTimer) { + clearTimeout(this.heartbeatManagerFetchRetryTimer); + this.heartbeatManagerFetchRetryTimer = undefined; + } + } + + private onHeartbeatCatalogFetchError(error: unknown): void { + // Stale-while-revalidate opens the manager even when the fetch fails; + // surface the failure in the view instead of silently showing a stale + // catalog, and keep a next retrieval scheduled. + this.heartbeatCatalogFetchError = error instanceof Error ? error.message : String(error); + this.armHeartbeatManagerFetchRetry(); + this.scheduleHeartbeatManagerRefresh(); + this.ui.requestRender(); + } + private async manageHeartbeat( heartbeat: AgentConnectionHeartbeat, action: AgentHeartbeatManagementAction, ): Promise { + // Sequence the action by issue order so an older in-flight fetch + // cannot later revert its result; the action itself is authoritative + // and always applies. + const actionSeq = ++this.heartbeatCatalogSeq; const updated = await this.agentConnection.manageHeartbeat( heartbeat.job.activeSessionId, heartbeat.job.id, @@ -9805,6 +9936,7 @@ export class InteractiveMode { if (updated.source === "heartbeat" && updated.activeSessionId === this.connectionState?.activeSessionId) { this.patchConnectionState({ heartbeat: action === "stop" ? null : updated }); } + this.heartbeatCatalogAppliedSeq = Math.max(this.heartbeatCatalogAppliedSeq, actionSeq); const remaining = this.heartbeatCatalog.filter((entry) => entry.job.id !== updated.id); this.applyHeartbeatCatalog( updated.status === "active" || updated.status === "paused" diff --git a/packages/coding-agent/test/interactive-heartbeat-management.test.ts b/packages/coding-agent/test/interactive-heartbeat-management.test.ts index 664779e00e4..d891331bd89 100644 --- a/packages/coding-agent/test/interactive-heartbeat-management.test.ts +++ b/packages/coding-agent/test/interactive-heartbeat-management.test.ts @@ -1,10 +1,19 @@ -import { describe, expect, it, vi } from "vitest"; +import { setKeybindings } from "@earendil-works/pi-tui"; +import stripAnsi from "strip-ansi"; +import { beforeAll, describe, expect, it, vi } from "vitest"; import type { AgentCronJob, AgentHeartbeatManagementAction } from "../src/core/cron-jobs.js"; +import { KeybindingsManager } from "../src/core/keybindings.js"; import type { AgentConnectionHeartbeat, AgentConnectionRlmChildAgentSnapshot, } from "../src/modes/agent-connection/types.js"; -import { InteractiveMode } from "../src/modes/interactive/interactive-mode.js"; +import type { HeartbeatManagerComponent } from "../src/modes/interactive/components/heartbeat-manager.js"; +import { + HEARTBEAT_REFRESH_FETCH_TIMEOUT_MS, + HEARTBEAT_REFRESH_RETRY_DELAY_MS, + InteractiveMode, +} from "../src/modes/interactive/interactive-mode.js"; +import { initTheme } from "../src/modes/interactive/theme/theme.js"; interface HeartbeatManagementHarness { heartbeatCatalog: AgentConnectionHeartbeat[]; @@ -246,3 +255,546 @@ describe("interactive heartbeat management", () => { } }); }); + +interface HeartbeatManagerOpenHarness { + heartbeatManager: HeartbeatManagerComponent | undefined; + heartbeatManagerHandle: { focus(): void; hide(): void } | undefined; + heartbeatCatalog: AgentConnectionHeartbeat[]; + heartbeatRefreshPromise: Promise | undefined; + heartbeatRefreshRequested: boolean; + heartbeatManagerFetchRetryTimer: ReturnType | undefined; + heartbeatManagerRefreshTimer: ReturnType | undefined; + heartbeatManagerRefreshAt: number | undefined; + heartbeatCatalogSeq: number; + heartbeatCatalogAppliedSeq: number; + connectionState: { activeSessionId: string; sessionId: string }; + subagentSnapshots: Map; + agentConnection: { + listHeartbeats(): Promise; + manageHeartbeat?( + activeSessionId: string, + jobId: string, + action: AgentHeartbeatManagementAction, + ): Promise; + }; + isShuttingDown: boolean; + isReturningToAgentsView: boolean; + ui: { + requestRender(): void; + terminal: { rows: number }; + showOverlay(): { focus(): void; hide(): void }; + }; + patchConnectionState(patch: { heartbeat: AgentCronJob | null }): void; + scheduleHeartbeatManagerRefresh(): void; + updateSubagentSummaryLine(): void; + refreshHeartbeatCatalog(): Promise; + showHeartbeatManager(): void; + manageHeartbeat(heartbeat: AgentConnectionHeartbeat, action: AgentHeartbeatManagementAction): Promise; +} + +interface HeartbeatCommandHarness { + defaultEditor: { onSubmit?: (text: string) => Promise }; + editor: { getText(): string; setText(text: string): void }; + agentConnection: { prompt(text: string): Promise }; + showError(message: string): void; + heartbeatManagerHandle: { focus(): void; hide(): void } | undefined; + heartbeatCatalog: AgentConnectionHeartbeat[]; + refreshHeartbeatCatalog(): Promise; + refreshHeartbeatCatalogCalled?: boolean; + ui: { + requestRender(): void; + terminal: { rows: number }; + showOverlay(): { focus(): void; hide(): void }; + }; + [key: string]: unknown; +} + +const setupEditorSubmitHandlerPrototype = ( + InteractiveMode.prototype as unknown as { + setupEditorSubmitHandler(this: unknown): void; + } +).setupEditorSubmitHandler; + +describe("interactive heartbeat manager open (stale-while-revalidate)", () => { + beforeAll(() => { + initTheme("dark"); + setKeybindings(new KeybindingsManager()); + }); + + const overlayHandle = { focus: vi.fn(), hide: vi.fn() }; + + function makeOpenHarness(options?: { + listHeartbeats?: () => Promise; + /** Keep the real nextRunAt poll scheduling instead of a mock. */ + realSchedule?: boolean; + }): { + harness: HeartbeatManagerOpenHarness; + requestRender: ReturnType; + } { + const requestRender = vi.fn(); + const harness = Object.create(InteractiveMode.prototype) as HeartbeatManagerOpenHarness; + harness.heartbeatManager = undefined; + harness.heartbeatManagerHandle = undefined; + harness.heartbeatCatalog = [{ job: heartbeat() }]; + harness.heartbeatRefreshPromise = undefined; + harness.heartbeatRefreshRequested = false; + harness.heartbeatManagerFetchRetryTimer = undefined; + harness.heartbeatManagerRefreshTimer = undefined; + harness.heartbeatManagerRefreshAt = undefined; + // Prototype-only harnesses skip the constructor, so sequence fields + // must be seeded explicitly. + harness.heartbeatCatalogSeq = 0; + harness.heartbeatCatalogAppliedSeq = 0; + harness.connectionState = { activeSessionId: "active-1", sessionId: "session-1" }; + harness.subagentSnapshots = new Map([ + [ + "child-1", + { + id: "child-1", + activeSessionId: "active-2", + label: "child", + status: "running", + sessionDir: "/tmp/child-1", + }, + ], + ]); + harness.agentConnection = { + listHeartbeats: options?.listHeartbeats ?? vi.fn(() => new Promise(() => {})), + }; + harness.isShuttingDown = false; + harness.isReturningToAgentsView = false; + harness.ui = { + requestRender, + terminal: { rows: 24 }, + showOverlay: vi.fn(() => overlayHandle), + }; + harness.patchConnectionState = vi.fn(); + if (options?.realSchedule) { + harness.scheduleHeartbeatManagerRefresh = ( + InteractiveMode.prototype as unknown as { + scheduleHeartbeatManagerRefresh(this: HeartbeatManagerOpenHarness): void; + } + ).scheduleHeartbeatManagerRefresh; + } else { + harness.scheduleHeartbeatManagerRefresh = vi.fn(); + } + harness.updateSubagentSummaryLine = vi.fn(); + return { harness, requestRender }; + } + + function renderedManager(manager: NonNullable): string { + return stripAnsi(manager.render(100).join("\n")); + } + + it("opens immediately with the cached catalog while the fetch hangs, then updates when data lands", async () => { + let resolveFetch: (heartbeats: AgentConnectionHeartbeat[]) => void = () => {}; + const listHeartbeats = vi.fn( + () => + new Promise((resolve) => { + resolveFetch = resolve; + }), + ); + const { harness, requestRender } = makeOpenHarness({ listHeartbeats }); + + harness.showHeartbeatManager(); + + // The overlay must open synchronously with the cached catalog, before the fetch settles. + expect(listHeartbeats).toHaveBeenCalledOnce(); + expect(harness.heartbeatManagerHandle).toBe(overlayHandle); + expect(harness.heartbeatManager).toBeDefined(); + expect(renderedManager(harness.heartbeatManager!)).toContain("1 heartbeat."); + + // Input keeps working while the fetch is still in flight. + harness.heartbeatManager!.handleInput("\x1b[B"); + expect(requestRender).toHaveBeenCalled(); + + // Fresh data lands: the still-open view updates through the live catalog getter. + resolveFetch([ + { job: heartbeat() }, + { + job: heartbeat({ + id: "heartbeat-2", + activeSessionId: "active-2", + sessionId: "session-2", + sessionFile: "/tmp/session-2.jsonl", + status: "paused", + nextRunAt: undefined, + }), + }, + ]); + const refreshed = Promise.race([ + harness.heartbeatRefreshPromise ?? Promise.resolve(), + new Promise((_, reject) => setTimeout(() => reject(new Error("refresh deadline")), 1_000)), + ]); + await expect(refreshed).resolves.toBeUndefined(); + expect(renderedManager(harness.heartbeatManager!)).toContain("2 heartbeats · 1 paused"); + }); + + it("keeps the stale catalog usable when the fetch deadline expires", { timeout: 10_000 }, async () => { + vi.useFakeTimers(); + try { + const { harness } = makeOpenHarness({ + listHeartbeats: () => new Promise(() => {}), + }); + + harness.showHeartbeatManager(); + expect(harness.heartbeatManager).toBeDefined(); + const refresh = harness.heartbeatRefreshPromise; + expect(refresh).toBeDefined(); + + await vi.advanceTimersByTimeAsync(HEARTBEAT_REFRESH_FETCH_TIMEOUT_MS + 10); + + // The deadline settles the refresh quietly: no rejection, no stale overwrite. + await expect(refresh).resolves.toBeUndefined(); + expect(harness.heartbeatCatalog).toEqual([{ job: heartbeat() }]); + expect(harness.updateSubagentSummaryLine).not.toHaveBeenCalled(); + expect(harness.heartbeatRefreshPromise).toBeUndefined(); + + // The open view stays usable with the stale catalog. + expect(renderedManager(harness.heartbeatManager!)).toContain("1 heartbeat."); + harness.heartbeatManager!.handleInput("\x1b[B"); + expect(harness.ui.requestRender).toHaveBeenCalled(); + } finally { + vi.useRealTimers(); + } + }); + + it("retries after the deadline when heartbeats_changed arrived mid-fetch", { timeout: 10_000 }, async () => { + vi.useFakeTimers(); + try { + const fresh: AgentConnectionHeartbeat[] = [ + { job: heartbeat() }, + { + job: heartbeat({ + id: "heartbeat-2", + activeSessionId: "active-2", + sessionId: "session-2", + sessionFile: "/tmp/session-2.jsonl", + status: "paused", + nextRunAt: undefined, + }), + }, + ]; + const listHeartbeats = vi.fn(() => new Promise(() => {})); + listHeartbeats.mockImplementationOnce(() => new Promise(() => {})); + listHeartbeats.mockImplementationOnce(() => Promise.resolve(fresh)); + const { harness } = makeOpenHarness({ listHeartbeats }); + + const first = harness.refreshHeartbeatCatalog(); + // A heartbeats_changed event coalesces into the in-flight refresh. + const joined = harness.refreshHeartbeatCatalog(); + await vi.advanceTimersByTimeAsync(HEARTBEAT_REFRESH_FETCH_TIMEOUT_MS + 10); + await expect(first).resolves.toBeUndefined(); + await expect(joined).resolves.toBeUndefined(); + + // The drained follow-up refresh fetches again and converges. + await vi.advanceTimersByTimeAsync(10); + expect(listHeartbeats).toHaveBeenCalledTimes(2); + expect(harness.heartbeatCatalog).toEqual(fresh); + } finally { + vi.useRealTimers(); + } + }); + + it("retries a timed-out initial fetch on a short cadence so the open view converges", { + timeout: 10_000, + }, async () => { + vi.useFakeTimers(); + try { + // No heartbeats_changed event arrives and the empty catalog has no + // nextRunAt for the poll to schedule from: only the retry chain can + // guarantee the open view a next retrieval. + const fresh: AgentConnectionHeartbeat[] = [{ job: heartbeat() }]; + const listHeartbeats = vi.fn(() => Promise.resolve(fresh)); + listHeartbeats.mockImplementationOnce(() => new Promise(() => {})); + const { harness } = makeOpenHarness({ listHeartbeats }); + harness.heartbeatCatalog = []; + + harness.showHeartbeatManager(); + expect(renderedManager(harness.heartbeatManager!)).toContain("No running or paused heartbeats"); + + // The deadline expires quietly: the stale catalog stays, but the + // view must not be left without a next scheduled retrieval. + await vi.advanceTimersByTimeAsync(HEARTBEAT_REFRESH_FETCH_TIMEOUT_MS + 10); + expect(listHeartbeats).toHaveBeenCalledTimes(1); + expect(harness.heartbeatManagerFetchRetryTimer).toBeDefined(); + expect(harness.heartbeatCatalog).toEqual([]); + + // The bounded retry fires and converges the open view. + await vi.advanceTimersByTimeAsync(HEARTBEAT_REFRESH_RETRY_DELAY_MS + 10); + expect(listHeartbeats).toHaveBeenCalledTimes(2); + expect(harness.heartbeatCatalog).toEqual(fresh); + expect(renderedManager(harness.heartbeatManager!)).toContain("1 heartbeat."); + expect(harness.updateSubagentSummaryLine).toHaveBeenCalled(); + } finally { + vi.useRealTimers(); + } + }); + + it("applies the late result to the open view after the fetch deadline expired", { timeout: 10_000 }, async () => { + vi.useFakeTimers(); + try { + let resolveFetch: (heartbeats: AgentConnectionHeartbeat[]) => void = () => {}; + const listHeartbeats = vi.fn( + () => + new Promise((resolve) => { + resolveFetch = resolve; + }), + ); + const { harness } = makeOpenHarness({ listHeartbeats }); + + harness.showHeartbeatManager(); + const refresh = harness.heartbeatRefreshPromise; + expect(refresh).toBeDefined(); + + await vi.advanceTimersByTimeAsync(HEARTBEAT_REFRESH_FETCH_TIMEOUT_MS + 10); + await expect(refresh).resolves.toBeUndefined(); + // The deadline kept the stale catalog... + expect(harness.heartbeatCatalog).toEqual([{ job: heartbeat() }]); + expect(harness.updateSubagentSummaryLine).not.toHaveBeenCalled(); + + // ...but the retained fetch must still land in the open view. + const fresh: AgentConnectionHeartbeat[] = [ + { job: heartbeat() }, + { + job: heartbeat({ + id: "heartbeat-2", + activeSessionId: "active-2", + sessionId: "session-2", + sessionFile: "/tmp/session-2.jsonl", + status: "paused", + nextRunAt: undefined, + }), + }, + ]; + resolveFetch(fresh); + await vi.advanceTimersByTimeAsync(10); + await Promise.resolve(); + + expect(listHeartbeats).toHaveBeenCalledTimes(1); + expect(harness.heartbeatCatalog).toEqual(fresh); + expect(harness.updateSubagentSummaryLine).toHaveBeenCalled(); + expect(renderedManager(harness.heartbeatManager!)).toContain("2 heartbeats · 1 paused"); + } finally { + vi.useRealTimers(); + } + }); + + it("opens with a visible failure state when the initial fetch rejects, then recovers on the retry", { + timeout: 10_000, + }, async () => { + vi.useFakeTimers(); + try { + const fresh: AgentConnectionHeartbeat[] = [{ job: heartbeat() }]; + const listHeartbeats = vi.fn(() => Promise.resolve(fresh)); + listHeartbeats.mockImplementationOnce(() => Promise.reject(new Error("worker recovering"))); + const { harness } = makeOpenHarness({ listHeartbeats }); + + harness.showHeartbeatManager(); + // The manager must still open without waiting for the fetch. + expect(harness.heartbeatManagerHandle).toBe(overlayHandle); + await vi.advanceTimersByTimeAsync(10); + + // The rejection surfaces in the view instead of a silent stale catalog. + const rendered = renderedManager(harness.heartbeatManager!); + expect(rendered).toContain("Heartbeat refresh failed: worker recovering"); + expect(rendered).toContain("1 heartbeat."); + // A failed fetch still leaves the open view a next scheduled retrieval. + expect(harness.heartbeatManagerFetchRetryTimer).toBeDefined(); + + // The retry recovers: fresh data lands and the failure notice clears. + await vi.advanceTimersByTimeAsync(HEARTBEAT_REFRESH_RETRY_DELAY_MS + 10); + expect(listHeartbeats).toHaveBeenCalledTimes(2); + expect(harness.heartbeatCatalog).toEqual(fresh); + const recovered = renderedManager(harness.heartbeatManager!); + expect(recovered).not.toContain("Heartbeat refresh failed"); + expect(recovered).toContain("1 heartbeat."); + } finally { + vi.useRealTimers(); + } + }); + + it("keeps a next retrieval scheduled after a timed-out poll", { timeout: 10_000 }, async () => { + vi.useFakeTimers(); + try { + // Overdue heartbeat: the nextRunAt poll falls back to the 5s cadence. + vi.setSystemTime(new Date("2026-01-01T00:10:00.000Z")); + const listHeartbeats = vi.fn(() => new Promise(() => {})); + const { harness } = makeOpenHarness({ listHeartbeats, realSchedule: true }); + + harness.showHeartbeatManager(); + expect(listHeartbeats).toHaveBeenCalledTimes(1); + + // The overdue poll fires and joins the in-flight initial fetch. + await vi.advanceTimersByTimeAsync(5_250); + expect(listHeartbeats).toHaveBeenCalledTimes(1); + + // That fetch times out: the poll is spent and nothing applied, yet + // the retry chain keeps a next retrieval armed and the follow-up + // refresh issued. The view must never go quiet from here. + await vi.advanceTimersByTimeAsync(4_760); + expect(listHeartbeats).toHaveBeenCalledTimes(2); + expect(harness.heartbeatManagerFetchRetryTimer).toBeDefined(); + + await vi.advanceTimersByTimeAsync(10_010); + expect(listHeartbeats).toHaveBeenCalledTimes(3); + await vi.advanceTimersByTimeAsync(10_010); + expect(listHeartbeats).toHaveBeenCalledTimes(4); + await vi.advanceTimersByTimeAsync(10_010); + expect(listHeartbeats).toHaveBeenCalledTimes(5); + + // All those fetches timed out: the view kept its stale catalog and + // still never lacked a next scheduled retrieval. + expect(harness.heartbeatCatalog).toEqual([{ job: heartbeat() }]); + expect(harness.updateSubagentSummaryLine).not.toHaveBeenCalled(); + } finally { + vi.useRealTimers(); + } + }); + + it("does not overwrite a newer catalog with a late timed-out fetch result", { timeout: 10_000 }, async () => { + vi.useFakeTimers(); + try { + const resolvers: Array<(heartbeats: AgentConnectionHeartbeat[]) => void> = []; + const listHeartbeats = vi.fn( + () => + new Promise((resolve) => { + resolvers.push(resolve); + }), + ); + const { harness } = makeOpenHarness({ listHeartbeats }); + + harness.showHeartbeatManager(); + // The initial fetch times out and is retained. + await vi.advanceTimersByTimeAsync(HEARTBEAT_REFRESH_FETCH_TIMEOUT_MS + 10); + + // The retry fetch lands first with fresh data. + await vi.advanceTimersByTimeAsync(HEARTBEAT_REFRESH_RETRY_DELAY_MS + 10); + expect(resolvers.length).toBe(2); + const fresh: AgentConnectionHeartbeat[] = [{ job: heartbeat({ id: "heartbeat-9" }) }]; + resolvers[1]?.(fresh); + await vi.advanceTimersByTimeAsync(10); + await Promise.resolve(); + expect(harness.heartbeatCatalog).toEqual(fresh); + + // The older timed-out fetch resolves last: it must not overwrite. + const stale: AgentConnectionHeartbeat[] = [{ job: heartbeat() }, { job: heartbeat({ id: "heartbeat-8" }) }]; + resolvers[0]?.(stale); + await vi.advanceTimersByTimeAsync(10); + await Promise.resolve(); + + expect(harness.heartbeatCatalog).toEqual(fresh); + expect(harness.updateSubagentSummaryLine).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); + + it("does not let a late timed-out fetch revert a local heartbeat action", { timeout: 10_000 }, async () => { + vi.useFakeTimers(); + try { + const resolvers: Array<(heartbeats: AgentConnectionHeartbeat[]) => void> = []; + const listHeartbeats = vi.fn( + () => + new Promise((resolve) => { + resolvers.push(resolve); + }), + ); + const current = heartbeat(); + const paused = { ...current, status: "paused" as const, nextRunAt: undefined }; + const { harness } = makeOpenHarness({ listHeartbeats }); + harness.agentConnection = { + listHeartbeats, + manageHeartbeat: vi.fn(async () => paused), + }; + + harness.showHeartbeatManager(); + // The user pauses while the initial fetch is still in flight. + await harness.manageHeartbeat({ job: current }, "pause"); + expect(harness.heartbeatCatalog).toEqual([{ job: paused }]); + + // The initial fetch times out and is retained. + await vi.advanceTimersByTimeAsync(HEARTBEAT_REFRESH_FETCH_TIMEOUT_MS + 10); + + // The pre-pause snapshot lands late: it must not revert the pause. + resolvers[0]?.([{ job: current }]); + await vi.advanceTimersByTimeAsync(10); + await Promise.resolve(); + + expect(harness.heartbeatCatalog).toEqual([{ job: paused }]); + expect(renderedManager(harness.heartbeatManager!)).toContain("1 paused"); + } finally { + vi.useRealTimers(); + } + }); +}); + +describe("interactive /heartbeats command", () => { + beforeAll(() => { + initTheme("dark"); + setKeybindings(new KeybindingsManager()); + }); + + it("opens the manager without waiting for the catalog fetch", async () => { + const overlayHandle = { focus: vi.fn(), hide: vi.fn() }; + let editorText = ""; + const refreshHeartbeatCatalog = vi.fn(() => new Promise(() => {})); + const context = Object.create(InteractiveMode.prototype) as HeartbeatCommandHarness; + Object.assign(context, { + defaultEditor: {} as { onSubmit?: (text: string) => Promise }, + editor: { + getText: () => editorText, + setText: (text: string) => { + editorText = text; + }, + }, + agentConnection: { prompt: vi.fn(async () => undefined) }, + // Only settingsManager is read on the /heartbeats path (telemetry + // capture is fire-and-forget), so a minimal uiServices suffices. + uiServices: { settingsManager: {} }, + showError: vi.fn(), + showStatus: vi.fn(), + echoLocalCommand: vi.fn(), + submittedInputBehavior: "steer", + inputSubmissionGeneration: 0, + inputSubmissionsPending: 0, + pendingPromptStashReleases: [], + promptStashState: {}, + pendingSubmittedPromptStash: undefined, + snapshotPromptStash: vi.fn(() => ({ text: "" })), + promptStash: undefined, + promptStashSessionId: "session-1", + sessionId: "session-1", + clearShortcutGuide: vi.fn(), + // Support for the real showHeartbeatManager: a fetch that never settles. + heartbeatManager: undefined, + heartbeatManagerHandle: undefined, + heartbeatCatalog: [{ job: heartbeat() }], + heartbeatRefreshPromise: undefined, + heartbeatRefreshRequested: false, + connectionState: { activeSessionId: "active-1", sessionId: "session-1" }, + subagentSnapshots: new Map(), + refreshHeartbeatCatalog, + ui: { + requestRender: vi.fn(), + terminal: { rows: 24 }, + showOverlay: vi.fn(() => overlayHandle), + }, + scheduleHeartbeatManagerRefresh: vi.fn(), + updateSubagentSummaryLine: vi.fn(), + }); + + setupEditorSubmitHandlerPrototype.call(context); + const submitted = context.defaultEditor.onSubmit?.("/heartbeats"); + expect(submitted).toBeDefined(); + + // The submit handler must settle without the fetch: bound the wait explicitly. + const outcome = await Promise.race([ + submitted!.then(() => "settled" as const), + new Promise<"deadline">((resolve) => setTimeout(() => resolve("deadline"), 1_000)), + ]); + expect(outcome).toBe("settled"); + expect(context.heartbeatManagerHandle).toBe(overlayHandle); + expect(refreshHeartbeatCatalog).toHaveBeenCalledOnce(); + }); +}); From f13eb23cba189d6288ce229402f0ccccbc5485c5 Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:44:22 -0400 Subject: [PATCH 06/26] [RSI, bug] fix(daemon): answer heartbeats_list from live state without queueing behind serialized work (#2378) --- .../heartbeats-list-response-latency.md | 1 + .../src/modes/daemon/daemon-supervisor.ts | 156 +++++++++++++++++- .../src/modes/daemon/rlm-ledger.ts | 90 +++++----- .../heartbeats-list-response-latency.test.ts | 83 ++++++++++ 4 files changed, 283 insertions(+), 47 deletions(-) create mode 100644 packages/coding-agent/.changes/heartbeats-list-response-latency.md create mode 100644 packages/coding-agent/test/suite/regressions/heartbeats-list-response-latency.test.ts diff --git a/packages/coding-agent/.changes/heartbeats-list-response-latency.md b/packages/coding-agent/.changes/heartbeats-list-response-latency.md new file mode 100644 index 00000000000..2717be099ce --- /dev/null +++ b/packages/coding-agent/.changes/heartbeats-list-response-latency.md @@ -0,0 +1 @@ +- Fixed `heartbeats_list` timing out on busy daemons: the supervisor now answers from a shared in-memory scheduled-jobs snapshot (one in-flight disk scan serves every concurrent catalog request, heartbeat mutations drop the snapshot, aged snapshots refresh in the background), and sibling name scans no longer hold the serialized ledger queue that scheduled-catalog topology reads wait on. diff --git a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts index 4acd587f1f3..bfab643baa8 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts @@ -235,6 +235,18 @@ const ROSTER_STALE_AFTER_MS = 3 * ROSTER_HEARTBEAT_INTERVAL_MS; const PEER_TRANSPORT_GRANT_TTL_MS = 10_000; const WORKER_REQUEST_TIMEOUT_MS = 24 * 60 * 60 * 1000; const SNAPSHOT_ATTEMPT_TIMEOUT_MS = 30_000; +// The daemon client's default request budget is 30s and each ready-worker heartbeat +// forward gets 5s, so waiting longer than this on in-flight launches would only +// surface as a client transport timeout instead of a bounded per-worker state error. +export const HEARTBEAT_LIST_LAUNCH_WAIT_MS = 15_000; +// Session-scoped listing must fail inside the client's request budget too; the +// worker request default (24h) would turn a stuck worker into a client transport +// timeout instead of a daemon-side failure. +export const HEARTBEAT_LIST_FORWARD_TIMEOUT_MS = 25_000; +// A served scheduled-jobs snapshot older than this refreshes in the background +// (stale-while-revalidate) so external artifact writes converge without ever +// blocking a catalog response behind a fresh saved-session disk scan. +const PASSIVE_SCHEDULED_JOBS_REFRESH_MS = 5_000; const INPUT_PAUSE_CLEANUP_TIMEOUT_MS = 5_000; function createPublicSnapshotTransferId(): string { @@ -606,6 +618,18 @@ interface WorkerMatch { summary: SessionSummary; } +/** A durable scheduled job whose owning session has no resident worker. */ +interface PassiveScheduledJob { + rootSessionFile: string; + job: AgentCronJob; + info: SessionInfo; +} + +/** Catalog rows keep the same status filter a fresh scan applies. */ +function activeScheduledJobs(rows: PassiveScheduledJob[]): PassiveScheduledJob[] { + return rows.filter((row) => row.job.status === "active" || row.job.status === "paused"); +} + interface WorkerAttachData { result: DaemonAttachResult; worker: ResidentWorker; @@ -1026,6 +1050,16 @@ export class DaemonSupervisor { private scheduledWakeRecompute?: Promise; private scheduledWakeRecomputeQueued = false; private readonly scheduledWakeFailures = new Map(); + /** Shared passive scheduled-jobs snapshot; undefined before the first scan and after each invalidation. */ + private passiveScheduledJobs?: { rows: PassiveScheduledJob[]; scannedAt: number }; + /** One in-flight scan shared by every concurrent catalog read. */ + private passiveScheduledJobsScan?: Promise; + /** + * Monotonic epoch claimed at each scan start and at each invalidation: only + * the newest claim may publish, so an older scan cannot overwrite a newer + * scan's rows and a scan that raced a mutation cannot republish stale rows. + */ + private passiveScheduledJobsEpoch = 0; constructor( private readonly socketPath: string, @@ -1216,10 +1250,109 @@ export class DaemonSupervisor { }); } - /** Durable truth: the ledger family (fork headers stripped) plus each session's scheduled-jobs artifact. */ - private async collectPassiveScheduledJobs( - includeInactive = false, - ): Promise> { + /** + * Durable truth: the ledger family (fork headers stripped) plus each session's + * scheduled-jobs artifact. Always scans fresh per caller: management + * decisions, wake recomputes, and the ephemeral-cancel retries they drive + * must see the latest disk state. Every completed scan opportunistically + * refreshes the snapshot that heartbeats_list serves, but only while it is + * still the newest scan: an older scan loses the publish race. + */ + private async collectPassiveScheduledJobs(includeInactive = false): Promise { + const epoch = this.claimPassiveScheduledJobsEpoch(); + const rows = await this.scanPassiveScheduledJobs(); + this.storePassiveScheduledJobs(epoch, rows); + return includeInactive ? rows : activeScheduledJobs(rows); + } + + /** + * Snapshot-served rows for heartbeats_list responses, where per-request + * freshness matters less than not queueing: with hundreds of saved sessions, + * N concurrent catalog requests used to enqueue N full scans and push the + * last response past the client's 30s transport timeout. Daemon-owned + * mutations drop the snapshot and detach the in-flight shared scan + * (broadcastHeartbeatsChanged and the saved-session delete/rename paths); + * the next list or the wake recompute armed by the same broadcast rescans, + * and concurrent cold lists share one in-flight scan. A served snapshot + * older than PASSIVE_SCHEDULED_JOBS_REFRESH_MS refreshes in the background + * without being dropped, so lists that arrive during the refresh keep + * serving bounded-stale rows and external artifact writes converge without + * blocking a response. + */ + private async catalogPassiveScheduledJobs(): Promise { + const cached = this.passiveScheduledJobs; + if (cached) { + if (Date.now() - cached.scannedAt >= PASSIVE_SCHEDULED_JOBS_REFRESH_MS && !this.shuttingDown) { + // Stale-while-revalidate: serve the snapshot now and refresh in the + // background without dropping it, so requests that arrive during the + // refresh keep answering from bounded-stale rows instead of queueing + // behind the scan. A failure only logs; the next read retries. + void this.sharedPassiveScheduledJobsScan().catch((error: unknown) => + this.log(`Passive scheduled-jobs refresh failed: ${String(error)}`), + ); + } + return activeScheduledJobs(cached.rows); + } + return activeScheduledJobs(await this.sharedPassiveScheduledJobsScan()); + } + + private sharedPassiveScheduledJobsScan(): Promise { + this.passiveScheduledJobsScan ??= this.startSharedPassiveScheduledJobsScan(); + return this.passiveScheduledJobsScan; + } + + /** + * Claims the publish epoch when the scan actually starts, so callers that + * join an existing in-flight scan keep its rows, while a caller that starts + * after an invalidation never adopts the detached pre-mutation scan. + */ + private startSharedPassiveScheduledJobsScan(): Promise { + const epoch = this.claimPassiveScheduledJobsEpoch(); + const scan = this.scanPassiveScheduledJobs() + .then((rows) => { + this.storePassiveScheduledJobs(epoch, rows); + return rows; + }) + .finally(() => { + // An invalidation may have detached this promise and let a newer + // scan take the slot; clear it only while it is still ours. + if (this.passiveScheduledJobsScan === scan) { + this.passiveScheduledJobsScan = undefined; + } + }); + return scan; + } + + /** + * Compare-and-swap publish: a scan may store only while it still owns the + * newest epoch, so an older scan cannot overwrite a newer snapshot and a + * scan that raced an invalidation cannot republish pre-mutation rows. + */ + private storePassiveScheduledJobs(epoch: number, rows: PassiveScheduledJob[]): void { + if (epoch === this.passiveScheduledJobsEpoch) { + this.passiveScheduledJobs = { rows, scannedAt: Date.now() }; + } + } + + private claimPassiveScheduledJobsEpoch(): number { + this.passiveScheduledJobsEpoch += 1; + return this.passiveScheduledJobsEpoch; + } + + /** + * Daemon-owned mutations invalidate: claim the epoch so in-flight scans can + * no longer publish, drop the cached snapshot, and detach the shared + * in-flight promise so a list that starts after the mutation rescans + * instead of awaiting pre-mutation rows. + */ + private invalidatePassiveScheduledJobs(): void { + this.claimPassiveScheduledJobsEpoch(); + this.passiveScheduledJobs = undefined; + this.passiveScheduledJobsScan = undefined; + } + + /** The uncached disk scan; all statuses load so one row set serves every caller's filter. */ + private async scanPassiveScheduledJobs(): Promise { const pendingCancelRoots = new Set(); for (const intent of this.collectEphemeralCancelIntents()) { const context = this.workerSessionArtifactContext(intent); @@ -1263,7 +1396,7 @@ export class DaemonSupervisor { } return current.path; }; - const results: Array<{ rootSessionFile: string; job: AgentCronJob; info: SessionInfo }> = []; + const results: PassiveScheduledJob[] = []; for (const [artifactSessionId, store] of storeBySessionId) { let jobs: AgentCronJob[]; try { @@ -1273,7 +1406,6 @@ export class DaemonSupervisor { continue; } for (const job of jobs) { - if (!includeInactive && job.status !== "active" && job.status !== "paused") continue; const info = infoBySessionId.get(job.sessionId); if (!info) continue; const rootSessionFile = uncoveredRootFor(info); @@ -3152,7 +3284,7 @@ export class DaemonSupervisor { } } // Passivated sessions keep their armed heartbeats; no worker can list them. - for (const { job, info } of await this.collectPassiveScheduledJobs()) { + for (const { job, info } of await this.catalogPassiveScheduledJobs()) { if (!isHeartbeatCronJob(job) || heartbeats.has(job.id)) continue; heartbeats.set(job.id, { job, @@ -3293,6 +3425,9 @@ export class DaemonSupervisor { if (entry) { this.writeRosterEntry({ ...entry, summary: { ...entry.summary, sessionName: target.name } }); } + // Heartbeat rows carry the owning session's display name; + // drop the snapshot so the next list shows the new one. + this.invalidatePassiveScheduledJobs(); return success(command.id, command.type); } const match = await this.findWorkerForClient(client, command.activeSessionId); @@ -3327,6 +3462,9 @@ export class DaemonSupervisor { if (result.ok && entry && this.roster().get(entry.agentId) === entry) { this.roster().delete(entry.agentId); } + // The deleted session's armed heartbeats must leave the catalog + // snapshot too; their rows would otherwise linger until the TTL. + if (result.ok) this.invalidatePassiveScheduledJobs(); return success(command.id, command.type, result); } break; @@ -9971,6 +10109,10 @@ export class DaemonSupervisor { } private broadcastHeartbeatsChanged(): void { + // Every daemon-owned scheduled-job mutation and worker residency change + // lands here: drop the shared snapshot so the next catalog read (or the + // recompute armed below) rescans instead of serving pre-mutation rows. + this.invalidatePassiveScheduledJobs(); this.scheduleScheduledSessionWakeRecompute(); for (const client of this.clients) { this.write(client, { type: "heartbeats_changed" }); diff --git a/packages/coding-agent/src/modes/daemon/rlm-ledger.ts b/packages/coding-agent/src/modes/daemon/rlm-ledger.ts index f8900ae23d1..3f614ccd615 100644 --- a/packages/coding-agent/src/modes/daemon/rlm-ledger.ts +++ b/packages/coding-agent/src/modes/daemon/rlm-ledger.ts @@ -391,45 +391,55 @@ export class RlmSpawnLedger { /** Same-parent rows for a child session path, including the child itself. */ siblings(sessionPath: string): Promise { - return this.enqueue(async () => { - const target = canonicalSessionPath(sessionPath); - const family = await this.familyUnlocked(); - const edges = [...this.replaySync().values()].filter((edge) => !edge.deleted); - const parentByChild = new Map( - edges.map((edge) => [canonicalSessionPath(edge.child), canonicalSessionPath(edge.parent)]), - ); - const parent = parentByChild.get(target); - if (parent !== undefined) { - const rows = family.filter((row) => parentByChild.get(canonicalSessionPath(row.path)) === parent); - // The target's edge can be reconciliation-dropped (parent file - // gone) while its own file still exists: fall back to presenting - // the survivor alone rather than an empty set the callers would - // read as "session not found". - if (!rows.some((row) => canonicalSessionPath(row.path) === target)) { - try { - if ((await stat(target)).isFile()) { - return [await this.sessionRow(target, 0, undefined, undefined)]; - } - } catch { - // fall through to the (possibly empty) sibling rows + // Capture ordered topology first; transcript metadata must not block + // spawn appends or scheduled-catalog topology reads behind a full + // saved-session scan (same split as family()). + return this.enqueue(async () => ({ + alive: await this.liveEdgesUnlocked(), + edges: [...this.replaySync().values()].filter((edge) => !edge.deleted), + })).then((topology) => this.siblingsUnlocked(sessionPath, topology)); + } + + private async siblingsUnlocked( + sessionPath: string, + topology: { alive: RlmLedgerEdge[]; edges: RlmLedgerEdge[] }, + ): Promise { + const target = canonicalSessionPath(sessionPath); + const family = await this.familyUnlocked(topology.alive); + const parentByChild = new Map( + topology.edges.map((edge) => [canonicalSessionPath(edge.child), canonicalSessionPath(edge.parent)]), + ); + const parent = parentByChild.get(target); + if (parent !== undefined) { + const rows = family.filter((row) => parentByChild.get(canonicalSessionPath(row.path)) === parent); + // The target's edge can be reconciliation-dropped (parent file + // gone) while its own file still exists: fall back to presenting + // the survivor alone rather than an empty set the callers would + // read as "session not found". + if (!rows.some((row) => canonicalSessionPath(row.path) === target)) { + try { + if ((await stat(target)).isFile()) { + return [await this.sessionRow(target, 0, undefined, undefined)]; } + } catch { + // fall through to the (possibly empty) sibling rows } - return rows; } - // Roots are siblings of the other roots. A session outside both the - // ledger and the sessions dir is presented alone (matching the - // registry-walking reader's behavior for parentless sessions). - const roots = family.filter((row) => row.rlmDepth === 0); - if (roots.some((row) => canonicalSessionPath(row.path) === target)) { - return roots; - } - try { - if (!(await stat(target)).isFile()) return []; - } catch { - return []; - } - return [await this.sessionRow(target, 0, undefined, undefined)]; - }); + return rows; + } + // Roots are siblings of the other roots. A session outside both the + // ledger and the sessions dir is presented alone (matching the + // registry-walking reader's behavior for parentless sessions). + const roots = family.filter((row) => row.rlmDepth === 0); + if (roots.some((row) => canonicalSessionPath(row.path) === target)) { + return roots; + } + try { + if (!(await stat(target)).isFile()) return []; + } catch { + return []; + } + return [await this.sessionRow(target, 0, undefined, undefined)]; } private enqueue(fn: () => Promise | T): Promise { @@ -513,12 +523,12 @@ export class RlmSpawnLedger { return alive; } - private async familyUnlocked(): Promise { + private async familyUnlocked(aliveEdges?: RlmLedgerEdge[]): Promise { // One replay, one stat snapshot: byChild comes from the same alive set that emits child rows, // so a child whose dead edge was reconciled away degrades to a root row instead of vanishing. - let alive: RlmLedgerEdge[] = await this.liveEdgesUnlocked( - [...this.replaySync().values()].filter((candidate) => !candidate.deleted), - ); + let alive: RlmLedgerEdge[] = + aliveEdges ?? + (await this.liveEdgesUnlocked([...this.replaySync().values()].filter((candidate) => !candidate.deleted))); const byChild = new Map(); for (const edge of alive) { byChild.set(canonicalSessionPath(edge.child), edge); diff --git a/packages/coding-agent/test/suite/regressions/heartbeats-list-response-latency.test.ts b/packages/coding-agent/test/suite/regressions/heartbeats-list-response-latency.test.ts new file mode 100644 index 00000000000..702dc39a3c1 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/heartbeats-list-response-latency.test.ts @@ -0,0 +1,83 @@ +import { realpathSync } from "node:fs"; +import { join } from "node:path"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { AgentCronJobStore } from "../../../src/core/cron-jobs.js"; +import * as sessionManager from "../../../src/core/session-manager.js"; +import { DaemonSupervisor } from "../../../src/modes/daemon/daemon-supervisor.js"; +import { createHarness } from "../harness.js"; +import { createDeferred } from "../scheduling.js"; + +const harnesses: Array<{ cleanup: () => void }> = []; +afterEach(() => { + vi.restoreAllMocks(); + for (const harness of harnesses.splice(0)) harness.cleanup(); +}); +async function createSupervisorHarness() { + const harness = await createHarness(); + harnesses.push(harness); + const directory = harness.tempDir; + return new DaemonSupervisor(join(directory, "daemon.sock"), { + defaultSessionConfig: { agentDir: directory, cwd: directory }, + descriptorDir: join(directory, "workers"), + }) as unknown as { + defaultSessionConfig: { agentDir: string }; + rlmSpawnLedger: DaemonSupervisor["rlmSpawnLedger"]; + handleCommand: DaemonSupervisor["handleCommand"]; + broadcastHeartbeatsChanged: DaemonSupervisor["broadcastHeartbeatsChanged"]; + passiveScheduledJobs?: { rows: unknown[]; scannedAt: number }; + }; +} +type Supervisor = Awaited>; +function listHeartbeats(supervisor: Supervisor, id: string) { + return supervisor.handleCommand({} as never, { id, type: "heartbeats_list" }); +} +function heartbeatIds(response: Awaited>): string[] { + const rows = (response as { data?: { heartbeats?: Array<{ job: { id: string } }> } })?.data?.heartbeats; + return (rows ?? []).map((h) => h.job.id); +} +describe("heartbeats_list response latency", () => { + it("serves heartbeats_list from one shared scan without queueing behind metadata reads", async () => { + const supervisor = await createSupervisorHarness(); + const directory = realpathSync(supervisor.defaultSessionConfig.agentDir); + const manager = sessionManager.SessionManager.create(directory, join(directory, "sessions")); + manager.newSession(); + manager.flushNow(); + const store = AgentCronJobStore.forSessionArtifacts(); + store.registerSessionArtifact(manager.getSessionId(), manager.getSessionArtifactDir()!); + const job = store.createHeartbeat({ + activeSessionId: manager.getSessionId(), + sessionId: manager.getSessionId(), + sessionFile: manager.getSessionFile()!, + cwd: directory, + scheduleText: "every 1h", + prompt: "continue", + }).id; + const family = vi.spyOn(supervisor.rlmSpawnLedger(), "family"); + const responses = await Promise.all(["1", "2", "3", "4", "5"].map((id) => listHeartbeats(supervisor, id))); + expect(responses.map(heartbeatIds)).toEqual([[job], [job], [job], [job], [job]]); // one scan served five lists + await listHeartbeats(supervisor, "6"); + expect(family).toHaveBeenCalledTimes(1); // all six lists shared one scan, then served from the snapshot + const readSessionInfo = sessionManager.readSessionInfo; + const [scanStarted, releaseScan] = [createDeferred(), createDeferred()]; + vi.spyOn(sessionManager, "readSessionInfo").mockImplementationOnce(async (...args) => { + scanStarted.resolve(); + await releaseScan.promise; + return readSessionInfo(...args); + }); + supervisor.passiveScheduledJobs = undefined; // force the next list to scan cold + const siblings = supervisor.rlmSpawnLedger().siblings(manager.getSessionFile()!); + await scanStarted.promise; + try { + expect(heartbeatIds(await listHeartbeats(supervisor, "cold"))).toEqual([job]); + } finally { + releaseScan.resolve(); + await siblings; + } + store.manageHeartbeat(manager.getSessionId(), job, "pause"); + supervisor.broadcastHeartbeatsChanged(); + expect(heartbeatIds(await listHeartbeats(supervisor, "7"))).toEqual([job]); // broadcast dropped the snapshot + store.manageHeartbeat(manager.getSessionId(), job, "stop"); + supervisor.broadcastHeartbeatsChanged(); + expect(heartbeatIds(await listHeartbeats(supervisor, "8"))).toEqual([]); + }); +}); From 2cd9521ef1eaffcd247ad2390a8c24879a4e19a2 Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:44:20 -0400 Subject: [PATCH 07/26] [RSI, bug] fix(daemon): let session opens wait through an update restart instead of failing (#2391) --- .../.changes/fix-update-restart-open-wait.md | 1 + .../src/modes/agents-view/agents-view-mode.ts | 191 +++++++++++++++++- .../src/modes/daemon/daemon-errors.ts | 39 ++++ .../src/modes/daemon/daemon-mode.ts | 19 +- .../src/modes/daemon/daemon-protocol.ts | 3 +- .../src/modes/daemon/daemon-supervisor.ts | 32 ++- .../test/agents-view-mode.test.ts | 188 ++++++++++++++++- .../coding-agent/test/daemon-errors.test.ts | 28 +++ .../test/daemon-supervisor-monitor.test.ts | 1 + 9 files changed, 477 insertions(+), 25 deletions(-) create mode 100644 packages/coding-agent/.changes/fix-update-restart-open-wait.md diff --git a/packages/coding-agent/.changes/fix-update-restart-open-wait.md b/packages/coding-agent/.changes/fix-update-restart-open-wait.md new file mode 100644 index 00000000000..74ecb7923bb --- /dev/null +++ b/packages/coding-agent/.changes/fix-update-restart-open-wait.md @@ -0,0 +1 @@ +- Fixed opening an agent from Agents View during a daemon auto-update failing with "Daemon is preparing an update restart": the open now waits through the update restart (bounded) and reconnects once the daemon returns, with a notice that it waited instead of a hard failure. diff --git a/packages/coding-agent/src/modes/agents-view/agents-view-mode.ts b/packages/coding-agent/src/modes/agents-view/agents-view-mode.ts index 7dd3e01e63f..ffac99c6726 100644 --- a/packages/coding-agent/src/modes/agents-view/agents-view-mode.ts +++ b/packages/coding-agent/src/modes/agents-view/agents-view-mode.ts @@ -29,7 +29,7 @@ import { ensureTool } from "../../utils/tools-manager.js"; import { DaemonAgentConnection } from "../agent-connection/daemon-agent-connection.js"; import type { AgentConnectionHeartbeat, AgentConnectionSavedSessionInfo } from "../agent-connection/types.js"; import { DaemonClient, getDaemonSocketCloseReason } from "../daemon/daemon-client.js"; -import { DaemonSessionRecoveringError } from "../daemon/daemon-errors.js"; +import { DaemonSessionRecoveringError, isDaemonUpdateRestartingError } from "../daemon/daemon-errors.js"; import { collectDaemonClientEnv, type DaemonClosingReason, @@ -37,6 +37,7 @@ import { type DaemonResponse, isUnknownDaemonCommandError, } from "../daemon/daemon-protocol.js"; +import { DaemonControlPlaneTransportError } from "../daemon/daemon-routed-client.js"; import { resolveAttachModelFallbackMessage, type SessionSummary } from "../daemon/daemon-session-list.js"; import { listDaemonHeartbeats } from "../daemon/heartbeat-catalog.js"; import { @@ -314,6 +315,157 @@ interface OpenedAgentsViewSession { connection: DaemonAgentConnection; summary: SessionSummary; cwdFallbackNotice?: string; + updateRestartWaitNotice?: string; +} + +/** + * Bounded wait budget for an open that arrives while the daemon is preparing + * an update restart. Mirrors the update coordinator's worst case (100s + * prepare + supervisor stop + 60s successor startup + session restore), the + * same budget attached sessions get to reconnect after an update. + */ +export const DAEMON_UPDATE_RESTART_OPEN_WAIT_MS = 240_000; +const DAEMON_UPDATE_RESTART_OPEN_RETRY_MS = 500; + +export interface DaemonUpdateRestartWaitResult { + result: T; + waitedForUpdateRestart: boolean; +} + +/** + * Transport-level error codes that mean the socket was down while the daemon + * exited or its successor had not finished booting. + */ +const DAEMON_UPDATE_RESTART_TRANSIENT_ERROR_CODES = new Set([ + "ECONNREFUSED", + "ECONNRESET", + "EPIPE", + "ENOENT", + "ETIMEDOUT", + "ECONNABORTED", +]); + +/** + * True when an open failure is part of the normal update-restart window + * rather than a permanent failure: the preparing-restart rejection itself, + * transport failures while the daemon exits and its successor boots (socket + * close, connect/handshake/request timeouts, routed control-plane transport + * failures), and session-not-restored-yet misses ("Unknown active session", a + * session still recovering). Permanent create and attach failures (e.g. a + * missing session import file) return false so the open fails immediately + * instead of hiding behind the bounded update wait. + */ +function isDaemonUpdateRestartTransientError(error: unknown): boolean { + if (isDaemonUpdateRestartingError(error)) return true; + if (!(error instanceof Error)) return false; + if (isUnknownActiveSessionError(error)) return true; + if (error instanceof DaemonSessionRecoveringError) return true; + // A routed session transport wraps any control-plane transport failure + // (socket close, connect, timeouts — the restart-window shapes above); + // capability errors stay unwrapped and permanent. + if (error instanceof DaemonControlPlaneTransportError) return true; + const code = (error as NodeJS.ErrnoException).code; + if (typeof code === "string" && DAEMON_UPDATE_RESTART_TRANSIENT_ERROR_CODES.has(code)) return true; + return ( + // Socket closed while the daemon exits for the restart. + error.message.startsWith("Connection to the Prime Agent daemon closed.") || + // Connect failures while the successor socket is not listening yet. + error.message.startsWith("Failed to connect to the Prime Agent daemon:") || + // A request attempted while the transport is down between daemon processes. + error.message.startsWith("Cannot send daemon command") || + // Transport timeouts while the daemon exits and its successor boots: + // connect, handshake, and in-flight requests (e.g. create) that cannot + // get a response until the successor is ready. + /^Timed out after \d+ms (connecting to the Prime Agent daemon|waiting for the Prime Agent daemon handshake|waiting for the Prime Agent daemon response to)/.test( + error.message, + ) + ); +} + +function daemonUpdateRestartDeadlineError(waitMs: number, lastError: unknown): Error { + const lastErrorText = + lastError === undefined ? "none yet" : lastError instanceof Error ? lastError.message : String(lastError); + return new Error( + `The Prime Agent daemon did not finish its update restart within ${Math.round(waitMs / 1000)} seconds. Try opening this agent again once the update finishes. Last error: ${lastErrorText}`, + ); +} + +/** + * Run an open attempt, retrying while the daemon is in the update-restart + * transient state instead of failing the open. The first + * "preparing an update restart" rejection arms the wait; once armed, only + * restart-transient failures of the restart itself (socket close while the + * daemon exits, connect errors while the successor boots, session-not-yet- + * restored misses) stay inside the same bounded loop, because they are all part + * of the same normal update restart, while permanent failures (e.g. a missing + * session file) propagate immediately instead of hiding behind the wait. A + * non-update error before any update-restart signal propagates unchanged. The + * wait budget bounds the whole wait: each attempt is raced against the + * remaining budget, so an in-flight attempt (e.g. a create request with its own + * 30-second timeout) cannot hold the open past the deadline, which fails with a + * clear actionable message that includes the last error. A losing attempt that + * still settles afterwards is not abandoned silently: its result goes to + * onAbandoned for disposal and its failure is swallowed. + */ +export async function waitThroughDaemonUpdateRestart( + attempt: () => Promise, + options: { + waitMs?: number; + retryMs?: number; + onWait?: (error: unknown) => void; + /** Called with the result of an attempt that resolves after the deadline already failed the open, so resources nobody receives can be disposed. */ + onAbandoned?: (result: T) => void; + } = {}, +): Promise> { + const waitMs = options.waitMs ?? DAEMON_UPDATE_RESTART_OPEN_WAIT_MS; + const retryMs = options.retryMs ?? DAEMON_UPDATE_RESTART_OPEN_RETRY_MS; + const deadline = Date.now() + waitMs; + let sawUpdateRestart = false; + let lastError: unknown; + let attemptAbandoned = false; + while (true) { + const deadlineError = daemonUpdateRestartDeadlineError(waitMs, lastError); + // An in-flight attempt (e.g. a create request with its own 30-second + // timeout) must not push the open past the deadline: race every attempt + // against the remaining budget so the bound holds. + let deadlineTimer: ReturnType | undefined; + const deadlineHit = new Promise((_, reject) => { + deadlineTimer = setTimeout(() => reject(deadlineError), Math.max(0, deadline - Date.now())); + }); + const attemptPromise = attempt(); + // The race cannot cancel the losing attempt: when the deadline wins, a + // late success must still be disposed (nobody receives its result), and + // a late failure is expected and must not surface as unhandled. + void attemptPromise.then( + (result) => { + if (attemptAbandoned) options.onAbandoned?.(result); + }, + () => undefined, + ); + try { + const result = await Promise.race([attemptPromise, deadlineHit]); + clearTimeout(deadlineTimer); + return { result, waitedForUpdateRestart: sawUpdateRestart }; + } catch (error) { + clearTimeout(deadlineTimer); + if (error === deadlineError) { + attemptAbandoned = true; + throw deadlineError; + } + if (!sawUpdateRestart) { + if (!isDaemonUpdateRestartingError(error)) throw error; + sawUpdateRestart = true; + options.onWait?.(error); + } else if (!isDaemonUpdateRestartTransientError(error)) { + throw error; + } + lastError = error; + if (Date.now() + retryMs > deadline) { + throw daemonUpdateRestartDeadlineError(waitMs, lastError); + } + await new Promise((resolve) => setTimeout(resolve, retryMs)); + } + } } export function resolveAgentsViewOpenCwd( @@ -474,13 +626,36 @@ async function runAgentsViewLoop( let opened: OpenedAgentsViewSession | undefined; try { - opened = await openAgentsViewSession(options, result.summary); + // An auto-update can fence session opens as "preparing an update + // restart" at the exact moment the user hits enter; the restart is a + // normal transient state, so wait through it instead of failing. + const openedThroughUpdate = await waitThroughDaemonUpdateRestart( + () => openAgentsViewSession(options, result.summary), + { + onWait: (error) => logClientError("Waiting for daemon update restart to finish before opening", error), + // An open that resolves after the deadline already failed the + // wait: dispose the connection nobody received instead of + // leaking it. + onAbandoned: (abandoned) => { + void abandoned.connection.dispose().catch(() => undefined); + }, + }, + ); + opened = openedThroughUpdate.result; persistentState.backSession = opened.summary; - if (opened.cwdFallbackNotice) { - persistentState.statusMessage = combineAgentsViewStartupNotices( - result.statusMessage, - opened.cwdFallbackNotice, - ); + if (openedThroughUpdate.waitedForUpdateRestart) { + opened = { + ...opened, + updateRestartWaitNotice: + "Waited for the Prime Agent daemon update restart to finish before opening this agent", + }; + } + const openStartupNotice = combineAgentsViewStartupNotices( + opened.cwdFallbackNotice, + opened.updateRestartWaitNotice, + ); + if (openStartupNotice) { + persistentState.statusMessage = combineAgentsViewStartupNotices(result.statusMessage, openStartupNotice); } const uiServices = await resolveAgentsViewSessionUiServices(options, opened.summary); const interactiveMode = new InteractiveMode({ @@ -492,7 +667,7 @@ async function runAgentsViewLoop( bindLocalSessionExtensions: false, migratedProviders: options.migratedProviders, modelFallbackMessage: resolveAttachModelFallbackMessage(opened.summary, options.modelFallbackMessage), - startupNotice: combineAgentsViewStartupNotices(result.statusMessage, opened.cwdFallbackNotice), + startupNotice: combineAgentsViewStartupNotices(result.statusMessage, openStartupNotice), verbose: options.verbose, returnToAgentsView: true, forceFullscreen: true, diff --git a/packages/coding-agent/src/modes/daemon/daemon-errors.ts b/packages/coding-agent/src/modes/daemon/daemon-errors.ts index 760ed0b3335..b85429d5830 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-errors.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-errors.ts @@ -3,6 +3,12 @@ import { SessionImportFileNotFoundError } from "../../core/session-import-errors import { SessionAlreadyActiveError } from "../../core/session-lease.js"; import type { DaemonErrorInfo, DaemonResponse } from "./daemon-protocol.js"; +/** Rejection message for commands fenced out while the daemon prepares an update restart. */ +export const UPDATE_RESTART_PREPARING_MESSAGE = "Daemon is preparing an update restart"; + +/** Machine-readable companion for {@link UPDATE_RESTART_PREPARING_MESSAGE} rejections. */ +export const UPDATE_RESTART_PREPARING_ERROR_INFO: DaemonErrorInfo = { code: "update_restarting" }; + /** A known session (a persisted descriptor names it) that cannot be routed to yet; retryable, unlike "Unknown active session". */ export class DaemonSessionRecoveringError extends Error { readonly code = "session_recovering" as const; @@ -13,6 +19,33 @@ export class DaemonSessionRecoveringError extends Error { } } +/** + * The daemon is preparing an update restart: mutations (including session + * opens) are rejected while the restart coordinator drains and checkpoints. + * This is a normal transient state, so clients should wait through it and + * retry, not surface it as a hard failure. + */ +export class DaemonUpdateRestartingError extends Error { + readonly code = "update_restarting" as const; + + constructor(message = "Daemon is preparing an update restart") { + super(message); + this.name = "DaemonUpdateRestartingError"; + } +} + +/** + * True when an open failure is the update-restart transient state. Typed + * errors come from current daemons; the exact-message fallback also recognizes + * older daemons that reject with the same plain string. + */ +export function isDaemonUpdateRestartingError(error: unknown): boolean { + return ( + error instanceof DaemonUpdateRestartingError || + (error instanceof Error && error.message === "Daemon is preparing an update restart") + ); +} + export function serializeDaemonError(error: unknown): DaemonErrorInfo | undefined { if (error instanceof MissingSessionCwdError) { return { code: "missing_session_cwd", issue: error.issue }; @@ -30,6 +63,9 @@ export function serializeDaemonError(error: unknown): DaemonErrorInfo | undefine if (error instanceof DaemonSessionRecoveringError) { return { code: "session_recovering", activeSessionId: error.activeSessionId }; } + if (error instanceof DaemonUpdateRestartingError) { + return { code: "update_restarting" }; + } return undefined; } @@ -61,5 +97,8 @@ export function deserializeDaemonError(response: Extract { return { ...actual, appendRotatingLog: vi.fn() }; }); -vi.mock("../src/modes/daemon/daemon-client.js", () => ({ - DaemonClient: class { - connect = vi.fn(async () => undefined); - close = vi.fn(); - request = modeMocks.clientRequest; - }, - getDaemonSocketCloseReason: vi.fn(), -})); +vi.mock("../src/modes/daemon/daemon-client.js", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + DaemonClient: class { + connect = vi.fn(async () => undefined); + close = vi.fn(); + request = modeMocks.clientRequest; + }, + }; +}); vi.mock("../src/modes/agent-connection/daemon-agent-connection.js", () => ({ DaemonAgentConnection: Object.assign(function DaemonAgentConnection() {}, { @@ -1424,3 +1430,169 @@ describe("agents view startup notices", () => { expect(runs).toBe(2); }); }); + +describe("agents view open during a daemon update restart", () => { + beforeEach(() => vi.clearAllMocks()); + + it.each([ + [ + "waits through the update restart and surfaces the wait notice", + true, + "Waited for the Prime Agent daemon update restart to finish", + ], + [ + "surfaces a permanent create failure unmasked by the wait", + false, + "Failed to open agent: File not found: /tmp/scope.jsonl", + ], + ] as const)("%s", async (_name, reachesSession, expectedMessage) => { + const saved = summary({ activeSessionId: undefined, lifecycle: "archived" }); + let runs = 0; + vi.spyOn(AgentsViewMode.prototype, "run").mockImplementation(async function (this: AgentsViewMode) { + runs += 1; + if (runs === 1) return { type: "open", summary: saved, hasChildren: false }; + expect(String(Reflect.get(this, "persistentState").statusMessage)).toContain(expectedMessage); + return { type: "exit" }; + }); + modeMocks.clientRequest + .mockResolvedValueOnce({ + success: false, + error: "Daemon is preparing an update restart", + }) + .mockResolvedValueOnce( + reachesSession + ? { success: true, data: { ...saved, activeSessionId: "resumed-after-update", lifecycle: "live" } } + : { success: false, error: "File not found: /tmp/scope.jsonl" }, + ); + modeMocks.interactiveRun.mockResolvedValue({ + type: "agents_view", + source: { activeSessionId: "resumed-after-update", sessionId: saved.sessionId, cwd: saved.cwd }, + } as never); + + await runAgentsViewMode({ + config: { cwd: process.cwd() }, + socketPath: "/tmp/agents-view-test.sock", + uiServices: createUiServices(), + }); + + expect(modeMocks.clientRequest).toHaveBeenCalledTimes(2); + expect(modeMocks.interactiveRun).toHaveBeenCalledTimes(reachesSession ? 1 : 0); + expect(runs).toBe(2); + }); +}); + +describe("agents view open during a daemon update restart", () => { + beforeEach(() => vi.clearAllMocks()); + + it.each([ + [ + "waits through the update restart and surfaces the wait notice", + true, + "Waited for the Prime Agent daemon update restart to finish", + ], + [ + "surfaces a permanent create failure unmasked by the wait", + false, + "Failed to open agent: File not found: /tmp/scope.jsonl", + ], + ] as const)("%s", async (_name, reachesSession, expectedMessage) => { + const saved = summary({ activeSessionId: undefined, lifecycle: "archived" }); + let runs = 0; + vi.spyOn(AgentsViewMode.prototype, "run").mockImplementation(async function (this: AgentsViewMode) { + runs += 1; + if (runs === 1) return { type: "open", summary: saved, hasChildren: false }; + expect(String(Reflect.get(this, "persistentState").statusMessage)).toContain(expectedMessage); + return { type: "exit" }; + }); + modeMocks.clientRequest + .mockResolvedValueOnce({ + success: false, + error: "Daemon is preparing an update restart", + }) + .mockResolvedValueOnce( + reachesSession + ? { success: true, data: { ...saved, activeSessionId: "resumed-after-update", lifecycle: "live" } } + : { success: false, error: "File not found: /tmp/scope.jsonl" }, + ); + modeMocks.interactiveRun.mockResolvedValue({ + type: "agents_view", + source: { activeSessionId: "resumed-after-update", sessionId: saved.sessionId, cwd: saved.cwd }, + } as never); + + await runAgentsViewMode({ + config: { cwd: process.cwd() }, + socketPath: "/tmp/agents-view-test.sock", + uiServices: createUiServices(), + }); + + expect(modeMocks.clientRequest).toHaveBeenCalledTimes(2); + expect(modeMocks.interactiveRun).toHaveBeenCalledTimes(reachesSession ? 1 : 0); + expect(runs).toBe(2); + }); +}); + +describe("waitThroughDaemonUpdateRestart", () => { + const updateRestartDeadline = + /The Prime Agent daemon did not finish its update restart within \d+ seconds\. Try opening this agent again once the update finishes\. Last error: Daemon is preparing an update restart/; + + // Post-arm transient shapes only; arming is pinned at the loop/deadline layers. + it("retries every restart-transient failure and reports that it waited", async () => { + const transientFailures = [ + () => new Error("Failed to connect to the Prime Agent daemon: connect ENOENT"), + () => Object.assign(new Error("connect ECONNREFUSED"), { code: "ECONNREFUSED" }), + () => new Error("Connection to the Prime Agent daemon closed."), + () => new Error('Timed out after 30000ms waiting for the Prime Agent daemon response to "create".'), + () => new DaemonControlPlaneTransportError(new Error("Connection to the Prime Agent daemon closed.")), + () => new Error("Unknown active session: update-restart-session"), + () => new DaemonSessionRecoveringError("update-restart-session"), + ]; + let attempts = 0; + const outcome = await waitThroughDaemonUpdateRestart( + async () => { + attempts += 1; + if (attempts === 1) throw new DaemonUpdateRestartingError(); + const failure = transientFailures[attempts - 2]; + if (failure) throw failure(); + return "opened"; + }, + { waitMs: 5_000, retryMs: 1 }, + ); + expect(outcome).toEqual({ result: "opened", waitedForUpdateRestart: true }); + }); + + it("propagates a non-update failure before any update-restart signal", async () => { + let attempts = 0; + await expect( + waitThroughDaemonUpdateRestart(async () => { + attempts += 1; + throw new Error("spawn EMFILE"); + }), + ).rejects.toThrow("spawn EMFILE"); + expect(attempts).toBe(1); + }); + + // The deadline races every attempt so an in-flight create cannot hold the open past the budget. + it("fails at the deadline even when an in-flight attempt would block past it", async () => { + vi.useFakeTimers(); + const inFlight = new Promise(() => {}); + let attempts = 0; + try { + const opening = waitThroughDaemonUpdateRestart( + async () => { + attempts += 1; + if (attempts === 1) throw new DaemonUpdateRestartingError(); + return inFlight; + }, + { waitMs: 60, retryMs: 5 }, + ).then( + () => "unexpectedly opened", + (error: Error) => error.message, + ); + await vi.advanceTimersByTimeAsync(60); + expect(await opening).toMatch(updateRestartDeadline); + expect(attempts).toBe(2); + } finally { + vi.useRealTimers(); + } + }); +}); diff --git a/packages/coding-agent/test/daemon-errors.test.ts b/packages/coding-agent/test/daemon-errors.test.ts index e1c17942f21..c148c891d59 100644 --- a/packages/coding-agent/test/daemon-errors.test.ts +++ b/packages/coding-agent/test/daemon-errors.test.ts @@ -3,8 +3,10 @@ import { SessionAlreadyActiveError } from "../src/core/session-lease.js"; import { DaemonSessionCreateError, DaemonSessionRecoveringError, + DaemonUpdateRestartingError, deserializeDaemonCreateError, deserializeDaemonError, + isDaemonUpdateRestartingError, serializeDaemonError, } from "../src/modes/daemon/daemon-errors.js"; @@ -58,3 +60,29 @@ describe("session_recovering wire round-trip", () => { expect(legacy).not.toBeInstanceOf(DaemonSessionRecoveringError); }); }); + +describe("update_restarting wire round-trip", () => { + it("serializes for old clients (readable message) and deserializes for new clients (typed, retryable)", () => { + const error = new DaemonUpdateRestartingError(); + const errorInfo = serializeDaemonError(error); + expect(errorInfo).toEqual({ code: "update_restarting" }); + expect(error.message).toBe("Daemon is preparing an update restart"); + const roundTripped = deserializeDaemonError({ + type: "response", + command: "create", + success: false, + error: error.message, + errorInfo, + }); + expect(roundTripped).toBeInstanceOf(DaemonUpdateRestartingError); + const legacy = deserializeDaemonError({ + type: "response", + command: "create", + success: false, + error: "Daemon is preparing an update restart", + }); + expect(legacy).not.toBeInstanceOf(DaemonUpdateRestartingError); + expect(isDaemonUpdateRestartingError(legacy)).toBe(true); + expect(isDaemonUpdateRestartingError(new Error("Unknown active session: active-gap"))).toBe(false); + }); +}); diff --git a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts index a3f64410816..3df17cd44ab 100644 --- a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts @@ -6304,6 +6304,7 @@ describe("daemon worker supervisor monitoring", () => { await prepare; await expect(client.request({ type: "abort", activeSessionId: "missing" })).resolves.toMatchObject({ error: "Daemon is preparing an update restart", + errorInfo: { code: "update_restarting" }, }); } finally { client.close(); From 406b46742d845f208aff6affd504e11f7fa98763 Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:44:41 -0400 Subject: [PATCH 08/26] [RSI, bug] Hold spawn name reservations until admission is durable (#2396) --- .../.changes/hold-spawn-name-reservation.md | 1 + .../coding-agent/src/core/agent-session.ts | 45 ++++++++++++----- .../src/modes/daemon/daemon-mode.ts | 27 ++++++++++ .../test/agent-session-recursion.test.ts | 23 +++++++++ .../coding-agent/test/daemon-mode.test.ts | 49 ++++++++++++------- 5 files changed, 114 insertions(+), 31 deletions(-) create mode 100644 packages/coding-agent/.changes/hold-spawn-name-reservation.md diff --git a/packages/coding-agent/.changes/hold-spawn-name-reservation.md b/packages/coding-agent/.changes/hold-spawn-name-reservation.md new file mode 100644 index 00000000000..bd2857e1964 --- /dev/null +++ b/packages/coding-agent/.changes/hold-spawn-name-reservation.md @@ -0,0 +1 @@ +- Fixed spawn name reservations releasing before admission was durable, which let two parallel same-name spawns both admit and write duplicate ledger edges that made child selectors ambiguous. diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index c08755cf531..a28617b1d70 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -11693,24 +11693,36 @@ export class AgentSession { } this._pendingRlmSubagentSessionNames.add(requestedSessionName); } + // The name stays reserved until the spawn admission settles: the + // detached runtime task releases it at admission completion (success + // or failure), and every pre-admission failure path releases it here. + const releaseReservedSessionName = () => { + if (requestedSessionName) this._pendingRlmSubagentSessionNames.delete(requestedSessionName); + }; let modelSelection: RlmSubagentModelSelection; + let childSessionDir: string; try { if (requestedSessionName) await this._assertRlmSubagentSessionNameAvailable(requestedSessionName, true); + // An unpinned spawn model resolves against the persisted subagent + // default; an unavailable default fails the spawn instead of silently + // inheriting the parent model. modelSelection = await this._resolveRlmSubagentModel(requestedModel); - } finally { - if (requestedSessionName) this._pendingRlmSubagentSessionNames.delete(requestedSessionName); - } - if (requestedThinkingLevel !== undefined) { - const supported = getSupportedThinkingLevels(modelSelection.model) as ThinkingLevel[]; - if (!supported.includes(requestedThinkingLevel)) { - throw new Error( - `Requested thinking level "${requestedThinkingLevel}" is not supported by model "${modelSelection.model.provider}/${modelSelection.model.id}"; supported levels: ${supported.join(", ")}`, - ); + if (requestedThinkingLevel !== undefined) { + const supported = getSupportedThinkingLevels(modelSelection.model) as ThinkingLevel[]; + if (!supported.includes(requestedThinkingLevel)) { + throw new Error( + `Requested thinking level "${requestedThinkingLevel}" is not supported by model "${modelSelection.model.provider}/${modelSelection.model.id}"; supported levels: ${supported.join(", ")}`, + ); + } } + if (this._disposed || this._disposing) { + throw new Error("Cannot spawn a subagent after its parent was disposed"); + } + childSessionDir = this._createChildRlmSessionDir(); + } catch (error) { + releaseReservedSessionName(); + throw error; } - if (this._disposed || this._disposing) throw new Error("Cannot spawn a subagent after its parent was disposed"); - - const childSessionDir = this._createChildRlmSessionDir(); const childNodeId = basename(childSessionDir); const sessionName = requestedSessionName ?? createDefaultRlmSubagentSessionName(prompt, childNodeId); if (!requestedSessionName) await this._assertRlmSubagentSessionNameAvailable(sessionName); @@ -11883,7 +11895,14 @@ export class AgentSession { void (async () => { let childRuntime: RlmSubagentRuntime | undefined; try { - childRuntime = await this._createRlmSubagentRuntime(subagentOptions); + try { + childRuntime = await this._createRlmSubagentRuntime(subagentOptions); + } finally { + // Admission settled: in daemon mode the spawn edge is now + // durable, so the name transfers from the pending reservation + // to the admitted run. A failed admission frees the name. + releaseReservedSessionName(); + } const child = childRuntime.session; if (run.status === "cancelled") throw new Error(run.error ?? "RLM child cancelled"); if (child.sessionName !== sessionName) child.setSessionName(sessionName); diff --git a/packages/coding-agent/src/modes/daemon/daemon-mode.ts b/packages/coding-agent/src/modes/daemon/daemon-mode.ts index 8ee3a08cb58..680b850d01a 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-mode.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-mode.ts @@ -2839,6 +2839,33 @@ export class AgentDaemon { private async createRlmSubagentRuntime( parentState: ActiveSessionState, options: CreateRlmSubagentRuntimeOptions, + ): Promise { + // The sibling name is held under a daemon-wide reservation for the + // whole admission and re-asserted at this boundary, so a same-name + // sibling that lands mid-admission fails closed before the durable + // ledger edge is appended. + const nameReservation = { + name: options.sessionName, + depth: options.rlmDepth, + parentSessionId: options.parentSession.sessionId, + ...(options.parentSession.sessionFile ? { parentSessionPath: options.parentSession.sessionFile } : {}), + }; + const reservationKey = sessionNameReservationKey(nameReservation); + if (this.pendingSessionNames.has(reservationKey)) { + throw new Error(formatAgentSessionNameUnavailable(options.sessionName, options.rlmDepth)); + } + this.pendingSessionNames.add(reservationKey); + try { + await this.assertFamilySessionNameAvailable(nameReservation, parentState, true); + return await this.admitRlmSubagentRuntime(parentState, options); + } finally { + this.pendingSessionNames.delete(reservationKey); + } + } + + private async admitRlmSubagentRuntime( + parentState: ActiveSessionState, + options: CreateRlmSubagentRuntimeOptions, ): Promise { const sessionManager = SessionManager.create(options.parentSession.sessionManager.getCwd(), options.sessionDir); sessionManager.newSession({ diff --git a/packages/coding-agent/test/agent-session-recursion.test.ts b/packages/coding-agent/test/agent-session-recursion.test.ts index 38b6456a618..a96418c5178 100644 --- a/packages/coding-agent/test/agent-session-recursion.test.ts +++ b/packages/coding-agent/test/agent-session-recursion.test.ts @@ -16,6 +16,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { type AgentSessionMessageController, createAgentSessionMessage, + formatAgentSessionNameUnavailable, isAgentSessionMessage, } from "../src/core/agent-messages.js"; import { AgentSession, type RlmChildAgentSnapshot } from "../src/core/agent-session.js"; @@ -568,6 +569,28 @@ describe("AgentSession rlm recursion", () => { await runPromise; }); + it("holds a spawn name reservation until admission settles, then frees it", async () => { + const releaseAdmission = deferred(); + const root = createSession({ + subagentRuntimeHost: { + createRlmSubagentRuntime: async () => { + await releaseAdmission.promise; + throw new Error("kernel startup failed"); + }, + deleteRlmSubagentRuntime: async () => {}, + }, + }); + const internals = root as unknown as InspectableRlmSession & { _pendingRlmSubagentSessionNames: Set }; + const unavailable = formatAgentSessionNameUnavailable("slow-worker", root.rlmDepth + 1); + const spawned = await root.runRlmChild("slow admitting child", { name: "slow-worker" }); + expect(internals._pendingRlmSubagentSessionNames.has("slow-worker")).toBe(true); + await expect(root.runRlmChild("racing spawn", { name: "slow-worker" })).rejects.toThrow(unavailable); + releaseAdmission.resolve(); + await internals._activeRlmChildRuns.get(spawned.rlm_child_id)!.settlement!.promise; + expect(internals._pendingRlmSubagentSessionNames.has("slow-worker")).toBe(false); + await expect(root.runRlmChild("respawn while retained", { name: "slow-worker" })).rejects.toThrow(unavailable); + }); + it("makes an externally restored retained child listable and deletable", async () => { const childId = "restored-child"; const childDir = join(tempDir, childId); diff --git a/packages/coding-agent/test/daemon-mode.test.ts b/packages/coding-agent/test/daemon-mode.test.ts index bd5eedd2320..9d9c8e3b44e 100644 --- a/packages/coding-agent/test/daemon-mode.test.ts +++ b/packages/coding-agent/test/daemon-mode.test.ts @@ -1021,6 +1021,7 @@ describe("daemon mode helpers", () => { saved: Awaited>, jobs: AgentCronJob[], ): Promise>; + rlmSpawnLedger(): { liveEdges(): Promise> }; }; const parentState = await internals.createRuntime({ type: "create", sessionPath: parentSessionFile }); Object.assign(parentState.runtime.session, { @@ -1033,24 +1034,30 @@ describe("daemon mode helpers", () => { hasRunningRlmChildren: () => false, getSessionActionSnapshot: () => ({ queuedCount: 0, steering: [], followUps: [] }), }); - const childRuntime = await internals.createRlmSubagentRuntime(parentState, { - parentSession: parentState.runtime.session, - id: "child-1", - prompt: "complete and persist", - sessionName: "real-worker", - sessionDir: childSessionDir, - model: { provider: "test", id: "model" } as Model, - thinkingLevel: "off", - serviceTier: null, - scopedModels: [], - activeToolNames: [], - customTools: [], - includeGoals: false, - includeCompactSkill: false, - rlmDepth: 1, - rlmMaxDepth: 4, - rlmParentNodeId: "child-1", - }); + const spawn = (id: string) => + internals.createRlmSubagentRuntime(parentState, { + parentSession: parentState.runtime.session, + id, + prompt: "complete and persist", + sessionName: "real-worker", + sessionDir: join(parentManager.getSessionArtifactDir()!, id), + model: { provider: "test", id: "model" } as Model, + thinkingLevel: "off" as const, + serviceTier: null, + scopedModels: [], + activeToolNames: [], + customTools: [], + includeGoals: false, + includeCompactSkill: false, + rlmDepth: 1, + rlmMaxDepth: 4, + rlmParentNodeId: id, + }); + const admission = spawn("child-1"); + await expect(spawn("child-2")).rejects.toThrow('Agent name "real-worker" is unavailable'); + const childRuntime = await admission; + const edges = await internals.rlmSpawnLedger().liveEdges(); + expect(edges.filter((edge) => edge.name === "real-worker")).toHaveLength(1); const childState = [...internals.sessions.values()].find( (state) => state.runtime.session === childRuntime.session, ); @@ -1088,6 +1095,12 @@ describe("daemon mode helpers", () => { status: "completed", prompt: "complete and persist", }); + await host.deleteRlmSubagentRuntime?.("child-1", childRuntime.session); + await spawn("child-3"); + const edgesAfter = await internals.rlmSpawnLedger().liveEdges(); + const namedEdges = edgesAfter.filter((edge) => edge.name === "real-worker"); + expect(namedEdges).toHaveLength(1); + expect(namedEdges[0]?.childId).toBe("child-3"); } finally { rmSync(tempDir, { recursive: true, force: true }); } From cf2e779aa82adae99b449f8b32e3e7ed948eaac7 Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Tue, 22 Sep 2026 05:59:45 -0600 Subject: [PATCH 09/26] [RSI, bug] Fix stale async-bash completion notice race at the turn boundary (#2372) --- prime-agent-runtime/src/rlm/bash.py | 41 ++++++++++++++------------- prime-agent-runtime/test/test_repl.py | 4 +++ 2 files changed, 26 insertions(+), 19 deletions(-) diff --git a/prime-agent-runtime/src/rlm/bash.py b/prime-agent-runtime/src/rlm/bash.py index e59be773c95..07ad1ec24f9 100644 --- a/prime-agent-runtime/src/rlm/bash.py +++ b/prime-agent-runtime/src/rlm/bash.py @@ -17,6 +17,7 @@ import sys import threading import time +import uuid from collections import deque from collections.abc import Callable, Generator from dataclasses import dataclass @@ -744,17 +745,7 @@ async def _notify_background_completion( def _arm_consumed_notice(self, command: str) -> None: # Armed only post-acceptance: the withdrawal can never overtake its notice. - loop = asyncio.get_running_loop() - - def dispatch() -> None: - def start() -> None: - task = loop.create_task(self._notify_result_consumed(command)) - task.add_done_callback(_consume_notice_task) - - try: - loop.call_soon_threadsafe(start) - except RuntimeError: - pass # notifying loop already closed + dispatch = functools.partial(self._notify_result_consumed, command) with self._callback_lock: if not self._result_consumed: @@ -762,22 +753,34 @@ def start() -> None: return dispatch() - async def _notify_result_consumed(self, command: str) -> None: + def _notify_result_consumed(self, command: str) -> None: + """Ship the withdrawal inside the read, ahead of the cell's done event. + + The host delivers a queued notice at the reading cell's turn boundary, + which begins when that cell's done event is processed: a withdrawal + frame that leaves the kernel after done arrives too late, and the stale + notice wakes the model anyway. Reads happen inside a live cell, so + writing the frame right here puts it ahead of done on the wire, where + the host must withdraw before it can dispatch. The reply never matters + (unknown reply ids are dropped), so the request is fire-and-forget: + no future to await, no event-loop hop that could run after the cell. + """ from . import repl if not repl.is_active(): return - try: - await repl.host_request( - { + repl._send( + { + "event": "host_request", + "id": uuid.uuid4().hex, + "data": { "type": "bash.consumed", "completionId": self._completion_id, "pid": self._pid, "command": command, - } - ) - except (OSError, RuntimeError): - return # bridge closed at teardown; old hosts error-reply — both fine + }, + } + ) async def _wait_reaped(self) -> None: loop = asyncio.get_running_loop() diff --git a/prime-agent-runtime/test/test_repl.py b/prime-agent-runtime/test/test_repl.py index 7bc22a49beb..9867d89549b 100644 --- a/prime-agent-runtime/test/test_repl.py +++ b/prime-agent-runtime/test/test_repl.py @@ -793,6 +793,10 @@ def test_result_read_after_notice_requests_one_withdrawal(self): reply_ok(self.repl, notice) read = self.repl.execute(f"withdraw-{label}-read", read_code) self.assertIn(f"withdrawn-{label}", one(read, "result")["text"]) + # The withdrawal must ship ahead of the reading cell's done, or the + # notice is delivered as a stale turn-boundary notice instead. + kinds = [event.get("event") for event in read] + self.assertLess(kinds.index("host_request"), kinds.index("done")) request = wait_for_host_request(self.repl, read) self.assertEqual( request["data"], {"type": "bash.consumed", "completionId": notice["data"]["completionId"], "pid": pid, "command": command} From 2035ceae929d5fcdc13238ad9e1bda7d55818f1e Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Tue, 22 Sep 2026 06:00:14 -0600 Subject: [PATCH 10/26] [RSI, bug] fix(coding-agent): re-park pending next-turn messages when an async-bash notice is withdrawn (#2386) --- .../coding-agent/src/core/agent-session.ts | 5 +- .../2068-shell-message-steering.test.ts | 53 +++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index a28617b1d70..0207237d08d 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -2255,7 +2255,10 @@ export class AgentSession { const restorable = payload.records .filter( (record): record is DeliveryRecord & { message: CustomMessage } => - (record.role === "next_turn" || (payload.acceptedAgentMessage && record.role === "prefix")) && + // Prefix records are parked next-turn context the action captured + // on admission; a cancelled turn hands them back, like the + // admission-rejection and dispatch-failure paths already do. + (record.role === "next_turn" || record.role === "prefix") && record.message.role === "custom" && record.message.customType !== HARNESS_DIGEST_CUSTOM_TYPE && !record.durable, diff --git a/packages/coding-agent/test/suite/regressions/2068-shell-message-steering.test.ts b/packages/coding-agent/test/suite/regressions/2068-shell-message-steering.test.ts index d640f9f9d19..89ee998b048 100644 --- a/packages/coding-agent/test/suite/regressions/2068-shell-message-steering.test.ts +++ b/packages/coding-agent/test/suite/regressions/2068-shell-message-steering.test.ts @@ -177,6 +177,59 @@ describe("#2068 shell message steering", () => { expect(harness.eventsOfType("agent_start")).toHaveLength(1); }); + it("re-parks pending next-turn messages when an async-bash notice is withdrawn", async () => { + const started = createDeferred(); + const release = createDeferred(); + let providerSawParkedContext = false; + const tool: AgentTool = { + name: "wait", + label: "Wait", + description: "Hold the current tool until released", + parameters: Type.Object({}), + execute: async () => { + started.resolve(); + await release.promise; + return { content: [{ type: "text", text: "released" }], details: {} }; + }, + }; + const harness = await createHarness({ tools: [tool] }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), + fauxAssistantMessage("Finished the original work."), + (context) => { + providerSawParkedContext = context.messages.some((message) => getMessageText(message) === "carry this"); + return fauxAssistantMessage("Follow-up turn complete."); + }, + ]); + const original = harness.session.prompt("Continue working."); + try { + await started.promise; + await harness.session.sendCustomMessage( + { customType: "next-turn", content: "carry this", display: true, details: {} }, + { deliverAs: "nextTurn" }, + ); + await completeShell(harness); + const [notice] = harness.session.getSessionActionRecoverySnapshot().actions; + const prefixes = + notice?.payload.kind === "turn" ? notice.payload.records.filter((record) => record.role === "prefix") : []; + expect(prefixes.map((record) => getMessageText(record.message))).toEqual(["carry this"]); + expect(harness.session.getSteeringMessages()).toHaveLength(1); + await readShellResult(harness); + expect(harness.session.getSteeringMessages()).toEqual([]); + expect(harness.session.getPendingNextTurnMessageSnapshots().map(getMessageText)).toEqual(["carry this"]); + } finally { + release.resolve(); + await original; + } + await harness.session.waitForIdle(); + expect(shellMessages(harness)).toEqual([]); + expect(harness.eventsOfType("agent_start")).toHaveLength(1); + await harness.session.prompt("next turn"); + await harness.session.waitForIdle(); + expect(providerSawParkedContext).toBe(true); + }); + it("renders shell-specific queue and transcript labels without generic delivery prefixes", () => { const message = createAsyncBashCompletionMessage(completion); const preview = `${ASYNC_BASH_COMPLETION_PREVIEW_LABEL}: pid 42, exit 0`; From 5311c2d5dbe4aaf3ad9f06a640beeca6471a09f9 Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:43:59 -0400 Subject: [PATCH 11/26] [RSI, bug] compaction: keep tool identity in serialized summarizer input (#2424) * fix(coding-agent): keep tool identity in compaction summarizer serialization Compaction serialized all tool results as anonymous "[Tool result]" blobs, so parallel tool calls could not be paired with their results and failures were indistinguishable from successes in the summarizer input. Label each result with its tool name and mark failed calls with an error tag so summaries keep failure attribution. Truncation behavior is unchanged. * fix(coding-agent): pair serialized tool calls and results by index serializeConversation labeled tool results with the tool name only, so when the same tool was called twice in one turn the summarizer could not reliably associate a result (or its error status) with the correct call, and skipped empty results shifted the positional pairing. Assign each serialized tool call a 1-based sequential index, prefix the call with it (`#1 ipython(...); #2 bash(...)`), and repeat the index in its result label (`[Tool result (ipython) #1]`, `[Tool result (ipython, error) #2]`). The short index stands in for the raw provider toolCallId, which can exceed 450 characters and would bloat the token-budgeted summarizer input. Results whose call was not serialized (branch-summary budget slices drop older entries) fall back to the name-only label. Update the compaction docs examples and the changelog, and add regression tests for same-tool pairing and the orphaned-result fallback. * docs(coding-agent): correct branch-summary note in compaction serialization docs Branch summarization drops toolResult entries in getMessageFromEntry before serializeConversation is ever called, so branch summaries show the `#N` call prefixes with no result lines at all. The fallback sentence added in 6e3aa4c9b wrongly attributed orphaned results to branch-summary budget slices. Re-attribute the name-only fallback to serialized inputs that lack the matching call, which extension callers can pass since serializeConversation is exported. Docs and in-code comment only; no behavior change and no test changes. (cherry picked from commit 5fb3d9acbafa06ae9228c49a4b8a4eca9abc0997) --- .../.changes/compaction-tool-result-labels.md | 3 + packages/coding-agent/docs/compaction.md | 11 ++- .../coding-agent/src/core/compaction/utils.ts | 46 ++++++++++-- packages/coding-agent/test/compaction.test.ts | 74 ++++++++++++++++++- 4 files changed, 122 insertions(+), 12 deletions(-) create mode 100644 packages/coding-agent/.changes/compaction-tool-result-labels.md diff --git a/packages/coding-agent/.changes/compaction-tool-result-labels.md b/packages/coding-agent/.changes/compaction-tool-result-labels.md new file mode 100644 index 00000000000..30e8484e282 --- /dev/null +++ b/packages/coding-agent/.changes/compaction-tool-result-labels.md @@ -0,0 +1,3 @@ +- Fixed: compaction, branch, and refine summaries now label each serialized tool result with its tool name, so parallel tool calls can be paired with their results. +- Fixed: failed tool results are marked as errors in the label, keeping failure attribution visible in summaries. +- Fixed: serialized tool calls now carry a sequential `#N` prefix and results repeat the matching index, so repeated calls of the same tool pair unambiguously. diff --git a/packages/coding-agent/docs/compaction.md b/packages/coding-agent/docs/compaction.md index 11f2837cfe3..69b3080d5f3 100644 --- a/packages/coding-agent/docs/compaction.md +++ b/packages/coding-agent/docs/compaction.md @@ -259,12 +259,15 @@ Before summarization, messages are serialized to text via [`serializeConversatio [User]: What they said [Assistant thinking]: Internal reasoning [Assistant]: Response text -[Assistant tool calls]: ipython(code="open('foo.ts').read()"); edit(path="bar.ts", ...) -[Tool result]: Output from tool +[Assistant tool calls]: #1 ipython(code="open('foo.ts').read()"); #2 bash(command="npm test") +[Tool result (ipython) #1]: Output from tool +[Tool result (bash, error) #2]: Output from a failed tool call ``` This prevents the model from treating it as a conversation to continue. +Each tool call carries a sequential `#N` prefix and its result repeats the matching index, so repeated calls of the same tool in one turn pair unambiguously. Branch summaries drop tool-result entries before serialization (their context remains attached to the assistant tool call), so they show the `#N` call prefixes without any result lines. When the serialized input lacks the matching call — possible for extension callers passing partial message lists, since `serializeConversation` is exported — the result falls back to the name-only label. + Tool results are truncated to 2000 characters during serialization. Content beyond that limit is replaced with a marker indicating how many characters were truncated. This keeps summarization requests within reasonable token budgets, since tool results, especially from `ipython` and optional `bash`, are typically the largest contributors to context size. ## Custom Summarization via Extensions @@ -323,8 +326,8 @@ pi.on("session_before_compact", async (event, ctx) => { // [User]: message text // [Assistant thinking]: thinking content // [Assistant]: response text - // [Assistant tool calls]: ipython(code="open('...').read()"); bash(command="...") - // [Tool result]: output text + // [Assistant tool calls]: #1 ipython(code="open('...').read()"); #2 bash(command="...") + // [Tool result (ipython) #1]: output text // Now send to your model for summarization const summary = await myModel.summarize(conversationText); diff --git a/packages/coding-agent/src/core/compaction/utils.ts b/packages/coding-agent/src/core/compaction/utils.ts index 72a764566d0..a582b66ba74 100644 --- a/packages/coding-agent/src/core/compaction/utils.ts +++ b/packages/coding-agent/src/core/compaction/utils.ts @@ -71,15 +71,27 @@ export function formatFileOperations(readFiles: string[], modifiedFiles: string[ } /** Maximum characters for a tool result in serialized summaries. */ const TOOL_RESULT_MAX_CHARS = 2000; +/** + * Characters kept from the end of a truncated tool result. Tool output is + * tail-heavy: exit errors, stack traces, and log tails appear at the end. + */ +const TOOL_RESULT_TAIL_CHARS = 500; /** * Truncate text to a maximum character length for summarization. - * Keeps the beginning and appends a truncation marker. + * Keeps the beginning and the end within the same total budget, marking + * the elided middle. */ function truncateForSummary(text: string, maxChars: number): string { if (text.length <= maxChars) return text; - const truncatedChars = text.length - maxChars; - return `${text.slice(0, maxChars)}\n\n[... ${truncatedChars} more characters truncated]`; + // The marker's digit counts are largest when the elided and kept sizes hit + // the text and budget maxima, so reserve space for that worst case to keep + // the result within maxChars. + const markerMaxLength = + `[... ${text.length} characters truncated; first ${maxChars} and last ${TOOL_RESULT_TAIL_CHARS} kept ...]`.length; + const headChars = maxChars - TOOL_RESULT_TAIL_CHARS - markerMaxLength - 4; + const elided = text.length - headChars - TOOL_RESULT_TAIL_CHARS; + return `${text.slice(0, headChars)}\n\n[... ${elided} characters truncated; first ${headChars} and last ${TOOL_RESULT_TAIL_CHARS} kept ...]\n\n${text.slice(text.length - TOOL_RESULT_TAIL_CHARS)}`; } /** @@ -89,9 +101,19 @@ function truncateForSummary(text: string, maxChars: number): string { * * Tool results are truncated to keep the summarization request within * reasonable token budgets. Full content is not needed for summarization. + * + * Tool calls are serialized with a sequential `#N` prefix and results repeat + * the matching index, so repeated calls of the same tool pair unambiguously. */ export function serializeConversation(messages: Message[]): string { const parts: string[] = []; + // Tool calls are serialized with a 1-based sequential index and results + // repeat the index of their call (matched by toolCallId), so repeated + // calls of the same tool pair unambiguously in the summarizer input. + // The short index stands in for the raw provider toolCallId, which can + // exceed 450 characters on some providers. + const toolCallIndices = new Map(); + let toolCallIndex = 0; for (const msg of messages) { if (msg.role === "user") { @@ -118,7 +140,9 @@ export function serializeConversation(messages: Message[]): string { const argsStr = Object.entries(args) .map(([k, v]) => `${k}=${JSON.stringify(v)}`) .join(", "); - toolCalls.push(`${block.name}(${argsStr})`); + toolCallIndex += 1; + toolCallIndices.set(block.id, toolCallIndex); + toolCalls.push(`#${toolCallIndex} ${block.name}(${argsStr})`); } } @@ -137,7 +161,19 @@ export function serializeConversation(messages: Message[]): string { .map((c) => c.text) .join(""); if (content) { - parts.push(`[Tool result]: ${truncateForSummary(content, TOOL_RESULT_MAX_CHARS)}`); + // Label the tool name, error status, and the index of the + // paired call so the summarizer can match each result to + // its `#N`-prefixed entry in the [Assistant tool calls] + // lines even when the same tool is called repeatedly in + // one turn. Results whose call is not part of the input + // (extension callers may pass partial message lists) + // fall back to the name-only label. + const callIndex = toolCallIndices.get(msg.toolCallId); + const indexSuffix = callIndex === undefined ? "" : ` #${callIndex}`; + const label = msg.isError + ? `[Tool result (${msg.toolName}, error)${indexSuffix}]` + : `[Tool result (${msg.toolName})${indexSuffix}]`; + parts.push(`${label}: ${truncateForSummary(content, TOOL_RESULT_MAX_CHARS)}`); } } } diff --git a/packages/coding-agent/test/compaction.test.ts b/packages/coding-agent/test/compaction.test.ts index cfc545c17e9..c3bb4848620 100644 --- a/packages/coding-agent/test/compaction.test.ts +++ b/packages/coding-agent/test/compaction.test.ts @@ -1,5 +1,5 @@ import type { AgentMessage } from "@earendil-works/pi-agent-core"; -import type { AssistantMessage, Usage } from "@earendil-works/pi-ai"; +import type { AssistantMessage, Message, ToolCall, Usage } from "@earendil-works/pi-ai"; import { getModel } from "@earendil-works/pi-ai"; import { readFileSync } from "fs"; import { join } from "path"; @@ -16,6 +16,7 @@ import { prepareCompaction, shouldCompact, } from "../src/core/compaction/index.js"; +import { serializeConversation } from "../src/core/compaction/utils.js"; import { buildSessionContext, type CompactionEntry, @@ -54,10 +55,10 @@ function createUserMessage(text: string): AgentMessage { return { role: "user", content: text, timestamp: Date.now() }; } -function createAssistantMessage(text: string, usage?: Usage): AssistantMessage { +function createAssistantMessage(content: AssistantMessage["content"] | string, usage?: Usage): AssistantMessage { return { role: "assistant", - content: [{ type: "text", text }], + content: typeof content === "string" ? [{ type: "text", text: content }] : content, usage: usage || createMockUsage(100, 50), stopReason: "stop", timestamp: Date.now(), @@ -629,3 +630,70 @@ describe.skipIf(!process.env.ANTHROPIC_OAUTH_TOKEN)("LLM summarization", () => { console.log("After compaction:", reloaded.messages.length); }, 60000); }); + +function toolResult(text: string, toolName = "ipython", isError = false, toolCallId = "tc1"): Message { + return { + role: "toolResult", + toolCallId, + toolName, + content: [{ type: "text", text }], + isError, + timestamp: Date.now(), + }; +} + +function toolCall(id: string, name: string, args: Record): ToolCall { + return { type: "toolCall", id, name, arguments: args }; +} + +describe("serializeConversation", () => { + it("truncates long tool results keeping head and tail within the summary budget", () => { + const head = "A".repeat(1431); + const tail = "T".repeat(500); + const result = serializeConversation([toolResult(head + "B".repeat(3069) + tail)]); + + expect(result).toContain("[Tool result (ipython)]:"); + expect(result).toContain(head); + expect(result).toContain(tail); + expect(result).toContain("[... 3069 characters truncated; first 1431 and last 500 kept ...]"); + expect(result).not.toContain("B".repeat(10)); + expect(result.length).toBeLessThanOrEqual("[Tool result (ipython)]: ".length + 2000); + }); + + it.each([ + ["labels short success results with the tool name", "bash", false, "[Tool result (bash)]"], + ["labels short error results as failed", "edit", true, "[Tool result (edit, error)]"], + ])("%s", (_label, toolName, isError, label) => { + const shortContent = "x".repeat(1500); + expect(serializeConversation([toolResult(shortContent, toolName, isError)])).toBe(`${label}: ${shortContent}`); + }); + + it.each([ + [ + "pairs repeated same-name tool calls with their results by index", + [toolCall("c1", "ipython", { code: "a" }), toolCall("c2", "ipython", { code: "b" })], + [toolResult("first output", "ipython", false, "c1"), toolResult("second output", "ipython", true, "c2")], + '[Assistant tool calls]: #1 ipython(code="a"); #2 ipython(code="b")\n\n' + + "[Tool result (ipython) #1]: first output\n\n" + + "[Tool result (ipython, error) #2]: second output", + ], + [ + "falls back to the name-only label when the result's call was not serialized", + [toolCall("c1", "bash", { command: "ls" })], + [toolResult("orphan output", "ipython", false, "tc-orphan")], + '[Assistant tool calls]: #1 bash(command="ls")\n\n[Tool result (ipython)]: orphan output', + ], + ])("%s", (_label, calls, results, expected) => { + expect(serializeConversation([createAssistantMessage(calls), ...results])).toBe(expected); + }); + + it("does not truncate user or assistant messages", () => { + const longText = "y".repeat(5000); + const result = serializeConversation([ + { role: "user", content: [{ type: "text", text: longText }], timestamp: Date.now() }, + createAssistantMessage(longText), + ]); + expect(result).not.toContain("truncated"); + expect(result).toContain(longText); + }); +}); From cd0f467a4aa06970f48dbfccbc8583f9ac5bb79d Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Mon, 21 Sep 2026 09:05:18 -0400 Subject: [PATCH 12/26] [RSI, bug] compaction: anchor summaries to kept-tail state and stop re-summarizing file lists (#2385) * fix(coding-agent): anchor compaction summaries to kept-tail state and stop re-summarizing file lists P1 (stale-summary drift): the summarizer only saw messages before the cut point, so post-compaction summaries described pre-tail state. prepareCompaction now extracts the newest kept-tail assistant text as a recency anchor, the summarizer prompt carries it as (retained messages below are authoritative), and the in-context [compaction-summary] prefix line states that the retained messages below are authoritative over the summary. P2 (compounding file lists): stored summaries end with mechanically appended / blocks that rode the previous summary back into the update prompt, so the model re-summarized lists that compound across repeated compactions. prepareCompaction strips the blocks before the update path (entry details plus the fresh append stay the single source of truth) and computeFileLists caps the combined lists at 6000 characters, dropping read-only entries first. * test(compaction): compress recency-anchor coverage to cross-layer essentials * test(compaction): execute owner-authorized cut menu to fit the test-line budget (cherry picked from commit 27f32ddb7e99803555c4654c9ce38edad3e90e22) --- .../.changes/compaction-anchor-file-lists.md | 2 + .../src/core/compaction/compaction.ts | 54 ++++++++++- .../coding-agent/src/core/compaction/utils.ts | 91 ++++++++++++++++++- packages/coding-agent/src/core/messages.ts | 3 +- packages/coding-agent/test/compaction.test.ts | 41 +++++++++ 5 files changed, 184 insertions(+), 7 deletions(-) create mode 100644 packages/coding-agent/.changes/compaction-anchor-file-lists.md diff --git a/packages/coding-agent/.changes/compaction-anchor-file-lists.md b/packages/coding-agent/.changes/compaction-anchor-file-lists.md new file mode 100644 index 00000000000..4d32d740e5f --- /dev/null +++ b/packages/coding-agent/.changes/compaction-anchor-file-lists.md @@ -0,0 +1,2 @@ +- Fixed compaction summaries drifting behind the retained conversation: the summarizer now receives the newest kept-tail assistant text as a `` and the in-context `[compaction-summary]` prefix states that the retained messages below are authoritative. +- Fixed compaction file lists compounding through repeated summaries: ``/`` blocks are stripped from the previous summary before the update prompt (entry details plus the fresh append remain the single source), and the combined file lists are capped at 6000 characters, dropping read-only entries first. diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index 8cb58096f34..4b4483c0387 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -27,6 +27,7 @@ import { formatFileOperations, SUMMARIZATION_SYSTEM_PROMPT, serializeConversation, + stripFileListBlocks, } from "./utils.js"; /** Details stored in CompactionEntry.details for file tracking */ export interface CompactionDetails { @@ -525,6 +526,9 @@ export function buildSummarizationPrompt(customInstructions?: string, previousSu * `onPayload` carries the owning session's provider identity so the call is * not anonymous at the provider; it is expected to be scoped, because the * summarization prompt is not the session's own conversation. + * If recentStateAnchor is provided (newest retained assistant text), the + * prompt marks it as the current state so the summary cannot lag behind the + * kept tail. */ export async function generateSummary( currentMessages: AgentMessage[], @@ -539,6 +543,7 @@ export async function generateSummary( retry?: ProviderRetryPolicy, onPayload?: ProviderPayloadHook, sessionId?: string, + recentStateAnchor?: string, ): Promise { const maxTokens = Math.floor(0.8 * reserveTokens); @@ -550,6 +555,9 @@ export async function generateSummary( if (previousSummary) { promptText += `\n${previousSummary}\n\n\n`; } + if (recentStateAnchor) { + promptText += `\nNewest assistant message that stays retained below the summary. The conversation to summarize is older than this anchor; the retained messages below are authoritative, so treat this anchor, not the conversation above, as the current state.\n\n${recentStateAnchor}\n\n\n`; + } promptText += basePrompt; const summarizationMessages = [ @@ -596,8 +604,10 @@ export interface CompactionPreparation { /** Whether this is a split turn (cut point in middle of turn) */ isSplitTurn: boolean; tokensBefore: number; - /** Summary from previous compaction, for iterative update */ + /** Summary from previous compaction, for iterative update (file-list blocks stripped) */ previousSummary?: string; + /** Newest retained assistant text; anchors the summary to kept-tail state */ + recentStateAnchor?: string; /** File operations extracted from messagesToSummarize */ fileOps: FileOperations; /** Compaction settions from settings.jsonl */ @@ -624,7 +634,10 @@ export function prepareCompaction( let boundaryStart = 0; if (prevCompactionIndex >= 0) { const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry; - previousSummary = prevCompaction.summary; + // File-list blocks never reach the update prompt: they are re-appended + // mechanically below and compound when the model re-summarizes them. + const strippedSummary = stripFileListBlocks(prevCompaction.summary); + previousSummary = strippedSummary.length > 0 ? strippedSummary : undefined; const firstKeptEntryIndex = pathEntries.findIndex((entry) => entry.id === prevCompaction.firstKeptEntryId); boundaryStart = firstKeptEntryIndex >= 0 ? firstKeptEntryIndex : prevCompactionIndex + 1; } @@ -653,6 +666,11 @@ export function prepareCompaction( } } + // Recency anchor: the summarizer sees only messages before the cut, so its + // summary would describe pre-tail state. The newest retained assistant + // text is the state the next turn actually sees; pass it to the summarizer. + const recentStateAnchor = extractRecentStateAnchor(pathEntries, cutPoint.firstKeptEntryIndex, pathEntries.length); + // Avoid a compaction that would summarize no history. if (messagesToSummarize.length === 0 && turnPrefixMessages.length === 0 && !previousSummary) { return undefined; @@ -672,10 +690,39 @@ export function prepareCompaction( isSplitTurn: cutPoint.isSplitTurn, tokensBefore, previousSummary, + recentStateAnchor, fileOps, settings, }; } + +/** + * Maximum characters kept from the retained tail for the recency anchor. + * The end of a message holds the newest state, so long text keeps its tail. + */ +const RECENT_STATE_ANCHOR_MAX_CHARS = 2000; + +/** + * Extract the newest retained assistant text (the recency anchor) from the + * kept tail [keptStart, keptEnd). Returns undefined when the tail has no + * assistant text; long text is tail-truncated to the anchor budget. + */ +function extractRecentStateAnchor(entries: SessionEntry[], keptStart: number, keptEnd: number): string | undefined { + for (let i = keptEnd - 1; i >= keptStart; i--) { + const msg = getMessageFromEntryForCompaction(entries[i]); + if (!msg || msg.role !== "assistant" || !("content" in msg) || !Array.isArray(msg.content)) continue; + const text = msg.content + .filter((block): block is { type: "text"; text: string } => block.type === "text") + .map((block) => block.text) + .join("\n") + .trim(); + if (!text) continue; + return text.length > RECENT_STATE_ANCHOR_MAX_CHARS + ? text.slice(text.length - RECENT_STATE_ANCHOR_MAX_CHARS) + : text; + } + return undefined; +} const TURN_PREFIX_SUMMARIZATION_PROMPT = `This is the PREFIX of a turn that was too large to keep. The SUFFIX (recent work) is retained. Summarize the prefix to provide context for the retained suffix: @@ -723,6 +770,7 @@ export async function compact( isSplitTurn, tokensBefore, previousSummary, + recentStateAnchor, fileOps, settings, } = preparation; @@ -747,6 +795,7 @@ export async function compact( retry, onPayload, sessionId, + recentStateAnchor, ), ) : Promise.resolve({ summary: "No prior history." }), @@ -782,6 +831,7 @@ export async function compact( retry, onPayload, sessionId, + recentStateAnchor, ), ); slices.push(result); diff --git a/packages/coding-agent/src/core/compaction/utils.ts b/packages/coding-agent/src/core/compaction/utils.ts index a582b66ba74..e953e617ac0 100644 --- a/packages/coding-agent/src/core/compaction/utils.ts +++ b/packages/coding-agent/src/core/compaction/utils.ts @@ -19,9 +19,14 @@ export function createFileOps(): FileOperations { } /** - * Extract file operations from tool calls in an assistant message. + * Extract file operations from tool calls in an assistant message and from + * structured tool results. */ export function extractFileOpsFromMessage(message: AgentMessage, fileOps: FileOperations): void { + if (message.role === "toolResult") { + extractFileOpsFromToolResult(message, fileOps); + return; + } if (message.role !== "assistant") return; if (!("content" in message) || !Array.isArray(message.content)) return; @@ -44,15 +49,93 @@ export function extractFileOpsFromMessage(message: AgentMessage, fileOps: FileOp } } +/** + * Extract file operations from a tool result message. + * + * The default toolset routes file edits through the ipython kernel: the + * kernel's edit skill reports structured diff displays (path, oldStr, newStr) + * that ride on the tool result's details, and no assistant-side tool call + * ever carries the path. Without this branch, compaction summaries never + * learn about kernel-performed edits and stays empty in the + * default configuration. + */ +function extractFileOpsFromToolResult(message: AgentMessage, fileOps: FileOperations): void { + if (message.role !== "toolResult" || message.toolName !== "ipython") return; + const details = + typeof message.details === "object" && message.details !== null && !Array.isArray(message.details) + ? (message.details as Record) + : {}; + const diffs = Array.isArray(details.diffs) ? details.diffs : []; + for (const diff of diffs) { + if (typeof diff !== "object" || diff === null || Array.isArray(diff)) continue; + const path = (diff as Record).path; + if (typeof path === "string" && path) fileOps.edited.add(path); + } +} + +/** + * Maximum files kept per summary block, so a single oversized kernel + * result cannot produce a file list larger than the model context limit. + */ +const FILE_LIST_MAX_ENTRIES = 200; + +/** + * Maximum combined characters the two file blocks may add to a summary. + * Repeated compactions merge lists carried in the previous entry's details, + * so without a character cap the appended block grows without bound. + */ +const FILE_LIST_MAX_COMBINED_CHARS = 6000; + +function fileListChars(readFiles: string[], modifiedFiles: string[]): number { + let chars = 0; + for (const file of readFiles) chars += file.length + 1; + for (const file of modifiedFiles) chars += file.length + 1; + return chars; +} + /** * Compute final file lists from file operations. * Returns readFiles (files only read, not modified) and modifiedFiles. + * Both lists are capped at FILE_LIST_MAX_ENTRIES (sorted, then truncated) and + * at FILE_LIST_MAX_COMBINED_CHARS combined characters: read-only entries are + * least valuable and drop first (from the alphabetical end), then modified + * entries drop only after the read-only list is empty. */ export function computeFileLists(fileOps: FileOperations): { readFiles: string[]; modifiedFiles: string[] } { const modified = new Set([...fileOps.edited, ...fileOps.written]); - const readOnly = [...fileOps.read].filter((f) => !modified.has(f)).sort(); - const modifiedFiles = [...modified].sort(); - return { readFiles: readOnly, modifiedFiles }; + const readOnly = [...fileOps.read] + .filter((f) => !modified.has(f)) + .sort() + .slice(0, FILE_LIST_MAX_ENTRIES); + const modifiedFiles = [...modified].sort().slice(0, FILE_LIST_MAX_ENTRIES); + const readFiles = readOnly.slice(); + while (readFiles.length > 0 && fileListChars(readFiles, modifiedFiles) > FILE_LIST_MAX_COMBINED_CHARS) { + readFiles.pop(); + } + while ( + modifiedFiles.length > 0 && + readFiles.length === 0 && + fileListChars(readFiles, modifiedFiles) > FILE_LIST_MAX_COMBINED_CHARS + ) { + modifiedFiles.pop(); + } + return { readFiles, modifiedFiles }; +} + +/** + * Remove / blocks from a stored summary. + * + * The blocks are re-appended mechanically after every summarization (see + * computeFileLists/formatFileOperations) and carried in the compaction entry's + * details. Feeding stale blocks back into the update prompt makes the model + * re-summarize them, so lists compound across repeated compactions. Strip + * them before a previous summary reaches the summarizer; the details plus the + * fresh append remain the single source of truth. + */ +const FILE_LIST_BLOCK_PATTERN = /(?:\n*)<(read-files|modified-files)>[\s\S]*?<\/\1>/g; + +export function stripFileListBlocks(summary: string): string { + return summary.replace(FILE_LIST_BLOCK_PATTERN, "").trimEnd(); } /** diff --git a/packages/coding-agent/src/core/messages.ts b/packages/coding-agent/src/core/messages.ts index 50f9598f554..892551247ca 100644 --- a/packages/coding-agent/src/core/messages.ts +++ b/packages/coding-agent/src/core/messages.ts @@ -18,7 +18,8 @@ import { isSessionSlashCommandName, parseSessionSlashCommand, type SessionSlashC export const COMPACTION_SUMMARY_PREFIX = `[compaction-summary] -The conversation history before this point was compacted into the following summary: +The conversation history before this point was compacted into the following summary. +The retained messages below are authoritative; this summary may lag behind them. `; diff --git a/packages/coding-agent/test/compaction.test.ts b/packages/coding-agent/test/compaction.test.ts index c3bb4848620..51877a9b6c4 100644 --- a/packages/coding-agent/test/compaction.test.ts +++ b/packages/coding-agent/test/compaction.test.ts @@ -541,6 +541,47 @@ describe("prepareCompaction with previous compaction", () => { }); }); +describe("prepareCompaction recency anchor", () => { + // A long user message crosses the tiny keep-recent budget, so the cut + // lands on it: a deterministic user-message cut with no split turn. + const longUserText = `user tail ${"x".repeat(400)}`; + const anchoredPreparation = (previousSummary: string, tailTexts: string[]) => { + const u2 = createMessageEntry(createUserMessage("user msg 2")); + const entries = [ + createMessageEntry(createUserMessage("user msg 1")), + createMessageEntry(createAssistantMessage("assistant msg 1", createMockUsage(5000, 1000))), + u2, + createMessageEntry(createAssistantMessage("assistant msg 2", createMockUsage(6000, 2000))), + createCompactionEntry(previousSummary, u2.id), + createMessageEntry(createUserMessage(longUserText)), + ...tailTexts.map((text) => createMessageEntry(createAssistantMessage(text, createMockUsage(7000, 3000)))), + ]; + return prepareCompaction(entries, { ...DEFAULT_COMPACTION_SETTINGS, keepRecentTokens: 10 }); + }; + + it.each([ + { + name: "anchors to the newest kept-tail assistant text, skipping empty ones", + previousSummary: "First summary", + tailTexts: ["", "older kept text", "newest kept text"], + anchor: "newest kept text", + keptSummary: "First summary", + }, + { + name: "drops the previous summary entirely when it contained only file blocks", + previousSummary: "\nold/a.ts\n\n\n\nold/b.ts\n", + tailTexts: [], + anchor: undefined, + keptSummary: undefined, + }, + ])("$name", ({ previousSummary, tailTexts, anchor, keptSummary }) => { + const preparation = anchoredPreparation(previousSummary, tailTexts); + expect(preparation).toBeDefined(); + expect(preparation!.recentStateAnchor).toBe(anchor); + expect(preparation!.previousSummary).toBe(keptSummary); + }); +}); + // ============================================================================ // Integration tests with real session data // ============================================================================ From 600bd30326ce3e568c95c9f446d874bd422ebf9d Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:02:43 -0400 Subject: [PATCH 13/26] fix(harness): validate refinement writes and skip malformed entries in the digest (#2463) * fix(harness): validate refinement writes and skip malformed entries in the digest A harness entry with non-string content (the incident: content persisted as a list by the kernel CRUD path) crashed the daemon digest builder on every session build with "text.replace is not a function", bricking all child spawns. Validate entry shape at the single write choke points (kernel HarnessState._upsert for rlm.harness CRUD, validateEdit for refine/rollback edits) so invalid entries are rejected with an error naming the entry and the field instead of being saved, and make every harness renderer (digest, refine overview, notice body) skip malformed entries and refinement events with a "harness: skipped malformed entry (...)" diagnostic so one corrupt store entry can never break session creation. * fix(harness): guard non-object refinement events and event ids Review round 1 (Bugbot persona) found two residual gaps in the same crash class: a non-object element in the persisted refinements array still crashed the digest and its fingerprint (event.trigger / event.id dereferenced null), and record_refinement persisted a non-string event id that the kernel loader then silently dropped. Skip non-object refinement events with an "event not an object" diagnostic, filter them out of the digest fingerprint, and reject non-string/empty refinement event ids at write time. * fix(harness): skip refinement events with non-string ids and change elements Review round 2 (Macroscope) found the digest still rendered refinement events whose id or change elements were not strings: a corrupt store event like {"id":null,"trigger":"t","changes":[{"a":1}]} passed harnessRefinementMalformation and rendered junk ("- [null] t: [object Object]") into every session's prompt digest, and malformedRefinementEventLabel interpolated the unvalidated id value. Persisted stores are JSON, so these shapes cannot throw (only non-JSON values like Symbols can, which JSON.parse never produces), but the skip-malformed contract requires string-validated rendering: extend harnessRefinementMalformation with "id not a string" and "changes contain a non-string", and label invalid ids by bounded type instead of interpolating the value. The digest and fingerprint share the predicate, so both skip paths stay render-equal. (cherry picked from commit 41e4e0e1868f5465cfccd4b02bc5e3141710309d) --- .../.changes/harness-refinement-validation.md | 2 + .../src/core/refinement/refinement.ts | 85 ++++++- packages/coding-agent/test/refinement.test.ts | 209 ++++++++++++++---- .../test/suite/agent-session-prompt.test.ts | 47 ++++ prime-agent-runtime/src/rlm/harness.py | 146 +++++++++++- prime-agent-runtime/test/test_harness.py | 53 +++++ 6 files changed, 491 insertions(+), 51 deletions(-) create mode 100644 packages/coding-agent/.changes/harness-refinement-validation.md diff --git a/packages/coding-agent/.changes/harness-refinement-validation.md b/packages/coding-agent/.changes/harness-refinement-validation.md new file mode 100644 index 00000000000..3151be60f35 --- /dev/null +++ b/packages/coding-agent/.changes/harness-refinement-validation.md @@ -0,0 +1,2 @@ +- Harness store writes now validate entry shape before persisting: `rlm.harness` create/update/upsert calls reject non-string or empty `title`/`content`, empty or non-string ids, malformed paths, non-dict `reference`/`arguments`/`metadata`, and skill entries without a valid Python reference, with a clear error naming the entry and the violated entry field; refinement events reject non-string triggers and invalid `changes`/ids. `/refine` and rollback edits reject the same shapes through `validateEdit`. +- The harness digest and refine overview now skip malformed persisted entries and refinement events (non-object events plus events with non-string ids, triggers, change elements, or outcomes, labeled by bounded type instead of value) with a `harness: skipped malformed entry (...)` diagnostic line instead of crashing session creation, so a single corrupt store entry can no longer brick every session and child spawn. diff --git a/packages/coding-agent/src/core/refinement/refinement.ts b/packages/coding-agent/src/core/refinement/refinement.ts index c1378f449c2..96ca727b4d4 100644 --- a/packages/coding-agent/src/core/refinement/refinement.ts +++ b/packages/coding-agent/src/core/refinement/refinement.ts @@ -252,7 +252,9 @@ function emptyHarnessState(): HarnessState { } function slug(raw: string, fallback: string): string { - const normalized = raw + // A malformed value (for example a non-string title) cannot be normalized; resolve + // to the fallback so apply-time validation can still reject the edit by id. + const normalized = (typeof raw === "string" ? raw : fallback) .trim() .toLowerCase() .replace(/[^a-z0-9]+/g, "_") @@ -441,6 +443,42 @@ export function mergeRefinementHistory( return [...byId.values()]; } +/** Why a persisted harness entry cannot be rendered safely: the field whose + * stored type violates the entry contract (write paths reject these shapes). + * Render paths skip such entries with a diagnostic instead of throwing, so one + * corrupt entry (from an older build or a hand-edited store) can never break + * session creation by crashing the harness digest. */ +export function harnessEntryMalformation(entry: HarnessEntry): string | undefined { + if (typeof entry.content !== "string") return "content not a string"; + if (typeof entry.title !== "string") return "title not a string"; + return undefined; +} + +/** Same contract for refinement events: the digest renders id, trigger, changes, + * and outcome with string operations, so a non-string id or trigger, non-array + * changes, non-string change elements, or non-string outcome must be skipped + * with a diagnostic rather than crash the digest or render junk. */ +export function harnessRefinementMalformation(event: HarnessRefinementEvent): string | undefined { + if (typeof event !== "object" || event === null) return "event not an object"; + if (typeof event.id !== "string") return "id not a string"; + if (typeof event.trigger !== "string") return "trigger not a string"; + if (!Array.isArray(event.changes)) return "changes not an array"; + if (!event.changes.every((change) => typeof change === "string")) return "changes contain a non-string"; + if (event.outcome !== undefined && typeof event.outcome !== "string") return "outcome not a string"; + return undefined; +} + +/** Bounded label for a skipped malformed refinement event. Non-object elements + * and invalid ids are labeled by type, never by value: a corrupt store element + * must not inject arbitrary unbounded text into every session's prompt digest. */ +function malformedRefinementEventLabel(event: HarnessRefinementEvent): string { + if (event === null) return "null"; + if (typeof event === "undefined") return "undefined"; + if (typeof event !== "object") return `a ${typeof event}`; + if (Array.isArray(event)) return "an array"; + return typeof event.id === "string" ? event.id : `a ${typeof event.id} id`; +} + function compactText(text: string, maxLength: number): string { const normalized = text.replace(/\s+/g, " ").trim(); if (normalized.length <= maxLength) { @@ -456,11 +494,12 @@ export function formatRefinementNoticeBody(result: RefinementResult): string { if (!edit.applied) continue; const entry = edit.after ?? edit.before; const scope = entry?.scope ?? result.scope ?? "local"; + const malformation = entry ? harnessEntryMalformation(entry) : undefined; lines.push( - `- ${edit.action} ${edit.kind} [${scope}:${edit.id}] ${entry?.title ?? edit.id}: ${compactText( - entry?.content ?? "", + `- ${edit.action} ${edit.kind} [${scope}:${edit.id}] ${malformation ? edit.id : (entry?.title ?? edit.id)}: ${compactText( + malformation ? "" : (entry?.content ?? ""), DEFAULT_OVERVIEW_CONTENT_LIMIT, - )}`, + )}${malformation ? ` (skipped malformed entry: ${malformation})` : ""}`, ); } return lines.join("\n"); @@ -519,6 +558,11 @@ export function formatHarnessStateForPrompt( lines.push(`${kind}: ${entries.length}`); } for (const entry of entries.slice(0, maxEntriesPerKind)) { + const malformation = harnessEntryMalformation(entry); + if (malformation) { + lines.push(`harness: skipped malformed entry ${entry.id} (${malformation})`); + continue; + } const argumentsText = entry.kind === "skill" && Object.keys(entry.arguments).length > 0 ? ` args=${compactText(JSON.stringify(entry.arguments), maxContentLength)}` @@ -547,6 +591,13 @@ export function formatHarnessStateForPrompt( lines.push(`recent refinements: ${state.refinements.length}`); for (const event of state.refinements.slice(-maxRefinements)) { + const malformation = harnessRefinementMalformation(event); + if (malformation) { + lines.push( + `harness: skipped malformed refinement event ${malformedRefinementEventLabel(event)} (${malformation})`, + ); + continue; + } const changes = event.changes.length > 0 ? event.changes.join(", ") : "no applied edits"; const outcome = event.outcome ? `; outcome: ${compactText(event.outcome, maxContentLength)}` : ""; lines.push(`- [${event.id}] ${compactText(event.trigger, maxContentLength)}: ${changes}${outcome}`); @@ -565,6 +616,11 @@ function overviewForPrompt(state: HarnessState): string { const entries = Object.values(state.entries[kind]); lines.push(`${kind}: ${entries.length}`); for (const entry of entries.slice(0, 40)) { + const malformation = harnessEntryMalformation(entry); + if (malformation) { + lines.push(`- harness: skipped malformed entry ${entry.id} (${malformation})`); + continue; + } const content = entry.content.replace(/\s+/g, " ").slice(0, 240); const argumentsText = entry.kind === "skill" && Object.keys(entry.arguments).length > 0 @@ -726,6 +782,27 @@ function validateEdit(edit: RefinementEdit, computedId?: string): string | undef if (edit.action !== "delete" && (!edit.title || !edit.content)) { return `${edit.action} requires title and content`; } + if (edit.id !== undefined && (typeof edit.id !== "string" || edit.id.length === 0)) { + return `${edit.action} requires id to be a non-empty string when provided`; + } + if (edit.path !== undefined && (typeof edit.path !== "string" || edit.path.length === 0)) { + return `${edit.action} requires path to be a non-empty string when provided`; + } + if ( + edit.action !== "delete" && + (typeof edit.title !== "string" || typeof edit.content !== "string" || !edit.title || !edit.content) + ) { + return `${edit.action} requires title and content to be non-empty strings`; + } + if (edit.reference !== undefined && objectRecord(edit.reference) === undefined) { + return `${edit.action} requires reference to be an object when provided`; + } + if (edit.arguments !== undefined && objectRecord(edit.arguments) === undefined) { + return `${edit.action} requires arguments to be an object when provided`; + } + if (edit.metadata !== undefined && objectRecord(edit.metadata) === undefined) { + return `${edit.action} requires metadata to be an object when provided`; + } if (edit.action !== "delete" && edit.kind === "skill" && edit.arguments === undefined) { return `${edit.action} skill requires arguments`; } diff --git a/packages/coding-agent/test/refinement.test.ts b/packages/coding-agent/test/refinement.test.ts index 26495159edf..2886ca2e916 100644 --- a/packages/coding-agent/test/refinement.test.ts +++ b/packages/coding-agent/test/refinement.test.ts @@ -67,6 +67,10 @@ const skillReference = { callable: "run", call_pattern: "await run(...)", }; +const skillContract = { + reference: skillReference, + arguments: { input: { type: "string", required: true, description: "Task input" } }, +}; function proposal(summary: string, edits: RefinementProposal["edits"]): RefinementProposal { return { @@ -501,49 +505,174 @@ describe("harness refinement", () => { }); }); - it("requires Python references for harness-created skills", () => { + // Every rejected edit must leave harness state and the refinement log untouched. + type InvalidCase = { + label: string; + edit: RefinementProposal["edits"][number]; + error: string; + seed?: RefinementKind; + }; + const skillFieldsFor = (kind: RefinementKind) => (kind === "skill" ? skillContract : {}); + // Raw (unnormalized) edit shapes for apply-time validation; overrides carry the violation. + const editWith = ( + action: RefinementAction, + kind: RefinementKind, + id: string, + overrides: Record = {}, + ): RefinementProposal["edits"][number] => + ({ + action, + kind, + id, + title: "t", + content: "c", + ...skillFieldsFor(kind), + ...overrides, + }) as RefinementProposal["edits"][number]; + it.each([ + ...kinds.map( + (kind): InvalidCase => ({ + label: `duplicate create for ${kind}`, + seed: kind, + edit: { action: "create", kind, id: `${kind}_entry`, title: "t", content: "c", ...skillFieldsFor(kind) }, + error: "entry already exists", + }), + ), + ...kinds.map( + (kind): InvalidCase => ({ + label: `update of a missing ${kind}`, + edit: { action: "update", kind, id: `${kind}_missing`, title: "t", content: "c", ...skillFieldsFor(kind) }, + error: "entry not found", + }), + ), + ...kinds.map( + (kind): InvalidCase => ({ + label: `delete of a missing ${kind}`, + edit: { action: "delete", kind, id: `${kind}_missing` }, + error: "entry not found", + }), + ), + { + label: "create without content", + edit: { action: "create", kind: "memory", id: "missing_fields", title: "t" }, + error: "create requires title and content", + }, + { + label: "update without content", + seed: "memory", + edit: { action: "update", kind: "memory", id: "memory_entry", title: "t" }, + error: "update requires title and content", + }, + { + label: "update without an id", + edit: { action: "update", kind: "skill", title: "t", content: "c" }, + error: "update requires id", + }, + { label: "delete without an id", edit: { action: "delete", kind: "skill" }, error: "delete requires id" }, + { + label: "an unsupported action", + edit: { action: "rename" as RefinementAction, kind: "memory", id: "bad_action", title: "t", content: "c" }, + error: "unsupported action rename", + }, + { + label: "an unsupported kind", + edit: { action: "create", kind: "tool" as RefinementKind, id: "bad_kind", title: "t", content: "c" }, + error: "unsupported kind tool", + }, + { + label: "a skill without an argument contract", + edit: { + action: "create", + kind: "skill", + id: "argumentless", + title: "t", + content: "c", + reference: skillReference, + }, + error: "create skill requires arguments", + }, + { + label: "a skill without a Python reference", + edit: { action: "create", kind: "skill", id: "unbacked", title: "t", content: "c", arguments: {} }, + error: "create skill requires python reference", + }, + { + label: "a skill backed by a non-Python reference", + edit: { + action: "create", + kind: "skill", + id: "shell_skill", + title: "t", + content: "c", + reference: { type: "shell", command: "edit" }, + arguments: {}, + }, + error: "create skill reference.type must be python", + }, + { + label: "an update of the base system prompt", + edit: { action: "update", kind: "prompt", id: "base_system_prompt", title: "t", content: "c" }, + error: "base system prompt", + }, + { + label: "a create whose title derives the base system prompt id", + edit: { action: "create", kind: "prompt", title: "Base System Prompt", content: "c" }, + error: "base system prompt", + }, + { + label: "a create with list content", + edit: editWith("create", "memory", "list_content", { content: ["one string"] }), + error: "create requires title and content to be non-empty strings", + }, + { + label: "a create with list title", + edit: editWith("create", "prompt", "list_title", { title: ["t"] }), + error: "create requires title and content to be non-empty strings", + }, + { + label: "an update restoring list content", + seed: "memory", + edit: editWith("update", "memory", "memory_entry", { content: ["one string"] }), + error: "update requires title and content to be non-empty strings", + }, + { + label: "a create with a numeric id", + edit: editWith("create", "memory", "numeric_id", { id: 7 }), + error: "create requires id to be a non-empty string when provided", + }, + { + label: "a skill with a list reference", + edit: editWith("create", "skill", "list_reference", { reference: ["bad"] }), + error: "create requires reference to be an object when provided", + }, + { + label: "a create with a numeric path", + edit: editWith("create", "memory", "bad_path", { path: 7 }), + error: "create requires path to be a non-empty string when provided", + }, + { + label: "a skill with list arguments", + edit: editWith("create", "skill", "list_arguments", { arguments: ["a"] }), + error: "create requires arguments to be an object when provided", + }, + { + label: "an update with list metadata", + seed: "memory", + edit: editWith("update", "memory", "memory_entry", { metadata: ["m"] }), + error: "update requires metadata to be an object when provided", + }, + ])("rejects $label without mutating state", ({ seed, edit, error }) => { const state = loadHarnessState(makeTempDir()); + if (seed) seedEntry(state, seed); + const before = structuredClone(state.entries); - const missingReference = applyRefinementProposal( - state, - proposal("Create skill without reference", [ - { - action: "create", - kind: "skill", - id: "unbacked_skill", - title: "Unbacked skill", - content: "This should not be accepted without a Python reference.", - arguments: {}, - }, - ]), - { id: "refine_missing_skill_reference" }, - ); - const nonPythonReference = applyRefinementProposal( - state, - proposal("Create skill with non-python reference", [ - { - action: "create", - kind: "skill", - id: "shell_skill", - title: "Shell skill", - content: "This should not be accepted as a harness skill.", - reference: { type: "shell", command: "edit" }, - arguments: {}, - }, - ]), - { id: "refine_non_python_skill_reference" }, - ); + const result = applyRefinementProposal(state, proposal("Invalid edit", [edit]), { id: "refine_invalid" }); - expect(missingReference.appliedEdits[0]).toMatchObject({ - applied: false, - error: "create skill requires python reference", - }); - expect(nonPythonReference.appliedEdits[0]).toMatchObject({ - applied: false, - error: "create skill reference.type must be python", - }); - expect(state.entries.skill.unbacked_skill).toBeUndefined(); - expect(state.entries.skill.shell_skill).toBeUndefined(); + expect(result.appliedEdits).toHaveLength(1); + expect(result.appliedEdits[0].applied).toBe(false); + expect(result.appliedEdits[0].error).toContain(error); + expect(state.entries).toEqual(before); + expect(state.refinements.at(-1)?.changes).toEqual([]); }); it("uses a global harness state directory under the agent dir by default", () => { diff --git a/packages/coding-agent/test/suite/agent-session-prompt.test.ts b/packages/coding-agent/test/suite/agent-session-prompt.test.ts index 93255ea85b2..986f2dcacc1 100644 --- a/packages/coding-agent/test/suite/agent-session-prompt.test.ts +++ b/packages/coding-agent/test/suite/agent-session-prompt.test.ts @@ -2212,6 +2212,19 @@ describe("Harness digest at cold boundaries", () => { expect(harness.session.messages[0]).toMatchObject({ role: "custom", customType: HARNESS_DIGEST_CUSTOM_TYPE }); }); + function isolatedAgentDir(prefix: string): string { + const previousAgentDir = process.env.PRIME_AGENT_CODING_AGENT_DIR; + const agentDir = join(tmpdir(), `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2)}`); + mkdirSync(agentDir, { recursive: true }); + tempDirs.push(agentDir); + process.env.PRIME_AGENT_CODING_AGENT_DIR = agentDir; + onTestFinished(() => { + if (previousAgentDir === undefined) delete process.env.PRIME_AGENT_CODING_AGENT_DIR; + else process.env.PRIME_AGENT_CODING_AGENT_DIR = previousAgentDir; + }); + return agentDir; + } + it("treats tree navigation as a cold boundary with digest dedupe", async () => { // Empty global store: digest content must reflect only the local test entry. const previousAgentDir = process.env.PRIME_AGENT_CODING_AGENT_DIR; @@ -2262,6 +2275,40 @@ describe("Harness digest at cold boundaries", () => { expect(getMessageText(digests.at(-1))).toContain("[local:nav_test_memory] Nav test memory"); }); + it("delivers a diagnostic digest instead of crashing on a malformed global entry", async () => { + // Regression for the fleet-wide incident: one entry with list content in + // the global store bricked all child spawn creation via the digest crash. + const agentDir = isolatedAgentDir("pi-digest-malformed"); + mkdirSync(join(agentDir, "harness"), { recursive: true }); + writeFileSync( + join(agentDir, "harness", "harness_state.json"), + '{"schema":1,"entries":{"prompt":{},"skill":{},"subagent":{},"memory":{"broken_memory":{"id":"broken_memory","kind":"memory","title":"Breaking memory","content":["one string"],"path":"arc","scope":"global","version":1},"valid_memory":{"id":"valid_memory","kind":"memory","title":"Valid memory","content":"Worktree workflow notes.","path":"general","scope":"global","version":1}}},"refinements":[{"id":"refine_bad","trigger":["not a string"],"changes":[],"evidence":"","outcome":""},null,"RAWLEAK-5f1e",{"id":null,"trigger":"t","changes":["update memory:m"]},{"id":"bad_changes","trigger":"t","changes":[7]}]}', + ); + + const harness = await createHarness({ persistSession: true }); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("hi")]); + await harness.session.prompt("hello"); + + const digests = digestMessages(harness); + expect(digests).toHaveLength(1); + const digest = getMessageText(digests[0]); + expect(digest).toContain("harness: skipped malformed entry broken_memory (content not a string)"); + expect(digest).toContain("harness: skipped malformed refinement event refine_bad (trigger not a string)"); + expect(digest).toContain("harness: skipped malformed refinement event null (event not an object)"); + // Non-object elements are labeled by type only: the raw value must not leak. + expect(digest).toContain("harness: skipped malformed refinement event a string (event not an object)"); + expect(digest).not.toContain("RAWLEAK-5f1e"); + // Non-string ids and non-string change elements are skipped by type label, not rendered. + expect(digest).toContain("harness: skipped malformed refinement event a object id (id not a string)"); + expect(digest).toContain( + "harness: skipped malformed refinement event bad_changes (changes contain a non-string)", + ); + expect(digest).toContain("[global:valid_memory]"); + // The malformed content itself must never leak into the digest. + expect(digest).not.toContain("one string"); + }); + it("prefers the newest digest by timestamp over a retained pre-compaction digest", async () => { const harness = await createHarness({ persistSession: true }); harnesses.push(harness); diff --git a/prime-agent-runtime/src/rlm/harness.py b/prime-agent-runtime/src/rlm/harness.py index 71452e59080..8e44f9f4519 100644 --- a/prime-agent-runtime/src/rlm/harness.py +++ b/prime-agent-runtime/src/rlm/harness.py @@ -128,19 +128,121 @@ class RefinementEvent: _REFINEMENT_FIELDS = {field.name for field in fields(RefinementEvent)} -def _validate_python_skill_reference(reference: dict[str, Any] | None) -> dict[str, Any]: +def _validate_python_skill_reference(reference: dict[str, Any] | None, entry_name: str = "") -> dict[str, Any]: + # Rejections name the entry so the caller can repair the right skill; the + # suffix keeps the historical message text greppable. + prefix = f"skill entry {entry_name!r} rejected: " if entry_name else "" + + def reject(message: str) -> None: + raise ValueError(f"{prefix}{message}") + if not isinstance(reference, dict): - raise ValueError("skill entries require a Python reference") + reject("skill entries require a Python reference") normalized = dict(reference) if normalized.get("type") != "python": - raise ValueError("skill reference.type must be 'python'") + reject("skill reference.type must be 'python'") if not any(isinstance(normalized.get(key), str) and normalized[key] for key in ("import", "python_import")): - raise ValueError("skill reference requires a Python import") + reject("skill reference requires a Python import") if not any(isinstance(normalized.get(key), str) and normalized[key] for key in ("callable", "call_pattern")): - raise ValueError("skill reference requires a callable or call_pattern") + reject("skill reference requires a callable or call_pattern") return normalized +def _type_name(value: Any) -> str: + if isinstance(value, list): + return "a list" + if value == "": + return "an empty string" + return type(value).__name__ + + +def _describe_entry(id: Any, title: Any) -> str: + """Best available entry name for rejection messages.""" + if isinstance(id, str) and id: + return id + if isinstance(title, str) and title: + return title + return "" + + +def _require_text(kind: str, entry_name: str, field: str, value: Any) -> None: + if not isinstance(value, str) or not value: + raise ValueError( + f"{kind} entry {entry_name!r} rejected: {field} must be a non-empty string, got {_type_name(value)}" + ) + + +def _require_optional_text(kind: str, entry_name: str, field: str, value: Any) -> None: + if value is not None: + _require_text(kind, entry_name, field, value) + + +def _require_optional_record(kind: str, entry_name: str, field: str, value: Any) -> None: + if value is not None and not isinstance(value, dict): + raise ValueError( + f"{kind} entry {entry_name!r} rejected: {field} must be a dict when provided, got {_type_name(value)}" + ) + + +def _validate_entry_shape( + kind: str, + entry_id: Any, + title: Any, + content: Any, + *, + path: Any, + reference: Any, + arguments: Any, + metadata: Any, + source: Any, + existing: "HarnessEntry | None", +) -> None: + """Reject an invalid harness entry before anything is persisted. + + Every create/update/upsert write funnels through here, so a malformed + entry (content as a list, title as a number) fails with an actionable + error naming the entry and the field instead of being saved and later + crashing the host digest that renders every session's system prompt. + """ + entry_name = _describe_entry(entry_id, title) + _require_text(kind, entry_name, "id", entry_id) + _require_text(kind, entry_name, "title", title) + _require_text(kind, entry_name, "content", content) + _require_optional_text(kind, entry_name, "path", path) + _require_optional_record(kind, entry_name, "reference", reference) + _require_optional_record(kind, entry_name, "arguments", arguments) + _require_optional_record(kind, entry_name, "metadata", metadata) + _require_text(kind, entry_name, "source", source) + if kind == "skill": + if reference is None: + # A new skill without a Python reference is invalid; an update that + # omits it preserves the existing reference instead. + if existing is None: + raise ValueError(f"skill entry {entry_name!r} rejected: skill entries require a Python reference") + else: + _validate_python_skill_reference(reference, entry_name) + + +def _validate_refinement_event(trigger: Any, changes: Any, *, evidence: Any, outcome: Any) -> None: + """Reject a refinement event whose persisted shape would break the digest.""" + if not isinstance(trigger, str) or not trigger: + raise ValueError(f"refinement event rejected: trigger must be a non-empty string, got {_type_name(trigger)}") + if isinstance(changes, str): + if not changes: + raise ValueError("refinement event rejected: changes must be a non-empty string or a list of strings") + elif isinstance(changes, list): + if not all(isinstance(change, str) and change for change in changes): + raise ValueError("refinement event rejected: changes must be a list of non-empty strings") + else: + raise ValueError( + f"refinement event rejected: changes must be a string or a list of strings, got {_type_name(changes)}" + ) + if not isinstance(evidence, str): + raise ValueError(f"refinement event rejected: evidence must be a string when provided, got {_type_name(evidence)}") + if not isinstance(outcome, str): + raise ValueError(f"refinement event rejected: outcome must be a string when provided, got {_type_name(outcome)}") + + class HarnessState: """CRUD store for reset-free harness refinement state.""" @@ -380,8 +482,27 @@ def _upsert( if kind not in self.entries: raise ValueError(f"unknown harness kind {kind!r}; expected one of {_KINDS}") + # Guard before the id slug and the dict lookup: a non-string title or a + # non-string id (falsy ids included, which the slug fallback would + # silently collapse) must fail with a clear rejection, not an + # AttributeError inside slug normalization or a TypeError from the lookup. + _require_text(kind, _describe_entry(id, title), "title", title) + if id is not None: + _require_text(kind, _describe_entry(id, title), "id", id) entry_id = id or _slug(title, kind) existing = self.entries[kind].get(entry_id) + _validate_entry_shape( + kind, + entry_id, + title, + content, + path=path, + reference=reference, + arguments=arguments, + metadata=metadata, + source=source, + existing=existing, + ) if existing: existing.title = title existing.content = content @@ -485,6 +606,9 @@ def create( self._sync_from_disk() if kind not in self.entries: raise ValueError(f"unknown harness kind {kind!r}; expected one of {_KINDS}") + _require_text(kind, _describe_entry(id, title), "title", title) + if id is not None: + _require_text(kind, _describe_entry(id, title), "id", id) entry_id = id or _slug(title, kind) if entry_id in self.entries[kind]: raise ValueError(f"{kind} entry {entry_id!r} already exists") @@ -532,6 +656,7 @@ def update( self._sync_from_disk() if kind not in self.entries: raise ValueError(f"unknown harness kind {kind!r}; expected one of {_KINDS}") + _require_text(kind, _describe_entry(id, title), "id", id) if id not in self.entries[kind]: raise ValueError(f"{kind} entry {id!r} does not exist") return self._upsert( @@ -623,7 +748,7 @@ def create_skill( content, id=id, path=path, - reference=_validate_python_skill_reference(reference), + reference=_validate_python_skill_reference(reference, _describe_entry(id, title)), arguments=arguments, metadata=metadata, global_=global_, @@ -646,7 +771,9 @@ def update_skill( # Only validate a reference when one is supplied; omitting it preserves the # existing reference (see _upsert) rather than forcing every title/content-only # update to re-send the full Python reference. - validated_reference = _validate_python_skill_reference(reference) if reference is not None else None + validated_reference = ( + _validate_python_skill_reference(reference, _describe_entry(id, title)) if reference is not None else None + ) return self.update( "skill", id, @@ -707,6 +834,11 @@ def record_refinement( return target.record_refinement(trigger, changes, evidence=evidence, outcome=outcome, id=id) self._ensure_local_writable() self._sync_from_disk() + _validate_refinement_event(trigger, changes, evidence=evidence, outcome=outcome) + if id is not None and (not isinstance(id, str) or not id): + raise ValueError( + f"refinement event rejected: id must be a non-empty string when provided, got {_type_name(id)}" + ) event_id = id or f"refine_{len(self.refinements) + 1:04d}" normalized_changes = [changes] if isinstance(changes, str) else list(changes) event = RefinementEvent( diff --git a/prime-agent-runtime/test/test_harness.py b/prime-agent-runtime/test/test_harness.py index 49cc7c06f11..e695a4a83f6 100644 --- a/prime-agent-runtime/test/test_harness.py +++ b/prime-agent-runtime/test/test_harness.py @@ -362,6 +362,59 @@ def test_skill_references_must_be_python(self) -> None: arguments={}, ) + def test_rejects_invalid_entry_fields_before_persisting(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + state_path = Path(temp_dir) / "harness_state.json" + state = HarnessState(state_path) + state.create_memory("Valid", "seed content", id="valid_entry") + seeded = state_path.read_text(encoding="utf-8") + + for label, kwargs, message in [ + ("content list", dict(title="T", content=["one string"]), "content must be a non-empty string, got a list"), + ("title list", dict(title=["T"], content="c"), "title must be a non-empty string, got a list"), + ("empty title", dict(title="", content="c"), "title must be a non-empty string, got an empty string"), + ("numeric id", dict(title="T", content="c", id=7), "id must be a non-empty string, got int"), + ("unhashable list id", dict(title="T", content="c", id=["x"]), "id must be a non-empty string, got a list"), + ("falsy numeric id", dict(title="Zero", content="c", id=0), "id must be a non-empty string, got int"), + ("numeric path", dict(title="T", content="c", path=7), "path must be a non-empty string, got int"), + ("list metadata", dict(title="T", content="c", metadata=["m"]), "metadata must be a dict when provided, got a list"), + ]: + with self.subTest(case=f"create {label}"): + with self.assertRaisesRegex(ValueError, message): + state.create_memory(**kwargs) + with self.subTest(case="subagent with list content"): + with self.assertRaisesRegex(ValueError, "content must be a non-empty string, got a list"): + state.create_subagent("T", ["one string"]) + with self.subTest(case="skill create without a reference"): + with self.assertRaisesRegex(ValueError, "skill entries require a Python reference"): + state.create("skill", "Skill", "content", id="orphan_skill") + with self.subTest(case="skill reference as a list"): + with self.assertRaisesRegex( + ValueError, "skill entry 'Skill' rejected: skill entries require a Python reference" + ): + state.create_skill("Skill", "content", reference=["bad"]) + with self.subTest(case="update with list content"): + with self.assertRaisesRegex(ValueError, "content must be a non-empty string, got a list"): + state.update_memory("valid_entry", "Valid", ["one string"]) + self.assertEqual(state_path.read_text(encoding="utf-8"), seeded) + with self.subTest(case="invalid refinement events"): + for label, kwargs, message in [ + ("trigger list", dict(trigger=["t"], changes="ok"), "trigger must be a non-empty string, got a list"), + ("event id list", dict(trigger="t", changes="ok", id=["x"]), "id must be a non-empty string when provided, got a list"), + ("changes int", dict(trigger="t", changes=7), "changes must be a string or a list of strings, got int"), + ("mixed changes", dict(trigger="t", changes=["ok", 2]), "changes must be a list of non-empty strings"), + ("evidence list", dict(trigger="t", changes="ok", evidence=["e"]), "evidence must be a string"), + ("outcome int", dict(trigger="t", changes="ok", outcome=7), "outcome must be a string"), + ]: + with self.subTest(case=label): + with self.assertRaisesRegex(ValueError, message): + state.record_refinement(**kwargs) + + reloaded = HarnessState(state_path) + self.assertEqual([entry.id for entry in reloaded.list("memory")], ["valid_entry"]) + self.assertEqual(reloaded.refinements, []) + self.assertIsNone(reloaded.get("skill", "orphan_skill")) + def test_load_tolerates_corrupt_or_non_object_state(self) -> None: for payload in ("not json at all", "null", "[]", '"a string"', "123"): with tempfile.TemporaryDirectory() as temp_dir: From cfe95e5a27868f656d294f37c5778624cfb022ae Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:01:59 -0400 Subject: [PATCH 14/26] [RSI, bug] fix(ai): recover stale Codex chains after metadata (#2374) * [RSI] fix(ai): recover stale Codex chains after metadata * test(ai): consolidate the paired open-loop SSE tests to fit the test-line budget Fold the two copy-paste "even when the SSE body stays open" tests into one table-driven case (both stop reasons keep their own row), drop their 1s wall-clock race in favour of the suite timeout, and pin the responseId reset on the chain-reset path by seeding the stale-metadata attempt in the retry-failure test. Net test lines drop to +2 against 13 added source lines. (cherry picked from commit e683fcb53db67c2e74fad65063fe914e6a01dba5) --- .../fix-codex-chain-reset-metadata.md | 1 + .../src/providers/openai-codex-responses.ts | 28 +++- packages/ai/test/openai-codex-stream.test.ts | 146 +++++++++--------- 3 files changed, 98 insertions(+), 77 deletions(-) create mode 100644 packages/ai/.changes/fix-codex-chain-reset-metadata.md diff --git a/packages/ai/.changes/fix-codex-chain-reset-metadata.md b/packages/ai/.changes/fix-codex-chain-reset-metadata.md new file mode 100644 index 00000000000..625d23d1d6d --- /dev/null +++ b/packages/ai/.changes/fix-codex-chain-reset-metadata.md @@ -0,0 +1 @@ +- Fixed Codex stale-continuation recovery when WebSocket lifecycle or metadata events arrive before the provider error. diff --git a/packages/ai/src/providers/openai-codex-responses.ts b/packages/ai/src/providers/openai-codex-responses.ts index a3af0629156..f373812c9f0 100644 --- a/packages/ai/src/providers/openai-codex-responses.ts +++ b/packages/ai/src/providers/openai-codex-responses.ts @@ -174,10 +174,13 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses" return; } catch (error) { const aborted = options?.signal?.aborted; - // Only reset the chain while nothing was streamed yet: after the - // first event the retry would duplicate "start"/content events. + // Only reset the chain before visible output starts; retrying after + // content events would duplicate streamed assistant output. if (!aborted && !websocketStarted && !chainResetRetried && isStaleCodexContinuationError(error)) { chainResetRetried = true; + // A failed attempt may have supplied response metadata before + // rejecting the continuation. Do not retain that dead anchor. + delete output.responseId; continue; } if (aborted || isCodexNonTransportError(error)) { @@ -1097,7 +1100,19 @@ function buildCachedWebSocketRequestBody(entry: CachedWebSocketConnection, body: }; } -async function* startWebSocketOutputOnFirstEvent( +function isCodexVisibleResponseEvent(event: ResponseStreamEvent): boolean { + return ( + event.type === "response.completed" || + event.type === "response.incomplete" || + event.type.startsWith("response.output_") || + event.type.startsWith("response.reasoning_") || + event.type.startsWith("response.content_") || + event.type.startsWith("response.refusal.") || + event.type.startsWith("response.function_") + ); +} + +async function* startWebSocketOutputOnFirstVisibleEvent( events: AsyncIterable, output: AssistantMessage, stream: AssistantMessageEventStream, @@ -1105,7 +1120,10 @@ async function* startWebSocketOutputOnFirstEvent( ): AsyncGenerator { let started = false; for await (const event of events) { - if (!started) { + // Codex can emit lifecycle, telemetry, or vendor metadata before rejecting + // a stale continuation. Keep the attempt retryable until an event can + // produce output consumed by processResponsesStream. + if (!started && isCodexVisibleResponseEvent(event)) { started = true; onStart(); stream.push({ type: "start", partial: output }); @@ -1152,7 +1170,7 @@ async function processWebSocketStream( try { socket.send(JSON.stringify({ type: "response.create", ...requestBody })); await processResponsesStream( - startWebSocketOutputOnFirstEvent( + startWebSocketOutputOnFirstVisibleEvent( mapCodexEvents(parseWebSocket(socket, options?.signal)), output, stream, diff --git a/packages/ai/test/openai-codex-stream.test.ts b/packages/ai/test/openai-codex-stream.test.ts index 2f1645357a0..79cb69df61b 100644 --- a/packages/ai/test/openai-codex-stream.test.ts +++ b/packages/ai/test/openai-codex-stream.test.ts @@ -8,7 +8,7 @@ import { streamOpenAICodexResponses, streamSimpleOpenAICodexResponses, } from "../src/providers/openai-codex-responses.js"; -import type { AssistantMessage, Context, Model } from "../src/types.js"; +import type { AssistantMessage, AssistantMessageEvent, Context, Model } from "../src/types.js"; const originalFetch = global.fetch; const originalWebSocket = globalThis.WebSocket; @@ -193,71 +193,23 @@ describe("openai-codex streaming", () => { expect(sawDone).toBe(true); }); - it("completes after response.completed even when the SSE body stays open", async () => { - const tempDir = mkdtempSync(join(tmpdir(), "pi-codex-stream-")); - process.env.PI_CODING_AGENT_DIR = tempDir; - const token = mockToken(); - const encoder = new TextEncoder(); - const sse = buildSSEPayload({ status: "completed", includeDone: true }); - - const stream = new ReadableStream({ - start(controller) { - controller.enqueue(encoder.encode(sse)); - }, - }); - - global.fetch = vi.fn(async (input: string | URL) => { - const url = typeof input === "string" ? input : input.toString(); - if (url === "https://api.github.com/repos/openai/codex/releases/latest") { - return new Response(JSON.stringify({ tag_name: "rust-v0.0.0" }), { status: 200 }); - } - if (url.startsWith("https://raw.githubusercontent.com/openai/codex/")) { - return new Response("PROMPT", { status: 200, headers: { etag: '"etag"' } }); - } - if (url === "https://chatgpt.com/backend-api/codex/responses") { - return new Response(stream, { - status: 200, - headers: { "content-type": "text/event-stream" }, - }); - } - return new Response("not found", { status: 404 }); - }) as typeof fetch; - - const model: Model<"openai-codex-responses"> = { - id: "gpt-5.1-codex", - name: "GPT-5.1 Codex", - api: "openai-codex-responses", - provider: "openai-codex", - baseUrl: "https://chatgpt.com/backend-api", - reasoning: true, - input: ["text"], - cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, - contextWindow: 400000, - maxTokens: 128000, - }; - - const context: Context = { - systemPrompt: "You are a helpful assistant.", - messages: [{ role: "user", content: "Say hello", timestamp: Date.now() }], - }; - - const result = await Promise.race([ - streamOpenAICodexResponses(model, context, { apiKey: token, transport: "sse" }).result(), - new Promise((_, reject) => { - setTimeout(() => reject(new Error("Timed out waiting for completed SSE stream")), 1000); - }), - ]); - - expect(result.content.find((c) => c.type === "text")?.text).toBe("Hello"); - expect(result.stopReason).toBe("stop"); - }); - - it("maps response.incomplete to stopReason length even when the SSE body stays open", async () => { + it.each([ + { + label: "completes after response.completed", + payload: { status: "completed" as const, includeDone: true }, + stopReason: "stop" as const, + }, + { + label: "maps response.incomplete to stopReason length", + payload: { status: "incomplete" as const }, + stopReason: "length" as const, + }, + ])("$label even when the SSE body stays open", async ({ payload, stopReason }) => { const tempDir = mkdtempSync(join(tmpdir(), "pi-codex-stream-")); process.env.PI_CODING_AGENT_DIR = tempDir; const token = mockToken(); const encoder = new TextEncoder(); - const sse = buildSSEPayload({ status: "incomplete" }); + const sse = buildSSEPayload(payload); const stream = new ReadableStream({ start(controller) { @@ -300,15 +252,12 @@ describe("openai-codex streaming", () => { messages: [{ role: "user", content: "Say hello", timestamp: Date.now() }], }; - const result = await Promise.race([ - streamOpenAICodexResponses(model, context, { apiKey: token, transport: "sse" }).result(), - new Promise((_, reject) => { - setTimeout(() => reject(new Error("Timed out waiting for incomplete SSE stream")), 1000); - }), - ]); + // The SSE body never closes: the terminal event must end the stream on its own, so + // the suite timeout is the only bound (no wall-clock race). + const result = await streamOpenAICodexResponses(model, context, { apiKey: token, transport: "sse" }).result(); expect(result.content.find((c) => c.type === "text")?.text).toBe("Hello"); - expect(result.stopReason).toBe("length"); + expect(result.stopReason).toBe(stopReason); }); it("sets session_id/x-client-request-id headers and prompt_cache_key when sessionId is provided", async () => { @@ -1170,14 +1119,66 @@ describe("openai-codex streaming", () => { }); }); + it("recovers a stale previous_response_id after lifecycle and metadata events without exposing an error", async () => { + const token = mockToken(); + const sentBodies = installScriptedCodexWebSocket([ + (socket) => socket.emit(codexResponseEvents({ responseId: "resp_1", messageId: "msg_1", text: "Hello" })), + (socket) => + socket.emit([ + { type: "response.created", response: { id: "resp_stale" } }, + { type: "response.in_progress", response: { id: "resp_stale" } }, + { type: "codex.response.metadata", headers: {} }, + { type: "responsesapi.websocket_timing", elapsed_ms: 1 }, + ...codexErrorEvents("previous_response_not_found", "Previous response with id 'resp_1' not found"), + ]), + (socket) => socket.emit(codexResponseEvents({ responseId: "resp_2", messageId: "msg_2", text: "Done" })), + ]); + + const model = codexTestModel(); + const firstContext: Context = { + systemPrompt: "You are a helpful assistant.", + messages: [{ role: "user", content: "Say hello", timestamp: 1 }], + }; + const first = await streamOpenAICodexResponses(model, firstContext, { + apiKey: token, + sessionId: "session-chain-reset-lifecycle", + transport: "websocket-cached", + }).result(); + + const secondContext: Context = { + systemPrompt: "You are a helpful assistant.", + messages: [...firstContext.messages, first, { role: "user", content: "Now finish", timestamp: 2 }], + }; + const secondStream = streamOpenAICodexResponses(model, secondContext, { + apiKey: token, + sessionId: "session-chain-reset-lifecycle", + transport: "websocket-cached", + }); + const secondEvents: AssistantMessageEvent[] = []; + for await (const event of secondStream) secondEvents.push(event); + const second = await secondStream.result(); + + expect(second.stopReason).toBe("stop"); + expect(second.content[0]).toMatchObject({ type: "text", text: "Done" }); + expect(second.responseId).toBe("resp_2"); + expect(sentBodies).toHaveLength(3); + expect(sentBodies[1].previous_response_id).toBe("resp_1"); + expect(sentBodies[2].previous_response_id).toBeUndefined(); + expect(secondEvents.filter((event) => event.type === "start")).toHaveLength(1); + expect(secondEvents.at(-1)?.type).toBe("done"); + expect(global.fetch).not.toHaveBeenCalled(); + }); + it("surfaces the error without further retries when the chain-reset retry fails again", async () => { const token = mockToken(); const sentBodies = installScriptedCodexWebSocket([ (socket) => socket.emit(codexResponseEvents({ responseId: "resp_1", messageId: "msg_1", text: "Hello" })), (socket) => - socket.emit( - codexErrorEvents("previous_response_not_found", "Previous response with id 'resp_1' not found"), - ), + socket.emit([ + // Metadata arriving before the stale rejection must not survive as the anchor. + { type: "response.created", response: { id: "resp_stale" } }, + ...codexErrorEvents("previous_response_not_found", "Previous response with id 'resp_1' not found"), + ]), (socket) => socket.emit( codexErrorEvents("previous_response_not_found", "Previous response with id 'resp_9' not found"), @@ -1207,6 +1208,7 @@ describe("openai-codex streaming", () => { expect(second.stopReason).toBe("error"); expect(second.errorMessage).toBe("Codex error: Previous response with id 'resp_9' not found"); + expect(second.responseId).toBeUndefined(); // Exactly one chain-reset retry, then the error surfaces unchanged. expect(sentBodies).toHaveLength(3); const retryBody = sentBodies[2] as { previous_response_id?: string }; From 36805739ef95f5b78346e2effc293117225692cd Mon Sep 17 00:00:00 2001 From: rasdani <73563550+rasdani@users.noreply.github.com> Date: Fri, 18 Sep 2026 20:48:36 -0700 Subject: [PATCH 15/26] fix(ai): stop sending enable_thinking to Prime Inference GLM models (#2459) (cherry picked from commit 63d88319bf5870cf609fef01f9042bbf431d7df2) --- .../prime-inference-glm-thinking-compat.md | 1 + packages/ai/scripts/generate-models.ts | 21 +++- packages/ai/src/models.generated.ts | 20 ++-- .../openai-completions-tool-choice.test.ts | 19 ++++ .../coding-agent/src/core/model-resolver.ts | 10 ++ ...reate-session-private-model-params.test.ts | 95 +++++++++++++++++++ .../coding-agent/test/model-resolver.test.ts | 1 - 7 files changed, 153 insertions(+), 14 deletions(-) create mode 100644 packages/ai/.changes/prime-inference-glm-thinking-compat.md create mode 100644 packages/coding-agent/test/create-session-private-model-params.test.ts diff --git a/packages/ai/.changes/prime-inference-glm-thinking-compat.md b/packages/ai/.changes/prime-inference-glm-thinking-compat.md new file mode 100644 index 00000000000..30f48e1e974 --- /dev/null +++ b/packages/ai/.changes/prime-inference-glm-thinking-compat.md @@ -0,0 +1 @@ +- Fixed Prime Inference GLM models failing every request with an unsupported `enable_thinking` parameter. diff --git a/packages/ai/scripts/generate-models.ts b/packages/ai/scripts/generate-models.ts index 07862276660..ba57bbaedff 100644 --- a/packages/ai/scripts/generate-models.ts +++ b/packages/ai/scripts/generate-models.ts @@ -117,6 +117,13 @@ const ZAI_THINKING_COMPAT: OpenAICompletionsCompat = { thinkingFormat: "zai", }; +// Prime Inference rejects `enable_thinking` on GLM routes with a 400 on every +// request. The routes think by default, so send no thinking parameter (and no +// reasoning_effort, whose support is unverified). +const PRIME_INFERENCE_ZAI_COMPAT: OpenAICompletionsCompat = { + supportsReasoningEffort: false, +}; + const PRIME_INFERENCE_BASE_URL = "https://api.pinference.ai/api/v1"; const PRIME_INFERENCE_COMPAT: OpenAICompletionsCompat = { supportsStore: false, @@ -452,6 +459,10 @@ function isPrimeInferenceReasoningModel(modelId: string, catalogReasoning?: bool ); } +function isPrimeInferenceZaiModelId(modelId: string): boolean { + return modelId.toLowerCase().startsWith("z-ai/glm-"); +} + function getPrimeInferenceCompat(modelId: string): OpenAICompletionsCompat { const id = modelId.toLowerCase(); if (id.includes("deepseek-v4")) { @@ -460,10 +471,10 @@ function getPrimeInferenceCompat(modelId: string): OpenAICompletionsCompat { ...DEEPSEEK_V4_COMPAT, }; } - if (id.startsWith("z-ai/glm-")) { + if (isPrimeInferenceZaiModelId(id)) { return { ...PRIME_INFERENCE_COMPAT, - ...ZAI_THINKING_COMPAT, + ...PRIME_INFERENCE_ZAI_COMPAT, }; } @@ -610,7 +621,11 @@ function createPrimeInferenceModel( ...(openRouter?.supportsReasoningEffort === false ? { supportsReasoningEffort: false, - ...(!compat.thinkingFormat ? { thinkingFormat: "openrouter" as const } : {}), + // GLM routes must not receive any thinking parameter; the + // `reasoning` object is not verified against the gateway. + ...(!compat.thinkingFormat && !isPrimeInferenceZaiModelId(entry.id) + ? { thinkingFormat: "openrouter" as const } + : {}), } : {}), }, diff --git a/packages/ai/src/models.generated.ts b/packages/ai/src/models.generated.ts index ff0f332401a..b5ff751ed7b 100644 --- a/packages/ai/src/models.generated.ts +++ b/packages/ai/src/models.generated.ts @@ -17924,7 +17924,7 @@ export const MODELS = { api: "openai-completions", provider: "prime-inference", baseUrl: "https://api.pinference.ai/api/v1", - compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"thinkingFormat":"zai"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false}, reasoning: true, thinkingLevelMap: {"minimal":null,"low":null,"medium":null,"high":"high","xhigh":null,"max":null}, input: ["text"], @@ -17943,7 +17943,7 @@ export const MODELS = { api: "openai-completions", provider: "prime-inference", baseUrl: "https://api.pinference.ai/api/v1", - compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"thinkingFormat":"zai"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false}, reasoning: true, thinkingLevelMap: {"minimal":null,"low":null,"medium":null,"high":"high","xhigh":null,"max":null}, input: ["text"], @@ -17962,7 +17962,7 @@ export const MODELS = { api: "openai-completions", provider: "prime-inference", baseUrl: "https://api.pinference.ai/api/v1", - compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"thinkingFormat":"zai"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false}, reasoning: true, thinkingLevelMap: {"minimal":null,"low":null,"medium":null,"high":"high","xhigh":null,"max":null}, input: ["text"], @@ -17981,7 +17981,7 @@ export const MODELS = { api: "openai-completions", provider: "prime-inference", baseUrl: "https://api.pinference.ai/api/v1", - compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"thinkingFormat":"zai"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false}, reasoning: true, thinkingLevelMap: {"minimal":null,"low":null,"medium":null,"high":"high","xhigh":null,"max":null}, input: ["text"], @@ -18000,7 +18000,7 @@ export const MODELS = { api: "openai-completions", provider: "prime-inference", baseUrl: "https://api.pinference.ai/api/v1", - compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"thinkingFormat":"zai"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false}, reasoning: true, thinkingLevelMap: {"minimal":null,"low":null,"medium":null,"high":"high","xhigh":null,"max":null}, input: ["text"], @@ -18019,7 +18019,7 @@ export const MODELS = { api: "openai-completions", provider: "prime-inference", baseUrl: "https://api.pinference.ai/api/v1", - compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"thinkingFormat":"zai"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false}, reasoning: true, thinkingLevelMap: {"minimal":null,"low":null,"medium":null,"high":"high","xhigh":null,"max":null}, input: ["text"], @@ -18039,7 +18039,7 @@ export const MODELS = { api: "openai-completions", provider: "prime-inference", baseUrl: "https://api.pinference.ai/api/v1", - compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"thinkingFormat":"zai"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false}, reasoning: true, thinkingLevelMap: {"minimal":null,"low":null,"medium":null,"high":"high","xhigh":null,"max":null}, input: ["text"], @@ -18059,7 +18059,7 @@ export const MODELS = { api: "openai-completions", provider: "prime-inference", baseUrl: "https://api.pinference.ai/api/v1", - compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"thinkingFormat":"zai"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false}, reasoning: true, thinkingLevelMap: {"minimal":null,"low":null,"medium":null,"high":"high","xhigh":"xhigh","max":null}, input: ["text"], @@ -18079,7 +18079,7 @@ export const MODELS = { api: "openai-completions", provider: "prime-inference", baseUrl: "https://api.pinference.ai/api/v1", - compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"thinkingFormat":"zai"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false}, reasoning: true, thinkingLevelMap: {"off":null,"minimal":null,"low":"low","medium":null,"high":"high","xhigh":null,"max":"max"}, input: ["text"], @@ -18098,7 +18098,7 @@ export const MODELS = { api: "openai-completions", provider: "prime-inference", baseUrl: "https://api.pinference.ai/api/v1", - compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"thinkingFormat":"zai"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false}, reasoning: true, thinkingLevelMap: {"off":null,"minimal":null,"low":"low","medium":null,"high":"high","xhigh":null,"max":"max"}, input: ["text", "image"], diff --git a/packages/ai/test/openai-completions-tool-choice.test.ts b/packages/ai/test/openai-completions-tool-choice.test.ts index 64dd360b759..fea8594e027 100644 --- a/packages/ai/test/openai-completions-tool-choice.test.ts +++ b/packages/ai/test/openai-completions-tool-choice.test.ts @@ -1127,6 +1127,25 @@ describe("openai-completions tool_choice", () => { expect((payload as { reasoning_effort?: unknown }).reasoning_effort).toBe("none"); }); + it("sends no thinking parameter to Prime Inference GLM routes but keeps the z.ai toggle", async () => { + const context = { messages: [{ role: "user" as const, content: "Hi", timestamp: Date.now() }] }; + const payloads = new Map>(); + for (const model of [getModel("prime-inference", "z-ai/glm-5.3")!, getModel("zai", "glm-5.3")!]) { + await streamSimple(model, context, { + apiKey: "test", + reasoning: "high", + onPayload: (params: unknown) => { + payloads.set(model.id, params as Record); + }, + }).result(); + } + const prime = payloads.get("z-ai/glm-5.3"); + expect(prime?.enable_thinking).toBeUndefined(); + expect(prime?.reasoning_effort).toBeUndefined(); + expect(prime?.reasoning).toBeUndefined(); + expect(payloads.get("glm-5.3")?.enable_thinking).toBe(true); + }); + it("serializes explicit off only for models that allow disabling reasoning", async () => { const baseModel = getModel("openrouter", "deepseek/deepseek-r1")!; const effortCompat = { ...baseModel.compat, supportsReasoningEffort: true }; diff --git a/packages/coding-agent/src/core/model-resolver.ts b/packages/coding-agent/src/core/model-resolver.ts index 21f0763ddd7..904abbb8551 100644 --- a/packages/coding-agent/src/core/model-resolver.ts +++ b/packages/coding-agent/src/core/model-resolver.ts @@ -158,6 +158,16 @@ function buildFallbackModel(provider: string, modelId: string, availableModels: const providerModels = availableModels.filter((m) => m.provider === provider); if (providerModels.length === 0) return undefined; + // Daemon-created sessions (rlm.create_session) re-resolve their model from + // provider/id strings in a registry that has not refreshed the team-authorized + // private catalog, so unknown private ids reach this fallback. They must + // inherit a private-route template: public-route limits and thinking-level + // maps do not describe private routes. + if (isPrivatePrimeInferenceModel({ provider, id: modelId })) { + const privateTemplate = providerModels.find((m) => isPrivatePrimeInferenceModel(m)); + if (!privateTemplate) return undefined; + return { ...privateTemplate, id: modelId, name: modelId }; + } const defaultId = defaultModelPerProvider[provider as KnownProvider]; const baseModel = defaultId ? (providerModels.find((m) => m.id === defaultId) ?? providerModels[0]) diff --git a/packages/coding-agent/test/create-session-private-model-params.test.ts b/packages/coding-agent/test/create-session-private-model-params.test.ts new file mode 100644 index 00000000000..05c222a1cbb --- /dev/null +++ b/packages/coding-agent/test/create-session-private-model-params.test.ts @@ -0,0 +1,95 @@ +import { getModels, type Model, streamSimple } from "@earendil-works/pi-ai"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { resolveCliModel } from "../src/core/model-resolver.js"; +import { getPrivatePrimeInferenceModels } from "../src/core/prime-inference-models.js"; + +const mockState = vi.hoisted(() => ({ + lastParams: undefined as unknown, +})); + +vi.mock("openai", () => { + class FakeOpenAI { + chat = { + completions: { + create: (params: unknown) => { + mockState.lastParams = params; + const stream = { + async *[Symbol.asyncIterator]() { + yield { + choices: [{ delta: {}, finish_reason: "stop" }], + usage: { + prompt_tokens: 1, + completion_tokens: 1, + prompt_tokens_details: { cached_tokens: 0, cache_write_tokens: 0 }, + completion_tokens_details: { reasoning_tokens: 0 }, + }, + }; + }, + }; + const promise = Promise.resolve(stream) as Promise & { + withResponse: () => Promise<{ + data: typeof stream; + response: { status: number; headers: Headers }; + }>; + }; + promise.withResponse = async () => ({ + data: stream, + response: { status: 200, headers: new Headers() }, + }); + return promise; + }, + }, + }; + } + + return { default: FakeOpenAI }; +}); + +describe("created-session private model request params", () => { + beforeEach(() => { + mockState.lastParams = undefined; + }); + + // The registry a daemon worker builds for a fresh session: bundled public + // catalog plus bundled private models, without the team-authorized private + // catalog that only loads after refreshAvailableModels(). This is what an + // rlm.create_session session resolves its provider/model strings against. + const freshWorkerRegistry = { + getAll: () => [...getModels("prime-inference"), ...getPrivatePrimeInferenceModels()], + } as unknown as Parameters[0]["modelRegistry"]; + + it("omits enable_thinking for an internal model resolved like a created session", async () => { + const resolved = resolveCliModel({ + cliProvider: "prime-inference", + cliModel: "internal/glm-5.3-fast", + modelRegistry: freshWorkerRegistry, + }); + expect(resolved.error).toBeUndefined(); + const model = resolved.model as Model<"openai-completions">; + + await streamSimple( + model, + { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] }, + { apiKey: "test", reasoning: "high" }, + ).result(); + + const params = mockState.lastParams as { model?: string; enable_thinking?: boolean }; + expect(params.model).toBe("internal/glm-5.3-fast"); + expect("enable_thinking" in params).toBe(false); + }); + + it("omits enable_thinking for public Prime Inference z-ai models", async () => { + const model = getModels("prime-inference").find( + (candidate) => candidate.id === "z-ai/glm-5.3", + ) as Model<"openai-completions">; + + await streamSimple( + model, + { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] }, + { apiKey: "test", reasoning: "high" }, + ).result(); + + const params = mockState.lastParams as { enable_thinking?: boolean }; + expect("enable_thinking" in params).toBe(false); + }); +}); diff --git a/packages/coding-agent/test/model-resolver.test.ts b/packages/coding-agent/test/model-resolver.test.ts index 716ba2c13b6..7deb98fc52b 100644 --- a/packages/coding-agent/test/model-resolver.test.ts +++ b/packages/coding-agent/test/model-resolver.test.ts @@ -274,7 +274,6 @@ describe("resolveCliModel", () => { cliModel: "zai/glm-5", modelRegistry: registry, }); - expect(result.error).toBeUndefined(); expect(result.model?.provider).toBe("zai"); expect(result.model?.id).toBe("glm-5"); From c6d7f399ee6a7857faeee2f91e3af8cf3fb09803 Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:57:11 -0400 Subject: [PATCH 16/26] [RSI, performance] Cache the branch array on the per-turn hot path (#2414) * perf(session-manager): cache the live leaf branch on the per-turn read path The per-turn compaction check and the context-usage state build call getBranch() on every assistant message end, and each call walked leaf-to-root and allocated a fresh array (thousands of entries on a long session). SessionManager now caches the live leaf's path. An append is always a child of the current leaf, so appends extend the cached array in place. Every other leafId write (branch, resetLeaf, branchWithSummary, append rollback, session load/rebuild) drops the cache, and getBranch(fromId) reads for a non-leaf entry are never cached. getBranch() returns the cached array itself, so the one caller that hands the branch to an awaited extension handler (session_before_compact) now passes a slice. That keeps the previous snapshot contract: an entry appended while the handler runs no longer grows the array the extension holds. docs/extensions.md states that contract next to session_before_compact. Compaction decisions are unchanged. On a synthetic 20k-entry branch the cached read costs 0.02us versus 297us for the leaf-to-root walk it replaces. Tests: test/session-manager/leaf-branch-cache.test.ts covers cache identity across reads, in-place extension on append, invalidation on branch/resetLeaf/branchWithSummary and on append rollback, non-leaf suffix reads, and agreement with a reopened session. test/suite/agent-session-compaction.test.ts pins the session_before_compact snapshot: it appends an entry from inside the awaited handler and asserts the handler's branchEntries stays unchanged while the live branch grows (that test fails without the slice). * chore(changelog): add fragment for the leaf branch cache * fix(session-manager): keep a rolled-back append out of the live leaf branch cache The cache was extended before _persist, so a failed persist or a failed post-append flush left the rolled-back entry in an array getBranch() had already handed out. - extend the cache only after _persist returns - pop the rolled-back entry from the live cached array before invalidating - trim the cache comments to current behavior and document the read-only live-array contract on getBranch() (a ReadonlySessionManager surface) - pin the rollback property and the in-place append identity in the cache test * test(session-manager): rely on the configured test timeout in the branch snapshot test * test(session-manager): compress leaf-branch-cache tests and prune tree-traversal duplicates Consolidate the new cache-contract tests around shared fixtures (merged leaf-move test, trimmed rollback phases, deleted the reopen-equivalence duplicate) and compact the session_before_compact snapshot test. Prune eleven tree-traversal vectors whose observables are already pinned by the new cache tests and kept siblings (probed: each twin fails under the matching mutation), fold the five append-operation tests into one table preserving every vector, and fold the in-memory fork result assertion into the kept branched-tree test. Net test additions fall from 283 to -7 against a 26-line budget. * test(session-manager): pin the third-sibling tree shape the compression dropped The pruned getTree multi-branch duplicate was the only test with 3+ siblings at one branch point; the surviving twin pinned only 2. Re-buy the shape by branching the same point twice in the kept test (scoped review hardening). * changelog: note getBranch() returns a live shared array --------- Co-authored-by: prime-agent (cherry picked from commit 0498deb439b450603d2c09cb3b87d749e50e4f9d) --- .../.changes/perf-branch-cache.md | 2 + packages/coding-agent/docs/extensions.md | 2 +- .../coding-agent/src/core/agent-session.ts | 4 +- .../coding-agent/src/core/session-manager.ts | 50 ++- .../session-manager/leaf-branch-cache.test.ts | 158 ++++++++ .../session-manager/tree-traversal.test.ts | 364 ++++-------------- .../suite/agent-session-compaction.test.ts | 51 ++- 7 files changed, 339 insertions(+), 292 deletions(-) create mode 100644 packages/coding-agent/.changes/perf-branch-cache.md create mode 100644 packages/coding-agent/test/session-manager/leaf-branch-cache.test.ts diff --git a/packages/coding-agent/.changes/perf-branch-cache.md b/packages/coding-agent/.changes/perf-branch-cache.md new file mode 100644 index 00000000000..31bc5dd9656 --- /dev/null +++ b/packages/coding-agent/.changes/perf-branch-cache.md @@ -0,0 +1,2 @@ +- Cached the active branch path in the session manager, so per-turn compaction checks and context-usage updates no longer rebuild the whole leaf-to-root path after every assistant message. +- `SessionManager.getBranch()` now returns the live shared branch array (public API): repeated reads return the same array object and straight-line appends extend it in place. Callers must treat it as read-only and take `.slice()` for a snapshot; `session_before_compact` already passes a snapshot. diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index 7254cc35b2a..8b38846f791 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -411,7 +411,7 @@ Do cleanup work in `session_shutdown`, then reestablish any in-memory state in ` #### session_before_compact / session_compact -Fired on compaction. See [compaction.md](compaction.md) for details. +Fired on compaction. See [compaction.md](compaction.md) for details. `branchEntries` is a read-only snapshot of the active branch taken before compaction, so it does not include entries appended while your handler runs and must not be mutated. ```typescript pi.on("session_before_compact", async (event, ctx) => { diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 0207237d08d..04b71a01068 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -8520,7 +8520,9 @@ export class AgentSession { const result = (await this._extensionRunner.emit({ type: "session_before_compact", preparation, - branchEntries: pathEntries, + // slice: getBranch() returns the live leaf-branch cache, which appends + // extend in place, so the awaited handler gets a snapshot. + branchEntries: pathEntries.slice(), customInstructions, signal, })) as SessionBeforeCompactResult | undefined; diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index 2403578d06f..b2e25e29fb4 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -1450,6 +1450,12 @@ export class SessionManager { private labelsById: Map = new Map(); private labelTimestampsById: Map = new Map(); private leafId: string | null = null; + // Cache of the live leaf's branch path, so the per-turn compaction check and the + // context-usage state build get an O(1) lookup instead of a leaf-to-root walk. + // Appends extend the cached array in place; every other leafId write drops the + // cache. The array is handed out as-is, so callers must treat it as read-only + // and must not hold it across an append. + private leafBranchCache: { leafId: string | null; entries: SessionEntry[] } | null = null; private persistListeners = new Set(); private constructor( @@ -1565,6 +1571,7 @@ export class SessionManager { this.labelsById.clear(); this.labelTimestampsById.clear(); this.leafId = null; + this.leafBranchCache = null; this.flushed = false; if (this.persist) { @@ -1578,6 +1585,7 @@ export class SessionManager { this.labelsById.clear(); this.labelTimestampsById.clear(); this.leafId = null; + this.leafBranchCache = null; for (const entry of this.fileEntries) { if (entry.type === "session") continue; this.byId.set(entry.id, entry); @@ -1736,6 +1744,21 @@ export class SessionManager { this.byId.set(entry.id, entry); this.leafId = entry.id; this._persist(entry); + // After _persist: a throwing persist rolls the append back, and the cache must + // not grow past a leaf that never survived it. + this._extendLeafBranchCache(entry); + } + + // Appends only ever extend the live branch path, so the leaf branch cache grows + // in place. Anything other than a straight append off the cached leaf drops it. + private _extendLeafBranchCache(entry: SessionEntry): void { + const cache = this.leafBranchCache; + if (cache && cache.leafId === entry.parentId) { + cache.leafId = entry.id; + cache.entries.push(entry); + } else { + this.leafBranchCache = null; + } } appendMessage(message: Message | CustomMessage | BashExecutionMessage): string { @@ -2042,9 +2065,17 @@ export class SessionManager { } catch (error) { // The append indexes the entry before persisting it; undo exactly that. if (this.leafId !== null && this.leafId !== previousLeafId) { - this.byId.delete(this.leafId); + const rolledBackId = this.leafId; + this.byId.delete(rolledBackId); this.fileEntries.pop(); + // A cache extended by the append is live, so drop the rolled-back entry from + // the array a caller may already hold before invalidating the cache. + const cache = this.leafBranchCache; + if (cache && cache.entries[cache.entries.length - 1]?.id === rolledBackId) { + cache.entries.pop(); + } this.leafId = previousLeafId; + this.leafBranchCache = null; // The failed append may have left a torn line on disk. Restore the file // from the rolled-back entries now; if that also fails (e.g. the disk is // still full), fall back to forcing the next persist to rewrite. @@ -2108,7 +2139,18 @@ export class SessionManager { return entry.id; } + /** + * Entries on the active branch, root to leaf. A leaf read returns the live + * leaf-branch cache: treat the result as read-only and do not hold it across + * appends. + */ getBranch(fromId?: string): SessionEntry[] { + // Leaf-path reads are the per-turn hot path (compaction checks, context + // usage); a read for another entry walks the chain instead. + const isLeafPath = fromId === undefined || fromId === this.leafId; + if (isLeafPath && this.leafBranchCache?.leafId === this.leafId) { + return this.leafBranchCache.entries; + } // push+reverse, not unshift-per-entry: unshift is O(n), which makes this O(n^2) on long sessions. const path: SessionEntry[] = []; const startId = fromId ?? this.leafId; @@ -2118,6 +2160,9 @@ export class SessionManager { current = current.parentId ? this.byId.get(current.parentId) : undefined; } path.reverse(); + if (isLeafPath) { + this.leafBranchCache = { leafId: this.leafId, entries: path }; + } return path; } @@ -2188,10 +2233,12 @@ export class SessionManager { throw new Error(`Entry ${branchFromId} not found`); } this.leafId = branchFromId; + this.leafBranchCache = null; } resetLeaf(): void { this.leafId = null; + this.leafBranchCache = null; } branchWithSummary( @@ -2205,6 +2252,7 @@ export class SessionManager { throw new Error(`Entry ${branchFromId} not found`); } this.leafId = branchFromId; + this.leafBranchCache = null; const entry: BranchSummaryEntry = { type: "branch_summary", id: generateId(this.byId), diff --git a/packages/coding-agent/test/session-manager/leaf-branch-cache.test.ts b/packages/coding-agent/test/session-manager/leaf-branch-cache.test.ts new file mode 100644 index 00000000000..4163038a8ca --- /dev/null +++ b/packages/coding-agent/test/session-manager/leaf-branch-cache.test.ts @@ -0,0 +1,158 @@ +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { SessionManager } from "../../src/core/session-manager.js"; +import { assistantMsg, userMsg } from "../utilities.js"; + +// SessionManager caches the live leaf's branch path (getBranch() reads) because +// agent-session re-reads it on every assistant message end. These tests pin the +// cache's contract: identical reads share the array, appends extend it, every +// other leaf move drops it, and non-leaf reads never poison it. +describe("SessionManager leaf branch cache", () => { + const tempDirs: string[] = []; + + afterEach(() => { + while (tempDirs.length > 0) { + rmSync(tempDirs.pop()!, { recursive: true, force: true }); + } + }); + + function persistedSession(): SessionManager { + const dir = mkdtempSync(join(tmpdir(), "pi-leaf-branch-cache-")); + tempDirs.push(dir); + return SessionManager.create(dir, join(dir, "sessions")); + } + + const ids = (session: SessionManager) => session.getBranch().map((entry) => entry.id); + + // Walks the served path and checks it against the byId parent chain, so a + // stale or half-extended cache cannot pass by returning the right ids only. + function expectChainMatches(session: SessionManager, expected: string[]): void { + const branch = session.getBranch(); + expect(branch.map((entry) => entry.id)).toEqual(expected); + for (let index = 1; index < branch.length; index++) { + expect(branch[index]!.parentId).toBe(branch[index - 1]!.id); + expect(session.getEntry(branch[index]!.id)).toBe(branch[index]); + } + } + + function seedTwo() { + const session = SessionManager.inMemory(); + const firstId = session.appendMessage(userMsg("one")); + const secondId = session.appendMessage(assistantMsg("two")); + return { session, firstId, secondId }; + } + + it("serves one shared cached array that appends extend in place", () => { + const { session, firstId, secondId } = seedTwo(); + const first = session.getBranch(); + expect(session.getBranch()).toBe(first); // repeated leaf reads share the array + expect(first.map((entry) => entry.id)).toEqual([firstId, secondId]); + + const thirdId = session.appendCustomMessageEntry("note", "hello", false); + expectChainMatches(session, [firstId, secondId, thirdId]); + // An append extends the same array the earlier read handed out. + expect(session.getBranch()).toBe(first); + + const fourthId = session.appendMessage(userMsg("four")); + expectChainMatches(session, [firstId, secondId, thirdId, fourthId]); + }); + + it("drops the cache on every leaf-changing operation and re-caches the new path", () => { + const { session, firstId, secondId } = seedTwo(); + expect(ids(session)).toEqual([firstId, secondId]); // populates the cache + + // branch(): the abandoned path must not leak, and a sibling append re-roots. + session.branch(firstId); + expectChainMatches(session, [firstId]); + const siblingId = session.appendMessage(userMsg("sibling")); + expectChainMatches(session, [firstId, siblingId]); + // Branching back to the abandoned entry must not resurrect the old cache. + session.branch(secondId); + expectChainMatches(session, [firstId, secondId]); + + // resetLeaf(): an empty branch, then the next append re-roots the tree. + session.resetLeaf(); + expect(session.getLeafId()).toBeNull(); + expect(session.getBranch()).toEqual([]); + const rootId = session.appendMessage(userMsg("root again")); + expect(session.getEntry(rootId)?.parentId).toBeNull(); + expectChainMatches(session, [rootId]); + + // branchWithSummary(): the summary entry becomes the new leaf, and the + // compaction appended after it lands on the summarized branch. + const summaryId = session.branchWithSummary(secondId, "summary of the abandoned path"); + expectChainMatches(session, [firstId, secondId, summaryId]); + const compactionId = session.appendCompaction("compaction summary", rootId, 1234, { + readFiles: [], + modifiedFiles: [], + }); + expectChainMatches(session, [firstId, secondId, summaryId, compactionId]); + }); + + it("returns the correct suffix for a mid-branch read without poisoning the leaf cache", () => { + const session = SessionManager.inMemory(); + const firstId = session.appendMessage(userMsg("one")); + const secondId = session.appendMessage(assistantMsg("two")); + const thirdId = session.appendMessage(userMsg("three")); + const fourthId = session.appendMessage(assistantMsg("four")); + + const leafPath = session.getBranch(); + expect(ids(session)).toEqual([firstId, secondId, thirdId, fourthId]); + + const suffix = session.getBranch(secondId); + expect(suffix.map((entry) => entry.id)).toEqual([firstId, secondId]); + expect(suffix).not.toBe(leafPath); + + // The mid-branch read must not have replaced the cached leaf path... + expect(session.getBranch()).toBe(leafPath); + // ...nor cached a suffix as if it were the leaf path. + const fifthId = session.appendMessage(userMsg("five")); + expectChainMatches(session, [firstId, secondId, thirdId, fourthId, fifthId]); + }); + + it("drops the cached branch when a failed append is rolled back", () => { + const session = persistedSession(); + const firstId = session.appendMessage(userMsg("one")); + expect(ids(session)).toEqual([firstId]); // populates the cache + + vi.spyOn(session, "_persist").mockImplementationOnce(() => { + throw new Error("disk full"); + }); + expect(() => session.appendCustomMessageEntryWithRollback("note", "unsaved", false)).toThrow("disk full"); + + expect(session.getLeafId()).toBe(firstId); + expectChainMatches(session, [firstId]); + // The next append starts from the rolled-back leaf and stays coherent. + const secondId = session.appendCustomMessageEntry("note", "saved", false); + expectChainMatches(session, [firstId, secondId]); + }); + + it("never leaves a rolled-back append in a branch array a caller already holds", () => { + const session = persistedSession(); + const firstId = session.appendMessage(userMsg("one")); + const held = session.getBranch(); // the live cached array + + // The append persists, so it reaches the cached array in place... + vi.spyOn(session, "flushNow").mockImplementationOnce(() => { + throw new Error("flush failed"); + }); + expect(() => session.appendCustomMessageEntryWithRollback("note", "unsaved", false)).toThrow("flush failed"); + + // ...and the rollback takes it back out of that same array. + expect(held.map((entry) => entry.id)).toEqual([firstId]); + expect(session.getLeafId()).toBe(firstId); + expectChainMatches(session, [firstId]); + + // A failing persist never extends the cache in the first place. + const heldAgain = session.getBranch(); + vi.spyOn(session, "_persist").mockImplementationOnce(() => { + throw new Error("disk full"); + }); + expect(() => session.appendCustomMessageEntryWithRollback("note", "unsaved again", false)).toThrow("disk full"); + + expect(heldAgain.map((entry) => entry.id)).toEqual([firstId]); + expectChainMatches(session, [firstId]); + }); +}); diff --git a/packages/coding-agent/test/session-manager/tree-traversal.test.ts b/packages/coding-agent/test/session-manager/tree-traversal.test.ts index 9b7206d4524..b3d016db192 100644 --- a/packages/coding-agent/test/session-manager/tree-traversal.test.ts +++ b/packages/coding-agent/test/session-manager/tree-traversal.test.ts @@ -2,156 +2,87 @@ import { existsSync, mkdirSync, readFileSync, rmSync } from "fs"; import { tmpdir } from "os"; import { join } from "path"; import { describe, expect, it } from "vitest"; -import { type CustomEntry, SessionManager } from "../../src/core/session-manager.js"; +import { type SessionEntry, SessionManager } from "../../src/core/session-manager.js"; import { assistantMsg, userMsg } from "../utilities.js"; describe("SessionManager append and tree traversal", () => { describe("append operations", () => { - it("appendMessage creates entry with correct parentId chain", () => { + const cases: Array<{ + name: string; + seed: (session: SessionManager) => string[]; + append: (session: SessionManager) => string; + type: string; + expectEntry?: (entry: SessionEntry, seedIds: string[]) => void; + }> = [ + { + name: "appendMessage creates entry with correct parentId chain", + seed: (session) => [session.appendMessage(userMsg("first"))], + append: (session) => session.appendMessage(assistantMsg("second")), + type: "message", + }, + { + name: "appendThinkingLevelChange integrates into tree", + seed: (session) => [session.appendMessage(userMsg("hello"))], + append: (session) => session.appendThinkingLevelChange("high"), + type: "thinking_level_change", + expectEntry: (entry) => expect(entry).toMatchObject({ thinkingLevel: "high" }), + }, + { + name: "appendModelChange integrates into tree", + seed: (session) => [session.appendMessage(userMsg("hello"))], + append: (session) => session.appendModelChange("openai", "gpt-4"), + type: "model_change", + expectEntry: (entry) => expect(entry).toMatchObject({ provider: "openai", modelId: "gpt-4" }), + }, + { + name: "appendCompaction integrates into tree and persists customInstructions", + seed: (session) => { + const firstId = session.appendMessage(userMsg("1")); + return [firstId, session.appendMessage(assistantMsg("2"))]; + }, + append: (session) => + session.appendCompaction( + "summary", + session.getEntries()[0]!.id, + 1000, + undefined, + undefined, + "focus on xyz", + ), + type: "compaction", + expectEntry: (entry, seedIds) => + expect(entry).toMatchObject({ + summary: "summary", + firstKeptEntryId: seedIds[0], + tokensBefore: 1000, + customInstructions: "focus on xyz", + }), + }, + { + name: "appendCustomEntry integrates into tree", + seed: (session) => [session.appendMessage(userMsg("hello"))], + append: (session) => session.appendCustomEntry("my_data", { key: "value" }), + type: "custom", + expectEntry: (entry) => expect(entry).toMatchObject({ customType: "my_data", data: { key: "value" } }), + }, + ]; + it.each(cases)("$name", ({ seed, append, type, expectEntry }) => { const session = SessionManager.inMemory(); - - const id1 = session.appendMessage(userMsg("first")); - const id2 = session.appendMessage(assistantMsg("second")); - const id3 = session.appendMessage(userMsg("third")); - - const entries = session.getEntries(); - expect(entries).toHaveLength(3); - - expect(entries[0].id).toBe(id1); - expect(entries[0].parentId).toBeNull(); - expect(entries[0].type).toBe("message"); - - expect(entries[1].id).toBe(id2); - expect(entries[1].parentId).toBe(id1); - - expect(entries[2].id).toBe(id3); - expect(entries[2].parentId).toBe(id2); - }); - - it("appendThinkingLevelChange integrates into tree", () => { - const session = SessionManager.inMemory(); - - const msgId = session.appendMessage(userMsg("hello")); - const thinkingId = session.appendThinkingLevelChange("high"); - const _msg2Id = session.appendMessage(assistantMsg("response")); - - const entries = session.getEntries(); - expect(entries).toHaveLength(3); - - const thinkingEntry = entries.find((e) => e.type === "thinking_level_change"); - expect(thinkingEntry).toBeDefined(); - expect(thinkingEntry!.id).toBe(thinkingId); - expect(thinkingEntry!.parentId).toBe(msgId); - - expect(entries[2].parentId).toBe(thinkingId); - }); - - it("appendModelChange integrates into tree", () => { - const session = SessionManager.inMemory(); - - const msgId = session.appendMessage(userMsg("hello")); - const modelId = session.appendModelChange("openai", "gpt-4"); - const _msg2Id = session.appendMessage(assistantMsg("response")); - - const entries = session.getEntries(); - const modelEntry = entries.find((e) => e.type === "model_change"); - expect(modelEntry).toBeDefined(); - expect(modelEntry?.id).toBe(modelId); - expect(modelEntry?.parentId).toBe(msgId); - if (modelEntry?.type === "model_change") { - expect(modelEntry.provider).toBe("openai"); - expect(modelEntry.modelId).toBe("gpt-4"); - } - - expect(entries[2].parentId).toBe(modelId); - }); - - it("appendCompaction integrates into tree", () => { - const session = SessionManager.inMemory(); - - const id1 = session.appendMessage(userMsg("1")); - const id2 = session.appendMessage(assistantMsg("2")); - const compactionId = session.appendCompaction("summary", id1, 1000); - const _id3 = session.appendMessage(userMsg("3")); + const seedIds = seed(session); + const appendedId = append(session); + const followingId = session.appendMessage(userMsg("after")); const entries = session.getEntries(); - const compactionEntry = entries.find((e) => e.type === "compaction"); - expect(compactionEntry).toBeDefined(); - expect(compactionEntry?.id).toBe(compactionId); - expect(compactionEntry?.parentId).toBe(id2); - if (compactionEntry?.type === "compaction") { - expect(compactionEntry.summary).toBe("summary"); - expect(compactionEntry.firstKeptEntryId).toBe(id1); - expect(compactionEntry.tokensBefore).toBe(1000); - } - - expect(entries[3].parentId).toBe(compactionId); - }); - - it("appendCompaction persists customInstructions", () => { - const session = SessionManager.inMemory(); - - const id1 = session.appendMessage(userMsg("1")); - session.appendMessage(assistantMsg("2")); - session.appendCompaction("summary", id1, 1000, undefined, undefined, "focus on xyz"); - - const compactionEntry = session.getEntries().find((e) => e.type === "compaction"); - if (compactionEntry?.type === "compaction") { - expect(compactionEntry.customInstructions).toBe("focus on xyz"); - } else { - expect.fail("compaction entry not found"); - } - }); - - it("appendCustomEntry integrates into tree", () => { - const session = SessionManager.inMemory(); - - const msgId = session.appendMessage(userMsg("hello")); - const customId = session.appendCustomEntry("my_data", { key: "value" }); - const _msg2Id = session.appendMessage(assistantMsg("response")); - - const entries = session.getEntries(); - const customEntry = entries.find((e) => e.type === "custom") as CustomEntry; - expect(customEntry).toBeDefined(); - expect(customEntry.id).toBe(customId); - expect(customEntry.parentId).toBe(msgId); - expect(customEntry.customType).toBe("my_data"); - expect(customEntry.data).toEqual({ key: "value" }); - - expect(entries[2].parentId).toBe(customId); - }); - - it("leaf pointer advances after each append", () => { - const session = SessionManager.inMemory(); - - expect(session.getLeafId()).toBeNull(); - - const id1 = session.appendMessage(userMsg("1")); - expect(session.getLeafId()).toBe(id1); - - const id2 = session.appendMessage(assistantMsg("2")); - expect(session.getLeafId()).toBe(id2); - - const id3 = session.appendThinkingLevelChange("high"); - expect(session.getLeafId()).toBe(id3); + const appended = entries.find((e) => e.id === appendedId)!; + expect(appended.type).toBe(type); + expect(appended.parentId).toBe(seedIds.at(-1) ?? null); + expectEntry?.(appended, seedIds); + // The entry after the appended one parents to it, pinning the chain. + expect(entries.find((e) => e.id === followingId)?.parentId).toBe(appendedId); }); }); describe("getPath", () => { - it("returns empty array for empty session", () => { - const session = SessionManager.inMemory(); - expect(session.getBranch()).toEqual([]); - }); - - it("returns single entry path", () => { - const session = SessionManager.inMemory(); - const id = session.appendMessage(userMsg("hello")); - - const path = session.getBranch(); - expect(path).toHaveLength(1); - expect(path[0].id).toBe(id); - }); - it("returns full path from root to leaf", () => { const session = SessionManager.inMemory(); @@ -164,19 +95,6 @@ describe("SessionManager append and tree traversal", () => { expect(path).toHaveLength(4); expect(path.map((e) => e.id)).toEqual([id1, id2, id3, id4]); }); - - it("returns path from specified entry to root", () => { - const session = SessionManager.inMemory(); - - const id1 = session.appendMessage(userMsg("1")); - const id2 = session.appendMessage(assistantMsg("2")); - const _id3 = session.appendMessage(userMsg("3")); - const _id4 = session.appendMessage(assistantMsg("4")); - - const path = session.getBranch(id2); - expect(path).toHaveLength(2); - expect(path.map((e) => e.id)).toEqual([id1, id2]); - }); }); describe("getTree", () => { @@ -185,25 +103,6 @@ describe("SessionManager append and tree traversal", () => { expect(session.getTree()).toEqual([]); }); - it("returns single root for linear session", () => { - const session = SessionManager.inMemory(); - - const id1 = session.appendMessage(userMsg("1")); - const id2 = session.appendMessage(assistantMsg("2")); - const id3 = session.appendMessage(userMsg("3")); - - const tree = session.getTree(); - expect(tree).toHaveLength(1); - - const root = tree[0]; - expect(root.entry.id).toBe(id1); - expect(root.children).toHaveLength(1); - expect(root.children[0].entry.id).toBe(id2); - expect(root.children[0].children).toHaveLength(1); - expect(root.children[0].children[0].entry.id).toBe(id3); - expect(root.children[0].children[0].children).toHaveLength(0); - }); - it("returns tree with branches after branch", () => { const session = SessionManager.inMemory(); @@ -213,6 +112,8 @@ describe("SessionManager append and tree traversal", () => { session.branch(id2); const id4 = session.appendMessage(userMsg("4-branch")); + session.branch(id2); + const id5 = session.appendMessage(userMsg("5-branch")); const tree = session.getTree(); expect(tree).toHaveLength(1); @@ -223,34 +124,10 @@ describe("SessionManager append and tree traversal", () => { const node2 = root.children[0]; expect(node2.entry.id).toBe(id2); - expect(node2.children).toHaveLength(2); // id3 and id4 are siblings + expect(node2.children).toHaveLength(3); // id3, id4, and id5 are siblings const childIds = node2.children.map((c) => c.entry.id).sort(); - expect(childIds).toEqual([id3, id4].sort()); - }); - - it("handles multiple branches at same point", () => { - const session = SessionManager.inMemory(); - - const _id1 = session.appendMessage(userMsg("root")); - const id2 = session.appendMessage(assistantMsg("response")); - - session.branch(id2); - const idA = session.appendMessage(userMsg("branch-A")); - - session.branch(id2); - const idB = session.appendMessage(userMsg("branch-B")); - - session.branch(id2); - const idC = session.appendMessage(userMsg("branch-C")); - - const tree = session.getTree(); - const node2 = tree[0].children[0]; - expect(node2.entry.id).toBe(id2); - expect(node2.children).toHaveLength(3); - - const branchIds = node2.children.map((c) => c.entry.id).sort(); - expect(branchIds).toEqual([idA, idB, idC].sort()); + expect(childIds).toEqual([id3, id4, id5].sort()); }); it("handles deep branching", () => { @@ -284,62 +161,15 @@ describe("SessionManager append and tree traversal", () => { }); describe("branch", () => { - it("moves leaf pointer to specified entry", () => { - const session = SessionManager.inMemory(); - - const id1 = session.appendMessage(userMsg("1")); - const _id2 = session.appendMessage(assistantMsg("2")); - const id3 = session.appendMessage(userMsg("3")); - - expect(session.getLeafId()).toBe(id3); - - session.branch(id1); - expect(session.getLeafId()).toBe(id1); - }); - it("throws for non-existent entry", () => { const session = SessionManager.inMemory(); session.appendMessage(userMsg("hello")); expect(() => session.branch("nonexistent")).toThrow("Entry nonexistent not found"); }); - - it("new appends become children of branch point", () => { - const session = SessionManager.inMemory(); - - const id1 = session.appendMessage(userMsg("1")); - const _id2 = session.appendMessage(assistantMsg("2")); - - session.branch(id1); - const id3 = session.appendMessage(userMsg("branched")); - - const entries = session.getEntries(); - const branchedEntry = entries.find((e) => e.id === id3)!; - expect(branchedEntry.parentId).toBe(id1); // sibling of id2 - }); }); describe("branchWithSummary", () => { - it("inserts branch summary and advances leaf", () => { - const session = SessionManager.inMemory(); - - const id1 = session.appendMessage(userMsg("1")); - const _id2 = session.appendMessage(assistantMsg("2")); - const _id3 = session.appendMessage(userMsg("3")); - - const summaryId = session.branchWithSummary(id1, "Summary of abandoned work"); - - expect(session.getLeafId()).toBe(summaryId); - - const entries = session.getEntries(); - const summaryEntry = entries.find((e) => e.type === "branch_summary"); - expect(summaryEntry).toBeDefined(); - expect(summaryEntry?.parentId).toBe(id1); - if (summaryEntry?.type === "branch_summary") { - expect(summaryEntry.summary).toBe("Summary of abandoned work"); - } - }); - it("throws for non-existent entry", () => { const session = SessionManager.inMemory(); session.appendMessage(userMsg("hello")); @@ -371,26 +201,6 @@ describe("SessionManager append and tree traversal", () => { const session = SessionManager.inMemory(); expect(session.getEntry("nonexistent")).toBeUndefined(); }); - - it("returns entry by id", () => { - const session = SessionManager.inMemory(); - - const id1 = session.appendMessage(userMsg("first")); - const id2 = session.appendMessage(assistantMsg("second")); - - const entry1 = session.getEntry(id1); - expect(entry1).toBeDefined(); - expect(entry1?.type).toBe("message"); - if (entry1?.type === "message" && entry1.message.role === "user") { - expect(entry1.message.content).toBe("first"); - } - - const entry2 = session.getEntry(id2); - expect(entry2).toBeDefined(); - if (entry2?.type === "message" && entry2.message.role === "assistant") { - expect((entry2.message.content as any)[0].text).toBe("second"); - } - }); }); describe("buildSessionContext with branches", () => { @@ -423,28 +233,6 @@ describe("createBranchedSession", () => { expect(() => session.createBranchedSession("nonexistent")).toThrow("Entry nonexistent not found"); }); - it("creates new session with path to specified leaf (in-memory)", () => { - const session = SessionManager.inMemory(); - - const id1 = session.appendMessage(userMsg("1")); - const id2 = session.appendMessage(assistantMsg("2")); - const id3 = session.appendMessage(userMsg("3")); - session.appendMessage(assistantMsg("4")); - - // Branch from 3: 3 -> 5 - session.branch(id3); - const _id5 = session.appendMessage(userMsg("5")); - - // Create branched session from id2 (should only have 1 -> 2) - const result = session.createBranchedSession(id2); - expect(result).toBeUndefined(); // in-memory returns null - - const entries = session.getEntries(); - expect(entries).toHaveLength(2); - expect(entries[0].id).toBe(id1); - expect(entries[1].id).toBe(id2); - }); - it("extracts correct path from branched tree", () => { const session = SessionManager.inMemory(); @@ -456,8 +244,8 @@ describe("createBranchedSession", () => { const id4 = session.appendMessage(userMsg("4")); const id5 = session.appendMessage(assistantMsg("5")); - // Create branched session from id5 (should have 1 -> 2 -> 4 -> 5) - session.createBranchedSession(id5); + // In-memory branching rebuilds this session (no new file): the entries become the path to id5. + expect(session.createBranchedSession(id5)).toBeUndefined(); const entries = session.getEntries(); expect(entries).toHaveLength(4); diff --git a/packages/coding-agent/test/suite/agent-session-compaction.test.ts b/packages/coding-agent/test/suite/agent-session-compaction.test.ts index 65afe62088a..a3b203f9beb 100644 --- a/packages/coding-agent/test/suite/agent-session-compaction.test.ts +++ b/packages/coding-agent/test/suite/agent-session-compaction.test.ts @@ -10,7 +10,7 @@ import { import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { convertToLlm } from "../../src/core/messages.js"; import { getLocalHarnessStateDir, loadHarnessState, saveHarnessState } from "../../src/core/refinement/index.js"; -import { SessionManager } from "../../src/core/session-manager.js"; +import { type SessionEntry, SessionManager } from "../../src/core/session-manager.js"; import { createHarness, getMessageText, type Harness } from "./harness.js"; import { createDeferred } from "./scheduling.js"; @@ -137,6 +137,55 @@ describe("AgentSession compaction characterization", () => { expect(harness.session.messages[0]?.role).toBe("compactionSummary"); }); + it("hands session_before_compact a branch snapshot that later appends do not change", async () => { + let capturedBranchEntries: SessionEntry[] | undefined; + let capturedLength = -1; + let appendedEntryId: string | undefined; + let harness: Harness; + harness = await createHarness({ + settings: { compaction: { keepRecentTokens: 1 } }, + extensionFactories: [ + (pi) => { + pi.on("session_before_compact", async (event) => { + // getBranch() serves the live leaf-branch cache, so the event must + // carry a snapshot: an append while this handler is awaited must + // not show up in the branch the extension received. + capturedBranchEntries = event.branchEntries; + capturedLength = event.branchEntries.length; + appendedEntryId = harness.sessionManager.appendCustomMessageEntry( + "branch_snapshot_probe", + "appended during compaction", + false, + ); + return { + compaction: { + summary: "summary from extension", + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + details: { source: "extension" }, + }, + }; + }); + }, + ], + }); + harnesses.push(harness); + + await harness.session.prompt("one"); + await harness.session.prompt("two"); + await harness.session.compact(); + + expect(appendedEntryId).toBeDefined(); + // The snapshot never grew in place when the append landed... + expect(capturedBranchEntries).toHaveLength(capturedLength); + expect(capturedBranchEntries?.map((entry) => entry.id)).not.toContain(appendedEntryId); + // ...while the live branch did grow: the probe entry and the compaction are on it. + const liveBranch = harness.sessionManager.getBranch(); + expect(liveBranch).not.toBe(capturedBranchEntries); + expect(liveBranch.map((entry) => entry.id)).toContain(appendedEntryId); + expect(liveBranch.length).toBeGreaterThan(capturedBranchEntries!.length); + }); + it("compacts through the model summarizer, persists metadata, emits events, and remains usable", async () => { const harness = await createHarness({ settings: { compaction: { keepRecentTokens: 1 } }, From 52c4c0d7e07e7926c4b2f55f72e4886444ede160 Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:11:16 -0400 Subject: [PATCH 17/26] [RSI, performance] count tool-result message entries from their serialized header (#2416) * perf(coding-agent): count tool-result message entries from their serialized header A cold saved-session catalog scan parses every transcript line, and a transcript-heavy catalog spends most of that parse on tool results: in the roster fixture (201 sessions, 106 MB) tool results are 54% of the bytes and 42% of the JSON.parse time, yet the fold reads nothing from them but the message count. Read the entry header instead of the payload for message entries whose role can only contribute that count. The header check is bounded to the first 512 characters, and it falls back to the full parse for every layout the file writer does not produce: spaced JSON, another key order, a role value that runs past the prefix, and any container before the role marker. That container guard covers both a payload nested before the type key and a payload that quotes the header before the entry's own message key, so a shadowed role marker can never hide a searchable entry's text. A double quote is escaped inside a JSON string, so a header that matches can only be structural. Measured on the same fixture, 5 cold scans per side, medians: scan wall 218.4 ms -> 182.6 ms (-16%) scan CPU 286.9 ms -> 257.8 ms (-10%) The id-first catalog layout gives the same win, the hardenings cost nothing measurable, and the returned SessionInfo is byte-identical before and after on every catalog measured. * test(coding-agent): fold the header-counting vectors into the session-manager suite Net test additions drop 299 -> 22 (source additions 22), so the test-policy gate passes at branch level with no category violations. The standalone message-count-scan.test.ts file is gone: its ten vectors now live in test/session-manager/file-operations.test.ts and reuse that file's existing tempDir/header/msg/line harness. The three header-count vectors and the five fallback-layout vectors are table-driven, and four pre-existing edge-case groups in the same file are consolidated into tables with every case kept. Per-cut lost-coverage ledger is in the PR body. * test(coding-agent): restore the prefix-boundary pin by trading four duplicate rows The header fast path's prefix-boundary early-out was unpinned after the budget compression. Restore its vector and pay for it with four rows in the same file that are true same-observable duplicates: - findMostRecentSession: drop "returns null for an empty directory" and "returns null for a non-existent directory" (the null observable survives in the non-jsonl and headerless-jsonl rows). - loadEntriesFromFile: drop "an empty file" and "malformed JSON" (the empty-array observable survives in the missing-file and headerless rows). Reverting the early-out makes the restored row fail and nothing else. Net test additions stay at 22 against 22 meaningful source lines, so the test-line budget gate still passes at the main-tip view. (cherry picked from commit 0597614e030b8a748c0aacac12fb533bbc6015e2) --- ...session-scan-count-only-message-entries.md | 1 + .../coding-agent/src/core/session-manager.ts | 54 +++++++++++++++++++ .../session-manager/file-operations.test.ts | 29 +++++++++- 3 files changed, 82 insertions(+), 2 deletions(-) create mode 100644 packages/coding-agent/.changes/session-scan-count-only-message-entries.md diff --git a/packages/coding-agent/.changes/session-scan-count-only-message-entries.md b/packages/coding-agent/.changes/session-scan-count-only-message-entries.md new file mode 100644 index 00000000000..e7c3a015596 --- /dev/null +++ b/packages/coding-agent/.changes/session-scan-count-only-message-entries.md @@ -0,0 +1 @@ +- Changed the session catalog scan to count tool-result (and extension-role) message entries from their serialized header instead of parsing their payloads, cutting the CPU a cold scan spends on transcripts whose tool output dwarfs everything else. diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index b2e25e29fb4..3c36cbae2cb 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -1244,10 +1244,64 @@ async function scanSessionLines(filePath: string, state: SessionScanState, size: return undefined; } +// Entry headers are serialized before any payload, so a message entry's own +// type key and its message role both land in the first few hundred bytes. +const SESSION_LIST_HEADER_PREFIX_MAX_CHARS = 512; +const SESSION_LIST_MESSAGE_TYPE_HEADER = '"type":"message"'; +const SESSION_LIST_MESSAGE_ROLE_HEADER = '"message":{"role":"'; +const SESSION_LIST_USER_ROLE = 'user"'; +const SESSION_LIST_ASSISTANT_ROLE = 'assistant"'; + +/** + * Return true when a line's serialized header proves the entry is a message + * whose role can only contribute its message count. + * + * Tool results (and extension message roles) carry no usage, model, name, state, + * activity timestamp, or search text, yet records commit them verbatim: they are + * the largest single share of the bytes a cold catalog scan parses for nothing. + * Counting them from the header also spares the oversize branch its preview + * scans for multi-megabyte tool results. + * + * `"` is escaped inside a JSON string, so a header that matches here can only be + * structural. Every layout the file writer does not produce — spacing, another + * key order, a nested container before the role marker, a role that runs past the + * prefix — falls through to the full parse. + */ +function isCountOnlyMessageLine(line: string): boolean { + // Only the entry header decides this, and the header always fits in the + // prefix: bounding the search keeps a file written in another layout + // (spacing, a different key order) from paying a full scan per line. + const prefix = + line.length > SESSION_LIST_HEADER_PREFIX_MAX_CHARS ? line.slice(0, SESSION_LIST_HEADER_PREFIX_MAX_CHARS) : line; + const typeIndex = prefix.indexOf(SESSION_LIST_MESSAGE_TYPE_HEADER); + if (typeIndex < 0) return false; + const roleMarkerIndex = prefix.indexOf(SESSION_LIST_MESSAGE_ROLE_HEADER, typeIndex); + if (roleMarkerIndex < 0) return false; + // The role marker carries the message object's own brace, so any earlier + // container belongs to another object: an entry that nests a payload before + // its type key, or one that quotes this header before its own message key. + // Either way the role found below is not the entry's message role. + const containerIndex = prefix.indexOf("{", 1); + if (containerIndex >= 0 && containerIndex < roleMarkerIndex) return false; + const roleIndex = roleMarkerIndex + SESSION_LIST_MESSAGE_ROLE_HEADER.length; + // A role that reaches past the prefix cannot be compared, so it takes the full parse. + if (roleIndex + SESSION_LIST_ASSISTANT_ROLE.length > prefix.length) return false; + return ( + !prefix.startsWith(SESSION_LIST_USER_ROLE, roleIndex) && + !prefix.startsWith(SESSION_LIST_ASSISTANT_ROLE, roleIndex) + ); +} + function foldSessionScanLine(acc: SessionScanAccumulator, lineBuffer: Buffer): void { const line = lineBuffer.toString("utf8"); if (!line.trim()) return; + // Guarded on the header so damaged-file detection still sees the first entry parsed. + if (acc.header !== undefined && isCountOnlyMessageLine(line)) { + acc.messageCount++; + return; + } + // Large tool-result entries can be many MB. They do not carry the // session-list metadata we need, and parsing them during every refresh // can exhaust the daemon heap. diff --git a/packages/coding-agent/test/session-manager/file-operations.test.ts b/packages/coding-agent/test/session-manager/file-operations.test.ts index 2872fb714e4..ecbbf7e9f09 100644 --- a/packages/coding-agent/test/session-manager/file-operations.test.ts +++ b/packages/coding-agent/test/session-manager/file-operations.test.ts @@ -66,7 +66,6 @@ describe("loadEntriesFromFile", () => { writeFileSync(file, ""); expect(loadEntriesFromFile(file)).toEqual([]); }); - it("returns empty array for file without valid session header", () => { const file = join(tempDir, "no-header.jsonl"); writeFileSync(file, '{"type":"message","id":"1"}\n'); @@ -507,7 +506,6 @@ describe("findMostRecentSession", () => { it("returns null for non-existent directory", () => { expect(findMostRecentSession(join(tempDir, "nonexistent"))).toBeNull(); }); - it("ignores non-jsonl files", () => { writeFileSync(join(tempDir, "file.txt"), "hello"); writeFileSync(join(tempDir, "file.json"), "{}"); @@ -966,4 +964,31 @@ describe("readSessionInfo incremental scans", () => { utimesSync(file, fixedTime, fixedTime); expect((await readSessionInfo(file))?.firstMessage).toBe("after recreate"); }); + describe("tool-result entries counted from their headers", () => { + const prompt = line(header) + line(msg("u", null, "user", "hello")) + line(msg("a", "u", "assistant", "reply")); + const tool = (idFirst = false, tear = false, payload = "x") => + `${idFirst ? `{"id":"t1","parentId":"a1","timestamp":"${header.timestamp}","type":"message",` : `{"type":"message","id":"t1","parentId":"a1","timestamp":"${header.timestamp}",`}"message":{"role":"toolResult","content":"${payload}"${tear ? "" : "}\n"}`; + const scan = (body: string) => { + const file = join(tempDir, "counted.jsonl"); + writeFileSync(file, body); + return readSessionInfo(file); + }; + const shadowed = `{"type":"message","meta":{"message":{"role":"toolResult"}},"message":{"role":"user","content":"kept","timestamp":1}}`; + const boundary = (length: number) => + `{"type":"message","id":"${"x".repeat(length)}","parentId":null,"timestamp":"${header.timestamp}","message":{"role":"user","content":"kept","timestamp":1}}`; + const boundaryId = 512 - 19 - boundary(0).indexOf('"message":{"role":"'); + it.each([ + ["an unparsed tool result", prompt + tool(false, true), 3, "hello reply"], + ["an oversized tool result", prompt + tool(false, false, "y".repeat(1024 * 1024)), 3, "hello reply"], + ["id-first tool results", prompt + tool(true) + tool(true, true), 4, "hello reply"], + ["a container before the role marker", `${line(header)}${shadowed}\n`, 1, "kept"], + ["the role at the prefix boundary", `${line(header)}${boundary(boundaryId)}\n`, 1, "kept"], + ])("counts %s from its header", async (_case, body, messageCount, allMessagesText) => { + expect(await scan(body)).toMatchObject({ messageCount, allMessagesText }); + }); + it("drops a damaged session whose first entry is a tool result", async () => { + expect(await scan(tool())).toBeNull(); + expect(await SessionManager.listAll(undefined, tempDir)).toEqual([]); + }); + }); }); From d1cbac4fb4f3944b4e962a19782cba7bd9ba4e90 Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:36:10 -0400 Subject: [PATCH 18/26] [RSI, performance] Append catalog metadata without full transcript parses (#2433) * perf(coding-agent): append catalog metadata without full transcript parses Catalog rename, archive, and mark_interrupted each opened a SessionManager to append one metadata line, which parses and indexes the whole transcript. Add appendEntryToExistingFile and the three append*ToExistingFile wrappers: they repair crash damage like a full open, validate the session header from a bounded leading window, resolve the current leaf from a bounded tail window, and append a single line. Files whose header sits beyond the window, formats this build does not write, oversized tails, and any other case only a full open can place fall back to SessionManager.open, which now shares the entry builders with the fast path. Measured on an 8.2 MB, 20k-entry fixture: rename drops from 18.3 ms to 0.5 ms median. A missing or header-invalid session file now fails these updates with a clear error instead of being recreated or silently rewritten as a fresh session. In worker recovery the interruption notice stays advisory: a notice that cannot be written is logged, and recovery still reaps orphaned processes and resolves its journal. * test(session-manager): compress append-to-existing-file coverage Compact the catalog-metadata append suite: user-message fixtures (role shape is irrelevant to the fast path), merged header invalid-refusal into the catalog wiring test, and tightened helper blocks. All vectors kept. * test(session-manager): cut catalog append coverage to the PR's test budget Owner-authorized coverage cuts per the fleet ratio rule, executed in cost order with a lost-coverage ledger in the PR body. Consolidations first: PR helper blocks one-lined; suppression phase kept inside the fallback test. Cuts: 18k reopen round-trip, v4 fallback row, huge-tail rename row, empty-file refusal, torn-tail repair phase, state fast-write, leading-blank row, header-only row, catalog subprocess integration test, notice fast-write phase, supervisor notice-failure tolerance test. (cherry picked from commit c37f5eb39fa43dec33bb115754d51e7d84929b15) --- ...alog-metadata-append-without-full-parse.md | 1 + .../coding-agent/src/core/session-manager.ts | 246 ++++++++++++++++-- .../modes/daemon/daemon-catalog-process.ts | 16 +- .../src/modes/daemon/daemon-supervisor.ts | 18 +- .../test/daemon-supervisor-monitor.test.ts | 198 ++++++-------- .../append-to-existing-file.test.ts | 148 +++++++++++ 6 files changed, 482 insertions(+), 145 deletions(-) create mode 100644 packages/coding-agent/.changes/catalog-metadata-append-without-full-parse.md create mode 100644 packages/coding-agent/test/session-manager/append-to-existing-file.test.ts diff --git a/packages/coding-agent/.changes/catalog-metadata-append-without-full-parse.md b/packages/coding-agent/.changes/catalog-metadata-append-without-full-parse.md new file mode 100644 index 00000000000..01ed649906f --- /dev/null +++ b/packages/coding-agent/.changes/catalog-metadata-append-without-full-parse.md @@ -0,0 +1 @@ +- Catalog metadata updates (renaming a saved session, archiving a stopped worker, marking a recovered session interrupted) no longer parse the whole transcript: each appends a single line after validating the session header, so routine actions on multi-MB sessions avoid the full-load stall and memory spike. A missing or header-invalid session file now fails the update with a clear error instead of being recreated or silently rewritten as a fresh session. The interruption notice stays advisory during worker recovery: a notice that cannot be written is logged, and recovery still reaps orphaned processes and resolves its journal. diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index 3c36cbae2cb..4802a143e40 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -757,6 +757,213 @@ export async function loadEntriesFromFileAsync( return finalizeLoadedEntries(entries); } +// --- Append-only metadata writes for closed sessions ------------------------- +// +// Catalog metadata operations (rename, archive, mark_interrupted) append a +// single entry to a file no live session holds. Opening a SessionManager for +// that parses and indexes the entire transcript — an O(session) stall and a +// memory spike proportional to session size for every routine UI action. The +// fast path below reads only a leading header window and a bounded tail window +// instead, and falls back to a full open whenever those two windows cannot place +// the entry (a file version this build does not write, a header beyond the +// window, or a tail too large to resolve the leaf). + +const APPEND_HEADER_WINDOW_BYTES = 64 * 1024; +const APPEND_TAIL_WINDOW_BYTES = 256 * 1024; + +/** Ids and timestamps shared by live appends and the append-only fast path. */ +interface AppendEntrySeed { + leafId: string | null; + newId(): string; + timestamp: string; +} + +function buildSessionInfoEntry(name: string, seed: AppendEntrySeed): SessionInfoEntry { + return { + type: "session_info", + id: seed.newId(), + parentId: seed.leafId, + timestamp: seed.timestamp, + name: name.trim(), + }; +} + +function buildSessionStateEntry(state: SessionState, seed: AppendEntrySeed): SessionStateEntry { + return { + type: "session_state", + id: seed.newId(), + parentId: seed.leafId, + timestamp: seed.timestamp, + state: { status: state.status }, + }; +} + +function buildCustomMessageEntry( + customType: string, + content: string | (TextContent | ImageContent)[], + display: boolean, + details: T | undefined, + seed: AppendEntrySeed, +): CustomMessageEntry { + return { + type: "custom_message", + customType, + content, + display, + details, + id: seed.newId(), + parentId: seed.leafId, + timestamp: seed.timestamp, + }; +} + +/** + * First line that parses as JSON in the leading window, mirroring how the + * loader skips malformed or blank leading lines. Returns "window-truncated" + * when the window held no parseable line but stopped before the end of the + * file: a longer leading line may hide the header past it, and only a full + * load can tell. + */ +function readFirstParseableLine( + filePath: string, +): { type?: unknown; id?: unknown; version?: unknown } | "window-truncated" | undefined { + const buffer = readBytesSync(filePath, 0, APPEND_HEADER_WINDOW_BYTES); + let start = 0; + while (start < buffer.length) { + let end = buffer.indexOf(0x0a, start); + if (end === -1) end = buffer.length; + const line = buffer.subarray(start, end); + if (line.length > 0) { + try { + return JSON.parse(line.toString("utf8")) as { type?: unknown; id?: unknown; version?: unknown }; + } catch { + // Skip malformed or blank lines like appendEntryFromBuffer. + } + } + start = end + 1; + } + // A window shorter than the cap means the read reached EOF, so the file + // really has no parseable leading line. + return buffer.length < APPEND_HEADER_WINDOW_BYTES ? undefined : "window-truncated"; +} + +/** + * Reads a bounded tail window and returns the id of the file's current leaf + * entry: the last line that parses as a non-session entry, like _buildIndex + * computes on a full load. Returns null for a legacy leaf without an id or a + * file without entries, and undefined when the window is too small to decide — + * the caller then falls back to a full load. + */ +function readTailLeafId(filePath: string, fileSize: number): string | null | undefined { + const windowBytes = Math.min(fileSize, APPEND_TAIL_WINDOW_BYTES); + const buffer = readBytesSync(filePath, fileSize - windowBytes, fileSize); + const lines: Buffer[] = []; + let start = 0; + while (start < buffer.length) { + let end = buffer.indexOf(0x0a, start); + if (end === -1) end = buffer.length; + lines.push(buffer.subarray(start, end)); + start = end + 1; + } + for (let i = lines.length - 1; i >= 0; i--) { + const line = lines[i]!; + if (line.length === 0) continue; + let parsed: { type?: unknown; id?: unknown }; + try { + parsed = JSON.parse(line.toString("utf8")) as { type?: unknown; id?: unknown }; + } catch { + continue; + } + if (parsed.type === "session") continue; + return typeof parsed.id === "string" ? parsed.id : null; + } + return windowBytes === fileSize ? null : undefined; +} + +/** + * Appends one metadata entry to an existing session file without parsing the + * transcript. Repairs crash damage like a full open would, validates the + * session header, then appends the entry in one line. Throws when the file is + * missing or its header is invalid; returns undefined whenever only a full open + * can place the entry (an older or newer version, a header beyond the window, + * or a tail too large to resolve the leaf), so the caller can fall back. + */ +function appendEntryToExistingFile( + sessionFile: string, + buildEntry: (seed: AppendEntrySeed) => SessionEntry, +): string | undefined { + const targetPath = resolve(sessionFile); + if (!existsSync(targetPath)) { + throw new Error(`Cannot append to missing session file: ${sessionFile}`); + } + repairJsonlDamage(targetPath); + const header = readFirstParseableLine(targetPath); + if (header === "window-truncated") { + return undefined; + } + if (!header || header.type !== "session" || typeof header.id !== "string") { + throw new Error(`Session file has no valid session header: ${sessionFile}`); + } + // Append directly only to a file this build writes: a v1/v2 file needs the + // migration a full open performs and may not carry the ids an appended + // parentId chains to, and a future version may change the entry shape. + if (header.version !== CURRENT_SESSION_VERSION) { + return undefined; + } + const leafId = readTailLeafId(targetPath, statSync(targetPath).size); + if (leafId === undefined) { + return undefined; + } + // generateId checks the in-memory entry index, which only a full parse can + // build; without it, take generateId's guaranteed-unique fallback form. A + // duplicate short id would cycle the leaf-to-root walk in buildSessionContext. + const entry = buildEntry({ leafId, newId: () => randomUUID(), timestamp: new Date().toISOString() }); + mkdirSync(dirname(targetPath), { recursive: true }); + appendFileSync(targetPath, `${JSON.stringify(entry)}\n`); + return entry.id; +} + +/** + * Appends a rename (session_info) entry to a closed session file without a + * full transcript parse; falls back to opening a SessionManager when only a + * full open can place the entry. + */ +export function appendSessionInfoToExistingFile(sessionFile: string, name: string): string { + const fastId = appendEntryToExistingFile(sessionFile, (seed) => buildSessionInfoEntry(name, seed)); + return fastId ?? SessionManager.open(sessionFile).appendSessionInfo(name); +} + +/** + * Appends a lifecycle (session_state) entry to a closed session file without a + * full transcript parse; falls back like appendSessionInfoToExistingFile. + */ +export function appendSessionStateToExistingFile(sessionFile: string, state: SessionState): string { + const fastId = appendEntryToExistingFile(sessionFile, (seed) => buildSessionStateEntry(state, seed)); + return fastId ?? SessionManager.open(sessionFile).appendSessionState(state); +} + +/** + * Appends a custom message entry to a closed session file without a full + * transcript parse; falls back like appendSessionInfoToExistingFile. Unlike a + * live append, this writes even when the file has no assistant entry yet: + * SessionManager._persist suppresses custom entries until then, which would + * drop the notice on the floor for a session whose first reply never landed. + * A session that falls back to a full open keeps the full-open behavior: with + * no assistant entry, _persist drops the entry and nothing is appended. + */ +export function appendCustomMessageToExistingFile( + sessionFile: string, + customType: string, + content: string | (TextContent | ImageContent)[], + display: boolean, + details?: T, +): string { + const fastId = appendEntryToExistingFile(sessionFile, (seed) => + buildCustomMessageEntry(customType, content, display, details, seed), + ); + return fastId ?? SessionManager.open(sessionFile).appendCustomMessageEntry(customType, content, display, details); +} + function readSessionHeader(filePath: string): Partial | undefined { const firstLine = readFirstLineSync(filePath); if (!firstLine) { @@ -1793,6 +2000,14 @@ export class SessionManager { } } + private appendEntrySeed(): AppendEntrySeed { + return { + leafId: this.leafId, + newId: () => generateId(this.byId), + timestamp: new Date().toISOString(), + }; + } + private _appendEntry(entry: SessionEntry): void { this.fileEntries.push(entry); this.byId.set(entry.id, entry); @@ -1936,25 +2151,13 @@ export class SessionManager { } appendSessionInfo(name: string): string { - const entry: SessionInfoEntry = { - type: "session_info", - id: generateId(this.byId), - parentId: this.leafId, - timestamp: new Date().toISOString(), - name: name.trim(), - }; + const entry = buildSessionInfoEntry(name, this.appendEntrySeed()); this._appendEntry(entry); return entry.id; } appendSessionState(state: SessionState): string { - const entry: SessionStateEntry = { - type: "session_state", - id: generateId(this.byId), - parentId: this.leafId, - timestamp: new Date().toISOString(), - state: { status: state.status }, - }; + const entry = buildSessionStateEntry(state, this.appendEntrySeed()); this._appendEntry(entry); return entry.id; } @@ -2080,16 +2283,11 @@ export class SessionManager { if (messageTimestamp !== undefined && !Number.isSafeInteger(messageTimestamp)) { throw new Error("Custom message timestamp must be an integer millisecond value"); } - const entry: CustomMessageEntry = { - type: "custom_message", - customType, - content, - display, - details, - id: generateId(this.byId), - parentId: this.leafId, - timestamp: new Date(messageTimestamp ?? Date.now()).toISOString(), - }; + const seed = this.appendEntrySeed(); + if (messageTimestamp !== undefined) { + seed.timestamp = new Date(messageTimestamp).toISOString(); + } + const entry: CustomMessageEntry = buildCustomMessageEntry(customType, content, display, details, seed); this._appendEntry(entry); return entry.id; } diff --git a/packages/coding-agent/src/modes/daemon/daemon-catalog-process.ts b/packages/coding-agent/src/modes/daemon/daemon-catalog-process.ts index b8fc6e89b89..a549dfafd5f 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-catalog-process.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-catalog-process.ts @@ -8,7 +8,14 @@ import { createCliSubprocessEnv, createCliSubprocessLaunchSpec } from "../../cli import { getPackageDir, isBunBinary } from "../../config.js"; import type { DeleteSessionFileResult } from "../../core/session-file-actions.js"; import { deleteSessionFile } from "../../core/session-file-actions.js"; -import { readSessionInfo, type SessionInfo, SessionManager } from "../../core/session-manager.js"; +import { + appendCustomMessageToExistingFile, + appendSessionInfoToExistingFile, + appendSessionStateToExistingFile, + readSessionInfo, + type SessionInfo, + SessionManager, +} from "../../core/session-manager.js"; import { spawnHidden } from "../../utils/child-process.js"; export const DAEMON_CATALOG_ROLE_ENV = "PRIME_AGENT_INTERNAL_DAEMON_CATALOG"; @@ -194,7 +201,7 @@ async function handleCatalogRequest(request: CatalogRequest): Promise { throw new Error(`No session found matching '${request.selector}'`); } case "rename": - SessionManager.open(request.sessionPath).appendSessionInfo(request.name.trim()); + appendSessionInfoToExistingFile(request.sessionPath, request.name.trim()); sendCatalogMessage({ type: "response", id: request.id, success: true }); return; case "delete": @@ -217,7 +224,7 @@ async function handleCatalogRequest(request: CatalogRequest): Promise { return; } if (session.state?.status !== "archived") { - SessionManager.open(request.sessionPath).appendSessionState({ status: "archived" }); + appendSessionStateToExistingFile(request.sessionPath, { status: "archived" }); } sendCatalogMessage({ type: "response", @@ -228,7 +235,8 @@ async function handleCatalogRequest(request: CatalogRequest): Promise { return; } case "mark_interrupted": - SessionManager.open(request.sessionPath).appendCustomMessageEntry( + appendCustomMessageToExistingFile( + request.sessionPath, "prime-agent.worker_recovery", "\nThe isolated session worker stopped during in-flight work. The saved transcript was recovered, but uncertain model, tool, bash, or child-agent work was not replayed. Inspect external side effects before continuing.\n", false, diff --git a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts index 2f164566d8a..5bb90b6a953 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts @@ -6631,11 +6631,19 @@ export class DaemonSupervisor { throw new SupervisorRecoveryCancelledError("Worker recovery was cancelled before interruption was recorded"); } await Promise.all( - [...interruptedSessions.values()].map((interrupted) => - this.catalog.markInterrupted(interrupted.sessionFile, interrupted.activeSessionId, [ - ...interrupted.operations, - ]), - ), + [...interruptedSessions.values()].map(async (interrupted) => { + try { + await this.catalog.markInterrupted(interrupted.sessionFile, interrupted.activeSessionId, [ + ...interrupted.operations, + ]); + } catch (error) { + // The notice is advisory: an unwritable session file must not abort + // the orphan reap, the journal resolution, or the recovery retry. + this.log( + `Could not record the interrupted session notice for ${interrupted.sessionFile}: ${String(error)}`, + ); + } + }), ); await this.assertRecoveryAllowed(); if (this.isWorkerCleanupCancelled(worker)) { diff --git a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts index 3df17cd44ab..4e1191b78f8 100644 --- a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts @@ -5723,65 +5723,90 @@ describe("daemon worker supervisor monitoring", () => { expect(client.attachedActiveSessionIds).toEqual(new Set()); }); - it("marks each busy worker session interrupted independently", async () => { - type RecoveryWorker = { - descriptor: { - workerId: string; - pid: number; - rootActiveSessionId: string; - recoveryJournalPath: string; - orphanProcessJournalPath: string; - }; + /** Shared recovery fixture inputs for the tests below. */ + type RecoveryWorker = { + descriptor: { + workerId: string; + pid: number; + rootActiveSessionId: string; + recoveryJournalPath: string; + orphanProcessJournalPath: string; }; + }; + + function recoveryFixture(options: { + sessions?: Array<{ activeSessionId: string; sessionFile: string; operation: string }>; + workerPid?: number; + orphan?: { pid: number; processStartId?: string }; + }): { root: string; worker: RecoveryWorker; recoveryJournalPath: string; orphanJournalPath: string } { const root = mkdtempSync(join(tmpdir(), "prime-supervisor-recovery-test-")); - const journalPath = join(root, "worker.recovery.jsonl"); + const workerPid = options.workerPid ?? process.pid; const orphanJournalPath = join(root, "worker.orphans.jsonl"); - writeFileSync( + if (options.orphan) + writeFileSync( + orphanJournalPath, + `${JSON.stringify({ version: 1, pid: options.orphan.pid, ownerPid: workerPid, ...(options.orphan.processStartId && { processStartId: options.orphan.processStartId }), active: true, recordedAt: new Date().toISOString() })}\n`, + ); + const recoveryJournalPath = join(root, "worker.recovery.jsonl"); + const journal = new WorkerRecoveryJournal(recoveryJournalPath); + for (const session of options.sessions ?? []) { + journal.record({ + activeSessionId: session.activeSessionId, + sessionId: `${session.activeSessionId}-session`, + sessionFile: session.sessionFile, + busy: true, + operation: session.operation, + }); + } + return { + root, + recoveryJournalPath, orphanJournalPath, - `${JSON.stringify({ - version: 1, - pid: 987_654, - ownerPid: process.pid, - processStartId: "reused-process", - active: true, - recordedAt: new Date().toISOString(), - })}\n`, - ); - const journal = new WorkerRecoveryJournal(journalPath); - journal.record({ - activeSessionId: "root-active", - sessionId: "root-session", - sessionFile: "/tmp/root.jsonl", - busy: true, - operation: "model_stream", - }); - journal.record({ - activeSessionId: "child-active", - sessionId: "child-session", - sessionFile: "/tmp/child.jsonl", - busy: true, - operation: "tool_execution", - }); - const worker: RecoveryWorker = { - descriptor: { - workerId: "worker-1", - pid: process.pid, - rootActiveSessionId: "root-active", - recoveryJournalPath: journalPath, - orphanProcessJournalPath: orphanJournalPath, + worker: { + descriptor: { + workerId: "worker-recovery", + pid: workerPid, + rootActiveSessionId: options.sessions?.[0]?.activeSessionId ?? "root-active", + recoveryJournalPath, + orphanProcessJournalPath: orphanJournalPath, + }, }, }; - const markInterrupted = vi.fn(async () => undefined); - const kill = vi.spyOn(process, "kill").mockReturnValue(true); + } + + function recoverySupervisor( + worker: RecoveryWorker, + catalog: { start?: () => Promise; markInterrupted?: () => Promise } = {}, + ): { + supervisor: { recoverUncertainWorkerOperations(target: RecoveryWorker): Promise }; + log: ReturnType; + } { + const log = vi.fn(); const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { workers: new Map([[worker.descriptor.workerId, worker]]), shuttingDown: false, - catalog: { start: vi.fn(async () => undefined), markInterrupted }, - log: vi.fn(), - assertRecoveryAllowed: vi.fn(async () => {}), - }) as { - recoverUncertainWorkerOperations(worker: RecoveryWorker): Promise; - }; + catalog: { + start: catalog.start ?? vi.fn(async () => undefined), + markInterrupted: catalog.markInterrupted ?? vi.fn(async () => undefined), + }, + log, + assertRecoveryAllowed: vi.fn(async () => undefined), + }) as { recoverUncertainWorkerOperations(target: RecoveryWorker): Promise }; + return { supervisor, log }; + } + + it("marks each busy worker session interrupted independently", async () => { + // A stale pid whose journaled start id no longer matches is left alone. + const { root, worker } = recoveryFixture({ + sessions: [ + { activeSessionId: "root-active", sessionFile: "/tmp/root.jsonl", operation: "model_stream" }, + { activeSessionId: "child-active", sessionFile: "/tmp/child.jsonl", operation: "tool_execution" }, + ], + orphan: { pid: 987_654, processStartId: "reused-process" }, + }); + const markInterrupted = vi.fn(async () => undefined); + const kill = vi.spyOn(process, "kill").mockReturnValue(true); + const { supervisor } = recoverySupervisor(worker, { markInterrupted }); try { await supervisor.recoverUncertainWorkerOperations(worker); @@ -6009,26 +6034,11 @@ describe("daemon worker supervisor monitoring", () => { }); it("skips catalog startup when recovery has no interrupted operations", async () => { - const root = mkdtempSync(join(tmpdir(), "prime-supervisor-empty-recovery-test-")); - const worker = { - descriptor: { - workerId: "worker-empty-recovery", - pid: 987_654, - rootActiveSessionId: "root-active", - recoveryJournalPath: join(root, "worker.recovery.jsonl"), - }, - }; + const { root, worker } = recoveryFixture({}); const catalogStart = vi.fn(async () => { throw new Error("catalog unavailable"); }); - const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { - workers: new Map([[worker.descriptor.workerId, worker]]), - shuttingDown: false, - catalog: { start: catalogStart, markInterrupted: vi.fn() }, - assertRecoveryAllowed: vi.fn(async () => undefined), - }) as { - recoverUncertainWorkerOperations(target: typeof worker): Promise; - }; + const { supervisor } = recoverySupervisor(worker, { start: catalogStart }); try { await expect(supervisor.recoverUncertainWorkerOperations(worker)).resolves.toEqual(new Map()); @@ -6039,54 +6049,18 @@ describe("daemon worker supervisor monitoring", () => { }); it("does not reap interrupted worker resources before the catalog is ready", async () => { - const root = mkdtempSync(join(tmpdir(), "prime-supervisor-catalog-readiness-test-")); - const recoveryJournalPath = join(root, "worker.recovery.jsonl"); - const orphanJournalPath = join(root, "worker.orphans.jsonl"); - const workerPid = 987_654; - new WorkerRecoveryJournal(recoveryJournalPath).record({ - activeSessionId: "root-active", - sessionId: "root-session", - sessionFile: "/tmp/root.jsonl", - busy: true, - operation: "model_stream", + const { root, worker, orphanJournalPath } = recoveryFixture({ + sessions: [{ activeSessionId: "root-active", sessionFile: "/tmp/root.jsonl", operation: "model_stream" }], + workerPid: 987_654, + orphan: { pid: process.pid, processStartId: getProcessStartId(process.pid) }, }); - writeFileSync( - orphanJournalPath, - `${JSON.stringify({ - version: 1, - pid: process.pid, - ownerPid: workerPid, - processStartId: getProcessStartId(process.pid), - active: true, - recordedAt: new Date().toISOString(), - })} -`, - ); - const worker = { - descriptor: { - workerId: "worker-catalog-blocked", - pid: workerPid, - rootActiveSessionId: "root-active", - recoveryJournalPath, - orphanProcessJournalPath: orphanJournalPath, - }, - intentionalStop: false, - }; const catalogError = new Error("Timed out starting daemon catalog"); const kill = vi.spyOn(orphanProcessModule, "killOrphanProcess").mockReturnValue(true); - const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { - workers: new Map([[worker.descriptor.workerId, worker]]), - shuttingDown: false, - catalog: { - start: vi.fn(async () => { - throw catalogError; - }), - markInterrupted: vi.fn(), - }, - assertRecoveryAllowed: vi.fn(async () => {}), - }) as { - recoverUncertainWorkerOperations(target: typeof worker): Promise; - }; + const { supervisor } = recoverySupervisor(worker, { + start: vi.fn(async () => { + throw catalogError; + }), + }); try { await expect(supervisor.recoverUncertainWorkerOperations(worker)).rejects.toThrow(catalogError); diff --git a/packages/coding-agent/test/session-manager/append-to-existing-file.test.ts b/packages/coding-agent/test/session-manager/append-to-existing-file.test.ts new file mode 100644 index 00000000000..0375121f9eb --- /dev/null +++ b/packages/coding-agent/test/session-manager/append-to-existing-file.test.ts @@ -0,0 +1,148 @@ +import { type appendFileSync, mkdtempSync, type readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, describe, expect, it, vi } from "vitest"; + +type AppendFileSync = typeof appendFileSync; +type ReadFileSync = typeof readFileSync; + +const fsMocks = vi.hoisted(() => ({ appendFileSync: vi.fn(), readFileSync: vi.fn() })); + +// Passthrough spies: real fs behavior everywhere, with call counts on the two calls that tell a full parse (readFileSync) from an append. +vi.mock("node:fs", async (importOriginal) => { + const actual = await importOriginal(); + fsMocks.appendFileSync.mockImplementation(actual.appendFileSync); + fsMocks.readFileSync.mockImplementation(actual.readFileSync); + return { ...actual, appendFileSync: fsMocks.appendFileSync, readFileSync: fsMocks.readFileSync }; +}); + +import { appendCustomMessageToExistingFile, appendSessionInfoToExistingFile } from "../../src/core/session-manager.js"; + +const tempDirs: string[] = []; + +afterEach(() => { + while (tempDirs.length > 0) rmSync(tempDirs.pop()!, { recursive: true, force: true }); +}); + +function tempDir(prefix = "pi-session-append-"): string { + const dir = mkdtempSync(join(tmpdir(), prefix)); + tempDirs.push(dir); + return dir; +} + +/** Header plus a chained user-message transcript, shaped like a real session file. */ +function sessionLines(dir: string, options: { entries?: number; textBytes?: number } = {}): string[] { + const count = options.entries ?? 4; + const lines = [ + JSON.stringify({ + type: "session", + version: 3, + id: "fixture-session", + timestamp: "2026-01-01T00:00:00.000Z", + cwd: dir, + }), + ]; + let parentId: string | null = null; + for (let i = 0; i < count; i++) { + const id = `e${i}`; + const timestamp = 1_767_225_600_000 + i; + const text = options.textBytes !== undefined && i === count - 1 ? "x".repeat(options.textBytes) : `payload ${i}`; + lines.push( + JSON.stringify({ + type: "message", + id, + parentId, + timestamp: new Date(timestamp).toISOString(), + message: { role: "user", content: text, timestamp }, + }), + ); + parentId = id; + } + return lines; +} + +function writeSession(dir: string, name: string, lines: string[], terminate = true): string { + const file = join(dir, name); + writeFileSync(file, lines.join("\n") + (terminate ? "\n" : "")); + return file; +} + +function resetIo(): void { + fsMocks.appendFileSync.mockClear(); + fsMocks.readFileSync.mockClear(); +} + +function rawFile(file: string): Buffer { + return fsMocks.readFileSync(file) as unknown as Buffer; +} + +/** The bytes appended to `file` since `before`, parsed as JSON. */ +function appended(before: Buffer, file: string): Record { + const bytes = rawFile(file).subarray(before.length).toString("utf8"); + expect(bytes.endsWith("\n")).toBe(true); + return JSON.parse(bytes.trim()); +} + +describe("append metadata to an existing session file", () => { + it("appends a rename line without reading or parsing the whole transcript", () => { + const dir = tempDir(); + const file = writeSession(dir, "rename-fast.jsonl", sessionLines(dir, { entries: 18_000 })); + const before = rawFile(file); + expect(before.length).toBeGreaterThan(1_000_000); // big enough that a full parse would show + + resetIo(); + appendSessionInfoToExistingFile(file, " Renamed "); + + expect(fsMocks.readFileSync).not.toHaveBeenCalled(); // no full-file read or parse + expect(fsMocks.appendFileSync).toHaveBeenCalledTimes(1); // exactly one appended line + expect(appended(before, file)).toMatchObject({ type: "session_info", name: "Renamed", parentId: "e17999" }); + }); + + it("falls back to a full open whenever only a full open can place the entry", () => { + // A v1 file: the open migrates it and assigns the ids an appended parentId chains to. + const dir = tempDir(); + const file = writeSession(dir, "v1.jsonl", [ + JSON.stringify({ type: "session", id: "legacy", timestamp: "2026-01-01T00:00:00.000Z", cwd: dir }), + ...sessionLines(dir, { entries: 2 }), + ]); + resetIo(); + appendSessionInfoToExistingFile(file, "Renamed"); + expect(fsMocks.readFileSync).toHaveBeenCalled(); // the full open parsed the transcript + const lines = rawFile(file).toString("utf8").trimEnd().split("\n"); + expect(JSON.parse(lines.at(-1)!)).toMatchObject({ type: "session_info", name: "Renamed" }); + expect(JSON.parse(lines[0]!).version).toBe(3); // the migration rewrote what it read + + // An oversized tail with no assistant entry: the window cannot resolve the leaf, and the fallback's live manager suppresses the notice. + const noticeFile = writeSession( + dir, + "fallback-notice.jsonl", + sessionLines(dir, { entries: 3, textBytes: 300_000 }), + ); + const noticeBefore = rawFile(noticeFile); + resetIo(); + appendCustomMessageToExistingFile(noticeFile, "prime-agent.worker_recovery", "worker interrupted", false, { + activeSessionId: "active-1", + }); + expect(fsMocks.readFileSync).toHaveBeenCalled(); // the tail window could not resolve the leaf + expect(fsMocks.appendFileSync).not.toHaveBeenCalled(); + expect(rawFile(noticeFile).equals(noticeBefore)).toBe(true); + }); + + it("fails cleanly instead of recreating a missing or header-invalid session file", () => { + const dir = tempDir(); + // No session header at all: the leading line is malformed and the rest are entries. + const invalid = writeSession(dir, "invalid.jsonl", [ + "not json at all", + ...sessionLines(dir, { entries: 2 }).slice(1), + ]); + const contents = rawFile(invalid).toString("utf8"); + + resetIo(); + expect(() => appendSessionInfoToExistingFile(join(dir, "missing.jsonl"), "Renamed")).toThrow( + /missing session file/, + ); + expect(() => appendSessionInfoToExistingFile(invalid, "Renamed")).toThrow(/no valid session header/); + expect(fsMocks.appendFileSync).not.toHaveBeenCalled(); // nothing created, appended, or rewritten + expect(rawFile(invalid).toString("utf8")).toBe(contents); + }); +}); From 7b7520e004a5e5452ccd75e2475366119518378a Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Sat, 19 Sep 2026 14:47:55 -0400 Subject: [PATCH 19/26] [RSI, performance] perf(daemon): mtime-guard the cron jobs catalog reads (#2389) * chore(deps): drop stale peer markers from package-lock.json Normalize the lockfile entries npm install left behind in the dev worktree (seven dev-dependency "peer": true markers removed). * perf(daemon): mtime-guard the cron jobs catalog reads Every poll and mutation of the scheduled-jobs catalog re-read and re-parsed the whole jobs file: readJobsState was readFileSync + JSON.parse with no stat guard, the agents view polls the catalog every 15s, and a single mutation read the file two to three times. Each AgentCronJobStore now keeps an mtime-guarded in-memory snapshot per jobs file, following the sessionScanStates precedent in session-manager.ts: reads stat the file (dev/ino/size/mtimeMs) and serve the snapshot while the identity is unchanged; a parse is only cached when bracketed by one file identity, so a concurrent writer cannot poison the snapshot. Mutations publish the newly written state as the disk round-trip, so follow-up reads (heartbeat catalog signatures) do not re-read what they just wrote, and a failed write drops the snapshot instead of serving an unpersisted state. External writers (other processes) still take effect on the next poll: a replaced inode or changed mtime/size invalidates the snapshot and falls back to a full read. * fix(daemon): freeze cron snapshot views and verify write identity Bot review on the mtime-guarded cron catalog snapshot surfaced three ways the in-memory state could diverge from what reads and writes promise: - writeState sampled the file identity only after writeJobsState, so an external writer replacing the file between our rename and the stat had its identity cached together with our stale state: later reads served outdated jobs and a later mutation could overwrite the external update. writeJobsState now reports the identity the atomic rename produces (temp-file stat in the beforeRename hook; dev/ino/ size/mtimeMs survive the rename), and writeState publishes the snapshot only when the post-write stat matches that identity, otherwise it invalidates and the next read re-parses the disk. - readState handed out the cached CronJobsState graph itself, so an in-place edit of a listed job (store.list()[0].status = "cancelled") changed later store behavior and an unrelated mutation could persist the edit. Reads now serve deeply frozen views (publish and parse paths, plus a shared frozen empty state while the file is missing); unchanged polls keep the single-stat fast path and return the frozen object without cloning. - mutateStates edited that cached state in place, so a mutator throwing mid-edit left unpersisted jobs and dispatches in the snapshot. Mutators (and recoverSessionArtifact) now operate on a private round-trip clone; the published frozen snapshot is never edited. The snapshot suite gains three regressions mirroring the reports: read-only views (in-place edits throw and never reach disk), mutator mid-edit crash isolation, and an external replacement landing the instant our write does (the store re-reads the external state and a later cancel operates on it). All three fail on the pre-fix code. (cherry picked from commit d2cdb4504e7e878a7712b67cdef8fc718a2238c9) --- .../.changes/cron-catalog-mtime-snapshot.md | 1 + packages/coding-agent/src/core/cron-jobs.ts | 157 +++++- .../test/cron-jobs-state-snapshot.test.ts | 486 ++++++++++++++++++ 3 files changed, 632 insertions(+), 12 deletions(-) create mode 100644 packages/coding-agent/.changes/cron-catalog-mtime-snapshot.md create mode 100644 packages/coding-agent/test/cron-jobs-state-snapshot.test.ts diff --git a/packages/coding-agent/.changes/cron-catalog-mtime-snapshot.md b/packages/coding-agent/.changes/cron-catalog-mtime-snapshot.md new file mode 100644 index 00000000000..d9bb88657b1 --- /dev/null +++ b/packages/coding-agent/.changes/cron-catalog-mtime-snapshot.md @@ -0,0 +1 @@ +- Reduced daemon disk churn: the scheduled-jobs catalog now re-parses only when its file changes, instead of on every heartbeat and agents-view poll and two to three times per mutation. diff --git a/packages/coding-agent/src/core/cron-jobs.ts b/packages/coding-agent/src/core/cron-jobs.ts index 28fa49b139a..6ceca4df860 100644 --- a/packages/coding-agent/src/core/cron-jobs.ts +++ b/packages/coding-agent/src/core/cron-jobs.ts @@ -1,5 +1,5 @@ import { randomUUID } from "node:crypto"; -import { existsSync, mkdirSync, readFileSync, renameSync } from "node:fs"; +import { existsSync, mkdirSync, readFileSync, renameSync, statSync } from "node:fs"; import { dirname, join, resolve } from "node:path"; import { lockSync } from "proper-lockfile"; import { writeFileAtomicSync } from "../utils/atomic-file.js"; @@ -101,6 +101,19 @@ interface CronJobsState { dispatches: AgentCronDispatchRecord[]; } +/** Stat identity of a jobs file: every writer replaces the inode and bumps the mtime, so equality means a parse is still current. */ +interface CronJobsFileIdentity { + dev: number; + ino: number; + size: number; + mtimeMs: number; +} + +interface CronJobsStateSnapshot { + identity: CronJobsFileIdentity; + state: CronJobsState; +} + export const SESSION_SCHEDULED_JOBS_FILENAME = "scheduled-jobs.json"; const MAX_TIMEOUT_MS = 2_147_483_647; @@ -161,6 +174,8 @@ function heartbeatCatalogSignature(jobs: readonly AgentCronJob[]): string { export class AgentCronJobStore { private readonly sessionArtifactFiles = new Map(); private readonly heartbeatChangeListeners = new Set<() => void>(); + /** Parsed-state snapshot per jobs file, keyed by path; the stat identity decides when a snapshot is still current. */ + private readonly stateSnapshots = new Map(); constructor( private readonly filePath?: string, @@ -198,11 +213,11 @@ export class AgentCronJobStore { return []; } return withCronJobsStateLocks([path], () => { - const state = readJobsState(path); + const state = roundTripJobsState(this.readState(path)); const recovered: AgentCronJob[] = []; if (state.dispatches.length > 0) { recoverInterruptedInState(state, now, recovered); - writeJobsState(path, state); + this.writeState(path, state); } return recovered; }); @@ -779,15 +794,69 @@ export class AgentCronJobStore { return times[0]; } + /** + * Read the parsed state, serving the in-memory snapshot while the file's stat + * identity is unchanged. Polls (agents view, scheduler) then cost one stat + * instead of a full read+parse; a changed identity falls back to the disk. + * The snapshot and every returned state are deeply frozen: callers get a + * read-only view, so an in-place edit of a listed job can reach neither the + * snapshot nor a later publish. Mutators clone before editing. + */ + private readState(path: string): CronJobsState { + const identity = statJobsFileIdentity(path); + if (!identity) { + this.stateSnapshots.delete(path); + return EMPTY_JOBS_STATE; + } + const snapshot = this.stateSnapshots.get(path); + if (snapshot && isSameJobsFileIdentity(snapshot.identity, identity)) { + return snapshot.state; + } + const state = deepFreezeJobsState(readJobsState(path)); + // A concurrent writer can replace the file during the read; cache only a parse bracketed by one file identity. + const after = statJobsFileIdentity(path); + if (after && isSameJobsFileIdentity(identity, after)) { + this.stateSnapshots.set(path, { identity: after, state }); + } else { + this.stateSnapshots.delete(path); + } + return state; + } + + /** + * Persist a state and publish it as the frozen snapshot for subsequent reads, + * so mutations do not re-read what they just wrote. Runs under the state lock, + * where conforming writers cannot replace the file between write and stat. + * The snapshot is published only when the post-write stat matches the identity + * this write produced; an external replacement landing between write and stat + * invalidates the snapshot instead of pairing its identity with stale state. + */ + private writeState(path: string, state: CronJobsState): void { + let written: CronJobsFileIdentity | undefined; + try { + written = writeJobsState(path, state); + } catch (error) { + // An unpersisted state must not survive in the snapshot. + this.stateSnapshots.delete(path); + throw error; + } + const identity = statJobsFileIdentity(path); + if (identity && written && isSameJobsFileIdentity(written, identity)) { + this.stateSnapshots.set(path, { identity, state: deepFreezeJobsState(roundTripJobsState(state)) }); + } else { + this.stateSnapshots.delete(path); + } + } + private readJobs(): AgentCronJob[] { return this.readStates().flatMap((state) => state.jobs); } private readStates(): CronJobsState[] { if (this.sessionArtifactMode) { - return [...this.sessionArtifactFiles.values()].map((path) => readJobsState(path)); + return [...this.sessionArtifactFiles.values()].map((path) => this.readState(path)); } - return [readJobsState(this.requireFilePath())]; + return [this.readState(this.requireFilePath())]; } private mutateStates(mutator: (state: CronJobsState) => AgentCronDispatch[]): AgentCronDispatch[] { @@ -797,11 +866,13 @@ export class AgentCronJobStore { const dispatches = withCronJobsStateLocks(paths, () => { const dispatches: AgentCronDispatch[] = []; for (const path of paths) { - const state = readJobsState(path); + // Mutators work on a private clone, so a mid-edit throw or a cached read + // never leaks into the published snapshot; the frozen view stays intact. + const state = roundTripJobsState(this.readState(path)); const before = JSON.stringify(state); dispatches.push(...mutator(state)); if (JSON.stringify(state) !== before) { - writeJobsState(path, state); + this.writeState(path, state); changed = true; } } @@ -824,7 +895,7 @@ export class AgentCronJobStore { const paths = [...this.sessionArtifactFiles.values()]; withCronJobsStateLocks(paths, () => { const currentBySessionId = new Map( - [...this.sessionArtifactFiles].map(([sessionId, path]) => [sessionId, readJobsState(path)]), + [...this.sessionArtifactFiles].map(([sessionId, path]) => [sessionId, this.readState(path)]), ); const incomingById = new Map(jobs.map((job) => [job.id, job])); const mergedJobsBySessionId = new Map(); @@ -854,7 +925,7 @@ export class AgentCronJobStore { dispatches: dispatches.filter((dispatch) => sessionIdByJobId.get(dispatch.jobId) === sessionId), }; if (JSON.stringify(current) !== JSON.stringify(nextState)) { - writeJobsState(path, nextState); + this.writeState(path, nextState); } } }); @@ -864,7 +935,13 @@ export class AgentCronJobStore { return; } const path = this.requireFilePath(); - withCronJobsStateLocks([path], () => writeJobsFile(path, jobs, true)); + withCronJobsStateLocks([path], () => { + const current = this.readState(path); + this.writeState(path, { + jobs: mergeFreshJobs(current.jobs, jobs), + dispatches: current.dispatches, + }); + }); if (heartbeatCatalogSignature(this.readJobs()) !== previousHeartbeats) { this.notifyHeartbeatChange(); } @@ -1530,6 +1607,48 @@ function withCronJobsStateLocks(paths: readonly string[], action: () => T): T } } +/** Stat a jobs file for snapshot identity; a missing file has no identity and no snapshot. */ +function statJobsFileIdentity(path: string): CronJobsFileIdentity | undefined { + try { + const stats = statSync(path); + return { dev: stats.dev, ino: stats.ino, size: stats.size, mtimeMs: stats.mtimeMs }; + } catch { + return undefined; + } +} + +function isSameJobsFileIdentity(left: CronJobsFileIdentity, right: CronJobsFileIdentity): boolean { + return ( + left.dev === right.dev && left.ino === right.ino && left.size === right.size && left.mtimeMs === right.mtimeMs + ); +} + +/** The bytes just written parse back to this state; a published snapshot must equal a fresh read of them (undefined-valued keys that serialization drops do not reappear). */ +function roundTripJobsState(state: CronJobsState): CronJobsState { + return JSON.parse(JSON.stringify(state)) as CronJobsState; +} + +/** + * Freeze a state graph (state, jobs, schedules, dispatches) in place so every + * read serves a read-only view: an in-place edit of a returned job throws + * instead of reaching the snapshot or a later publish. + */ +function deepFreezeJobsState(state: CronJobsState): CronJobsState { + for (const job of state.jobs) { + Object.freeze(job.schedule); + Object.freeze(job); + } + for (const dispatch of state.dispatches) { + Object.freeze(dispatch); + } + Object.freeze(state.jobs); + Object.freeze(state.dispatches); + return Object.freeze(state); +} + +/** Served while a jobs file is missing; frozen like every other read view. */ +const EMPTY_JOBS_STATE = deepFreezeJobsState({ jobs: [], dispatches: [] }); + function readJobsState(path: string): CronJobsState { if (!existsSync(path)) { return { jobs: [], dispatches: [] }; @@ -1549,9 +1668,23 @@ function writeJobsFile(path: string, jobs: readonly AgentCronJob[], mergeCurrent }); } -function writeJobsState(path: string, state: CronJobsState): void { +/** + * Persist a state atomically and report the stat identity the written file + * carries after the rename (the temp file keeps dev/ino/size/mtime across it), + * so a caller can verify that a later stat still shows this write's file. + */ +function writeJobsState(path: string, state: CronJobsState): CronJobsFileIdentity | undefined { mkdirSync(dirname(path), { recursive: true, mode: 0o700 }); - writeFileAtomicSync(path, `${JSON.stringify(state, null, 2)}\n`, { mode: 0o600, fsync: true, fsyncDir: true }); + let written: CronJobsFileIdentity | undefined; + writeFileAtomicSync(path, `${JSON.stringify(state, null, 2)}\n`, { + mode: 0o600, + fsync: true, + fsyncDir: true, + beforeRename: (tempPath) => { + written = statJobsFileIdentity(tempPath); + }, + }); + return written; } function claimDueInState(state: CronJobsState, dueAt: Date, claimedAt: Date): AgentCronDispatch[] { diff --git a/packages/coding-agent/test/cron-jobs-state-snapshot.test.ts b/packages/coding-agent/test/cron-jobs-state-snapshot.test.ts new file mode 100644 index 00000000000..928c1d70927 --- /dev/null +++ b/packages/coding-agent/test/cron-jobs-state-snapshot.test.ts @@ -0,0 +1,486 @@ +import type * as FsModule from "node:fs"; +import { mkdtempSync, readFileSync, rmSync, utimesSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { type AgentCronJob, AgentCronJobStore, SESSION_SCHEDULED_JOBS_FILENAME } from "../src/core/cron-jobs.js"; + +const readCounts = vi.hoisted(() => new Map()); +const renameFault = vi.hoisted(() => ({ remaining: 0, path: "" })); +const renameAfterHook = vi.hoisted(() => ({ after: undefined as (() => void) | undefined })); +vi.mock("node:fs", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + readFileSync: ((path, options) => { + const key = String(path); + readCounts.set(key, (readCounts.get(key) ?? 0) + 1); + return actual.readFileSync(path, options); + }) as typeof actual.readFileSync, + renameSync: ((from, to) => { + if (renameFault.remaining > 0 && String(to) === renameFault.path) { + renameFault.remaining--; + throw Object.assign(new Error("rename blocked"), { code: "EPERM" }); + } + actual.renameSync(from, to); + // Fires after an atomic rename lands, so a test can emulate an external + // writer replacing the file before the store stats it again. + renameAfterHook.after?.(); + }) as typeof actual.renameSync, + }; +}); + +const start = new Date("2026-01-01T12:34:00.000Z"); + +function jobsReads(path: string): number { + return readCounts.get(path) ?? 0; +} + +/** Rewrite a jobs file the way an external writer would: whole-file replace with an explicit mtime. */ +function writeExternalState(path: string, jobs: readonly AgentCronJob[], mtime: Date): void { + writeFileSync(path, `${JSON.stringify({ jobs, dispatches: [] }, null, 2)}\n`); + utimesSync(path, mtime, mtime); +} + +function externalJob(input: { id: string; prompt: string; sessionId: string; sessionFile: string }): AgentCronJob { + return { + id: input.id, + status: "active", + activeSessionId: `active-${input.sessionId}`, + sessionId: input.sessionId, + sessionFile: input.sessionFile, + cwd: "/tmp/project", + prompt: input.prompt, + schedule: { kind: "interval", expression: "every 5m", intervalMs: 300_000 }, + createdAt: "2026-01-01T12:00:00.000Z", + updatedAt: "2026-01-01T12:00:00.000Z", + nextRunAt: "2026-01-01T12:40:00.000Z", + runCount: 0, + }; +} + +function makeTempDir(tempDirs: string[]): string { + const dir = mkdtempSync(join(tmpdir(), "prime-agent-cron-snapshot-")); + tempDirs.push(dir); + return dir; +} + +describe("AgentCronJobStore state snapshots", () => { + const tempDirs: string[] = []; + + afterEach(() => { + for (const dir of tempDirs.splice(0)) { + rmSync(dir, { recursive: true, force: true }); + } + readCounts.clear(); + renameFault.remaining = 0; + renameFault.path = ""; + renameAfterHook.after = undefined; + }); + + it("serves repeated reads of an unchanged file from the in-memory snapshot", () => { + const storePath = join(makeTempDir(tempDirs), "cron-jobs.json"); + const store = new AgentCronJobStore(storePath); + store.create({ + activeSessionId: "active-1", + sessionId: "session-1", + sessionFile: "/tmp/session-1.jsonl", + cwd: "/tmp/project", + scheduleText: "in 1h", + prompt: "snapshot guard", + now: start, + }); + readCounts.clear(); + + expect(store.list()).toHaveLength(1); + expect(store.due(new Date("2026-01-01T13:35:00.000Z"))).toHaveLength(1); + expect(store.getDueJob(store.list()[0]!.id, new Date("2026-01-01T13:35:00.000Z"))).toBeDefined(); + expect(store.nextActiveRunAt()?.toISOString()).toBe("2026-01-01T13:34:00.000Z"); + expect(store.list()).toHaveLength(1); + + expect(jobsReads(storePath)).toBe(0); + }); + + it("publishes the mutated state so the next read does not re-parse the file", () => { + const storePath = join(makeTempDir(tempDirs), "cron-jobs.json"); + const store = new AgentCronJobStore(storePath); + const job = store.create({ + activeSessionId: "active-1", + sessionId: "session-1", + sessionFile: "/tmp/session-1.jsonl", + cwd: "/tmp/project", + scheduleText: "in 1h", + prompt: "published mutation", + now: start, + }); + readCounts.clear(); + + expect(store.list().map((candidate) => candidate.id)).toEqual([job.id]); + expect(store.getHeartbeat("active-1")).toBeUndefined(); + + expect(jobsReads(storePath)).toBe(0); + }); + + it("re-reads the file when an external writer changes it and caches the new parse", () => { + const storePath = join(makeTempDir(tempDirs), "cron-jobs.json"); + const store = new AgentCronJobStore(storePath); + store.create({ + activeSessionId: "active-1", + sessionId: "session-1", + sessionFile: "/tmp/session-1.jsonl", + cwd: "/tmp/project", + scheduleText: "in 1h", + prompt: "aaaa", + now: start, + }); + // Same byte size, different content: only the mtime moves. + writeExternalState( + storePath, + [ + externalJob({ + id: "external-1", + prompt: "bbbb", + sessionId: "session-1", + sessionFile: "/tmp/session-1.jsonl", + }), + ], + new Date("2026-01-02T00:00:00.000Z"), + ); + readCounts.clear(); + + expect(store.list().map((candidate) => candidate.prompt)).toEqual(["bbbb"]); + expect(jobsReads(storePath)).toBe(1); + + expect(store.list().map((candidate) => candidate.prompt)).toEqual(["bbbb"]); + expect(jobsReads(storePath)).toBe(1); + }); + + it("re-reads when the file is replaced by a fresh inode with the same size and mtime", () => { + const storePath = join(makeTempDir(tempDirs), "cron-jobs.json"); + const store = new AgentCronJobStore(storePath); + store.create({ + activeSessionId: "active-1", + sessionId: "session-1", + sessionFile: "/tmp/session-1.jsonl", + cwd: "/tmp/project", + scheduleText: "in 1h", + prompt: "inode guard", + now: start, + }); + const pinned = new Date("2026-01-03T00:00:00.000Z"); + utimesSync(storePath, pinned, pinned); + expect(store.list()).toHaveLength(1); + const bytes = readFileSync(storePath, "utf-8"); + readCounts.clear(); + + rmSync(storePath); + writeFileSync(storePath, bytes); + utimesSync(storePath, pinned, pinned); + + expect(store.list()).toHaveLength(1); + expect(jobsReads(storePath)).toBe(1); + }); + + it("keeps concurrent store instances on one file consistent", () => { + const storePath = join(makeTempDir(tempDirs), "cron-jobs.json"); + const first = new AgentCronJobStore(storePath); + const second = new AgentCronJobStore(storePath); + const jobA = first.create({ + activeSessionId: "active-1", + sessionId: "session-1", + sessionFile: "/tmp/session-1.jsonl", + cwd: "/tmp/project", + scheduleText: "in 1h", + prompt: "from first", + now: start, + }); + expect(second.list().map((candidate) => candidate.id)).toEqual([jobA.id]); + + const jobB = second.create({ + activeSessionId: "active-2", + sessionId: "session-2", + sessionFile: "/tmp/session-2.jsonl", + cwd: "/tmp/project", + scheduleText: "in 2h", + prompt: "from second", + now: start, + }); + expect( + first + .list() + .map((candidate) => candidate.id) + .sort(), + ).toEqual([jobA.id, jobB.id].sort()); + + writeExternalState(storePath, [], new Date("2026-01-04T00:00:00.000Z")); + expect(first.list()).toEqual([]); + expect(second.list()).toEqual([]); + + writeExternalState( + storePath, + [ + externalJob({ + id: "external-1", + prompt: "external writer", + sessionId: "session-1", + sessionFile: "/tmp/session-1.jsonl", + }), + ], + new Date("2026-01-05T00:00:00.000Z"), + ); + expect(first.list().map((candidate) => candidate.id)).toEqual(["external-1"]); + expect(second.list().map((candidate) => candidate.id)).toEqual(["external-1"]); + }); + + it("drops the snapshot when a mutation fails to persist and re-reads the disk", () => { + const storePath = join(makeTempDir(tempDirs), "cron-jobs.json"); + const store = new AgentCronJobStore(storePath); + const first = store.create({ + activeSessionId: "active-1", + sessionId: "session-1", + sessionFile: "/tmp/session-1.jsonl", + cwd: "/tmp/project", + scheduleText: "in 1h", + prompt: "survives", + now: start, + }); + readCounts.clear(); + expect(store.list().map((candidate) => candidate.id)).toEqual([first.id]); + expect(jobsReads(storePath)).toBe(0); + + renameFault.remaining = 1; + renameFault.path = storePath; + expect(() => + store.create({ + activeSessionId: "active-2", + sessionId: "session-2", + sessionFile: "/tmp/session-2.jsonl", + cwd: "/tmp/project", + scheduleText: "in 1h", + prompt: "never lands", + now: start, + }), + ).toThrow(); + + readCounts.clear(); + expect(store.list().map((candidate) => candidate.id)).toEqual([first.id]); + expect(jobsReads(storePath)).toBe(1); + }); + + it("serves an empty state while the file is missing and reloads it when it appears", () => { + const storePath = join(makeTempDir(tempDirs), "cron-jobs.json"); + const store = new AgentCronJobStore(storePath); + expect(store.list()).toEqual([]); + writeExternalState( + storePath, + [ + externalJob({ + id: "external-1", + prompt: "appears later", + sessionId: "session-1", + sessionFile: "/tmp/session-1.jsonl", + }), + ], + new Date("2026-01-06T00:00:00.000Z"), + ); + expect(store.list().map((candidate) => candidate.id)).toEqual(["external-1"]); + readCounts.clear(); + expect(store.list()).toHaveLength(1); + expect(jobsReads(storePath)).toBe(0); + }); + + it("caches per registered session artifact and re-reads only the changed file", () => { + const root = makeTempDir(tempDirs); + const store = AgentCronJobStore.forSessionArtifacts(); + const firstDir = join(root, "artifacts-1"); + const secondDir = join(root, "artifacts-2"); + store.registerSessionArtifact("session-1", firstDir); + store.registerSessionArtifact("session-2", secondDir); + store.createHeartbeat({ + activeSessionId: "active-1", + sessionId: "session-1", + sessionFile: join(root, "session-1.jsonl"), + cwd: root, + scheduleText: "every 5m", + prompt: "first heartbeat", + now: start, + }); + store.createHeartbeat({ + activeSessionId: "active-2", + sessionId: "session-2", + sessionFile: join(root, "session-2.jsonl"), + cwd: root, + scheduleText: "every 5m", + prompt: "second heartbeat", + now: start, + }); + const firstPath = join(firstDir, SESSION_SCHEDULED_JOBS_FILENAME); + const secondPath = join(secondDir, SESSION_SCHEDULED_JOBS_FILENAME); + readCounts.clear(); + + expect(store.list()).toHaveLength(2); + expect(jobsReads(firstPath)).toBe(0); + expect(jobsReads(secondPath)).toBe(0); + + writeExternalState( + firstPath, + [ + externalJob({ + id: "external-1", + prompt: "rewritten externally", + sessionId: "session-1", + sessionFile: join(root, "session-1.jsonl"), + }), + ], + new Date("2026-01-07T00:00:00.000Z"), + ); + expect(store.list().map((candidate) => candidate.prompt)).toEqual( + expect.arrayContaining(["rewritten externally", "second heartbeat"]), + ); + expect(jobsReads(firstPath)).toBe(1); + expect(jobsReads(secondPath)).toBe(0); + }); + + it("serves read-only views: an in-place edit of a listed job cannot change store behavior or reach disk", () => { + const storePath = join(makeTempDir(tempDirs), "cron-jobs.json"); + const store = new AgentCronJobStore(storePath); + const job = store.create({ + activeSessionId: "active-1", + sessionId: "session-1", + sessionFile: "/tmp/session-1.jsonl", + cwd: "/tmp/project", + scheduleText: "in 1h", + prompt: "read-only view", + now: start, + }); + readCounts.clear(); + + const listed = store.list(); + const dueJob = store.getDueJob(job.id, new Date("2026-01-01T13:35:00.000Z")); + expect(listed).toHaveLength(1); + expect(dueJob).toBeDefined(); + expect(Object.isFrozen(listed[0])).toBe(true); + expect(Object.isFrozen(dueJob)).toBe(true); + + expect(() => setStatus(listed[0], "cancelled")).toThrow(TypeError); + expect(() => setStatus(dueJob, "cancelled")).toThrow(TypeError); + + expect(store.list()[0]?.status).toBe("active"); + expect(jobsReads(storePath)).toBe(0); + + // An unrelated mutation must persist the true state, not a leaked in-place edit. + store.create({ + activeSessionId: "active-2", + sessionId: "session-2", + sessionFile: "/tmp/session-2.jsonl", + cwd: "/tmp/project", + scheduleText: "in 1h", + prompt: "unrelated mutation", + now: start, + }); + const persisted = JSON.parse(readFileSync(storePath, "utf-8")) as { jobs: AgentCronJob[] }; + expect(persisted.jobs.find((candidate) => candidate.id === job.id)?.status).toBe("active"); + expect(store.list()).toHaveLength(2); + }); + + it("keeps the published snapshot intact when a mutator throws mid-edit", () => { + const storePath = join(makeTempDir(tempDirs), "cron-jobs.json"); + const store = new AgentCronJobStore(storePath); + const job = store.create({ + activeSessionId: "active-1", + sessionId: "session-1", + sessionFile: "/tmp/session-1.jsonl", + cwd: "/tmp/project", + scheduleText: "in 1h", + prompt: "survives the mutator crash", + now: start, + }); + readCounts.clear(); + + expect(() => + ( + store as unknown as { mutateStates: (mutator: (state: MutableJobsState) => unknown[]) => unknown[] } + ).mutateStates((state) => { + state.jobs = state.jobs.map((candidate) => ({ ...candidate, status: "cancelled" as const })); + throw new Error("mutator crashed mid-edit"); + }), + ).toThrow("mutator crashed mid-edit"); + + // The unpersisted partial edit never reaches reads... + expect(store.list().map((candidate) => candidate.status)).toEqual(["active"]); + expect(jobsReads(storePath)).toBe(0); + // ...nor does a later unrelated mutation resurrect it. + store.create({ + activeSessionId: "active-2", + sessionId: "session-2", + sessionFile: "/tmp/session-2.jsonl", + cwd: "/tmp/project", + scheduleText: "in 1h", + prompt: "after the crash", + now: start, + }); + const persisted = JSON.parse(readFileSync(storePath, "utf-8")) as { jobs: AgentCronJob[] }; + expect(persisted.jobs.find((candidate) => candidate.id === job.id)?.status).toBe("active"); + expect(store.list().map((candidate) => candidate.status)).toEqual(["active", "active"]); + }); + + it("re-reads instead of caching an external replacement that lands right after our own write", () => { + const storePath = join(makeTempDir(tempDirs), "cron-jobs.json"); + const store = new AgentCronJobStore(storePath); + store.create({ + activeSessionId: "active-1", + sessionId: "session-1", + sessionFile: "/tmp/session-1.jsonl", + cwd: "/tmp/project", + scheduleText: "in 1h", + prompt: "our job", + now: start, + }); + readCounts.clear(); + // An external writer replaces the file the instant our atomic rename lands: + // the post-write stat must not pair the external identity with our state. + renameAfterHook.after = () => { + renameAfterHook.after = undefined; + writeExternalState( + storePath, + [ + externalJob({ + id: "external-raced", + prompt: "external wins", + sessionId: "session-1", + sessionFile: "/tmp/session-1.jsonl", + }), + ], + new Date("2026-01-08T00:00:00.000Z"), + ); + }; + store.create({ + activeSessionId: "active-2", + sessionId: "session-2", + sessionFile: "/tmp/session-2.jsonl", + cwd: "/tmp/project", + scheduleText: "in 1h", + prompt: "raced write", + now: start, + }); + + expect(store.list().map((candidate) => candidate.id)).toEqual(["external-raced"]); + expect(jobsReads(storePath)).toBe(1); + expect(store.list().map((candidate) => candidate.id)).toEqual(["external-raced"]); + expect(jobsReads(storePath)).toBe(1); + + // A later mutation must operate on the external state, not our stale write. + const cancelled = store.cancel("external-raced"); + expect(cancelled?.status).toBe("cancelled"); + expect(store.list()[0]?.status).toBe("cancelled"); + }); +}); + +function setStatus(target: unknown, status: string): void { + (target as { status: string }).status = status; +} + +interface MutableJobsState { + jobs: AgentCronJob[]; + dispatches: Array<{ id: string; jobId: string; claimedAt: string; scheduledFor: string }>; +} From 37274a0a70c799d3133c70ec867421478ff8148c Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Sat, 19 Sep 2026 14:47:59 -0400 Subject: [PATCH 20/26] perf(coding-agent): skip roster re-composition for unchanged sessions (#2393) Roster heartbeat frames re-composed every active session summary on every flush: latestMessageActivityAt walked the whole message array per session, each summary was rebuilt and re-serialized, and listHeartbeats built a full summary per registered job on every poll. Daemon CPU therefore grew linearly with messages x sessions x flushes, dominated by long-lived 18k-message sessions that had not changed at all. - summaryForActiveSession now memoizes the last composed summary per session and returns the same object when a cheap fingerprint of every summary input (message count, latest activity, busy bits, usage, verdict, registrations, action snapshot, identity fields) is unchanged. The fingerprint compares the full input set rather than leaf id + message count alone because streaming, bash, compaction, tool, attach, and verdict edges all change the summary without a single append; the memoized summary stays byte-identical to a fresh compose, so the wire delta is unchanged. - The latest-activity timestamp scan is incremental: the memo folds only messages appended since the last scan and falls back to a full walk when the array is replaced, shrinks, or shifts (mid-array insert detected by the boundary tail reference). - flushRoster reuses the previous roster entry and its serialization when the session composed to the same memoized summary, so unchanged sessions skip both the compose and the JSON.stringify per flush; freshly composed entries still stringify and compare exactly as before. - listHeartbeats takes its sessionName/firstMessage from a new sessionDisplayLabels helper instead of composing a full summary per heartbeat job per poll. Measured on an 18k-message idle session: 59us -> 9us per compose cycle when unchanged, 44us -> 9us after an append. (cherry picked from commit e311d6495124cf0bdc629c813fc97a39a9a3054d) --- .../.changes/roster-compose-skip.md | 1 + .../src/modes/daemon/daemon-mode.ts | 37 +- .../src/modes/daemon/daemon-session-list.ts | 325 ++++++++++++++++-- .../test/daemon-agent-roster.test.ts | 65 ++++ .../test/daemon-session-list.test.ts | 212 ++++++++++++ .../test/daemon-supervisor-monitor.test.ts | 1 + 6 files changed, 610 insertions(+), 31 deletions(-) create mode 100644 packages/coding-agent/.changes/roster-compose-skip.md diff --git a/packages/coding-agent/.changes/roster-compose-skip.md b/packages/coding-agent/.changes/roster-compose-skip.md new file mode 100644 index 00000000000..c84accbb0f2 --- /dev/null +++ b/packages/coding-agent/.changes/roster-compose-skip.md @@ -0,0 +1 @@ +- Fixed worker CPU growing with session size: roster heartbeat flushes no longer re-walk every message for the latest-activity timestamp, no longer re-compose and re-serialize unchanged session summaries each cycle, and the heartbeats list no longer builds a full session summary per registered job per poll. Unchanged sessions reuse their last composed roster entry; any live-state change (new message, streaming, bash, compaction, verdict, or registration flag) still recomposes and publishes the same wire delta as before. diff --git a/packages/coding-agent/src/modes/daemon/daemon-mode.ts b/packages/coding-agent/src/modes/daemon/daemon-mode.ts index 680b850d01a..1f54c065ad2 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-mode.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-mode.ts @@ -196,6 +196,7 @@ import { inactiveLifecycleForSession, type SessionSummary, scheduledJobRegistrations, + sessionDisplayLabels, summaryForActiveSession, } from "./daemon-session-list.js"; import { DaemonSessionSummarizer } from "./daemon-session-summarizer.js"; @@ -642,6 +643,7 @@ export class AgentDaemon { private readonly recoveryJournal?: WorkerRecoveryJournal; private readonly rosterReporter: WorkerRosterReporterState = { lastComposed: new Map(), + lastComposedSource: new Map(), lastComposedJson: new Map(), queuedChildren: new Map(), removedAgentIds: new Map(), @@ -2268,11 +2270,13 @@ export class AgentDaemon { .filter((job) => isHeartbeatCronJob(job) && (job.status === "active" || job.status === "paused")) .map((job) => { const state = this.sessions.get(job.activeSessionId); - const summary = state ? summaryForActiveSession(state) : undefined; + // Only the session's display labels are needed; a full summary compose + // would re-walk the transcript once per heartbeat job on every poll. + const labels = state ? sessionDisplayLabels(state) : undefined; return { job, - ...(summary?.sessionName ? { sessionName: summary.sessionName } : {}), - ...(summary?.firstMessage ? { firstMessage: summary.firstMessage } : {}), + ...(labels?.sessionName ? { sessionName: labels.sessionName } : {}), + ...(labels?.firstMessage ? { firstMessage: labels.firstMessage } : {}), }; }); } @@ -7736,10 +7740,23 @@ export class AgentDaemon { private flushRoster(): void { const reporter = this.rosterReporter; const entries = new Map(); + // Summary objects are memoized per session: an unchanged session composes + // to the same summary reference, so its roster entry and serialization can + // be reused instead of re-composed and re-stringified on every flush. + const composedSources = new Map(); const scheduledJobs = this.cronStore.list(); for (const summary of buildSessionList([...this.sessions.values()], [], scheduledJobs)) { - const entry = workerRosterEntryFromSummary(summary); - entries.set(entry.agentId, entry); + const agentId = rosterAgentIdForSummary(summary); + if (reporter.lastComposedSource.get(agentId) === summary) { + const entry = reporter.lastComposed.get(agentId); + if (entry) { + entries.set(agentId, entry); + composedSources.set(agentId, summary); + continue; + } + } + entries.set(agentId, workerRosterEntryFromSummary(summary)); + composedSources.set(agentId, summary); } for (const [agentId, queued] of reporter.queuedChildren) { if (entries.has(agentId)) { @@ -7792,12 +7809,18 @@ export class AgentDaemon { const changed: WorkerRosterEntry[] = []; const nextJson = new Map(); for (const entry of entries.values()) { - const json = JSON.stringify(entry); + // An entry reused from the last flush is unchanged by construction: + // skip its serialization and keep the previous json for the delta compare. + const json = + reporter.lastComposed.get(entry.agentId) === entry + ? (reporter.lastComposedJson.get(entry.agentId) ?? JSON.stringify(entry)) + : JSON.stringify(entry); nextJson.set(entry.agentId, json); if (reporter.lastComposedJson.get(entry.agentId) !== json) changed.push(entry); } const removedAgentIds = [...reporter.removedAgentIds.keys()]; reporter.lastComposed = new Map(entries); + reporter.lastComposedSource = composedSources; reporter.lastComposedJson = nextJson; if (!this.hasAuthenticatedSupervisorClient()) { if (changed.length > 0 || removedAgentIds.length > 0) reporter.snapshotPending = true; @@ -8264,6 +8287,8 @@ export class AgentDaemon { interface WorkerRosterReporterState { lastComposed: Map; + /** Summary each lastComposed entry was composed from; an unchanged session reuses its entry. */ + lastComposedSource: Map; lastComposedJson: Map; queuedChildren: Map; /** Pending removals: agentId -> removed sessionId; a new incarnation of the id cancels it. */ 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 21e0794e324..73a5c7c3980 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-session-list.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-session-list.ts @@ -6,7 +6,7 @@ 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"; -import type { AgentTaskState, SessionInfo } from "../../core/session-manager.js"; +import type { AgentStatus, AgentTaskState, SessionInfo } from "../../core/session-manager.js"; import type { SessionUsageSummary } from "../../core/usage.js"; import type { AgentConnectionRlmChildAgentSnapshot } from "../agent-connection/types.js"; import type { ActiveSessionState } from "./active-session-state.js"; @@ -216,6 +216,77 @@ export function buildSessionList( return entries; } +// Compose memos, keyed weakly by session state so entries die with the session. +// Roster flushes re-compose every active session summary on every event; with +// large sessions that meant a full message walk plus a full compose plus a +// serialize per session per flush. The activity memo keeps the timestamp scan +// incremental, and the fingerprint memo returns the last composed summary +// (same object) whenever every input that feeds it is unchanged. +const messageActivityMemos = new WeakMap(); +const summaryComposeMemos = new WeakMap< + ActiveSessionState, + { fingerprint: SummaryComposeFingerprint; summary: SessionSummary } +>(); + +/** + * Cheap snapshot of every summary input; any difference forces a fresh compose. + * + * Leaf id and message count alone are not a sufficient key: isStreaming, + * isBashRunning, isCompacting, attachment counts, the summary verdict, and the + * heartbeat/cron flags all change without a single append, and the daemon + * schedules roster flushes at exactly those edges. Comparing the full input set + * keeps the memoized summary byte-identical to a fresh compose. + */ +interface SummaryComposeFingerprint { + // Caller inputs: registration flags vary per compose site, and the saved + // catalog entry (when present) feeds created/modified/firstMessage/parentSessionPath. + hasActiveHeartbeat: boolean; + hasRegisteredHeartbeat: boolean; + hasRegisteredCronJob: boolean; + savedSession: SessionInfo | undefined; + // Busy-state bits flip without appends (turn, bash, compaction, tool edges). + isStreaming: boolean; + isCompacting: boolean; + isBashRunning: boolean; + pendingToolCallsSize: number; + isSessionActive: boolean; + hasRunningRlmChildren: boolean; + unfinishedActionCount: number; + attachedClients: number; + directAttachedClients: number; + messageCount: number; + // References: memoized upstream (usage), replaced on change (model, + // summaryState, streamingMessage), or structural (sessionActions, + // diagnostics: post-construction changes replace the array or append). + usage: SessionUsageSummary | undefined; + model: Model | undefined; + thinkingLevel: ThinkingLevel | undefined; + streamingMessage: AgentMessage | undefined; + summaryState: AgentStatus | undefined; + diagnostics: readonly AgentSessionRuntimeDiagnostic[]; + repliedSinceTask: boolean | undefined; + sessionActions: SessionActionSnapshot; + // Identity and display inputs: the metadata getter returns a fresh copy on + // every read, so its summary-relevant fields compare by value. + metadataKind: string; + metadataParentActiveSessionId: string | undefined; + metadataParentSessionId: string | undefined; + metadataParentSessionFile: string | undefined; + metadataRlmChildId: string | undefined; + metadataRlmParentNodeId: string | undefined; + sessionName: string | undefined; + sessionId: string; + sessionFile: string | undefined; + cwd: string; + rlmDepth: number | undefined; + modelFallbackMessage: string | undefined; + // Values computed once per attempt and shared with the compose. + headerTimestamp: string | undefined; + modified: string | undefined; + lastActivityAt: string | undefined; + firstMessage: string | undefined; +} + export function summaryForActiveSession( activeSession: ActiveSessionState, savedSession?: SessionInfo, @@ -237,7 +308,71 @@ export function summaryForActiveSession( const directAttachedClients = [...activeSession.clients].filter( (client) => client.authenticationRole === "session_client", ).length; - return { + // The activity memo folds only messages appended since the last scan. + let activityMemo = messageActivityMemos.get(activeSession); + if (activityMemo === undefined) { + activityMemo = { source: undefined, scannedLength: 0, tailRef: undefined, latest: undefined }; + messageActivityMemos.set(activeSession, activityMemo); + } + const headerTimestamp = session.sessionManager.getHeader?.()?.timestamp; + const lastActivityAt = latestMessageActivityAt(session.messages, activityMemo) ?? modified ?? headerTimestamp; + // Subagent sessions live in artifact dirs that the saved-session scan + // never sees; their spawn prompt is the most identifying title we have. + // A freshly created top-level session has neither yet — its jsonl is not + // scanned until it flushes — so derive from the live first user message to + // avoid titling the chat with its session ID until the file lands. + const firstMessage = + savedSession?.firstMessage ?? + (metadata.prompt ? compactRlmText(metadata.prompt, 120) : undefined) ?? + firstUserMessageText(session); + + const fingerprint: SummaryComposeFingerprint = { + hasActiveHeartbeat, + hasRegisteredHeartbeat, + hasRegisteredCronJob, + savedSession, + isStreaming: session.isStreaming, + isCompacting: session.isCompacting, + isBashRunning: session.isBashRunning, + pendingToolCallsSize: session.state.pendingToolCalls.size, + isSessionActive: session.isSessionActive, + hasRunningRlmChildren: session.hasRunningRlmChildren(), + unfinishedActionCount: session.unfinishedActionCount, + attachedClients: activeSession.clients.size, + directAttachedClients, + messageCount: session.messages.length, + usage: session.getOwnUsageSummary?.(), + model: session.model as Model | undefined, + thinkingLevel: session.thinkingLevel, + streamingMessage: session.state.streamingMessage, + summaryState: activeSession.summaryState, + diagnostics: activeSession.runtime.diagnostics, + repliedSinceTask: metadata.kind === "subagent" ? session.repliedToParentSinceTask : undefined, + sessionActions: session.getSessionActionSnapshot(), + metadataKind: metadata.kind, + metadataParentActiveSessionId: metadata.parentActiveSessionId, + metadataParentSessionId: metadata.parentSessionId, + metadataParentSessionFile: metadata.parentSessionFile, + metadataRlmChildId: metadata.rlmChildId, + metadataRlmParentNodeId: metadata.rlmParentNodeId, + sessionName: session.sessionName, + sessionId: session.sessionId, + sessionFile: session.sessionFile, + cwd: session.sessionManager.getCwd(), + rlmDepth: session.rlmDepth, + modelFallbackMessage: activeSession.runtime.modelFallbackMessage, + headerTimestamp, + modified, + lastActivityAt, + firstMessage, + }; + + const memo = summaryComposeMemos.get(activeSession); + if (memo && summaryComposeFingerprintsEqual(memo.fingerprint, fingerprint)) { + return memo.summary; + } + + const summary: SessionSummary = { id: activeSession.activeSessionId, lifecycle: activeLifecycleForSession(activeSession), activity: activeActivityForSession(activeSession), @@ -245,8 +380,7 @@ export function summaryForActiveSession( hasActiveHeartbeat: hasActiveHeartbeat || undefined, hasRegisteredHeartbeat: hasRegisteredHeartbeat || undefined, hasRegisteredCronJob: hasRegisteredCronJob || undefined, - lastActivityAt: - latestMessageActivityAt(session.messages) ?? modified ?? session.sessionManager.getHeader?.()?.timestamp, + lastActivityAt, runtimeKind: metadata.kind, rlmDepth: session.rlmDepth, activeSessionId: activeSession.activeSessionId, @@ -268,17 +402,9 @@ export function summaryForActiveSession( unfinishedActionCount: session.unfinishedActionCount, sessionActions: session.getSessionActionSnapshot(), streamingMessage: session.state.streamingMessage, - created: savedSession?.created.toISOString() ?? session.sessionManager.getHeader?.()?.timestamp, + created: savedSession?.created.toISOString() ?? headerTimestamp, modified, - // Subagent sessions live in artifact dirs that the saved-session scan - // never sees; their spawn prompt is the most identifying title we have. - // A freshly created top-level session has neither yet — its jsonl is not - // scanned until it flushes — so derive from the live first user message to - // avoid titling the chat with its session ID until the file lands. - firstMessage: - savedSession?.firstMessage ?? - (metadata.prompt ? compactRlmText(metadata.prompt, 120) : undefined) ?? - firstUserMessageText(session), + firstMessage, parentActiveSessionId: metadata.parentActiveSessionId, parentSessionId: metadata.parentSessionId, parentSessionPath: savedSession?.parentSessionPath ?? metadata.parentSessionFile, @@ -298,20 +424,169 @@ export function summaryForActiveSession( summary: activeSession.summaryState?.summary, ...(isSummaryCurrent(activeSession) ? { taskState: activeSession.summaryState?.taskState } : {}), }; + summaryComposeMemos.set(activeSession, { fingerprint, summary }); + return summary; +} + +function summaryComposeFingerprintsEqual(left: SummaryComposeFingerprint, right: SummaryComposeFingerprint): boolean { + return ( + left.hasActiveHeartbeat === right.hasActiveHeartbeat && + left.hasRegisteredHeartbeat === right.hasRegisteredHeartbeat && + left.hasRegisteredCronJob === right.hasRegisteredCronJob && + left.savedSession === right.savedSession && + left.isStreaming === right.isStreaming && + left.isCompacting === right.isCompacting && + left.isBashRunning === right.isBashRunning && + left.pendingToolCallsSize === right.pendingToolCallsSize && + left.isSessionActive === right.isSessionActive && + left.hasRunningRlmChildren === right.hasRunningRlmChildren && + left.unfinishedActionCount === right.unfinishedActionCount && + left.attachedClients === right.attachedClients && + left.directAttachedClients === right.directAttachedClients && + left.messageCount === right.messageCount && + left.usage === right.usage && + left.model === right.model && + left.thinkingLevel === right.thinkingLevel && + left.streamingMessage === right.streamingMessage && + left.summaryState === right.summaryState && + left.repliedSinceTask === right.repliedSinceTask && + left.metadataKind === right.metadataKind && + left.metadataParentActiveSessionId === right.metadataParentActiveSessionId && + left.metadataParentSessionId === right.metadataParentSessionId && + left.metadataParentSessionFile === right.metadataParentSessionFile && + left.metadataRlmChildId === right.metadataRlmChildId && + left.metadataRlmParentNodeId === right.metadataRlmParentNodeId && + left.sessionName === right.sessionName && + left.sessionId === right.sessionId && + left.sessionFile === right.sessionFile && + left.cwd === right.cwd && + left.rlmDepth === right.rlmDepth && + left.modelFallbackMessage === right.modelFallbackMessage && + left.headerTimestamp === right.headerTimestamp && + left.modified === right.modified && + left.lastActivityAt === right.lastActivityAt && + left.firstMessage === right.firstMessage && + diagnosticsEqual(left.diagnostics, right.diagnostics) && + sessionActionSnapshotsEqual(left.sessionActions, right.sessionActions) + ); +} + +// Runtime diagnostics change by wholesale replacement or by append; a stable +// length and tail element covers both without a deep compare. +function diagnosticsEqual( + left: readonly AgentSessionRuntimeDiagnostic[], + right: readonly AgentSessionRuntimeDiagnostic[], +): boolean { + if (left === right) return true; + if (left.length !== right.length) return false; + return left.length === 0 || left[left.length - 1] === right[right.length - 1]; +} + +function sessionActionSnapshotsEqual(left: SessionActionSnapshot, right: SessionActionSnapshot): boolean { + if (left === right) return true; + if (left.queuedCount !== right.queuedCount) return false; + if ( + left.active?.kind !== right.active?.kind || + left.active?.phase !== right.active?.phase || + left.active?.label !== right.active?.label + ) { + return false; + } + return stringArraysEqual(left.steering, right.steering) && stringArraysEqual(left.followUps, right.followUps); +} + +function stringArraysEqual(left: readonly string[], right: readonly string[]): boolean { + if (left === right) return true; + if (left.length !== right.length) return false; + for (let index = 0; index < left.length; index += 1) { + if (left[index] !== right[index]) return false; + } + return true; +} + +/** + * Cheap display labels for callers that need only the session's identity (the + * heartbeats list polls once per registered job); avoids the full summary + * compose. Mirrors the sessionName/firstMessage fields summaryForActiveSession + * produces for a session without a saved catalog entry. + */ +export function sessionDisplayLabels(activeSession: ActiveSessionState): { + sessionName: string | undefined; + firstMessage: string | undefined; +} { + const session = activeSession.runtime.session; + const metadata = activeSession.runtime.metadata ?? { kind: "top-level" as const }; + return { + sessionName: session.sessionName, + firstMessage: + (metadata.prompt ? compactRlmText(metadata.prompt, 120) : undefined) ?? firstUserMessageText(session), + }; +} + +/** + * Incremental max-message-timestamp tracker for one active session. + * + * Sessions only append to their message array in place (full reassignments + * replace the array object), so the memo folds just the messages appended + * since the last scan. It falls back to a full walk whenever the array is + * replaced, shrinks, or shifts: a mid-array insert changes the element at the + * previous scan boundary, which the tail reference detects. + */ +export interface MessageActivityMemo { + /** Live messages array the memo last scanned; an identity change forces a full walk. */ + source: readonly AgentMessage[] | undefined; + /** Number of leading elements already folded into `latest`. */ + scannedLength: number; + /** Element at `scannedLength - 1` at the last scan; detects mid-array inserts. */ + tailRef: AgentMessage | undefined; + /** Largest valid message timestamp seen so far, or undefined before the first valid one. */ + latest: number | undefined; } -function latestMessageActivityAt(messages: readonly AgentMessage[]): string | undefined { +function messageActivityTimestamp(message: AgentMessage): number | undefined { + // Tool results and custom messages are real session activity too. Looking at + // every timestamp also keeps this correct for future AgentMessage variants. + if ( + typeof message.timestamp === "number" && + Number.isFinite(message.timestamp) && + Math.abs(message.timestamp) <= MAX_DATE_TIMESTAMP_MS + ) { + return message.timestamp; + } + return undefined; +} + +function foldMessageActivity( + messages: readonly AgentMessage[], + from: number, + latest: number | undefined, +): number | undefined { + for (let index = from; index < messages.length; index += 1) { + const timestamp = messageActivityTimestamp(messages[index]!); + if (timestamp === undefined) continue; + latest = latest === undefined ? timestamp : Math.max(latest, timestamp); + } + return latest; +} + +export function latestMessageActivityAt( + messages: readonly AgentMessage[], + memo?: MessageActivityMemo, +): string | undefined { let latest: number | undefined; - for (const message of messages) { - // Tool results and custom messages are real session activity too. Looking at - // every timestamp also keeps this correct for future AgentMessage variants. - if ( - typeof message.timestamp === "number" && - Number.isFinite(message.timestamp) && - Math.abs(message.timestamp) <= MAX_DATE_TIMESTAMP_MS - ) { - latest = latest === undefined ? message.timestamp : Math.max(latest, message.timestamp); - } + if (memo === undefined || memo.source !== messages || messages.length < memo.scannedLength) { + latest = foldMessageActivity(messages, 0, undefined); + } else if (memo.scannedLength > 0 && messages[memo.scannedLength - 1] !== memo.tailRef) { + // A mid-array insert shifted the scanned prefix; rescan everything. + latest = foldMessageActivity(messages, 0, undefined); + } else { + latest = foldMessageActivity(messages, memo.scannedLength, memo.latest); + } + if (memo !== undefined) { + memo.source = messages; + memo.scannedLength = messages.length; + memo.tailRef = messages.length > 0 ? messages[messages.length - 1] : undefined; + memo.latest = latest; } return latest === undefined ? undefined : new Date(latest).toISOString(); } diff --git a/packages/coding-agent/test/daemon-agent-roster.test.ts b/packages/coding-agent/test/daemon-agent-roster.test.ts index e9a13c48c90..4ca45de549d 100644 --- a/packages/coding-agent/test/daemon-agent-roster.test.ts +++ b/packages/coding-agent/test/daemon-agent-roster.test.ts @@ -34,6 +34,7 @@ interface WorkerReporterFixture { flushRoster(): void; rosterReporter: { lastComposed: Map; + lastComposedSource: Map; lastComposedJson: Map; queuedChildren: Map; removedAgentIds: Map; @@ -56,6 +57,7 @@ function makeWorkerReporter(connected = true): WorkerReporterFixture { snapshotPayloadGenerations: new Map(), rosterReporter: { lastComposed: new Map(), + lastComposedSource: new Map(), lastComposedJson: new Map(), queuedChildren: new Map(), removedAgentIds: new Map(), @@ -484,6 +486,69 @@ describe("worker roster reporter", () => { await new Promise((resolveSettle) => setImmediate(resolveSettle)); expect(internals.rosterReporter.lastComposed.get(agentId)?.summary.model).toMatchObject({ id: "m2" }); }); + + it("reuses the composed entry for an unchanged session and recomposes on append", () => { + const { daemon, sentDeltas } = makeWorkerReporter(); + const messages: AgentMessage[] = []; + const state = makeState({ activeSessionId: "steady-active", sessionId: "steady-session", messages }); + daemon.sessions.set(state.activeSessionId, state); + + daemon.flushRoster(); + expect(sentDeltas).toHaveLength(1); + const firstEntry = daemon.rosterReporter.lastComposed.get("steady-session"); + expect(firstEntry?.summary.messageCount).toBe(0); + + // Nothing changed: no new delta, and the same entry object is reused + // instead of being recomposed and re-serialized from the same inputs. + daemon.flushRoster(); + expect(sentDeltas).toHaveLength(1); + expect(daemon.rosterReporter.lastComposed.get("steady-session")).toBe(firstEntry); + + // A new message recomposes the row and publishes the updated activity. + messages.push({ + role: "user", + content: "next turn", + timestamp: Date.parse("2026-05-02T00:00:00.000Z"), + } as AgentMessage); + daemon.flushRoster(); + expect(sentDeltas).toHaveLength(2); + const secondEntry = daemon.rosterReporter.lastComposed.get("steady-session"); + expect(secondEntry).not.toBe(firstEntry); + expect(secondEntry?.summary).toMatchObject({ + messageCount: 1, + lastActivityAt: "2026-05-02T00:00:00.000Z", + }); + + // Unchanged again: the recomposed entry is reused across later flushes. + daemon.flushRoster(); + expect(sentDeltas).toHaveLength(2); + expect(daemon.rosterReporter.lastComposed.get("steady-session")).toBe(secondEntry); + }); + + it("republishes busy-state flips that arrive without any append", () => { + const { daemon, sentDeltas } = makeWorkerReporter(); + const state = makeState({ activeSessionId: "busy-active", sessionId: "busy-session", messages: [] }); + daemon.sessions.set(state.activeSessionId, state); + + daemon.flushRoster(); + expect(sentDeltas.at(-1)?.entries.some((entry) => entry.summary.isStreaming === false)).toBe(true); + + // turn_start/bash_start/compaction_start schedule flushes with no message + // appended: the entry must recompose, not reuse the idle snapshot. + const session = state.runtime.session as unknown as { isStreaming: boolean; isSessionActive: boolean }; + session.isStreaming = true; + session.isSessionActive = true; + daemon.flushRoster(); + const delta = sentDeltas.at(-1); + expect(delta?.entries).toHaveLength(1); + expect(delta?.entries[0]?.summary).toMatchObject({ isStreaming: true, activity: "working" }); + + // Back to idle: the row recomposes again even though no append happened. + session.isStreaming = false; + session.isSessionActive = false; + daemon.flushRoster(); + expect(sentDeltas.at(-1)?.entries[0]?.summary).toMatchObject({ isStreaming: false, activity: "idle" }); + }); }); // --- Supervisor-side roster ledger --- diff --git a/packages/coding-agent/test/daemon-session-list.test.ts b/packages/coding-agent/test/daemon-session-list.test.ts index 61a9a2cd604..664b53bf68e 100644 --- a/packages/coding-agent/test/daemon-session-list.test.ts +++ b/packages/coding-agent/test/daemon-session-list.test.ts @@ -3,6 +3,7 @@ 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 { SessionActionSnapshot } from "../src/core/session-action-store.js"; import type { SessionInfo } from "../src/core/session-manager.js"; import type { SessionUsageSummary } from "../src/core/usage.js"; import type { ActiveSessionState, DaemonSocketClient } from "../src/modes/daemon/active-session-state.js"; @@ -10,8 +11,11 @@ import { passivatedWorkerRosterEntry, workerRosterEntryFromSummary } from "../sr import { buildRlmChildSnapshots, buildSessionList, + latestMessageActivityAt, + type MessageActivityMemo, resolveAttachModelFallbackMessage, type SessionSummary, + sessionDisplayLabels, summaryForActiveSession, } from "../src/modes/daemon/daemon-session-list.js"; @@ -546,6 +550,214 @@ describe("summaryForActiveSession recap currency", () => { }); }); +describe("summary compose memoization", () => { + const messageAt = (iso: string, content = "hello") => + ({ role: "user", content, timestamp: Date.parse(iso) }) as AgentMessage; + + it("returns the same summary object when nothing changed", () => { + const state = makeState({ activeSessionId: "steady", messages: [messageAt("2026-05-01T00:00:00.000Z")] }); + const first = summaryForActiveSession(state); + expect(first).toMatchObject({ messageCount: 1, lastActivityAt: "2026-05-01T00:00:00.000Z" }); + // Roster flushes recompose every session each cycle; an unchanged session + // must reuse the previous summary instead of recomposing it. + expect(summaryForActiveSession(state)).toBe(first); + }); + + it("recomposes when a message is appended and folds only the new tail", () => { + const messages = [messageAt("2026-05-01T00:00:00.000Z")]; + const state = makeState({ activeSessionId: "append", messages }); + const first = summaryForActiveSession(state); + messages.push(messageAt("2026-05-02T00:00:00.000Z")); + const second = summaryForActiveSession(state); + expect(second).not.toBe(first); + expect(second).toMatchObject({ messageCount: 2, lastActivityAt: "2026-05-02T00:00:00.000Z" }); + // A stale appended timestamp must not unseat the existing latest activity. + messages.push(messageAt("2026-05-01T12:00:00.000Z")); + const third = summaryForActiveSession(state); + expect(third).not.toBe(second); + expect(third.lastActivityAt).toBe("2026-05-02T00:00:00.000Z"); + }); + + it("recomposes when busy-state bits flip without an append", () => { + const state = makeState({ activeSessionId: "busy", messages: [messageAt("2026-05-01T00:00:00.000Z")] }); + const first = summaryForActiveSession(state); + expect(first.isStreaming).toBe(false); + (state.runtime.session as unknown as { isStreaming: boolean }).isStreaming = true; + (state.runtime.session as unknown as { isSessionActive: boolean }).isSessionActive = true; + const second = summaryForActiveSession(state); + expect(second).not.toBe(first); + expect(second).toMatchObject({ isStreaming: true, activity: "working" }); + }); + + it("recomposes when the summarizer verdict is replaced", () => { + const state = makeState({ + activeSessionId: "verdict", + messages: [messageAt("2026-05-01T00:00:00.000Z")], + }); + const first = summaryForActiveSession(state); + expect(first.summary).toBeUndefined(); + (state as unknown as { summaryState?: object }).summaryState = { + summary: "Editing the router", + taskState: "completed", + basedOnMessageCount: 1, + }; + const second = summaryForActiveSession(state); + expect(second).not.toBe(first); + expect(second).toMatchObject({ summary: "Editing the router", taskState: "completed" }); + }); + + it("recomposes when the visible action snapshot changes", () => { + const state = makeState({ activeSessionId: "queued" }); + const first = summaryForActiveSession(state); + expect(first.sessionActions).toMatchObject({ queuedCount: 0 }); + ( + state.runtime.session as unknown as { + getSessionActionSnapshot: () => SessionActionSnapshot; + } + ).getSessionActionSnapshot = () => ({ + queuedCount: 1, + steering: ["revised plan"], + followUps: [], + }); + const second = summaryForActiveSession(state); + expect(second).not.toBe(first); + expect(second.sessionActions).toMatchObject({ queuedCount: 1, steering: ["revised plan"] }); + }); + + it("recomposes when heartbeat registration flags differ per call site", () => { + const state = makeState({ activeSessionId: "flags" }); + const unflagged = summaryForActiveSession(state); + expect(unflagged.hasActiveHeartbeat).toBeUndefined(); + const flagged = summaryForActiveSession(state, undefined, true); + expect(flagged).not.toBe(unflagged); + expect(flagged.hasActiveHeartbeat).toBe(true); + expect(summaryForActiveSession(state, undefined, true)).toBe(flagged); + }); +}); + +describe("latestMessageActivityAt memo", () => { + const messageAt = (iso: string, content = "hello") => + ({ role: "user", content, timestamp: Date.parse(iso) }) as AgentMessage; + + function countingArray(messages: AgentMessage[]): { reads: number[]; array: AgentMessage[] } { + const reads: number[] = []; + const array = new Proxy(messages, { + get(target, property, receiver) { + if (typeof property === "string" && /^\d+$/.test(property)) reads.push(Number(property)); + return Reflect.get(target, property, receiver); + }, + }); + return { reads, array }; + } + + function freshMemo(): MessageActivityMemo { + return { source: undefined, scannedLength: 0, tailRef: undefined, latest: undefined }; + } + + it("scans the whole array once and then only appended messages", () => { + // Twenty messages stand in for a large transcript: the memo must never + // re-walk the old prefix once it has scanned it. + const messages = Array.from({ length: 20 }, (_, index) => + messageAt(new Date(Date.UTC(2026, 4, 1, 0, 0, index)).toISOString()), + ); + const { reads, array } = countingArray(messages); + const memo = freshMemo(); + + const latestIso = new Date(Date.UTC(2026, 4, 1, 0, 0, 19)).toISOString(); + expect(latestMessageActivityAt(array, memo)).toBe(latestIso); + expect(reads).toContain(0); + expect(memo.scannedLength).toBe(20); + + // Unchanged array: O(1) boundary/tail reads only, no message walk. + reads.length = 0; + expect(latestMessageActivityAt(array, memo)).toBe(latestIso); + expect(reads.length).toBeGreaterThan(0); + expect(reads.every((index) => index >= 19)).toBe(true); + + // One append: the boundary check plus the new element, nothing older. + messages.push(messageAt("2026-06-01T00:00:00.000Z")); + reads.length = 0; + expect(latestMessageActivityAt(array, memo)).toBe("2026-06-01T00:00:00.000Z"); + expect(reads.every((index) => index >= 19)).toBe(true); + expect(reads).toContain(20); + expect(memo.scannedLength).toBe(21); + }); + + it("rescans everything after a mid-array insert shifts the boundary", () => { + const messages = [messageAt("2026-05-01T00:00:00.000Z"), messageAt("2026-05-04T00:00:00.000Z")]; + const memo = freshMemo(); + expect(latestMessageActivityAt(messages, memo)).toBe("2026-05-04T00:00:00.000Z"); + // Insert before the tail: the pre-error splice path in agent-session. + messages.splice(1, 0, messageAt("2026-05-03T00:00:00.000Z")); + expect(latestMessageActivityAt(messages, memo)).toBe("2026-05-04T00:00:00.000Z"); + expect(memo.scannedLength).toBe(3); + expect(memo.latest).toBe(Date.parse("2026-05-04T00:00:00.000Z")); + }); + + it("rescans when the array is replaced or shrinks", () => { + const messages = [messageAt("2026-05-01T00:00:00.000Z"), messageAt("2026-05-02T00:00:00.000Z")]; + const memo = freshMemo(); + expect(latestMessageActivityAt(messages, memo)).toBe("2026-05-02T00:00:00.000Z"); + // Compaction and filtering reassign the array; the memo must not trust the old scan. + const replacement = [...messages]; + expect(latestMessageActivityAt(replacement, memo)).toBe("2026-05-02T00:00:00.000Z"); + replacement.pop(); + expect(latestMessageActivityAt(replacement, memo)).toBe("2026-05-01T00:00:00.000Z"); + expect(memo.scannedLength).toBe(1); + }); + + it("keeps ignoring invalid timestamps on appended messages", () => { + const messages = [messageAt("2026-05-01T00:00:00.000Z")]; + const memo = freshMemo(); + expect(latestMessageActivityAt(messages, memo)).toBe("2026-05-01T00:00:00.000Z"); + messages.push({ + role: "assistant", + content: "corrupt", + timestamp: 8.64e15 + 1, + } as unknown as AgentMessage); + expect(latestMessageActivityAt(messages, memo)).toBe("2026-05-01T00:00:00.000Z"); + expect(memo.scannedLength).toBe(2); + }); + + it("matches a full rescan on every mutation pattern", () => { + const build = (count: number) => + Array.from({ length: count }, (_, index) => messageAt(new Date(2026, 4, 1, 0, 0, index).toISOString())); + const messages = build(25); + const memo = freshMemo(); + const fullScan = () => latestMessageActivityAt(messages); + for (const mutation of [ + () => messages.push(messageAt("2026-06-01T00:00:00.000Z")), + () => messages.splice(10, 0, messageAt("2026-05-15T00:00:00.000Z")), + () => messages.pop(), + () => messages.push(messageAt("2026-04-01T00:00:00.000Z")), + ]) { + mutation(); + expect(latestMessageActivityAt(messages, memo)).toBe(fullScan()); + } + }); +}); + +describe("sessionDisplayLabels", () => { + it("mirrors the summary's session name and first message", () => { + const state = makeState({ + activeSessionId: "labels", + messages: [{ role: "user", content: "first prompt" } as AgentMessage], + }); + const labels = sessionDisplayLabels(state); + const summary = summaryForActiveSession(state); + expect(labels.sessionName).toBe(summary.sessionName); + expect(labels.firstMessage).toBe(summary.firstMessage); + }); + + it("uses the spawn prompt as the subagent title", () => { + const state = makeState({ + activeSessionId: "spawned", + metadata: { kind: "subagent", createdAt: 1, prompt: "Audit the retry\nlogic", rlmChildId: "c1" }, + }); + expect(sessionDisplayLabels(state).firstMessage).toBe("Audit the retry logic"); + }); +}); + describe("buildRlmChildSnapshots", () => { it("uses the AgentSession projection and adds resident active session ids", () => { const queued = { diff --git a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts index 4e1191b78f8..7a14ccbde37 100644 --- a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts @@ -510,6 +510,7 @@ describe("daemon worker supervisor monitoring", () => { cronStore: { list: () => [] }, rosterReporter: { lastComposed: new Map(), + lastComposedSource: new Map(), lastComposedJson: new Map(), queuedChildren: new Map(), removedAgentIds: new Map(), From f355141e28f67b35903c7c7fb29fd8047be296cb Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:43:51 -0400 Subject: [PATCH 21/26] [RSI, performance] perf(kernel): defer the event-loop import stack past the ready event (#2379) * perf(kernel): defer the event-loop import stack past the ready event The Python kernel imported asyncio (plus the bash tool's secrets, shutil, datetime, selectors, struct, fcntl/termios, atexit, and tempfile) before sending the ready event. python -X importtime shows the asyncio subtree is the heaviest part of the boot chain (ssl, concurrent.futures, and logging ride along). Defer them: rlm.bash binds asyncio on first BashHandle construction (its only entry point), repl imports asyncio in main() after the ready event and in the handlers that reference it, and the remaining stdlib modules import at their first real use. No protocol, API, or behavior changes; the kernel serves requests and interrupts exactly as before. Local A/B (macOS, 16 interleaved trials, same harness phases as scripts/benchmarks): - kernel_start: 31.4 -> 24.4 ms (-22.3%) - kernel_rss: 28.6 -> 26.8 MB (-6.2%) - kernel_exec, bash, git_status, output, mixed, interrupt, snapshot, loaded_rss: unchanged within noise - restore: 90.1 -> 97.3 ms (+8%): the one-time loop import now lands in the first request of a fresh kernel instead of its boot; after one warm-up cell restore returns to parity (90.0 vs 89.4 ms baseline) Regression tests: import rlm must keep asyncio/secrets off the boot path; a serving kernel loads asyncio by the first cell; BashHandle construction binds asyncio with no event loop or prior bash() call. * fix(kernel): keep default SIGINT handler through deferred event-loop boot _sigint_handler has no task to target before serving starts, so the PR-head ordering (install before the ready send) silently swallowed a Ctrl-C during the deferred asyncio import and event-loop startup. Keep the default handler until the loop and _serve_task exist, then install the custom handler. Adds a regression test that parks the kernel inside the post-ready deferred import (fake asyncio on PYTHONPATH) and asserts a SIGINT there terminates the kernel; it fails on the pre-fix ordering and passes with the fix. * test(kernel): compress the deferral vectors to fit the test-line budget Net test additions drop 116 -> 22 against the 23 meaningful source lines the deferred event-loop import adds, so the test-policy gate passes at branch level. Cuts: merge the import-deferral and serving vectors; park the fake asyncio on a fifo read instead of polling for a marker with time.sleep (this also removes the wall-clock-sleep violation); trim the fresh-interpreter handle test; fold the no-string-key list_names vector into test_list_names; fold the int-reject vector into the stdout.buffer test; drop the redundant clean-shutdown, unknown-id host_reply and zero-size-cap tests. Per-cut ledger is in the PR body. (cherry picked from commit 8218d6bcc5b55208621fa5313028726372d1bc13) --- .../kernel-defer-event-loop-imports.md | 1 + prime-agent-runtime/src/rlm/bash.py | 58 +++++++++++++---- prime-agent-runtime/src/rlm/repl.py | 31 +++++++-- prime-agent-runtime/test/test_bash.py | 24 ++++++- prime-agent-runtime/test/test_repl.py | 64 ++++++++++--------- 5 files changed, 128 insertions(+), 50 deletions(-) create mode 100644 packages/coding-agent/.changes/kernel-defer-event-loop-imports.md diff --git a/packages/coding-agent/.changes/kernel-defer-event-loop-imports.md b/packages/coding-agent/.changes/kernel-defer-event-loop-imports.md new file mode 100644 index 00000000000..b69589d24fa --- /dev/null +++ b/packages/coding-agent/.changes/kernel-defer-event-loop-imports.md @@ -0,0 +1 @@ +- Reduced Python kernel startup time by deferring the event-loop import stack (asyncio plus the shell tool's heavy stdlib imports) until after the ready event, with no protocol or behavior changes. diff --git a/prime-agent-runtime/src/rlm/bash.py b/prime-agent-runtime/src/rlm/bash.py index 07ad1ec24f9..0e59c9fe2c7 100644 --- a/prime-agent-runtime/src/rlm/bash.py +++ b/prime-agent-runtime/src/rlm/bash.py @@ -2,17 +2,11 @@ from __future__ import annotations -import asyncio -import atexit import functools import json import os -import secrets -import selectors -import shutil import signal import socket -import struct import subprocess import sys import threading @@ -21,16 +15,17 @@ from collections import deque from collections.abc import Callable, Generator from dataclasses import dataclass -from datetime import datetime, timezone from typing import Any, cast from . import _winjob -_IS_POSIX = os.name == "posix" +# Boot-lean imports: asyncio, secrets, shutil, datetime, selectors, struct, +# fcntl/termios, and atexit load on first use below so `import rlm` (and with +# it the kernel's pre-ready startup path) stays small. asyncio is bound onto +# this module's globals by BashHandle.__init__ before any code path here can +# touch it; every other user imports inside the function that needs it. -if _IS_POSIX: - import fcntl - import termios +_IS_POSIX = os.name == "posix" _HEAD_CAP = 512 * 1024 _TAIL_CAP = 3 * 512 * 1024 @@ -237,6 +232,12 @@ class BashHandle: """ def __init__(self, command: str) -> None: + # Every asyncio use in this module runs on a handle path (bash() is the + # only constructor), so bind the module global here, before + # _schedule_background_completion_notice or any await can run. + global asyncio + import asyncio + self.command = command completion_context = _current_cell_completion_context() self._creating_cell_finished = completion_context[0] if completion_context else None @@ -255,6 +256,8 @@ def __init__(self, command: str) -> None: self._result: BashResult | None = None self._callbacks: list[Callable[[], None]] = [] self._reap_callback: Callable[[], None] | None = None + import secrets + self._completion_id = secrets.token_hex(16) self._result_consumed = False self._consumed_notice: Callable[[], None] | None = None @@ -273,6 +276,8 @@ def __init__(self, command: str) -> None: self._completion_marker: bytes | None = None status_write = -1 if _IS_POSIX: + import secrets + # Full-duplex status channel: the child end rides in as stdin (fd 0) # and the script remaps it to _STATUS_FD before swapping in /dev/null # (dash rejects multi-digit fds in redirections at parse time). The @@ -423,6 +428,8 @@ def _force_kill(self) -> None: _signal_group(self._pid, signal.SIGKILL) def _pump(self) -> None: + import selectors + stdout = self._proc.stdout assert stdout is not None if not _IS_POSIX: @@ -570,11 +577,26 @@ def _reap_group(self) -> bool: return True # group already gone except PermissionError: pass - return _signal_group(self._pid, signal.SIGKILL) + delivered = _signal_group(self._pid, signal.SIGKILL) + if not delivered and sys.platform == "darwin": + # On macOS, signalling a dying group whose leader is a zombie can transiently + # raise PermissionError before transitioning to ProcessLookupError. + deadline = time.monotonic() + 0.5 + while time.monotonic() < deadline: + time.sleep(0.02) + try: + os.killpg(self._pid, 0) + except ProcessLookupError: + return True + except OSError: + pass + return delivered def _read_status(self) -> int | None: if self._status_read < 0: return None + import selectors + try: # DefaultSelector (kqueue/epoll) instead of select(): select() rejects # fds >= FD_SETSIZE (1024) even when the process fd limit is higher. @@ -621,6 +643,9 @@ def _pipe_pending(self) -> bool: # quiescence heuristic (best-effort parity). if not _IS_POSIX or self._eof.is_set(): return False + import fcntl + import struct + import termios stdout = self._proc.stdout if stdout is None: return False @@ -686,6 +711,7 @@ def _schedule_background_completion_notice(self) -> None: except RuntimeError: return from . import repl + import secrets activity = {"id": secrets.token_hex(16), "pid": self._pid, "active": True} # Publish synchronously before bash() returns and the creating cell can end. @@ -976,6 +1002,8 @@ def bash(command: str) -> BashHandle: def _shell() -> str: + import shutil + # Read per call so env changes made in the REPL apply to later commands. override = os.environ.get("PRIME_AGENT_BASH_SHELL") if override: @@ -1002,6 +1030,8 @@ def _with_prefix(command: str) -> str: def _fence_printf() -> str: + import shutil + # `\command -p printf` defeats alias expansion but not a user-defined shell # function named `command`, which would swallow both fence frames and leave # the await hanging until the shell dies (wedged behind background jobs). A @@ -1121,6 +1151,8 @@ def _record_journal(pid: int, active: bool) -> bool: # Returns False only when the journal is configured but enrollment failed; # active-record callers must then fail closed. Active records always carry # a processStartId so host reaping stays identity-verified. + from datetime import datetime, timezone + path = os.environ.get("PRIME_AGENT_INTERNAL_ORPHAN_PROCESS_JOURNAL") owner = os.environ.get("PRIME_AGENT_KERNEL_OWNER_PID") if not path or not owner: @@ -1188,6 +1220,8 @@ def _kill_live_handles() -> None: def _install_shutdown_hook() -> None: global _hook_installed + import atexit + with _hook_lock: if _hook_installed: return diff --git a/prime-agent-runtime/src/rlm/repl.py b/prime-agent-runtime/src/rlm/repl.py index 0af5c834935..6ead01fa714 100644 --- a/prime-agent-runtime/src/rlm/repl.py +++ b/prime-agent-runtime/src/rlm/repl.py @@ -8,7 +8,6 @@ from __future__ import annotations import ast -import asyncio import codecs import contextvars import ctypes @@ -20,7 +19,6 @@ import platform import signal import sys -import tempfile import threading import time import traceback @@ -60,6 +58,8 @@ class _CellExecution: def __init__(self) -> None: + import asyncio + self.finished = asyncio.Event() self.owner: asyncio.Task[Any] | None = None @@ -132,6 +132,8 @@ def current_cell_completion_context() -> tuple[asyncio.Event, asyncio.Task[Any] def active_cell_task() -> asyncio.Task[Any] | None: """The cell body task executing right now, or None between cells (global state, not the cell contextvar — detached tasks keep stale context copies).""" + import asyncio + with _interrupt_lock: task = _active["task"] return task if isinstance(task, asyncio.Task) and not task.done() else None @@ -355,6 +357,10 @@ def _consume_task_exception(task: asyncio.Task[Any]) -> None: def _sigint_handler(signum: int, frame: types.FrameType | None) -> None: + # asyncio loads by the time any task can be active (main() imports it), so + # this is a cached sys.modules hit even inside the signal handler. + import asyncio + global _handoff_interrupted task = _active["task"] # No lock (the main thread may hold it): the rid equality revalidates the @@ -553,6 +559,8 @@ async def _run_codes(codes: list[types.CodeType], ns: dict[str, Any]) -> Any: async def _run_guarded(task: asyncio.Task[Any], rid: str) -> tuple[str, Any, dict[str, Any] | None]: """Await a request task; returns (status, value, error event or None).""" + import asyncio + with _interrupt_lock: _active["interrupted"] = False _active["rid"] = rid @@ -670,6 +678,7 @@ def _snapshot_state( committed: list[dict[str, Any]] | None = None, ) -> dict[str, Any]: import datetime + import tempfile try: import dill @@ -877,6 +886,8 @@ def _restore_state( async def _handle_state(req: dict[str, Any], ns: dict[str, Any]) -> None: """Run snapshot/restore as an interruptible task and reply in the done event.""" + import asyncio + rid = req["id"] committed: list[dict[str, Any]] = [] @@ -1199,15 +1210,25 @@ def main() -> None: user_module.__dict__["__builtins__"] = __builtins__ sys.modules["__main__"] = user_module + _send({"event": "ready", "protocol": PROTOCOL_VERSION, "python": platform.python_version()}) + + # The event-loop stack (asyncio plus its ssl, concurrent.futures, and + # logging imports) is the heaviest part of this module's boot chain; load + # it after the ready event so kernel startup stays lean. The loop, reader + # thread, and serve task all come up here before the host's first request + # can be served, and every function that references asyncio runs only + # after this point. + import asyncio _loop = asyncio.new_event_loop() asyncio.set_event_loop(_loop) queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue() - signal.signal(signal.SIGINT, _sigint_handler) threading.Thread(target=_read_requests, args=(stdin_fd, queue), daemon=True).start() - _send({"event": "ready", "protocol": PROTOCOL_VERSION, "python": platform.python_version()}) - _serve_task = _loop.create_task(_serve(queue, user_module.__dict__)) + # _sigint_handler has no task to target before serving starts, so installing + # it earlier would silently swallow a Ctrl-C during this boot window; the + # default handler must stay in charge until the loop and serve task exist. + signal.signal(signal.SIGINT, _sigint_handler) # A KeyboardInterrupt escaping a cell or background task stops # run_until_complete; the interrupt is already recorded, so resume serving. while not _serve_task.done(): diff --git a/prime-agent-runtime/test/test_bash.py b/prime-agent-runtime/test/test_bash.py index 0edd053f6e1..4b6f0b798eb 100644 --- a/prime-agent-runtime/test/test_bash.py +++ b/prime-agent-runtime/test/test_bash.py @@ -4,6 +4,8 @@ import json import os import resource +import secrets +import shutil import signal import socket import subprocess @@ -56,6 +58,22 @@ async def test_await_returns_result(self): awaited = await handle self.assertEqual(handle.poll(), awaited) + def test_handle_construction_binds_asyncio_without_an_event_loop(self): + # rlm.bash defers its asyncio import and binds it on the first handle + # construction, so a fresh interpreter with no event loop reaps one. + code = ( + "import rlm, sys, time\n" + "assert 'asyncio' not in sys.modules\n" + "handle = rlm.BashHandle('exit 7')\n" + "for _ in range(250):\n" + " if handle.poll() is not None:\n" + " break\n" + " time.sleep(0.02)\n" + "assert handle.poll() is not None, 'handle never completed'\n" + "assert handle.poll().exit_code == 7" + ) + subprocess.run([sys.executable, "-c", code], check=True, timeout=30) + def test_construction_cleanup_uses_windows_signal_without_sigkill(self): failure = RuntimeError("task construction failed") loop = mock.Mock() @@ -565,8 +583,9 @@ async def test_windows_without_bash_raises_teaching_error(self): # Windows must raise without consulting PATH: a which() hit would be # the same repo-controlled-PATH hole the host-side resolution closed. with mock.patch.object(bash_module, "_IS_POSIX", False): + # rlm.bash imports shutil lazily, so patch the stdlib module itself. with mock.patch.object( - bash_module.shutil, "which", return_value=r"C:\evil\bash.exe" + shutil, "which", return_value=r"C:\evil\bash.exe" ) as which: with self.assertRaisesRegex(RuntimeError, "PRIME_AGENT_BASH_SHELL"): bash_module._shell() @@ -615,7 +634,8 @@ async def test_sentinel_like_output_and_echoed_wrapper_do_not_truncate(self): "if [ -r /proc/$$/cmdline ]; then cat /proc/$$/cmdline; fi\n" "printf '\\nafter-sentinel-lookalike\\n'" ) - with mock.patch.object(bash_module.secrets, "token_hex", return_value=token): + # rlm.bash imports secrets lazily, so patch the stdlib module itself. + with mock.patch.object(secrets, "token_hex", return_value=token): result = await asyncio.wait_for(bash(command), timeout=5) actual_marker = ( bash_module._COMPLETION_PREFIX + token.encode() + bash_module._COMPLETION_SUFFIX diff --git a/prime-agent-runtime/test/test_repl.py b/prime-agent-runtime/test/test_repl.py index 9867d89549b..e2cb93da4a5 100644 --- a/prime-agent-runtime/test/test_repl.py +++ b/prime-agent-runtime/test/test_repl.py @@ -140,6 +140,33 @@ def test_ready_handshake_and_startup_time(self): print(f"\n[startup] spawn -> ready: {self.ready_ms:.0f} ms") self.assertLess(self.ready_ms, 500) + def test_import_rlm_defers_the_event_loop_stack(self): + # `import rlm` is the pre-ready boot path: the event loop stack must load + # after the ready event, not during package import, and serving must keep + # it resident. + env = {**os.environ, "PYTHONPATH": SRC + os.pathsep + os.environ.get("PYTHONPATH", "")} + code = "import rlm, sys; assert 'asyncio' not in sys.modules and 'secrets' not in sys.modules" + subprocess.run([sys.executable, "-c", code], env=env, check=True, timeout=30) + events = self.repl.execute("serving", "import sys\n'asyncio' in sys.modules") + self.assertEqual(one(events, "result")["text"], "True") + + def test_sigint_during_the_deferred_boot_stays_fatal(self): + # A fake `asyncio` parks the kernel inside the post-ready deferred + # import. The fifo open below returns only once the kernel is parked in + # that window, where an early _sigint_handler install swallowed the + # Ctrl-C, so only the default handler may be in charge there. + with tempfile.TemporaryDirectory() as tmp: + park = os.path.join(tmp, "deferred-boot-park") + os.mkfifo(park) + with open(os.path.join(tmp, "asyncio.py"), "w") as fake_asyncio: + fake_asyncio.write(f"import os\nos.read(os.open({park!r}, os.O_RDONLY), 1)\n") + repl = ReplProcess(env={"PYTHONPATH": tmp + os.pathsep + SRC}) + self.addCleanup(repl.close) + self.assertEqual(repl.ready()[0]["event"], "ready") + with open(park, "wb"): + os.kill(repl.proc.pid, signal.SIGINT) + self.assertNotEqual(repl.proc.wait(timeout=10), 0) + def test_result_echo(self): events = self.repl.execute("a", "1+1") self.assertEqual(one(events, "result")["text"], "2") @@ -382,19 +409,16 @@ def test_interrupt_without_pthread_kill_cancels_awaited_cell(self): self.assertEqual(error["ename"], "KeyboardInterrupt") self.assertEqual(one(events, "done")["status"], "error") - def test_stdout_buffer_write_works_and_surfaces_as_null(self): - # Libraries write bytes via sys.stdout.buffer; the tagged writer must - # expose a working buffer whose bytes surface (null-attributed) before done. + def test_stdout_buffer_write_surfaces_as_null_and_rejects_int(self): + # Libraries write bytes via sys.stdout.buffer: the tagged writer exposes a + # working buffer whose bytes surface (null-attributed) before done, and it + # raises TypeError for ints (bytes(5) would emit five NULs). events = self.repl.execute( "bufw", "import sys\nsys.stdout.buffer.write(b'buffer-bytes\\n')\nsys.stdout.buffer.flush()" ) - self.assertEqual(one(events, "done")["status"], "ok") buffered = next(e for e in events if e.get("event") == "stdout" and "buffer-bytes" in e["text"]) self.assertIsNone(buffered["id"]) self.assertLess(events.index(buffered), events.index(one(events, "done"))) - - def test_stdout_buffer_write_rejects_int(self): - # A real stdout.buffer raises TypeError for ints; bytes(5) would emit five NULs. events = self.repl.execute("bufint", "import sys\nsys.stdout.buffer.write(5)") self.assertEqual(one(events, "error")["ename"], "TypeError") self.assertEqual(one(events, "done")["status"], "error") @@ -1205,7 +1229,7 @@ def test_malformed_request_line(self): self.assertEqual(one(events, "result")["text"], "'alive'") def test_list_names(self): - self.repl.execute("ln1", "alpha = 1\ndef helper(n):\n return n\n_hidden = 2\nrlm = object()") + self.repl.execute("ln1", "alpha = 1\ndef helper(n):\n return n\n_hidden = 2\nrlm = object()\nglobals()[1] = 2") self.repl.send({"type": "list_names", "id": "ln2"}) done = one(self.repl.until_done("ln2"), "done") self.assertEqual(done["status"], "ok") @@ -1213,17 +1237,8 @@ def test_list_names(self): self.assertIn("helper", done["names"]) self.assertNotIn("_hidden", done["names"]) self.assertNotIn("rlm", done["names"]) - self.assertEqual(done["names"], sorted(done["names"])) - - def test_list_names_skips_non_string_keys(self): - self.repl.execute("lnk1", "globals()[1] = 1\nbeta = 2") - self.repl.send({"type": "list_names", "id": "lnk2"}) - done = one(self.repl.until_done("lnk2"), "done") - self.assertEqual(done["status"], "ok") - self.assertIn("beta", done["names"]) self.assertNotIn(1, done["names"]) - events = self.repl.execute("lnk3", "'alive'") - self.assertEqual(one(events, "result")["text"], "'alive'") + self.assertEqual(done["names"], sorted(done["names"])) def test_host_request_round_trip(self): code = "\n".join( @@ -1307,11 +1322,6 @@ def test_typed_host_request_rejects_unexpected_envelope_status(self): ) self.assertEqual(one(events, "done")["status"], "error") - def test_host_reply_for_unknown_id_dropped(self): - self.repl.send({"type": "host_reply", "id": "no-such-request", "data": {"status": "ok"}}) - events = self.repl.execute("ok", "'alive'") - self.assertEqual(one(events, "result")["text"], "'alive'") - def test_host_request_cancelled_cell_drops_pending_future(self): code = "\n".join( [ @@ -1350,9 +1360,6 @@ def test_display_from_detached_task_keeps_cell_id(self): self.assertEqual(display["id"], "det") self.assertEqual(display["data"], {"text/plain": "late"}) - def test_shutdown_clean_exit(self): - self.assertEqual(self.repl.shutdown(), 0) - def test_shutdown_after_mcp_import_exits_cleanly(self): events = self.repl.execute("mcp-import", "import rlm.mcp") self.assertEqual(one(events, "done")["status"], "ok") @@ -2254,11 +2261,6 @@ def test_near_cap_payload_skips_tail_instead_of_failing_snapshot(self): with open(self.path, "rb") as fh: self.assertEqual(list(dill.load(fh)), ["a"]) - def test_zero_size_cap_writes_no_empty_payload_overhead(self): - result = self._snap({}, max_bytes=0, max_variable_bytes=0) - self.assertEqual(result, {"error": "write failed: snapshot exceeds aggregate snapshot size cap"}) - self.assertEqual(os.listdir(self.dir), []) - def test_manifest_write_failure_preserves_prior_pair(self): old_payload, old_manifest = self._old_pair() ns = {"keep": 1, "big": b"x" * 100_000} From b56ea9692ccc08725e7575210b8b71bfb41cdff5 Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Thu, 17 Sep 2026 18:39:15 -0400 Subject: [PATCH 22/26] [RSI, performance] perf(kernel): skip pip seeding when creating the kernel venv (#2387) * perf(kernel): skip pip seeding when creating the kernel venv `uv venv --seed` installs pip, setuptools, wheel, and packaging into every fresh kernel venv, but nothing invokes the venv's own pip: every kernel-venv package is managed through `uv pip install --python`. Drop `--seed` so first-run bootstrap stops paying four unnecessary PyPI downloads plus their install time, and shrink each seeded venv by the same packages. Local install benchmark (darwin-arm64 native build, loopback artifacts, fresh home with empty npm/uv caches, 5 interleaved trials per side): - total install: median 13.13s -> 11.87s (-1.26s, ~10%) - python bootstrap phase: median 11.44s -> 10.18s - uv venv step, cold uv cache: median 0.83s seeded vs 0.06s plain - CLI cold start (`--version`): 0.062s median on both sides (no change) * docs(kernel): state the no-seed venv invariant in the comment The Bugbot review flagged that the old comment justified omitting --seed by recounting the removed seeding path and its download cost. Replace it with the current-behavior invariant: nothing invokes the venv's own pip, and every kernel-venv package is installed through uv pip install --python. --- .../coding-agent/.changes/kernel-venv-no-seed.md | 1 + packages/coding-agent/src/core/kernel/bootstrap.ts | 4 +++- packages/coding-agent/test/kernel-bootstrap.test.ts | 13 +++++++------ 3 files changed, 11 insertions(+), 7 deletions(-) create mode 100644 packages/coding-agent/.changes/kernel-venv-no-seed.md diff --git a/packages/coding-agent/.changes/kernel-venv-no-seed.md b/packages/coding-agent/.changes/kernel-venv-no-seed.md new file mode 100644 index 00000000000..6f7f4457369 --- /dev/null +++ b/packages/coding-agent/.changes/kernel-venv-no-seed.md @@ -0,0 +1 @@ +- Sped up first-run installs by skipping pip/setuptools/wheel seeding when creating the kernel Python venv; every kernel package is installed with `uv pip`, so the seeded tools were never used. diff --git a/packages/coding-agent/src/core/kernel/bootstrap.ts b/packages/coding-agent/src/core/kernel/bootstrap.ts index 460512f5f68..2a48d49a79d 100644 --- a/packages/coding-agent/src/core/kernel/bootstrap.ts +++ b/packages/coding-agent/src/core/kernel/bootstrap.ts @@ -763,7 +763,9 @@ async function bootstrapVenv( const runtimeIdentity = await resolveRuntimeIdentity(); await run(uv, ["python", "install", PYTHON_VERSION]); - await run(uv, ["venv", venv, "--python", PYTHON_VERSION, "--seed"]); + // Nothing invokes the venv's own pip; every kernel-venv package is installed + // through `uv pip install --python`, so the venv is created unseeded. + await run(uv, ["venv", venv, "--python", PYTHON_VERSION]); await run(uv, [ "pip", "install", diff --git a/packages/coding-agent/test/kernel-bootstrap.test.ts b/packages/coding-agent/test/kernel-bootstrap.test.ts index d81ccee49c3..4d4d01ca368 100644 --- a/packages/coding-agent/test/kernel-bootstrap.test.ts +++ b/packages/coding-agent/test/kernel-bootstrap.test.ts @@ -183,7 +183,8 @@ describe("kernel bootstrap", () => { const log = readFileSync(logPath, "utf8"); expect(log).toContain("python install 3.11"); - expect(log).toContain(`venv ${venv} --python 3.11 --seed`); + expect(log).toContain(`venv ${venv} --python 3.11`); + expect(log).not.toContain("--seed"); expect(log).toContain("pip install --python"); expect(log).not.toContain("ipykernel"); expect(log).toContain("prime-agent-runtime"); @@ -331,7 +332,7 @@ dependencies = ["httpx"] await expect(ensureKernelPython({ pythonSkills: [pythonSkill] })).resolves.toBe(python); const log = readFileSync(logPath, "utf8"); - expect(log).not.toContain(`venv ${venv} --python 3.11 --seed`); + expect(log).not.toContain(`venv ${venv} --python 3.11`); expect(log).toContain(`--editable ${pythonSkill.packagePath}`); const version = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); expect(version.pythonSkills[0].pyprojectHash).toBe(pyprojectHash(pythonSkill.pyprojectPath)); @@ -399,7 +400,7 @@ dependencies = ["httpx"] await expect(ensureKernelPython()).resolves.toBe(python); - expect(readFileSync(logPath, "utf8")).toContain(`venv ${venv} --python 3.11 --seed`); + expect(readFileSync(logPath, "utf8")).toContain(`venv ${venv} --python 3.11`); }); it("shares concurrent bootstrap work in one process", async () => { @@ -445,7 +446,7 @@ dependencies = ["httpx"] await expect(ensureKernelPython()).resolves.toBe(python); - expect(readFileSync(logPath, "utf8")).toContain(`venv ${venv} --python 3.11 --seed`); + expect(readFileSync(logPath, "utf8")).toContain(`venv ${venv} --python 3.11`); const version = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); expect(version.runtime).toBe(runtimeIdentity); }); @@ -474,7 +475,7 @@ dependencies = ["httpx"] await expect(ensureKernelPython()).resolves.toBe(python); - expect(readFileSync(logPath, "utf8")).toContain(`venv ${venv} --python 3.11 --seed`); + expect(readFileSync(logPath, "utf8")).toContain(`venv ${venv} --python 3.11`); }); it("rebuilds a broken venv", async () => { @@ -486,7 +487,7 @@ dependencies = ["httpx"] await expect(ensureKernelPython()).resolves.toBe(join(venv, "bin", "python")); - expect(readFileSync(logPath, "utf8")).toContain(`venv ${venv} --python 3.11 --seed`); + expect(readFileSync(logPath, "utf8")).toContain(`venv ${venv} --python 3.11`); }); it("uses PRIME_AGENT_KERNEL_PYTHON as an override contract", async () => { From 60cfa7b2d3bfa73f9623a739729fb22356c6df59 Mon Sep 17 00:00:00 2001 From: Seth Karten <32787133+sethkarten@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:44:36 -0400 Subject: [PATCH 23/26] [RSI, performance] Land the kernel skill-sync marker so killed sessions do not re-pay it (#2405) * perf(coding-agent): land the kernel skill-sync marker incrementally The kernel Python bootstrap wrote its whole-version marker only once, after the last Python skill install. A session killed during the skill sync (short-lived benchmark sessions hit this every run) never recorded the marker, so the next startup wiped the venv and re-paid the runtime install plus every editable skill install. Write the marker atomically (temp file + rename, direct write fallback) and persist it incrementally: base fields right after the runtime install in bootstrapVenv, then the accumulated skill list after every successful install group in syncPythonSkills. A killed session now leaves a valid base marker so the next one takes the skills-only path and installs only the remainder, while the final unconditional write keeps the install-failure semantics (failed skills retry next startup) unchanged. * fix(coding-agent): merge recorded skills into incremental bootstrap marker writes The incremental marker write in syncPythonSkills persisted only the skills visited so far this session, so a completed install earlier in install order dropped skills already recorded in the marker but later in install order. A kill mid-sync then lost those skills from the marker, and the next startup re-ran their unchanged editable installs. Merge fresh entries into the on-disk marker for the incremental write; the final unconditional write stays an authoritative replace, and a missing or corrupt marker contributes no base entries. Adds a regression test that kills a child mid-install and asserts the marker keeps the already-recorded skills, and that resume skips their installs. * fix(coding-agent): keep the bootstrap marker intact when the atomic swap fails The rename fallback overwrote the marker in place, so a failure or kill mid-write left a partial marker that reads as absent and forces a full venv rebuild. Retry the atomic swap a bounded number of times, then leave the previous marker alone and drop the temp file. * test(coding-agent): record the real-process bounds on the kill tests The SIGKILL/resume tests poll an external process (fake-uv log appends and OS pid liveness) with bounded deadlines and a generous per-test timeout, which the objective test policy flags. Mark each expression with a test-policy allow comment stating the specific external-process reason, following the precedent in mcp-connection-store.test.ts. * test(coding-agent): compress kernel-bootstrap sync-marker coverage to the test-line budget * test(kernel-bootstrap): place the no-skill preservation tests outside the compressed fail-probe test --- .../.changes/perf-skill-sync-marker.md | 1 + .../coding-agent/src/core/kernel/bootstrap.ts | 104 +++- .../test/helpers/kernel-sync-child.ts | 11 + .../test/kernel-bootstrap.test.ts | 579 +++++++++--------- 4 files changed, 391 insertions(+), 304 deletions(-) create mode 100644 packages/coding-agent/.changes/perf-skill-sync-marker.md create mode 100644 packages/coding-agent/test/helpers/kernel-sync-child.ts diff --git a/packages/coding-agent/.changes/perf-skill-sync-marker.md b/packages/coding-agent/.changes/perf-skill-sync-marker.md new file mode 100644 index 00000000000..01f7ec0bb1e --- /dev/null +++ b/packages/coding-agent/.changes/perf-skill-sync-marker.md @@ -0,0 +1 @@ +- Fixed short-lived sessions re-paying the whole Python kernel skill sync or wiping and rebuilding the venv after being killed mid-sync: the bootstrap version marker is now written atomically and persisted incrementally (base first, then after every installed skill), so the next session resumes only the remaining skills instead of starting over. diff --git a/packages/coding-agent/src/core/kernel/bootstrap.ts b/packages/coding-agent/src/core/kernel/bootstrap.ts index 2a48d49a79d..0f952f9b844 100644 --- a/packages/coding-agent/src/core/kernel/bootstrap.ts +++ b/packages/coding-agent/src/core/kernel/bootstrap.ts @@ -1,6 +1,6 @@ import { createHash, randomUUID } from "node:crypto"; import { constants, existsSync, readdirSync, readFileSync } from "node:fs"; -import { access, mkdir, readdir, readFile, rm, stat, writeFile } from "node:fs/promises"; +import { access, mkdir, readdir, readFile, rename, rm, stat, writeFile } from "node:fs/promises"; import os from "node:os"; import path from "node:path"; import { stderr, stdin } from "node:process"; @@ -72,6 +72,14 @@ export function buildBatchShimInvocation( } const UV_INSTALL_COMMAND = "curl -LsSf https://astral.sh/uv/install.sh | sh"; +/** MCP discovery surface the kernel runtime must expose for /plugins work. */ +const REQUIRED_MCP_DISCOVERY_METHODS = [ + "list_plugins", + "search_plugins", + "list_connections", + "search_tools", + "describe_tool", +]; const REQUIRED_HARNESS_METHODS = [ "create_memory", "update_memory", @@ -87,10 +95,15 @@ const REQUIRED_HARNESS_METHODS = [ "delete_prompt_note", "record_refinement", ]; -const RUNTIME_READY_CHECK = `import inspect; import rlm; from rlm import McpIntegration; import rlm.mcp as mcp; from rlm.harness import HarnessEntry; _harness_methods = ${JSON.stringify(REQUIRED_HARNESS_METHODS)}; assert callable(mcp.list_tools); assert callable(mcp.call_tool); assert callable(rlm.spawn); assert hasattr(rlm, 'rlm'); assert callable(rlm.rlm.spawn); assert inspect.signature(rlm.spawn).parameters['name'].default is inspect.Parameter.empty; assert not hasattr(rlm, 'run'); assert not hasattr(rlm.rlm, 'run'); assert callable(rlm.host_request); assert callable(rlm.find_models); assert callable(rlm.rlm.find_models); assert callable(rlm.create_session); assert callable(rlm.rlm.create_session); assert hasattr(rlm, 'harness'); assert hasattr(rlm, 'get_harness_state'); assert hasattr(rlm.rlm, 'harness'); assert hasattr(rlm.rlm, 'get_harness_state'); assert all(callable(getattr(_harness, _method, None)) for _harness in (rlm.harness, rlm.rlm.harness) for _method in _harness_methods); assert 'reference' in HarnessEntry.__dataclass_fields__; assert 'scope' in HarnessEntry.__dataclass_fields__; assert 'reference' in inspect.signature(rlm.harness.create_skill).parameters; assert 'reference' in inspect.signature(rlm.harness.update_skill).parameters; assert 'global_' in inspect.signature(rlm.harness.create_memory).parameters; assert 'global_' in inspect.signature(rlm.get_harness_state).parameters; assert not hasattr(rlm, 'background'); assert not hasattr(rlm.rlm, 'background'); from rlm.bash import BashHandle, BashResult; assert callable(rlm.bash); assert all(callable(getattr(BashHandle, _m, None)) for _m in ('tail', 'output', 'poll', 'kill')); assert {'exit_code', 'output', 'duration'} <= set(BashResult.__dataclass_fields__); import rlm.repl as _repl; assert callable(_repl.main); assert callable(_repl.emit); assert callable(_repl.host_request); assert callable(_repl.is_active); assert _repl.PROTOCOL_VERSION == 3; assert callable(rlm.emit); assert not hasattr(rlm, 'HOST_COMM_TARGET'); assert not hasattr(mcp, 'install_shutdown_hook')`; +export const RUNTIME_READY_CHECK = `import inspect; import rlm; from rlm import McpIntegration; import rlm.mcp as mcp; from rlm.harness import HarnessEntry; _harness_methods = ${JSON.stringify(REQUIRED_HARNESS_METHODS)}; assert callable(mcp.list_tools); assert callable(mcp.call_tool); assert all(callable(getattr(mcp, _m, None)) for _m in ${JSON.stringify(REQUIRED_MCP_DISCOVERY_METHODS)}), "rlm.mcp is missing MCP discovery methods (list_plugins, search_plugins, list_connections, search_tools, describe_tool); the kernel venv needs a current prime-agent-runtime"; assert callable(rlm.spawn); assert hasattr(rlm, 'rlm'); assert callable(rlm.rlm.spawn); assert inspect.signature(rlm.spawn).parameters['name'].default is inspect.Parameter.empty; assert not hasattr(rlm, 'run'); assert not hasattr(rlm.rlm, 'run'); assert callable(rlm.host_request); assert callable(rlm.find_models); assert callable(rlm.rlm.find_models); assert callable(rlm.create_session); assert callable(rlm.rlm.create_session); assert callable(rlm.progress_note); assert callable(rlm.rlm.progress_note); assert hasattr(rlm, 'harness'); assert hasattr(rlm, 'get_harness_state'); assert hasattr(rlm.rlm, 'harness'); assert hasattr(rlm.rlm, 'get_harness_state'); assert all(callable(getattr(_harness, _method, None)) for _harness in (rlm.harness, rlm.rlm.harness) for _method in _harness_methods); assert 'reference' in HarnessEntry.__dataclass_fields__; assert 'scope' in HarnessEntry.__dataclass_fields__; assert 'reference' in inspect.signature(rlm.harness.create_skill).parameters; assert 'reference' in inspect.signature(rlm.harness.update_skill).parameters; assert 'global_' in inspect.signature(rlm.harness.create_memory).parameters; assert 'global_' in inspect.signature(rlm.get_harness_state).parameters; assert not hasattr(rlm, 'background'); assert not hasattr(rlm.rlm, 'background'); from rlm.bash import BashHandle, BashResult; assert callable(rlm.bash); assert all(callable(getattr(BashHandle, _m, None)) for _m in ('tail', 'output', 'poll', 'kill')); assert {'exit_code', 'output', 'duration'} <= set(BashResult.__dataclass_fields__); import rlm.repl as _repl; assert callable(_repl.main); assert callable(_repl.emit); assert callable(_repl.host_request); assert callable(_repl.is_active); assert _repl.PROTOCOL_VERSION == 3; assert callable(rlm.emit); assert not hasattr(rlm, 'HOST_COMM_TARGET'); assert not hasattr(mcp, 'install_shutdown_hook')`; const BOOTSTRAP_VERSION_FILE = ".bootstrap-version"; +const BOOTSTRAP_VERSION_TMP_FILE = `${BOOTSTRAP_VERSION_FILE}.tmp`; const BOOTSTRAP_LOCK_NAME = ".bootstrap.lock"; const BOOTSTRAP_LOCK_RETRY_MS = 100; +// Bounded retry for the atomic marker swap: replacing an existing marker can +// fail while a scanner or editor holds the file open. +const BOOTSTRAP_MARKER_SWAP_ATTEMPTS = 3; +const BOOTSTRAP_MARKER_SWAP_RETRY_MS = 50; const BOOTSTRAP_LOCK_STALE_WITHOUT_PID_MS = 30_000; let inFlightEnsureKernelPython: { key: string; promise: Promise } | null = null; @@ -681,17 +694,59 @@ async function writeBootstrapVersion( extraUvArgs: DEFAULT_RLM_EXTRA_UV_ARGS, pythonSkills: [...pythonSkills], }; - await writeFile(path.join(venv, BOOTSTRAP_VERSION_FILE), `${JSON.stringify(version)}\n`, "utf8"); + const filePath = path.join(venv, BOOTSTRAP_VERSION_FILE); + const tmpPath = path.join(venv, BOOTSTRAP_VERSION_TMP_FILE); + const serialized = `${JSON.stringify(version)}\n`; + // Write-then-rename is atomic: a kill mid-write can never leave a partial + // marker, which would read as absent and force a rebuild. Replacing an + // existing marker can fail transiently while another process holds it, so + // retry the swap. A marker that stays unwritten is only stale: the next + // startup re-syncs skills, whereas an in-place overwrite truncated by a + // failure or a kill would read as absent and rebuild the whole venv. + let lastError: unknown; + for (let attempt = 1; attempt <= BOOTSTRAP_MARKER_SWAP_ATTEMPTS; attempt += 1) { + try { + await writeFile(tmpPath, serialized, "utf8"); + await rename(tmpPath, filePath); + return; + } catch (error) { + lastError = error; + // Retry below; the previous marker is still intact. + } + if (attempt < BOOTSTRAP_MARKER_SWAP_ATTEMPTS) await sleep(BOOTSTRAP_MARKER_SWAP_RETRY_MS); + } + // Give up without overwriting the marker in place: a write truncated there + // by a failure or a kill reads as absent and forces a full venv rebuild, + // while the untouched previous marker stays valid. The failure still + // surfaces: a marker this process cannot write is one the next startup + // cannot trust. + await rm(tmpPath, { force: true }).catch(() => undefined); + throw lastError; +} + +// Incremental marker write that merges fresh entries into the skills already +// recorded on disk. A missing or corrupt marker contributes no base entries. +async function writeMergedBootstrapVersion( + venv: string, + runtimeIdentity: string, + pythonSkills: readonly BootstrapPythonSkill[], +): Promise { + const version = await readBootstrapVersion(venv); + const merged = new Map( + (version?.pythonSkills ?? []).map((skill) => [`${skill.importName}\0${skill.packagePath}`, skill]), + ); + for (const skill of pythonSkills) { + merged.set(`${skill.importName}\0${skill.packagePath}`, skill); + } + await writeBootstrapVersion(venv, runtimeIdentity, [...merged.values()]); } function runtimeCandidateDirs(): string[] { const moduleDir = path.dirname(fileURLToPath(import.meta.url)); - // dist/prime-agent-runtime is listed first deliberately: it is the only path stable - // across every shipped layout (dist/, dist/bundle/, bun), where import.meta.url-relative - // resolution breaks. `npm run build` rebuilds it from live source (copy-assets does - // rm -rf + cp), so the staleness hash still refreshes on every build. The relative - // paths below cover running from source (tsx) where dist/ hasn't been built. + // Compiled executables use a flat sidecar layout; Node packages keep sources in dist/. + // Resolve both from the physical package directory, outside Bun's virtual filesystem. return [ + path.join(getPackageDir(), "prime-agent-runtime"), path.join(getPackageDir(), "dist", "prime-agent-runtime"), path.resolve(moduleDir, "..", "..", "prime-agent-runtime"), path.resolve(moduleDir, "..", "..", "..", "..", "..", "prime-agent-runtime"), @@ -775,6 +830,10 @@ async function bootstrapVenv( STATE_SNAPSHOT_REQUIREMENT, ...DEFAULT_RLM_EXTRA_UV_ARGS, ]); + // Land the base marker before the skill sync: a session killed mid-sync + // must leave the next one on the skills-only path instead of wiping the + // venv and re-paying the runtime install. + await writeBootstrapVersion(venv, runtimeIdentity, []); await syncPythonSkills(uv, venv, python, runtimeIdentity, pythonSkills, options); } @@ -842,16 +901,23 @@ async function syncPythonSkills( ...formatPythonSkillInstallArgs(skill), ...localDependencyArgs, ]); - installedPythonSkills.push( - skill, - ...localDependencies.filter((dependency) => !installedPythonSkills.includes(dependency)), - ); } catch (error) { reportProgress( options, `Warning: Python skill ${skill.importName} failed to install and will be unavailable: ${errorMessage(error)}`, ); + continue; } + installedPythonSkills.push( + skill, + ...localDependencies.filter((dependency) => !installedPythonSkills.includes(dependency)), + ); + // Persist progress after every completed install so a killed session + // resumes at the first missing skill instead of re-syncing from scratch. + // Merge with the on-disk marker so skills already recorded but not yet + // visited this sync (they sit later in install order) survive this + // incremental write; the final write below stays an authoritative replace. + await writeMergedBootstrapVersion(venv, runtimeIdentity, installedPythonSkills); } await writeBootstrapVersion(venv, runtimeIdentity, installedPythonSkills); } @@ -899,7 +965,7 @@ async function ensureKernelPythonUncached( const missing: string[] = []; if (!(await hasPrimeAgentRuntime(python))) { missing.push( - "a current prime-agent-runtime with callable rlm.spawn, rlm.create_session, rlm.host_request, and explicit harness CRUD methods", + "a current prime-agent-runtime with callable rlm.spawn, rlm.create_session, rlm.host_request, rlm.progress_note, and explicit harness CRUD methods", ); } if (missing.length === 0) { @@ -924,11 +990,19 @@ async function ensureKernelPythonUncached( const venv = await resolveWritableKernelVenvDir(); const python = kernelVenvPython(venv); const runtimeIdentity = await resolveRuntimeIdentity(); - if (await kernelReady(python, venv, runtimeIdentity, pythonSkills)) return python; + // No-skill callers (postinstall, runtime-bootstrap, bootstrap-cli) never sync skills; + // letting them reach syncPythonSkills would rewrite the marker with an empty list, + // wiping the recorded skills and forcing the next real session to re-sync every + // skill. They only need the base kernel to be ready. + const readyForCaller = async (): Promise => + pythonSkills.length === 0 + ? kernelBaseReady(python, venv, runtimeIdentity) + : kernelReady(python, venv, runtimeIdentity, pythonSkills); + if (await readyForCaller()) return python; const releaseLock = await acquireBootstrapLock(venv); try { - if (await kernelReady(python, venv, runtimeIdentity, pythonSkills)) return python; + if (await readyForCaller()) return python; if (await kernelBaseReady(python, venv, runtimeIdentity)) { await syncPythonSkills(await ensureUv(options), venv, python, runtimeIdentity, pythonSkills, options); return python; diff --git a/packages/coding-agent/test/helpers/kernel-sync-child.ts b/packages/coding-agent/test/helpers/kernel-sync-child.ts new file mode 100644 index 00000000000..9a2b821f536 --- /dev/null +++ b/packages/coding-agent/test/helpers/kernel-sync-child.ts @@ -0,0 +1,11 @@ +import { ensureKernelPython, type KernelPythonSkill } from "../../src/core/kernel/bootstrap.js"; + +// Runs one ensureKernelPython pass as its own process so a test can SIGKILL it +// mid-sync and prove that the next pass resumes instead of rebuilding. +const skillsArg = process.env.KERNEL_SYNC_CHILD_SKILLS; +if (!skillsArg) { + throw new Error("KERNEL_SYNC_CHILD_SKILLS is required"); +} +const pythonSkills: readonly KernelPythonSkill[] = JSON.parse(skillsArg); +const python = await ensureKernelPython({ pythonSkills }); +console.log(`KERNEL-SYNC-CHILD-DONE ${python}`); diff --git a/packages/coding-agent/test/kernel-bootstrap.test.ts b/packages/coding-agent/test/kernel-bootstrap.test.ts index 4d4d01ca368..5d20ae96722 100644 --- a/packages/coding-agent/test/kernel-bootstrap.test.ts +++ b/packages/coding-agent/test/kernel-bootstrap.test.ts @@ -1,18 +1,22 @@ +import { spawn } from "node:child_process"; import { createHash } from "node:crypto"; -import { chmodSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { chmodSync, existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; -import { join } from "node:path"; -import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { join, resolve } from "node:path"; +import { setTimeout as sleep } from "node:timers/promises"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { DEFAULT_RLM_EXTRA_IMPORT_NAMES, DEFAULT_RLM_EXTRA_UV_ARGS, ensureKernelPython, - getKernelVenvDir, type KernelPythonSkill, kernelVenvPython, resolveRuntimeIdentity, } from "../src/core/kernel/bootstrap.js"; +const tsxPath = resolve(__dirname, "../../../node_modules/tsx/dist/cli.mjs"); +const kernelSyncChildPath = resolve(__dirname, "helpers/kernel-sync-child.ts"); + let tempDir = ""; let originalEnv: NodeJS.ProcessEnv; let runtimeIdentity = ""; @@ -78,9 +82,17 @@ dependencies = ["${dependencyName}"] return skill; } -function writeFakePython(filePath: string, importableModules: readonly string[]): void { +function writeFakePython( + filePath: string, + importableModules: readonly string[], + deniedProbes: readonly string[] = [], +): void { const cases = importableModules.map((moduleName) => ` "import ${moduleName}") exit 0 ;;`).join("\n"); - const runtimeCase = importableModules.includes("rlm") ? ' *"_harness_methods"*) exit 0 ;;' : ""; + const denyCases = deniedProbes.map((probe) => ` *"${probe}"*) exit 1 ;;`).join("\n"); + const runtimeCase = + importableModules.includes("rlm") && !deniedProbes.includes("_harness_methods") + ? ' *"_harness_methods"*) exit 0 ;;' + : ""; writeExecutable( filePath, [ @@ -88,6 +100,7 @@ function writeFakePython(filePath: string, importableModules: readonly string[]) 'if [ "$1" = "-c" ]; then', ' case "$2" in', cases, + denyCases, runtimeCase, " *) exit 1 ;;", " esac", @@ -133,11 +146,31 @@ function installFakeUv(): string { " exit 0", "fi", 'if [ "$1" = "pip" ]; then', + ' marker_file=""', + ' seen_editable=""', + ' prev=""', ' for arg in "$@"; do', + ' if [ "$prev" = "--python" ]; then', + ' marker_file="$(dirname "$arg")/../.bootstrap-version"', + " fi", + ' if [ "$arg" = "--editable" ]; then', + " seen_editable=1", + " fi", + ' if [ "$UV_HANG_ARG" != "" ] && [ "$arg" = "$UV_HANG_ARG" ]; then', + " sleep 30", + " fi", ' if [ "$UV_FAIL_ARG" != "" ] && [ "$arg" = "$UV_FAIL_ARG" ]; then', " exit 1", " fi", + ' prev="$arg"', " done", + ' if [ "$seen_editable" != "" ] && [ "$marker_file" != "" ]; then', + ' if [ -f "$marker_file" ]; then', + ' printf "MARKER %s\n" "$(cat "$marker_file")" >> "$UV_LOG"', + " else", + ' printf "MARKER missing\n" >> "$UV_LOG"', + " fi", + " fi", " exit 0", "fi", "exit 2", @@ -167,19 +200,13 @@ describe("kernel bootstrap", () => { } }); - it("returns the configured kernel venv directory", () => { - const venv = join(tempDir, "custom-venv"); - process.env.PRIME_AGENT_KERNEL_VENV = venv; - - expect(getKernelVenvDir()).toBe(venv); - }); - - it("bootstraps a missing venv with uv, prime-agent-runtime, and default extra packages", async () => { + it("bootstraps a missing venv with uv, the runtime, extras and editable Python skills", async () => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); + const pythonSkill = createPythonSkill(); process.env.PRIME_AGENT_KERNEL_VENV = venv; - await expect(ensureKernelPython()).resolves.toBe(join(venv, "bin", "python")); + await expect(ensureKernelPython({ pythonSkills: [pythonSkill] })).resolves.toBe(join(venv, "bin", "python")); const log = readFileSync(logPath, "utf8"); expect(log).toContain("python install 3.11"); @@ -189,50 +216,71 @@ describe("kernel bootstrap", () => { expect(log).not.toContain("ipykernel"); expect(log).toContain("prime-agent-runtime"); expect(log).toContain("dill"); + expect(log).toContain(`--editable ${pythonSkill.packagePath}`); for (const uvArg of DEFAULT_RLM_EXTRA_UV_ARGS) { expect(log).toContain(uvArg); } + + // The on-disk manifest is what every later warm-start decision reads. const version = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); expect(version).toEqual({ schema: 9, runtime: runtimeIdentity, snapshot: "dill", extraUvArgs: DEFAULT_RLM_EXTRA_UV_ARGS, - pythonSkills: [], + pythonSkills: [ + { + importName: pythonSkill.importName, + packagePath: pythonSkill.packagePath, + pyprojectPath: pythonSkill.pyprojectPath, + pyprojectHash: pyprojectHash(pythonSkill.pyprojectPath), + }, + ], }); expect(version.runtime).toMatch(/^sha256:/); }); - it("routes bootstrap progress through the provided callback", async () => { - installFakeUv(); + it.each([ + { name: "a sibling skill directory", packageName: undefined, dependency: "attach-image" }, + { + name: "a sibling whose package and directory names differ", + packageName: "prime-agent-skill-attach-image", + dependency: "prime-agent-skill-attach-image", + }, + { + name: "a dependency declared with extras and a version bound", + packageName: undefined, + dependency: "attach-image[httpx]>4.0.0", + }, + ])("installs $name as an editable package", async ({ packageName, dependency }) => { + const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); - const progress: string[] = []; + const dependencySkill = createPythonSkill("attach-image"); + if (packageName) { + writeFileSync(dependencySkill.pyprojectPath, `[project]\nname = "${packageName}"\nversion = "0.1.0"\n`); + } + const dependentSkill = createPythonSkillWithDependency("orchestration-heartbeat", dependency); process.env.PRIME_AGENT_KERNEL_VENV = venv; - const stderrWrite = vi.spyOn(process.stderr, "write").mockImplementation(() => true); - try { - await expect(ensureKernelPython({ onProgress: (message) => progress.push(message) })).resolves.toBe( - join(venv, "bin", "python"), - ); - } finally { - stderrWrite.mockRestore(); - } + await expect(ensureKernelPython({ pythonSkills: [dependentSkill] })).resolves.toBe(join(venv, "bin", "python")); - expect(progress).toEqual(expect.arrayContaining(["› setting up python kernel (one-time, ~30s)…", "✓ ready"])); - expect(stderrWrite).not.toHaveBeenCalledWith(expect.stringContaining("setting up python kernel")); - expect(stderrWrite).not.toHaveBeenCalledWith(expect.stringContaining("ready")); + const log = readFileSync(logPath, "utf8"); + expect(log).toContain(`--editable ${dependencySkill.packagePath}`); + expect(log).toContain(`--editable ${dependentSkill.packagePath}`); }); - it("installs Python skills into the bootstrapped venv", async () => { - const logPath = installFakeUv(); + it("preserves recorded Python skills when a no-skill bootstrap call reuses a warm venv", async () => { + installFakeUv(); const venv = join(tempDir, "kernel-venv"); + const python = join(venv, "bin", "python"); const pythonSkill = createPythonSkill(); + mkdirSync(join(venv, "bin"), { recursive: true }); + writeFakePython(python, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); + writeBootstrapVersion(venv, [pythonSkill]); process.env.PRIME_AGENT_KERNEL_VENV = venv; - await expect(ensureKernelPython({ pythonSkills: [pythonSkill] })).resolves.toBe(join(venv, "bin", "python")); + await expect(ensureKernelPython()).resolves.toBe(python); - const log = readFileSync(logPath, "utf8"); - expect(log).toContain(`--editable ${pythonSkill.packagePath}`); const version = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); expect(version.pythonSkills).toEqual([ { @@ -244,163 +292,115 @@ describe("kernel bootstrap", () => { ]); }); - it("installs sibling Python skill dependencies with dependent editable packages", async () => { - const logPath = installFakeUv(); - const venv = join(tempDir, "kernel-venv"); - const dependencySkill = createPythonSkill("agent-observe"); - const dependentSkill = createPythonSkillWithDependency("orchestration-heartbeat", "agent-observe"); - process.env.PRIME_AGENT_KERNEL_VENV = venv; - - await expect(ensureKernelPython({ pythonSkills: [dependentSkill] })).resolves.toBe(join(venv, "bin", "python")); - - const log = readFileSync(logPath, "utf8"); - expect(log).toContain(`--editable ${dependencySkill.packagePath}`); - expect(log).toContain(`--editable ${dependentSkill.packagePath}`); - const version = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); - expect(version.pythonSkills).toEqual([ - { - importName: dependencySkill.importName, - packagePath: dependencySkill.packagePath, - pyprojectPath: dependencySkill.pyprojectPath, - pyprojectHash: pyprojectHash(dependencySkill.pyprojectPath), - }, - { - importName: dependentSkill.importName, - packagePath: dependentSkill.packagePath, - pyprojectPath: dependentSkill.pyprojectPath, - pyprojectHash: pyprojectHash(dependentSkill.pyprojectPath), - }, - ]); - }); - - it("installs sibling Python skill dependencies when package and directory names differ", async () => { - const logPath = installFakeUv(); - const venv = join(tempDir, "kernel-venv"); - const dependencySkill = createPythonSkill("attach-image"); - writeFileSync( - dependencySkill.pyprojectPath, - `[project] -name = "prime-agent-skill-attach-image" -version = "0.1.0" -`, - ); - const dependentSkill = createPythonSkillWithDependency( - "orchestration-heartbeat", - "prime-agent-skill-attach-image", - ); - process.env.PRIME_AGENT_KERNEL_VENV = venv; - - await expect(ensureKernelPython({ pythonSkills: [dependentSkill] })).resolves.toBe(join(venv, "bin", "python")); - - const log = readFileSync(logPath, "utf8"); - expect(log).toContain(`--editable ${dependencySkill.packagePath}`); - expect(log).toContain(`--editable ${dependentSkill.packagePath}`); - }); - - it("parses Python skill dependencies with extras", async () => { + it("keeps a skill-synced venv fast for real sessions after a no-skill bootstrap call", async () => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); - const dependencySkill = createPythonSkill("gidgethub"); - const dependentSkill = createPythonSkillWithDependency("orchestration-heartbeat", "gidgethub[httpx]>4.0.0"); - process.env.PRIME_AGENT_KERNEL_VENV = venv; - - await expect(ensureKernelPython({ pythonSkills: [dependentSkill] })).resolves.toBe(join(venv, "bin", "python")); - - const log = readFileSync(logPath, "utf8"); - expect(log).toContain(`--editable ${dependencySkill.packagePath}`); - expect(log).toContain(`--editable ${dependentSkill.packagePath}`); - }); - - it("syncs a warm venv when a Python skill pyproject changes", async () => { - const logPath = installFakeUv(); - const venv = join(tempDir, "kernel-venv"); - const python = join(venv, "bin", "python"); const pythonSkill = createPythonSkill(); - mkdirSync(join(venv, "bin"), { recursive: true }); - writeFakePython(python, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); - writeBootstrapVersion(venv, [pythonSkill]); - writeFileSync( - pythonSkill.pyprojectPath, - `[project] -name = "${pythonSkill.name}" -version = "0.1.0" -dependencies = ["httpx"] -`, - ); process.env.PRIME_AGENT_KERNEL_VENV = venv; - await expect(ensureKernelPython({ pythonSkills: [pythonSkill] })).resolves.toBe(python); + await expect(ensureKernelPython({ pythonSkills: [pythonSkill] })).resolves.toBe(join(venv, "bin", "python")); + const syncedLog = readFileSync(logPath, "utf8"); - const log = readFileSync(logPath, "utf8"); - expect(log).not.toContain(`venv ${venv} --python 3.11`); - expect(log).toContain(`--editable ${pythonSkill.packagePath}`); - const version = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); - expect(version.pythonSkills[0].pyprojectHash).toBe(pyprojectHash(pythonSkill.pyprojectPath)); + await expect(ensureKernelPython()).resolves.toBe(join(venv, "bin", "python")); + await expect(ensureKernelPython({ pythonSkills: [pythonSkill] })).resolves.toBe(join(venv, "bin", "python")); + + expect(readFileSync(logPath, "utf8")).toBe(syncedLog); }); - it("continues when a Python skill editable install fails and retries it next startup", async () => { + it("writes the base marker before the first install and keeps a failed skill out of the record", async () => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); - const goodSkill = createPythonSkill("good-skill"); - const brokenSkill = createPythonSkill("broken-skill"); + const installedSkill = createPythonSkill("agent-a"); + const brokenSkill = createPythonSkill("agent-b"); process.env.PRIME_AGENT_KERNEL_VENV = venv; process.env.UV_FAIL_ARG = brokenSkill.packagePath; - await expect(ensureKernelPython({ pythonSkills: [goodSkill, brokenSkill] })).resolves.toBe( + await expect(ensureKernelPython({ pythonSkills: [installedSkill, brokenSkill] })).resolves.toBe( join(venv, "bin", "python"), ); - - const log = readFileSync(logPath, "utf8"); - expect(log).toContain(`--editable ${goodSkill.packagePath}`); - expect(log).toContain(`--editable ${brokenSkill.packagePath}`); + const markers = readFileSync(logPath, "utf8") + .split("\n") + .filter((line) => line.startsWith("MARKER ")) + .map((probe) => probe.slice("MARKER ".length)); + expect(JSON.parse(markers[0]).pythonSkills).toEqual([]); const version = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); - expect(version.pythonSkills).toEqual([ - { - importName: goodSkill.importName, - packagePath: goodSkill.packagePath, - pyprojectPath: goodSkill.pyprojectPath, - pyprojectHash: pyprojectHash(goodSkill.pyprojectPath), - }, + expect(version.pythonSkills.map((skill: { importName: string }) => skill.importName)).toEqual([ + installedSkill.importName, ]); - - await expect(ensureKernelPython({ pythonSkills: [goodSkill, brokenSkill] })).resolves.toBe( - join(venv, "bin", "python"), - ); - - const retryLog = readFileSync(logPath, "utf8"); - expect(retryLog.split("\n").filter((line) => line.startsWith(`venv ${venv} `))).toHaveLength(1); - expect( - retryLog.split("\n").filter((line) => line.includes(`--editable ${brokenSkill.packagePath}`)), - ).toHaveLength(2); }); - it("rebuilds a warm venv with legacy unhashed Python skill manifest entries", async () => { + // test-policy: allow explicit-test-timeout -- bounds real killed tsx respawn and resume variance, not the assertion + it("resumes a real killed mid-sync process without rebuilding the venv", { timeout: 60_000 }, async () => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); const python = join(venv, "bin", "python"); - const pythonSkill = createPythonSkill(); + const staleSkill = createPythonSkill("agent-a"); + const recordedSkill = createPythonSkill("agent-b"); + const hangingSkill = createPythonSkill("agent-c"); mkdirSync(join(venv, "bin"), { recursive: true }); writeFakePython(python, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); - writeFileSync( - join(venv, ".bootstrap-version"), - `${JSON.stringify({ - schema: 4, - runtime: "prime-agent-runtime", - extraUvArgs: DEFAULT_RLM_EXTRA_UV_ARGS, - pythonSkills: [ - { - importName: pythonSkill.importName, - packagePath: pythonSkill.packagePath, - pyprojectPath: pythonSkill.pyprojectPath, - }, - ], - })}\n`, - ); + writeBootstrapVersion(venv, [staleSkill, recordedSkill]); + writeFileSync(staleSkill.pyprojectPath, `[project]\nname = "${staleSkill.name}"\nversion = "0.2.0"\n`); process.env.PRIME_AGENT_KERNEL_VENV = venv; - await expect(ensureKernelPython()).resolves.toBe(python); + const child = spawn(process.execPath, [tsxPath, kernelSyncChildPath], { + env: { + ...process.env, + TMPDIR: "/tmp", + KERNEL_SYNC_CHILD_SKILLS: JSON.stringify([staleSkill, recordedSkill, hangingSkill]), + TSX_TSCONFIG_PATH: resolve(__dirname, "../../../tsconfig.json"), + UV_HANG_ARG: hangingSkill.packagePath, + }, + stdio: "ignore", + detached: true, + }); + try { + const deadline = Date.now() + 30_000; + while (Date.now() < deadline) { + if (child.exitCode !== null || child.signalCode !== null) break; + const log = existsSync(logPath) ? readFileSync(logPath, "utf8") : ""; + if ( + log.includes(`--editable ${staleSkill.packagePath}`) && + log.includes(`--editable ${hangingSkill.packagePath}`) + ) { + break; + } + // test-policy: allow wall-clock-sleep -- polls the external fake-uv log appends; no in-process signal exists + await sleep(50); + } + } finally { + const pid = child.pid; + if (pid !== undefined) { + try { + process.kill(-pid, "SIGKILL"); + } catch { + child.kill("SIGKILL"); + } + } + await new Promise((resolve) => child.once("exit", resolve)); + } + const versionAtKill = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); + expect(versionAtKill.pythonSkills.map((skill: { importName: string }) => skill.importName)).toEqual([ + staleSkill.importName, + recordedSkill.importName, + ]); + expect(versionAtKill.pythonSkills[0].pyprojectHash).toBe(pyprojectHash(staleSkill.pyprojectPath)); + + await expect(ensureKernelPython({ pythonSkills: [staleSkill, recordedSkill, hangingSkill] })).resolves.toBe( + python, + ); - expect(readFileSync(logPath, "utf8")).toContain(`venv ${venv} --python 3.11`); + const lines = readFileSync(logPath, "utf8").split("\n"); + expect(lines.some((line) => line.startsWith(`venv ${venv} `))).toBe(false); + expect(lines.filter((line) => line.includes(`--editable ${staleSkill.packagePath}`))).toHaveLength(1); + expect(lines.filter((line) => line.includes(`--editable ${recordedSkill.packagePath}`))).toHaveLength(0); + expect(lines.filter((line) => line.includes(`--editable ${hangingSkill.packagePath}`))).toHaveLength(2); + const version = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); + expect(version.pythonSkills.map((skill: { importName: string }) => skill.importName)).toEqual([ + staleSkill.importName, + recordedSkill.importName, + hangingSkill.importName, + ]); }); it("shares concurrent bootstrap work in one process", async () => { @@ -415,143 +415,144 @@ dependencies = ["httpx"] expect(log.split("\n").filter((line) => line.startsWith(`venv ${venv} `))).toHaveLength(1); }); - it("reuses a current warm venv without invoking uv", async () => { - const venv = join(tempDir, "kernel-venv"); - const python = join(venv, "bin", "python"); - mkdirSync(join(venv, "bin"), { recursive: true }); - writeFakePython(python, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); - writeBootstrapVersion(venv); - process.env.PRIME_AGENT_KERNEL_VENV = venv; - - await expect(ensureKernelPython()).resolves.toBe(python); - }); - - it("rebuilds a warm venv whose recorded runtime hash no longer matches local source", async () => { - const logPath = installFakeUv(); - const venv = join(tempDir, "kernel-venv"); - const python = join(venv, "bin", "python"); - mkdirSync(join(venv, "bin"), { recursive: true }); - writeFakePython(python, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); - writeFileSync( - join(venv, ".bootstrap-version"), - `${JSON.stringify({ - schema: 9, - runtime: "sha256:stale", - snapshot: "dill", - extraUvArgs: DEFAULT_RLM_EXTRA_UV_ARGS, - pythonSkills: [], - })}\n`, - ); - process.env.PRIME_AGENT_KERNEL_VENV = venv; - - await expect(ensureKernelPython()).resolves.toBe(python); - - expect(readFileSync(logPath, "utf8")).toContain(`venv ${venv} --python 3.11`); - const version = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); - expect(version.runtime).toBe(runtimeIdentity); - }); - - it("rebuilds a warm venv with a stale rlm runtime", async () => { + it.each([ + { + name: "reuses a current warm venv without invoking uv", + rebuilds: false, + prepare: (venv: string, python: string) => { + writeFakePython(python, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); + writeBootstrapVersion(venv); + }, + }, + { + name: "rebuilds when the recorded runtime hash no longer matches local source", + rebuilds: true, + prepare: (venv: string, python: string) => { + writeFakePython(python, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); + writeFileSync( + join(venv, ".bootstrap-version"), + `${JSON.stringify({ + schema: 9, + runtime: "sha256:stale", + snapshot: "dill", + extraUvArgs: DEFAULT_RLM_EXTRA_UV_ARGS, + pythonSkills: [], + })}\n`, + ); + }, + }, + { + name: "rebuilds a legacy unhashed Python skill manifest", + rebuilds: true, + prepare: (venv: string, python: string) => { + const pythonSkill = createPythonSkill(); + writeFakePython(python, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); + writeFileSync( + join(venv, ".bootstrap-version"), + `${JSON.stringify({ + schema: 4, + runtime: "prime-agent-runtime", + extraUvArgs: DEFAULT_RLM_EXTRA_UV_ARGS, + pythonSkills: [ + { + importName: pythonSkill.importName, + packagePath: pythonSkill.packagePath, + pyprojectPath: pythonSkill.pyprojectPath, + }, + ], + })}\n`, + ); + }, + }, + { + name: "rebuilds a venv whose rlm runtime is stale", + rebuilds: true, + prepare: (venv: string, python: string) => { + // Imports rlm but fails the runtime-capability probe. + writeFakePython(python, ["rlm"], ["_harness_methods"]); + writeBootstrapVersion(venv); + }, + }, + { + name: "rebuilds a venv whose interpreter is missing entirely", + rebuilds: true, + prepare: (venv: string) => { + writeBootstrapVersion(venv); + }, + }, + ])("$name", async ({ rebuilds, prepare }) => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); const python = join(venv, "bin", "python"); mkdirSync(join(venv, "bin"), { recursive: true }); - writeExecutable( - python, - [ - "#!/bin/sh", - 'if [ "$1" = "-c" ]; then', - ' case "$2" in', - ' "import rlm") exit 0 ;;', - " *) exit 1 ;;", - " esac", - "fi", - "exit 0", - "", - ].join("\n"), - ); - writeBootstrapVersion(venv); + prepare(venv, python); process.env.PRIME_AGENT_KERNEL_VENV = venv; await expect(ensureKernelPython()).resolves.toBe(python); - expect(readFileSync(logPath, "utf8")).toContain(`venv ${venv} --python 3.11`); - }); - - it("rebuilds a broken venv", async () => { - const logPath = installFakeUv(); - const venv = join(tempDir, "kernel-venv"); - mkdirSync(join(venv, "bin"), { recursive: true }); - writeBootstrapVersion(venv); - process.env.PRIME_AGENT_KERNEL_VENV = venv; - - await expect(ensureKernelPython()).resolves.toBe(join(venv, "bin", "python")); - - expect(readFileSync(logPath, "utf8")).toContain(`venv ${venv} --python 3.11`); - }); - - it("uses PRIME_AGENT_KERNEL_PYTHON as an override contract", async () => { - const overridePython = join(tempDir, "override-python"); - writeFakePython(overridePython, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); - process.env.PRIME_AGENT_KERNEL_PYTHON = overridePython; - - await expect(ensureKernelPython()).resolves.toBe(overridePython); - }); - - it("allows PRIME_AGENT_KERNEL_PYTHON missing Python skill imports", async () => { - const overridePython = join(tempDir, "override-python"); - const pythonSkill = createPythonSkill(); - writeFakePython(overridePython, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); - process.env.PRIME_AGENT_KERNEL_PYTHON = overridePython; - - await expect(ensureKernelPython({ pythonSkills: [pythonSkill] })).resolves.toBe(overridePython); - }); - - it("rejects PRIME_AGENT_KERNEL_PYTHON missing default extra packages", async () => { - const overridePython = join(tempDir, "override-python"); - writeFakePython(overridePython, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES.filter((name) => name !== "yaml")]); - process.env.PRIME_AGENT_KERNEL_PYTHON = overridePython; - - await expect(ensureKernelPython()).rejects.toThrow(/default Python packages \(yaml \(PyYAML\)\)/); - }); - - it("rejects PRIME_AGENT_KERNEL_PYTHON with a stale rlm runtime", async () => { - const overridePython = join(tempDir, "override-python"); - writeFakePython(overridePython, ["dill"]); - process.env.PRIME_AGENT_KERNEL_PYTHON = overridePython; - - await expect(ensureKernelPython()).rejects.toThrow(/current prime-agent-runtime with callable rlm\.spawn/); - }); - - it("rejects PRIME_AGENT_KERNEL_PYTHON with a legacy harness API", async () => { - const overridePython = join(tempDir, "override-python"); - writeExecutable( - overridePython, - [ - "#!/bin/sh", - 'if [ "$1" = "-c" ]; then', - ' case "$2" in', - ' "import rlm") exit 0 ;;', - ' *"_harness_methods"*) exit 1 ;;', - " *\"assert not hasattr(rlm.rlm, 'background')\"*) exit 0 ;;", - " *) exit 1 ;;", - " esac", - "fi", - "exit 0", - "", - ].join("\n"), - ); - process.env.PRIME_AGENT_KERNEL_PYTHON = overridePython; - - await expect(ensureKernelPython()).rejects.toThrow(/current prime-agent-runtime with callable rlm\.spawn/); + // A warm venv that needs no work never runs uv, so the log may not exist at all. + const log = existsSync(logPath) ? readFileSync(logPath, "utf8") : ""; + expect(log.includes(`venv ${venv} --python 3.11`)).toBe(rebuilds); + if (rebuilds) { + expect(JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")).runtime).toBe(runtimeIdentity); + } }); - it("fails an invalid PRIME_AGENT_KERNEL_PYTHON without bootstrapping", async () => { + it.each([ + { + name: "accepts an interpreter carrying the runtime and every default extra", + write: (path: string) => writeFakePython(path, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]), + withSkill: false, + error: undefined, + }, + { + name: "accepts an interpreter that cannot import the Python skills", + write: (path: string) => writeFakePython(path, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]), + withSkill: true, + error: undefined, + }, + { + name: "rejects an interpreter missing a default extra package", + write: (path: string) => + writeFakePython(path, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES.filter((module) => module !== "yaml")]), + withSkill: false, + error: /default Python packages \(yaml \(PyYAML\)\)/, + }, + { + name: "rejects an interpreter with a stale rlm runtime", + write: (path: string) => writeFakePython(path, ["dill"]), + withSkill: false, + error: /current prime-agent-runtime with callable rlm\.spawn/, + }, + { + name: "rejects an interpreter exposing only the legacy harness API", + write: (path: string) => writeFakePython(path, ["rlm"], ["_harness_methods"]), + withSkill: false, + error: /current prime-agent-runtime with callable rlm\.spawn/, + }, + { + name: "rejects an interpreter with a pre-progress-note rlm runtime", + write: (path: string) => writeFakePython(path, ["rlm"], ["progress_note"]), + withSkill: false, + error: /current prime-agent-runtime with callable rlm\.spawn, rlm\.create_session, rlm\.host_request, rlm\.progress_note/, + }, + { + name: "rejects an interpreter missing the runtime, without bootstrapping a venv", + write: (path: string) => writeFakePython(path, []), + withSkill: false, + error: /PRIME_AGENT_KERNEL_PYTHON points to a Python missing/, + }, + ])("PRIME_AGENT_KERNEL_PYTHON $name", async ({ write, withSkill, error }) => { const overridePython = join(tempDir, "override-python"); - writeFakePython(overridePython, []); + write(overridePython); process.env.PRIME_AGENT_KERNEL_PYTHON = overridePython; + const options = withSkill ? { pythonSkills: [createPythonSkill()] } : {}; - await expect(ensureKernelPython()).rejects.toThrow(/PRIME_AGENT_KERNEL_PYTHON points to a Python missing/); + if (error) { + await expect(ensureKernelPython(options)).rejects.toThrow(error); + } else { + await expect(ensureKernelPython(options)).resolves.toBe(overridePython); + } }); it("resolves the venv python under Scripts\\python.exe on win32 (uv layout)", () => { From 05ee4fc635d29d767369d408fcbc781198e4551c Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Tue, 22 Sep 2026 06:29:14 -0600 Subject: [PATCH 24/26] fix(daemon): update DAEMON_SCHEMA_ID for update_restarting wire variant --- packages/coding-agent/src/modes/daemon/daemon-protocol.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/coding-agent/src/modes/daemon/daemon-protocol.ts b/packages/coding-agent/src/modes/daemon/daemon-protocol.ts index 4ad30dcfa6a..0cc674054a9 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-protocol.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-protocol.ts @@ -92,7 +92,7 @@ export const DAEMON_COMMAND_ENVELOPE_MIN_PROTOCOL_VERSION = 7; // Revision 33 preserves optional worker replay metadata on runtime snapshots under event_sequence. export const DAEMON_SCHEMA_REVISION = 33; export const DAEMON_RUNTIME_SNAPSHOT_REPLAY_MIN_SCHEMA_REVISION = 33; -export const DAEMON_SCHEMA_ID = "protocol-7-schema-33-5924c5b19b8b"; +export const DAEMON_SCHEMA_ID = "protocol-7-schema-33-df7072089924"; export type DaemonProtocolName = typeof DAEMON_PROTOCOL_NAME; export type DaemonProtocolVersion = number; From 9a96243df3ccea4b4b97024914c2f06c392dd015 Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Tue, 22 Sep 2026 06:30:23 -0600 Subject: [PATCH 25/26] test(compaction): update compaction-serialization expectations for tool labels (#2424) --- .../coding-agent/test/compaction-serialization.test.ts | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/coding-agent/test/compaction-serialization.test.ts b/packages/coding-agent/test/compaction-serialization.test.ts index 75ceea1b167..1c3423ef1d0 100644 --- a/packages/coding-agent/test/compaction-serialization.test.ts +++ b/packages/coding-agent/test/compaction-serialization.test.ts @@ -18,10 +18,10 @@ describe("serializeConversation", () => { const result = serializeConversation(messages); - expect(result).toContain("[Tool result]:"); - expect(result).toContain("[... 3000 more characters truncated]"); + expect(result).toContain("[Tool result (ipython)]:"); + expect(result).toContain("truncated"); expect(result).not.toContain("x".repeat(3000)); - expect(result).toContain("x".repeat(2000)); + expect(result).toContain("x".repeat(1000)); }); it("should not truncate short tool results", () => { @@ -39,7 +39,7 @@ describe("serializeConversation", () => { const result = serializeConversation(messages); - expect(result).toBe(`[Tool result]: ${shortContent}`); + expect(result).toBe(`[Tool result (ipython)]: ${shortContent}`); expect(result).not.toContain("truncated"); }); From 6195f701b22e46dceff8a6d6dacb9c2a42553cc6 Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Fri, 25 Sep 2026 13:20:59 -0600 Subject: [PATCH 26/26] test(ai): align GLM thinking format expectations with #2459 --- packages/ai/test/prime-inference-models.test.ts | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/packages/ai/test/prime-inference-models.test.ts b/packages/ai/test/prime-inference-models.test.ts index 58c3a03b8fe..1c26d83f969 100644 --- a/packages/ai/test/prime-inference-models.test.ts +++ b/packages/ai/test/prime-inference-models.test.ts @@ -171,16 +171,12 @@ describe("Prime Inference models", () => { }); const glm51 = getModel("prime-inference", "z-ai/glm-5.1"); expect(glm51.reasoning).toBe(true); - expect(glm51.compat).toMatchObject({ - supportsReasoningEffort: false, - thinkingFormat: "zai", - }); + expect(glm51.compat?.supportsReasoningEffort).toBe(false); + expect(glm51.compat?.thinkingFormat).toBeUndefined(); const glm52 = getModel("prime-inference", "z-ai/glm-5.2"); expect(glm52.reasoning).toBe(true); - expect(glm52.compat).toMatchObject({ - supportsReasoningEffort: false, - thinkingFormat: "zai", - }); + expect(glm52.compat?.supportsReasoningEffort).toBe(false); + expect(glm52.compat?.thinkingFormat).toBeUndefined(); expect(getModel("prime-inference", "qwen/qwen3-coder-next").reasoning).toBe(false); expect(getModel("prime-inference", "x-ai/grok-4.20").reasoning).toBe(true); expect(getModel("prime-inference", "minimax/minimax-m3").reasoning).toBe(true);