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

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 }),
 }));
-