@marimo-team/islands 0.20.5-dev19 → 0.20.5-dev21

package/dist/main.js

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

package/package.json

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

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

@@ -411,6 +411,85 @@ describe("RunStaleCellsTool", () => {
     });
   });
 
+  describe("output truncation", () => {
+    it("should summarize text/html output instead of dumping raw content", async () => {
+      const notebook = MockNotebook.notebookState({
+        cellData: {
+          [cellId1]: { code: "fig.show()", 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 largeHtml = `<div>${"x".repeat(2_000_000)}</div>`;
+      vi.mocked(getCellContextData).mockReturnValue({
+        cellOutput: {
+          outputType: "text",
+          processedContent: null,
+          imageUrl: null,
+          output: { mimetype: "text/html", data: largeHtml },
+        },
+        consoleOutputs: null,
+        cellName: "cell1",
+      } as never);
+
+      const result = await tool.handler({}, toolContext as never);
+
+      expect(result.status).toBe("success");
+      const output = result.cellsToOutput?.[cellId1]?.cellOutput ?? "";
+      expect(output).toContain("HTML Output:");
+      expect(output).toContain("text/html");
+      expect(output.length).toBeLessThan(200);
+      expect(output).not.toContain(largeHtml);
+    });
+
+    it("should truncate large text output to MAX_TEXT_OUTPUT_CHARS", async () => {
+      const notebook = MockNotebook.notebookState({
+        cellData: {
+          [cellId1]: { code: "print(big_string)", 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 largeText = "a".repeat(10_000);
+      vi.mocked(getCellContextData).mockReturnValue({
+        cellOutput: {
+          outputType: "text",
+          processedContent: largeText,
+          imageUrl: null,
+          output: { mimetype: "text/plain", data: largeText },
+        },
+        consoleOutputs: null,
+        cellName: "cell1",
+      } as never);
+
+      const result = await tool.handler({}, toolContext as never);
+
+      const output = result.cellsToOutput?.[cellId1]?.cellOutput ?? "";
+      expect(output).toContain("[TRUNCATED:");
+      expect(output).toContain("Full output visible in the notebook UI.");
+      // Output should be capped (2000 chars content + "Output:\n" prefix + truncation message)
+      expect(output.length).toBeLessThan(2200);
+    });
+  });
+
   describe("cell execution completion", () => {
     it("should complete immediately if cells are already idle", async () => {
       const notebook = MockNotebook.notebookState({

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

@@ -24,6 +24,11 @@ import type { CopilotMode } from "./registry";
 const POST_EXECUTION_DELAY = 200;
 const WAIT_FOR_CELLS_TIMEOUT = 30_000;
 
+// Output size limits to prevent exceeding LLM token limits.
+const MAX_TEXT_OUTPUT_CHARS = 2000;
+const MAX_ERROR_OUTPUT_CHARS = 3000;
+const MAX_TOOL_OUTPUT_CHARS = 40_000;
+
 interface CellOutput {
   consoleOutput?: string;
   cellOutput?: string;
@@ -118,42 +123,60 @@ export class RunStaleCellsTool
     const cellsToOutput = new Map<CellId, CellOutput | null>();
     let resultMessage = "";
     let outputHasErrors = false;
+    let totalOutputChars = 0;
 
     for (const cellId of staleCells) {
       const cellContextData = getCellContextData(cellId, updatedNotebook, {
         includeConsoleOutput: true,
       });
 
-      let cellOutputString: string | undefined;
-      let consoleOutputString: string | undefined;
-
       const cellOutput = cellContextData.cellOutput;
       const consoleOutputs = cellContextData.consoleOutputs;
       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))
+      ) {
+        outputHasErrors = true;
+      }
+
       if (!cellOutput && !hasConsoleOutput) {
-        // Set null to show no output
         cellsToOutput.set(cellId, null);
         continue;
       }
 
+      // If total budget exceeded, summarize remaining cells
+      if (totalOutputChars >= MAX_TOOL_OUTPUT_CHARS) {
+        cellsToOutput.set(cellId, {
+          cellOutput: "Cell executed (output omitted due to context limits).",
+        });
+        continue;
+      }
+
+      let cellOutputString: string | undefined;
+      let consoleOutputString: string | undefined;
+
       if (cellOutput) {
         cellOutputString = this.formatOutputString(cellOutput);
-        if (this.outputHasErrors(cellOutput)) {
-          outputHasErrors = true;
-        }
+        totalOutputChars += cellOutputString.length;
       }
 
       if (hasConsoleOutput) {
         consoleOutputString = consoleOutputs
           .map((output) => this.formatOutputString(output))
           .join("\n");
+        consoleOutputString = this.truncateString(
+          consoleOutputString,
+          MAX_TEXT_OUTPUT_CHARS,
+        );
+        totalOutputChars += consoleOutputString.length;
         resultMessage +=
           "Console output represents the stdout or stderr of the cell (eg. print statements).";
-
-        if (consoleOutputs.some((output) => this.outputHasErrors(output))) {
-          outputHasErrors = true;
-        }
       }
 
       cellsToOutput.set(cellId, {
@@ -199,13 +222,29 @@ export class RunStaleCellsTool
     let outputString = "";
     const { outputType, processedContent, imageUrl, output } = cellOutput;
     if (outputType === "text") {
-      outputString += "Output:\n";
-      if (processedContent) {
-        outputString += processedContent;
-      } else if (typeof output.data === "string") {
-        outputString += output.data;
+      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 {
-        outputString += JSON.stringify(output.data);
+        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`;
@@ -216,6 +255,13 @@ export class RunStaleCellsTool
     return outputString;
   }
 
+  private truncateString(str: string, maxLength: number): string {
+    if (str.length <= maxLength) {
+      return str;
+    }
+    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