@marimo-team/frontend 0.17.7 → 0.17.8

package/src/core/cells/__tests__/collapseConsoleOutputs.test.ts

@@ -1,8 +1,12 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { describe, expect, it } from "vitest";
+import { StatefulOutputMessage } from "@/components/editor/output/ansi-reduce";
 import type { OutputMessage } from "@/core/kernel/messages";
-import { collapseConsoleOutputs } from "../collapseConsoleOutputs";
+import {
+  collapseConsoleOutputs,
+  maybeMakeOutputStateful,
+} from "../collapseConsoleOutputs";
 
 describe("collapseConsoleOutputs", () => {
   it("should collapse last two text/plain outputs on the same channel", () => {
@@ -236,4 +240,397 @@ describe("collapseConsoleOutputs", () => {
     // heuristic: non-text counts as 1 line ...
     expect(result[2].data).toBe("<pre>E\nF\nG\nH\n</pre>");
   });
+
+  describe("ANSI escape sequences", () => {
+    it("should handle cursor movement with collapse", () => {
+      const consoleOutputs: OutputMessage[] = [
+        {
+          mimetype: "text/plain",
+          channel: "output",
+          data: "Hello",
+          timestamp: 0,
+        },
+        {
+          mimetype: "text/plain",
+          channel: "output",
+          data: "\u001B[5DWorld", // Move cursor back 5, write World
+          timestamp: 0,
+        },
+      ];
+      const result = collapseConsoleOutputs(consoleOutputs);
+      expect(result[0].data).toBe("World");
+    });
+
+    it("should handle progress bar simulation", () => {
+      // Simulate streaming: collapse is called after each new message
+      let consoleOutputs: OutputMessage[] = [
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "Progress: 0%\r",
+          timestamp: 0,
+        },
+      ];
+
+      consoleOutputs = collapseConsoleOutputs([
+        ...consoleOutputs,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "Progress: 50%\r",
+          timestamp: 0,
+        },
+      ]);
+
+      consoleOutputs = collapseConsoleOutputs([
+        ...consoleOutputs,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "Progress: 100%",
+          timestamp: 0,
+        },
+      ]);
+
+      expect(consoleOutputs[0].data).toBe("Progress: 100%");
+    });
+
+    it("should handle tqdm-like progress bars", () => {
+      // Simulate streaming: collapse is called after each new message
+      let consoleOutputs: OutputMessage[] = [
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "Processing: |          | 0/100\r",
+          timestamp: 0,
+        },
+      ];
+
+      consoleOutputs = collapseConsoleOutputs([
+        ...consoleOutputs,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "Processing: |█████     | 50/100\r",
+          timestamp: 0,
+        },
+      ]);
+
+      consoleOutputs = collapseConsoleOutputs([
+        ...consoleOutputs,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "Processing: |██████████| 100/100",
+          timestamp: 0,
+        },
+      ]);
+
+      expect(consoleOutputs[0].data).toBe("Processing: |██████████| 100/100");
+    });
+
+    it("should handle cursor up/down movements", () => {
+      // Test that cursor movements within a single message are handled
+      const consoleOutputs: OutputMessage[] = [
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "Line 1\nLine 2\n\u001B[1AModified",
+          timestamp: 0,
+        },
+      ];
+
+      const result = collapseConsoleOutputs(consoleOutputs);
+      // After moving up 1 line and writing "Modified", Line 2 gets overwritten
+      expect(result[0].data).toBe("Line 1\nModified");
+    });
+
+    it("should handle clear screen sequence", () => {
+      let consoleOutputs: OutputMessage[] = [
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "Old content\nMore old content\n",
+          timestamp: 0,
+        },
+      ];
+
+      consoleOutputs = collapseConsoleOutputs([
+        ...consoleOutputs,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "\u001B[2J", // Clear screen
+          timestamp: 0,
+        },
+      ]);
+
+      consoleOutputs = collapseConsoleOutputs([
+        ...consoleOutputs,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "New content",
+          timestamp: 0,
+        },
+      ]);
+
+      expect(consoleOutputs[0].data).toBe("New content");
+    });
+
+    it("should handle erase line sequences", () => {
+      let consoleOutputs: OutputMessage[] = [
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "Hello World",
+          timestamp: 0,
+        },
+      ];
+
+      consoleOutputs = collapseConsoleOutputs([
+        ...consoleOutputs,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "\u001B[2K", // Erase entire line
+          timestamp: 0,
+        },
+      ]);
+
+      consoleOutputs = collapseConsoleOutputs([
+        ...consoleOutputs,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "Goodbye",
+          timestamp: 0,
+        },
+      ]);
+
+      expect(consoleOutputs[0].data).toBe("           Goodbye");
+    });
+
+    it("should handle cursor positioning", () => {
+      let consoleOutputs: OutputMessage[] = [
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "abc",
+          timestamp: 0,
+        },
+      ];
+
+      consoleOutputs = collapseConsoleOutputs([
+        ...consoleOutputs,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "\u001B[1;2H", // Move to position (1,2)
+          timestamp: 0,
+        },
+      ]);
+
+      consoleOutputs = collapseConsoleOutputs([
+        ...consoleOutputs,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "XY",
+          timestamp: 0,
+        },
+      ]);
+
+      expect(consoleOutputs[0].data).toBe("aXY");
+    });
+
+    it("should handle complex multi-line progress updates", () => {
+      // Test multi-line updates with ANSI sequences all in one message
+      const consoleOutputs: OutputMessage[] = [
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "Task 1: 0%\nTask 2: 0%\n\u001B[2A\rTask 1: 100%\n\rTask 2: 100%\n",
+          timestamp: 0,
+        },
+      ];
+      const result = collapseConsoleOutputs(consoleOutputs);
+      expect(result[0].data).toBe("Task 1: 100%\nTask 2: 100%\n");
+    });
+
+    it("should preserve ANSI state across multiple collapsed messages", () => {
+      // Simulate streaming collapses
+      let consoleOutputs: OutputMessage[] = [
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "Loading",
+          timestamp: 0,
+        },
+      ];
+
+      consoleOutputs = collapseConsoleOutputs([
+        ...consoleOutputs,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "\rLoading.",
+          timestamp: 0,
+        },
+      ]);
+
+      consoleOutputs = collapseConsoleOutputs([
+        ...consoleOutputs,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "\rLoading..",
+          timestamp: 0,
+        },
+      ]);
+
+      consoleOutputs = collapseConsoleOutputs([
+        ...consoleOutputs,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "\rLoading...",
+          timestamp: 0,
+        },
+      ]);
+
+      consoleOutputs = collapseConsoleOutputs([
+        ...consoleOutputs,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "\rComplete!  ",
+          timestamp: 0,
+        },
+      ]);
+
+      expect(consoleOutputs[0].data).toBe("Complete!  ");
+    });
+  });
+
+  describe("StatefulOutputMessage", () => {
+    it("should preserve StatefulOutputMessage through collapse", () => {
+      const statefulMessage = StatefulOutputMessage.create({
+        mimetype: "text/plain",
+        channel: "stdout",
+        data: "Hello\rWorld",
+        timestamp: 0,
+      });
+
+      const consoleOutputs: OutputMessage[] = [
+        statefulMessage,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: " Complete",
+          timestamp: 0,
+        },
+      ];
+
+      const result = collapseConsoleOutputs(consoleOutputs);
+      expect(result.length).toBe(1);
+      expect(result[0]).toBeInstanceOf(StatefulOutputMessage);
+      expect(result[0].data).toBe("World Complete");
+    });
+
+    it("should handle mixing StatefulOutputMessage and regular messages", () => {
+      let consoleOutputs: OutputMessage[] = [
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "Line 1\n",
+          timestamp: 0,
+        },
+      ];
+
+      consoleOutputs = collapseConsoleOutputs([
+        ...consoleOutputs,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "Progress: 0%",
+          timestamp: 0,
+        },
+      ]);
+
+      consoleOutputs = collapseConsoleOutputs([
+        ...consoleOutputs,
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "\rProgress: 100%",
+          timestamp: 0,
+        },
+      ]);
+
+      expect(consoleOutputs.length).toBe(1);
+      expect(consoleOutputs[0]).toBeInstanceOf(StatefulOutputMessage);
+      expect(consoleOutputs[0].data).toBe("Line 1\nProgress: 100%");
+    });
+
+    it("should handle StatefulOutputMessage with empty array input", () => {
+      const consoleOutputs: OutputMessage[] = [];
+      const result = collapseConsoleOutputs(consoleOutputs);
+      expect(result).toEqual([]);
+    });
+
+    it("should handle StatefulOutputMessage with single message", () => {
+      const statefulMessage = StatefulOutputMessage.create({
+        mimetype: "text/plain",
+        channel: "stdout",
+        data: "Hello\rWorld",
+        timestamp: 0,
+      });
+
+      const result = collapseConsoleOutputs([statefulMessage]);
+      expect(result.length).toBe(1);
+      expect(result[0]).toBeInstanceOf(StatefulOutputMessage);
+      expect(result[0].data).toBe("World");
+    });
+  });
+
+  describe("maybeMakeOutputStateful", () => {
+    it("should convert regular string output to StatefulOutputMessage", () => {
+      const output: OutputMessage = {
+        mimetype: "text/plain",
+        channel: "stdout",
+        data: "Hello",
+        timestamp: 0,
+      };
+
+      const result = maybeMakeOutputStateful(output);
+      expect(result).toBeInstanceOf(StatefulOutputMessage);
+      expect(result.data).toBe("Hello");
+    });
+
+    it("should preserve StatefulOutputMessage as-is", () => {
+      const statefulMessage = StatefulOutputMessage.create({
+        mimetype: "text/plain",
+        channel: "stdout",
+        data: "Hello",
+        timestamp: 0,
+      });
+
+      const result = maybeMakeOutputStateful(statefulMessage);
+      expect(result).toBe(statefulMessage);
+    });
+
+    it("should return non-string output as-is", () => {
+      const output: OutputMessage = {
+        mimetype: "application/json",
+        channel: "output",
+        data: { key: "value" },
+        timestamp: 0,
+      };
+
+      const result = maybeMakeOutputStateful(output);
+      expect(result).toBe(output);
+      expect(result).not.toBeInstanceOf(StatefulOutputMessage);
+    });
+  });
 });

package/src/core/cells/cell.ts

@@ -126,15 +126,14 @@ export function transitionCell(
     });
   }
 
-  if (message.console !== null) {
+  const messageConsole = message.console;
+  if (messageConsole != null) {
     // The kernel sends an empty array to clear the console; otherwise,
-    // message.console is an output that needs to be appended to the
+    // messageConsole is an output that needs to be appended to the
     // existing console outputs.
-    consoleOutputs = Array.isArray(message.console)
-      ? message.console
-      : collapseConsoleOutputs(
-          [...consoleOutputs, message.console].filter(Boolean),
-        );
+    consoleOutputs = Array.isArray(messageConsole)
+      ? collapseConsoleOutputs(messageConsole)
+      : collapseConsoleOutputs([...consoleOutputs, messageConsole]);
   }
   nextCell.consoleOutputs = consoleOutputs;
   // Derive outline from output

package/src/core/cells/collapseConsoleOutputs.tsx

@@ -1,4 +1,8 @@
 /* Copyright 2024 Marimo. All rights reserved. */
+import {
+  StatefulOutputMessage,
+  type StringOutputMessage,
+} from "@/components/editor/output/ansi-reduce";
 import type { OutputMessage } from "@/core/kernel/messages";
 import { invariant } from "@/utils/invariant";
 
@@ -13,21 +17,27 @@ export function collapseConsoleOutputs(
   const newConsoleOutputs = [...consoleOutputs];
 
   if (newConsoleOutputs.length < 2) {
-    return truncateHead(handleCarriageReturns(newConsoleOutputs), maxLines);
+    return truncateHead(
+      newConsoleOutputs.map(maybeMakeOutputStateful),
+      maxLines,
+    );
   }
 
   const lastOutput = newConsoleOutputs[newConsoleOutputs.length - 1];
   const secondLastOutput = newConsoleOutputs[newConsoleOutputs.length - 2];
 
   if (shouldCollapse(lastOutput, secondLastOutput)) {
-    invariant(typeof lastOutput.data === "string", "expected string");
-    invariant(typeof secondLastOutput.data === "string", "expected string");
+    assertStringOutputMessage(lastOutput);
+    assertStringOutputMessage(secondLastOutput);
 
-    secondLastOutput.data += lastOutput.data;
+    newConsoleOutputs[newConsoleOutputs.length - 2] = mergeLeft(
+      secondLastOutput,
+      lastOutput,
+    );
     newConsoleOutputs.pop();
   }
 
-  return truncateHead(handleCarriageReturns(newConsoleOutputs), maxLines);
+  return truncateHead(newConsoleOutputs, maxLines);
 }
 
 function shouldCollapse(
@@ -46,51 +56,31 @@ function shouldCollapse(
   return isTextPlain && isSameChannel && isNotStdin;
 }
 
-function handleCarriageReturns(
-  consoleOutputs: OutputMessage[],
-): OutputMessage[] {
-  const newConsoleOutputs = [...consoleOutputs];
-  if (newConsoleOutputs.length === 0) {
-    return newConsoleOutputs;
+export function maybeMakeOutputStateful(output: OutputMessage): OutputMessage {
+  if (output instanceof StatefulOutputMessage) {
+    return output;
   }
-
-  const lastOutput = newConsoleOutputs[newConsoleOutputs.length - 1];
-  if (lastOutput.mimetype !== "text/plain") {
-    return newConsoleOutputs;
+  if (typeof output.data === "string") {
+    return StatefulOutputMessage.create(output as StringOutputMessage);
   }
+  return output;
+}
 
-  // eslint-disable-next-line no-control-regex
-  const carriagePattern = /\r[^\n]/g;
-  // collapse carriage returns in the final output's data
-  let text = lastOutput.data;
-  invariant(typeof text === "string", "expected string");
-  let carriageIdx = text.search(carriagePattern);
-  while (carriageIdx > -1) {
-    // find the newline character preceding the carriage return, if any
-    let newlineIdx = -1;
-    for (let i = carriageIdx - 1; i >= 0; i--) {
-      if (text.at(i) === "\n") {
-        newlineIdx = i;
-        break;
-      }
-    }
-    const postCarriageText: string = text.slice(carriageIdx + 1);
-    const prefix = text.slice(0, newlineIdx + 1);
-    const intermediateText: string = text.slice(newlineIdx + 1, carriageIdx);
-    text =
-      intermediateText.length <= postCarriageText.length
-        ? prefix + postCarriageText
-        : prefix +
-          postCarriageText +
-          intermediateText.slice(postCarriageText.length);
-    carriageIdx = text.search(carriagePattern);
+function mergeLeft(
+  first: StringOutputMessage | StatefulOutputMessage,
+  second: StringOutputMessage,
+): StatefulOutputMessage {
+  if (first instanceof StatefulOutputMessage) {
+    return first.appendData(second.data);
   }
 
-  newConsoleOutputs[newConsoleOutputs.length - 1] = {
-    ...lastOutput,
-    data: text,
-  };
-  return newConsoleOutputs;
+  return StatefulOutputMessage.create(first).appendData(second.data);
+}
+
+function assertStringOutputMessage(
+  output: OutputMessage,
+): asserts output is StringOutputMessage {
+  invariant(typeof output.data === "string", "expected string output");
 }
 
 function truncateHead(consoleOutputs: OutputMessage[], limit: number) {

package/src/core/export/hooks.ts

@@ -17,9 +17,10 @@ export function useAutoExport() {
   const htmlEnabled = appConfig.auto_download.includes("html");
   const ipynbEnabled = appConfig.auto_download.includes("ipynb");
 
-  const markdownDisabled = !markdownEnabled || state !== WebSocketState.OPEN;
-  const htmlDisabled = !htmlEnabled || state !== WebSocketState.OPEN;
-  const ipynbDisabled = !ipynbEnabled || state !== WebSocketState.OPEN;
+  const isConnected = state === WebSocketState.OPEN;
+  const markdownDisabled = !markdownEnabled || !isConnected;
+  const htmlDisabled = !htmlEnabled || !isConnected;
+  const ipynbDisabled = !ipynbEnabled || !isConnected;
   const { autoExportAsHTML, autoExportAsIPYNB, autoExportAsMarkdown } =
     useRequestClient();
 

package/src/core/static/__tests__/virtual-file-tracker.test.ts

@@ -56,6 +56,48 @@ describe("findVirtualFiles", () => {
     const result = findVirtualFiles(input);
     expect(result).toEqual(expected);
   });
+
+  it("should match virtual files with HTML entity endings", () => {
+    const input = '"/@file/938947-3587525-4XcHZISt.csv&quot"';
+    const expected = new Set(["/@file/938947-3587525-4XcHZISt.csv"]);
+    const result = findVirtualFiles(input);
+    expect(result).toEqual(expected);
+  });
+
+  it("should not match virtual files with mismatched quotes", () => {
+    const input = "'/@file/938947-3587525-4XcHZISt.csv\"'";
+    const expected = new Set(["/@file/938947-3587525-4XcHZISt.csv"]);
+    const result = findVirtualFiles(input);
+    expect(result).toEqual(expected);
+  });
+
+  it("should match virtual files followed by quotes", () => {
+    const input = '/@file/test-file.csv"';
+    const expected = new Set(["/@file/test-file.csv"]);
+    const result = findVirtualFiles(input);
+    expect(result).toEqual(expected);
+  });
+
+  it("should match valid virtual files even when surrounded by proper quotes", () => {
+    const input = '"/@file/test-file.csv"';
+    const expected = new Set(["/@file/test-file.csv"]);
+    const result = findVirtualFiles(input);
+    expect(result).toEqual(expected);
+  });
+
+  it("should extremely long greedy matching", () => {
+    const input = `/@file/test-file.csv","format":{"type":"csv"}},"mark":{"type":"circle","size":4},"encoding":{"color":{"field":"digit","type":"nominal"},"x":{"field":"x","scale":{"domain":[-2.5,2.5]},"type":"quantitative"},"y":{"field":"y","scale":{"domain":[-2.5,2.5. Error: [Errno 63] File name too long: '/3587525-4XcHZISt.csv","format":{"type":"csv"}},"mark":{"type":"circle","size":4},"encoding":{"color":{"field":"digit","type":"nominal"},"x":{"field":"x","scale":{"domain":[-2.5,2.5]},"type":"quantitative"},"y":{"field":"y","scale":{"domain":[-2.5,2.5'`;
+    const expected = new Set(["/@file/test-file.csv"]);
+    const result = findVirtualFiles(input);
+    expect(result).toEqual(expected);
+  });
+
+  it("should not match virtual files with quotes in the filename", () => {
+    const input = '/@file/file".csv';
+    const expected = new Set();
+    const result = findVirtualFiles(input);
+    expect(result).toEqual(expected);
+  });
 });
 
 describe("VirtualFileTracker", () => {

package/src/core/static/virtual-file-tracker.ts

@@ -4,7 +4,7 @@ import type { CellId } from "../cells/ids";
 import type { CellMessage } from "../kernel/messages";
 
 // Virtual files are of the form /@file/<file-name>.<extension>
-const VIRTUAL_FILE_REGEX = /\/@file\/([^\s/]+)\.([\dA-Za-z]+)/g;
+const VIRTUAL_FILE_REGEX = /\/@file\/([^\s"&'/]+)\.([\dA-Za-z]+)/g;
 
 /**
  * Tracks virtual files that are present on the page.

package/src/plugins/impl/DataTablePlugin.tsx

@@ -690,6 +690,8 @@ export const LoadingDataTableComponent = memo(
           <TablePanel
             displayHeader={displayHeader}
             data={data?.rows || []}
+            columns={props.totalColumns}
+            totalRows={props.totalRows}
             dataTable={dataTable}
             getDataUrl={props.get_data_url}
             fieldTypes={props.fieldTypes}