@gajae-code/coding-agent 0.5.2 → 0.5.4

package/src/runtime-mcp/transports/stdio.ts

@@ -5,7 +5,8 @@
  * Messages are newline-delimited JSON.
  */
 
-import { getProjectDir, ptree, readJsonl, Snowflake } from "@gajae-code/utils";
+import { getProjectDir, readJsonl, Snowflake } from "@gajae-code/utils";
+import { type OwnedProcess, spawnOwnedProcess } from "../../runtime/process-lifecycle";
 import type {
 	JsonRpcError,
 	JsonRpcMessage,
@@ -24,7 +25,7 @@ import { toJsonRpcError } from "../../runtime-mcp/types";
 const CLOSE_WAIT_MS = 1_000;
 
 export class StdioTransport implements MCPTransport {
-	#process: ptree.ChildProcess<"pipe"> | null = null;
+	#process: OwnedProcess | null = null;
 	#pendingRequests = new Map<
 		string | number,
 		{
@@ -34,6 +35,8 @@ export class StdioTransport implements MCPTransport {
 	>();
 	#connected = false;
 	#readLoop: Promise<void> | null = null;
+	#stderrLoop: Promise<void> | null = null;
+	#closePromise: Promise<void> | null = null;
 
 	onClose?: () => void;
 	onError?: (error: Error) => void;
@@ -46,10 +49,17 @@ export class StdioTransport implements MCPTransport {
 		return this.#connected;
 	}
 
+	get closeBeforeReconnect(): true {
+		return true;
+	}
+
 	/**
 	 * Start the subprocess and begin reading.
 	 */
 	async connect(): Promise<void> {
+		if (this.#closePromise) {
+			throw new Error("Transport is closing");
+		}
 		if (this.#connected) return;
 
 		const args = this.config.args ?? [];
@@ -58,11 +68,12 @@ export class StdioTransport implements MCPTransport {
 			...this.config.env,
 		};
 
-		this.#process = ptree.spawn([this.config.command, ...args], {
+		this.#process = spawnOwnedProcess([this.config.command, ...args], {
 			cwd: this.config.cwd ?? getProjectDir(),
 			env,
 			stdin: "pipe",
-			stderr: "full",
+			gracefulMs: CLOSE_WAIT_MS,
+			name: `mcp-stdio:${this.config.command}`,
 		});
 
 		this.#connected = true;
@@ -71,13 +82,13 @@ export class StdioTransport implements MCPTransport {
 		this.#readLoop = this.#startReadLoop();
 
 		// Log stderr for debugging
-		this.#startStderrLoop();
+		this.#stderrLoop = this.#startStderrLoop();
 	}
 
 	async #startReadLoop(): Promise<void> {
-		if (!this.#process?.stdout) return;
+		if (!this.#process?.child.stdout) return;
 		try {
-			for await (const line of readJsonl(this.#process.stdout)) {
+			for await (const line of readJsonl(this.#process.child.stdout)) {
 				if (!this.#connected) break;
 				try {
 					this.#handleMessage(line as JsonRpcMessage);
@@ -95,9 +106,9 @@ export class StdioTransport implements MCPTransport {
 	}
 
 	async #startStderrLoop(): Promise<void> {
-		if (!this.#process?.stderr) return;
+		if (!this.#process?.child.stderr) return;
 
-		const reader = this.#process.stderr.getReader();
+		const reader = this.#process.child.stderr.getReader();
 		const decoder = new TextDecoder();
 
 		try {
@@ -168,26 +179,23 @@ export class StdioTransport implements MCPTransport {
 		}
 	}
 
+	#getStdin(): Bun.FileSink | null {
+		const stdin = this.#process?.child.stdin;
+		return typeof stdin === "object" && stdin !== null ? stdin : null;
+	}
+
 	#sendResponse(id: string | number, result?: unknown, error?: JsonRpcError): void {
-		if (!this.#connected || !this.#process?.stdin) return;
+		const stdin = this.#getStdin();
+		if (!this.#connected || !stdin) return;
 		const response = error
 			? { jsonrpc: "2.0" as const, id, error }
 			: { jsonrpc: "2.0" as const, id, result: result ?? {} };
-		this.#process.stdin.write(`${JSON.stringify(response)}\n`);
-		this.#process.stdin.flush();
+		stdin.write(`${JSON.stringify(response)}\n`);
+		stdin.flush();
 	}
 
 	#handleClose(): void {
-		if (!this.#connected) return;
-		this.#connected = false;
-
-		// Reject all pending requests
-		for (const [, pending] of this.#pendingRequests) {
-			pending.reject(new Error("Transport closed"));
-		}
-		this.#pendingRequests.clear();
-
-		this.onClose?.();
+		void this.#closeInternal(true);
 	}
 
 	async request<T = unknown>(
@@ -195,7 +203,8 @@ export class StdioTransport implements MCPTransport {
 		params?: Record<string, unknown>,
 		options?: MCPRequestOptions,
 	): Promise<T> {
-		if (!this.#connected || !this.#process?.stdin) {
+		const stdin = this.#getStdin();
+		if (!this.#connected || !stdin) {
 			throw new Error("Transport not connected");
 		}
 
@@ -261,8 +270,8 @@ export class StdioTransport implements MCPTransport {
 		const message = `${JSON.stringify(request)}\n`;
 		try {
 			// Bun's FileSink has write() method directly
-			this.#process.stdin.write(message);
-			this.#process.stdin.flush();
+			stdin.write(message);
+			stdin.flush();
 		} catch (error: unknown) {
 			cleanup();
 			reject(error instanceof Error ? error : new Error(String(error)));
@@ -272,7 +281,8 @@ export class StdioTransport implements MCPTransport {
 	}
 
 	async notify(method: string, params?: Record<string, unknown>): Promise<void> {
-		if (!this.#connected || !this.#process?.stdin) {
+		const stdin = this.#getStdin();
+		if (!this.#connected || !stdin) {
 			throw new Error("Transport not connected");
 		}
 
@@ -284,35 +294,51 @@ export class StdioTransport implements MCPTransport {
 
 		const message = `${JSON.stringify(notification)}\n`;
 		// Bun's FileSink has write() method directly
-		this.#process.stdin.write(message);
-		this.#process.stdin.flush();
+		stdin.write(message);
+		stdin.flush();
 	}
 
 	async close(): Promise<void> {
-		if (!this.#connected) return;
+		await this.#closeInternal(false);
+	}
+
+	#closeInternal(fromReadLoop: boolean): Promise<void> {
+		if (this.#closePromise) return this.#closePromise;
+		this.#closePromise = this.#finishClose(fromReadLoop).finally(() => {
+			this.#closePromise = null;
+		});
+		return this.#closePromise;
+	}
+
+	async #finishClose(fromReadLoop: boolean): Promise<void> {
+		const wasConnected = this.#connected;
 		this.#connected = false;
 
-		// Reject pending requests
 		for (const [, pending] of this.#pendingRequests) {
 			pending.reject(new Error("Transport closed"));
 		}
 		this.#pendingRequests.clear();
 
-		// Terminate the subprocess tree and keep the handle until exit is observed.
+		const stdin = this.#getStdin();
 		const process = this.#process;
+		this.#process = null;
 		if (process) {
-			process.kill();
-			await Promise.race([process.exited.catch(() => {}), Bun.sleep(CLOSE_WAIT_MS)]);
-			this.#process = null;
+			stdin?.end();
+			await process.dispose().catch(() => {});
+			await process.awaitExit({ timeoutMs: CLOSE_WAIT_MS }).catch(() => ({ exited: false, code: null }));
 		}
 
-		// Wait for read loop to finish
-		if (this.#readLoop) {
+		if (!fromReadLoop && this.#readLoop) {
 			await this.#readLoop.catch(() => {});
-			this.#readLoop = null;
+		}
+		this.#readLoop = null;
+
+		if (this.#stderrLoop) {
+			await this.#stderrLoop.catch(() => {});
+			this.#stderrLoop = null;
 		}
 
-		this.onClose?.();
+		if (wasConnected) this.onClose?.();
 	}
 }
 

package/src/runtime-mcp/types.ts

@@ -225,6 +225,9 @@ export interface MCPTransport {
 	/** Close the transport */
 	close(): Promise<void>;
 
+	/** Whether close must finish before reconnect can safely spawn a replacement. */
+	readonly closeBeforeReconnect?: boolean;
+
 	/** Whether the transport is connected */
 	readonly connected: boolean;
 

package/src/sdk.ts

@@ -52,6 +52,7 @@ import { resolveConfigValue } from "./config/resolve-config-value";
 import { getEmbeddedDefaultGjcSkills } from "./defaults/gjc-defaults";
 import { BUNDLED_GROK_BUILD_EXTENSION_ID, getBundledGrokBuildExtensionFactory } from "./defaults/gjc-grok-cli";
 import { initializeWithSettings } from "./discovery";
+import { disposeAllVmContexts, disposeVmContextsByOwner } from "./eval/js/context-manager";
 import { disposeAllKernelSessions, disposeKernelSessionsByOwner } from "./eval/py/executor";
 import { TtsrManager } from "./export/ttsr";
 import type { CustomCommandsLoadResult, LoadedCustomCommand } from "./extensibility/custom-commands";
@@ -414,6 +415,7 @@ function getDefaultAgentDir(): string {
  */
 export async function discoverAuthStorage(agentDir: string = getDefaultAgentDir()): Promise<AuthStorage> {
 	const brokerConfig = await resolveAuthBrokerConfig();
+	const credentialRankingMode = resolveCredentialRankingMode();
 	if (brokerConfig) {
 		const client = new AuthBrokerClient({ url: brokerConfig.url, token: brokerConfig.token });
 		const initialResult = await client.fetchSnapshot();
@@ -424,6 +426,7 @@ export async function discoverAuthStorage(agentDir: string = getDefaultAgentDir(
 		const storage = new AuthStorage(store, {
 			configValueResolver: resolveConfigValue,
 			sourceLabel: `broker ${brokerConfig.url}`,
+			credentialRankingMode,
 		});
 		await storage.reload();
 		return storage;
@@ -432,11 +435,25 @@ export async function discoverAuthStorage(agentDir: string = getDefaultAgentDir(
 	const storage = await AuthStorage.create(dbPath, {
 		configValueResolver: resolveConfigValue,
 		sourceLabel: `local ${dbPath}`,
+		credentialRankingMode,
 	});
 	await storage.reload();
 	return storage;
 }
 
+/**
+ * Opt-in multi-account credential ranking mode, read from the
+ * `GJC_CREDENTIAL_RANKING_MODE` env var. Unset/unknown → `undefined`, leaving
+ * {@link AuthStorage}'s default (`balanced`) untouched. `earliest-reset`
+ * switches to earliest-expiry-first selection so soon-to-reset tumbling-window
+ * quota is drained before it is lost.
+ */
+function resolveCredentialRankingMode(): "balanced" | "earliest-reset" | undefined {
+	const raw = process.env.GJC_CREDENTIAL_RANKING_MODE?.trim();
+	if (raw === "balanced" || raw === "earliest-reset") return raw;
+	return undefined;
+}
+
 /**
  * Discover extensions from cwd.
  */
@@ -570,6 +587,14 @@ function registerPythonCleanup(): void {
 	postmortem.register("python-cleanup", disposeAllKernelSessions);
 }
 
+let jsVmCleanupRegistered = false;
+
+function registerJsVmCleanup(): void {
+	if (jsVmCleanupRegistered) return;
+	jsVmCleanupRegistered = true;
+	postmortem.register("js-vm-cleanup", disposeAllVmContexts);
+}
+
 /**
  * Resolve whether to enable append-only context mode based on the setting and provider.
  *
@@ -806,6 +831,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 
 	registerSshCleanup();
 	registerPythonCleanup();
+	registerJsVmCleanup();
 
 	// Pin authStorage to modelRegistry.authStorage: ModelRegistry.getApiKey() routes refresh
 	// failures through that instance, so any divergent storage handed to the bridge / mcpManager
@@ -2200,6 +2226,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 			} else {
 				if (hasRegistered) agentRegistry.unregister(resolvedAgentId);
 				await disposeKernelSessionsByOwner(evalKernelOwnerId);
+				await disposeVmContextsByOwner(evalKernelOwnerId);
 			}
 		} catch (cleanupError) {
 			logger.warn("Failed to clean up createAgentSession resources after startup error", {