feat: add relay reconnect with grace period and branch suggestions

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

@@ -287,6 +287,65 @@ describe("DaemonClient", () => {
     );
   });
 
+  test("requests branch suggestions via RPC", async () => {
+    const logger = createMockLogger();
+    const mock = createMockTransport();
+
+    const client = new DaemonClient({
+      url: "ws://test",
+      logger,
+      reconnect: { enabled: false },
+      transportFactory: () => mock.transport,
+    });
+    clients.push(client);
+
+    const connectPromise = client.connect();
+    mock.triggerOpen();
+    await connectPromise;
+
+    const promise = client.getBranchSuggestions(
+      { cwd: "/tmp/project", query: "mai", limit: 5 },
+      "req-branches"
+    );
+
+    expect(mock.sent).toHaveLength(1);
+    const request = JSON.parse(mock.sent[0]) as {
+      type: "session";
+      message: {
+        type: "branch_suggestions_request";
+        cwd: string;
+        query?: string;
+        limit?: number;
+        requestId: string;
+      };
+    };
+    expect(request.message.type).toBe("branch_suggestions_request");
+    expect(request.message.cwd).toBe("/tmp/project");
+    expect(request.message.query).toBe("mai");
+    expect(request.message.limit).toBe(5);
+    expect(request.message.requestId).toBe("req-branches");
+
+    mock.triggerMessage(
+      JSON.stringify({
+        type: "session",
+        message: {
+          type: "branch_suggestions_response",
+          payload: {
+            branches: ["main"],
+            error: null,
+            requestId: "req-branches",
+          },
+        },
+      })
+    );
+
+    await expect(promise).resolves.toEqual({
+      branches: ["main"],
+      error: null,
+      requestId: "req-branches",
+    });
+  });
+
   test("resubscribes checkout diff streams after reconnect", async () => {
     const logger = createMockLogger();
     const mock = createMockTransport();

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

@@ -24,6 +24,7 @@ import type {
   CheckoutPrCreateResponse,
   CheckoutPrStatusResponse,
   ValidateBranchResponse,
+  BranchSuggestionsResponse,
   PaseoWorktreeListResponse,
   PaseoWorktreeArchiveResponse,
   ProjectIconResponse,
@@ -191,6 +192,7 @@ type CheckoutPushPayload = CheckoutPushResponse["payload"];
 type CheckoutPrCreatePayload = CheckoutPrCreateResponse["payload"];
 type CheckoutPrStatusPayload = CheckoutPrStatusResponse["payload"];
 type ValidateBranchPayload = ValidateBranchResponse["payload"];
+type BranchSuggestionsPayload = BranchSuggestionsResponse["payload"];
 type PaseoWorktreeListPayload = PaseoWorktreeListResponse["payload"];
 type PaseoWorktreeArchivePayload = PaseoWorktreeArchiveResponse["payload"];
 type FileExplorerPayload = FileExplorerResponse["payload"];
@@ -1936,6 +1938,35 @@ export class DaemonClient {
     });
   }
 
+  async getBranchSuggestions(
+    options: { cwd: string; query?: string; limit?: number },
+    requestId?: string
+  ): Promise<BranchSuggestionsPayload> {
+    const resolvedRequestId = this.createRequestId(requestId);
+    const message = SessionInboundMessageSchema.parse({
+      type: "branch_suggestions_request",
+      cwd: options.cwd,
+      query: options.query,
+      limit: options.limit,
+      requestId: resolvedRequestId,
+    });
+    return this.sendRequest({
+      requestId: resolvedRequestId,
+      message,
+      timeout: 10000,
+      options: { skipQueue: true },
+      select: (msg) => {
+        if (msg.type !== "branch_suggestions_response") {
+          return null;
+        }
+        if (msg.payload.requestId !== resolvedRequestId) {
+          return null;
+        }
+        return msg.payload;
+      },
+    });
+  }
+
   // ============================================================================
   // File Explorer
   // ============================================================================

packages/server/src/server/bootstrap.ts

@@ -538,11 +538,11 @@ export async function createPaseoDaemon(
               relayTransport?.stop().catch(() => undefined);
               relayTransport = startRelayTransport({
                 logger,
-                attachSocket: (ws) => {
+                attachSocket: (ws, metadata) => {
                   if (!wsServer) {
                     throw new Error("WebSocket server not initialized");
                   }
-                  return wsServer.attachExternalSocket(ws);
+                  return wsServer.attachExternalSocket(ws, metadata);
                 },
                 relayEndpoint,
                 serverId,

packages/server/src/server/relay-transport.test.ts

@@ -188,4 +188,36 @@ describe("relay-transport control lifecycle", () => {
     expect(hasLogMessage(logger.warn, "relay_control_stale_terminating")).toBe(true);
     expect(control.terminateCalls).toBe(1);
   });
+
+  test("passes stable relay external session metadata when attaching data socket", async () => {
+    const logger = createMockLogger();
+    const attachSocket = vi.fn(async () => {});
+    const controller = startRelayTransport({
+      logger: logger as any,
+      attachSocket,
+      relayEndpoint: "relay.paseo.sh:443",
+      serverId: "srv_test",
+    });
+    controllers.push(controller);
+
+    const control = MockWebSocket.instances[0];
+    control.open();
+    control.message(JSON.stringify({ type: "pong", ts: Date.now() }));
+    control.message(JSON.stringify({ type: "client_connected", clientId: "clt_test" }));
+
+    const dataSocket = MockWebSocket.instances[1];
+    expect(dataSocket).toBeDefined();
+    dataSocket.open();
+
+    await Promise.resolve();
+
+    expect(attachSocket).toHaveBeenCalledTimes(1);
+    expect(attachSocket).toHaveBeenCalledWith(
+      dataSocket,
+      {
+        transport: "relay",
+        externalSessionKey: "relay:clt_test",
+      }
+    );
+  });
 });

packages/server/src/server/relay-transport.ts

@@ -9,10 +9,11 @@ import {
   type KeyPair,
 } from "@getpaseo/relay/e2ee";
 import { buildRelayWebSocketUrl } from "../shared/daemon-endpoints.js";
+import type { ExternalSocketMetadata } from "./websocket-server.js";
 
 type RelayTransportOptions = {
   logger: pino.Logger;
-  attachSocket: (ws: RelaySocketLike) => Promise<void>;
+  attachSocket: (ws: RelaySocketLike, metadata?: ExternalSocketMetadata) => Promise<void>;
   relayEndpoint: string; // "host:port"
   serverId: string;
   daemonKeyPair?: KeyPair;
@@ -321,10 +322,20 @@ export function startRelayTransport({
       relayLogger.info({ url, clientId }, "relay_data_connected");
       if (attached) return;
       attached = true;
+      const externalMetadata: ExternalSocketMetadata = {
+        transport: "relay",
+        externalSessionKey: `relay:${clientId}`,
+      };
       if (daemonKeyPair) {
-        void attachEncryptedSocket(socket, daemonKeyPair, relayLogger.child({ clientId }), attachSocket);
+        void attachEncryptedSocket(
+          socket,
+          daemonKeyPair,
+          relayLogger.child({ clientId }),
+          attachSocket,
+          externalMetadata
+        );
       } else {
-        void attachSocket(socket);
+        void attachSocket(socket, externalMetadata);
       }
     });
 
@@ -353,7 +364,8 @@ async function attachEncryptedSocket(
   socket: WebSocket,
   daemonKeyPair: KeyPair,
   logger: pino.Logger,
-  attachSocket: (ws: RelaySocketLike) => Promise<void>
+  attachSocket: (ws: RelaySocketLike, metadata?: ExternalSocketMetadata) => Promise<void>,
+  metadata?: ExternalSocketMetadata
 ): Promise<void> {
   try {
     const relayTransport = createRelayTransportAdapter(socket);
@@ -367,7 +379,7 @@ async function attachEncryptedSocket(
       },
     });
     const encryptedSocket = createEncryptedSocket(channel, emitter);
-    await attachSocket(encryptedSocket);
+    await attachSocket(encryptedSocket, metadata);
   } catch (error) {
     logger.warn({ err: error }, "relay_e2ee_handshake_failed");
     try {

packages/server/src/server/session.ts

@@ -103,6 +103,7 @@ import {
   getCheckoutDiff,
   getCheckoutStatus,
   getCheckoutStatusLite,
+  listBranchSuggestions,
   NotGitRepoError,
   MergeConflictError,
   MergeFromBaseConflictError,
@@ -1232,6 +1233,10 @@ export class Session {
           await this.handleValidateBranchRequest(msg);
           break;
 
+        case "branch_suggestions_request":
+          await this.handleBranchSuggestionsRequest(msg);
+          break;
+
         case "subscribe_checkout_diff_request":
           await this.handleSubscribeCheckoutDiffRequest(msg);
           break;
@@ -3589,6 +3594,34 @@ export class Session {
     }
   }
 
+  private async handleBranchSuggestionsRequest(
+    msg: Extract<SessionInboundMessage, { type: "branch_suggestions_request" }>
+  ): Promise<void> {
+    const { cwd, query, limit, requestId } = msg;
+
+    try {
+      const resolvedCwd = expandTilde(cwd);
+      const branches = await listBranchSuggestions(resolvedCwd, { query, limit });
+      this.emit({
+        type: "branch_suggestions_response",
+        payload: {
+          branches,
+          error: null,
+          requestId,
+        },
+      });
+    } catch (error) {
+      this.emit({
+        type: "branch_suggestions_response",
+        payload: {
+          branches: [],
+          error: error instanceof Error ? error.message : String(error),
+          requestId,
+        },
+      });
+    }
+  }
+
   private normalizeCheckoutDiffCompare(compare: CheckoutDiffCompareInput): CheckoutDiffCompareInput {
     if (compare.mode === "uncommitted") {
       return { mode: "uncommitted" };

packages/server/src/server/websocket-server.relay-reconnect.test.ts

@@ -0,0 +1,199 @@
+import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
+
+const wsModuleMock = vi.hoisted(() => {
+  class MockWebSocketServer {
+    static instances: MockWebSocketServer[] = [];
+    readonly handlers = new Map<string, (...args: any[]) => void>();
+
+    constructor(_options: unknown) {
+      MockWebSocketServer.instances.push(this);
+    }
+
+    on(event: string, handler: (...args: any[]) => void) {
+      this.handlers.set(event, handler);
+      return this;
+    }
+
+    close() {
+      // no-op
+    }
+  }
+
+  return { MockWebSocketServer };
+});
+
+const sessionMock = vi.hoisted(() => {
+  const instances: MockSession[] = [];
+
+  class MockSession {
+    cleanup = vi.fn(async () => {});
+    handleMessage = vi.fn(async () => {});
+    getClientActivity = vi.fn(() => null);
+    readonly args: Record<string, unknown>;
+
+    constructor(args: Record<string, unknown>) {
+      this.args = args;
+      instances.push(this);
+    }
+  }
+
+  return { MockSession, instances };
+});
+
+vi.mock("ws", () => ({
+  WebSocketServer: wsModuleMock.MockWebSocketServer,
+}));
+
+vi.mock("./session.js", () => ({
+  Session: sessionMock.MockSession,
+}));
+
+vi.mock("./push/token-store.js", () => ({
+  PushTokenStore: class {
+    getAllTokens(): string[] {
+      return [];
+    }
+  },
+}));
+
+vi.mock("./push/push-service.js", () => ({
+  PushService: class {
+    async sendPush(): Promise<void> {
+      // no-op
+    }
+  },
+}));
+
+import {
+  VoiceAssistantWebSocketServer,
+  type ExternalSocketMetadata,
+} from "./websocket-server";
+
+class MockSocket {
+  readyState = 1;
+  sent: string[] = [];
+  private listeners = new Map<string, Array<(...args: any[]) => void>>();
+
+  on(event: "message" | "close" | "error", listener: (...args: any[]) => void): void {
+    const handlers = this.listeners.get(event) ?? [];
+    handlers.push(listener);
+    this.listeners.set(event, handlers);
+  }
+
+  once(event: "close" | "error", listener: (...args: any[]) => void): void {
+    const wrapped = (...args: any[]) => {
+      this.off(event, wrapped);
+      listener(...args);
+    };
+    this.on(event, wrapped);
+  }
+
+  send(data: string): void {
+    this.sent.push(data);
+  }
+
+  close(code?: number, reason?: string): void {
+    this.readyState = 3;
+    this.emit("close", code ?? 1000, reason ?? "");
+  }
+
+  emit(event: "message" | "close" | "error", ...args: any[]): void {
+    const handlers = this.listeners.get(event) ?? [];
+    for (const handler of [...handlers]) {
+      handler(...args);
+    }
+  }
+
+  private off(event: "close" | "error", listener: (...args: any[]) => void): void {
+    const handlers = this.listeners.get(event) ?? [];
+    this.listeners.set(
+      event,
+      handlers.filter((handler) => handler !== listener)
+    );
+  }
+}
+
+function createLogger() {
+  const logger = {
+    child: vi.fn(() => logger),
+    trace: vi.fn(),
+    debug: vi.fn(),
+    info: vi.fn(),
+    warn: vi.fn(),
+    error: vi.fn(),
+  };
+  return logger;
+}
+
+function createServer() {
+  return new VoiceAssistantWebSocketServer(
+    {} as any,
+    createLogger() as any,
+    "srv_test",
+    {
+      setAgentAttentionCallback: vi.fn(),
+      getAgent: vi.fn(() => null),
+    } as any,
+    {} as any,
+    {} as any,
+    "/tmp/paseo-test",
+    async () => ({} as any),
+    { allowedOrigins: new Set() }
+  );
+}
+
+describe("relay external socket reconnect behavior", () => {
+  beforeEach(() => {
+    sessionMock.instances.length = 0;
+    vi.useFakeTimers();
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  test("keeps the same session when relay reconnects within grace window", async () => {
+    const server = createServer();
+    const metadata: ExternalSocketMetadata = {
+      transport: "relay",
+      externalSessionKey: "relay:client-1",
+    };
+
+    const socket1 = new MockSocket();
+    await server.attachExternalSocket(socket1, metadata);
+    expect(sessionMock.instances).toHaveLength(1);
+    const session = sessionMock.instances[0]!;
+
+    socket1.emit("close", 1006, "");
+    await vi.advanceTimersByTimeAsync(1_000);
+    expect(session.cleanup).not.toHaveBeenCalled();
+
+    const socket2 = new MockSocket();
+    await server.attachExternalSocket(socket2, metadata);
+    expect(sessionMock.instances).toHaveLength(1);
+
+    await vi.advanceTimersByTimeAsync(20_000);
+    expect(session.cleanup).not.toHaveBeenCalled();
+
+    await server.close();
+  });
+
+  test("cleans up relay session when reconnect grace expires", async () => {
+    const server = createServer();
+    const metadata: ExternalSocketMetadata = {
+      transport: "relay",
+      externalSessionKey: "relay:client-2",
+    };
+
+    const socket1 = new MockSocket();
+    await server.attachExternalSocket(socket1, metadata);
+    expect(sessionMock.instances).toHaveLength(1);
+    const session = sessionMock.instances[0]!;
+
+    socket1.emit("close", 1006, "");
+    await vi.advanceTimersByTimeAsync(90_000);
+    expect(session.cleanup).toHaveBeenCalledTimes(1);
+
+    await server.close();
+  });
+});

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

@@ -29,6 +29,10 @@ import type {
 } from "./voice-types.js";
 
 export type AgentMcpTransportFactory = () => Promise<Transport>;
+export type ExternalSocketMetadata = {
+  transport: "relay";
+  externalSessionKey: string;
+};
 
 type WebSocketServerConfig = {
   allowedOrigins: Set<string>;
@@ -56,13 +60,25 @@ type WebSocketLike = {
   once: (event: "close" | "error", listener: (...args: any[]) => void) => void;
 };
 
+type SessionConnection = {
+  session: Session;
+  clientId: string;
+  connectionLogger: pino.Logger;
+  socketRef: { current: WebSocketLike };
+  externalSessionKey: string | null;
+  externalDisconnectCleanupTimeout: ReturnType<typeof setTimeout> | null;
+};
+
+const EXTERNAL_SESSION_DISCONNECT_GRACE_MS = 90_000;
+
 /**
  * WebSocket server that only accepts sockets + parses/forwards messages to the session layer.
  */
 export class VoiceAssistantWebSocketServer {
   private readonly logger: pino.Logger;
   private readonly wss: WebSocketServer;
-  private readonly sessions: Map<WebSocketLike, Session> = new Map();
+  private readonly sessions: Map<WebSocketLike, SessionConnection> = new Map();
+  private readonly externalSessionsByKey: Map<string, SessionConnection> = new Map();
   private clientIdCounter = 0;
   private readonly serverId: string;
   private readonly agentManager: AgentManager;
@@ -191,15 +207,27 @@ export class VoiceAssistantWebSocketServer {
   }
 
   public async attachExternalSocket(
-    ws: WebSocketLike
+    ws: WebSocketLike,
+    metadata?: ExternalSocketMetadata
   ): Promise<void> {
-    await this.attachSocket(ws);
+    await this.attachSocket(ws, undefined, metadata);
   }
 
   public async close(): Promise<void> {
+    const uniqueConnections = new Set<SessionConnection>([
+      ...this.sessions.values(),
+      ...this.externalSessionsByKey.values(),
+    ]);
+
     const cleanupPromises: Promise<void>[] = [];
-    for (const [ws, session] of this.sessions) {
-      cleanupPromises.push(session.cleanup());
+    for (const connection of uniqueConnections) {
+      if (connection.externalDisconnectCleanupTimeout) {
+        clearTimeout(connection.externalDisconnectCleanupTimeout);
+        connection.externalDisconnectCleanupTimeout = null;
+      }
+
+      const ws = connection.socketRef.current;
+      cleanupPromises.push(connection.session.cleanup());
       cleanupPromises.push(
         new Promise<void>((resolve) => {
           // WebSocket.CLOSED = 3
@@ -214,6 +242,7 @@ export class VoiceAssistantWebSocketServer {
     }
     await Promise.all(cleanupPromises);
     this.sessions.clear();
+    this.externalSessionsByKey.clear();
     this.wss.close();
   }
 
@@ -224,14 +253,53 @@ export class VoiceAssistantWebSocketServer {
     }
   }
 
-  private async attachSocket(ws: WebSocketLike, _request?: unknown): Promise<void> {
+  private async attachSocket(
+    ws: WebSocketLike,
+    _request?: unknown,
+    metadata?: ExternalSocketMetadata
+  ): Promise<void> {
+    const externalSessionKey =
+      metadata?.transport === "relay" && metadata.externalSessionKey.trim().length > 0
+        ? metadata.externalSessionKey
+        : null;
+
+    if (externalSessionKey) {
+      const existing = this.externalSessionsByKey.get(externalSessionKey);
+      if (existing) {
+        if (existing.externalDisconnectCleanupTimeout) {
+          clearTimeout(existing.externalDisconnectCleanupTimeout);
+          existing.externalDisconnectCleanupTimeout = null;
+        }
+
+        const previousSocket = existing.socketRef.current;
+        if (previousSocket !== ws) {
+          this.sessions.delete(previousSocket);
+          existing.socketRef.current = ws;
+        }
+
+        this.sessions.set(ws, existing);
+        this.sendServerInfo(ws);
+        existing.connectionLogger.trace(
+          {
+            clientId: existing.clientId,
+            externalSessionKey,
+            totalSessions: this.sessions.size,
+          },
+          "Client reconnected"
+        );
+        this.bindSocketHandlers(ws, existing);
+        return;
+      }
+    }
+
     const clientId = `client-${++this.clientIdCounter}`;
     const connectionLogger = this.logger.child({ clientId });
+    const socketRef = { current: ws };
 
     const session = new Session({
       clientId,
       onMessage: (msg) => {
-        this.sendToClient(ws, wrapSessionMessage(msg));
+        this.sendToClient(socketRef.current, wrapSessionMessage(msg));
       },
       logger: connectionLogger.child({ module: "session" }),
       downloadTokenStore: this.downloadTokenStore,
@@ -264,8 +332,31 @@ export class VoiceAssistantWebSocketServer {
       agentProviderRuntimeSettings: this.agentProviderRuntimeSettings,
     });
 
-    this.sessions.set(ws, session);
+    const connection: SessionConnection = {
+      session,
+      clientId,
+      connectionLogger,
+      socketRef,
+      externalSessionKey,
+      externalDisconnectCleanupTimeout: null,
+    };
 
+    this.sessions.set(ws, connection);
+    if (externalSessionKey) {
+      this.externalSessionsByKey.set(externalSessionKey, connection);
+    }
+
+    this.sendServerInfo(ws);
+
+    connectionLogger.trace(
+      { clientId, externalSessionKey, totalSessions: this.sessions.size },
+      "Client connected"
+    );
+
+    this.bindSocketHandlers(ws, connection);
+  }
+
+  private sendServerInfo(ws: WebSocketLike): void {
     // Advertise stable server identity immediately on connect (used for URL/shareable IDs).
     this.sendToClient(
       ws,
@@ -278,24 +369,27 @@ export class VoiceAssistantWebSocketServer {
         },
       })
     );
+  }
 
-    connectionLogger.trace(
-      { clientId, totalSessions: this.sessions.size },
-      "Client connected"
-    );
-
+  private bindSocketHandlers(
+    ws: WebSocketLike,
+    connection: SessionConnection
+  ): void {
     ws.on("message", (data) => {
       void this.handleRawMessage(ws, data);
     });
 
-    ws.on("close", async () => {
-      await this.detachSocket(ws, connectionLogger, clientId);
+    ws.on("close", async (code: number, reason: unknown) => {
+      await this.detachSocket(ws, connection, {
+        code: typeof code === "number" ? code : undefined,
+        reason,
+      });
     });
 
     ws.on("error", async (error) => {
       const err = error instanceof Error ? error : new Error(String(error));
-      connectionLogger.error({ err }, "Client error");
-      await this.detachSocket(ws, connectionLogger, clientId);
+      connection.connectionLogger.error({ err }, "Client error");
+      await this.detachSocket(ws, connection, { error: err });
     });
   }
 
@@ -313,19 +407,72 @@ export class VoiceAssistantWebSocketServer {
 
   private async detachSocket(
     ws: WebSocketLike,
-    connectionLogger: pino.Logger,
-    clientId: string
+    connection: SessionConnection,
+    details: {
+      code?: number;
+      reason?: unknown;
+      error?: Error;
+    }
   ): Promise<void> {
-    const session = this.sessions.get(ws);
-    if (!session) return;
-
-    connectionLogger.trace(
-      { clientId, totalSessions: this.sessions.size - 1 },
-      "Client disconnected"
-    );
-
-    await session.cleanup();
+    const activeConnection = this.sessions.get(ws);
+    if (activeConnection !== connection) return;
     this.sessions.delete(ws);
+
+    if (
+      connection.externalSessionKey &&
+      connection.socketRef.current === ws
+    ) {
+      if (connection.externalDisconnectCleanupTimeout) {
+        clearTimeout(connection.externalDisconnectCleanupTimeout);
+      }
+      const timeout = setTimeout(() => {
+        if (connection.externalDisconnectCleanupTimeout !== timeout) {
+          return;
+        }
+        connection.externalDisconnectCleanupTimeout = null;
+        void this.cleanupConnection(connection, "Client disconnected (grace timeout)");
+      }, EXTERNAL_SESSION_DISCONNECT_GRACE_MS);
+      connection.externalDisconnectCleanupTimeout = timeout;
+
+      connection.connectionLogger.trace(
+        {
+          clientId: connection.clientId,
+          externalSessionKey: connection.externalSessionKey,
+          code: details.code,
+          reason: stringifyCloseReason(details.reason),
+          reconnectGraceMs: EXTERNAL_SESSION_DISCONNECT_GRACE_MS,
+        },
+        "Client disconnected; waiting for reconnect"
+      );
+      return;
+    }
+
+    await this.cleanupConnection(connection, "Client disconnected");
+  }
+
+  private async cleanupConnection(
+    connection: SessionConnection,
+    logMessage: string
+  ): Promise<void> {
+    if (connection.externalDisconnectCleanupTimeout) {
+      clearTimeout(connection.externalDisconnectCleanupTimeout);
+      connection.externalDisconnectCleanupTimeout = null;
+    }
+
+    const currentSocket = connection.socketRef.current;
+    this.sessions.delete(currentSocket);
+    if (connection.externalSessionKey) {
+      const existing = this.externalSessionsByKey.get(connection.externalSessionKey);
+      if (existing === connection) {
+        this.externalSessionsByKey.delete(connection.externalSessionKey);
+      }
+    }
+
+    connection.connectionLogger.trace(
+      { clientId: connection.clientId, totalSessions: this.sessions.size },
+      logMessage
+    );
+    await connection.session.cleanup();
   }
 
   private async handleRawMessage(
@@ -395,14 +542,14 @@ export class VoiceAssistantWebSocketServer {
         return;
       }
 
-      const session = this.sessions.get(ws);
-      if (!session) {
+      const connection = this.sessions.get(ws);
+      if (!connection) {
         this.logger.error("No session found for client");
         return;
       }
 
       if (message.type === "session") {
-        await session.handleMessage(message.message);
+        await connection.session.handleMessage(message.message);
       }
     } catch (error) {
       const err = error instanceof Error ? error : new Error(String(error));
@@ -555,10 +702,10 @@ export class VoiceAssistantWebSocketServer {
       };
     }> = [];
 
-    for (const [ws, session] of this.sessions) {
+    for (const [ws, connection] of this.sessions) {
       clientEntries.push({
         ws,
-        state: this.getClientActivityState(session),
+        state: this.getClientActivityState(connection.session),
       });
     }
 
@@ -626,6 +773,21 @@ export class VoiceAssistantWebSocketServer {
   }
 }
 
+function stringifyCloseReason(reason: unknown): string | null {
+  if (typeof reason === "string") {
+    return reason.length > 0 ? reason : null;
+  }
+  if (Buffer.isBuffer(reason)) {
+    const text = reason.toString();
+    return text.length > 0 ? text : null;
+  }
+  if (reason == null) {
+    return null;
+  }
+  const text = String(reason);
+  return text.length > 0 ? text : null;
+}
+
 function extractRequestInfoFromUnknownWsInbound(
   payload: unknown
 ): { requestId: string; requestType?: string } | null {

packages/server/src/shared/messages.ts

@@ -791,6 +791,14 @@ export const ValidateBranchRequestSchema = z.object({
   requestId: z.string(),
 });
 
+export const BranchSuggestionsRequestSchema = z.object({
+  type: z.literal("branch_suggestions_request"),
+  cwd: z.string(),
+  query: z.string().optional(),
+  limit: z.number().int().min(1).max(200).optional(),
+  requestId: z.string(),
+});
+
 export const PaseoWorktreeListRequestSchema = z.object({
   type: z.literal("paseo_worktree_list_request"),
   cwd: z.string().optional(),
@@ -1015,6 +1023,7 @@ export const SessionInboundMessageSchema = z.discriminatedUnion("type", [
   CheckoutPrCreateRequestSchema,
   CheckoutPrStatusRequestSchema,
   ValidateBranchRequestSchema,
+  BranchSuggestionsRequestSchema,
   PaseoWorktreeListRequestSchema,
   PaseoWorktreeArchiveRequestSchema,
   FileExplorerRequestSchema,
@@ -1565,6 +1574,15 @@ export const ValidateBranchResponseSchema = z.object({
   }),
 });
 
+export const BranchSuggestionsResponseSchema = z.object({
+  type: z.literal("branch_suggestions_response"),
+  payload: z.object({
+    branches: z.array(z.string()),
+    error: z.string().nullable(),
+    requestId: z.string(),
+  }),
+});
+
 const PaseoWorktreeSchema = z.object({
   worktreePath: z.string(),
   branchName: z.string().nullable().optional(),
@@ -1835,6 +1853,7 @@ export const SessionOutboundMessageSchema = z.discriminatedUnion("type", [
   CheckoutPrCreateResponseSchema,
   CheckoutPrStatusResponseSchema,
   ValidateBranchResponseSchema,
+  BranchSuggestionsResponseSchema,
   PaseoWorktreeListResponseSchema,
   PaseoWorktreeArchiveResponseSchema,
   FileExplorerResponseSchema,
@@ -1966,6 +1985,8 @@ export type CheckoutPrStatusRequest = z.infer<typeof CheckoutPrStatusRequestSche
 export type CheckoutPrStatusResponse = z.infer<typeof CheckoutPrStatusResponseSchema>;
 export type ValidateBranchRequest = z.infer<typeof ValidateBranchRequestSchema>;
 export type ValidateBranchResponse = z.infer<typeof ValidateBranchResponseSchema>;
+export type BranchSuggestionsRequest = z.infer<typeof BranchSuggestionsRequestSchema>;
+export type BranchSuggestionsResponse = z.infer<typeof BranchSuggestionsResponseSchema>;
 export type PaseoWorktreeListRequest = z.infer<typeof PaseoWorktreeListRequestSchema>;
 export type PaseoWorktreeListResponse = z.infer<typeof PaseoWorktreeListResponseSchema>;
 export type PaseoWorktreeArchiveRequest = z.infer<typeof PaseoWorktreeArchiveRequestSchema>;

packages/server/src/utils/checkout-git.test.ts

@@ -9,6 +9,7 @@ import {
   getPullRequestStatus,
   getCheckoutStatus,
   getCheckoutStatusLite,
+  listBranchSuggestions,
   mergeToBase,
   mergeFromBase,
   MergeConflictError,
@@ -402,6 +403,51 @@ describe("checkout git utilities", () => {
     execSync(`git --git-dir ${remoteDir} show-ref --verify refs/heads/feature`);
   });
 
+  it("lists merged local and remote branch suggestions without duplicates", async () => {
+    const remoteDir = join(tempDir, "remote.git");
+    execSync(`git init --bare -b main ${remoteDir}`);
+    execSync(`git remote add origin ${remoteDir}`, { cwd: repoDir });
+    execSync("git push -u origin main", { cwd: repoDir });
+
+    execSync("git checkout -b local-only", { cwd: repoDir });
+    execSync("git checkout main", { cwd: repoDir });
+
+    const otherClone = join(tempDir, "other-clone");
+    execSync(`git clone ${remoteDir} ${otherClone}`);
+    execSync("git config user.email 'test@test.com'", { cwd: otherClone });
+    execSync("git config user.name 'Test'", { cwd: otherClone });
+    execSync("git checkout -b remote-only", { cwd: otherClone });
+    writeFileSync(join(otherClone, "remote-only.txt"), "remote-only\n");
+    execSync("git add remote-only.txt", { cwd: otherClone });
+    execSync("git -c commit.gpgsign=false commit -m 'remote only branch'", { cwd: otherClone });
+    execSync("git push -u origin remote-only", { cwd: otherClone });
+    execSync("git fetch origin", { cwd: repoDir });
+
+    const branches = await listBranchSuggestions(repoDir, { limit: 50 });
+    expect(branches).toContain("main");
+    expect(branches).toContain("local-only");
+    expect(branches).toContain("remote-only");
+    expect(branches.filter((name) => name === "main")).toHaveLength(1);
+    expect(branches).not.toContain("HEAD");
+    expect(branches.some((name) => name.startsWith("origin/"))).toBe(false);
+  });
+
+  it("filters branch suggestions by query and enforces result limit", async () => {
+    execSync("git checkout -b feature/alpha", { cwd: repoDir });
+    execSync("git checkout main", { cwd: repoDir });
+    execSync("git checkout -b feature/beta", { cwd: repoDir });
+    execSync("git checkout main", { cwd: repoDir });
+    execSync("git checkout -b chore/docs", { cwd: repoDir });
+    execSync("git checkout main", { cwd: repoDir });
+
+    const branches = await listBranchSuggestions(repoDir, {
+      query: "FEATURE/",
+      limit: 1,
+    });
+    expect(branches).toHaveLength(1);
+    expect(branches[0]?.toLowerCase()).toContain("feature/");
+  });
+
   it("disables GitHub features when gh is unavailable", async () => {
     execSync("git remote add origin https://github.com/getpaseo/paseo.git", { cwd: repoDir });
 

packages/server/src/utils/checkout-git.ts

@@ -104,6 +104,128 @@ type CheckoutFileChange = {
   isUntracked?: boolean;
 };
 
+type BranchSuggestionRefOrigin = "local" | "remote";
+
+function normalizeBranchSuggestionName(raw: string): string | null {
+  const trimmed = raw.trim();
+  if (!trimmed) {
+    return null;
+  }
+
+  let normalized = trimmed;
+  if (normalized.startsWith("refs/heads/")) {
+    normalized = normalized.slice("refs/heads/".length);
+  } else if (normalized.startsWith("refs/remotes/")) {
+    normalized = normalized.slice("refs/remotes/".length);
+  }
+
+  if (normalized.startsWith("origin/")) {
+    normalized = normalized.slice("origin/".length);
+  }
+
+  if (!normalized || normalized === "HEAD") {
+    return null;
+  }
+
+  return normalized;
+}
+
+async function listGitRefs(cwd: string, refPrefix: string): Promise<string[]> {
+  const { stdout } = await execGit(
+    `git for-each-ref --format="%(refname:short)" ${refPrefix}`,
+    {
+      cwd,
+      env: READ_ONLY_GIT_ENV,
+    }
+  );
+  return stdout
+    .split("\n")
+    .map((line) => line.trim())
+    .filter((line) => line.length > 0);
+}
+
+function sortBranchSuggestions(
+  branchNames: string[],
+  localBranchNames: Set<string>,
+  query: string
+): string[] {
+  const normalizedQuery = query.trim().toLowerCase();
+  const hasQuery = normalizedQuery.length > 0;
+  return branchNames.sort((a, b) => {
+    const aLower = a.toLowerCase();
+    const bLower = b.toLowerCase();
+
+    if (hasQuery) {
+      const aPrefix = aLower.startsWith(normalizedQuery);
+      const bPrefix = bLower.startsWith(normalizedQuery);
+      if (aPrefix !== bPrefix) {
+        return aPrefix ? -1 : 1;
+      }
+    }
+
+    const aIsLocal = localBranchNames.has(a);
+    const bIsLocal = localBranchNames.has(b);
+    if (aIsLocal !== bIsLocal) {
+      return aIsLocal ? -1 : 1;
+    }
+
+    return a.localeCompare(b);
+  });
+}
+
+export async function listBranchSuggestions(
+  cwd: string,
+  options?: { query?: string; limit?: number }
+): Promise<string[]> {
+  await requireGitRepo(cwd);
+
+  const requestedLimit = options?.limit ?? 50;
+  const limit = Math.max(1, Math.min(200, requestedLimit));
+  const query = options?.query?.trim().toLowerCase() ?? "";
+
+  const [localRefs, remoteRefs] = await Promise.all([
+    listGitRefs(cwd, "refs/heads"),
+    listGitRefs(cwd, "refs/remotes/origin"),
+  ]);
+
+  const merged = new Map<string, Set<BranchSuggestionRefOrigin>>();
+  for (const localRef of localRefs) {
+    const normalized = normalizeBranchSuggestionName(localRef);
+    if (!normalized) {
+      continue;
+    }
+    const origins = merged.get(normalized) ?? new Set<BranchSuggestionRefOrigin>();
+    origins.add("local");
+    merged.set(normalized, origins);
+  }
+  for (const remoteRef of remoteRefs) {
+    const normalized = normalizeBranchSuggestionName(remoteRef);
+    if (!normalized) {
+      continue;
+    }
+    const origins = merged.get(normalized) ?? new Set<BranchSuggestionRefOrigin>();
+    origins.add("remote");
+    merged.set(normalized, origins);
+  }
+
+  const filteredNames = Array.from(merged.keys()).filter((name) =>
+    query ? name.toLowerCase().includes(query) : true
+  );
+  if (filteredNames.length === 0) {
+    return [];
+  }
+
+  const localBranchNames = new Set<string>();
+  for (const [name, origins] of merged) {
+    if (origins.has("local")) {
+      localBranchNames.add(name);
+    }
+  }
+
+  const ordered = sortBranchSuggestions(filteredNames, localBranchNames, query);
+  return ordered.slice(0, limit);
+}
+
 async function listCheckoutFileChanges(cwd: string, ref: string): Promise<CheckoutFileChange[]> {
   const changes: CheckoutFileChange[] = [];
 

packages/server/src/utils/worktree.test.ts

@@ -171,7 +171,8 @@ describe("createWorktree", () => {
     const paseoConfig = {
       worktree: {
         setup: [
-          'echo "root=$PASEO_ROOT_PATH" > setup.log',
+          'echo "source=$PASEO_SOURCE_CHECKOUT_PATH" > setup.log',
+          'echo "root_alias=$PASEO_ROOT_PATH" >> setup.log',
           'echo "worktree=$PASEO_WORKTREE_PATH" >> setup.log',
           'echo "branch=$PASEO_BRANCH_NAME" >> setup.log',
           'echo "port=$PASEO_WORKTREE_PORT" >> setup.log',
@@ -193,7 +194,8 @@ describe("createWorktree", () => {
 
     // Verify setup ran and env vars were available
     const setupLog = readFileSync(join(result.worktreePath, "setup.log"), "utf8");
-    expect(setupLog).toContain(`root=${repoDir}`);
+    expect(setupLog).toContain(`source=${repoDir}`);
+    expect(setupLog).toContain(`root_alias=${repoDir}`);
     expect(setupLog).toContain(`worktree=${result.worktreePath}`);
     expect(setupLog).toContain("branch=setup-test");
     const portLine = setupLog
@@ -335,9 +337,10 @@ describe("paseo worktree manager", () => {
     const paseoConfig = {
       worktree: {
         destroy: [
-          'echo "root=$PASEO_ROOT_PATH" > "$PASEO_ROOT_PATH/destroy.log"',
-          'echo "worktree=$PASEO_WORKTREE_PATH" >> "$PASEO_ROOT_PATH/destroy.log"',
-          'echo "branch=$PASEO_BRANCH_NAME" >> "$PASEO_ROOT_PATH/destroy.log"',
+          'echo "source=$PASEO_SOURCE_CHECKOUT_PATH" > "$PASEO_SOURCE_CHECKOUT_PATH/destroy.log"',
+          'echo "root_alias=$PASEO_ROOT_PATH" >> "$PASEO_SOURCE_CHECKOUT_PATH/destroy.log"',
+          'echo "worktree=$PASEO_WORKTREE_PATH" >> "$PASEO_SOURCE_CHECKOUT_PATH/destroy.log"',
+          'echo "branch=$PASEO_BRANCH_NAME" >> "$PASEO_SOURCE_CHECKOUT_PATH/destroy.log"',
         ],
       },
     };
@@ -359,7 +362,8 @@ describe("paseo worktree manager", () => {
     expect(existsSync(created.worktreePath)).toBe(false);
 
     const destroyLog = readFileSync(join(repoDir, "destroy.log"), "utf8");
-    expect(destroyLog).toContain(`root=${repoDir}`);
+    expect(destroyLog).toContain(`source=${repoDir}`);
+    expect(destroyLog).toContain(`root_alias=${repoDir}`);
     expect(destroyLog).toContain(`worktree=${created.worktreePath}`);
     expect(destroyLog).toContain("branch=destroy-branch");
   });
@@ -368,7 +372,7 @@ describe("paseo worktree manager", () => {
     const paseoConfig = {
       worktree: {
         destroy: [
-          'echo "started" > "$PASEO_ROOT_PATH/destroy-start.log"',
+          'echo "started" > "$PASEO_SOURCE_CHECKOUT_PATH/destroy-start.log"',
           "echo boom 1>&2; exit 9",
         ],
       },

packages/server/src/utils/worktree.ts

@@ -204,8 +204,11 @@ export async function runWorktreeSetupCommands(options: {
 
   const setupEnv = {
     ...process.env,
-    // Root is the original git repo root (shared across worktrees), not the worktree itself.
-    // This allows setup scripts to copy uncommitted local files (e.g. .env) from the main checkout.
+    // Source checkout path is the original git repo root (shared across worktrees), not the
+    // worktree itself. This allows setup scripts to copy local files (e.g. .env) from the
+    // source checkout.
+    PASEO_SOURCE_CHECKOUT_PATH: repoRootPath,
+    // Backward-compatible alias.
     PASEO_ROOT_PATH: repoRootPath,
     PASEO_WORKTREE_PATH: options.worktreePath,
     PASEO_BRANCH_NAME: options.branchName,
@@ -275,8 +278,11 @@ export async function runWorktreeDestroyCommands(options: {
 
   const destroyEnv = {
     ...process.env,
-    // Root is the original git repo root (shared across worktrees), not the worktree itself.
-    // This allows destroy scripts to clean resources using paths from the main checkout.
+    // Source checkout path is the original git repo root (shared across worktrees), not the
+    // worktree itself. This allows destroy scripts to clean resources using paths from the
+    // source checkout.
+    PASEO_SOURCE_CHECKOUT_PATH: repoRootPath,
+    // Backward-compatible alias.
     PASEO_ROOT_PATH: repoRootPath,
     PASEO_WORKTREE_PATH: options.worktreePath,
     PASEO_BRANCH_NAME: branchName,