Fix Cmd+K focus restore and focus input

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

@@ -28,6 +28,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;
@@ -98,6 +99,8 @@ export function AgentInputArea({
   const [selectedImages, setSelectedImages] = useState<ImageAttachment[]>([]);
   const [isCancellingAgent, setIsCancellingAgent] = useState(false);
   const [commandSelectedIndex, setCommandSelectedIndex] = useState(0);
+  const focusChatInputRequest = useKeyboardNavStore((s) => s.focusChatInputRequest);
+  const clearFocusChatInputRequest = useKeyboardNavStore((s) => s.clearFocusChatInputRequest);
 
   // Command autocomplete logic
   const showCommandAutocomplete = userInput.startsWith("/") && !userInput.includes(" ");
@@ -398,6 +401,19 @@ export function AgentInputArea({
     queueMessage(payload.text, payload.images);
   }, []);
 
+  useEffect(() => {
+    const req = focusChatInputRequest;
+    if (!req) return;
+    if (req.agentKey !== `${serverId}:${agentId}`) return;
+
+    requestAnimationFrame(() => {
+      requestAnimationFrame(() => {
+        messageInputRef.current?.focus();
+        clearFocusChatInputRequest();
+      });
+    });
+  }, [agentId, clearFocusChatInputRequest, focusChatInputRequest, serverId]);
+
   // Handle command selection from autocomplete
   const handleCommandSelect = useCallback(
     (cmd: { name: string; description: string; argumentHint: string }) => {

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 takeFocusRestoreElement = useKeyboardNavStore((s) => s.takeFocusRestoreElement);
   const inputRef = useRef<TextInput>(null);
+  const didNavigateRef = useRef(false);
+  const prevOpenRef = useRef(open);
   const [query, setQuery] = useState("");
   const [activeIndex, setActiveIndex] = useState(0);
 
@@ -59,17 +63,39 @@ export function CommandCenter() {
   }, [agents, query]);
 
   useEffect(() => {
+    const prevOpen = prevOpenRef.current;
+    prevOpenRef.current = open;
+
     if (!open) {
       setQuery("");
       setActiveIndex(0);
+
+      if (prevOpen && !didNavigateRef.current) {
+        const el = takeFocusRestoreElement();
+        if (el && el.isConnected) {
+          // Modal unmount can steal focus; restore on next tick.
+          requestAnimationFrame(() => {
+            requestAnimationFrame(() => {
+              try {
+                el.focus();
+              } catch {
+                // ignore
+              }
+            });
+          });
+        }
+      }
+
       return;
     }
 
+    didNavigateRef.current = false;
+
     const id = setTimeout(() => {
       inputRef.current?.focus();
     }, 0);
     return () => clearTimeout(id);
-  }, [open]);
+  }, [open, takeFocusRestoreElement]);
 
   useEffect(() => {
     if (!open) return;
@@ -84,16 +110,20 @@ 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}`);
+      // Don't restore focus back to the prior element after we navigate.
+      takeFocusRestoreElement();
       setOpen(false);
       navigate(`/agent/${agent.serverId}/${agent.id}` as any);
     },
-    [pathname, setOpen]
+    [pathname, requestFocusChatInput, setOpen, takeFocusRestoreElement]
   );
 
   useEffect(() => {