@gajae-code/coding-agent 0.5.0 → 0.5.2

package/src/modes/rpc/rpc-mode.ts

@@ -10,8 +10,10 @@
  * - Events: AgentSessionEvent objects streamed as they occur
  * - Extension UI: Extension UI requests are emitted, client responds with extension_ui_response
  */
+
+import * as fs from "node:fs/promises";
 import * as path from "node:path";
-import { $env, readLines, Snowflake } from "@gajae-code/utils";
+import { $pickenv, logger, readLines, Snowflake } from "@gajae-code/utils";
 import type {
 	ExtensionUIContext,
 	ExtensionUIDialogOptions,
@@ -23,8 +25,9 @@ import { initializeExtensions } from "../runtime-init";
 import { dispatchRpcCommand } from "../shared/agent-wire/command-dispatch";
 import { AgentWireFrameSequencer, toAgentWireEventFrame } from "../shared/agent-wire/event-envelope";
 import { rpcError as error } from "../shared/agent-wire/responses";
+import { registerRpcSession, unregisterRpcSession } from "../shared/agent-wire/session-registry";
 import { defaultAuditPath, UnattendedAuditLog } from "../shared/agent-wire/unattended-audit";
-import { UnattendedSessionControlPlane } from "../shared/agent-wire/unattended-session";
+import { modelSupportsTokenCostMetrics, UnattendedSessionControlPlane } from "../shared/agent-wire/unattended-session";
 import { FileGateStore } from "../shared/agent-wire/workflow-gate-broker";
 import { isRpcHostToolResult, isRpcHostToolUpdate, RpcHostToolBridge } from "./host-tools";
 import { isRpcHostUriResult, RpcHostUriBridge } from "./host-uris";
@@ -70,13 +73,91 @@ function parseValueDialogResponse(
 	return undefined;
 }
 
-function shouldEmitRpcTitles(): boolean {
-	const raw = $env.PI_RPC_EMIT_TITLE;
+export function shouldEmitRpcTitlesForTest(): boolean {
+	const raw = $pickenv("GJC_RPC_EMIT_TITLE", "PI_RPC_EMIT_TITLE");
 	if (!raw) return false;
 	const normalized = raw.trim().toLowerCase();
 	return normalized === "1" || normalized === "true" || normalized === "yes" || normalized === "on";
 }
 
+const shouldEmitRpcTitles = shouldEmitRpcTitlesForTest;
+
+/**
+ * Cancellation commands bypass the ordered serial chain because they must
+ * interrupt in-flight work — they cannot wait behind the very command they are
+ * meant to abort.
+ */
+export const RPC_CANCELLATION_COMMANDS: ReadonlySet<RpcCommand["type"]> = new Set<RpcCommand["type"]>([
+	"abort",
+	"abort_bash",
+	"abort_retry",
+]);
+
+/**
+ * Safe read-only commands that bypass the ordered serial chain so they never
+ * head-of-line-block behind a long-running ordered command like
+ * `bash`/`compact`/`handoff`/`login` (#606, issue 13 — the partial fix only
+ * fast-laned cancellation).
+ *
+ * Every command listed here has a dispatch handler that is **fully synchronous
+ * and side-effect-free**: on the single-threaded event loop it runs to
+ * completion between the await points of any in-flight ordered command, reading
+ * live state without mutating it. Because such a read performs no causal write,
+ * jumping ahead of an earlier *queued* ordered command is observably harmless —
+ * there is no state change to reorder. Read payloads are additionally
+ * snapshotted inside the handler (e.g. `get_messages` returns a shallow copy of
+ * `session.messages`) so a fast-lane read can never serialize a half-mutated
+ * array that an ordered turn/compaction is rewriting in place.
+ *
+ * Deliberately excluded (kept ordered): every async/long command and every
+ * mutating command. In particular the control-flag setters (`set_thinking_level`,
+ * `cycle_thinking_level`, `set_steering_mode`, `set_follow_up_mode`,
+ * `set_interrupt_mode`, `set_auto_compaction`, `set_auto_retry`) stay ordered.
+ * Their handlers are synchronous, so fast-laning one ahead of an already-queued
+ * `prompt`/`bash` would apply the new mode *before* that earlier command runs —
+ * the earlier command would then observe the later setter's value, a
+ * causal-order (arrival-order) regression. Mutations therefore stay on the
+ * chain, and new command types default to ordered (fail-safe).
+ */
+export const RPC_SAFE_READ_CONTROL_COMMANDS: ReadonlySet<RpcCommand["type"]> = new Set<RpcCommand["type"]>([
+	// Pure synchronous reads — snapshot live state at processing time, never mutate.
+	"get_state",
+	"get_session_stats",
+	"get_available_models",
+	"get_branch_messages",
+	"get_last_assistant_text",
+	"get_messages",
+	"get_login_providers",
+]);
+
+/** True when a command may bypass the ordered serial chain and run immediately. */
+export function isFastLaneRpcCommand(type: RpcCommand["type"]): boolean {
+	return RPC_CANCELLATION_COMMANDS.has(type) || RPC_SAFE_READ_CONTROL_COMMANDS.has(type);
+}
+
+/**
+ * Schedules inbound RPC commands: fast-lane commands run immediately while
+ * everything else runs through a serial chain so causal order is preserved. The
+ * read loop never blocks, which is what lets a fast-lane command reach a
+ * long-running ordered command instead of being head-of-line-blocked behind it.
+ */
+export function createRpcCommandScheduler(
+	run: (command: RpcCommand) => Promise<void>,
+	track: (task: Promise<void>) => void,
+): { dispatch: (command: RpcCommand) => void } {
+	let orderedChain: Promise<void> = Promise.resolve();
+	return {
+		dispatch(command: RpcCommand): void {
+			if (isFastLaneRpcCommand(command.type)) {
+				track(run(command));
+				return;
+			}
+			orderedChain = orderedChain.then(() => run(command));
+			track(orderedChain);
+		},
+	};
+}
+
 function auditOutcomeFor(event: string): "accepted" | "rejected" | "denied" | "exceeded" | "aborted" | "info" {
 	if (event.includes("denied")) return "denied";
 	if (event.includes("exceeded")) return "exceeded";
@@ -86,6 +167,43 @@ function auditOutcomeFor(event: string): "accepted" | "rejected" | "denied" | "e
 	return "info";
 }
 
+export class RpcListenRefusedError extends Error {
+	constructor(socketPath: string) {
+		super(
+			`RPC --listen refused: a live server is already listening on ${socketPath}. ` +
+				"Stop it first or choose a different --listen path.",
+		);
+		this.name = "RpcListenRefusedError";
+	}
+}
+
+/**
+ * Probe whether a unix-domain socket path has a live server accepting
+ * connections. Returns `true` when a connection succeeds (a previous owner is
+ * still alive), and returns `false` only for known missing/stale endpoints
+ * (ENOENT / ECONNREFUSED). Unexpected probe failures fail closed as "alive" so
+ * `--listen` startup refuses to unlink a path it could not safely classify.
+ */
+export async function isUnixSocketAlive(socketPath: string): Promise<boolean> {
+	try {
+		const socket = await Bun.connect({
+			unix: socketPath,
+			socket: { data() {}, open() {}, error() {}, close() {} },
+		});
+		socket.end();
+		return true;
+	} catch (err) {
+		const code = err && typeof err === "object" ? (err as { code?: unknown }).code : undefined;
+		if (code === "ENOENT" || code === "ECONNREFUSED") return false;
+		logger.warn("RPC --listen socket probe failed closed", {
+			socketPath,
+			code: typeof code === "string" ? code : undefined,
+			error: err instanceof Error ? err.message : String(err),
+		});
+		return true;
+	}
+}
+
 export function requestRpcEditor(
 	pendingRequests: Map<string, PendingExtensionRequest>,
 	output: RpcOutput,
@@ -156,6 +274,7 @@ export function requestRpcEditor(
 export async function runRpcMode(
 	session: AgentSession,
 	setToolUIContext?: (uiContext: ExtensionUIContext, hasUI: boolean) => void,
+	options?: { listen?: string },
 ): Promise<never> {
 	// Signal to RPC clients that the server is ready to accept commands
 	// Suppress terminal notifications: they write \x07 (BEL) or OSC sequences directly to
@@ -164,10 +283,18 @@ export async function runRpcMode(
 	// may write there.
 	process.env.PI_NOTIFICATIONS = "off";
 
-	process.stdout.write(`${JSON.stringify({ type: "ready" })}\n`);
+	// Frames go to a swappable sink: stdout for stdio, the active client socket for a
+	// persistent --listen (UDS) server. Defaults to stdout, so the stdio path is unchanged.
+	let frameSink = (line: string): void => {
+		process.stdout.write(line);
+	};
 	const output = (obj: RpcResponse | RpcExtensionUIRequest | object) => {
-		process.stdout.write(`${JSON.stringify(obj)}\n`);
+		frameSink(`${JSON.stringify(obj)}\n`);
 	};
+	// stdio announces readiness immediately; the UDS server announces it per client connection.
+	if (!options?.listen) {
+		output({ type: "ready" });
+	}
 	const emitRpcTitles = shouldEmitRpcTitles();
 	const decodeError = (err: unknown): string => (err instanceof Error ? err.message : String(err));
 
@@ -216,6 +343,7 @@ export async function runRpcMode(
 		emitFrame: gate => output(gate),
 		store: gateStore,
 		audit: recordAudit,
+		providerSupportsTokenCostMetrics: modelSupportsTokenCostMetrics(session.model),
 		getUsageSnapshot: () => {
 			const stats = session.getSessionStats();
 			return { tokens: stats.tokens.total, costUsd: stats.cost };
@@ -231,11 +359,20 @@ export async function runRpcMode(
 	// Shutdown request flag (wrapped in object to allow mutation with const)
 	const shutdownState = { requested: false };
 	let shutdownStarted = false;
+	// Tracks in-flight non-blocking command handlers so shutdown can drain them.
+	const inFlightCommands = new Set<Promise<void>>();
 	async function shutdown(exitCode: number, reason: string): Promise<never> {
 		if (shutdownStarted) {
 			process.exit(exitCode);
 		}
 		shutdownStarted = true;
+		// Let in-flight non-blocking commands (bash/compact/handoff) finish and emit
+		// their responses before teardown, bounded so a never-resolving login cannot
+		// wedge shutdown (issue 13).
+		if (inFlightCommands.size > 0) {
+			await Promise.race([Promise.allSettled([...inFlightCommands]), Bun.sleep(5000)]);
+		}
+		await unregisterRpcSession(session.sessionId).catch(() => {});
 		hostToolBridge.rejectAllPending(`${reason} before host tool execution completed`);
 		hostUriBridge.clear(`${reason} before host URI request completed`);
 		try {
@@ -399,7 +536,7 @@ export async function runRpcMode(
 		}
 
 		setTitle(title: string): void {
-			// Title updates are low-value noise for most RPC hosts; opt in via PI_RPC_EMIT_TITLE=1.
+			// Title updates are low-value noise for most RPC hosts; opt in via GJC_RPC_EMIT_TITLE=1.
 			if (!emitRpcTitles) return;
 			this.output({
 				type: "extension_ui_request",
@@ -514,6 +651,26 @@ export async function runRpcMode(
 			unattendedControlPlane,
 		});
 
+	// Fast-lane commands (cancellation + safe read/control, see
+	// isFastLaneRpcCommand) bypass the ordered serial chain and run immediately;
+	// everything else runs through a serial chain so causal order is preserved
+	// (e.g. an ordered `set_model` after `bash` still applies after the bash
+	// result) while the read loop itself never blocks — that is what lets a
+	// fast-lane command reach a long-running `bash`/`compact`/`handoff`/`login`
+	// instead of being head-of-line-blocked behind it (issue 13).
+	const runCommand = async (command: RpcCommand): Promise<void> => {
+		try {
+			output(await handleCommand(command));
+		} catch (err) {
+			output(error(command.id, command.type, decodeError(err)));
+		}
+	};
+	const trackCommand = (task: Promise<void>): void => {
+		inFlightCommands.add(task);
+		void task.finally(() => inFlightCommands.delete(task));
+	};
+	const { dispatch: dispatchCommand } = createRpcCommandScheduler(runCommand, trackCommand);
+
 	/**
 	 * Check if shutdown was requested and perform shutdown if so.
 	 * Called after handling each command when waiting for the next command.
@@ -523,59 +680,134 @@ export async function runRpcMode(
 		await shutdown(0, "RPC shutdown requested");
 	}
 
-	// Listen for JSONL input using Bun's stdin. Parse frame-by-frame so a malformed
-	// command reports a parse error without poisoning the whole long-lived RPC session.
+	// Parse + route a single inbound JSONL frame. Shared by the stdio reader and the
+	// persistent UDS server so both transports use the same command surface.
 	const inputDecoder = new TextDecoder("utf-8", { fatal: false });
-	for await (const line of readLines(Bun.stdin.stream())) {
-		const text = inputDecoder.decode(line).trim();
-		if (!text) continue;
-
+	async function handleInboundLine(text: string): Promise<void> {
 		let parsed: unknown;
 		try {
 			parsed = JSON.parse(text);
 		} catch (err) {
 			output(error(undefined, "parse", `Failed to parse command: ${decodeError(err)}`));
-			continue;
+			return;
 		}
-
 		try {
-			// Handle extension UI responses
 			if ((parsed as RpcExtensionUIResponse).type === "extension_ui_response") {
 				const response = parsed as RpcExtensionUIResponse;
-				const pending = pendingExtensionRequests.get(response.id);
-				if (pending) {
-					pending.resolve(response);
-				}
-				continue;
+				pendingExtensionRequests.get(response.id)?.resolve(response);
+				return;
 			}
-
 			if (isRpcHostToolResult(parsed)) {
 				hostToolBridge.handleResult(parsed);
-				continue;
+				return;
 			}
-
 			if (isRpcHostToolUpdate(parsed)) {
 				hostToolBridge.handleUpdate(parsed);
-				continue;
+				return;
 			}
-
 			if (isRpcHostUriResult(parsed)) {
 				hostUriBridge.handleResult(parsed);
-				continue;
+				return;
 			}
-
-			// Handle regular commands
-			const command = parsed as RpcCommand;
-			const response = await handleCommand(command);
-			output(response);
-
-			// Check for deferred shutdown request (idle between commands)
+			// Ordered commands run through a serial chain to preserve causal order; the
+			// reader never blocks, so cancellation commands stay responsive even while a
+			// long command is in flight (issue 13).
+			dispatchCommand(parsed as RpcCommand);
 			await checkShutdownRequested();
 		} catch (err) {
 			output(error(undefined, "parse", `Failed to parse command: ${decodeError(err)}`));
 		}
 	}
 
+	// Persistent UDS server (issue 09): keep the AgentSession alive across client
+	// reconnects instead of exiting on stdin EOF. Frames route to the active client
+	// socket; while no client is connected they are dropped (clients resync via
+	// get_state/get_messages on reconnect).
+	if (options?.listen) {
+		const socketPath = options.listen;
+		await fs.mkdir(path.dirname(socketPath), { recursive: true }).catch(() => {});
+		// Refuse to clobber a live previous owner: probe the path first and only
+		// unlink a stale endpoint. A second `--listen` on the same path must not
+		// remove the socket another running server is still serving (#606).
+		// Unexpected probe failures are treated as alive, so this also refuses
+		// rather than unlinking a socket path we could not safely classify.
+		if (await isUnixSocketAlive(socketPath)) {
+			throw new RpcListenRefusedError(socketPath);
+		}
+		await fs.rm(socketPath, { force: true }).catch(() => {});
+		await registerRpcSession({
+			sessionId: session.sessionId,
+			pid: process.pid,
+			transport: "socket",
+			cwd: session.sessionManager.getCwd(),
+			model: session.model?.id,
+			startedAt: new Date().toISOString(),
+			endpoint: socketPath,
+		}).catch(() => {});
+
+		const noopSink = (_line: string): void => {};
+		let currentSocket: object | undefined;
+		let buf = "";
+		Bun.listen({
+			unix: socketPath,
+			socket: {
+				open(socket) {
+					currentSocket = socket;
+					buf = "";
+					frameSink = (line: string) => {
+						socket.write(line);
+					};
+					output({ type: "ready" });
+				},
+				data(socket, data) {
+					if (socket !== currentSocket) return;
+					buf += inputDecoder.decode(data);
+					while (true) {
+						const nl = buf.indexOf("\n");
+						if (nl < 0) break;
+						const text = buf.slice(0, nl).trim();
+						buf = buf.slice(nl + 1);
+						if (text) void handleInboundLine(text);
+					}
+				},
+				close(socket) {
+					if (socket === currentSocket) {
+						currentSocket = undefined;
+						frameSink = noopSink;
+					}
+				},
+				error() {},
+			},
+		});
+
+		const onSignal = (): void => {
+			void shutdown(0, "RPC socket server signal");
+		};
+		process.on("SIGINT", onSignal);
+		process.on("SIGTERM", onSignal);
+		// Block until an explicit shutdown (signal/extension) calls process.exit.
+		await new Promise<never>(() => {});
+		throw new Error("RPC socket server returned unexpectedly");
+	}
+
+	// Register this stdio RPC session so other processes can discover it (issue 10).
+	await registerRpcSession({
+		sessionId: session.sessionId,
+		pid: process.pid,
+		transport: "stdio",
+		cwd: session.sessionManager.getCwd(),
+		model: session.model?.id,
+		startedAt: new Date().toISOString(),
+	}).catch(() => {});
+
+	// Listen for JSONL input using Bun's stdin. Parse frame-by-frame so a malformed
+	// command reports a parse error without poisoning the whole long-lived RPC session.
+	for await (const line of readLines(Bun.stdin.stream())) {
+		const text = inputDecoder.decode(line).trim();
+		if (!text) continue;
+		await handleInboundLine(text);
+	}
+
 	// stdin closed — RPC client is gone, flush durable state and exit cleanly
 	await shutdown(0, "RPC client disconnected");
 	throw new Error("RPC shutdown returned unexpectedly");