Fix Cmd+K focus restore and focus input after agent switch

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

@@ -374,9 +374,11 @@ export function AgentInputArea({
     }
     lastHandledFocusRequestIdRef.current = focusChatInputRequest.id;
 
-    setTimeout(() => {
-      messageInputRef.current?.focus();
-    }, 0);
+    requestAnimationFrame(() => {
+      requestAnimationFrame(() => {
+        messageInputRef.current?.focus();
+      });
+    });
     clearFocusChatInputRequest(focusChatInputRequest.id);
   }, [
     agentId,

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

@@ -49,8 +49,8 @@ export function CommandCenter() {
   const open = useKeyboardNavStore((s) => s.commandCenterOpen);
   const setOpen = useKeyboardNavStore((s) => s.setCommandCenterOpen);
   const requestFocusChatInput = useKeyboardNavStore((s) => s.requestFocusChatInput);
+  const takeFocusRestoreElement = useKeyboardNavStore((s) => s.takeFocusRestoreElement);
   const inputRef = useRef<TextInput>(null);
-  const previouslyFocusedRef = useRef<HTMLElement | null>(null);
   const didNavigateRef = useRef(false);
   const prevOpenRef = useRef(open);
   const [query, setQuery] = useState("");
@@ -71,16 +71,18 @@ export function CommandCenter() {
       setActiveIndex(0);
 
       if (prevOpen && !didNavigateRef.current) {
-        const el = previouslyFocusedRef.current;
+        const el = takeFocusRestoreElement();
         if (el && el.isConnected) {
           // Modal unmount can steal focus; restore on next tick.
-          setTimeout(() => {
-            try {
-              el.focus();
-            } catch {
-              // ignore
-            }
-          }, 0);
+          requestAnimationFrame(() => {
+            requestAnimationFrame(() => {
+              try {
+                el.focus();
+              } catch {
+                // ignore
+              }
+            });
+          });
         }
       }
 
@@ -88,8 +90,6 @@ export function CommandCenter() {
     }
 
     didNavigateRef.current = false;
-    previouslyFocusedRef.current =
-      typeof document !== "undefined" ? (document.activeElement as HTMLElement | null) : null;
 
     const id = setTimeout(() => {
       inputRef.current?.focus();
@@ -118,10 +118,12 @@ export function CommandCenter() {
       const navigate = shouldReplace ? router.replace : router.push;
 
       requestFocusChatInput(`${agent.serverId}:${agent.id}`);
+      // Don't restore focus back to the prior element after we navigate.
+      takeFocusRestoreElement();
       setOpen(false);
       navigate(`/agent/${agent.serverId}/${agent.id}` as any);
     },
-    [pathname, requestFocusChatInput, setOpen]
+    [pathname, requestFocusChatInput, setOpen, takeFocusRestoreElement]
   );
 
   useEffect(() => {

packages/app/src/hooks/use-global-keyboard-nav.ts

@@ -127,7 +127,19 @@ export function useGlobalKeyboardNav({
       if ((event.metaKey || event.ctrlKey) && lowerKey === "k") {
         event.preventDefault();
         const s = useKeyboardNavStore.getState();
-        s.setCommandCenterOpen(!s.commandCenterOpen);
+        if (!s.commandCenterOpen) {
+          const target = event.target as unknown;
+          const el =
+            typeof HTMLElement !== "undefined" && target instanceof HTMLElement
+              ? target
+              : (typeof document !== "undefined"
+                  ? (document.activeElement as HTMLElement | null)
+                  : null);
+          s.setFocusRestoreElement(el);
+          s.setCommandCenterOpen(true);
+        } else {
+          s.setCommandCenterOpen(false);
+        }
         return;
       }
 

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

@@ -6,12 +6,16 @@ interface KeyboardNavState {
   cmdOrCtrlDown: boolean;
   /** Sidebar-visible agent keys (up to 9), in top-to-bottom visual order. */
   sidebarShortcutAgentKeys: string[];
+  /** Captured focus target when opening the command center (web only). */
+  focusRestoreElement: HTMLElement | null;
   focusChatInputRequest: { id: number; agentKey: string | null } | null;
 
   setCommandCenterOpen: (open: boolean) => void;
   setAltDown: (down: boolean) => void;
   setCmdOrCtrlDown: (down: boolean) => void;
   setSidebarShortcutAgentKeys: (keys: string[]) => void;
+  setFocusRestoreElement: (el: HTMLElement | null) => void;
+  takeFocusRestoreElement: () => HTMLElement | null;
   requestFocusChatInput: (agentKey: string | null) => void;
   clearFocusChatInputRequest: (id: number) => void;
   resetModifiers: () => void;
@@ -22,12 +26,22 @@ export const useKeyboardNavStore = create<KeyboardNavState>((set) => ({
   altDown: false,
   cmdOrCtrlDown: false,
   sidebarShortcutAgentKeys: [],
+  focusRestoreElement: null,
   focusChatInputRequest: null,
 
   setCommandCenterOpen: (open) => set({ commandCenterOpen: open }),
   setAltDown: (down) => set({ altDown: down }),
   setCmdOrCtrlDown: (down) => set({ cmdOrCtrlDown: down }),
   setSidebarShortcutAgentKeys: (keys) => set({ sidebarShortcutAgentKeys: keys }),
+  setFocusRestoreElement: (el) => set({ focusRestoreElement: el }),
+  takeFocusRestoreElement: () => {
+    let value: HTMLElement | null = null;
+    set((state) => {
+      value = state.focusRestoreElement;
+      return { focusRestoreElement: null };
+    });
+    return value;
+  },
   requestFocusChatInput: (agentKey) =>
     set((state) => ({
       focusChatInputRequest: {