@marimo-team/frontend 0.23.9-dev8 → 0.23.9

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