@marimo-team/islands 0.23.17-dev33 → 0.23.17-dev35

package/src/core/ai/staged-cells.ts

@@ -176,14 +176,17 @@ export function useStagedCells(store: JotaiStore) {
       case "message-metadata":
       case "tool-input-available":
       case "tool-output-available":
+      case "tool-approval-response":
       case "reasoning-start":
       case "reasoning-delta":
       case "reasoning-end":
       case "file":
+      case "reasoning-file":
       case "source-document":
       case "source-url":
       case "tool-input-start":
       case "tool-input-delta":
+      case "custom":
         Logger.debug(chunk.type, { chunk });
         break;
       default:

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

@@ -26,6 +26,11 @@ import {
 } from "lucide-react";
 import React, { useEffect, useRef, useState } from "react";
 import { z } from "zod";
+import {
+  describeChatAbortReason,
+  getAbortReasonFromChunk,
+  resolveChatAbortReason,
+} from "@/components/chat/chat-abort";
 import { renderUIMessage } from "@/components/chat/chat-display";
 import {
   convertToFileUIPart,
@@ -168,6 +173,12 @@ export const Chatbot: React.FC<Props> = (props) => {
   // 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 lastAbortReasonRef = useRef<string | null>(null);
+  const [abortNotice, setAbortNotice] = useState<string | null>(null);
+  const clearAbortState = () => {
+    lastAbortReasonRef.current = null;
+    setAbortNotice(null);
+  };
 
   const { data: backendMessages } = useAsyncData(async () => {
     const response = await props.get_chat_history({});
@@ -231,6 +242,7 @@ 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();
+          clearAbortState();
 
           const stream = new ReadableStream<UIMessageChunk>({
             start(controller) {
@@ -297,15 +309,30 @@ export const Chatbot: React.FC<Props> = (props) => {
       },
     }),
     messages: initialMessages,
-    onFinish: (message) => {
+    onFinish: ({ messages, isAbort, finishReason }) => {
       setFiles(undefined);
 
       if (fileInputRef.current) {
         fileInputRef.current.value = "";
       }
-      Logger.debug("Finished streaming message:", message);
+      Logger.debug("Finished streaming message:", { finishReason, isAbort });
+
+      if (isAbort) {
+        const reason = resolveChatAbortReason({
+          isAbort: true,
+          streamReason: lastAbortReasonRef.current,
+        });
+        clearAbortState();
+        if (reason != null) {
+          const description = describeChatAbortReason(reason);
+          Logger.debug("Chat stream aborted", { reason, finishReason });
+          setAbortNotice(description);
+        }
+      } else {
+        clearAbortState();
+      }
 
-      props.setValue(message.messages);
+      props.setValue(messages);
     },
     onError: (error) => {
       Logger.error("An error occurred:", error);
@@ -323,6 +350,17 @@ export const Chatbot: React.FC<Props> = (props) => {
       if (!parsedMessage.success) {
         return;
       }
+      // Only record abort reasons for the active request. Late chunks from a
+      // cancelled run must not overwrite the next turn's reason
+      if (
+        parsedMessage.data.content &&
+        parsedMessage.data.message_id === activeRequestIdRef.current
+      ) {
+        const abortReason = getAbortReasonFromChunk(parsedMessage.data.content);
+        if (abortReason !== undefined) {
+          lastAbortReasonRef.current = abortReason;
+        }
+      }
       routeIncomingChatChunk(parsedMessage.data, {
         controllerRef: frontendStreamControllerRef,
         activeRequestIdRef,
@@ -340,6 +378,9 @@ export const Chatbot: React.FC<Props> = (props) => {
       setMessages(newMessages);
 
       props.setValue(newMessages);
+      if (newMessages.length === 0) {
+        clearAbortState();
+      }
     }
   };
 
@@ -389,6 +430,7 @@ export const Chatbot: React.FC<Props> = (props) => {
     props.setValue([]);
     props.delete_chat_history({});
     clearError();
+    clearAbortState();
   };
 
   return (
@@ -483,6 +525,12 @@ export const Chatbot: React.FC<Props> = (props) => {
             </Button>
           </div>
         )}
+
+        {abortNotice && !isLoading && !error && (
+          <div className="flex justify-center py-2">
+            <span className="text-xs text-muted-foreground">{abortNotice}</span>
+          </div>
+        )}
       </div>
       <form
         onSubmit={async (evt) => {