import { WebSocketServer, WebSocket } from "ws"; import { createServer as createHttpServer } from "node:http"; import type { AgentHandler, AgentMessage, BrowserToRelayMessage, HandlerMessage, RelayToHandlerMessage, RelayToBrowserMessage, AgentContext, } from "./protocol.js"; import { DEFAULT_RELAY_PORT, RELAY_TOKEN_PARAM } from "./protocol.js"; interface RegisteredHandler { agentId: string; handler: AgentHandler; socket?: WebSocket; } interface SessionMessageQueue { push: (message: AgentMessage) => void; close: () => void; [Symbol.asyncIterator]: () => AsyncIterator; } const createSessionMessageQueue = (): SessionMessageQueue => { const pendingMessages: AgentMessage[] = []; let resolveNextMessage: | ((result: IteratorResult) => void) | null = null; let isClosed = false; return { push: (message) => { if (isClosed) return; if (resolveNextMessage) { resolveNextMessage({ value: message, done: false }); resolveNextMessage = null; } else { pendingMessages.push(message); } }, close: () => { isClosed = true; if (resolveNextMessage) { resolveNextMessage({ value: undefined, done: true, } as IteratorResult); resolveNextMessage = null; } }, [Symbol.asyncIterator]: () => ({ next: () => { if (pendingMessages.length > 0) { return Promise.resolve({ value: pendingMessages.shift()!, done: false, }); } if (isClosed) { return Promise.resolve({ value: undefined, done: true, } as IteratorResult); } return new Promise((resolve) => { resolveNextMessage = resolve; }); }, }), }; }; interface ActiveSession { sessionId: string; agentId: string; browserSocket: WebSocket; handlerSocket?: WebSocket; abortController: AbortController; } interface RelayServerOptions { port?: number; token?: string; } export interface RelayServer { start: () => Promise; stop: () => Promise; registerHandler: (handler: AgentHandler) => void; unregisterHandler: (agentId: string) => void; getRegisteredHandlerIds: () => string[]; } export const createRelayServer = ( options: RelayServerOptions = {}, ): RelayServer => { const port = options.port ?? DEFAULT_RELAY_PORT; const token = options.token; const registeredHandlers = new Map(); const activeSessions = new Map(); const browserSockets = new Set(); const handlerSockets = new Map(); const sessionMessageQueues = new Map(); const undoRedoSessionOwners = new Map(); let httpServer: ReturnType | null = null; let webSocketServer: WebSocketServer | null = null; const broadcastHandlerList = () => { const handlerIds = Array.from(registeredHandlers.keys()); const message: RelayToBrowserMessage = { type: "handlers", handlers: handlerIds, }; const serializedMessage = JSON.stringify(message); for (const socket of browserSockets) { if (socket.readyState === WebSocket.OPEN) { socket.send(serializedMessage); } } }; const sendToBrowser = (socket: WebSocket, message: RelayToBrowserMessage) => { if (socket.readyState === WebSocket.OPEN) { socket.send(JSON.stringify(message)); } }; const handleBrowserMessage = async ( socket: WebSocket, message: BrowserToRelayMessage, ) => { const { type, agentId, sessionId, context } = message; if (type === "health") { sendToBrowser(socket, { type: "health" }); return; } const registered = registeredHandlers.get(agentId); if (!registered) { sendToBrowser(socket, { type: "agent-error", agentId, sessionId, content: `Agent "${agentId}" is not registered`, }); return; } const effectiveSessionId = sessionId ?? `session-${Date.now()}-${Math.random().toString(36).slice(2)}`; if (type === "agent-request") { if ( !context || typeof context.prompt !== "string" || !Array.isArray(context.content) ) { sendToBrowser(socket, { type: "agent-error", agentId, sessionId: effectiveSessionId, content: "Invalid context: missing or malformed prompt/content", }); return; } const abortController = new AbortController(); activeSessions.set(effectiveSessionId, { sessionId: effectiveSessionId, agentId, browserSocket: socket, handlerSocket: registered.socket, abortController, }); runAgentTask( registered.handler, context, effectiveSessionId, socket, abortController.signal, ); } else if (type === "agent-abort") { const session = activeSessions.get(effectiveSessionId); if (session) { session.abortController.abort(); registered.handler.abort?.(effectiveSessionId); activeSessions.delete(effectiveSessionId); } } else if (type === "agent-undo") { try { await registered.handler.undo?.(); sendToBrowser(socket, { type: "agent-done", agentId, sessionId: effectiveSessionId, }); } catch (error) { const errorMessage = error instanceof Error ? error.message : "Unknown error"; sendToBrowser(socket, { type: "agent-error", agentId, sessionId: effectiveSessionId, content: `Undo failed: ${errorMessage}`, }); } } else if (type === "agent-redo") { try { await registered.handler.redo?.(); sendToBrowser(socket, { type: "agent-done", agentId, sessionId: effectiveSessionId, }); } catch (error) { const errorMessage = error instanceof Error ? error.message : "Unknown error"; sendToBrowser(socket, { type: "agent-error", agentId, sessionId: effectiveSessionId, content: `Redo failed: ${errorMessage}`, }); } } }; const runAgentTask = async ( handler: AgentHandler, context: AgentContext, sessionId: string, browserSocket: WebSocket, signal: AbortSignal, ) => { const combinedPrompt = `${context.prompt}\n\n${context.content.join("\n\n")}`; try { let didComplete = false; for await (const message of handler.run(combinedPrompt, { sessionId, signal, })) { if (signal.aborted) break; const getBrowserMessageType = ( messageType: string, ): "agent-status" | "agent-error" | "agent-done" => { if (messageType === "status") return "agent-status"; if (messageType === "error") return "agent-error"; return "agent-done"; }; sendToBrowser(browserSocket, { type: getBrowserMessageType(message.type), agentId: handler.agentId, sessionId, content: message.content, }); if (message.type === "done" || message.type === "error") { didComplete = true; break; } } if (!signal.aborted && !didComplete) { sendToBrowser(browserSocket, { type: "agent-done", agentId: handler.agentId, sessionId, }); } } catch (error) { if (!signal.aborted) { const errorMessage = error instanceof Error ? error.message : "Unknown error"; sendToBrowser(browserSocket, { type: "agent-error", agentId: handler.agentId, sessionId, content: errorMessage, }); } } finally { activeSessions.delete(sessionId); } }; const createRemoteHandler = ( agentId: string, socket: WebSocket, ): AgentHandler => { return { agentId, run: async function* (userPrompt, options) { const sessionId = options?.sessionId ?? `session-${Date.now()}-${Math.random().toString(36).slice(2)}`; const signal = options?.signal; const messageQueue = createSessionMessageQueue(); sessionMessageQueues.set(sessionId, messageQueue); const cleanupQueue = () => { messageQueue.close(); sessionMessageQueues.delete(sessionId); }; if (signal) { signal.addEventListener("abort", cleanupQueue, { once: true }); if (signal.aborted) { cleanupQueue(); return; } } const invokeMessage: RelayToHandlerMessage = { type: "invoke-handler", method: "run", sessionId, payload: { prompt: userPrompt }, }; if (socket.readyState !== WebSocket.OPEN) { yield { type: "error", content: "Handler disconnected" }; cleanupQueue(); return; } socket.send(JSON.stringify(invokeMessage)); try { for await (const message of messageQueue) { if (signal?.aborted) break; yield message; if (message.type === "done" || message.type === "error") { break; } } } finally { if (signal) { signal.removeEventListener("abort", cleanupQueue); } cleanupQueue(); } }, abort: (sessionId) => { if (socket.readyState === WebSocket.OPEN) { const abortMessage: RelayToHandlerMessage = { type: "invoke-handler", method: "abort", sessionId, }; socket.send(JSON.stringify(abortMessage)); } }, undo: async () => { if (socket.readyState !== WebSocket.OPEN) { throw new Error("Handler disconnected"); } const sessionId = `undo-${agentId}-${Date.now()}-${Math.random().toString(36).slice(2)}`; const messageQueue = createSessionMessageQueue(); sessionMessageQueues.set(sessionId, messageQueue); undoRedoSessionOwners.set(sessionId, socket); const undoMessage: RelayToHandlerMessage = { type: "invoke-handler", method: "undo", sessionId, }; socket.send(JSON.stringify(undoMessage)); try { for await (const message of messageQueue) { if (message.type === "error") { throw new Error(message.content); } if (message.type === "done") { return; } } } finally { messageQueue.close(); sessionMessageQueues.delete(sessionId); undoRedoSessionOwners.delete(sessionId); } }, redo: async () => { if (socket.readyState !== WebSocket.OPEN) { throw new Error("Handler disconnected"); } const sessionId = `redo-${agentId}-${Date.now()}-${Math.random().toString(36).slice(2)}`; const messageQueue = createSessionMessageQueue(); sessionMessageQueues.set(sessionId, messageQueue); undoRedoSessionOwners.set(sessionId, socket); const redoMessage: RelayToHandlerMessage = { type: "invoke-handler", method: "redo", sessionId, }; socket.send(JSON.stringify(redoMessage)); try { for await (const message of messageQueue) { if (message.type === "error") { throw new Error(message.content); } if (message.type === "done") { return; } } } finally { messageQueue.close(); sessionMessageQueues.delete(sessionId); undoRedoSessionOwners.delete(sessionId); } }, }; }; const cleanupSessionsByHandlerSocket = (socket: WebSocket) => { for (const [sessionId, session] of activeSessions) { if (session.handlerSocket === socket) { const queue = sessionMessageQueues.get(sessionId); if (queue) { queue.push({ type: "error", content: "Handler disconnected", }); queue.close(); sessionMessageQueues.delete(sessionId); } session.abortController.abort(); activeSessions.delete(sessionId); } } for (const [sessionId, ownerSocket] of undoRedoSessionOwners) { if (ownerSocket === socket) { const queue = sessionMessageQueues.get(sessionId); if (queue) { queue.push({ type: "error", content: "Handler disconnected", }); queue.close(); sessionMessageQueues.delete(sessionId); } undoRedoSessionOwners.delete(sessionId); } } }; const handleHandlerMessage = (socket: WebSocket, message: HandlerMessage) => { if (message.type === "register-handler") { const remoteHandler = createRemoteHandler(message.agentId, socket); registeredHandlers.set(message.agentId, { agentId: message.agentId, handler: remoteHandler, socket, }); handlerSockets.set(socket, message.agentId); broadcastHandlerList(); } else if (message.type === "unregister-handler") { const agentId = handlerSockets.get(socket); if (agentId) { const registeredHandler = registeredHandlers.get(agentId); const isCurrentHandler = registeredHandler?.socket === socket; cleanupSessionsByHandlerSocket(socket); if (isCurrentHandler) { registeredHandlers.delete(agentId); broadcastHandlerList(); } handlerSockets.delete(socket); } } else if ( message.type === "agent-status" || message.type === "agent-done" || message.type === "agent-error" ) { const messageQueue = sessionMessageQueues.get(message.sessionId); if (messageQueue) { const getQueueMessageType = ( handlerMessageType: string, ): "status" | "done" | "error" => { if (handlerMessageType === "agent-status") return "status"; if (handlerMessageType === "agent-done") return "done"; return "error"; }; messageQueue.push({ type: getQueueMessageType(message.type), content: message.content ?? "", }); if (message.type === "agent-done" || message.type === "agent-error") { messageQueue.close(); } } } }; const start = async (): Promise => { return new Promise((resolve, reject) => { httpServer = createHttpServer((req, res) => { const requestUrl = new URL(req.url ?? "", `http://localhost:${port}`); if (requestUrl.pathname === "/health") { if (token) { const clientToken = requestUrl.searchParams.get(RELAY_TOKEN_PARAM); if (clientToken !== token) { res.writeHead(401, { "Content-Type": "application/json" }); res.end(JSON.stringify({ error: "Unauthorized" })); return; } } res.writeHead(200, { "Content-Type": "application/json" }); res.end( JSON.stringify({ status: "ok", handlers: getRegisteredHandlerIds(), }), ); return; } res.writeHead(404); res.end(); }); httpServer.on("error", (error) => { reject(error); }); webSocketServer = new WebSocketServer({ server: httpServer }); webSocketServer.on("error", (error) => { reject(error); }); webSocketServer.on("connection", (socket, request) => { if (token) { const connectionUrl = new URL( request.url ?? "", `http://localhost:${port}`, ); const clientToken = connectionUrl.searchParams.get(RELAY_TOKEN_PARAM); if (clientToken !== token) { socket.close(4001, "Unauthorized"); return; } } const isHandlerConnection = request.headers["x-relay-handler"] === "true"; if (isHandlerConnection) { socket.on("message", (data) => { try { const message = JSON.parse(data.toString()) as HandlerMessage; handleHandlerMessage(socket, message); } catch {} }); const cleanupHandlerSocket = () => { const agentId = handlerSockets.get(socket); if (agentId) { const registeredHandler = registeredHandlers.get(agentId); const isCurrentHandler = registeredHandler?.socket === socket; cleanupSessionsByHandlerSocket(socket); if (isCurrentHandler) { registeredHandlers.delete(agentId); broadcastHandlerList(); } handlerSockets.delete(socket); } }; socket.on("close", cleanupHandlerSocket); socket.on("error", () => { cleanupHandlerSocket(); }); } else { browserSockets.add(socket); sendToBrowser(socket, { type: "handlers", handlers: getRegisteredHandlerIds(), }); socket.on("message", (data) => { try { const message = JSON.parse( data.toString(), ) as BrowserToRelayMessage; handleBrowserMessage(socket, message); } catch {} }); socket.on("close", () => { browserSockets.delete(socket); for (const [sessionId, session] of activeSessions) { if (session.browserSocket === socket) { session.abortController.abort(); const handler = registeredHandlers.get(session.agentId); handler?.handler.abort?.(sessionId); activeSessions.delete(sessionId); } } }); } }); httpServer.listen(port, () => { resolve(); }); }); }; const stop = async (): Promise => { for (const session of activeSessions.values()) { session.abortController.abort(); } activeSessions.clear(); for (const queue of sessionMessageQueues.values()) { queue.close(); } sessionMessageQueues.clear(); undoRedoSessionOwners.clear(); for (const socket of browserSockets) { socket.close(); } browserSockets.clear(); for (const socket of handlerSockets.keys()) { socket.close(); } handlerSockets.clear(); webSocketServer?.close(); httpServer?.close(); }; const registerHandler = (handler: AgentHandler) => { registeredHandlers.set(handler.agentId, { agentId: handler.agentId, handler, }); broadcastHandlerList(); }; const unregisterHandler = (agentId: string) => { registeredHandlers.delete(agentId); broadcastHandlerList(); }; const getRegisteredHandlerIds = (): string[] => { return Array.from(registeredHandlers.keys()); }; return { start, stop, registerHandler, unregisterHandler, getRegisteredHandlerIds, }; };