@marimo-team/frontend 0.23.9-dev4 → 0.23.9-dev40

package/src/plugins/impl/chat/chat-ui.tsx

@@ -20,13 +20,17 @@ import {
   RotateCwIcon,
   SendHorizontalIcon,
   SettingsIcon,
+  SquareIcon,
   Trash2Icon,
   X,
 } from "lucide-react";
 import React, { useEffect, useRef, useState } from "react";
 import { z } from "zod";
 import { renderUIMessage } from "@/components/chat/chat-display";
-import { convertToFileUIPart } from "@/components/chat/chat-utils";
+import {
+  convertToFileUIPart,
+  hasPendingToolCalls,
+} from "@/components/chat/chat-utils";
 import {
   type AdditionalCompletions,
   PromptInput,
@@ -60,6 +64,7 @@ import { cn } from "@/utils/cn";
 import { Logger } from "@/utils/Logger";
 import { Objects } from "@/utils/objects";
 import { Strings } from "@/utils/strings";
+import { generateUUID } from "@/utils/uuid";
 import { ErrorBanner } from "../common/error-banner";
 import type { PluginFunctions } from "./ChatPlugin";
 import type { ChatConfig } from "./types";
@@ -86,6 +91,48 @@ const ChatMessageIncomingSchema = z.object({
   is_final: z.boolean().optional(),
 });
 
+type ChatMessageIncoming = z.infer<typeof ChatMessageIncomingSchema>;
+
+export interface IncomingChatChunkRefs {
+  controllerRef: {
+    current: ReadableStreamDefaultController<UIMessageChunk> | null;
+  };
+  activeRequestIdRef: { current: string | null };
+}
+
+/**
+ * Route a single incoming chunk to the active stream controller, dropping it
+ * if it belongs to a stale (aborted-but-not-yet-cancelled) backend run.
+ */
+export function routeIncomingChatChunk(
+  message: ChatMessageIncoming,
+  refs: IncomingChatChunkRefs,
+): "enqueued" | "closed" | "dropped-no-controller" | "dropped-stale" {
+  const { controllerRef, activeRequestIdRef } = refs;
+  const controller = controllerRef.current;
+  if (controller === null) {
+    return "dropped-no-controller";
+  }
+  const activeRequestId = activeRequestIdRef.current;
+  if (activeRequestId !== null && message.message_id !== activeRequestId) {
+    Logger.debug("Dropping stale chat chunk", {
+      chunkRequestId: message.message_id,
+      activeRequestId,
+    });
+    return "dropped-stale";
+  }
+  if (message.content) {
+    controller.enqueue(message.content);
+  }
+  if (message.is_final) {
+    controller.close();
+    controllerRef.current = null;
+    activeRequestIdRef.current = null;
+    return "closed";
+  }
+  return "enqueued";
+}
+
 export const Chatbot: React.FC<Props> = (props) => {
   const [input, setInput] = useState("");
   const [config, setConfig] = useState<ChatConfig>(props.config);
@@ -113,16 +160,15 @@ export const Chatbot: React.FC<Props> = (props) => {
   const configRef = useRef<ChatConfig>(config);
   configRef.current = config;
 
-  // Track streaming state - maps backend message_id to frontend message index
-  const streamingStateRef = useRef<{
-    backendMessageId: string | null;
-    frontendMessageIndex: number | null;
-  }>({ backendMessageId: null, frontendMessageIndex: null });
-
   // For frontend-managed streaming, create a controller to enqueue chunks to.
   const frontendStreamControllerRef =
     useRef<ReadableStreamDefaultController<UIMessageChunk> | null>(null);
 
+  // The request_id of the currently-active prompt run. Chunks arriving with a
+  // different message_id are stale (from an aborted-but-not-yet-cancelled run
+  // on the kernel) and must be dropped
+  const activeRequestIdRef = useRef<string | null>(null);
+
   const { data: backendMessages } = useAsyncData(async () => {
     const response = await props.get_chat_history({});
     return response.messages;
@@ -143,7 +189,9 @@ export const Chatbot: React.FC<Props> = (props) => {
     error,
     regenerate,
     clearError,
+    addToolApprovalResponse,
   } = useChat({
+    sendAutomaticallyWhen: ({ messages }) => hasPendingToolCalls(messages),
     transport: new DefaultChatTransport({
       fetch: async (
         request: RequestInfo | URL,
@@ -180,17 +228,33 @@ export const Chatbot: React.FC<Props> = (props) => {
             };
           });
 
+          // Client-generated id used to (a) route chunks back to this stream
+          // and (b) ask the kernel to cancel just this run on Stop.
+          const requestId = generateUUID();
+
           const stream = new ReadableStream<UIMessageChunk>({
             start(controller) {
               frontendStreamControllerRef.current = controller;
+              activeRequestIdRef.current = requestId;
 
               const abortHandler = () => {
+                // Close the local controller first so the chat status flips to
+                // "ready" immediately and any racing chunks are dropped; then
+                // fire-and-forget the backend cancel so the kernel stops the
+                // model and we don't waste tokens / leak chunks to the next
+                // run.
                 try {
                   controller.close();
                 } catch (error) {
                   Logger.debug("Controller may already be closed", { error });
                 }
                 frontendStreamControllerRef.current = null;
+                activeRequestIdRef.current = null;
+                void props
+                  .cancel_prompt({ request_id: requestId })
+                  .catch((error: Error) => {
+                    Logger.debug("cancel_prompt failed", { error });
+                  });
               };
               signal?.addEventListener("abort", abortHandler);
 
@@ -200,28 +264,25 @@ export const Chatbot: React.FC<Props> = (props) => {
             },
             cancel() {
               frontendStreamControllerRef.current = null;
+              activeRequestIdRef.current = null;
             },
           });
 
           // Start the prompt, chunks will be sent via events
           void props
             .send_prompt({
+              request_id: requestId,
               messages: messages,
               config: chatConfig,
             })
             .catch((error: Error) => {
               frontendStreamControllerRef.current?.error(error);
               frontendStreamControllerRef.current = null;
+              activeRequestIdRef.current = null;
             });
 
           return createUIMessageStreamResponse({ stream });
         } catch (error: unknown) {
-          // Clear streaming state on error
-          streamingStateRef.current = {
-            backendMessageId: null,
-            frontendMessageIndex: null,
-          };
-
           // Handle abort gracefully without showing an error
           if (error instanceof Error && error.name === "AbortError") {
             return new Response("Aborted", { status: 499 });
@@ -244,21 +305,10 @@ export const Chatbot: React.FC<Props> = (props) => {
       }
       Logger.debug("Finished streaming message:", message);
 
-      // Clear streaming state
-      streamingStateRef.current = {
-        backendMessageId: null,
-        frontendMessageIndex: null,
-      };
-
       props.setValue(message.messages);
     },
     onError: (error) => {
       Logger.error("An error occurred:", error);
-      // Clear streaming state on error
-      streamingStateRef.current = {
-        backendMessageId: null,
-        frontendMessageIndex: null,
-      };
     },
   });
 
@@ -273,23 +323,10 @@ export const Chatbot: React.FC<Props> = (props) => {
       if (!parsedMessage.success) {
         return;
       }
-      const message = parsedMessage.data;
-
-      // Push to the stream for useChat to process
-      const controller = frontendStreamControllerRef.current;
-      if (!controller) {
-        return;
-      }
-
-      if (message.content) {
-        controller.enqueue(message.content);
-      }
-      if (message.is_final) {
-        controller.close();
-        frontendStreamControllerRef.current = null;
-      }
-
-      return;
+      routeIncomingChatChunk(parsedMessage.data, {
+        controllerRef: frontendStreamControllerRef,
+        activeRequestIdRef,
+      });
     },
   );
 
@@ -408,6 +445,9 @@ export const Chatbot: React.FC<Props> = (props) => {
                   message,
                   isStreamingReasoning: status === "streaming",
                   isLast,
+                  addToolApprovalResponse: isLast
+                    ? addToolApprovalResponse
+                    : undefined,
                 })}
               </div>
               <div className="flex justify-end text-xs gap-2 invisible group-hover:visible">
@@ -429,16 +469,8 @@ export const Chatbot: React.FC<Props> = (props) => {
         })}
 
         {isLoading && (
-          <div className="flex items-center justify-center space-x-2 mb-4">
+          <div className="flex items-center justify-center mb-4">
             <Spinner size="small" />
-            <Button
-              variant="link"
-              size="sm"
-              onClick={() => stop()}
-              className="text-(--red-9) hover:text-(--red-11)"
-            >
-              Stop
-            </Button>
           </div>
         )}
 
@@ -569,15 +601,30 @@ export const Chatbot: React.FC<Props> = (props) => {
             />
           </>
         )}
-        <Button
-          type="submit"
-          disabled={isLoading || !input}
-          variant="outline"
-          size="xs"
-          className="text-(--slate-11)"
-        >
-          <SendHorizontalIcon className="h-4 w-4" />
-        </Button>
+        {isLoading ? (
+          <Tooltip content="Stop generating">
+            <Button
+              type="button"
+              variant="link"
+              size="xs"
+              onClick={() => stop()}
+              className="text-(--red-9) hover:text-(--red-11)"
+            >
+              <SquareIcon className="h-4 w-4 fill-current" />
+            </Button>
+          </Tooltip>
+        ) : (
+          <Button
+            type="submit"
+            disabled={!input}
+            variant="outline"
+            size="xs"
+            className="text-(--slate-11)"
+            aria-label="Send message"
+          >
+            <SendHorizontalIcon className="h-4 w-4" />
+          </Button>
+        )}
       </form>
     </div>
   );

package/src/plugins/impl/chat/types.ts

@@ -22,6 +22,11 @@ export interface ChatConfig {
 }
 
 export interface SendMessageRequest {
+  request_id: string;
   messages: ChatMessage[];
   config: ChatConfig;
 }
+
+export interface CancelPromptRequest {
+  request_id: string;
+}

package/src/plugins/impl/data-frames/DataFramePlugin.tsx

@@ -64,7 +64,6 @@ type PluginFunctions = {
     column_types_per_step: FieldTypesWithExternalType[];
     python_code?: string | null;
     sql_code?: string | null;
-    size_bytes?: number | null;
   }>;
   get_column_values: (req: { column: string }) => Promise<{
     values: unknown[];
@@ -82,9 +81,11 @@ type PluginFunctions = {
   }) => Promise<{
     data: TableData<T>;
     total_rows: number;
-    size_bytes?: number | null;
   }>;
   download_as: DownloadAsArgs;
+  get_size_bytes: (opts: Record<string, never>) => Promise<{
+    size_bytes?: number | null;
+  }>;
 };
 
 // Value is selection, but it is not currently exposed to the user
@@ -120,7 +121,6 @@ export const DataFramePlugin = createPlugin<S>("marimo-dataframe")
         column_types_per_step: z.array(columnToFieldTypesSchema),
         python_code: z.string().nullish(),
         sql_code: z.string().nullish(),
-        size_bytes: z.number().nullish(),
       }),
     ),
     get_column_values: rpc.input(z.object({ column: z.string() })).output(
@@ -150,10 +150,12 @@ export const DataFramePlugin = createPlugin<S>("marimo-dataframe")
         z.object({
           data: z.union([z.string(), z.array(z.object({}).passthrough())]),
           total_rows: z.number(),
-          size_bytes: z.number().nullish(),
         }),
       ),
     download_as: DownloadAsSchema,
+    get_size_bytes: rpc
+      .input(z.object({}))
+      .output(z.object({ size_bytes: z.number().nullish() })),
   })
   .renderer((props) => (
     <TableProviders>
@@ -192,6 +194,7 @@ export const DataFrameComponent = memo(
     get_column_values,
     search,
     download_as,
+    get_size_bytes,
     host,
   }: DataTableProps): JSX.Element => {
     const { data, error, isPending } = useAsyncData(
@@ -207,7 +210,6 @@ export const DataFrameComponent = memo(
       column_types_per_step,
       python_code,
       sql_code,
-      size_bytes,
     } = data || {};
 
     const totalColumns = field_types?.length;
@@ -327,7 +329,6 @@ export const DataFrameComponent = memo(
           data={url || ""}
           hasStableRowId={false}
           totalRows={total_rows ?? 0}
-          sizeBytes={size_bytes ?? null}
           totalColumns={totalColumns ?? 0}
           maxColumns="all"
           pageSize={pageSize}
@@ -336,6 +337,7 @@ export const DataFrameComponent = memo(
           rowHeaders={row_headers || Arrays.EMPTY}
           showDownload={showDownload}
           download_as={download_as}
+          get_size_bytes={get_size_bytes}
           enableSearch={false}
           showFilters={false}
           search={search}

package/src/stories/dataframe.stories.tsx

@@ -42,6 +42,7 @@ export const DataFrame: StoryObj = {
         host={document.body}
         showDownload={false}
         download_as={async () => ({ url: "", filename: "" })}
+        get_size_bytes={async () => ({ size_bytes: null })}
         lazy={false}
       />
     );

package/src/utils/__tests__/json-parser.test.ts

@@ -1,10 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { expect, it } from "vitest";
-import {
-  jsonParseWithSpecialChar,
-  jsonToMarkdown,
-  jsonToTSV,
-} from "../json/json-parser";
+import { jsonParseWithSpecialChar, jsonToMarkdown } from "../json/json-parser";
 
 it("can jsonParseWithSpecialChar happy path", () => {
   expect(jsonParseWithSpecialChar('"hello"')).toEqual("hello");
@@ -72,70 +68,6 @@ it("can parse bigInts", () => {
   });
 });
 
-it("can convert json to tsv with en-US locale", () => {
-  const locale = "en-US";
-
-  expect(jsonToTSV([], locale)).toEqual("");
-
-  expect(jsonToTSV([{ a: 1, b: 2 }], locale)).toEqual("a\tb\n1\t2");
-
-  expect(
-    jsonToTSV(
-      [
-        { a: 1, b: 2 },
-        { a: 3, b: 4 },
-      ],
-      locale,
-    ),
-  ).toEqual("a\tb\n1\t2\n3\t4");
-
-  // Does not handle sparse arrays
-  expect(jsonToTSV([{ a: 1 }, { a: 2, b: 3 }], locale)).toMatchInlineSnapshot(
-    '"a\n1\n2"',
-  );
-
-  // Handles special characters
-  expect(
-    jsonToTSV([{ a: "hello\tworld", b: "new\nline" }], locale),
-  ).toMatchInlineSnapshot('"a\tb\nhello\tworld\tnew\nline"');
-
-  // Handles floats with en-US locale (uses . as decimal separator)
-  expect(jsonToTSV([{ a: 1.5, b: 2.7 }], locale)).toEqual("a\tb\n1.5\t2.7");
-});
-
-it("can convert json to tsv with de-DE locale", () => {
-  const locale = "de-DE";
-
-  // Handles floats with de-DE locale (uses , as decimal separator)
-  expect(jsonToTSV([{ a: 1.5, b: 2.7 }], locale)).toEqual("a\tb\n1,5\t2,7");
-
-  // Handles integers (no change)
-  expect(jsonToTSV([{ a: 1, b: 2 }], locale)).toEqual("a\tb\n1\t2");
-});
-
-it("can convert json to tsv with fr-FR locale", () => {
-  const locale = "fr-FR";
-
-  // Handles floats with fr-FR locale (uses , as decimal separator)
-  expect(jsonToTSV([{ a: 3.14, b: 2.123_45 }], locale)).toEqual(
-    "a\tb\n3,14\t2,12345",
-  );
-});
-
-it("handles null and undefined values in TSV", () => {
-  const locale = "en-US";
-
-  expect(jsonToTSV([{ a: null, b: undefined, c: 1 }], locale)).toEqual(
-    "a\tb\tc\n\t\t1",
-  );
-});
-
-it("handles NaN values in TSV", () => {
-  const locale = "en-US";
-
-  expect(jsonToTSV([{ a: Number.NaN, b: 1 }], locale)).toEqual("a\tb\nNaN\t1");
-});
-
 it("can convert json to markdown - basic table", () => {
   expect(jsonToMarkdown([])).toMatchInlineSnapshot(`""`);
 

package/src/utils/__tests__/local-variables.test.ts

@@ -0,0 +1,132 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { describe, expect, it } from "vitest";
+import {
+  containsMangledLocal,
+  splitMangledLocals,
+  unmangleLocal,
+} from "../local-variables";
+
+describe("unmangleLocal", () => {
+  it("extracts the cell id and original name", () => {
+    expect(unmangleLocal("_cell_Hbol_a")).toEqual({
+      cellId: "Hbol",
+      name: "_a",
+    });
+  });
+
+  it("handles names with multiple underscore segments", () => {
+    expect(unmangleLocal("_cell_Hbol_a_b")).toEqual({
+      cellId: "Hbol",
+      name: "_a_b",
+    });
+  });
+
+  it("returns null for non-mangled strings", () => {
+    expect(unmangleLocal("just_a_variable")).toBeNull();
+    expect(unmangleLocal("_private")).toBeNull();
+  });
+
+  it("handles the single-underscore local", () => {
+    // `_` is a valid local name (variables.py:62-63), mangling to
+    // `_cell_<id>_` with no trailing suffix.
+    expect(unmangleLocal("_cell_Hbol_")).toEqual({
+      cellId: "Hbol",
+      name: "_",
+    });
+  });
+
+  it("handles UUID-style cell ids", () => {
+    // External / VSCode notebooks use `external_prefix()` which is a
+    // `uuid4()` (hyphenated).
+    expect(
+      unmangleLocal("_cell_c9bf9e57-1685-4c89-bafb-ff5af830be8a_a"),
+    ).toEqual({
+      cellId: "c9bf9e57-1685-4c89-bafb-ff5af830be8a",
+      name: "_a",
+    });
+  });
+
+  it("does not match marimo cell file paths", () => {
+    // The compiled cell file is `__marimo__cell_<id>_.py` (two leading
+    // underscores, trailing `_` with no name); it must not be confused with a
+    // mangled local.
+    expect(unmangleLocal("__marimo__cell_Hbol_.py")).toBeNull();
+    expect(unmangleLocal("/tmp/marimo_42/__marimo__cell_Hbol_.py")).toBeNull();
+  });
+});
+
+describe("splitMangledLocals", () => {
+  it("returns a single text segment when nothing matches", () => {
+    expect(splitMangledLocals("plain text")).toEqual(["plain text"]);
+  });
+
+  it("splits a NameError message", () => {
+    expect(splitMangledLocals("name '_cell_Hbol_a' is not defined")).toEqual([
+      "name '",
+      { cellId: "Hbol", name: "_a" },
+      "' is not defined",
+    ]);
+  });
+
+  it("handles multiple mangled names in one string", () => {
+    expect(splitMangledLocals("_cell_AAAA_x and _cell_BBBB_y")).toEqual([
+      { cellId: "AAAA", name: "_x" },
+      " and ",
+      { cellId: "BBBB", name: "_y" },
+    ]);
+  });
+
+  it("leaves the cell file path alone", () => {
+    const path = "/tmp/marimo_42/__marimo__cell_Hbol_.py";
+    expect(splitMangledLocals(path)).toEqual([path]);
+  });
+
+  it("ignores `_cell_...` substrings preceded by `_`", () => {
+    // Mirrors `(?<!_)` in `variables.py`: a leading `_` (e.g. inside
+    // `__marimo__cell_<id>_<...>`) means this is not a mangle the compiler
+    // produced, so we must not demangle it.
+    const text = "see __marimo__cell_Hbol_a for details";
+    expect(splitMangledLocals(text)).toEqual([text]);
+  });
+
+  it("splits a UUID-style cell id", () => {
+    expect(
+      splitMangledLocals(
+        "name '_cell_c9bf9e57-1685-4c89-bafb-ff5af830be8a_a' is not defined",
+      ),
+    ).toEqual([
+      "name '",
+      {
+        cellId: "c9bf9e57-1685-4c89-bafb-ff5af830be8a",
+        name: "_a",
+      },
+      "' is not defined",
+    ]);
+  });
+
+  it("splits the single-underscore local", () => {
+    expect(splitMangledLocals("name '_cell_Hbol_' is not defined")).toEqual([
+      "name '",
+      { cellId: "Hbol", name: "_" },
+      "' is not defined",
+    ]);
+  });
+});
+
+describe("containsMangledLocal", () => {
+  it("detects mangled names", () => {
+    expect(containsMangledLocal("name '_cell_Hbol_a' is not defined")).toBe(
+      true,
+    );
+  });
+
+  it("ignores the cell file path", () => {
+    expect(containsMangledLocal("/tmp/marimo_42/__marimo__cell_Hbol_.py")).toBe(
+      false,
+    );
+  });
+
+  it("ignores `_cell_...` substrings preceded by `_`", () => {
+    expect(containsMangledLocal("see __marimo__cell_Hbol_a")).toBe(false);
+  });
+});

package/src/utils/json/json-parser.ts

@@ -66,36 +66,6 @@ export function jsonParseWithSpecialChar<T = unknown>(
   }
 }
 
-/**
- * Formats a value for TSV export, respecting user's locale for numbers
- */
-function formatValueForTSV(value: unknown, locale: string): string {
-  if (value === null || value === undefined) {
-    return "";
-  }
-  if (typeof value === "number" && !Number.isNaN(value)) {
-    // Use toLocaleString to format numbers according to user's locale
-    // This will use the appropriate decimal separator (e.g., "," in European locales)
-    return value.toLocaleString(locale, {
-      useGrouping: false,
-      maximumFractionDigits: 20,
-    });
-  }
-  return String(value);
-}
-
-export function jsonToTSV(json: Record<string, unknown>[], locale: string) {
-  if (json.length === 0) {
-    return "";
-  }
-
-  const keys = Object.keys(json[0]);
-  const values = json.map((row) =>
-    keys.map((key) => formatValueForTSV(row[key], locale)).join("\t"),
-  );
-  return `${keys.join("\t")}\n${values.join("\n")}`;
-}
-
 /**
  * Converts JSON data to a Markdown table format
  * Detects URLs and converts them to markdown links [url](url)