Notebook controls with special LiveStore events for bulk actions (#600)

markmiro · web-flow · commit c4c3ad70022e · 2025-10-29T10:12:22.000-07:00
Depends on #657 Added controls (new LiveStore events behind a feature flag): - Run all cells - Restart and run all cells (placeholder for now, shows toast) - Stop all cells - Clear all outputs - Delete all cells Commits: * WIP * WIP first pass dropdown menu * Move hide AI cells button * Add toasts for stopping cells * Implement hide AI cells * Show badge if cells running or AI cells are hidden * Use standard function for generating a queue id * Fix stop all cells * Do filtering on sqlite side * Show message if user adds AI cell if AI cells are hidden * Add warning if hiding AI cells * Use a single event for running all cells * Use single event for clearing all outputs and stopping execution * Update queue id generation to resist collisions better * Rename stop to cancel * Cancel all without arguments * Use past tense * Clear all with almost no arguments * Update toast * Don't return ops * Fix getting cells * Remove unused * Move runnableCellsWithIndices$ * Remove hide AI cells NB control * Change menu titles * Fix type error * Show live running status * Remove badge * merge: main to nb-controls-2 * Put bulk nb controls behind a feature flag * Remove unused * Move new queries into one place * Add extra checks * Undo schema change * Remove unimportant changes
diff --git a/src/components/notebook/cell/ExecutableCell.tsx b/src/components/notebook/cell/ExecutableCell.tsx
@@ -43,6 +43,7 @@ import { MaybeCellOutputs } from "@/components/outputs/MaybeCellOutputs.js";
 import { useToolApprovals } from "@/hooks/useToolApprovals.js";
 import { AiToolApprovalOutput } from "../../outputs/shared-with-iframe/AiToolApprovalOutput.js";
 import { cn } from "@/lib/utils.js";
+import { generateQueueId } from "@/util/queue-id.js";
 import { useTrpc } from "@/components/TrpcProvider.js";
 import { cycleCellType } from "@/util/cycle-cell-type.js";
 import { useFeatureFlag } from "@/contexts/FeatureFlagContext.js";
@@ -246,15 +247,12 @@ export const ExecutableCell: React.FC<ExecutableCellProps> = ({
       );
 
       // Generate unique queue ID
-      const queueId = `exec-${Date.now()}-${Math.random()
-        .toString(36)
-        .slice(2)}`;
       const executionCount = (cell.executionCount || 0) + 1;
 
       // Add to execution queue - runtimes will pick this up
       store.commit(
         events.executionRequested({
-          queueId,
+          queueId: generateQueueId(),
           cellId: cell.id,
           executionCount,
           requestedBy: userId,
diff --git a/src/components/notebooks/notebook/NotebookControls.tsx b/src/components/notebooks/notebook/NotebookControls.tsx
@@ -16,30 +16,157 @@ import { toast } from "sonner";
 import { useCreateNotebookAndNavigate } from "../dashboard/helpers";
 import type { NotebookProcessed } from "../types";
 
+import { useAuthenticatedUser } from "@/auth/index.js";
+import { useDebug } from "@/components/debug/debug-mode";
+import { Spinner } from "@/components/ui/Spinner";
+import { useFeatureFlag } from "@/contexts/FeatureFlagContext";
+import { useRuntimeHealth } from "@/hooks/useRuntimeHealth";
+import { runnableCellsWithIndices$, runningCells$ } from "@/queries";
+import { generateQueueId } from "@/util/queue-id";
+import { useQuery, useStore } from "@livestore/react";
+import { CellData, events, queries } from "@runtimed/schema";
+import { Eraser, Play, Square, Undo2 } from "lucide-react";
+import { useCallback } from "react";
+
 export function NotebookControls({
   notebook,
 }: {
   notebook: NotebookProcessed;
 }) {
+  const allowBulkNotebookControls = useFeatureFlag("bulk-notebook-controls");
+  const { store } = useStore();
+  const userId = useAuthenticatedUser();
+  const debug = useDebug();
+
+  const cellQueue = useQuery(runningCells$);
+
+  const handleCancelAll = useCallback(() => {
+    if (!allowBulkNotebookControls) {
+      toast.info("Bulk notebook controls are not enabled");
+      return;
+    }
+    if (cellQueue.length === 0) {
+      toast.info("No cells to stop");
+      return;
+    }
+    store.commit(events.allExecutionsCancelled());
+    toast.info("Cancelled all executions");
+  }, [store, cellQueue, allowBulkNotebookControls]);
+
+  const handleClearAllOutputs = useCallback(() => {
+    if (!allowBulkNotebookControls) {
+      toast.info("Bulk notebook controls are not enabled");
+      return;
+    }
+    store.commit(events.allOutputsCleared({ clearedBy: userId }));
+  }, [store, userId, allowBulkNotebookControls]);
+
   return (
-    <div className="flex items-center gap-1">
+    <div className="flex items-center gap-2">
+      {allowBulkNotebookControls && (
+        <ActiveBulkNotebookActions
+          cellQueue={cellQueue}
+          onCancelAll={handleCancelAll}
+        />
+      )}
       <DropdownMenu>
         <DropdownMenuTrigger asChild>
           <Button variant="ghost" size="sm" className="relative">
             <MoreHorizontal className="h-4 w-4" />
           </Button>
         </DropdownMenuTrigger>
         <DropdownMenuContent align="end">
+          {allowBulkNotebookControls && (
+            <BulkNotebookActions
+              cellQueue={cellQueue}
+              onCancelAll={handleCancelAll}
+              onClearAllOutputs={handleClearAllOutputs}
+            />
+          )}
           <CreateNotebookAction />
           <DuplicateAction notebook={notebook} />
           <DropdownMenuSeparator />
+          {debug.enabled && <DeleteAllCellsAction />}
           <DeleteAction notebook={notebook} />
         </DropdownMenuContent>
       </DropdownMenu>
     </div>
   );
 }
 
+function ActiveBulkNotebookActions({
+  cellQueue,
+  onCancelAll,
+}: {
+  cellQueue: readonly CellData[];
+  onCancelAll: () => void;
+}) {
+  if (cellQueue.length === 0) {
+    return null;
+  }
+
+  return (
+    <>
+      <span className="flex items-center gap-1">
+        <Spinner size="md" />
+        <span className="text-muted-foreground text-sm">
+          Running {cellQueue.length} cells
+        </span>
+      </span>
+      <Button variant="outline" size="sm" onClick={onCancelAll}>
+        <Square />
+        Stop All
+      </Button>
+    </>
+  );
+}
+
+function BulkNotebookActions({
+  cellQueue,
+  onCancelAll,
+  onClearAllOutputs,
+}: {
+  cellQueue: readonly CellData[];
+  onCancelAll: () => void;
+  onClearAllOutputs: () => void;
+}) {
+  const { runAllCells } = useRunAllCells();
+  const { restartAndRunAllCells } = useRestartAndRunAllCells();
+
+  return (
+    <>
+      {cellQueue.length > 0 ? (
+        <DropdownMenuItem onClick={onCancelAll}>
+          <Square />
+          Stop All
+        </DropdownMenuItem>
+      ) : (
+        <>
+          <DropdownMenuItem onClick={runAllCells}>
+            <Play />
+            Run All Code Cells
+          </DropdownMenuItem>
+          <DropdownMenuItem onClick={restartAndRunAllCells}>
+            <span className="relative">
+              <Play className="h-4 w-4" />
+              <Undo2
+                className="absolute bottom-0 left-0 size-3 -translate-x-[3px] translate-y-[3px] rounded-full bg-white p-[1px] text-gray-700"
+                strokeWidth={3}
+              />
+            </span>
+            Restart and Run All Code Cells
+          </DropdownMenuItem>
+        </>
+      )}
+      <DropdownMenuItem onClick={onClearAllOutputs}>
+        <Eraser />
+        Clear All Outputs
+      </DropdownMenuItem>
+      <DropdownMenuSeparator />
+    </>
+  );
+}
+
 function CreateNotebookAction() {
   const createNotebookAndNavigate = useCreateNotebookAndNavigate();
 
@@ -126,3 +253,74 @@ function DeleteAction({ notebook }: { notebook: NotebookProcessed }) {
     </DropdownMenuItem>
   );
 }
+
+function DeleteAllCellsAction() {
+  const { deleteAllCells } = useDeleteAllCells();
+  const { confirm } = useConfirm();
+
+  const handleDeleteAllCells = async () => {
+    confirm({
+      title: "Delete All Cells",
+      description: "Are you sure you want to delete all cells?",
+      onConfirm: deleteAllCells,
+      actionButtonText: "Delete All Cells",
+    });
+  };
+  return (
+    <DropdownMenuItem variant="destructive" onSelect={handleDeleteAllCells}>
+      <Trash2 />
+      DEBUG: Delete All Cells
+    </DropdownMenuItem>
+  );
+}
+
+function useDeleteAllCells() {
+  const { store } = useStore();
+  const cells = useQuery(queries.cellsWithIndices$);
+  const userId = useAuthenticatedUser();
+
+  const deleteAllCells = useCallback(() => {
+    cells.forEach((cell) => {
+      store.commit(events.cellDeleted({ id: cell.id, actorId: userId }));
+    });
+  }, [cells, store, userId]);
+
+  return { deleteAllCells };
+}
+
+function useRunAllCells() {
+  const { store } = useStore();
+  const cells = useQuery(runnableCellsWithIndices$);
+  const userId = useAuthenticatedUser();
+
+  const runAllCells = useCallback(() => {
+    store.commit(
+      events.multipleExecutionRequested({
+        requestedBy: userId,
+        cellsInfo: [
+          ...cells.map((cell) => ({
+            id: cell.id,
+            executionCount: (cell.executionCount || 0) + 1,
+            queueId: generateQueueId(),
+          })),
+        ],
+      })
+    );
+  }, [store, userId, cells]);
+
+  return { runAllCells };
+}
+
+function useRestartAndRunAllCells() {
+  const { hasActiveRuntime } = useRuntimeHealth();
+
+  const restartAndRunAllCells = useCallback(() => {
+    if (hasActiveRuntime) {
+      toast.info("Restart your runtime manually and run all cells");
+    } else {
+      toast.error("No active runtime found");
+    }
+  }, [hasActiveRuntime]);
+
+  return { restartAndRunAllCells };
+}
diff --git a/src/contexts/FeatureFlagContext.tsx b/src/contexts/FeatureFlagContext.tsx
@@ -6,6 +6,8 @@ export interface FeatureFlags {
   "test-flag": boolean;
   "ipynb-export": boolean;
   "file-upload": boolean;
+  /** Whether to enable the notebook controls */
+  "bulk-notebook-controls": boolean;
   /** Show AI capabilities in the AI cell dropdown. We'd enable this by default if we support vision or allow choosing models that don't have tool support. */
   "show-ai-capabilities": boolean;
   "user-saved-prompt": boolean;
@@ -16,7 +18,9 @@ const DEFAULT_FLAGS: FeatureFlags = {
   "test-flag": false,
   "ipynb-export": false,
   "file-upload": false,
+  "bulk-notebook-controls": false,
   "show-ai-capabilities": false,
+
   "user-saved-prompt": false,
 } as const;
 
diff --git a/src/queries/index.ts b/src/queries/index.ts
@@ -24,3 +24,19 @@ export const availableFiles$ = queryDb(
   tables.files.select().where({ deletedAt: null }),
   { label: "files.availableFiles" }
 );
+
+export const runnableCellsWithIndices$ = queryDb(
+  tables.cells
+    .select("id", "fractionalIndex", "cellType", "executionCount")
+    .where({ cellType: { op: "IN", value: ["code", "sql"] } })
+    .orderBy("fractionalIndex", "asc"),
+  { label: "cells.withIndices.runnable" }
+);
+
+export const runningCells$ = queryDb(
+  tables.cells
+    .select()
+    .where({ executionState: { op: "IN", value: ["running", "queued"] } })
+    .orderBy("fractionalIndex", "asc"),
+  { label: "cells.runningCells" }
+);
diff --git a/src/util/queue-id.ts b/src/util/queue-id.ts
@@ -0,0 +1,9 @@
+/**
+ * Generates a unique queue ID for an execution request
+ */
+export function generateQueueId() {
+  // Date should prevent most of the collisions, but not as reliable when making bulk requests
+  // 6 char length + alphabet of 36 chars = 6K IDs to have 1% chance of collision
+  // See: https://zelark.github.io/nano-id-cc/
+  return `exec-${Date.now()}-${Math.random().toString(36).slice(6)}`;
+}
