feat: add relay reconnect with grace period and branch suggestions

2026-07-29 12:01:31 +00:00 · 2026-02-13 09:46:40 +07:00
parent 837d0fb387
commit bf432b3d3a
23 changed files with 1137 additions and 103 deletions
--- a/packages/app/src/components/agent-input-area.tsx
+++ b/packages/app/src/components/agent-input-area.tsx
@@ -276,13 +276,11 @@ export function AgentInputArea({
    imageAttachments?: ImageAttachment[],
    forceSend?: boolean
  ) {
-    const socketConnected = isConnected;
    const trimmedMessage = message.trim();
    if (!trimmedMessage) return;
    // When the parent controls submission (e.g. draft agent creation), let it
    // decide what to do even if the socket is currently disconnected (so we
    // don't no-op and lose deterministic error handling in the UI/tests).
-    if (!onSubmitMessageRef.current && !socketConnected) return;
    if (!sendAgentMessageRef.current && !onSubmitMessageRef.current) return;

    if (agent?.status === "running" && !forceSend) {
@@ -485,7 +483,7 @@ export function AgentInputArea({

  async function handleSendQueuedNow(id: string) {
    const item = queuedMessages.find((q) => q.id === id);
-    if (!item || !isConnected) return;
+    if (!item) return;
    if (!sendAgentMessageRef.current && !onSubmitMessageRef.current) return;

    updateQueue((current) => current.filter((q) => q.id !== id));
--- a/packages/app/src/components/message-input.tsx
+++ b/packages/app/src/components/message-input.tsx
@@ -731,7 +731,7 @@ export const MessageInput = forwardRef<MessageInputRef, MessageInputProps>(
          scrollEnabled={IS_WEB ? inputHeight >= MAX_INPUT_HEIGHT : true}
          onContentSizeChange={handleContentSizeChange}
          editable={
-            !isDictating && !isRealtimeVoiceForCurrentAgent && isConnected && !disabled
+            !isDictating && !isRealtimeVoiceForCurrentAgent && !disabled
          }
          onKeyPress={
            shouldHandleDesktopSubmit ? handleDesktopKeyPress : undefined
@@ -836,24 +836,22 @@ export const MessageInput = forwardRef<MessageInputRef, MessageInputProps>(
            )}
            {shouldShowSendButton && (
              <Tooltip delayDuration={0} enabledOnDesktop enabledOnMobile={false}>
-                <TooltipTrigger
-                  onPress={handleSendMessage}
-                  disabled={
-                    !isConnected ||
+              <TooltipTrigger
+                onPress={handleSendMessage}
+                disabled={
                    isSubmitDisabled ||
                    isSubmitLoading ||
                    disabled
                  }
-                  accessibilityLabel={isAgentRunning ? "Send and interrupt" : "Send message"}
-                  accessibilityRole="button"
-                  style={[
-                    styles.sendButton,
-                    (!isConnected ||
-                      isSubmitDisabled ||
+                accessibilityLabel={isAgentRunning ? "Send and interrupt" : "Send message"}
+                accessibilityRole="button"
+                style={[
+                  styles.sendButton,
+                    (isSubmitDisabled ||
                      isSubmitLoading ||
                      disabled) &&
                      styles.buttonDisabled,
-                  ]}
+                ]}
                >
                  {isSubmitLoading ? (
                    <ActivityIndicator size="small" color="white" />
--- a/packages/app/src/components/multi-daemon-session-host.tsx
+++ b/packages/app/src/components/multi-daemon-session-host.tsx
@@ -183,7 +183,7 @@ function ManagedDaemonSession({ daemon }: { daemon: HostProfile }) {

  return (
    <SessionProvider
-      key={`${daemon.serverId}:${activeUrl}`}
+      key={daemon.serverId}
      serverUrl={activeUrl}
      serverId={daemon.serverId}
      activeConnection={active?.activeConnection ?? null}
--- a/packages/app/src/components/ui/combobox.tsx
+++ b/packages/app/src/components/ui/combobox.tsx
@@ -35,6 +35,7 @@ export interface ComboboxProps {
  options: ComboboxOption[];
  value: string;
  onSelect: (id: string) => void;
+  onSearchQueryChange?: (query: string) => void;
  placeholder?: string;
  searchPlaceholder?: string;
  emptyText?: string;
@@ -156,6 +157,7 @@ export function Combobox({
  options,
  value,
  onSelect,
+  onSearchQueryChange,
  placeholder = "Search...",
  searchPlaceholder,
  emptyText = "No options match your search.",
@@ -191,16 +193,24 @@ export function Combobox({
    [isControlled, onOpenChange]
  );

+  const setSearchQueryWithCallback = useCallback(
+    (nextQuery: string) => {
+      setSearchQuery(nextQuery);
+      onSearchQueryChange?.(nextQuery);
+    },
+    [onSearchQueryChange]
+  );
+
  const handleClose = useCallback(() => {
    setOpen(false);
-    setSearchQuery("");
-  }, [setOpen]);
+    setSearchQueryWithCallback("");
+  }, [setOpen, setSearchQueryWithCallback]);

  useEffect(() => {
    if (isOpen) {
-      setSearchQuery("");
+      setSearchQueryWithCallback("");
    }
-  }, [isOpen]);
+  }, [isOpen, setSearchQueryWithCallback]);

  const collisionPadding = useMemo(() => {
    const basePadding = 16;
@@ -426,7 +436,7 @@ export function Combobox({
    <SearchInput
      placeholder={searchPlaceholder ?? placeholder}
      value={searchQuery}
-      onChangeText={setSearchQuery}
+      onChangeText={setSearchQueryWithCallback}
      onSubmitEditing={handleSubmitSearch}
      autoFocus={!isMobile}
    />
--- a/packages/app/src/screens/agent/agent-ready-screen.tsx
+++ b/packages/app/src/screens/agent/agent-ready-screen.tsx
@@ -45,7 +45,6 @@ import {
 import { usePanelStore } from "@/stores/panel-store";
 import { useDaemonConnections } from "@/contexts/daemon-connections-context";
 import type { ConnectionStatus } from "@/contexts/daemon-connections-context";
-import { formatConnectionStatus } from "@/utils/daemons";
 import { useSessionStore } from "@/stores/session-store";
 import { useCreateFlowStore } from "@/stores/create-flow-store";
 import type { Agent } from "@/contexts/session-context";
@@ -74,6 +73,8 @@ import {

 const DROPDOWN_WIDTH = 220;
 const EMPTY_STREAM_ITEMS: StreamItem[] = [];
+const RECONNECT_NOTICE_DELAY_MS = 10_000;
+const CONNECTED_NOTICE_DURATION_MS = 2_500;

 export function AgentReadyScreen({
  serverId,
@@ -101,7 +102,6 @@ export function AgentReadyScreen({
  const isUnknownDaemon = Boolean(connectionServerId && !connection);
  const connectionStatus =
    connection?.status ?? (isUnknownDaemon ? "offline" : "idle");
-  const connectionStatusLabel = formatConnectionStatus(connectionStatus);
  const lastConnectionError = connection?.lastError ?? null;

  const handleBackToHome = useCallback(() => {
@@ -152,7 +152,6 @@ export function AgentReadyScreen({
        onBack={handleBackToHome}
        serverLabel={serverLabel}
        connectionStatus={connectionStatus}
-        connectionStatusLabel={connectionStatusLabel}
        lastError={lastConnectionError}
        isUnknownDaemon={isUnknownDaemon}
      />
@@ -164,6 +163,7 @@ export function AgentReadyScreen({
      <AgentScreenContent
        serverId={resolvedServerId}
        agentId={resolvedAgentId}
+        connectionStatus={connectionStatus}
      />
    </ExplorerSidebarAnimationProvider>
  );
@@ -172,6 +172,7 @@ export function AgentReadyScreen({
 type AgentScreenContentProps = {
  serverId: string;
  agentId?: string;
+  connectionStatus: ConnectionStatus;
 };

 type MissingAgentState =
@@ -194,6 +195,7 @@ function isNotFoundErrorMessage(message: string): boolean {
 function AgentScreenContent({
  serverId,
  agentId,
+  connectionStatus,
 }: AgentScreenContentProps) {
  const { theme } = useUnistyles();
  const toast = useToast();
@@ -383,6 +385,11 @@ function AgentScreenContent({
  const [missingAgentState, setMissingAgentState] = useState<MissingAgentState>({
    kind: "idle",
  });
+  const [showReconnectNotice, setShowReconnectNotice] = useState(false);
+  const [dismissedReconnectNotice, setDismissedReconnectNotice] = useState(false);
+  const [showConnectedNotice, setShowConnectedNotice] = useState(false);
+  const reconnectNoticeTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+  const connectedNoticeTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);
  const initAttemptTokenRef = useRef(0);
  const setFocusedAgentId = useCallback(
    (agentId: string | null) => {
@@ -410,6 +417,56 @@ function AgentScreenContent({
  const agentModel = extractAgentModel(agent);
  const modelDisplayValue = agentModel ?? "Unknown";

+  useEffect(() => {
+    if (reconnectNoticeTimeoutRef.current) {
+      clearTimeout(reconnectNoticeTimeoutRef.current);
+      reconnectNoticeTimeoutRef.current = null;
+    }
+
+    if (connectionStatus === "online") {
+      if (showReconnectNotice || dismissedReconnectNotice) {
+        setShowConnectedNotice(true);
+      }
+      setShowReconnectNotice(false);
+      setDismissedReconnectNotice(false);
+      return;
+    }
+
+    setShowConnectedNotice(false);
+    if (!showReconnectNotice && !dismissedReconnectNotice) {
+      reconnectNoticeTimeoutRef.current = setTimeout(() => {
+        setShowReconnectNotice(true);
+      }, RECONNECT_NOTICE_DELAY_MS);
+    }
+
+    return () => {
+      if (reconnectNoticeTimeoutRef.current) {
+        clearTimeout(reconnectNoticeTimeoutRef.current);
+        reconnectNoticeTimeoutRef.current = null;
+      }
+    };
+  }, [connectionStatus, dismissedReconnectNotice, showReconnectNotice]);
+
+  useEffect(() => {
+    if (!showConnectedNotice) {
+      if (connectedNoticeTimeoutRef.current) {
+        clearTimeout(connectedNoticeTimeoutRef.current);
+        connectedNoticeTimeoutRef.current = null;
+      }
+      return;
+    }
+    connectedNoticeTimeoutRef.current = setTimeout(() => {
+      setShowConnectedNotice(false);
+      connectedNoticeTimeoutRef.current = null;
+    }, CONNECTED_NOTICE_DURATION_MS);
+    return () => {
+      if (connectedNoticeTimeoutRef.current) {
+        clearTimeout(connectedNoticeTimeoutRef.current);
+        connectedNoticeTimeoutRef.current = null;
+      }
+    };
+  }, [showConnectedNotice]);
+
  // Checkout status for header subtitle
  const checkoutStatusQuery = useCheckoutStatusQuery({
    serverId,
@@ -902,6 +959,47 @@ function AgentScreenContent({
          }
        />

+          {(showReconnectNotice || showConnectedNotice) && (
+            <View
+              style={[
+                styles.connectionNotice,
+                showReconnectNotice
+                  ? styles.connectionNoticeReconnecting
+                  : styles.connectionNoticeConnected,
+              ]}
+            >
+              {showConnectedNotice ? (
+                <CheckCircle2
+                  size={14}
+                  color={theme.colors.palette.green[600]}
+                />
+              ) : null}
+              <Text
+                style={[
+                  styles.connectionNoticeText,
+                  showReconnectNotice
+                    ? styles.connectionNoticeTextReconnecting
+                    : styles.connectionNoticeTextConnected,
+                ]}
+              >
+                {showReconnectNotice ? "Reconnecting..." : "Connected"}
+              </Text>
+              {showReconnectNotice ? (
+                <Pressable
+                  onPress={() => {
+                    setShowReconnectNotice(false);
+                    setDismissedReconnectNotice(true);
+                  }}
+                  style={styles.connectionNoticeDismiss}
+                  accessibilityRole="button"
+                  accessibilityLabel="Dismiss reconnecting notice"
+                >
+                  <Text style={styles.connectionNoticeDismissText}>Dismiss</Text>
+                </Pressable>
+              ) : null}
+            </View>
+          )}
+
          {/* Content Area with Keyboard Animation */}
          <View style={styles.contentContainer}>
            {shouldBlockForHistorySync ? (
@@ -967,14 +1065,12 @@ function AgentSessionUnavailableState({
  onBack,
  serverLabel,
  connectionStatus,
-  connectionStatusLabel,
  lastError,
  isUnknownDaemon = false,
 }: {
  onBack: () => void;
  serverLabel: string;
  connectionStatus: ConnectionStatus;
-  connectionStatusLabel: string;
  lastError: string | null;
  isUnknownDaemon?: boolean;
 }) {
@@ -1015,11 +1111,10 @@ function AgentSessionUnavailableState({
        ) : (
          <>
            <Text style={styles.offlineTitle}>
-              {serverLabel} is currently {connectionStatusLabel.toLowerCase()}.
+              Reconnecting to {serverLabel}...
            </Text>
            <Text style={styles.offlineDescription}>
-              We'll reconnect automatically and show this agent as soon as the
-              host comes back online.
+              We will show this agent again as soon as the host is reachable.
            </Text>
            {lastError ? (
              <Text style={styles.offlineDetails}>{lastError}</Text>
@@ -1049,6 +1144,49 @@ const styles = StyleSheet.create((theme) => ({
    flex: 1,
    overflow: "hidden",
  },
+  connectionNotice: {
+    marginHorizontal: theme.spacing[4],
+    marginTop: theme.spacing[1],
+    marginBottom: theme.spacing[2],
+    minHeight: 32,
+    borderRadius: theme.borderRadius.full,
+    paddingHorizontal: theme.spacing[3],
+    flexDirection: "row",
+    alignItems: "center",
+    gap: theme.spacing[2],
+  },
+  connectionNoticeReconnecting: {
+    backgroundColor: `${theme.colors.palette.yellow[400]}22`,
+    borderWidth: theme.borderWidth[1],
+    borderColor: theme.colors.palette.yellow[400],
+  },
+  connectionNoticeConnected: {
+    backgroundColor: theme.colors.palette.green[100],
+    borderWidth: theme.borderWidth[1],
+    borderColor: theme.colors.palette.green[400],
+  },
+  connectionNoticeText: {
+    fontSize: theme.fontSize.sm,
+    fontWeight: theme.fontWeight.medium,
+  },
+  connectionNoticeTextReconnecting: {
+    color: theme.colors.foreground,
+  },
+  connectionNoticeTextConnected: {
+    color: theme.colors.palette.green[800],
+  },
+  connectionNoticeDismiss: {
+    marginLeft: "auto",
+    paddingVertical: theme.spacing[1],
+    paddingHorizontal: theme.spacing[2],
+    borderRadius: theme.borderRadius.full,
+    backgroundColor: theme.colors.surface0,
+  },
+  connectionNoticeDismissText: {
+    fontSize: theme.fontSize.xs,
+    fontWeight: theme.fontWeight.medium,
+    color: theme.colors.foregroundMuted,
+  },
  content: {
    flex: 1,
  },
--- a/packages/app/src/screens/agent/draft-agent-screen.tsx
+++ b/packages/app/src/screens/agent/draft-agent-screen.tsx
@@ -30,6 +30,7 @@ import {
 import { useDaemonConnections } from "@/contexts/daemon-connections-context";
 import { useDaemonRegistry } from "@/contexts/daemon-registry-context";
 import { formatConnectionStatus } from "@/utils/daemons";
+import { buildBranchComboOptions, normalizeBranchOptionName } from "@/utils/branch-suggestions";
 import { shortenPath } from "@/utils/shorten-path";
 import { usePanelStore } from "@/stores/panel-store";
 import { useSessionStore } from "@/stores/session-store";
@@ -249,6 +250,8 @@ export function DraftAgentScreen({
  const [isWorkingDirOpen, setIsWorkingDirOpen] = useState(false);
  const [isWorktreePickerOpen, setIsWorktreePickerOpen] = useState(false);
  const [isBranchOpen, setIsBranchOpen] = useState(false);
+  const [branchSearchQuery, setBranchSearchQuery] = useState("");
+  const [debouncedBranchSearchQuery, setDebouncedBranchSearchQuery] = useState("");
  const hostAnchorRef = useRef<View>(null);
  const workingDirAnchorRef = useRef<View>(null);
  const worktreeAnchorRef = useRef<View>(null);
@@ -258,6 +261,12 @@ export function DraftAgentScreen({
  const updatePendingAgentId = useCreateFlowStore((state) => state.updateAgentId);
  const clearPendingCreateAttempt = useCreateFlowStore((state) => state.clear);

+  useEffect(() => {
+    const trimmed = branchSearchQuery.trim();
+    const timer = setTimeout(() => setDebouncedBranchSearchQuery(trimmed), 180);
+    return () => clearTimeout(timer);
+  }, [branchSearchQuery]);
+
  type CreateAttempt = {
    messageId: string;
    text: string;
@@ -489,6 +498,40 @@ export function DraftAgentScreen({
      ? "Select a worktree to attach"
      : null;

+  const branchSuggestionsQuery = useQuery({
+    queryKey: [
+      "branchSuggestions",
+      selectedServerId,
+      trimmedWorkingDir,
+      debouncedBranchSearchQuery,
+    ],
+    queryFn: async () => {
+      const client = sessionClient;
+      if (!client) {
+        throw new Error("Daemon client unavailable");
+      }
+      const payload = await client.getBranchSuggestions({
+        cwd: trimmedWorkingDir || ".",
+        query: debouncedBranchSearchQuery || undefined,
+        limit: 50,
+      });
+      if (payload.error) {
+        throw new Error(payload.error);
+      }
+      return payload.branches ?? [];
+    },
+    enabled:
+      isCreateWorktree &&
+      isGitDirectory &&
+      !isNonGitDirectory &&
+      Boolean(trimmedWorkingDir) &&
+      !repoAvailabilityError &&
+      Boolean(sessionClient) &&
+      isConnected,
+    retry: false,
+    staleTime: 15_000,
+  });
+
  const validateWorktreeName = useCallback(
    (name: string): { valid: boolean; error?: string } => {
      if (!name) {
@@ -658,21 +701,36 @@ export function DraftAgentScreen({
  );

  const branchComboOptions = useMemo(() => {
-    const branchSet = new Set<string>();
-    const currentBranch = checkout?.isGit ? checkout.currentBranch?.trim() : null;
-    if (currentBranch && currentBranch !== "HEAD") {
-      branchSet.add(currentBranch);
+    const options = buildBranchComboOptions({
+      suggestedBranches: branchSuggestionsQuery.data ?? [],
+      currentBranch: checkout?.isGit ? checkout.currentBranch : null,
+      baseRef: checkout?.isGit ? checkout.baseRef : null,
+      typedBaseBranch: baseBranch,
+      worktreeBranchLabels: worktreeOptions.map((option) => option.label),
+    });
+
+    const normalizedQuery = normalizeBranchOptionName(branchSearchQuery)?.toLowerCase() ?? "";
+    if (!normalizedQuery) {
+      return options;
    }
-    if (baseBranch.trim()) {
-      branchSet.add(baseBranch.trim());
-    }
-    for (const option of worktreeOptions) {
-      if (option.label) {
-        branchSet.add(option.label);
+
+    return options.sort((a, b) => {
+      const aLower = a.label.toLowerCase();
+      const bLower = b.label.toLowerCase();
+      const aPrefix = aLower.startsWith(normalizedQuery);
+      const bPrefix = bLower.startsWith(normalizedQuery);
+      if (aPrefix !== bPrefix) {
+        return aPrefix ? -1 : 1;
      }
-    }
-    return Array.from(branchSet).map((name) => ({ id: name, label: name }));
-  }, [baseBranch, checkout, worktreeOptions]);
+      return aLower.localeCompare(bLower);
+    });
+  }, [
+    baseBranch,
+    branchSearchQuery,
+    branchSuggestionsQuery.data,
+    checkout,
+    worktreeOptions,
+  ]);

  const createAgentClient = useSessionStore((state) =>
    selectedServerId ? state.sessions[selectedServerId]?.client ?? null : null
@@ -1130,13 +1188,19 @@ export function DraftAgentScreen({
            options={branchComboOptions}
            value={baseBranch}
            onSelect={handleBaseBranchChange}
+            onSearchQueryChange={setBranchSearchQuery}
            searchPlaceholder="Choose a base branch..."
            allowCustomValue
            customValuePrefix="Use"
            customValueDescription="Use this branch name"
            title="Select base branch"
            open={isBranchOpen}
-            onOpenChange={setIsBranchOpen}
+            onOpenChange={(nextOpen) => {
+              setIsBranchOpen(nextOpen);
+              if (!nextOpen) {
+                setBranchSearchQuery("");
+              }
+            }}
            anchorRef={branchAnchorRef}
          />

--- a/packages/app/src/utils/branch-suggestions.test.ts
+++ b/packages/app/src/utils/branch-suggestions.test.ts
@@ -0,0 +1,36 @@
+import { describe, expect, it } from "vitest";
+import { buildBranchComboOptions, normalizeBranchOptionName } from "./branch-suggestions";
+
+describe("normalizeBranchOptionName", () => {
+  it("normalizes local and origin-prefixed refs", () => {
+    expect(normalizeBranchOptionName("refs/heads/main")).toBe("main");
+    expect(normalizeBranchOptionName("refs/remotes/origin/main")).toBe("main");
+    expect(normalizeBranchOptionName("origin/feature/test")).toBe("feature/test");
+    expect(normalizeBranchOptionName("feature/test")).toBe("feature/test");
+  });
+
+  it("filters out empty values and HEAD", () => {
+    expect(normalizeBranchOptionName("")).toBeNull();
+    expect(normalizeBranchOptionName("   ")).toBeNull();
+    expect(normalizeBranchOptionName("HEAD")).toBeNull();
+    expect(normalizeBranchOptionName("origin/HEAD")).toBeNull();
+  });
+});
+
+describe("buildBranchComboOptions", () => {
+  it("merges branch sources and de-duplicates normalized names", () => {
+    const options = buildBranchComboOptions({
+      suggestedBranches: ["origin/main", "refs/remotes/origin/main", "feature/a"],
+      currentBranch: "refs/heads/feature/a",
+      baseRef: "origin/main",
+      typedBaseBranch: "main",
+      worktreeBranchLabels: ["refs/heads/release/next"],
+    });
+
+    expect(options).toEqual([
+      { id: "main", label: "main" },
+      { id: "feature/a", label: "feature/a" },
+      { id: "release/next", label: "release/next" },
+    ]);
+  });
+});
--- a/packages/app/src/utils/branch-suggestions.ts
+++ b/packages/app/src/utils/branch-suggestions.ts
@@ -0,0 +1,51 @@
+export type BranchComboOption = {
+  id: string;
+  label: string;
+};
+
+export function normalizeBranchOptionName(input: string | null | undefined): string | null {
+  const trimmed = input?.trim();
+  if (!trimmed || trimmed === "HEAD") {
+    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);
+  }
+
+  return normalized.length > 0 && normalized !== "HEAD" ? normalized : null;
+}
+
+export function buildBranchComboOptions(input: {
+  suggestedBranches?: string[];
+  currentBranch?: string | null;
+  baseRef?: string | null;
+  typedBaseBranch?: string | null;
+  worktreeBranchLabels?: string[];
+}): BranchComboOption[] {
+  const branchSet = new Set<string>();
+  const addBranch = (name: string | null | undefined) => {
+    const normalized = normalizeBranchOptionName(name);
+    if (normalized) {
+      branchSet.add(normalized);
+    }
+  };
+
+  for (const branch of input.suggestedBranches ?? []) {
+    addBranch(branch);
+  }
+  addBranch(input.currentBranch ?? null);
+  addBranch(input.baseRef ?? null);
+  addBranch(input.typedBaseBranch ?? null);
+  for (const label of input.worktreeBranchLabels ?? []) {
+    addBranch(label);
+  }
+
+  return Array.from(branchSet).map((name) => ({ id: name, label: name }));
+}
--- a/packages/server/src/client/daemon-client.test.ts
+++ b/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();
--- a/packages/server/src/client/daemon-client.ts
+++ b/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
  // ============================================================================
--- a/packages/server/src/server/bootstrap.ts
+++ b/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,
--- a/packages/server/src/server/relay-transport.test.ts
+++ b/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",
+      }
+    );
+  });
 });
--- a/packages/server/src/server/relay-transport.ts
+++ b/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 {
--- a/packages/server/src/server/session.ts
+++ b/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" };
--- a/packages/server/src/server/websocket-server.relay-reconnect.test.ts
+++ b/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();
+  });
+});
--- a/packages/server/src/server/websocket-server.ts
+++ b/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 {
--- a/packages/server/src/shared/messages.ts
+++ b/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>;
--- a/packages/server/src/utils/checkout-git.test.ts
+++ b/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 });

--- a/packages/server/src/utils/checkout-git.ts
+++ b/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[] = [];

--- a/packages/server/src/utils/worktree.test.ts
+++ b/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",
        ],
      },
--- a/packages/server/src/utils/worktree.ts
+++ b/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,
--- a/packages/website/src/routes/docs/worktrees.tsx
+++ b/packages/website/src/routes/docs/worktrees.tsx
@@ -110,7 +110,7 @@ function Worktrees() {
  "worktree": {
    "setup": [
      "npm ci",
-      "cp \\"$PASEO_ROOT_PATH/.env\\" \\"$PASEO_WORKTREE_PATH/.env\\""
+      "cp \\"$PASEO_SOURCE_CHECKOUT_PATH/.env\\" \\"$PASEO_WORKTREE_PATH/.env\\""
    ]
  }
 }`}</pre>
@@ -120,6 +120,13 @@ function Worktrees() {
          after the worktree is created. Use it to install dependencies, copy local config
          files, or run any other initialization.
        </p>
+        <div className="bg-yellow-500/10 border border-yellow-500/30 rounded-lg p-4 text-white/80">
+          <strong>Important:</strong> Setup commands come from{' '}
+          <code className="font-mono">paseo.json</code> in the selected base branch. If you pick{' '}
+          <code className="font-mono">main</code>, Paseo reads the committed file on{' '}
+          <code className="font-mono">main</code>. Local or uncommitted changes in another branch
+          are not used for that worktree.
+        </div>
      </section>

      {/* Environment variables */}
@@ -130,7 +137,12 @@ function Worktrees() {
        </p>
        <ul className="text-white/60 space-y-2 list-disc list-inside">
          <li>
-            <code className="font-mono">$PASEO_ROOT_PATH</code> — your original repository root
+            <code className="font-mono">$PASEO_SOURCE_CHECKOUT_PATH</code> — your source checkout
+            path (original repository root)
+          </li>
+          <li>
+            <code className="font-mono">$PASEO_ROOT_PATH</code> — legacy alias of{' '}
+            <code className="font-mono">$PASEO_SOURCE_CHECKOUT_PATH</code>
          </li>
          <li>
            <code className="font-mono">$PASEO_WORKTREE_PATH</code> — the new worktree directory
@@ -143,9 +155,9 @@ function Worktrees() {
          </li>
        </ul>
        <p className="text-white/60 leading-relaxed">
-          Use <code className="font-mono">$PASEO_ROOT_PATH</code> to copy files that shouldn't
-          be in git (like <code className="font-mono">.env</code>) from your main checkout to
-          the worktree.
+          Use <code className="font-mono">$PASEO_SOURCE_CHECKOUT_PATH</code> to copy files that
+          shouldn't be in git (like <code className="font-mono">.env</code>) from your source
+          checkout to the worktree.
        </p>
      </section>

@@ -177,8 +189,8 @@ function Worktrees() {
  "worktree": {
    "setup": [
      "npm ci",
-      "cp \\"$PASEO_ROOT_PATH/.env\\" \\"$PASEO_WORKTREE_PATH/.env\\"",
-      "cp \\"$PASEO_ROOT_PATH/.env.local\\" \\"$PASEO_WORKTREE_PATH/.env.local\\""
+      "cp \\"$PASEO_SOURCE_CHECKOUT_PATH/.env\\" \\"$PASEO_WORKTREE_PATH/.env\\"",
+      "cp \\"$PASEO_SOURCE_CHECKOUT_PATH/.env.local\\" \\"$PASEO_WORKTREE_PATH/.env.local\\""
    ]
  }
 }`}</pre>
@@ -190,7 +202,7 @@ function Worktrees() {
  "worktree": {
    "setup": [
      "npm ci",
-      "cp \\"$PASEO_ROOT_PATH/.env\\" \\"$PASEO_WORKTREE_PATH/.env\\"",
+      "cp \\"$PASEO_SOURCE_CHECKOUT_PATH/.env\\" \\"$PASEO_WORKTREE_PATH/.env\\"",
      "npm run db:migrate"
    ]
  }
--- a/paseo.json
+++ b/paseo.json
@@ -2,7 +2,7 @@
  "worktree": {
    "setup": [
      "npm ci",
-      "cp \"$PASEO_ROOT_PATH/packages/server/.env\" \"$PASEO_WORKTREE_PATH/packages/server/.env\""
+      "cp \"$PASEO_SOURCE_CHECKOUT_PATH/packages/server/.env\" \"$PASEO_WORKTREE_PATH/packages/server/.env\""
    ]
  }
 }