move store to tool context

Author: Light2Dark

frontend/src/components/chat/chat-panel.tsx

@@ -5,7 +5,7 @@ import { useChat } from "@ai-sdk/react";
 import { storePrompt } from "@marimo-team/codemirror-ai";
 import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
 import { DefaultChatTransport, type ToolUIPart } from "ai";
-import { useAtom, useAtomValue, useSetAtom } from "jotai";
+import { useAtom, useAtomValue, useSetAtom, useStore } from "jotai";
 import {
   AtSignIcon,
   BotMessageSquareIcon,
@@ -528,6 +528,7 @@ const ChatPanelBody = () => {
   const { invokeAiTool, sendRun } = useRequestClient();
 
   const activeChatId = activeChat?.id;
+  const store = useStore();
 
   const { addStagedCell } = useStagedAICellsActions();
   const { createNewCell, prepareForRun } = useCellActions();
@@ -536,6 +537,7 @@ const ChatPanelBody = () => {
     createNewCell,
     prepareForRun,
     sendRun,
+    store,
   };
 
   const {

frontend/src/components/editor/ai/add-cell-with-ai.tsx

@@ -92,6 +92,7 @@ export const AddCellWithAI: React.FC<{
 
   const { createNewCell, prepareForRun } = useCellActions();
   const toolContext: ToolNotebookContext = {
+    store,
     addStagedCell,
     createNewCell,
     prepareForRun,

frontend/src/core/ai/tools/__tests__/edit-notebook-tool.test.ts

@@ -70,6 +70,7 @@ describe("EditNotebookTool", () => {
     createNewCell: ReturnType<typeof vi.fn>;
     prepareForRun: ReturnType<typeof vi.fn>;
     sendRun: ReturnType<typeof vi.fn>;
+    store: ReturnType<typeof getDefaultStore>;
   };
 
   beforeEach(() => {
@@ -86,8 +87,9 @@ describe("EditNotebookTool", () => {
       createNewCell: vi.fn(),
       prepareForRun: vi.fn(),
       sendRun: vi.fn().mockResolvedValue(null),
+      store,
     };
-    tool = new EditNotebookTool(store);
+    tool = new EditNotebookTool();
 
     cellId1 = "cell-1" as CellId;
     cellId2 = "cell-2" as CellId;

frontend/src/core/ai/tools/__tests__/run-cells-tool.test.ts

@@ -31,6 +31,7 @@ describe("RunStaleCellsTool", () => {
     createNewCell: ReturnType<typeof vi.fn>;
     prepareForRun: ReturnType<typeof vi.fn>;
     sendRun: ReturnType<typeof vi.fn>;
+    store: ReturnType<typeof getDefaultStore>;
   };
 
   beforeEach(() => {
@@ -40,9 +41,10 @@ describe("RunStaleCellsTool", () => {
       createNewCell: vi.fn(),
       prepareForRun: vi.fn(),
       sendRun: vi.fn().mockResolvedValue(null),
+      store,
     };
 
-    tool = new RunStaleCellsTool(store);
+    tool = new RunStaleCellsTool();
 
     cellId1 = "cell-1" as CellId;
     cellId2 = "cell-2" as CellId;

frontend/src/core/ai/tools/base.ts

@@ -4,6 +4,7 @@ import { z } from "zod";
 import type { CreateNewCellAction } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import type { RunRequest } from "@/core/network/types";
+import type { JotaiStore } from "@/core/state/jotai";
 import type { Edit } from "../staged-cells";
 import type { CopilotMode } from "./registry";
 
@@ -101,6 +102,7 @@ export interface ToolDescription {
 
 /** Utility functions for tools to interact with the notebook */
 export interface ToolNotebookContext {
+  store: JotaiStore;
   addStagedCell: (payload: { cellId: CellId; edit: Edit }) => void;
   createNewCell: (payload: CreateNewCellAction) => void;
   prepareForRun: (payload: { cellId: CellId }) => void;

frontend/src/core/ai/tools/edit-notebook-tool.ts

@@ -10,7 +10,6 @@ import {
 } from "@/core/cells/cells";
 import { CellId } from "@/core/cells/ids";
 import { updateEditorCodeFromPython } from "@/core/codemirror/language/utils";
-import type { JotaiStore } from "@/core/state/jotai";
 import type { CellColumnId } from "@/utils/id-tree";
 import {
   type AiTool,
@@ -87,28 +86,23 @@ export type EditType = EditOperation["type"];
 export class EditNotebookTool
   implements AiTool<EditNotebookInput, ToolOutputBase>
 {
-  private store: JotaiStore;
   readonly name = "edit_notebook_tool";
   readonly description = description;
   readonly schema = editNotebookSchema;
   readonly outputSchema = toolOutputBaseSchema;
   readonly mode: CopilotMode[] = ["agent"];
 
-  constructor(store: JotaiStore) {
-    this.store = store;
-  }
-
   handler = async (
     { edit }: EditNotebookInput,
     toolContext: ToolNotebookContext,
   ): Promise<ToolOutputBase> => {
-    const { addStagedCell, createNewCell } = toolContext;
+    const { addStagedCell, createNewCell, store } = toolContext;
 
     switch (edit.type) {
       case "update_cell": {
         const { cellId, code } = edit;
 
-        const notebook = this.store.get(notebookAtom);
+        const notebook = store.get(notebookAtom);
         this.validateCellIdExists(cellId, notebook);
         const editorView = this.getCellEditorView(cellId, notebook);
 
@@ -133,7 +127,7 @@ export class EditNotebookTool
         const newCellId = CellId.create();
 
         if (typeof position === "object") {
-          const notebook = this.store.get(notebookAtom);
+          const notebook = store.get(notebookAtom);
           if ("cellId" in position) {
             this.validateCellIdExists(position.cellId, notebook);
             notebookPosition = position.cellId;
@@ -164,7 +158,7 @@ export class EditNotebookTool
       case "delete_cell": {
         const { cellId } = edit;
 
-        const notebook = this.store.get(notebookAtom);
+        const notebook = store.get(notebookAtom);
         this.validateCellIdExists(cellId, notebook);
 
         const editorView = this.getCellEditorView(cellId, notebook);

frontend/src/core/ai/tools/registry.ts

@@ -3,7 +3,6 @@
 import type { components } from "@marimo-team/marimo-api";
 import { Memoize } from "typescript-memoize";
 import { type ZodObject, z } from "zod";
-import { store } from "@/core/state/jotai";
 import {
   type AiTool,
   ToolExecutionError,
@@ -131,6 +130,6 @@ export class FrontendToolRegistry {
 }
 
 export const FRONTEND_TOOL_REGISTRY = new FrontendToolRegistry([
-  new EditNotebookTool(store),
-  new RunStaleCellsTool(store),
+  new EditNotebookTool(),
+  new RunStaleCellsTool(),
 ]);

frontend/src/core/ai/tools/run-cells-tool.ts

@@ -66,19 +66,14 @@ export class RunStaleCellsTool
       .optional(),
   }) satisfies z.ZodType<RunStaleCellsOutput>;
   readonly mode: CopilotMode[] = ["agent"];
-  private store: JotaiStore;
-
-  constructor(store: JotaiStore) {
-    this.store = store;
-  }
 
   handler = async (
     _args: EmptyToolInput,
     toolContext: ToolNotebookContext,
   ): Promise<RunStaleCellsOutput> => {
-    const { prepareForRun, sendRun } = toolContext;
+    const { prepareForRun, sendRun, store } = toolContext;
 
-    const notebook = this.store.get(notebookAtom);
+    const notebook = store.get(notebookAtom);
     const staleCells = staleCellIds(notebook);
 
     if (staleCells.length === 0) {
@@ -96,7 +91,7 @@ export class RunStaleCellsTool
     });
 
     // Wait for all cells to finish executing
-    const allCellsFinished = await this.waitForCellsToFinish(staleCells);
+    const allCellsFinished = await this.waitForCellsToFinish(store, staleCells);
     if (!allCellsFinished) {
       return {
         status: "success",
@@ -106,7 +101,7 @@ export class RunStaleCellsTool
     }
 
     // Get notebook state after cells have finished
-    const updatedNotebook = this.store.get(notebookAtom);
+    const updatedNotebook = store.get(notebookAtom);
 
     const cellsToOutput = new Map<CellId, CellOutput | null>();
     let resultMessage = "";
@@ -192,6 +187,7 @@ export class RunStaleCellsTool
    * Returns true if all cells finished executing, false if the timeout was reached
    */
   private async waitForCellsToFinish(
+    store: JotaiStore,
     cellIds: CellId[],
     timeout = 30_000,
   ): Promise<boolean> {
@@ -207,7 +203,7 @@ export class RunStaleCellsTool
     };
 
     // If already finished, return immediately
-    if (checkAllFinished(this.store.get(notebookAtom))) {
+    if (checkAllFinished(store.get(notebookAtom))) {
       return true;
     }