@dbx-tools/ui-mastra 0.9.161 → 0.9.165

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 { closeOnTerminalChunk, waitForTerminalChunk, useChatStream } from "./src/react/chat-stream.ts";
+export { readMlflowTracePart, 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";

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.161";
+export const PACKAGE_VERSION = "0.9.165";

package/package.json

@@ -32,12 +32,12 @@
   },
   "dependencies": {
     "@databricks/appkit-ui": "0.81.0",
-    "@dbx-tools/shared-auth": "^0.9.161",
-    "@dbx-tools/shared-core": "^0.9.161",
-    "@dbx-tools/shared-genie": "^0.9.161",
-    "@dbx-tools/shared-mastra": "^0.9.161",
-    "@dbx-tools/shared-model": "^0.9.161",
-    "@dbx-tools/ui": "^0.9.161",
+    "@dbx-tools/shared-auth": "^0.9.165",
+    "@dbx-tools/shared-core": "^0.9.165",
+    "@dbx-tools/shared-genie": "^0.9.165",
+    "@dbx-tools/shared-mastra": "^0.9.165",
+    "@dbx-tools/shared-model": "^0.9.165",
+    "@dbx-tools/ui": "^0.9.165",
     "@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.161",
+  "version": "0.9.165",
   "types": "index.ts",
   "dbxToolsConfig": {
     "tags": [

package/src/react/chat-stream.ts

@@ -1,5 +1,10 @@
 import { log, toError, trimToUndefined } from "@dbx-tools/shared-core";
-import { MLFLOW_TRACE_ID_HEADER, isStaleMastraResumeError } from "@dbx-tools/shared-mastra";
+import {
+  MLFLOW_TRACE_ID_HEADER,
+  MLFLOW_TRACE_PART_TYPE,
+  MlflowTracePartDataSchema,
+  isStaleMastraResumeError,
+} from "@dbx-tools/shared-mastra";
 import { readUIMessageStream, type UIMessage, type UIMessageChunk } from "ai";
 import { useCallback } from "react";
 import type {
@@ -17,12 +22,20 @@ const readMlflowTraceId = (stream: unknown): string | undefined => {
   return trimToUndefined(headers?.get?.(MLFLOW_TRACE_ID_HEADER));
 };
 
+/** Read the trace-id fallback carried by one transient AI SDK data part. */
+export const readMlflowTracePart = (chunk: UIMessageChunk): string | undefined => {
+  if ((chunk as { type?: string }).type !== MLFLOW_TRACE_PART_TYPE) return undefined;
+  const parsed = MlflowTracePartDataSchema.safeParse((chunk as { data?: unknown }).data);
+  return parsed.success ? parsed.data.traceId : undefined;
+};
+
 class StreamAborted extends Error {}
 
 /** Stop reading a background run once the AI SDK has delivered its terminal chunk. */
 export function closeOnTerminalChunk(
   stream: ReadableStream<UIMessageChunk>,
   onTerminal?: () => void,
+  onChunk?: (chunk: UIMessageChunk) => void,
 ): ReadableStream<UIMessageChunk> {
   let reader: ReadableStreamDefaultReader<UIMessageChunk> | undefined;
   let cancelled = false;
@@ -36,6 +49,7 @@ export function closeOnTerminalChunk(
             controller.close();
             return;
           }
+          onChunk?.(value);
           controller.enqueue(value);
           if (value.type === "finish" || value.type === "abort") {
             onTerminal?.();
@@ -95,8 +109,8 @@ export function useChatStream({ getSession, updateSession, writeMessages }: UseC
       runIdRef: { current: string | null },
       signal: AbortSignal,
     ) => {
-      const traceId = readMlflowTraceId(stream);
-      if (traceId) {
+      const recordTraceId = (traceId: string | undefined): void => {
+        if (!traceId) return;
         updateSession(threadId, (session) =>
           session.feedbackByMessage[assistantId]?.traceId === traceId
             ? session
@@ -111,7 +125,8 @@ export function useChatStream({ getSession, updateSession, writeMessages }: UseC
                 },
               },
         );
-      }
+      };
+      recordTraceId(readMlflowTraceId(stream));
 
       try {
         let terminalReceived = false;
@@ -120,9 +135,13 @@ export function useChatStream({ getSession, updateSession, writeMessages }: UseC
         );
         for await (const message of readUIMessageStream({
           ...(existing ? { message: existing } : {}),
-          stream: closeOnTerminalChunk(stream.stream, () => {
-            terminalReceived = true;
-          }),
+          stream: closeOnTerminalChunk(
+            stream.stream,
+            () => {
+              terminalReceived = true;
+            },
+            (chunk) => recordTraceId(readMlflowTracePart(chunk)),
+          ),
           terminateOnError: true,
         })) {
           if (signal.aborted) throw new StreamAborted();