Add custom headers to direct connections (#2922)

* Add custom headers to direct connections

* Address custom header review feedback

* Fix direct connection header integration

* Fix desktop relay E2E transport setup

---------

Co-authored-by: Claude <claude@Claudes-Mac-mini.local>
Co-authored-by: Mohamed Boudra <boudra.moha@gmail.com>

packages/desktop/src/daemon/daemon-manager.ts

@@ -26,8 +26,10 @@ import {
 import { createSkillsCommandHandlers, getSkillsController } from "../integrations/skills/index.js";
 import {
   openLocalTransportSession,
+  openWebSocketTransportSession,
   sendLocalTransportMessage,
   closeLocalTransportSession,
+  type WebSocketTransportTarget,
 } from "./local-transport.js";
 import { createNodeEntrypointInvocation, resolveDaemonRunnerEntrypoint } from "./runtime-paths.js";
 import { runExternalCliJsonCommand, runExternalCliTextCommand } from "./cli/external.js";
@@ -537,6 +539,9 @@ export function createDaemonCommandHandlers(): Record<string, DesktopCommandHand
       const target = args as { transportType: "socket" | "pipe"; transportPath: string };
       return await openLocalTransportSession(target);
     },
+    open_websocket_daemon_transport: async (args) => {
+      return await openWebSocketTransportSession(args as unknown as WebSocketTransportTarget);
+    },
     send_local_daemon_transport_message: async (args) => {
       await sendLocalTransportMessage(
         args as { sessionId: string; text?: string; binaryBase64?: string },

packages/desktop/src/daemon/local-transport.test.ts

@@ -0,0 +1,86 @@
+import { createServer } from "node:http";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { WebSocketServer } from "ws";
+
+const send = vi.fn();
+
+vi.mock("electron", () => ({
+  BrowserWindow: {
+    getAllWindows: () => [{ webContents: { send } }],
+  },
+}));
+
+import { closeAllTransportSessions, openWebSocketTransportSession } from "./local-transport";
+
+afterEach(() => {
+  closeAllTransportSessions();
+  send.mockReset();
+});
+
+describe("desktop WebSocket transport", () => {
+  it("attaches custom headers when opening a remote daemon session", async () => {
+    const server = createServer();
+    const webSocketServer = new WebSocketServer({ server });
+    const requestHeaders = new Promise<Record<string, string | string[] | undefined>>((resolve) => {
+      webSocketServer.once("connection", (_socket, request) => resolve(request.headers));
+    });
+
+    await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
+    const address = server.address();
+    if (!address || typeof address === "string") {
+      throw new Error("Expected TCP test server address.");
+    }
+
+    try {
+      const sessionId = await openWebSocketTransportSession({
+        transportType: "websocket",
+        url: `ws://127.0.0.1:${address.port}/ws`,
+        headers: { "X-Tenant": "acme" },
+      });
+
+      expect((await requestHeaders)["x-tenant"]).toBe("acme");
+      expect(sessionId).toMatch(/^local-session-/);
+    } finally {
+      closeAllTransportSessions();
+      await new Promise<void>((resolve, reject) => {
+        webSocketServer.close((error) => (error ? reject(error) : resolve()));
+      });
+      await new Promise<void>((resolve, reject) => {
+        server.close((error) => {
+          if (error) {
+            reject(error);
+            return;
+          }
+          resolve();
+        });
+      });
+    }
+  });
+
+  it("rejects non-WebSocket URLs", () => {
+    expect(() =>
+      openWebSocketTransportSession({
+        transportType: "websocket",
+        url: "https://daemon.example/ws",
+      }),
+    ).toThrow("Unsupported WebSocket transport URL");
+  });
+
+  it("rejects malformed custom headers before opening a session", () => {
+    expect(() =>
+      openWebSocketTransportSession({
+        transportType: "websocket",
+        url: "ws://daemon.example/ws",
+        headers: { "Bad Header": "value" },
+      }),
+    ).toThrow("Invalid WebSocket transport header");
+
+    expect(() =>
+      openWebSocketTransportSession({
+        transportType: "websocket",
+        url: "ws://daemon.example/ws",
+        headers: { "X-Tenant": "acme\r\nX-Injected: true" },
+      }),
+    ).toThrow("Invalid WebSocket transport header");
+  });
+});

packages/desktop/src/daemon/local-transport.ts

@@ -6,6 +6,13 @@ interface TransportTarget {
   transportPath: string;
 }
 
+export interface WebSocketTransportTarget {
+  transportType: "websocket";
+  url: string;
+  headers?: Record<string, string>;
+  protocols?: string[];
+}
+
 interface TransportEventPayload {
   sessionId: string;
   kind: "open" | "message" | "close" | "error";
@@ -23,6 +30,7 @@ interface Session {
 }
 
 const WS_ENDPOINT_PATH = "/ws";
+const HTTP_HEADER_NAME_PATTERN = /^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$/;
 
 let nextSessionId = 0;
 const sessions = new Map<string, Session>();
@@ -52,6 +60,28 @@ function describeTransportTarget(target: TransportTarget): string {
   return target.transportType === "pipe" ? "local daemon pipe" : "local daemon socket";
 }
 
+function normalizeWebSocketTarget(target: WebSocketTransportTarget): WebSocketTransportTarget {
+  const parsed = new URL(target.url);
+  if (parsed.protocol !== "ws:" && parsed.protocol !== "wss:") {
+    throw new Error(`Unsupported WebSocket transport URL: ${target.url}`);
+  }
+
+  const headers: Record<string, string> = {};
+  for (const [name, value] of Object.entries(target.headers ?? {})) {
+    if (!HTTP_HEADER_NAME_PATTERN.test(name) || typeof value !== "string" || /[\r\n]/.test(value)) {
+      throw new Error("Invalid WebSocket transport header.");
+    }
+    headers[name] = value;
+  }
+
+  return {
+    transportType: "websocket",
+    url: parsed.toString(),
+    ...(Object.keys(headers).length > 0 ? { headers } : {}),
+    ...(target.protocols?.length ? { protocols: [...target.protocols] } : {}),
+  };
+}
+
 function decodeTransportMessage(input: { text?: string; binaryBase64?: string }): string | Buffer {
   if (typeof input.text === "string") {
     return input.text;
@@ -65,11 +95,28 @@ function decodeTransportMessage(input: { text?: string; binaryBase64?: string })
 }
 
 export function openLocalTransportSession(target: TransportTarget): Promise<string> {
+  return openTransportSession(buildLocalWebSocketUrl(target), describeTransportTarget(target));
+}
+
+export function openWebSocketTransportSession(target: WebSocketTransportTarget): Promise<string> {
+  const normalized = normalizeWebSocketTarget(target);
+  return openTransportSession(normalized.url, "remote daemon", {
+    headers: normalized.headers,
+    protocols: normalized.protocols,
+  });
+}
+
+function openTransportSession(
+  url: string,
+  targetDescription: string,
+  options?: { headers?: Record<string, string>; protocols?: string[] },
+): Promise<string> {
   const sessionId = `local-session-${++nextSessionId}`;
-  const url = buildLocalWebSocketUrl(target);
 
   return new Promise((resolve, reject) => {
-    const ws = new WebSocket(url);
+    const ws = options?.protocols
+      ? new WebSocket(url, options.protocols, { headers: options.headers })
+      : new WebSocket(url, { headers: options?.headers });
     const session: Session = {
       id: sessionId,
       ws,
@@ -121,9 +168,7 @@ export function openLocalTransportSession(target: TransportTarget): Promise<stri
       sessions.delete(sessionId);
 
       if (!openSettled) {
-        finalizeOpenFailure(
-          `${describeTransportTarget(target)} closed before the session became ready.`,
-        );
+        finalizeOpenFailure(`${targetDescription} closed before the session became ready.`);
         return;
       }
 
@@ -139,9 +184,7 @@ export function openLocalTransportSession(target: TransportTarget): Promise<stri
 
     ws.on("error", (err: Error) => {
       if (!openSettled) {
-        finalizeOpenFailure(
-          `Failed to connect to ${describeTransportTarget(target)}: ${err.message}`,
-        );
+        finalizeOpenFailure(`Failed to connect to ${targetDescription}: ${err.message}`);
         return;
       }