Restore input focus after Cmd+K and focus input on agent switch

packages/app/src/components/agent-input-area.tsx

@@ -3,6 +3,7 @@ import {
   Pressable,
   Text,
   ActivityIndicator,
+  Platform,
 } from "react-native";
 import { useState, useEffect, useRef, useCallback, useMemo } from "react";
 import { StyleSheet, useUnistyles } from "react-native-unistyles";
@@ -28,6 +29,7 @@ import { Theme } from "@/styles/theme";
 import { CommandAutocomplete } from "./command-autocomplete";
 import { useAgentCommandsQuery } from "@/hooks/use-agent-commands-query";
 import { encodeImages } from "@/utils/encode-images";
+import { useKeyboardNavStore } from "@/stores/keyboard-nav-store";
 
 type QueuedMessage = {
   id: string;
@@ -68,6 +70,8 @@ export function AgentInputArea({
   const insets = useSafeAreaInsets();
   const { height: keyboardHeight } = useReanimatedKeyboardAnimation();
   const isScreenFocused = useIsFocused();
+  const focusChatInputRequest = useKeyboardNavStore((s) => s.focusChatInputRequest);
+  const clearFocusChatInputRequest = useKeyboardNavStore((s) => s.clearFocusChatInputRequest);
 
   const client = useSessionStore(
     (state) => state.sessions[serverId]?.client ?? null
@@ -98,6 +102,7 @@ export function AgentInputArea({
   const [selectedImages, setSelectedImages] = useState<ImageAttachment[]>([]);
   const [isCancellingAgent, setIsCancellingAgent] = useState(false);
   const [commandSelectedIndex, setCommandSelectedIndex] = useState(0);
+  const lastHandledFocusRequestIdRef = useRef<number | null>(null);
 
   // Command autocomplete logic
   const showCommandAutocomplete = userInput.startsWith("/") && !userInput.includes(" ");
@@ -352,6 +357,35 @@ export function AgentInputArea({
     saveDraftInput(agentId, { text: userInput, images: selectedImages });
   }, [agentId, userInput, selectedImages, getDraftInput, saveDraftInput]);
 
+  // When switching agents from the command center, auto-focus the input on web.
+  useEffect(() => {
+    if (Platform.OS !== "web") return;
+    if (!isScreenFocused) return;
+    if (!focusChatInputRequest) return;
+
+    const target = focusChatInputRequest.agentKey;
+    const currentKey = `${serverId}:${agentId}`;
+    if (target !== null && target !== currentKey) {
+      return;
+    }
+
+    if (lastHandledFocusRequestIdRef.current === focusChatInputRequest.id) {
+      return;
+    }
+    lastHandledFocusRequestIdRef.current = focusChatInputRequest.id;
+
+    setTimeout(() => {
+      messageInputRef.current?.focus();
+    }, 0);
+    clearFocusChatInputRequest(focusChatInputRequest.id);
+  }, [
+    agentId,
+    clearFocusChatInputRequest,
+    focusChatInputRequest,
+    isScreenFocused,
+    serverId,
+  ]);
+
   const keyboardAnimatedStyle = useAnimatedStyle(() => {
     "worklet";
     const absoluteHeight = Math.abs(keyboardHeight.value);

packages/app/src/components/command-center.tsx

@@ -48,7 +48,11 @@ export function CommandCenter() {
   const { agents } = useAggregatedAgents();
   const open = useKeyboardNavStore((s) => s.commandCenterOpen);
   const setOpen = useKeyboardNavStore((s) => s.setCommandCenterOpen);
+  const requestFocusChatInput = useKeyboardNavStore((s) => s.requestFocusChatInput);
   const inputRef = useRef<TextInput>(null);
+  const previouslyFocusedRef = useRef<HTMLElement | null>(null);
+  const didNavigateRef = useRef(false);
+  const prevOpenRef = useRef(open);
   const [query, setQuery] = useState("");
   const [activeIndex, setActiveIndex] = useState(0);
 
@@ -59,12 +63,34 @@ export function CommandCenter() {
   }, [agents, query]);
 
   useEffect(() => {
+    const prevOpen = prevOpenRef.current;
+    prevOpenRef.current = open;
+
     if (!open) {
       setQuery("");
       setActiveIndex(0);
+
+      if (prevOpen && !didNavigateRef.current) {
+        const el = previouslyFocusedRef.current;
+        if (el && el.isConnected) {
+          // Modal unmount can steal focus; restore on next tick.
+          setTimeout(() => {
+            try {
+              el.focus();
+            } catch {
+              // ignore
+            }
+          }, 0);
+        }
+      }
+
       return;
     }
 
+    didNavigateRef.current = false;
+    previouslyFocusedRef.current =
+      typeof document !== "undefined" ? (document.activeElement as HTMLElement | null) : null;
+
     const id = setTimeout(() => {
       inputRef.current?.focus();
     }, 0);
@@ -84,16 +110,18 @@ export function CommandCenter() {
 
   const handleSelect = useCallback(
     (agent: AggregatedAgent) => {
+      didNavigateRef.current = true;
       const session = useSessionStore.getState().sessions[agent.serverId];
       session?.client?.clearAgentAttention(agent.id);
 
       const shouldReplace = pathname.startsWith("/agent/");
       const navigate = shouldReplace ? router.replace : router.push;
 
+      requestFocusChatInput(`${agent.serverId}:${agent.id}`);
       setOpen(false);
       navigate(`/agent/${agent.serverId}/${agent.id}` as any);
     },
-    [pathname, setOpen]
+    [pathname, requestFocusChatInput, setOpen]
   );
 
   useEffect(() => {

packages/app/src/stores/keyboard-nav-store.ts

@@ -6,11 +6,14 @@ interface KeyboardNavState {
   cmdOrCtrlDown: boolean;
   /** Sidebar-visible agent keys (up to 9), in top-to-bottom visual order. */
   sidebarShortcutAgentKeys: string[];
+  focusChatInputRequest: { id: number; agentKey: string | null } | null;
 
   setCommandCenterOpen: (open: boolean) => void;
   setAltDown: (down: boolean) => void;
   setCmdOrCtrlDown: (down: boolean) => void;
   setSidebarShortcutAgentKeys: (keys: string[]) => void;
+  requestFocusChatInput: (agentKey: string | null) => void;
+  clearFocusChatInputRequest: (id: number) => void;
   resetModifiers: () => void;
 }
 
@@ -19,11 +22,25 @@ export const useKeyboardNavStore = create<KeyboardNavState>((set) => ({
   altDown: false,
   cmdOrCtrlDown: false,
   sidebarShortcutAgentKeys: [],
+  focusChatInputRequest: null,
 
   setCommandCenterOpen: (open) => set({ commandCenterOpen: open }),
   setAltDown: (down) => set({ altDown: down }),
   setCmdOrCtrlDown: (down) => set({ cmdOrCtrlDown: down }),
   setSidebarShortcutAgentKeys: (keys) => set({ sidebarShortcutAgentKeys: keys }),
+  requestFocusChatInput: (agentKey) =>
+    set((state) => ({
+      focusChatInputRequest: {
+        id: (state.focusChatInputRequest?.id ?? 0) + 1,
+        agentKey,
+      },
+    })),
+  clearFocusChatInputRequest: (id) =>
+    set((state) => {
+      if (state.focusChatInputRequest?.id !== id) {
+        return state;
+      }
+      return { focusChatInputRequest: null };
+    }),
   resetModifiers: () => set({ altDown: false, cmdOrCtrlDown: false }),
 }));
-