@bastani/atomic 0.9.11-alpha.8 → 0.9.11

package/dist/modes/rpc/rpc-client.js

@@ -1,57 +1,17 @@
-import { QueuedWriter } from "./queued-writer.js";
-import { attachJsonlLineReader, serializeJsonLine } from "./jsonl.js";
-import { InteractiveEngineMonitor } from "../interactive-engine/engine-monitor.js";
-import { serializeInteractiveEngineFrame } from "../interactive-engine/protocol.js";
 import { sleep } from "../../utils/sleep.js";
-import { createInteractiveJsonlOptions, spawnRpcClientProcess, terminateRpcClientProcess } from "./rpc-client-process.js";
+import { InteractiveEngineMonitor } from "../interactive-engine/engine-monitor.js";
+import { serializeInteractiveEngineFrame, } from "../interactive-engine/protocol.js";
+import { attachJsonlLineReader, serializeJsonLine } from "./jsonl.js";
+import { QueuedWriter } from "./queued-writer.js";
 import { RpcClientApi } from "./rpc-client-api.js";
+import { appendBoundedStderr, createInteractiveJsonlOptions, restartCliArgs, spawnRpcClientProcess, terminateRpcClientProcess, } from "./rpc-client-process.js";
+import { collectRpcEvents, runUserBashWithUpdates, waitForRpcIdle } from "./rpc-client-waits.js";
+import { DEFAULT_REQUEST_TIMEOUT_MS, LONG_LIVED_COMMANDS, RESTART_CANCELLED_MESSAGE } from "./rpc-command-timeouts.js";
 import { RpcEventBuffer } from "./rpc-event-buffer.js";
-import { collectRpcEvents, waitForRpcIdle } from "./rpc-client-waits.js";
-function restartArgs(args, sessionFile) {
-    const result = [];
-    for (let index = 0; index < (args?.length ?? 0); index += 1) {
-        const value = args[index];
-        if (value === "--no-session")
-            continue;
-        if (value === "--session") {
-            index += 1;
-            continue;
-        }
-        result.push(value);
-    }
-    result.push(sessionFile ? "--session" : "--no-session");
-    if (sessionFile)
-        result.push(sessionFile);
-    return result;
-}
-/**
- * Commands whose response is legitimately gated on human/agent interaction or a
- * turn boundary and therefore must not be subject to the generic request
- * timeout: interactive prompts and custom-UI pickers (routed through `prompt`),
- * queued turn-boundary sends (`steer`/`follow_up`), long-running shell and
- * compaction work, and session-tree navigation/mutation that can open pickers.
- *
- * Failure detection is unaffected: engine exit, transport violations, aborts,
- * generation replacement, and explicit stop all reject pending requests via
- * rejectPendingRequests/failTransport independently of any timer.
- */
-const LONG_LIVED_COMMANDS = new Set([
-    "prompt",
-    "login_provider",
-    "steer",
-    "follow_up",
-    "bash",
-    "user_bash",
-    "compact",
-    "fork",
-    "clone",
-    "switch_session",
-    "new_session",
-    "import_session",
-    "navigate_tree",
-    "invoke_shortcut",
-]);
-const DEFAULT_REQUEST_TIMEOUT_MS = 30_000;
+import { GenerationBuffer } from "./rpc-generation-buffer.js";
+import { RpcPendingRequests } from "./rpc-pending-requests.js";
+import { RpcTerminalDrain } from "./rpc-terminal-drain.js";
+import { markRpcTransportFailure, rpcTransportError } from "./rpc-transport-error.js";
 export class RpcClient extends RpcClientApi {
     constructor(options = {}) {
         super();
@@ -59,14 +19,20 @@ export class RpcClient extends RpcClientApi {
         this.stopReadingStdout = null;
         this.eventListeners = [];
         this.extensionUIListeners = [];
-        this.pendingExtensionUIRequests = [];
+        this.pendingExtensionUIRequests = new GenerationBuffer();
         this.engineMessageListeners = [];
-        this.pendingEngineMessages = [];
-        this.pendingRequests = new Map();
+        this.pendingEngineMessages = new GenerationBuffer();
+        this.pendingRequests = new RpcPendingRequests();
+        /** Bumped by every explicit stop; a restart holds a permit across its own stop. */
+        this.restartRevision = 0;
+        this.terminalDrain = new RpcTerminalDrain();
+        this.stdoutDrained = Promise.resolve();
         this.requestId = 0;
         this.stderr = "";
         this.exitError = null;
         this.generation = 0;
+        this.lastEndedGeneration = 0;
+        this.generationEndedListeners = [];
         this.activeActivityIds = new Set();
         this.options = options;
     }
@@ -91,7 +57,7 @@ export class RpcClient extends RpcClientApi {
         this.activeActivityIds.clear();
         this.enginePid = undefined;
         this.engineMonitor = this.options.interactiveEngine
-            ? new InteractiveEngineMonitor(this.options.interactiveEngine.onDiagnostic, (message) => this.observeInteractiveEngineMessage(message))
+            ? new InteractiveEngineMonitor(this.options.interactiveEngine.onDiagnostic, (message) => this.observeInteractiveEngineMessage(message, generation))
             : undefined;
         this.eventBuffer = this.engineMonitor ? new RpcEventBuffer((event) => this.emitEvent(event)) : undefined;
         const childProcess = spawnRpcClientProcess({
@@ -102,74 +68,81 @@ export class RpcClient extends RpcClientApi {
             runtimeExecutable: this.options.runtimeExecutable,
             runtimeArgs: this.options.runtimeArgs,
             interactiveEngine: this.engineMonitor !== undefined,
+            ...(this.options.interactiveEngine?.apiKey
+                ? { interactiveEngineApiKey: this.options.interactiveEngine.apiKey }
+                : {}),
         });
         this.process = childProcess;
         this.stdinWriter = new QueuedWriter(childProcess.stdin);
-        childProcess.once("exit", (code, signal) => {
-            if (generation !== this.generation)
-                return;
-            const error = this.createProcessExitError(code, signal);
-            this.exitError = error;
-            this.stdinWriter?.close(error);
-            this.engineMonitor?.fail(error);
-            this.activeActivityIds.clear();
-            this.options.interactiveEngine?.onActivityChange?.(false);
-            this.rejectPendingRequests(error);
-        });
-        childProcess.once("error", (error) => {
-            if (generation !== this.generation)
-                return;
-            this.exitError = error;
-            this.stdinWriter?.close(error);
-            this.engineMonitor?.fail(error);
-            this.rejectPendingRequests(error);
-        });
-        childProcess.stdin?.on("error", (error) => {
-            if (generation !== this.generation)
-                return;
-            this.exitError = error;
-            this.stdinWriter?.close(error);
-            this.engineMonitor?.fail(error);
-            this.rejectPendingRequests(error);
-        });
+        childProcess.once("exit", (code, signal) => this.failGeneration(this.createProcessExitError(code, signal), generation, "exit"));
+        childProcess.once("error", (rawError) => this.failGeneration(rawError, generation, "process-error"));
+        childProcess.stdin?.on("error", (rawError) => this.failGeneration(rawError, generation, "stdin-error"));
         childProcess.stderr?.on("data", (data) => {
             if (generation !== this.generation)
                 return;
-            const next = this.stderr + data.toString();
-            this.stderr = Buffer.byteLength(next, "utf8") <= 256 * 1024
-                ? next
-                : `${Buffer.from(next).subarray(-256 * 1024).toString("utf8")}\n[stderr truncated]`;
+            this.stderr = appendBoundedStderr(this.stderr, data.toString());
             process.stderr.write(data);
         });
         const readerOptions = createInteractiveJsonlOptions(this.engineMonitor !== undefined);
+        let markStdoutDrained;
+        this.stdoutDrained = new Promise((resolve) => {
+            markStdoutDrained = resolve;
+        });
         this.stopReadingStdout = attachJsonlLineReader(childProcess.stdout, (line) => {
             if (generation === this.generation)
-                this.handleLine(line);
-        }, readerOptions);
+                this.handleLine(line, generation);
+        }, { ...readerOptions, onDrained: () => markStdoutDrained() });
         if (this.engineMonitor)
             await this.engineMonitor.waitUntilReady();
         else
             await sleep(100);
         if (generation !== this.generation || this.process?.exitCode !== null) {
-            throw new Error(`Agent process exited immediately. Stderr: ${this.stderr}`);
+            throw rpcTransportError(`Agent process exited immediately. Stderr: ${this.stderr}`);
         }
     }
+    /**
+     * Explicit stop. Voids any in-flight restart permit, so a replacement caught
+     * between its own stop and start can never spawn after this returns.
+     */
     async stop() {
+        this.restartRevision += 1;
+        await this.stopCurrentGeneration();
+    }
+    /**
+     * Tear down the current child without touching the restart permit. The reader
+     * stays attached and the generation stays current until the dead child's
+     * stdout settles, so an ownership frame it already sent still reaches its
+     * request; only then is the generation retired (see RpcTerminalDrain).
+     */
+    async stopCurrentGeneration() {
         const child = this.process;
         if (!child)
             return;
+        const endedGeneration = this.generation;
+        const stopError = rpcTransportError("Agent process stopped");
         this.invalidateEngineShortcuts();
         const terminateTree = this.engineMonitor !== undefined;
-        this.stopReadingStdout?.();
         this.engineMonitor?.stop();
-        this.stdinWriter?.close(new Error("Agent process stopped"));
+        this.stdinWriter?.close(stopError);
+        // Unblock an in-flight start()/waitUntilBound() for the generation being
+        // replaced. Readiness has no deadline by design, so without this an
+        // explicit stop of a hung replacement child could never be observed.
+        this.engineMonitor?.fail(stopError);
         this.engineMonitor = undefined;
-        this.stopReadingStdout = null;
         this.stdinWriter = undefined;
         this.process = null;
-        this.generation += 1;
+        // Publish before terminating so host controllers tear down this generation's
+        // UI while the writer is detached: disposal cannot write into a replacement.
+        this.publishGenerationEnded(endedGeneration, stopError, "explicit-stop", true);
+        // Claim before terminating so the exit this provokes does not relabel a
+        // deliberate stop; an exit that already claimed keeps its own error.
+        this.terminalDrain.claim(endedGeneration, stopError);
         await terminateRpcClientProcess(child, terminateTree);
-        this.rejectPendingRequests(new Error("Agent process stopped"));
+        await this.beginTerminalDrain(endedGeneration, stopError);
+        this.stopReadingStdout?.();
+        this.stopReadingStdout = null;
+        this.generation += 1;
+        this.pendingRequests.rejectAll(stopError);
         this.eventBuffer?.dispose();
         this.eventBuffer = undefined;
         this.activeActivityIds.clear();
@@ -187,7 +160,7 @@ export class RpcClient extends RpcClientApi {
     }
     onExtensionUIRequest(listener) {
         this.extensionUIListeners.push(listener);
-        for (const request of this.pendingExtensionUIRequests.splice(0))
+        for (const request of this.pendingExtensionUIRequests.drain(this.generation))
             listener(request);
         return () => {
             const index = this.extensionUIListeners.indexOf(listener);
@@ -200,7 +173,7 @@ export class RpcClient extends RpcClientApi {
     }
     onInteractiveEngineMessage(listener) {
         this.engineMessageListeners.push(listener);
-        for (const message of this.pendingEngineMessages.splice(0))
+        for (const message of this.pendingEngineMessages.drain(this.generation))
             listener(message);
         return () => {
             const index = this.engineMessageListeners.indexOf(listener);
@@ -222,41 +195,67 @@ export class RpcClient extends RpcClientApi {
                 this.engineMessageListeners.splice(index, 1);
         };
     }
+    /**
+     * Subscribe to host-local generation death. Fires at most once per
+     * generation, synchronously, before any restart work. A death that has
+     * already happened for the current generation replays immediately, so a
+     * listener attached after an exit still sees it.
+     */
+    onGenerationEnded(listener) {
+        this.generationEndedListeners.push(listener);
+        const retained = this.retainedGenerationEnd;
+        // Only while it still describes the live generation: once stop()/start()
+        // has moved on, that death belongs to a replaced child.
+        if (retained && retained.generation === this.generation)
+            listener(retained);
+        return () => {
+            const index = this.generationEndedListeners.indexOf(listener);
+            if (index !== -1)
+                this.generationEndedListeners.splice(index, 1);
+        };
+    }
     sendInteractiveEngineCommand(command) {
         const writer = this.stdinWriter;
         if (!writer)
             return;
         const frame = serializeInteractiveEngineFrame(command);
-        if (command.type === "engine_custom_render" || command.type === "engine_tool_render" || command.type === "engine_message_render") {
+        if (command.type === "engine_custom_render" ||
+            command.type === "engine_tool_render" ||
+            command.type === "engine_message_render") {
             writer.offerLatest(`render:${command.componentId}`, frame);
             return;
         }
         this.bestEffort(writer.write(frame), `engine command ${command.type}`);
     }
-    waitForInteractiveEngineBound() { return this.engineMonitor?.waitUntilBound() ?? Promise.resolve(); }
-    getEnginePid() { return this.enginePid; }
-    getGeneration() { return this.generation; }
+    waitForInteractiveEngineBound() {
+        return this.engineMonitor?.waitUntilBound() ?? Promise.resolve();
+    }
+    getEnginePid() {
+        return this.enginePid;
+    }
+    getGeneration() {
+        return this.generation;
+    }
+    /**
+     * Replace the engine child, unless an explicit stop supersedes this attempt.
+     * The permit is checked immediately before spawning, so a `stop()` in that
+     * window is always observed; cancellation throws rather than returning
+     * quietly, because callers go on to initialize against the new child.
+     */
     async restart(sessionFile) {
-        await this.stop();
-        this.options = { ...this.options, args: restartArgs(this.options.args, sessionFile) };
+        const permit = this.restartRevision;
+        await this.stopCurrentGeneration();
+        if (permit !== this.restartRevision)
+            throw rpcTransportError(RESTART_CANCELLED_MESSAGE);
+        this.options = { ...this.options, args: restartCliArgs(this.options.args, sessionFile) };
         await this.start();
         await this.waitForInteractiveEngineBound();
     }
     async requestInternal(command) {
         return this.data(await this.request(command));
     }
-    async userBashWithUpdates(command, onUpdate, options) {
-        let requestId;
-        const unsubscribe = this.onEvent((event) => {
-            if (event.type === "bash_execution_update" && event.id === requestId)
-                onUpdate(event.delta, event.channel);
-        });
-        try {
-            return this.data(await this.request({ type: "user_bash", command, excludeFromContext: options?.excludeFromContext }, (id) => { requestId = id; options?.onRequestId?.(id); }));
-        }
-        finally {
-            unsubscribe();
-        }
+    userBashWithUpdates(command, onUpdate, options) {
+        return runUserBashWithUpdates(this, command, onUpdate, options, (body, onId) => this.request(body, onId));
     }
     getStderr() {
         return this.stderr;
@@ -272,7 +271,9 @@ export class RpcClient extends RpcClientApi {
         await this.prompt(message, images);
         return eventsPromise;
     }
-    handleLine(line) {
+    handleLine(line, generation) {
+        if (this.terminalDrain.observe(generation, line, (id) => this.pendingRequests.markAccepted(id)))
+            return;
         if (this.engineMonitor?.handleLine(line))
             return;
         try {
@@ -280,17 +281,13 @@ export class RpcClient extends RpcClientApi {
             if (data.type === "extension_ui_request") {
                 const request = data;
                 if (this.extensionUIListeners.length === 0)
-                    this.pendingExtensionUIRequests.push(request);
+                    this.pendingExtensionUIRequests.push(generation, request);
                 for (const listener of this.extensionUIListeners)
                     listener(request);
                 return;
             }
-            if (data.type === "response" && data.id && this.pendingRequests.has(data.id)) {
-                const pending = this.pendingRequests.get(data.id);
-                this.pendingRequests.delete(data.id);
-                pending.resolve(data);
+            if (data.type === "response" && data.id && this.pendingRequests.resolve(data.id, data))
                 return;
-            }
             const event = data;
             if (this.eventBuffer)
                 this.eventBuffer.enqueue(event);
@@ -299,31 +296,44 @@ export class RpcClient extends RpcClientApi {
         }
         catch {
             if (this.engineMonitor)
-                this.failTransport(new Error("Interactive engine emitted malformed JSONL"));
+                this.failTransport(rpcTransportError("Interactive engine emitted malformed JSONL"));
         }
     }
     emitEvent(event) {
         for (const listener of this.eventListeners)
             listener(event);
     }
-    emitInteractiveEngineMessage(message) {
+    emitInteractiveEngineMessage(message, generation) {
         if (this.engineMessageListeners.length === 0 && message.type.startsWith("engine_custom_")) {
-            this.pendingEngineMessages.push(message);
+            this.pendingEngineMessages.push(generation, message);
         }
         for (const listener of this.engineMessageListeners)
             listener(message);
     }
     createProcessExitError(code, signal) {
-        return new Error(`Agent process exited (code=${code} signal=${signal}). Stderr: ${this.stderr}`);
+        return rpcTransportError(`Agent process exited (code=${code} signal=${signal}). Stderr: ${this.stderr}`);
     }
     invalidateEngineShortcuts() {
         if (!this.latestEngineKeybindingState || this.latestEngineKeybindingState.shortcuts.length === 0)
             return;
         const state = { ...this.latestEngineKeybindingState, shortcuts: [] };
         this.latestEngineKeybindingState = state;
-        this.emitInteractiveEngineMessage({ type: "engine_keybindings_reloaded", state });
+        this.emitInteractiveEngineMessage({ type: "engine_keybindings_reloaded", state }, this.generation);
     }
-    observeInteractiveEngineMessage(message) {
+    observeInteractiveEngineMessage(message, generation) {
+        // Ownership is the one thing a dead generation may still report: request
+        // ids never repeat, and the alternative is classifying work it already
+        // started as never sent. Nothing else from it may touch host state — a
+        // stale frame would revive a PID, an activity, or a mount that is gone.
+        if (message.type === "engine_request_accepted") {
+            if (generation === this.generation)
+                this.pendingRequests.markAccepted(message.requestId);
+            return;
+        }
+        if (generation !== this.generation || generation <= this.lastEndedGeneration)
+            return;
+        if (message.type === "engine_heartbeat")
+            this.options.interactiveEngine?.onHeartbeat?.();
         if (message.type === "engine_ready")
             this.enginePid = message.pid;
         if (message.type === "engine_keybindings_reloaded")
@@ -333,25 +343,59 @@ export class RpcClient extends RpcClientApi {
         else if (message.type === "engine_activity_finished")
             this.activeActivityIds.delete(message.activityId);
         this.options.interactiveEngine?.onActivityChange?.(this.activeActivityIds.size > 0);
-        this.emitInteractiveEngineMessage(message);
-    }
-    rejectPendingRequests(error) {
-        for (const { reject } of this.pendingRequests.values())
-            reject(error);
-        this.pendingRequests.clear();
+        this.emitInteractiveEngineMessage(message, generation);
     }
-    failTransport(error) {
+    /**
+     * Shared teardown for every unexpected end of `generation`. Death publishes
+     * at once for the host UI; its requests wait for stdout (RpcTerminalDrain).
+     */
+    failGeneration(rawError, generation, kind) {
+        if (generation !== this.generation)
+            return;
+        const error = markRpcTransportFailure(rawError);
         this.exitError = error;
         this.stdinWriter?.close(error);
         this.engineMonitor?.fail(error);
-        this.rejectPendingRequests(error);
+        if (kind === "exit") {
+            this.activeActivityIds.clear();
+            this.options.interactiveEngine?.onActivityChange?.(false);
+        }
+        this.beginTerminalDrain(generation, error);
+        this.publishGenerationEnded(generation, error, kind, false);
+    }
+    /** Hold this generation's requests until its stdout drains (RpcTerminalDrain). */
+    beginTerminalDrain(generation, error) {
+        return this.terminalDrain.begin(generation, this.stdoutDrained, error, (settled) => {
+            if (generation === this.generation)
+                this.pendingRequests.rejectAll(settled);
+        });
+    }
+    publishGenerationEnded(generation, error, kind, expected) {
+        if (generation <= this.lastEndedGeneration)
+            return;
+        this.lastEndedGeneration = generation;
+        const event = { generation, error, kind, expected };
+        this.retainedGenerationEnd = event;
+        // Death is a strict boundary: drop this generation's buffered frames
+        // BEFORE teardown runs, so no listener can drain a stale mount frame
+        // while it is closing that very generation's components.
+        this.pendingEngineMessages.dropGeneration(generation);
+        this.pendingExtensionUIRequests.dropGeneration(generation);
+        for (const listener of [...this.generationEndedListeners])
+            listener(event);
+    }
+    failTransport(rawError) {
         this.options.interactiveEngine?.onDiagnostic({
-            activity: undefined, elapsedMs: 0, level: "unresponsive", message: error.message,
+            activity: undefined,
+            elapsedMs: 0,
+            level: "unresponsive",
+            message: rawError.message,
         });
+        this.failGeneration(rawError, this.generation, "transport-error");
     }
     requireWriter() {
         if (!this.stdinWriter)
-            throw new Error("Agent process stdin is not writable");
+            throw rpcTransportError("Agent process stdin is not writable");
         return this.stdinWriter;
     }
     bestEffort(operation, label) {
@@ -359,7 +403,9 @@ export class RpcClient extends RpcClientApi {
             if (!this.process || this.exitError)
                 return;
             this.options.interactiveEngine?.onDiagnostic({
-                activity: undefined, elapsedMs: 0, level: "unresponsive",
+                activity: undefined,
+                elapsedMs: 0,
+                level: "unresponsive",
                 message: `Interactive engine ${label} failed: ${error.message}`,
             });
         });
@@ -367,7 +413,7 @@ export class RpcClient extends RpcClientApi {
     async request(command, onRequestId) {
         const childProcess = this.process;
         if (!childProcess)
-            throw new Error("Client not started");
+            throw rpcTransportError("Client not started");
         if (this.exitError)
             throw this.exitError;
         if (childProcess.exitCode !== null) {
@@ -382,11 +428,17 @@ export class RpcClient extends RpcClientApi {
         let rejectResponse;
         const response = new Promise((resolve, reject) => {
             rejectResponse = reject;
-            this.pendingRequests.set(id, {
-                resolve: (value) => { if (timeout)
-                    clearTimeout(timeout); resolve(value); },
-                reject: (error) => { if (timeout)
-                    clearTimeout(timeout); reject(error); },
+            this.pendingRequests.add(id, {
+                resolve: (value) => {
+                    if (timeout)
+                        clearTimeout(timeout);
+                    resolve(value);
+                },
+                reject: (error) => {
+                    if (timeout)
+                        clearTimeout(timeout);
+                    reject(error);
+                },
             });
         });
         try {
@@ -394,7 +446,10 @@ export class RpcClient extends RpcClientApi {
         }
         catch (error) {
             this.pendingRequests.delete(id);
-            rejectResponse(error instanceof Error ? error : new Error(String(error)));
+            // Idempotent: the writer already classified this, but a caller-supplied
+            // stream can reject with anything, and a send that never left the host
+            // is a transport failure whatever it threw.
+            rejectResponse(markRpcTransportFailure(error));
             return response;
         }
         if (this.pendingRequests.has(id) && !LONG_LIVED_COMMANDS.has(command.type)) {