@agentuity/coder 1.0.40 → 1.0.41

package/src/remote-session.ts

@@ -13,6 +13,15 @@
  */
 
 import type { ExtensionAPI, ExtensionContext } from '@mariozechner/pi-coding-agent';
+import {
+	applyRemoteLifecycleEvent,
+	clearRemoteLifecycleWorkingMessage,
+	createRemoteLifecycleState,
+	getRemoteLifecycleActivityLabel,
+	getRemoteLifecycleLabel,
+	syncRemoteLifecycleWorkingMessage,
+	type RemoteLifecycleState,
+} from './remote-lifecycle.ts';
 
 const DEBUG = !!process.env['AGENTUITY_DEBUG'];
 
@@ -54,6 +63,7 @@ export type RemoteUiHandler = (request: RpcUiRequest) => Promise<unknown>;
 export type RemoteConnectionHandler = (
 	state: 'connected' | 'reconnecting' | 'disconnected'
 ) => void;
+export type RemoteLifecycleHandler = (state: RemoteLifecycleState) => void;
 
 // ── Remote Session Client ──
 
@@ -72,6 +82,9 @@ export class RemoteSession {
 	private uiHandler: RemoteUiHandler | null = null;
 	private responseHandlers: RemoteResponseHandler[] = [];
 	private connectionHandlers: RemoteConnectionHandler[] = [];
+	private lifecycleHandlers: RemoteLifecycleHandler[] = [];
+	private lifecycleState: RemoteLifecycleState;
+	private replaySettledTimer: ReturnType<typeof setTimeout> | null = null;
 
 	/** Session ID this client is connected to */
 	public sessionId: string;
@@ -84,6 +97,7 @@ export class RemoteSession {
 
 	constructor(sessionId: string) {
 		this.sessionId = sessionId;
+		this.lifecycleState = createRemoteLifecycleState(sessionId);
 	}
 
 	private dispatchEvent(event: RpcEvent): void {
@@ -106,6 +120,88 @@ export class RemoteSession {
 		}
 	}
 
+	private applyLifecycle(event: Parameters<typeof applyRemoteLifecycleEvent>[1]): void {
+		const next = applyRemoteLifecycleEvent(this.lifecycleState, event);
+		if (next === this.lifecycleState) return;
+		this.lifecycleState = next;
+		for (const handler of this.lifecycleHandlers) {
+			try {
+				handler(this.lifecycleState);
+			} catch (err) {
+				log(`Lifecycle handler error: ${err instanceof Error ? err.message : String(err)}`);
+			}
+		}
+	}
+
+	private clearReplaySettledTimer(): void {
+		if (!this.replaySettledTimer) return;
+		clearTimeout(this.replaySettledTimer);
+		this.replaySettledTimer = null;
+	}
+
+	private scheduleReplaySettled(): void {
+		this.clearReplaySettledTimer();
+		this.replaySettledTimer = setTimeout(() => {
+			this.replaySettledTimer = null;
+			this.applyLifecycle({ type: 'replay_idle' });
+		}, 400);
+	}
+
+	private observeLiveSignal(eventType: string, isStreaming?: boolean): void {
+		const liveEvents = new Set([
+			'agent_start',
+			'agent_end',
+			'message_start',
+			'message_update',
+			'message_end',
+			'thinking_start',
+			'thinking_update',
+			'thinking_end',
+			'tool_call',
+			'tool_result',
+			'tool_execution_start',
+			'tool_execution_end',
+			'task_start',
+			'task_complete',
+			'task_error',
+			'turn_start',
+			'turn_end',
+			'rpc_response',
+			'rpc_ui_request',
+		]);
+		if (!liveEvents.has(eventType)) return;
+
+		this.clearReplaySettledTimer();
+		this.applyLifecycle({ type: 'live_signal', isStreaming });
+	}
+
+	private getLiveSignalStreamingState(eventType: string): boolean | undefined {
+		if (
+			eventType === 'agent_start' ||
+			eventType === 'message_start' ||
+			eventType === 'message_update' ||
+			eventType === 'thinking_start' ||
+			eventType === 'thinking_update' ||
+			eventType === 'tool_execution_start' ||
+			eventType === 'turn_start' ||
+			eventType === 'task_start'
+		) {
+			return true;
+		}
+		if (eventType === 'agent_end' || eventType === 'turn_end') {
+			return false;
+		}
+		return undefined;
+	}
+
+	private shouldMarkResuming(commandType: string): boolean {
+		return commandType === 'prompt' || commandType === 'follow_up' || commandType === 'steer';
+	}
+
+	private shouldObserveRpcResponseAsLive(): boolean {
+		return this.lifecycleState.phase !== 'paused' && this.lifecycleState.phase !== 'replaying';
+	}
+
 	/** Register a handler for RPC events from the sandbox */
 	onEvent(handler: RemoteEventHandler): void {
 		this.eventHandlers.push(handler);
@@ -126,6 +222,16 @@ export class RemoteSession {
 		this.connectionHandlers.push(handler);
 	}
 
+	/** Register a lifecycle state handler for remote attach/replay/live transitions. */
+	onLifecycleChange(handler: RemoteLifecycleHandler): void {
+		this.lifecycleHandlers.push(handler);
+		handler(this.lifecycleState);
+	}
+
+	getLifecycleState(): RemoteLifecycleState {
+		return this.lifecycleState;
+	}
+
 	/** Connect to the Hub WebSocket as a controller for the remote session */
 	async connect(hubWsUrl: string): Promise<void> {
 		this.hubWsUrl = hubWsUrl;
@@ -137,6 +243,7 @@ export class RemoteSession {
 	private doConnect(): Promise<void> {
 		return new Promise((resolve, reject) => {
 			const isReconnect = this.reconnectAttempts > 0;
+			this.applyLifecycle({ type: 'connect_start', reconnect: isReconnect });
 
 			// Build URL with controller params
 			const url = new URL(this.hubWsUrl);
@@ -180,6 +287,11 @@ export class RemoteSession {
 					this.connected = true;
 					this.reconnectAttempts = 0;
 					if (data.sessionId) this.sessionId = data.sessionId as string;
+					this.applyLifecycle({
+						type: 'init',
+						sessionId: typeof data.sessionId === 'string' ? data.sessionId : undefined,
+						label: typeof data.label === 'string' ? data.label : undefined,
+					});
 					log(`Connected to session ${this.sessionId}`);
 					this.notifyConnectionChange('connected');
 					resolve();
@@ -190,10 +302,58 @@ export class RemoteSession {
 				if (type === 'connection_rejected') {
 					clearTimeout(connectTimeout);
 					const msg = (data.message as string) || 'Connection rejected';
+					this.applyLifecycle({
+						type: 'rpc_command_error',
+						error: msg,
+						paused: false,
+					});
 					reject(new Error(msg));
 					return;
 				}
 
+				if (type === 'session_resume') {
+					this.applyLifecycle({
+						type: 'session_resume',
+						streamId: typeof data.streamId === 'string' ? data.streamId : null,
+						streamUrl: typeof data.streamUrl === 'string' ? data.streamUrl : null,
+					});
+					this.dispatchEvent({
+						type: 'session_resume',
+						...data,
+						_source: 'hub',
+					} as RpcEvent);
+					return;
+				}
+
+				if (type === 'session_stream_ready') {
+					this.applyLifecycle({
+						type: 'stream_ready',
+						streamId: typeof data.streamId === 'string' ? data.streamId : null,
+						streamUrl: typeof data.streamUrl === 'string' ? data.streamUrl : null,
+					});
+					this.dispatchEvent({
+						type: 'session_stream_ready',
+						...data,
+						_source: 'hub',
+					} as RpcEvent);
+					return;
+				}
+
+				if (type === 'rpc_command_error') {
+					const error = typeof data.error === 'string' ? data.error : 'Remote command failed';
+					this.applyLifecycle({
+						type: 'rpc_command_error',
+						error,
+						paused: /sandbox .*not connected|resume/i.test(error),
+					});
+					this.dispatchEvent({
+						type: 'rpc_command_error',
+						...data,
+						_source: 'hub',
+					} as RpcEvent);
+					return;
+				}
+
 				// Broadcast-wrapped messages from Hub (LIVE events)
 				// Format: { type: 'broadcast', event: '<name>', data: { ...payload } }
 				if (type === 'broadcast') {
@@ -201,11 +361,23 @@ export class RemoteSession {
 					const broadcastData = (data.data as Record<string, unknown>) ?? {};
 					if (broadcastEvent === 'rpc_event') {
 						const rpcEvent = broadcastData.event as RpcEvent;
-						if (rpcEvent) this.dispatchEvent({ ...rpcEvent, _source: 'live' } as RpcEvent);
+						if (rpcEvent) {
+							this.observeLiveSignal(
+								rpcEvent.type,
+								this.getLiveSignalStreamingState(rpcEvent.type)
+							);
+							this.dispatchEvent({ ...rpcEvent, _source: 'live' } as RpcEvent);
+						}
 					} else if (broadcastEvent === 'rpc_response') {
 						const response = broadcastData.response as RpcResponse;
-						if (response) this.dispatchResponse(response);
+						if (response) {
+							if (this.shouldObserveRpcResponseAsLive()) {
+								this.observeLiveSignal('rpc_response');
+							}
+							this.dispatchResponse(response);
+						}
 					} else if (broadcastEvent === 'rpc_ui_request') {
+						this.observeLiveSignal('rpc_ui_request');
 						this.handleUiRequest({
 							id: broadcastData.id as string,
 							method: broadcastData.method as string,
@@ -215,6 +387,10 @@ export class RemoteSession {
 						// Lifecycle event broadcasts (agent_start, message_end, turn_start, etc.)
 						// The broadcastData IS the event payload with a `type` field matching broadcastEvent.
 						// Dispatch as a regular event so the TUI can render agent activity.
+						this.observeLiveSignal(
+							broadcastEvent,
+							this.getLiveSignalStreamingState(broadcastEvent)
+						);
 						this.dispatchEvent({
 							type: broadcastEvent,
 							...broadcastData,
@@ -227,17 +403,27 @@ export class RemoteSession {
 				// Raw RPC messages (from Durable Stream replay — historical, not live)
 				if (type === 'rpc_event') {
 					const rpcEvent = data.event as RpcEvent;
-					if (rpcEvent) this.dispatchEvent({ ...rpcEvent, _source: 'replay' } as RpcEvent);
+					if (rpcEvent) {
+						this.applyLifecycle({ type: 'replay_event' });
+						this.scheduleReplaySettled();
+						this.dispatchEvent({ ...rpcEvent, _source: 'replay' } as RpcEvent);
+					}
 					return;
 				}
 
 				if (type === 'rpc_response') {
 					const response = data.response as RpcResponse;
-					if (response) this.dispatchResponse(response);
+					if (response) {
+						if (this.shouldObserveRpcResponseAsLive()) {
+							this.observeLiveSignal('rpc_response');
+						}
+						this.dispatchResponse(response);
+					}
 					return;
 				}
 
 				if (type === 'rpc_ui_request') {
+					this.observeLiveSignal('rpc_ui_request');
 					this.handleUiRequest({
 						id: data.id as string,
 						method: data.method as string,
@@ -248,6 +434,16 @@ export class RemoteSession {
 
 				// Session hydration (conversation entries + task states from observer hydration)
 				if (type === 'session_hydration') {
+					this.applyLifecycle({
+						type: 'hydration',
+						leadConnected:
+							typeof data.leadConnected === 'boolean' ? data.leadConnected : undefined,
+						isStreaming:
+							typeof (data.streamingState as { isStreaming?: unknown } | undefined)
+								?.isStreaming === 'boolean'
+								? Boolean((data.streamingState as { isStreaming?: boolean }).isStreaming)
+								: undefined,
+					});
 					// Pass through as an event so the extension can render it
 					for (const handler of this.eventHandlers) {
 						try {
@@ -276,10 +472,12 @@ export class RemoteSession {
 				clearTimeout(connectTimeout);
 				const wasConnected = this.connected;
 				this.connected = false;
+				this.clearReplaySettledTimer();
 				if (!this.intentionallyClosed) {
 					if (wasConnected) {
 						log('WebSocket closed unexpectedly — scheduling reconnect');
 						this.notifyConnectionChange('reconnecting');
+						this.applyLifecycle({ type: 'connection_change', state: 'reconnecting' });
 						this.scheduleReconnect();
 					} else if (!isReconnect) {
 						// Failed initial connect and not already in reconnect loop
@@ -295,6 +493,7 @@ export class RemoteSession {
 		if (this.reconnectAttempts >= MAX_RECONNECT_ATTEMPTS) {
 			log(`Max reconnect attempts (${MAX_RECONNECT_ATTEMPTS}) reached — giving up`);
 			this.notifyConnectionChange('disconnected');
+			this.applyLifecycle({ type: 'connection_change', state: 'disconnected' });
 			return;
 		}
 
@@ -334,6 +533,9 @@ export class RemoteSession {
 			log('Cannot send command — not connected');
 			return;
 		}
+		if (this.shouldMarkResuming(command.type) && this.lifecycleState.phase === 'paused') {
+			this.applyLifecycle({ type: 'local_resume_requested' });
+		}
 		this.ws.send(
 			JSON.stringify({
 				type: 'rpc_command',
@@ -379,12 +581,14 @@ export class RemoteSession {
 	/** Close the connection */
 	close(): void {
 		this.intentionallyClosed = true;
+		this.clearReplaySettledTimer();
 		if (this.reconnectTimer) {
 			clearTimeout(this.reconnectTimer);
 			this.reconnectTimer = null;
 		}
 		this.ws?.close();
 		this.ws = null;
+		this.applyLifecycle({ type: 'connection_change', state: 'disconnected' });
 	}
 
 	get isConnected(): boolean {
@@ -442,20 +646,18 @@ export async function setupRemoteMode(
 ): Promise<RemoteSession> {
 	const remote = new RemoteSession(sessionId);
 
-	// Connect to Hub
-	await remote.connect(hubWsUrl);
-	log(`Remote mode active — session ${sessionId}`);
-
 	// ── Track streaming state for widget rendering ──
 	let messageBuffer = '';
 	let thinkingBuffer = '';
 	let isStreaming = false;
 	let currentTool: string | null = null;
 	let extensionCtxRef: ExtensionContext | null = null;
+	let lifecycleOwnsWorkingMessage = false;
 
 	// Called by the extension setup to provide the rendering context
 	(remote as RemoteSessionInternal)._setExtensionCtx = (ctx: ExtensionContext) => {
 		extensionCtxRef = ctx;
+		applyLifecycleUi(remote.getLifecycleState());
 	};
 
 	// ── Render streaming output as a widget ──
@@ -477,6 +679,48 @@ export async function setupRemoteMode(
 		extensionCtxRef.ui.setWidget('remote_stream', undefined);
 	}
 
+	function applyLifecycleUi(state: RemoteLifecycleState): void {
+		if (!extensionCtxRef?.hasUI) return;
+		const shortSession = state.sessionId.slice(0, 16);
+		extensionCtxRef.ui.setStatus(
+			'remote_connection',
+			`Remote: ${shortSession}${shortSession.length < state.sessionId.length ? '...' : ''} ${getRemoteLifecycleLabel(state)}`
+		);
+		const activity = getRemoteLifecycleActivityLabel(state);
+		if (activity) {
+			extensionCtxRef.ui.setStatus('remote_activity', activity);
+		} else {
+			extensionCtxRef.ui.setStatus(
+				'remote_activity',
+				state.isStreaming ? 'agent working...' : 'idle'
+			);
+		}
+		lifecycleOwnsWorkingMessage = syncRemoteLifecycleWorkingMessage(
+			state,
+			extensionCtxRef.ui,
+			lifecycleOwnsWorkingMessage
+		);
+	}
+
+	function setNonLifecycleWorkingMessage(message?: string): void {
+		if (!extensionCtxRef?.hasUI) return;
+		extensionCtxRef.ui.setWorkingMessage(message);
+		lifecycleOwnsWorkingMessage = false;
+	}
+
+	function clearWorkingMessage(): void {
+		if (!extensionCtxRef?.hasUI) return;
+		if (lifecycleOwnsWorkingMessage) {
+			lifecycleOwnsWorkingMessage = clearRemoteLifecycleWorkingMessage(
+				extensionCtxRef.ui,
+				lifecycleOwnsWorkingMessage
+			);
+			return;
+		}
+		extensionCtxRef.ui.setWorkingMessage();
+		lifecycleOwnsWorkingMessage = false;
+	}
+
 	// ── Set up UI handler (wired to Pi's UI later in setupRemoteModeExtension) ──
 	// Default handler — overridden by setupRemoteModeExtension once ctx is available
 	remote.setUiHandler(async (request) => {
@@ -582,12 +826,39 @@ export async function setupRemoteMode(
 		const eventType = event.type as string;
 
 		switch (eventType) {
+			case 'session_resume':
+				log(
+					`Session resume signaled (${typeof (event as { streamId?: string }).streamId === 'string' ? (event as { streamId?: string }).streamId : 'no stream id'})`
+				);
+				break;
+
+			case 'session_stream_ready':
+				log(
+					`Durable stream ready (${typeof (event as { streamId?: string }).streamId === 'string' ? (event as { streamId?: string }).streamId : 'no stream id'})`
+				);
+				break;
+
+			case 'rpc_command_error': {
+				const error =
+					typeof (event as { error?: string }).error === 'string'
+						? (event as { error?: string }).error!
+						: 'Remote command failed';
+				if (extensionCtxRef?.hasUI) {
+					extensionCtxRef.ui.notify(error, 'warning');
+					clearWorkingMessage();
+				}
+				isStreaming = false;
+				clearStreamWidget();
+				log(`Remote command error: ${error}`);
+				break;
+			}
+
 			case 'message_start':
 				messageBuffer = '';
 				thinkingBuffer = '';
 				isStreaming = true;
 				if (extensionCtxRef?.hasUI) {
-					extensionCtxRef.ui.setWorkingMessage('Responding...');
+					setNonLifecycleWorkingMessage('Responding...');
 				}
 				break;
 
@@ -601,9 +872,7 @@ export async function setupRemoteMode(
 			case 'message_end': {
 				isStreaming = false;
 				clearStreamWidget();
-				if (extensionCtxRef?.hasUI) {
-					extensionCtxRef.ui.setWorkingMessage();
-				}
+				clearWorkingMessage();
 
 				// Extract content — prefer streamed buffer, fall back to message_end payload
 				let finalContent = messageBuffer.trim();
@@ -688,7 +957,7 @@ export async function setupRemoteMode(
 				const tool = (event as { toolName?: string }).toolName ?? 'tool';
 				currentTool = tool;
 				if (extensionCtxRef?.hasUI) {
-					extensionCtxRef.ui.setWorkingMessage(`Running ${tool}...`);
+					setNonLifecycleWorkingMessage(`Running ${tool}...`);
 					extensionCtxRef.ui.setStatus('remote_activity', `Running ${tool}...`);
 				}
 				log(`Tool: ${tool}`);
@@ -699,7 +968,7 @@ export async function setupRemoteMode(
 				const tool = (event as { toolName?: string }).toolName ?? currentTool ?? 'tool';
 				currentTool = null;
 				if (extensionCtxRef?.hasUI) {
-					extensionCtxRef.ui.setWorkingMessage();
+					clearWorkingMessage();
 					extensionCtxRef.ui.setStatus('remote_activity', 'agent working...');
 				}
 				log(`Tool done: ${tool}`);
@@ -716,8 +985,8 @@ export async function setupRemoteMode(
 			case 'turn_end':
 				if (extensionCtxRef?.hasUI) {
 					extensionCtxRef.ui.setStatus('remote_activity', 'idle');
-					extensionCtxRef.ui.setWorkingMessage();
 				}
+				clearWorkingMessage();
 				clearStreamWidget();
 				log('Turn ended');
 				break;
@@ -754,39 +1023,25 @@ export async function setupRemoteMode(
 
 			case 'auto_compaction_start':
 				if (extensionCtxRef?.hasUI) {
-					extensionCtxRef.ui.setWorkingMessage('Compacting context...');
+					setNonLifecycleWorkingMessage('Compacting context...');
 				}
 				break;
 
 			case 'auto_compaction_end':
-				if (extensionCtxRef?.hasUI) {
-					extensionCtxRef.ui.setWorkingMessage();
-				}
+				clearWorkingMessage();
 				break;
 		}
 	});
 
-	// ── Connection state handling ──
-	remote.onConnectionChange((state) => {
-		if (!extensionCtxRef?.hasUI) return;
-		switch (state) {
-			case 'connected':
-				extensionCtxRef.ui.setStatus(
-					'remote_connection',
-					`Remote: ${sessionId.slice(0, 16)}...`
-				);
-				break;
-			case 'reconnecting':
-				extensionCtxRef.ui.setStatus('remote_connection', 'Reconnecting...');
-				extensionCtxRef.ui.setWorkingMessage('Connection lost — reconnecting...');
-				break;
-			case 'disconnected':
-				extensionCtxRef.ui.setStatus('remote_connection', 'Disconnected');
-				extensionCtxRef.ui.setWorkingMessage('Connection lost');
-				break;
-		}
+	// ── Connection/lifecycle state handling ──
+	remote.onLifecycleChange((state) => {
+		applyLifecycleUi(state);
 	});
 
+	// Connect to Hub after all listeners are attached so hydration/replay frames are not dropped.
+	await remote.connect(hubWsUrl);
+	log(`Remote mode active — session ${sessionId}`);
+
 	// Request initial state from the sandbox
 	remote.getState();
 	remote.getMessages();

package/src/remote-tui.ts

@@ -38,6 +38,13 @@ import {
 	setNativeRemoteExtensionContext,
 	waitForNativeRemoteExtensionContext,
 } from './native-remote-ui-context.ts';
+import {
+	clearRemoteLifecycleWorkingMessage,
+	getRemoteLifecycleActivityLabel,
+	getRemoteLifecycleLabel,
+	syncRemoteLifecycleWorkingMessage,
+	type RemoteLifecycleState,
+} from './remote-lifecycle.ts';
 import { RemoteSession } from './remote-session.ts';
 import type { RpcEvent } from './remote-session.ts';
 import { agentuityCoderHub } from './index.ts';
@@ -78,6 +85,7 @@ export async function runRemoteTui(options: {
 	// TODO: Remove/Change when we get Agentuity service level auth enabled, this is just temporary
 	remote.apiKey = process.env.AGENTUITY_CODER_API_KEY || null;
 	let hydrationStreamingDetected = false;
+	let sessionResumeSeen = false;
 
 	// ── 2. Create AgentSession with coder extension loaded ──
 	// The extension provides Hub UI (footer, /hub overlay, commands, titlebar).
@@ -103,6 +111,41 @@ export async function runRemoteTui(options: {
 
 	// Access the Agent instance (typed as `any` for monkey-patching)
 	const agent: any = session.agent;
+	let lifecycleState = remote.getLifecycleState();
+	let lifecycleOwnsWorkingMessage = false;
+
+	function applyLifecycleUi(state: RemoteLifecycleState): void {
+		const ctx = getNativeRemoteExtensionContext();
+		if (!ctx?.hasUI) return;
+
+		const shortSession = state.sessionId.slice(0, 16);
+		ctx.ui.setStatus(
+			'remote_connection',
+			`Remote: ${shortSession}${shortSession.length < state.sessionId.length ? '...' : ''} ${getRemoteLifecycleLabel(state)}`
+		);
+
+		const activity = getRemoteLifecycleActivityLabel(state);
+		if (activity) {
+			ctx.ui.setStatus('remote_activity', activity);
+		} else {
+			ctx.ui.setStatus('remote_activity', state.isStreaming ? 'agent working...' : 'idle');
+		}
+
+		lifecycleOwnsWorkingMessage = syncRemoteLifecycleWorkingMessage(
+			state,
+			ctx.ui,
+			lifecycleOwnsWorkingMessage
+		);
+	}
+
+	remote.onLifecycleChange((state) => {
+		lifecycleState = state;
+		applyLifecycleUi(state);
+	});
+	void waitForNativeRemoteExtensionContext(10_000).then((ctx) => {
+		if (!ctx) return;
+		applyLifecycleUi(lifecycleState);
+	});
 
 	// ── 3. Patch Agent to be remote-backed ──
 	// Track the running prompt promise so InteractiveMode waits correctly
@@ -281,6 +324,43 @@ export async function runRemoteTui(options: {
 			(rpcEvent as any).isReplay === true;
 		log(`Event received: ${rpcEvent.type} (source=${source})`);
 
+		if (rpcEvent.type === 'session_resume') {
+			sessionResumeSeen = true;
+			log(
+				`Session resume signaled (${typeof (rpcEvent as any).streamId === 'string' ? (rpcEvent as any).streamId : 'no stream id'})`
+			);
+			return;
+		}
+
+		if (rpcEvent.type === 'session_stream_ready') {
+			log(
+				`Durable stream ready (${typeof (rpcEvent as any).streamId === 'string' ? (rpcEvent as any).streamId : 'no stream id'})`
+			);
+			return;
+		}
+
+		if (rpcEvent.type === 'rpc_command_error') {
+			const error =
+				typeof (rpcEvent as any).error === 'string'
+					? (rpcEvent as any).error
+					: 'Remote command failed';
+			const ctx = getNativeRemoteExtensionContext();
+			if (ctx?.hasUI) {
+				ctx.ui.notify(error, 'warning');
+				lifecycleOwnsWorkingMessage = clearRemoteLifecycleWorkingMessage(
+					ctx.ui,
+					lifecycleOwnsWorkingMessage
+				);
+			}
+			agent._state.error = error;
+			seenAgentStart = false;
+			seenMessageStart = false;
+			resolveRunningPrompt();
+			assistantStreamActive = false;
+			log(`Remote command error: ${error}`);
+			return;
+		}
+
 		// session_hydration is handled separately below — skip it here
 		if (rpcEvent.type === 'session_hydration') return;
 
@@ -493,8 +573,13 @@ export async function runRemoteTui(options: {
 	// InteractiveMode (which calls renderInitialMessages from SessionManager).
 	const sm = session.sessionManager;
 	let resolveHydration: () => void;
+	let hydrationComplete = false;
 	const hydrationReady = new Promise<void>((resolve) => {
-		resolveHydration = resolve;
+		resolveHydration = () => {
+			if (hydrationComplete) return;
+			hydrationComplete = true;
+			resolve();
+		};
 	});
 
 	let hydrationCount = 0;
@@ -673,15 +758,25 @@ export async function runRemoteTui(options: {
 
 	// Wait for hydration message (arrives right after init), with a timeout
 	// in case this is the first connection and there's nothing to hydrate.
-	const HYDRATION_TIMEOUT_MS = 2000;
 	await Promise.race([
 		hydrationReady,
-		new Promise<void>((resolve) =>
-			setTimeout(() => {
-				log('Hydration timeout — no session_hydration received');
-				resolve();
-			}, HYDRATION_TIMEOUT_MS)
-		),
+		new Promise<void>((resolve) => {
+			const waitStartedAt = Date.now();
+			const poll = (): void => {
+				if (hydrationComplete) {
+					resolve();
+					return;
+				}
+				const timeoutMs = sessionResumeSeen ? 5000 : 2000;
+				if (Date.now() - waitStartedAt >= timeoutMs) {
+					log('Hydration timeout — no session_hydration received');
+					resolve();
+					return;
+				}
+				setTimeout(poll, 50);
+			};
+			poll();
+		}),
 	]);
 	const smEntries = sm.getEntries?.() ?? [];
 	log(`SessionManager has ${smEntries.length} entries after hydration`);