@linkshell/gateway 0.3.9 → 0.3.10

package/src/relay.ts

@@ -2,12 +2,13 @@ import type WebSocket from "ws";
 import {
   parseEnvelope,
   parseTypedPayload,
+  protocolMessageSchemas,
   serializeEnvelope,
   createEnvelope,
 } from "@linkshell/protocol";
 import type { Envelope, ProtocolMessageType } from "@linkshell/protocol";
 import { ZodError } from "zod";
-import type { DeviceManager, ConnectedDevice } from "./sessions.js";
+import type { SessionManager, ConnectedDevice } from "./sessions.js";
 import {
   handleTunnelResponse,
   handleTunnelWsData,
@@ -15,88 +16,40 @@ import {
 } from "./tunnel.js";
 import { resolveAgentPermissionHttpAck } from "./agent-permission-http.js";
 
-const AGENT_SNAPSHOT_WARN_BYTES = Number(
-  process.env.AGENT_SNAPSHOT_WARN_BYTES ?? 1024 * 1024,
-);
-const DEFAULT_WS_BUFFERED_AMOUNT_LIMIT = 16 * 1024 * 1024;
-
-function wsBufferedAmountLimit(): number {
-  return Number(process.env.WS_BUFFERED_AMOUNT_LIMIT ?? DEFAULT_WS_BUFFERED_AMOUNT_LIMIT);
-}
-
 export function handleSocketMessage(
   socket: WebSocket,
   raw: string,
   role: "host" | "client",
-  hostDeviceId: string,
+  sessionId: string,
   deviceId: string,
-  sessions: DeviceManager,
+  sessions: SessionManager,
 ): void {
-  const codexHeader = tryParseCodexRpcHeader(raw);
-  if (codexHeader && codexHeader.type === "agent.codex.rpc") {
-    if (codexHeader.hostDeviceId !== hostDeviceId) {
-      sendSessionError(
-        socket,
-        hostDeviceId,
-        "invalid_message",
-        "Envelope hostDeviceId does not match connection hostDeviceId",
-      );
-      return;
-    }
-
-    const session = sessions.get(hostDeviceId);
-    if (!session) {
-      sendSessionError(socket, hostDeviceId, "device_not_found", "Device not found");
-      return;
-    }
-
-    if (role === "client" && !clientHasControl(session, deviceId)) {
-      sendControlConflict(socket, hostDeviceId);
-      return;
-    }
-
-    const routeEnvelope: Envelope = {
-      id: "",
-      type: "agent.codex.rpc",
-      hostDeviceId,
-      timestamp: new Date().toISOString(),
-      payload: {},
-    };
-
-    if (role === "host") {
-      broadcastToClients(session, routeEnvelope, raw);
-    } else {
-      sendToHost(session, routeEnvelope, raw);
-    }
-    return;
-  }
-
   let envelope: Envelope;
   try {
     envelope = parseEnvelope(raw);
   } catch {
-    sendSessionError(socket, hostDeviceId, "invalid_message", "Failed to parse envelope");
+    sendSessionError(socket, sessionId, "invalid_message", "Failed to parse envelope");
     return;
   }
 
-  if (envelope.hostDeviceId !== hostDeviceId) {
+  if (envelope.sessionId !== sessionId) {
     sendSessionError(
       socket,
-      hostDeviceId,
+      sessionId,
       "invalid_message",
-      "Envelope hostDeviceId does not match connection hostDeviceId",
+      "Envelope sessionId does not match connection sessionId",
     );
     return;
   }
 
-  const session = sessions.get(hostDeviceId);
+  const session = sessions.get(sessionId);
   if (!session) {
-    sendSessionError(socket, hostDeviceId, "device_not_found", "Device not found");
+    sendSessionError(socket, sessionId, "session_not_found", "Session not found");
     return;
   }
 
   try {
-    if (shouldValidatePayload(envelope.type)) {
+    if (isProtocolMessageType(envelope.type)) {
       envelope = {
         ...envelope,
         payload: parseTypedPayload(envelope.type, envelope.payload),
@@ -104,73 +57,31 @@ export function handleSocketMessage(
     }
 
     if (role === "host") {
-      handleHostMessage(envelope, raw, session, sessions);
+      handleHostMessage(envelope, session, sessions);
     } else {
-      handleClientMessage(envelope, raw, socket, session, deviceId, sessions);
+      handleClientMessage(envelope, socket, session, deviceId, sessions);
     }
   } catch (error) {
     if (error instanceof ZodError) {
       sendSessionError(
         socket,
-        hostDeviceId,
+        sessionId,
         "invalid_message",
-        error.issues[0]?.message ?? "Invalid message payload",
+        error.errors[0]?.message ?? "Invalid message payload",
       );
       return;
     }
-    sendSessionError(socket, hostDeviceId, "invalid_message", "Failed to handle message");
+    sendSessionError(socket, sessionId, "invalid_message", "Failed to handle message");
   }
 }
 
-function tryParseCodexRpcHeader(raw: string): { type: "agent.codex.rpc"; hostDeviceId: string } | undefined {
-  const trimmed = raw.trimStart();
-  if (!trimmed.startsWith("{")) return;
-  const snippet = trimmed.slice(0, 8192);
-  const hasType = /"type"\s*:\s*"agent\.codex\.rpc"/.test(snippet);
-  if (!hasType) return;
-  const hostMatch = /"hostDeviceId"\s*:\s*"([^"\\]*?)"/.exec(snippet);
-  if (!hostMatch?.[1]) return;
-  return { type: "agent.codex.rpc", hostDeviceId: hostMatch[1] };
-}
-
-function clientHasControl(session: ReturnType<DeviceManager["get"]> & {}, deviceId: string): boolean {
-  return session.controllerId === deviceId;
-}
-
-function sendControlConflict(socket: WebSocket, hostDeviceId: string): void {
-  safeSend(
-    socket,
-    serializeEnvelope(
-      createEnvelope({
-        type: "device.error",
-        hostDeviceId,
-        payload: {
-          code: "control_conflict",
-          message: "Not the controller",
-        },
-      }),
-    ),
-    hostDeviceId,
-  );
-}
-
-function shouldValidatePayload(type: string): type is ProtocolMessageType {
-  return (
-    type === "device.connect" ||
-    type === "device.ack" ||
-    type === "device.resume" ||
-    type === "terminal.input" ||
-    type === "terminal.resize" ||
-    type === "permission.decision" ||
-    type === "permission.decision.result" ||
-    type === "control.claim" ||
-    type === "control.release"
-  );
+function isProtocolMessageType(type: string): type is ProtocolMessageType {
+  return Object.prototype.hasOwnProperty.call(protocolMessageSchemas, type);
 }
 
 function sendSessionError(
   socket: WebSocket,
-  hostDeviceId: string,
+  sessionId: string,
   code: string,
   message: string,
 ): void {
@@ -178,8 +89,8 @@ function sendSessionError(
   socket.send(
     serializeEnvelope(
       createEnvelope({
-        type: "device.error",
-        hostDeviceId,
+        type: "session.error",
+        sessionId,
         payload: { code, message },
       }),
     ),
@@ -188,44 +99,42 @@ function sendSessionError(
 
 function handleHostMessage(
   envelope: Envelope,
-  raw: string,
-  session: ReturnType<DeviceManager["get"]> & {},
-  sessions: DeviceManager,
+  session: ReturnType<SessionManager["get"]> & {},
+  sessions: SessionManager,
 ): void {
   switch (envelope.type) {
-    case "device.connect": {
+    case "session.connect": {
       // Extract metadata from host's connect message
-      const p = parseTypedPayload("device.connect", envelope.payload);
-      if (p.machineId || p.hostname || p.platform || p.cwd || p.capabilities) {
+      const p = parseTypedPayload("session.connect", envelope.payload);
+      if (p.provider || p.machineId || p.hostname || p.platform || p.cwd || p.projectName) {
         sessions.setMetadata(
           session.id,
-          undefined,
+          p.provider ?? undefined,
           p.machineId ?? undefined,
           p.hostname ?? undefined,
           p.platform ?? undefined,
           p.cwd ?? undefined,
-          undefined,
-          p.capabilities ?? undefined,
+          p.projectName ?? undefined,
         );
       }
       break;
     }
     case "terminal.output": {
       sessions.bufferOutput(session.id, envelope);
-      broadcastToClients(session, envelope, raw);
+      broadcastToClients(session, envelope);
       break;
     }
     case "terminal.exit": {
       // Don't terminate session — other terminals may still be running
-      broadcastToClients(session, envelope, raw);
+      broadcastToClients(session, envelope);
       break;
     }
-    case "device.heartbeat":
+    case "session.heartbeat":
       break;
     case "permission.decision.result": {
       const p = parseTypedPayload("permission.decision.result", envelope.payload);
       resolveAgentPermissionHttpAck({
-        hostDeviceId: session.id,
+        sessionId: session.id,
         ack: {
           requestId: p.requestId,
           decision: p.decision,
@@ -235,7 +144,7 @@ function handleHostMessage(
           message: p.message,
         },
       });
-      broadcastToClients(session, envelope, raw);
+      broadcastToClients(session, envelope);
       break;
     }
     // Tunnel: host → gateway (not broadcast to clients)
@@ -256,80 +165,84 @@ function handleHostMessage(
     }
     case "control.grant":
     case "control.reject":
-      broadcastToClients(session, envelope, raw);
+      broadcastToClients(session, envelope);
       break;
     // Screen sharing: host → clients
     case "screen.frame":
     case "screen.status":
     case "screen.offer":
     case "screen.ice":
-    // Codex app-server JSON-RPC: host → clients.
-    case "agent.codex.rpc":
-    // Agent Workspace: host → clients
+    // Agent GUI: host → clients
+    case "agent.capabilities":
+    case "agent.update":
+    case "agent.permission.request":
+    case "agent.snapshot":
     case "agent.v2.capabilities":
     case "agent.v2.conversation.opened":
     case "agent.v2.conversation.list.result":
     case "agent.v2.event":
     case "agent.v2.snapshot":
-      if (envelope.type === "agent.v2.snapshot" && Buffer.byteLength(raw, "utf8") > AGENT_SNAPSHOT_WARN_BYTES) {
-        process.stderr.write(`[gateway:warn] oversized agent snapshot host=${session.id} bytes=${Buffer.byteLength(raw, "utf8")}\n`);
-      }
-      broadcastToClients(session, envelope, raw);
-      break;
-    case "agent.v2.history.page":
-    case "agent.v2.delta":
-    case "agent.v2.running_state":
     case "agent.v2.permission.request":
     // Multi-terminal: host → clients
     case "terminal.spawned":
     case "terminal.list":
     case "terminal.browse.result":
     case "terminal.file.read.result":
-      broadcastToClients(session, envelope, raw);
+      broadcastToClients(session, envelope);
       break;
     // Structured status from hooks
     case "terminal.status":
       sessions.cacheStatus(session.id, envelope);
-      broadcastToClients(session, envelope, raw);
+      broadcastToClients(session, envelope);
       break;
     default:
-      broadcastToClients(session, envelope, raw);
+      broadcastToClients(session, envelope);
       break;
   }
 }
 
 function handleClientMessage(
   envelope: Envelope,
-  raw: string,
   socket: WebSocket,
-  session: ReturnType<DeviceManager["get"]> & {},
+  session: ReturnType<SessionManager["get"]> & {},
   deviceId: string,
-  sessions: DeviceManager,
+  sessions: SessionManager,
 ): void {
   const requireController = (): boolean => {
-    if (clientHasControl(session, deviceId)) return true;
-    sendControlConflict(socket, session.id);
+    if (session.controllerId === deviceId) return true;
+    socket.send(
+      serializeEnvelope(
+        createEnvelope({
+          type: "session.error",
+          sessionId: session.id,
+          payload: {
+            code: "control_conflict",
+            message: "Not the controller",
+          },
+        }),
+      ),
+    );
     return false;
   };
 
   switch (envelope.type) {
     case "terminal.input": {
       if (!requireController()) return;
-      sendToHost(session, envelope, raw);
+      sendToHost(session, envelope);
       break;
     }
     case "terminal.resize": {
       if (!requireController()) return;
-      sendToHost(session, envelope, raw);
+      sendToHost(session, envelope);
       break;
     }
-    case "device.ack": {
+    case "session.ack": {
       // Forward ACK to host
-      sendToHost(session, envelope, raw);
+      sendToHost(session, envelope);
       break;
     }
-    case "device.resume": {
-      const p = parseTypedPayload("device.resume", envelope.payload);
+    case "session.resume": {
+      const p = parseTypedPayload("session.resume", envelope.payload);
       // Replay from gateway buffer first
       const replay = sessions.getReplayFrom(
         session.id,
@@ -338,24 +251,22 @@ function handleClientMessage(
       );
       for (const msg of replay) {
         const payload = msg.payload as Record<string, unknown>;
-        safeSend(
-          socket,
+        socket.send(
           serializeEnvelope(
             createEnvelope({
               type: "terminal.output",
-              hostDeviceId: session.id,
+              sessionId: session.id,
               terminalId: msg.terminalId,
               seq: msg.seq,
               payload: { ...payload, isReplay: true },
             }),
           ),
-          session.id,
         );
       }
       // Replay last terminal.status for each terminal
       const statusReplay = sessions.getStatusReplay(session.id);
       for (const statusMsg of statusReplay) {
-        safeSend(socket, serializeEnvelope(statusMsg), session.id);
+        socket.send(serializeEnvelope(statusMsg));
       }
       // Also forward resume to host so it can fill gaps beyond gateway buffer.
       sendToHost(session, session.machineId
@@ -373,7 +284,7 @@ function handleClientMessage(
       sessions.claimControl(session.id, deviceId);
       const grantMsg = createEnvelope({
         type: "control.grant",
-        hostDeviceId: session.id,
+        sessionId: session.id,
         payload: { deviceId },
       });
       // Broadcast to ALL clients so previous controller updates its state
@@ -385,20 +296,25 @@ function handleClientMessage(
       sessions.releaseControl(session.id, deviceId);
       const releaseMsg = createEnvelope({
         type: "control.release",
-        hostDeviceId: session.id,
+        sessionId: session.id,
         payload: { deviceId },
       });
       broadcastToClients(session, releaseMsg);
       sendToHost(session, releaseMsg);
       break;
     }
-    case "device.heartbeat":
+    case "session.heartbeat":
       break;
     // Screen sharing: client → host
     case "screen.start":
     case "screen.stop":
     case "screen.answer":
     case "screen.ice":
+    case "agent.session.new":
+    case "agent.session.load":
+    case "agent.prompt":
+    case "agent.cancel":
+    case "agent.permission.response":
     case "agent.v2.conversation.open":
     case "agent.v2.prompt":
     case "agent.v2.command.execute":
@@ -416,55 +332,41 @@ function handleClientMessage(
     case "file.upload":
     case "permission.decision":
       if (!requireController()) return;
-      sendToHost(session, envelope, raw);
+      sendToHost(session, envelope);
       break;
+    case "agent.initialize":
+    case "agent.session.list":
     case "agent.v2.capabilities.request":
     case "agent.v2.conversation.list":
     case "agent.v2.snapshot.request":
-    case "agent.v2.history.request":
-    case "agent.v2.delta.request":
-    // Codex app-server JSON-RPC: client → host.
-    case "agent.codex.rpc":
-      if (!requireController()) return;
-      sendToHost(session, envelope, raw);
+      sendToHost(session, envelope);
       break;
     default:
-      sendToHost(session, envelope, raw);
+      sendToHost(session, envelope);
       break;
   }
 }
 
-function safeSend(socket: WebSocket, data: string, hostDeviceId: string): boolean {
-  if (socket.readyState !== socket.OPEN) return false;
-  const limit = wsBufferedAmountLimit();
-  const frameBytes = Buffer.byteLength(data);
-  const queuedBytes = socket.bufferedAmount + frameBytes;
-  if (queuedBytes > limit) {
-    process.stderr.write(`[gateway:warn] closing slow socket host=${hostDeviceId} buffered=${socket.bufferedAmount} frame=${frameBytes} limit=${limit}\n`);
-    socket.close(1013, "client too slow");
-    return false;
-  }
-  socket.send(data);
-  return true;
-}
-
 function broadcastToClients(
-  session: ReturnType<DeviceManager["get"]> & {},
+  session: ReturnType<SessionManager["get"]> & {},
   envelope: Envelope,
-  raw?: string,
 ): void {
-  const data = raw ?? serializeEnvelope(envelope);
+  const data = serializeEnvelope(envelope);
   for (const [, client] of session.clients) {
-    safeSend(client.socket, data, session.id);
+    if (client.socket.readyState === client.socket.OPEN) {
+      client.socket.send(data);
+    }
   }
 }
 
 function sendToHost(
-  session: ReturnType<DeviceManager["get"]> & {},
+  session: ReturnType<SessionManager["get"]> & {},
   envelope: Envelope,
-  raw?: string,
 ): void {
-  if (session.host) {
-    safeSend(session.host.socket, raw ?? serializeEnvelope(envelope), session.id);
+  if (
+    session.host &&
+    session.host.socket.readyState === session.host.socket.OPEN
+  ) {
+    session.host.socket.send(serializeEnvelope(envelope));
   }
 }