path: root/packages/transport-ws/src/extension.ts

/**
 * Shell — thin imperative layer that owns the Bun.serve WebSocket server.
 *
 * All decision logic lives in router.ts (pure, unit-tested).
 * This file handles I/O only: WS accept, JSON parse/stringify,
 * provider.subscribe wiring, orchestrator turn-streaming, server lifecycle.
 */

import type { Extension, HostAPI } from "@dispatch/kernel";
import type { SessionOrchestrator } from "@dispatch/session-orchestrator";
import {
	conversationCompacted,
	conversationOpened,
	conversationStatusChanged,
	sessionOrchestratorHandle,
} from "@dispatch/session-orchestrator";
import type { SurfaceContext, SurfaceProvider, SurfaceRegistry } from "@dispatch/surface-registry";
import { surfaceRegistryHandle } from "@dispatch/surface-registry";
import type { WsClientMessage, WsServerMessage } from "@dispatch/transport-contract";
import { manifest } from "./manifest.js";
import { catalogMessage, routeClientMessage, subKey } from "./router.js";

/** Active provider subscriptions + chat subscription disposers for a single WS connection. */
interface ConnectionState {
	readonly subs: Set<string>;
	readonly providerDisposers: Map<string, () => void>;
	/** Per-conversation chat subscription disposers (orchestrator.subscribe). */
	readonly chatSubscriptions: Map<string, () => void>;
}

type Ws = Bun.ServerWebSocket<ConnectionState>;

export function createTransportWsExtension(): Extension {
	let server: ReturnType<typeof Bun.serve<ConnectionState>> | undefined;
	/** Every currently-connected WS client — used for global fan-out broadcasts. */
	const connections = new Set<Ws>();
	/** Disposers for host hook subscriptions (drained on deactivate). */
	const disposers: Array<() => void> = [];

	return {
		manifest,
		async activate(host: HostAPI) {
			const registry: SurfaceRegistry = host.getService(surfaceRegistryHandle);
			const orchestrator: SessionOrchestrator = host.getService(sessionOrchestratorHandle);
			const logger = host.logger;
			const port = host.config.get<number>("surfaceWsPort") ?? 24205;

			function send(ws: Ws, msg: WsServerMessage): void {
				try {
					ws.send(JSON.stringify(msg));
				} catch {
					// Connection may have been dropped; swallow.
				}
			}

			/** Broadcast a message to EVERY connected WS client (global fan-out). */
			function broadcast(msg: WsServerMessage): void {
				for (const ws of connections) {
					send(ws, msg);
				}
			}

			/**
			 * Ensure this connection is subscribed to a conversation's chat events.
			 * Idempotent — no-op if already subscribed. The orchestrator replays
			 * buffered events to new subscribers (late-join), then streams live.
			 */
			function ensureChatSubscribed(ws: Ws, state: ConnectionState, conversationId: string): void {
				if (state.chatSubscriptions.has(conversationId)) {
					return;
				}
				const unsubscribe = orchestrator.subscribe(conversationId, (event) => {
					send(ws, { type: "chat.delta", event });
				});
				state.chatSubscriptions.set(conversationId, unsubscribe);
			}

			function subscribeToProvider(
				ws: Ws,
				provider: SurfaceProvider,
				surfaceId: string,
				conversationId: string | undefined,
				state: ConnectionState,
			): void {
				const key = subKey(surfaceId, conversationId);
				if (!provider.subscribe || state.providerDisposers.has(key)) {
					return;
				}
				const context: SurfaceContext | undefined =
					conversationId !== undefined ? { conversationId } : undefined;
				const dispose = provider.subscribe(() => {
					try {
						const spec = provider.getSpec(context);
						if (spec instanceof Promise) {
							spec
								.then((s) =>
									send(ws, {
										type: "update",
										update: {
											surfaceId,
											spec: s,
											...(conversationId !== undefined ? { conversationId } : {}),
										},
									}),
								)
								.catch(() => {});
						} else {
							send(ws, {
								type: "update",
								update: {
									surfaceId,
									spec,
									...(conversationId !== undefined ? { conversationId } : {}),
								},
							});
						}
					} catch {
						// Provider threw — log but don't kill the connection.
					}
				});
				state.providerDisposers.set(key, dispose);
			}

			function unsubscribeFromProvider(state: ConnectionState, key: string): void {
				const dispose = state.providerDisposers.get(key);
				if (dispose) {
					dispose();
					state.providerDisposers.delete(key);
				}
			}

			// Broadcast a `conversation.open` WS message to ALL connected clients
			// whenever the orchestrator signals a conversation was opened (e.g. the
			// CLI `--open` flag). The frontend decides whether to open/focus a tab —
			// the backend just signals. This is a GLOBAL fan-out (like the catalog),
			// NOT a per-conversation chat broadcast. The payload's `workspaceId`
			// is the conversation's actual persisted workspace (resolved by the
			// orchestrator from the store), so a frontend opens/focuses the tab in
			// the correct workspace.
			disposers.push(
				host.on(conversationOpened, ({ conversationId, workspaceId }) => {
					broadcast({ type: "conversation.open", conversationId, workspaceId });
				}),
			);

			// Broadcast `conversation.statusChanged` to all connected clients so
			// tabs sync across devices in real time. `workspaceId` is the
			// conversation's actual persisted workspace (resolved by the
			// orchestrator from the store), forwarded so a frontend syncs the tab
			// in the correct workspace.
			disposers.push(
				host.on(conversationStatusChanged, ({ conversationId, status, workspaceId }) => {
					broadcast({
						type: "conversation.statusChanged",
						conversationId,
						status,
						workspaceId,
					});
				}),
			);

			// Broadcast `conversation.compacted` to all connected clients so
			// the FE reloads the conversation history after compaction.
			disposers.push(
				host.on(
					conversationCompacted,
					({ conversationId, newConversationId, messagesSummarized, messagesKept }) => {
						broadcast({
							type: "conversation.compacted",
							conversationId,
							newConversationId,
							messagesSummarized,
							messagesKept,
						});
					},
				),
			);

			server = Bun.serve<ConnectionState>({
				port,
				fetch(req, srv) {
					const initial: ConnectionState = {
						subs: new Set(),
						providerDisposers: new Map(),
						chatSubscriptions: new Map(),
					};
					if (srv.upgrade(req, { data: initial })) return;
					return new Response("expected websocket", { status: 426 });
				},
				websocket: {
					open(ws) {
						connections.add(ws);
						logger.debug("transport-ws: connection open");
						send(ws, catalogMessage(registry));
					},

					message(ws, message) {
						const state = ws.data;
						if (!state) return;

						let parsed: WsClientMessage;
						try {
							parsed = JSON.parse(String(message)) as WsClientMessage;
						} catch {
							send(ws, { type: "error", message: "Invalid JSON" });
							return;
						}

						const result = routeClientMessage(registry, state.subs, parsed);

						switch (result.kind) {
							case "surface": {
								// Log surface-op errors (unknown surface or invoke failure).
								for (const reply of result.replies) {
									if (reply.type === "error") {
										logger.warn?.("transport-ws: surface-op error", {
											...(reply.surfaceId !== undefined ? { surfaceId: reply.surfaceId } : {}),
											reason: reply.message,
										});
									}
								}

								// Apply sub change.
								if (result.subChange) {
									const key = subKey(result.subChange.surfaceId, result.subChange.conversationId);
									if (result.subChange.op === "add") {
										state.subs.add(key);
										const provider = registry.getSurface(result.subChange.surfaceId);
										if (provider) {
											subscribeToProvider(
												ws,
												provider,
												result.subChange.surfaceId,
												result.subChange.conversationId,
												state,
											);
										}
									} else {
										state.subs.delete(key);
										unsubscribeFromProvider(state, key);
									}
								}

								// Send replies.
								for (const reply of result.replies) {
									send(ws, reply);
								}

								// Perform invoke if signalled.
								if (result.invoke) {
									const provider = registry.getSurface(result.invoke.surfaceId);
									if (provider) {
										const context: SurfaceContext | undefined =
											result.invoke.conversationId !== undefined
												? { conversationId: result.invoke.conversationId }
												: undefined;
										try {
											const r = provider.invoke(
												result.invoke.actionId,
												result.invoke.payload,
												context,
											);
											if (r instanceof Promise) {
												r.catch(() => {});
											}
										} catch (err: unknown) {
											const reason = err instanceof Error ? err.message : "invoke failed";
											logger.warn?.("transport-ws: surface-op error", {
												surfaceId: result.invoke.surfaceId,
												actionId: result.invoke.actionId,
												reason,
											});
										}
									}
								}
								break;
							}

							case "chat": {
								const resolvedId = result.conversationId ?? crypto.randomUUID();
								// Auto-subscribe the sender so it sees the turn's events.
								ensureChatSubscribed(ws, state, resolvedId);
								// Start the turn detached from this connection.
								const startResult = orchestrator.startTurn({
									conversationId: resolvedId,
									text: result.message,
									...(result.model !== undefined ? { modelName: result.model } : {}),
									...(result.cwd !== undefined ? { cwd: result.cwd } : {}),
									...(result.reasoningEffort !== undefined
										? { reasoningEffort: result.reasoningEffort }
										: {}),
									...(result.workspaceId !== undefined ? { workspaceId: result.workspaceId } : {}),
									...(result.computerId !== undefined ? { computerId: result.computerId } : {}),
								});
								if (!startResult.started) {
									send(ws, {
										type: "chat.error",
										conversationId: resolvedId,
										message: "a turn is already generating for this conversation",
									});
								} else {
									logger.info?.("transport-ws: chat.send accepted", {
										conversationId: resolvedId,
										model: result.model ?? null,
									});
								}
								break;
							}

							case "chat-subscribe": {
								ensureChatSubscribed(ws, state, result.conversationId);
								break;
							}

							case "chat-unsubscribe": {
								const dispose = state.chatSubscriptions.get(result.conversationId);
								if (dispose) {
									dispose();
									state.chatSubscriptions.delete(result.conversationId);
								}
								break;
							}

							case "chat-queue": {
								// Fire-and-forget: success is confirmed by the message-queue
								// SURFACE updating (startedTurn:false) or by streaming
								// chat.deltas (startedTurn:true), NOT by a reply here. On
								// startedTurn:true the sender is auto-subscribed so the new
								// turn's events stream to it (same as chat.send); on
								// startedTurn:false (queued for steering) we emit NOTHING
								// back and do not auto-subscribe.
								const enqueueResult = orchestrator.enqueue({
									conversationId: result.conversationId,
									text: result.text,
									...(result.workspaceId !== undefined ? { workspaceId: result.workspaceId } : {}),
								});
								if (enqueueResult.startedTurn) {
									ensureChatSubscribed(ws, state, result.conversationId);
								}
								logger.info?.("transport-ws: chat.queue accepted", {
									conversationId: result.conversationId,
									startedTurn: enqueueResult.startedTurn,
								});
								break;
							}

							case "chat-error": {
								logger.warn?.("transport-ws: malformed chat.send", {
									reason: result.errorMessage,
									...(result.conversationId !== undefined
										? { conversationId: result.conversationId }
										: {}),
								});
								send(ws, {
									type: "chat.error",
									...(result.conversationId !== undefined
										? { conversationId: result.conversationId }
										: {}),
									message: result.errorMessage,
								});
								break;
							}
						}
					},

					close(ws) {
						connections.delete(ws);
						const state = ws.data;
						if (state) {
							// Dispose all chat subscriptions (does NOT abort turns).
							for (const dispose of state.chatSubscriptions.values()) {
								dispose();
							}
							state.chatSubscriptions.clear();
							// Dispose surface provider subscriptions.
							for (const dispose of state.providerDisposers.values()) {
								dispose();
							}
						}
						logger.debug("transport-ws: connection close");
					},
				},
			});

			logger.info?.("transport-ws: surface WebSocket listening", { port });
		},

		deactivate() {
			for (const dispose of disposers) {
				dispose();
			}
			disposers.length = 0;
			connections.clear();
			if (server) {
				server.stop();
				server = undefined;
			}
		},
	};
}