Improve connection selection and agent hydration reliability

packages/server/src/client/daemon-client-relay-e2ee-transport.ts

@@ -104,9 +104,11 @@ export function createEncryptedTransport(
         onerror: (error) => emitError(error),
       });
     } catch (error) {
-      logger.warn({ err: error }, "relay_e2ee_handshake_failed");
+      logger.warn({ err: normalizeTransportError(error) }, "relay_e2ee_handshake_failed");
       emitError(error);
-      base.close(1011, "E2EE handshake failed");
+      // Browser WebSocket.close only accepts 1000 or 3000-4999.
+      // Use an app-defined code so this path works in browser and Node runtimes.
+      base.close(4001, "E2EE handshake failed");
     }
   };
 
@@ -189,3 +191,14 @@ function invokeHandler<TArgs extends unknown[]>(
     // no-op
   }
 }
+
+function normalizeTransportError(error: unknown): Record<string, string> {
+  if (error instanceof Error) {
+    return {
+      name: error.name,
+      message: error.message,
+      ...(typeof error.stack === "string" ? { stack: error.stack } : {}),
+    };
+  }
+  return { message: String(error) };
+}

packages/server/src/client/daemon-client-transport.test.ts

@@ -1,5 +1,6 @@
 import { describe, expect, test, vi } from "vitest";
 import {
+  createEncryptedTransport,
   createWebSocketTransportFactory,
   decodeMessageData,
   describeTransportClose,
@@ -8,7 +9,48 @@ import {
   extractRelayMessageData,
 } from "./daemon-client-transport.js";
 
+const createClientChannelMock = vi.hoisted(() => vi.fn());
+
+vi.mock("@getpaseo/relay/e2ee", () => ({
+  createClientChannel: createClientChannelMock,
+}));
+
 describe("daemon-client transport helpers", () => {
+  test("createEncryptedTransport closes handshake failures with browser-safe code", async () => {
+    createClientChannelMock.mockReset();
+    createClientChannelMock.mockRejectedValueOnce(new Error("handshake failed"));
+
+    let openHandler: (() => void) | null = null;
+    const close = vi.fn();
+
+    createEncryptedTransport(
+      {
+        send: vi.fn(),
+        close,
+        onOpen: (handler) => {
+          openHandler = handler;
+          return () => {
+            if (openHandler === handler) {
+              openHandler = null;
+            }
+          };
+        },
+        onClose: () => () => {},
+        onError: () => () => {},
+        onMessage: () => () => {},
+      },
+      "daemon-public-key",
+      { warn: vi.fn() }
+    );
+
+    expect(openHandler).not.toBeNull();
+    openHandler?.();
+
+    await vi.waitFor(() => {
+      expect(close).toHaveBeenCalledWith(4001, "E2EE handshake failed");
+    });
+  });
+
   test("createWebSocketTransportFactory forwards sends when socket is open", () => {
     const addEventListener = vi.fn();
     const removeEventListener = vi.fn();

packages/server/src/server/websocket-server.ts

@@ -185,10 +185,14 @@ export class VoiceAssistantWebSocketServer {
       server,
       path: "/ws",
       verifyClient: ({ req }, callback) => {
-        const origin = req.headers.origin;
-        const requestHost = typeof req.headers.host === "string" ? req.headers.host : null;
+        const requestMetadata = extractSocketRequestMetadata(req);
+        const origin = requestMetadata.origin;
+        const requestHost = requestMetadata.host ?? null;
         if (requestHost && !isHostAllowed(requestHost, allowedHosts)) {
-          this.logger.warn({ host: requestHost }, "Rejected connection from disallowed host");
+          this.logger.warn(
+            { ...requestMetadata, host: requestHost },
+            "Rejected connection from disallowed host"
+          );
           callback(false, 403, "Host not allowed");
           return;
         }
@@ -200,7 +204,10 @@ export class VoiceAssistantWebSocketServer {
         if (!origin || allowedOrigins.has(origin) || sameOrigin) {
           callback(true);
         } else {
-          this.logger.warn({ origin }, "Rejected connection from origin");
+          this.logger.warn(
+            { ...requestMetadata, origin },
+            "Rejected connection from origin"
+          );
           callback(false, 403, "Origin not allowed");
         }
       },
@@ -282,7 +289,7 @@ export class VoiceAssistantWebSocketServer {
 
   private async attachSocket(
     ws: WebSocketLike,
-    _request?: unknown,
+    request?: unknown,
     metadata?: ExternalSocketMetadata
   ): Promise<void> {
     const externalSessionKey =
@@ -320,7 +327,24 @@ export class VoiceAssistantWebSocketServer {
     }
 
     const clientId = `client-${++this.clientIdCounter}`;
-    const connectionLogger = this.logger.child({ clientId });
+    const requestMetadata = extractSocketRequestMetadata(request);
+    const connectionLoggerFields: Record<string, string> = {
+      clientId,
+      transport: externalSessionKey ? "relay" : "direct",
+    };
+    if (requestMetadata.host) {
+      connectionLoggerFields.host = requestMetadata.host;
+    }
+    if (requestMetadata.origin) {
+      connectionLoggerFields.origin = requestMetadata.origin;
+    }
+    if (requestMetadata.userAgent) {
+      connectionLoggerFields.userAgent = requestMetadata.userAgent;
+    }
+    if (requestMetadata.remoteAddress) {
+      connectionLoggerFields.remoteAddress = requestMetadata.remoteAddress;
+    }
+    const connectionLogger = this.logger.child(connectionLoggerFields);
     const socketRef = { current: ws };
 
     const session = new Session({
@@ -510,17 +534,17 @@ export class VoiceAssistantWebSocketServer {
     data: Buffer | ArrayBuffer | Buffer[] | string
   ): Promise<void> {
     try {
+      const activeConnection = this.sessions.get(ws);
       const buffer = bufferFromWsData(data);
       const asBytes = asUint8Array(buffer);
       if (asBytes) {
         const frame = decodeBinaryMuxFrame(asBytes);
         if (frame) {
-          const connection = this.sessions.get(ws);
-          if (!connection) {
+          if (!activeConnection) {
             this.logger.error("No session found for client");
             return;
           }
-          connection.session.handleBinaryFrame(frame);
+          activeConnection.session.handleBinaryFrame(frame);
           return;
         }
       }
@@ -535,8 +559,10 @@ export class VoiceAssistantWebSocketServer {
           "type" in parsed &&
           (parsed as { type?: unknown }).type === "session";
 
-        this.logger.warn(
+        const log = activeConnection?.connectionLogger ?? this.logger;
+        log.warn(
           {
+            clientId: activeConnection?.clientId,
             requestId: requestInfo?.requestId,
             requestType: requestInfo?.requestType,
             error: parsedMessage.error.message,
@@ -585,14 +611,13 @@ export class VoiceAssistantWebSocketServer {
         return;
       }
 
-      const connection = this.sessions.get(ws);
-      if (!connection) {
+      if (!activeConnection) {
         this.logger.error("No session found for client");
         return;
       }
 
       if (message.type === "session") {
-        await connection.session.handleMessage(message.message);
+        await activeConnection.session.handleMessage(message.message);
       }
     } catch (error) {
       const err = error instanceof Error ? error : new Error(String(error));
@@ -751,6 +776,49 @@ export class VoiceAssistantWebSocketServer {
   }
 }
 
+type SocketRequestMetadata = {
+  host?: string;
+  origin?: string;
+  userAgent?: string;
+  remoteAddress?: string;
+};
+
+function extractSocketRequestMetadata(request: unknown): SocketRequestMetadata {
+  if (!request || typeof request !== "object") {
+    return {};
+  }
+
+  const record = request as {
+    headers?: {
+      host?: unknown;
+      origin?: unknown;
+      "user-agent"?: unknown;
+    };
+    socket?: {
+      remoteAddress?: unknown;
+    };
+  };
+
+  const host = typeof record.headers?.host === "string" ? record.headers.host : undefined;
+  const origin =
+    typeof record.headers?.origin === "string" ? record.headers.origin : undefined;
+  const userAgent =
+    typeof record.headers?.["user-agent"] === "string"
+      ? record.headers["user-agent"]
+      : undefined;
+  const remoteAddress =
+    typeof record.socket?.remoteAddress === "string"
+      ? record.socket.remoteAddress
+      : undefined;
+
+  return {
+    ...(host ? { host } : {}),
+    ...(origin ? { origin } : {}),
+    ...(userAgent ? { userAgent } : {}),
+    ...(remoteAddress ? { remoteAddress } : {}),
+  };
+}
+
 function stringifyCloseReason(reason: unknown): string | null {
   if (typeof reason === "string") {
     return reason.length > 0 ? reason : null;

packages/server/src/shared/messages.stream-parsing.test.ts

@@ -78,6 +78,28 @@ describe("shared messages stream parsing", () => {
     }
   });
 
+  it("rejects removed initialize_agent_request inbound payload", () => {
+    const parsed = SessionInboundMessageSchema.safeParse({
+      type: "initialize_agent_request",
+      agentId: "agent-legacy",
+      requestId: "req-legacy-1",
+    });
+    expect(parsed.success).toBe(false);
+  });
+
+  it("rejects removed initialize_agent_request outbound payload", () => {
+    const parsed = SessionOutboundMessageSchema.safeParse({
+      type: "initialize_agent_request",
+      payload: {
+        requestId: "req-legacy-1",
+        agentId: "agent-legacy",
+        agentStatus: "running",
+        timelineSize: 12,
+      },
+    });
+    expect(parsed.success).toBe(false);
+  });
+
   it("rejects websocket envelope for removed agent_stream_snapshot message type", () => {
     const fixture = {
       type: "agent_stream_snapshot",