@marimo-team/frontend 0.23.9-dev3 → 0.23.9-dev32

package/src/plugins/impl/chat/__tests__/chat-ui.test.ts

@@ -0,0 +1,278 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { UIMessageChunk } from "ai";
+import { describe, expect, it, vi } from "vitest";
+import { routeIncomingChatChunk } from "../chat-ui";
+
+/**
+ * The stale-chunk filter prevents chunks from an aborted run from being enqueued into a new run's stream.
+ *
+ * It triggers when:
+ *   1. User sends a prompt (request_id = OLD), kernel starts emitting chunks
+ *   2. User clicks Stop — frontend tears down its controller, fires cancel_prompt
+ *   3. Kernel hasn't received the cancel yet and is still emitting chunks
+ *   4. User sends a new prompt (request_id = NEW), new controller opens
+ *   5. Late chunks tagged OLD arrive after NEW's controller is in place
+ */
+
+const makeChunk = (opts: {
+  messageId: string;
+  content: unknown;
+  isFinal?: boolean;
+}): Parameters<typeof routeIncomingChatChunk>[0] => ({
+  type: "stream_chunk",
+  message_id: opts.messageId,
+  content: opts.content as UIMessageChunk | null,
+  is_final: opts.isFinal ?? false,
+});
+
+const makeRefs = () => ({
+  controllerRef: {
+    current: null as ReadableStreamDefaultController<UIMessageChunk> | null,
+  },
+  activeRequestIdRef: { current: null as string | null },
+});
+
+const makeMockController = () => {
+  return {
+    enqueue: vi.fn(),
+    close: vi.fn(),
+    error: vi.fn(),
+    desiredSize: 0,
+  } as unknown as ReadableStreamDefaultController<UIMessageChunk> & {
+    enqueue: ReturnType<typeof vi.fn>;
+    close: ReturnType<typeof vi.fn>;
+  };
+};
+
+describe("routeIncomingChatChunk", () => {
+  it("drops chunks when there is no active controller", () => {
+    const refs = makeRefs();
+
+    const result = routeIncomingChatChunk(
+      makeChunk({
+        messageId: "req-A",
+        content: { type: "text-delta", id: "t1", delta: "hi" },
+      }),
+      refs,
+    );
+
+    expect(result).toBe("dropped-no-controller");
+  });
+
+  it("enqueues chunks that match the active request_id", () => {
+    const refs = makeRefs();
+    const controller = makeMockController();
+    refs.controllerRef.current = controller;
+    refs.activeRequestIdRef.current = "req-A";
+
+    const chunk = { type: "text-delta", id: "t1", delta: "hi" } as const;
+    const result = routeIncomingChatChunk(
+      makeChunk({ messageId: "req-A", content: chunk }),
+      refs,
+    );
+
+    expect(result).toBe("enqueued");
+    expect(controller.enqueue).toHaveBeenCalledWith(chunk);
+    expect(controller.close).not.toHaveBeenCalled();
+  });
+
+  it("closes the controller and clears refs on is_final", () => {
+    const refs = makeRefs();
+    const controller = makeMockController();
+    refs.controllerRef.current = controller;
+    refs.activeRequestIdRef.current = "req-A";
+
+    const result = routeIncomingChatChunk(
+      makeChunk({ messageId: "req-A", content: null, isFinal: true }),
+      refs,
+    );
+
+    expect(result).toBe("closed");
+    expect(controller.close).toHaveBeenCalledTimes(1);
+    expect(refs.controllerRef.current).toBeNull();
+    expect(refs.activeRequestIdRef.current).toBeNull();
+  });
+
+  it("drops chunks whose message_id does not match the active run", () => {
+    // Simulates the bug: kernel hasn't received cancel for OLD yet but the
+    // user has already started a NEW run. A reasoning-delta for OLD arrives
+    // here; it must not be enqueued into NEW's stream.
+    const refs = makeRefs();
+    const controller = makeMockController();
+    refs.controllerRef.current = controller;
+    refs.activeRequestIdRef.current = "req-NEW";
+
+    const staleChunk = {
+      type: "reasoning-delta",
+      id: "r-old",
+      delta: "...",
+    } as const;
+    const result = routeIncomingChatChunk(
+      makeChunk({ messageId: "req-OLD", content: staleChunk }),
+      refs,
+    );
+
+    expect(result).toBe("dropped-stale");
+    expect(controller.enqueue).not.toHaveBeenCalled();
+    expect(controller.close).not.toHaveBeenCalled();
+    expect(refs.activeRequestIdRef.current).toBe("req-NEW");
+  });
+
+  it("drops is_final from a stale run without closing the active stream", () => {
+    // Belt-and-suspenders: an `is_final` for OLD that races in after NEW
+    // started must not tear down NEW's controller.
+    const refs = makeRefs();
+    const controller = makeMockController();
+    refs.controllerRef.current = controller;
+    refs.activeRequestIdRef.current = "req-NEW";
+
+    const result = routeIncomingChatChunk(
+      makeChunk({ messageId: "req-OLD", content: null, isFinal: true }),
+      refs,
+    );
+
+    expect(result).toBe("dropped-stale");
+    expect(controller.close).not.toHaveBeenCalled();
+    expect(refs.controllerRef.current).toBe(controller);
+    expect(refs.activeRequestIdRef.current).toBe("req-NEW");
+  });
+
+  it("forwards reasoning-start/delta/end sequences when ids match", () => {
+    // Walks the canonical happy path for a reasoning stream end-to-end.
+    const refs = makeRefs();
+    const controller = makeMockController();
+    refs.controllerRef.current = controller;
+    refs.activeRequestIdRef.current = "req-A";
+
+    const sequence = [
+      { type: "reasoning-start", id: "r1" },
+      { type: "reasoning-delta", id: "r1", delta: "thinking" },
+      { type: "reasoning-end", id: "r1" },
+    ] as const;
+    for (const chunk of sequence) {
+      const result = routeIncomingChatChunk(
+        makeChunk({ messageId: "req-A", content: chunk }),
+        refs,
+      );
+      expect(result).toBe("enqueued");
+    }
+
+    expect(controller.enqueue).toHaveBeenCalledTimes(3);
+    expect(controller.enqueue).toHaveBeenNthCalledWith(1, sequence[0]);
+    expect(controller.enqueue).toHaveBeenNthCalledWith(2, sequence[1]);
+    expect(controller.enqueue).toHaveBeenNthCalledWith(3, sequence[2]);
+  });
+
+  it(
+    "drops stale reasoning-delta after Stop → new run sequence " +
+      "(regression for missing reasoning part error)",
+    () => {
+      // Full scenario: A runs, A is stopped, B starts, A's late chunk arrives.
+      const refs = makeRefs();
+
+      // 1. Run A starts: controller A active.
+      const controllerA = makeMockController();
+      refs.controllerRef.current = controllerA;
+      refs.activeRequestIdRef.current = "req-A";
+
+      // First reasoning chunks for A flow through.
+      routeIncomingChatChunk(
+        makeChunk({
+          messageId: "req-A",
+          content: { type: "reasoning-start", id: "rA" },
+        }),
+        refs,
+      );
+      routeIncomingChatChunk(
+        makeChunk({
+          messageId: "req-A",
+          content: {
+            type: "reasoning-delta",
+            id: "rA",
+            delta: "thinking",
+          },
+        }),
+        refs,
+      );
+      expect(controllerA.enqueue).toHaveBeenCalledTimes(2);
+
+      // 2. User clicks Stop: abort handler clears refs (simulated).
+      refs.controllerRef.current = null;
+      refs.activeRequestIdRef.current = null;
+
+      // A late chunk for A arrives in this window — must be a no-op.
+      const between = routeIncomingChatChunk(
+        makeChunk({
+          messageId: "req-A",
+          content: {
+            type: "reasoning-delta",
+            id: "rA",
+            delta: "leftover",
+          },
+        }),
+        refs,
+      );
+      expect(between).toBe("dropped-no-controller");
+
+      // 3. User sends Run B: new controller, new active id.
+      const controllerB = makeMockController();
+      refs.controllerRef.current = controllerB;
+      refs.activeRequestIdRef.current = "req-B";
+
+      // 4. Another late chunk for A arrives AFTER B opened. This is the
+      // case that previously threw `Received reasoning-delta for missing
+      // reasoning part with ID "rA"` in the SDK parser.
+      const stale = routeIncomingChatChunk(
+        makeChunk({
+          messageId: "req-A",
+          content: {
+            type: "reasoning-delta",
+            id: "rA",
+            delta: "still leaking",
+          },
+        }),
+        refs,
+      );
+      expect(stale).toBe("dropped-stale");
+      expect(controllerB.enqueue).not.toHaveBeenCalled();
+
+      // 5. B's own chunks flow normally.
+      routeIncomingChatChunk(
+        makeChunk({
+          messageId: "req-B",
+          content: { type: "reasoning-start", id: "rB" },
+        }),
+        refs,
+      );
+      routeIncomingChatChunk(
+        makeChunk({
+          messageId: "req-B",
+          content: { type: "reasoning-delta", id: "rB", delta: "fresh" },
+        }),
+        refs,
+      );
+      expect(controllerB.enqueue).toHaveBeenCalledTimes(2);
+    },
+  );
+
+  it("enqueues content alongside is_final and then closes", () => {
+    // Sanity: a single chunk that carries both `content` and `is_final` (rare
+    // but legal — backend may bundle final content with the terminator)
+    // should enqueue then close.
+    const refs = makeRefs();
+    const controller = makeMockController();
+    refs.controllerRef.current = controller;
+    refs.activeRequestIdRef.current = "req-A";
+
+    const chunk = { type: "text-delta", id: "t1", delta: "bye" } as const;
+    const result = routeIncomingChatChunk(
+      makeChunk({ messageId: "req-A", content: chunk, isFinal: true }),
+      refs,
+    );
+
+    expect(result).toBe("closed");
+    expect(controller.enqueue).toHaveBeenCalledWith(chunk);
+    expect(controller.close).toHaveBeenCalledTimes(1);
+  });
+});

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}
       />
     );