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/app/src/desktop/daemon/desktop-daemon-transport.test.ts

@@ -1,5 +1,8 @@
 import { describe, expect, it, vi } from "vitest";
-import { createDesktopLocalDaemonTransportFactory } from "./desktop-daemon-transport";
+import {
+  createDesktopLocalDaemonTransportFactory,
+  createDesktopWebSocketTransportFactory,
+} from "./desktop-daemon-transport";
 import { createFakeLocalDaemonTransportRpc } from "./test-local-daemon-transport-rpc";
 
 const LOCAL_URL = "paseo+local://socket?path=%2Ftmp%2Fpaseo.sock";
@@ -43,4 +46,25 @@ describe("desktop-daemon-transport", () => {
     expect(rpc.closedSessions).toEqual(["local-session-2"]);
     expect(cleanup).toHaveBeenCalledTimes(1);
   });
+
+  it("forwards remote WebSocket headers and protocols through the desktop bridge", () => {
+    const rpc = createFakeLocalDaemonTransportRpc();
+    const transportFactory = createDesktopWebSocketTransportFactory(rpc);
+    expect(transportFactory).not.toBeNull();
+
+    transportFactory!({
+      url: "wss://daemon.example/ws",
+      headers: { "X-Tenant": "acme", Authorization: "Bearer secret" },
+      protocols: ["paseo.bearer.secret"],
+    });
+
+    expect(rpc.openCalls).toEqual([
+      {
+        transportType: "websocket",
+        url: "wss://daemon.example/ws",
+        headers: { "X-Tenant": "acme", Authorization: "Bearer secret" },
+        protocols: ["paseo.bearer.secret"],
+      },
+    ]);
+  });
 });

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

@@ -2,11 +2,13 @@ import type {
   DaemonTransport,
   DaemonTransportFactory,
 } from "@getpaseo/client/internal/daemon-client";
-import type { LocalTransportTarget } from "./desktop-daemon";
+import type { DesktopDaemonTransportTarget, LocalTransportTarget } from "./desktop-daemon";
 import {
   defaultLocalDaemonTransportRpc,
+  defaultWebSocketDaemonTransportRpc,
   type LocalDaemonTransportRpc,
 } from "./local-daemon-transport-rpc";
+import { isElectronRuntime } from "@/desktop/host";
 
 const LOCAL_TRANSPORT_SCHEME = "paseo+local:";
 
@@ -50,11 +52,12 @@ function parseLocalDaemonTransportUrl(url: string): LocalTransportTarget {
   };
 }
 
-export function createDesktopLocalDaemonTransportFactory(
-  rpc: LocalDaemonTransportRpc = defaultLocalDaemonTransportRpc,
-): DaemonTransportFactory | null {
-  return ({ url }) => {
-    const target = parseLocalDaemonTransportUrl(url);
+function createDesktopDaemonTransportFactory(
+  resolveTarget: (options: Parameters<DaemonTransportFactory>[0]) => DesktopDaemonTransportTarget,
+  rpc: LocalDaemonTransportRpc,
+): DaemonTransportFactory {
+  return (options) => {
+    const target = resolveTarget(options);
     let sessionId: string | null = null;
     let unlisten: (() => void) | null = null;
     let disposed = false;
@@ -187,3 +190,26 @@ export function createDesktopLocalDaemonTransportFactory(
     return transport;
   };
 }
+
+export function createDesktopLocalDaemonTransportFactory(
+  rpc: LocalDaemonTransportRpc = defaultLocalDaemonTransportRpc,
+): DaemonTransportFactory | null {
+  return createDesktopDaemonTransportFactory(({ url }) => parseLocalDaemonTransportUrl(url), rpc);
+}
+
+export function createDesktopWebSocketTransportFactory(
+  rpc?: LocalDaemonTransportRpc,
+): DaemonTransportFactory | null {
+  if (!rpc && !isElectronRuntime()) {
+    return null;
+  }
+  return createDesktopDaemonTransportFactory(
+    ({ url, headers, protocols }) => ({
+      transportType: "websocket" as const,
+      url,
+      ...(headers && Object.keys(headers).length > 0 ? { headers } : {}),
+      ...(protocols && protocols.length > 0 ? { protocols } : {}),
+    }),
+    rpc ?? defaultWebSocketDaemonTransportRpc,
+  );
+}

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

@@ -46,6 +46,16 @@ export interface LocalTransportTarget {
   transportPath: string;
 }
 
+export interface WebSocketTransportTarget {
+  [key: string]: unknown;
+  transportType: "websocket";
+  url: string;
+  headers?: Record<string, string>;
+  protocols?: string[];
+}
+
+export type DesktopDaemonTransportTarget = LocalTransportTarget | WebSocketTransportTarget;
+
 interface LocalTransportEventPayload {
   sessionId: string;
   kind: "open" | "message" | "close" | "error";
@@ -193,6 +203,16 @@ export async function openLocalTransportSession(target: LocalTransportTarget): P
   return raw;
 }
 
+export async function openWebSocketTransportSession(
+  target: WebSocketTransportTarget,
+): Promise<string> {
+  const raw = await invokeDesktopCommand<unknown>("open_websocket_daemon_transport", target);
+  if (typeof raw !== "string" || raw.trim().length === 0) {
+    throw new Error("Unexpected WebSocket transport session response.");
+  }
+  return raw;
+}
+
 export async function sendLocalTransportMessage(input: {
   sessionId: string;
   text?: string;

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

@@ -2,8 +2,10 @@ import {
   closeLocalTransportSession,
   listenToLocalTransportEvents,
   openLocalTransportSession,
+  openWebSocketTransportSession,
   sendLocalTransportMessage,
-  type LocalTransportTarget,
+  type DesktopDaemonTransportTarget,
+  type WebSocketTransportTarget,
 } from "./desktop-daemon";
 
 export interface LocalDaemonTransportEvent {
@@ -17,7 +19,7 @@ export interface LocalDaemonTransportEvent {
 }
 
 export interface LocalDaemonTransportRpc {
-  openSession(target: LocalTransportTarget): Promise<string>;
+  openSession(target: DesktopDaemonTransportTarget): Promise<string>;
   listenToEvents(handler: (event: LocalDaemonTransportEvent) => void): Promise<() => void>;
   sendMessage(input: { sessionId: string; text?: string; binaryBase64?: string }): Promise<void>;
   closeSession(sessionId: string): Promise<void>;
@@ -29,3 +31,10 @@ export const defaultLocalDaemonTransportRpc: LocalDaemonTransportRpc = {
   sendMessage: sendLocalTransportMessage,
   closeSession: closeLocalTransportSession,
 };
+
+export const defaultWebSocketDaemonTransportRpc: LocalDaemonTransportRpc = {
+  openSession: (target) => openWebSocketTransportSession(target as WebSocketTransportTarget),
+  listenToEvents: listenToLocalTransportEvents,
+  sendMessage: sendLocalTransportMessage,
+  closeSession: closeLocalTransportSession,
+};

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

@@ -1,4 +1,4 @@
-import type { LocalTransportTarget } from "./desktop-daemon";
+import type { DesktopDaemonTransportTarget } from "./desktop-daemon";
 import type {
   LocalDaemonTransportEvent,
   LocalDaemonTransportRpc,
@@ -11,7 +11,7 @@ export interface RecordedSend {
 }
 
 export interface FakeLocalDaemonTransportRpc extends LocalDaemonTransportRpc {
-  readonly openCalls: LocalTransportTarget[];
+  readonly openCalls: DesktopDaemonTransportTarget[];
   readonly recordedSends: RecordedSend[];
   readonly closedSessions: string[];
   resolveOpen(sessionId: string): void;
@@ -22,7 +22,7 @@ export interface FakeLocalDaemonTransportRpc extends LocalDaemonTransportRpc {
 }
 
 export function createFakeLocalDaemonTransportRpc(): FakeLocalDaemonTransportRpc {
-  const openCalls: LocalTransportTarget[] = [];
+  const openCalls: DesktopDaemonTransportTarget[] = [];
   const recordedSends: RecordedSend[] = [];
   const closedSessions: string[] = [];
   let eventHandler: ((event: LocalDaemonTransportEvent) => void) | null = null;