Drain cloud sandbox exec streams to EOF after completion

packages/cli/src/cmd/cloud/sandbox/exec.ts

@@ -11,6 +11,8 @@ import { streamUrlToWritable } from '../../../utils/stream-url.ts';
 const EXECUTION_WAIT_DURATION = '5m';
 const EMPTY_STREAM_FAST_POLL_MS = 100;
 const EMPTY_STREAM_FAST_TIMEOUT_MS = 2000;
+const STREAM_DRAIN_TIMEOUT_MS = 30_000;
+const STDOUT_DRAIN_GRACE_MS = 2000;
 
 const SandboxExecResponseSchema = z.object({
 	executionId: z.string().describe('Unique execution identifier'),
@@ -296,40 +298,68 @@ export const execSubcommand = createCommand({
 				finalExecution.exitCode ?? 'undefined'
 			);
 
+			let streamDrainTimedOut = false;
 			if (streamPromises.length > 0) {
 				logger.debug('[exec] waiting for %d stream(s) to EOF', streamPromises.length);
 				const streamWaitStart = Date.now();
-				let graceTriggered = false;
 				if (!streamResults) {
-					const streamGraceMs = 500;
-					const streamGrace = setTimeout(() => {
-						graceTriggered = true;
-						logger.debug(
-							'[exec] stream grace period (%dms) expired after execution complete — aborting streams',
-							streamGraceMs
-						);
-						streamAbortController.abort();
-					}, streamGraceMs);
+					const streamDrainPromise = Promise.all(streamPromises);
+					let drainTimeout: ReturnType<typeof setTimeout> | undefined;
 					try {
-						streamResults = await Promise.all(streamPromises);
+						const drainResult = await Promise.race([
+							streamDrainPromise.then((results) => ({ type: 'streams' as const, results })),
+							new Promise<{ type: 'timeout' }>((resolve) => {
+								drainTimeout = setTimeout(
+									() => resolve({ type: 'timeout' }),
+									STREAM_DRAIN_TIMEOUT_MS
+								);
+							}),
+						]);
+
+						if (drainResult.type === 'streams') {
+							streamResults = drainResult.results;
+						} else {
+							streamDrainTimedOut = true;
+							tui.warning(
+								`Timed out after ${STREAM_DRAIN_TIMEOUT_MS / 1000}s waiting for execution output streams to finish; output may be incomplete`
+							);
+							streamAbortController.abort();
+							void streamDrainPromise.catch((err) => {
+								logger.debug('[exec] stream failed after drain timeout: %s', err);
+							});
+						}
 					} finally {
-						clearTimeout(streamGrace);
+						if (drainTimeout) {
+							clearTimeout(drainTimeout);
+						}
 					}
 				}
-				logger.debug(
-					'[exec] all streams done in %dms (graceTriggered=%s)',
-					Date.now() - streamWaitStart,
-					graceTriggered
-				);
+				logger.debug('[exec] stream wait finished in %dms', Date.now() - streamWaitStart);
 			}
 
 			if (!options.json && process.stdout.writable) {
 				await new Promise<void>((resolve) => {
-					if (process.stdout.writableNeedDrain) {
-						process.stdout.once('drain', () => resolve());
-					} else {
+					if (!process.stdout.writableNeedDrain) {
+						resolve();
+						return;
+					}
+					// After a stream-drain timeout the output is already known
+					// incomplete, so a stuck stdout consumer must not hold the
+					// command open indefinitely.
+					let drainWait: ReturnType<typeof setTimeout> | undefined;
+					const onDrain = () => {
+						if (drainWait) {
+							clearTimeout(drainWait);
+						}
 						resolve();
+					};
+					if (streamDrainTimedOut) {
+						drainWait = setTimeout(() => {
+							process.stdout.removeListener('drain', onDrain);
+							resolve();
+						}, STDOUT_DRAIN_GRACE_MS);
 					}
+					process.stdout.once('drain', onDrain);
 				});
 			}
 

packages/cli/test/cmd/cloud/sandbox/exec.test.ts

@@ -0,0 +1,187 @@
+import { afterAll, afterEach, beforeAll, describe, expect, test } from 'bun:test';
+import { APIClient, createMinimalLogger } from '@agentuity/server';
+import { mkdtemp, rm } from 'node:fs/promises';
+import { tmpdir } from 'node:os';
+import { join } from 'node:path';
+import { closeDatabase } from '../../../../src/cache/index.ts';
+import { execSubcommand } from '../../../../src/cmd/cloud/sandbox/exec.ts';
+import type { CommandContext } from '../../../../src/types.ts';
+
+// Anchored to the moment the execution result is served, so the stream's EOF
+// always lands well past the old 500ms post-completion grace window that the
+// drain fix removed, regardless of scheduling skew before completion.
+const STREAM_DELAY_AFTER_EXECUTION_MS = 1200;
+const STREAM_OUTPUT = 'delayed output\n';
+const ORIGINAL_STDOUT_WRITE = process.stdout.write;
+const ORIGINAL_FETCH = globalThis.fetch;
+const ORIGINAL_CONFIG_DIR = process.env.AGENTUITY_CONFIG_DIR;
+
+let configDir: string;
+
+beforeAll(async () => {
+	configDir = await mkdtemp(join(tmpdir(), 'sandbox-exec-command-test-'));
+	process.env.AGENTUITY_CONFIG_DIR = configDir;
+});
+
+afterEach(() => {
+	process.stdout.write = ORIGINAL_STDOUT_WRITE;
+	globalThis.fetch = ORIGINAL_FETCH;
+});
+
+afterAll(async () => {
+	closeDatabase();
+	if (ORIGINAL_CONFIG_DIR === undefined) {
+		delete process.env.AGENTUITY_CONFIG_DIR;
+	} else {
+		process.env.AGENTUITY_CONFIG_DIR = ORIGINAL_CONFIG_DIR;
+	}
+	await rm(configDir, { recursive: true, force: true });
+});
+
+function delayedStream(
+	signal: AbortSignal,
+	executionServed: Promise<void>
+): ReadableStream<Uint8Array> {
+	let cancelled = false;
+	let delivered = false;
+
+	return new ReadableStream({
+		start(controller) {
+			const abort = () => {
+				cancelled = true;
+				controller.error(new DOMException('The operation was aborted', 'AbortError'));
+			};
+			if (signal.aborted) {
+				abort();
+			} else {
+				signal.addEventListener('abort', abort, { once: true });
+			}
+		},
+		async pull(controller) {
+			if (delivered) {
+				return;
+			}
+			delivered = true;
+			await executionServed;
+			await new Promise((resolve) => setTimeout(resolve, STREAM_DELAY_AFTER_EXECUTION_MS));
+			if (!cancelled) {
+				controller.enqueue(new TextEncoder().encode(STREAM_OUTPUT));
+				controller.close();
+			}
+		},
+		cancel() {
+			cancelled = true;
+		},
+	});
+}
+
+function installFetchFake(): void {
+	// Definite assignment: a Promise executor runs synchronously.
+	let resolveExecutionServed!: () => void;
+	const executionServed = new Promise<void>((resolve) => {
+		resolveExecutionServed = resolve;
+	});
+	globalThis.fetch = (async (input: string | URL | Request, init?: RequestInit) => {
+		const request = new Request(input, init);
+		const url = new URL(request.url);
+
+		if (request.method === 'GET' && url.pathname === '/cli/sandbox/sbx_test') {
+			return Response.json({
+				success: true,
+				data: {
+					id: 'sbx_test',
+					name: 'test',
+					region: 'local',
+					status: 'idle',
+					orgId: 'org_test',
+					projectId: null,
+				},
+			});
+		}
+
+		if (request.method === 'POST' && url.pathname === '/sandbox/sbx_test/execute') {
+			return Response.json({
+				success: true,
+				data: {
+					executionId: 'exec_test',
+					status: 'running',
+					stdoutStreamUrl: 'https://pulse.test/stdout',
+				},
+			});
+		}
+
+		if (request.method === 'GET' && url.pathname === '/sandbox/execution/exec_test') {
+			resolveExecutionServed();
+			return Response.json({
+				success: true,
+				data: {
+					executionId: 'exec_test',
+					sandboxId: 'sbx_test',
+					status: 'completed',
+					exitCode: 0,
+					durationMs: 10,
+				},
+			});
+		}
+
+		if (request.method === 'GET' && url.hostname === 'pulse.test') {
+			return new Response(delayedStream(request.signal, executionServed));
+		}
+
+		return Response.json({ success: false, message: 'unexpected request' }, { status: 404 });
+	}) as typeof fetch;
+}
+
+function makeContext(json: boolean, catalystUrl: string): CommandContext {
+	const logger = createMinimalLogger();
+	const context = {
+		args: { sandboxId: 'sbx_test', command: ['echo', 'hello'] },
+		opts: { timestamps: false, quiet: false },
+		options: { json, quiet: false, logLevel: 'error' },
+		auth: {
+			apiKey: 'ag_test',
+			userId: 'usr_test',
+			expires: new Date(Date.now() + 60_000),
+		},
+		apiClient: new APIClient(catalystUrl, logger, 'ag_test'),
+		config: { name: `exec-test-${json ? 'json' : 'plain'}`, preferences: {} },
+		logger,
+		getExecutingAgent: () => undefined,
+	};
+
+	return context as unknown as CommandContext;
+}
+
+async function runExec(json: boolean): Promise<unknown> {
+	installFetchFake();
+	const catalystUrl = 'https://catalyst.test';
+
+	const handler = execSubcommand.handler;
+	if (!handler) {
+		throw new Error('sandbox exec handler is missing');
+	}
+	return handler(makeContext(json, catalystUrl));
+}
+
+describe('cloud sandbox exec command', () => {
+	test('drains delayed stdout after execution completes in plain mode', async () => {
+		let stdout = '';
+		process.stdout.write = ((chunk: string | Uint8Array) => {
+			stdout += typeof chunk === 'string' ? chunk : Buffer.from(chunk).toString('utf-8');
+			return true;
+		}) as typeof process.stdout.write;
+
+		await runExec(false);
+
+		expect(stdout).toContain(STREAM_OUTPUT);
+	});
+
+	test('drains delayed stdout after execution completes in JSON mode', async () => {
+		const result = await runExec(true);
+
+		expect(result).toMatchObject({
+			stdout: STREAM_OUTPUT,
+			output: STREAM_OUTPUT,
+		});
+	});
+});