@marimo-team/islands 0.20.5-dev33 → 0.20.5-dev35

package/dist/main.js

@@ -70710,7 +70710,7 @@ Image URL: ${r.imageUrl}`)), contextToXml({
       return Logger.warn("Failed to get version from mount config"), null;
     }
   }
-  const marimoVersionAtom = atom(getVersionFromMountConfig() || "0.20.5-dev33"), showCodeInRunModeAtom = atom(true);
+  const marimoVersionAtom = atom(getVersionFromMountConfig() || "0.20.5-dev35"), showCodeInRunModeAtom = atom(true);
   atom(null);
   var import_compiler_runtime$89 = require_compiler_runtime();
   function useKeydownOnElement(e, r) {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@marimo-team/islands",
-  "version": "0.20.5-dev33",
+  "version": "0.20.5-dev35",
   "main": "dist/main.js",
   "types": "dist/index.d.ts",
   "type": "module",

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

@@ -488,6 +488,133 @@ describe("RunStaleCellsTool", () => {
       // Output should be capped (2000 chars content + "Output:\n" prefix + truncation message)
       expect(output.length).toBeLessThan(2200);
     });
+
+    it("should omit output for cells that exceed total output budget", async () => {
+      const cellIds = Array.from(
+        { length: 25 },
+        (_, i) => `budget-cell-${i}` as CellId,
+      );
+      const cellData: Record<string, { code: string; edited: boolean }> = {};
+      for (const id of cellIds) {
+        cellData[id] = { code: "x = 1", edited: true };
+      }
+
+      const notebook = MockNotebook.notebookState({ cellData });
+      store.set(notebookAtom, notebook);
+
+      vi.mocked(runCells).mockImplementation(async () => {
+        const updatedNotebook = store.get(notebookAtom);
+        for (const id of cellIds) {
+          updatedNotebook.cellRuntime[id] = {
+            ...updatedNotebook.cellRuntime[id],
+            status: "idle",
+          };
+        }
+        store.set(notebookAtom, updatedNotebook);
+      });
+
+      // Each cell produces ~2008 chars of formatted output ("Output:\n" + 2000 chars).
+      // After 20 cells the running total exceeds MAX_TOOL_OUTPUT_CHARS (40,000).
+      const content = "a".repeat(2000);
+      vi.mocked(getCellContextData).mockReturnValue({
+        cellOutput: {
+          outputType: "text",
+          processedContent: content,
+          imageUrl: null,
+          output: { mimetype: "text/plain", data: content },
+        },
+        consoleOutputs: null,
+        cellName: "cell",
+      } as never);
+
+      const result = await tool.handler({}, toolContext as never);
+
+      expect(result.cellsToOutput?.[cellIds[0]]?.cellOutput).toContain(
+        "Output:",
+      );
+      expect(result.cellsToOutput?.[cellIds[24]]?.cellOutput).toBe(
+        "Cell executed (output omitted due to context limits).",
+      );
+    });
+
+    it("should use higher truncation limit for error outputs", async () => {
+      const notebook = MockNotebook.notebookState({
+        cellData: {
+          [cellId1]: { code: "raise Exception()", edited: true },
+        },
+      });
+      store.set(notebookAtom, notebook);
+
+      vi.mocked(runCells).mockImplementation(async () => {
+        const updatedNotebook = store.get(notebookAtom);
+        updatedNotebook.cellRuntime[cellId1] = {
+          ...updatedNotebook.cellRuntime[cellId1],
+          status: "idle",
+        };
+        store.set(notebookAtom, updatedNotebook);
+      });
+
+      // 2500 chars sits between MAX_TEXT_OUTPUT_CHARS (2000) and MAX_ERROR_OUTPUT_CHARS (3000)
+      const errorContent = "E".repeat(2500);
+      vi.mocked(getCellContextData).mockReturnValue({
+        cellOutput: {
+          outputType: "text",
+          processedContent: errorContent,
+          imageUrl: null,
+          output: {
+            mimetype: "application/vnd.marimo+error",
+            data: errorContent,
+          },
+        },
+        consoleOutputs: null,
+        cellName: "cell1",
+      } as never);
+
+      const result = await tool.handler({}, toolContext as never);
+
+      const output = result.cellsToOutput?.[cellId1]?.cellOutput ?? "";
+      expect(output).not.toContain("[TRUNCATED:");
+      expect(output).toContain(errorContent);
+    });
+
+    it("should truncate large console output", async () => {
+      const notebook = MockNotebook.notebookState({
+        cellData: {
+          [cellId1]: { code: 'print("x" * 10000)', edited: true },
+        },
+      });
+      store.set(notebookAtom, notebook);
+
+      vi.mocked(runCells).mockImplementation(async () => {
+        const updatedNotebook = store.get(notebookAtom);
+        updatedNotebook.cellRuntime[cellId1] = {
+          ...updatedNotebook.cellRuntime[cellId1],
+          status: "idle",
+        };
+        store.set(notebookAtom, updatedNotebook);
+      });
+
+      const largeConsoleText = "x".repeat(10_000);
+      vi.mocked(getCellContextData).mockReturnValue({
+        cellOutput: null,
+        consoleOutputs: [
+          {
+            outputType: "text",
+            processedContent: largeConsoleText,
+            imageUrl: null,
+            output: { mimetype: "text/plain", data: largeConsoleText },
+          },
+        ],
+        cellName: "cell1",
+      } as never);
+
+      const result = await tool.handler({}, toolContext as never);
+
+      const consoleOutput =
+        result.cellsToOutput?.[cellId1]?.consoleOutput ?? "";
+      expect(consoleOutput).toContain("[TRUNCATED:");
+      expect(consoleOutput.length).toBeLessThan(2200);
+    });
   });
 
   describe("cell execution completion", () => {

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

@@ -97,7 +97,7 @@ export class RunStaleCellsTool
 
     await runCells({
       cellIds: staleCells,
-      sendRun: sendRun,
+      sendRun,
       prepareForRun,
       notebook,
     });
@@ -121,8 +121,8 @@ export class RunStaleCellsTool
     const updatedNotebook = store.get(notebookAtom);
 
     const cellsToOutput = new Map<CellId, CellOutput | null>();
-    let resultMessage = "";
     let outputHasErrors = false;
+    let hasAnyConsoleOutput = false;
     let totalOutputChars = 0;
 
     for (const cellId of staleCells) {
@@ -135,12 +135,10 @@ export class RunStaleCellsTool
       const hasConsoleOutput = consoleOutputs && consoleOutputs.length > 0;
 
       // Track errors regardless of budget
-      if (cellOutput && this.outputHasErrors(cellOutput)) {
-        outputHasErrors = true;
-      }
       if (
-        hasConsoleOutput &&
-        consoleOutputs.some((output) => this.outputHasErrors(output))
+        (cellOutput && this.outputHasErrors(cellOutput)) ||
+        (hasConsoleOutput &&
+          consoleOutputs.some((output) => this.outputHasErrors(output)))
       ) {
         outputHasErrors = true;
       }
@@ -167,6 +165,7 @@ export class RunStaleCellsTool
       }
 
       if (hasConsoleOutput) {
+        hasAnyConsoleOutput = true;
         consoleOutputString = consoleOutputs
           .map((output) => this.formatOutputString(output))
           .join("\n");
@@ -175,8 +174,6 @@ export class RunStaleCellsTool
           MAX_TEXT_OUTPUT_CHARS,
         );
         totalOutputChars += consoleOutputString.length;
-        resultMessage +=
-          "Console output represents the stdout or stderr of the cell (eg. print statements).";
       }
 
       cellsToOutput.set(cellId, {
@@ -202,57 +199,50 @@ export class RunStaleCellsTool
     return {
       status: "success",
       cellsToOutput: Object.fromEntries(cellsToOutput),
-      message: resultMessage === "" ? undefined : resultMessage,
+      message: hasAnyConsoleOutput
+        ? "Console output represents the stdout or stderr of the cell (eg. print statements)."
+        : undefined,
       next_steps: nextSteps,
     };
   };
 
   private outputHasErrors(cellOutput: BaseOutput): boolean {
-    const { output } = cellOutput;
-    if (
-      output.mimetype === "application/vnd.marimo+error" ||
-      output.mimetype === "application/vnd.marimo+traceback"
-    ) {
-      return true;
-    }
-    return false;
+    return (
+      cellOutput.output.mimetype === "application/vnd.marimo+error" ||
+      cellOutput.output.mimetype === "application/vnd.marimo+traceback"
+    );
   }
 
   private formatOutputString(cellOutput: BaseOutput): string {
-    let outputString = "";
     const { outputType, processedContent, imageUrl, output } = cellOutput;
-    if (outputType === "text") {
-      if (output.mimetype === "text/html") {
-        // text/html (e.g. plotly figures, rich dataframes) can be millions of
-        // chars and is not interpretable by LLMs — summarize instead
-        const dataLength =
-          typeof output.data === "string"
-            ? output.data.length
-            : JSON.stringify(output.data).length;
-        outputString += `HTML Output: text/html content (${dataLength.toLocaleString()} chars). Full output visible in notebook UI.`;
-      } else {
-        const isError = this.outputHasErrors(cellOutput);
-        const maxChars = isError
-          ? MAX_ERROR_OUTPUT_CHARS
-          : MAX_TEXT_OUTPUT_CHARS;
-        outputString += "Output:\n";
-        let content: string;
-        if (processedContent) {
-          content = processedContent;
-        } else if (typeof output.data === "string") {
-          content = output.data;
-        } else {
-          content = JSON.stringify(output.data);
-        }
-        outputString += this.truncateString(content, maxChars);
-      }
-    } else if (outputType === "media") {
-      outputString += `Media Output: Contains ${output.mimetype} content`;
-      if (imageUrl) {
-        outputString += `\nImage URL: ${imageUrl}`;
-      }
+
+    if (outputType === "media") {
+      const base = `Media Output: Contains ${output.mimetype} content`;
+      return imageUrl ? `${base}\nImage URL: ${imageUrl}` : base;
+    }
+
+    if (output.mimetype === "text/html") {
+      // text/html (e.g. plotly figures, rich dataframes) can be millions of
+      // chars and is not interpretable by LLMs — summarize instead
+      const dataLength =
+        typeof output.data === "string"
+          ? output.data.length
+          : JSON.stringify(output.data).length;
+      return `HTML Output: text/html content (${dataLength.toLocaleString()} chars). Full output visible in notebook UI.`;
     }
-    return outputString;
+
+    const maxChars = this.outputHasErrors(cellOutput)
+      ? MAX_ERROR_OUTPUT_CHARS
+      : MAX_TEXT_OUTPUT_CHARS;
+
+    let content = processedContent;
+    if (!content) {
+      content =
+        typeof output.data === "string"
+          ? output.data
+          : JSON.stringify(output.data);
+    }
+    return `Output:\n${this.truncateString(content, maxChars)}`;
   }
 
   private truncateString(str: string, maxLength: number): string {
@@ -261,7 +251,6 @@ export class RunStaleCellsTool
     }
     return `${str.slice(0, maxLength)}\n\n[TRUNCATED: ${str.length.toLocaleString()} → ${maxLength.toLocaleString()} chars. Full output visible in the notebook UI.]`;
   }
-
   /**
    * Wait for cells to finish executing (status becomes "idle")
    * Returns true if all cells finished executing, false if the timeout was reached