Improve review comment UI

packages/app/src/workspace/focus.tsx

@@ -0,0 +1,75 @@
+import React, {
+  createContext,
+  useCallback,
+  useContext,
+  useEffect,
+  useMemo,
+  useRef,
+  type ReactNode,
+} from "react";
+import { useWorkspaceLayoutStore } from "@/stores/workspace-layout-store";
+
+interface WorkspaceFocusContextValue {
+  workspaceKey: string | null;
+  unfocusPane: (workspaceKey: string) => string | null;
+  restorePaneFocus: (workspaceKey: string, token: string) => void;
+}
+
+export interface WorkspaceFocusRestoration {
+  unfocus: () => void;
+  restore: () => void;
+}
+
+const WorkspaceFocusContext = createContext<WorkspaceFocusContextValue | null>(null);
+
+const noopFocusRestoration: WorkspaceFocusRestoration = {
+  unfocus: () => {},
+  restore: () => {},
+};
+
+export function WorkspaceFocusProvider({
+  workspaceKey,
+  children,
+}: {
+  workspaceKey: string | null;
+  children: ReactNode;
+}) {
+  const unfocusPane = useWorkspaceLayoutStore((state) => state.unfocusPane);
+  const restorePaneFocus = useWorkspaceLayoutStore((state) => state.restorePaneFocus);
+  const value = useMemo<WorkspaceFocusContextValue>(
+    () => ({
+      workspaceKey,
+      unfocusPane,
+      restorePaneFocus,
+    }),
+    [restorePaneFocus, unfocusPane, workspaceKey],
+  );
+
+  return <WorkspaceFocusContext.Provider value={value}>{children}</WorkspaceFocusContext.Provider>;
+}
+
+export function useWorkspaceFocusRestoration(): WorkspaceFocusRestoration {
+  const context = useContext(WorkspaceFocusContext);
+  const tokenRef = useRef<string | null>(null);
+
+  const restore = useCallback(() => {
+    const token = tokenRef.current;
+    if (!context?.workspaceKey || !token) {
+      tokenRef.current = null;
+      return;
+    }
+    tokenRef.current = null;
+    context.restorePaneFocus(context.workspaceKey, token);
+  }, [context]);
+
+  const unfocus = useCallback(() => {
+    if (!context?.workspaceKey || tokenRef.current) {
+      return;
+    }
+    tokenRef.current = context.unfocusPane(context.workspaceKey);
+  }, [context]);
+
+  useEffect(() => restore, [restore]);
+
+  return context ? { unfocus, restore } : noopFocusRestoration;
+}