@dbx-tools/ui-mastra 0.9.175 → 0.9.177

package/index.react.ts

@@ -33,7 +33,7 @@ export type { GenieFeedbackTarget } from "./src/react/chat-feedback.ts";
 export { useChatHistory } from "./src/react/chat-history.ts";
 export { useChatSessions } from "./src/react/chat-sessions.ts";
 export type { ThreadSessionReader, ThreadSessionUpdater, ThreadMessageWriter } from "./src/react/chat-sessions.ts";
-export { readMlflowTracePart, closeOnTerminalChunk, waitForTerminalChunk, useChatStream } from "./src/react/chat-stream.ts";
+export { readMlflowTracePart, chatStatusForChunk, closeOnTerminalChunk, waitForTerminalChunk, useChatStream } from "./src/react/chat-stream.ts";
 export { ChatThreadLayout } from "./src/react/chat-thread-layout.tsx";
 export { useChatTranscriptController, ChatTranscript } from "./src/react/chat-transcript.tsx";
 export type { ChatTranscriptController } from "./src/react/chat-transcript.tsx";
@@ -57,5 +57,6 @@ export { ThreadActions, ThreadTabs } from "./src/react/thread-tabs.tsx";
 export type { ThreadTabsProps, ThreadActionsProps } from "./src/react/thread-tabs.tsx";
 export { humanizeToolName, formatRawToolPayload, structuredToolPayload, detectToolSourceLanguage, toolInputPresentation, toolOutputPresentation, webSearchProgressGroups, webSearchQueryLabel, webSearchResultLabel, ToolSessionPill } from "./src/react/tool-pill.tsx";
 export type { ToolInputPresentation, WebSearchProgressGroup } from "./src/react/tool-pill.tsx";
+export { isActiveChatStatus, isRespondingChatStatus } from "./src/react/types.ts";
 export type { ChatStatus, ToolEvent, ToolProgress, ChatModelOption, ChatInputMessage, QueuedSteer, FeedbackValue, FeedbackSubmission, MessageFeedback, ThreadSummary, ThreadPlacement, ChatViewProps, ApprovalDecision, PendingApproval } from "./src/react/types.ts";
 export { useChatThreadQuery } from "./src/react/use-chat-thread-query.ts";

package/index.ts

@@ -3,4 +3,4 @@
 // Hand edits are overwritten on the next watch; this file is read-only.
 
 export const PACKAGE_IDENTIFIER = "@dbx-tools/ui-mastra";
-export const PACKAGE_VERSION = "0.9.175";
+export const PACKAGE_VERSION = "0.9.177";

package/package.json

@@ -32,12 +32,12 @@
   },
   "dependencies": {
     "@databricks/appkit-ui": "0.81.0",
-    "@dbx-tools/shared-auth": "^0.9.175",
-    "@dbx-tools/shared-core": "^0.9.175",
-    "@dbx-tools/shared-genie": "^0.9.175",
-    "@dbx-tools/shared-mastra": "^0.9.175",
-    "@dbx-tools/shared-model": "^0.9.175",
-    "@dbx-tools/ui": "^0.9.175",
+    "@dbx-tools/shared-auth": "^0.9.177",
+    "@dbx-tools/shared-core": "^0.9.177",
+    "@dbx-tools/shared-genie": "^0.9.177",
+    "@dbx-tools/shared-mastra": "^0.9.177",
+    "@dbx-tools/shared-model": "^0.9.177",
+    "@dbx-tools/ui": "^0.9.177",
     "@mastra/client-js": "1.50.0",
     "@mastra/core": "1.71.0",
     "@tanstack/react-table": "^8.21.3",
@@ -55,7 +55,7 @@
   "publishConfig": {
     "access": "public"
   },
-  "version": "0.9.175",
+  "version": "0.9.177",
   "types": "index.ts",
   "dbxToolsConfig": {
     "tags": [

package/src/react/_assistant-response-status.tsx

@@ -1,13 +1,13 @@
 import { Spinner } from "@databricks/appkit-ui/react";
 
-import type { ChatStatus } from "./types.ts";
+import { isRespondingChatStatus, type ChatStatus } from "./types.ts";
 
 /** Whether an assistant bubble should show its active-response status. */
 export const shouldShowAssistantResponseStatus = (
   isLast: boolean,
   status: ChatStatus,
   hasVisibleResponse: boolean,
-): boolean => isLast && hasVisibleResponse && (status === "submitted" || status === "streaming");
+): boolean => isLast && hasVisibleResponse && isRespondingChatStatus(status);
 
 /** Compact progress affordance for an active assistant response. */
 export const AssistantResponseSpinner = () => (

package/src/react/bubbles.tsx

@@ -36,13 +36,14 @@ import { ExportMenu } from "./export-menu.tsx";
 import { FeedbackControls } from "./feedback-controls.tsx";
 import { ToolMarkdown } from "./markdown.tsx";
 import { ToolSessionPill, humanizeToolName } from "./tool-pill.tsx";
-import type {
-  ApprovalDecision,
-  ChatStatus,
-  FeedbackSubmission,
-  FeedbackValue,
-  PendingApproval,
-  ToolEvent,
+import {
+  isRespondingChatStatus,
+  type ApprovalDecision,
+  type ChatStatus,
+  type FeedbackSubmission,
+  type FeedbackValue,
+  type PendingApproval,
+  type ToolEvent,
 } from "./types.ts";
 import { copyText } from "../support/clipboard.ts";
 import type { ExportFormat } from "../support/export.ts";
@@ -390,7 +391,7 @@ export const AssistantBubble = ({
   // pop in mid-stream. A bubble is "settled" if it isn't the active
   // streaming target - either the agent has returned to `ready` /
   // `error`, or a newer message has taken over the `isLast` slot.
-  const isStreamingThisBubble = isLast && (status === "streaming" || status === "submitted");
+  const isStreamingThisBubble = isLast && isRespondingChatStatus(status);
   const showResponseStatus = shouldShowAssistantResponseStatus(
     isLast,
     status,

package/src/react/chat-composer.tsx

@@ -34,7 +34,7 @@ import { imageFileToUIPart, MAX_COMPOSER_IMAGES } from "./_image-attachments.ts"
 import { ModelSelector } from "./_model-selector.tsx";
 import { ExportMenu } from "./export-menu.tsx";
 import { SuggestionPills } from "./suggestion-pills.tsx";
-import type { ChatViewProps } from "./types.ts";
+import { isActiveChatStatus, isRespondingChatStatus, type ChatViewProps } from "./types.ts";
 
 type QueuedSteerListProps = {
   queuedSteers: NonNullable<ChatViewProps["queuedSteers"]>;
@@ -316,7 +316,8 @@ export const ChatComposer = ({
     return observeComposerWidth(group, resizeTextarea);
   }, [resizeTextarea]);
 
-  const isRunning = status === "submitted" || status === "streaming";
+  const isRunning = isActiveChatStatus(status);
+  const canStop = isRespondingChatStatus(status);
   const hasDraft = Boolean(input.trim() || files.length > 0);
   const submit = () => {
     const text = input.trim();
@@ -560,7 +561,7 @@ export const ChatComposer = ({
                     )}
                   </span>
                 ))}
-              {isRunning && onStop && !hasDraft ? (
+              {canStop && onStop && !hasDraft ? (
                 <InputGroupButton
                   type="button"
                   size="icon-sm"

package/src/react/chat-stream.ts

@@ -3,6 +3,7 @@ import {
   MLFLOW_TRACE_ID_HEADER,
   MLFLOW_TRACE_PART_TYPE,
   MlflowTracePartDataSchema,
+  RESPONSE_COMPLETE_PART_TYPE,
   isStaleMastraResumeError,
 } from "@dbx-tools/shared-mastra";
 import { readUIMessageStream, type UIMessage, type UIMessageChunk } from "ai";
@@ -12,6 +13,7 @@ import type {
   ThreadSessionReader,
   ThreadSessionUpdater,
 } from "./chat-sessions.ts";
+import type { ChatStatus } from "./types.ts";
 import type { MastraStreamResponse } from "../support/mastra-client.ts";
 
 const logger = log.logger("ui-mastra/chat");
@@ -29,6 +31,13 @@ export const readMlflowTracePart = (chunk: UIMessageChunk): string | undefined =
   return parsed.success ? parsed.data.traceId : undefined;
 };
 
+/** Map lifecycle-only chunks onto the UI's active-turn state. */
+export const chatStatusForChunk = (chunk: UIMessageChunk): ChatStatus | undefined => {
+  if ((chunk as { type?: string }).type === RESPONSE_COMPLETE_PART_TYPE) return "finishing";
+  if (chunk.type === "finish" || chunk.type === "abort") return "ready";
+  return undefined;
+};
+
 class StreamAborted extends Error {}
 
 /** Stop reading a background run once the AI SDK has delivered its terminal chunk. */
@@ -46,7 +55,7 @@ export function closeOnTerminalChunk(
         while (true) {
           const { done, value } = await reader.read();
           if (done) {
-            controller.close();
+            controller.error(new Error("Chat stream ended before a terminal AI SDK chunk"));
             return;
           }
           onChunk?.(value);
@@ -140,7 +149,15 @@ export function useChatStream({ getSession, updateSession, writeMessages }: UseC
             () => {
               terminalReceived = true;
             },
-            (chunk) => recordTraceId(readMlflowTracePart(chunk)),
+            (chunk) => {
+              recordTraceId(readMlflowTracePart(chunk));
+              const status = chatStatusForChunk(chunk);
+              if (!status) return;
+              updateSession(threadId, (current) => ({
+                ...current,
+                status,
+              }));
+            },
           ),
           terminateOnError: true,
         })) {
@@ -150,7 +167,12 @@ export function useChatStream({ getSession, updateSession, writeMessages }: UseC
           updateSession(threadId, (current) => ({
             ...current,
             runId: runIdRef.current,
-            status: terminalReceived ? "ready" : "streaming",
+            status:
+              terminalReceived || current.status === "finishing"
+                ? terminalReceived
+                  ? "ready"
+                  : "finishing"
+                : "streaming",
           }));
         }
       } catch (error) {

package/src/react/chat-transcript.tsx

@@ -17,7 +17,7 @@ import React, { useCallback, useEffect, useLayoutEffect, useRef, useState } from
 import { AssistantBubble, UserBubble } from "./bubbles.tsx";
 import { genieFeedbackTarget } from "./chat-feedback.ts";
 import { SuggestionPills } from "./suggestion-pills.tsx";
-import type { ChatViewProps } from "./types.ts";
+import { isRespondingChatStatus, type ChatViewProps } from "./types.ts";
 import { messageText } from "../support/export.ts";
 
 const BOTTOM_THRESHOLD_PX = 24;
@@ -186,7 +186,7 @@ export const ChatTranscript = ({
   onFeedback,
   mlflowExperimentUrl,
 }: ChatTranscriptProps) => {
-  const isRunning = status === "submitted" || status === "streaming";
+  const isRunning = isRespondingChatStatus(status);
   const lastMessage = messages.at(-1);
   const lastEvents = lastMessage ? toolEventsByMessage[lastMessage.id] : undefined;
   const lastAssistantParts = lastMessage?.role === "assistant" ? lastMessage.parts : [];

package/src/react/chat-view.tsx

@@ -2,7 +2,7 @@ import { ChatComposer } from "./chat-composer.tsx";
 import { ChatThreadLayout } from "./chat-thread-layout.tsx";
 import { ChatTranscript, useChatTranscriptController } from "./chat-transcript.tsx";
 import { collectSuggestions } from "./suggestions.ts";
-import type { ChatViewProps } from "./types.ts";
+import { isActiveChatStatus, type ChatViewProps } from "./types.ts";
 
 /**
  * Controlled chat facade for hosts that own message and transport state.
@@ -81,7 +81,7 @@ export const ChatView = ({
   });
   const latestMessage = messages[messages.length - 1];
   const latestAssistant = latestMessage?.role === "assistant" ? latestMessage : undefined;
-  const running = status === "submitted" || status === "streaming";
+  const running = isActiveChatStatus(status);
   const followUpQuestions =
     latestAssistant && !running
       ? collectSuggestions(

package/src/react/types.ts

@@ -10,7 +10,15 @@ import type { ExportFormat } from "../support/export.ts";
 // layer and the `useMastraChat` driver can share them.
 
 /** User-visible lifecycle state of the active chat turn. */
-export type ChatStatus = "submitted" | "streaming" | "ready" | "error";
+export type ChatStatus = "submitted" | "streaming" | "finishing" | "ready" | "error";
+
+/** Whether a turn still owns its session and queues new messages. */
+export const isActiveChatStatus = (status: ChatStatus): boolean =>
+  status === "submitted" || status === "streaming" || status === "finishing";
+
+/** Whether the assistant is still visibly producing its response. */
+export const isRespondingChatStatus = (status: ChatStatus): boolean =>
+  status === "submitted" || status === "streaming";
 
 /**
  * Lifecycle of a single tool invocation surfaced inline in the assistant

package/src/support/thread-sessions.ts

@@ -1,6 +1,12 @@
 import type { RequestContext } from "@mastra/client-js";
 import type { UIMessage } from "ai";
-import type { ChatStatus, MessageFeedback, PendingApproval, QueuedSteer } from "../react/types.ts";
+import {
+  isActiveChatStatus,
+  type ChatStatus,
+  type MessageFeedback,
+  type PendingApproval,
+  type QueuedSteer,
+} from "../react/types.ts";
 
 /** A queued steer plus the application context captured when it was submitted. */
 export type SessionQueuedSteer = QueuedSteer & {
@@ -60,9 +66,9 @@ export function createThreadSession(): ThreadSession {
   };
 }
 
-/** Return whether a thread has a submitted or actively streaming run. */
+/** Return whether a thread still owns an active run. */
 export function isSessionRunning(session: ThreadSession): boolean {
-  return session.status === "submitted" || session.status === "streaming";
+  return isActiveChatStatus(session.status);
 }
 
 /** Return the newest user message in a transcript or turn. */