@marimo-team/islands 0.23.9-dev2 → 0.23.9-dev20

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__/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/dates.ts

@@ -134,6 +134,45 @@ export function timeAgo(
   return value.toString();
 }
 
+function pad2(n: number): string {
+  return n.toString().padStart(2, "0");
+}
+
+function pad4(n: number): string {
+  return n.toString().padStart(4, "0");
+}
+
+/**
+ * Format a Date as `YYYY-MM-DD` using the date's local-time fields.
+ *
+ * The output reflects what the user sees in their own timezone (the calendar
+ * day on their clock), not the UTC day. Use this when round-tripping values
+ * that originated from local-time inputs — date pickers, "filter on this
+ * day", calendar UI — so the displayed and serialized days agree.
+ *
+ * Not suitable for cross-timezone storage; use `Date.toISOString()` for that.
+ */
+export function dateToLocalISODate(d: Date): string {
+  return `${pad4(d.getFullYear())}-${pad2(d.getMonth() + 1)}-${pad2(d.getDate())}`;
+}
+
+/**
+ * Format a Date as `HH:MM:SS` using the date's local-time fields.
+ *
+ * See `dateToLocalISODate` for the rationale on local vs UTC.
+ */
+export function dateToLocalISOTime(d: Date): string {
+  return `${pad2(d.getHours())}:${pad2(d.getMinutes())}:${pad2(d.getSeconds())}`;
+}
+
+/**
+ * Format a Date as `YYYY-MM-DDTHH:MM:SS` (no timezone suffix) using local
+ * fields. See `dateToLocalISODate` for the rationale on local vs UTC.
+ */
+export function dateToLocalISODateTime(d: Date): string {
+  return `${dateToLocalISODate(d)}T${dateToLocalISOTime(d)}`;
+}
+
 export const supportedDateFormats = ["yyyy", "yyyy-MM", "yyyy-MM-dd"] as const;
 export type DateFormat = (typeof supportedDateFormats)[number];
 

package/src/utils/local-variables.ts

@@ -0,0 +1,67 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import type { CellId } from "@/core/cells/ids";
+
+/**
+ * The marimo compiler rewrites underscore-prefixed references inside a cell
+ * (which Python would treat as module-private) into `_cell_<cell_id>_<name>`
+ * so each cell gets its own private namespace. When such a reference fails at
+ * runtime (`NameError: name '_cell_Hbol_a' is not defined`), the mangled name
+ * leaks into the error UI. These helpers undo that mangling for display.
+ *
+ * Mirrors `marimo/_ast/variables.py`.
+ */
+
+// Matches `_cell_<cell_id><name>` for normal ids and UUIDs. The `[\w-]`
+// id class admits hyphens; the `_\w*` name group admits the bare `_`
+// local; the `(?<!_)` lookbehind skips `__marimo__cell_...` paths.
+// Mirrors `_MANGLED_LOCAL_IN_TEXT_RE` in `variables.py`.
+const MANGLED_LOCAL_BODY = String.raw`_cell_([^\W_][\w-]*?)(_\w*)`;
+const MANGLED_LOCAL_PATTERN = String.raw`(?<!_)${MANGLED_LOCAL_BODY}`;
+// Strict (whole-string) form for `unmangleLocal`; the leading `^` makes the
+// lookbehind trivially satisfied, so use the bare body.
+const ANCHORED_RE = new RegExp(`^${MANGLED_LOCAL_BODY}$`);
+const UNANCHORED_RE = new RegExp(MANGLED_LOCAL_PATTERN);
+const GLOBAL_RE = new RegExp(MANGLED_LOCAL_PATTERN, "g");
+
+export interface UnmangledLocal {
+  cellId: CellId;
+  /** The original underscore-prefixed name, e.g. "_a". */
+  name: string;
+}
+
+export function unmangleLocal(mangled: string): UnmangledLocal | null {
+  const match = ANCHORED_RE.exec(mangled);
+  if (!match) {
+    return null;
+  }
+  return { cellId: match[1] as CellId, name: match[2] };
+}
+
+export type MangledSegment = string | UnmangledLocal;
+
+/**
+ * Split a plain text string into alternating literal text and unmangled-local
+ * segments, so callers can render mixed React content.
+ */
+export function splitMangledLocals(text: string): MangledSegment[] {
+  const segments: MangledSegment[] = [];
+  GLOBAL_RE.lastIndex = 0;
+  let lastIndex = 0;
+  let match: RegExpExecArray | null = GLOBAL_RE.exec(text);
+  while (match !== null) {
+    if (match.index > lastIndex) {
+      segments.push(text.slice(lastIndex, match.index));
+    }
+    segments.push({ cellId: match[1] as CellId, name: match[2] });
+    lastIndex = match.index + match[0].length;
+    match = GLOBAL_RE.exec(text);
+  }
+  if (lastIndex < text.length) {
+    segments.push(text.slice(lastIndex));
+  }
+  return segments;
+}
+
+export function containsMangledLocal(text: string): boolean {
+  return UNANCHORED_RE.test(text);
+}