@dbx-tools/ui-mastra 0.6.211 → 0.6.212

package/src/react/chat-stream-reducer.ts

@@ -0,0 +1,241 @@
+import { GenieWriterEventSchema, type MastraStreamChunk } from "@dbx-tools/shared-mastra";
+import type { UIMessage } from "ai";
+import type { PendingApproval, ToolEvent } from "./types.ts";
+
+/** Pure accumulated state for one assistant response stream. */
+export type ChatStreamState = {
+  textSegments: string[];
+  reasoning: string;
+  toolEvents: ToolEvent[];
+  pendingApprovals: PendingApproval[];
+  runId: string | null;
+  streaming: boolean;
+};
+
+/** Session fields changed by one stream reduction. */
+export type ChatStreamChanges = {
+  runId: boolean;
+  status: boolean;
+  toolEvents: boolean;
+  pendingApprovals: boolean;
+};
+
+/** Result of reducing one validated stream chunk. */
+export type ChatStreamReduction = {
+  state: ChatStreamState;
+  changes: ChatStreamChanges;
+  assistantChanged: boolean;
+  error?: string;
+};
+
+/** Seed pure stream state from an existing assistant message and session data. */
+export function createChatStreamState(options: {
+  existing?: UIMessage;
+  pendingApprovals: PendingApproval[];
+  runId: string | null;
+  toolEvents: ToolEvent[];
+}): ChatStreamState {
+  const textSegments: string[] = [];
+  let reasoning = "";
+  for (const part of options.existing?.parts ?? []) {
+    if (part.type === "text") textSegments.push(part.text);
+    else if (part.type === "reasoning") reasoning += (part as { text?: string }).text ?? "";
+  }
+  return {
+    textSegments,
+    reasoning,
+    toolEvents: options.toolEvents,
+    pendingApprovals: options.pendingApprovals,
+    runId: options.runId,
+    streaming: false,
+  };
+}
+
+/** Materialize the assistant message represented by accumulated stream state. */
+export function chatStreamAssistantMessage(assistantId: string, state: ChatStreamState): UIMessage {
+  const parts: UIMessage["parts"] = [];
+  if (state.reasoning) parts.push({ type: "reasoning", text: state.reasoning });
+  for (const segment of state.textSegments) {
+    if (segment.length > 0) parts.push({ type: "text", text: segment });
+  }
+  return {
+    id: assistantId,
+    role: "assistant",
+    parts: parts.length > 0 ? parts : [{ type: "text", text: "" }],
+  };
+}
+
+const withRunId = (
+  state: ChatStreamState,
+  chunk: MastraStreamChunk,
+): { state: ChatStreamState; changed: boolean } => {
+  if (!chunk.runId || state.runId) return { state, changed: false };
+  return { state: { ...state, runId: chunk.runId }, changed: true };
+};
+
+const reduction = (
+  state: ChatStreamState,
+  options: {
+    runIdChanged: boolean;
+    assistantChanged?: boolean;
+    statusChanged?: boolean;
+    toolEventsChanged?: boolean;
+    pendingApprovalsChanged?: boolean;
+    error?: string;
+  },
+): ChatStreamReduction => ({
+  state,
+  changes: {
+    runId: options.runIdChanged,
+    status: options.statusChanged ?? false,
+    toolEvents: options.toolEventsChanged ?? false,
+    pendingApprovals: options.pendingApprovalsChanged ?? false,
+  },
+  assistantChanged: options.assistantChanged ?? false,
+  ...(options.error ? { error: options.error } : {}),
+});
+
+/** Reduce one validated chunk into immutable assistant and session state. */
+export function reduceChatStreamChunk(
+  previous: ChatStreamState,
+  chunk: MastraStreamChunk,
+): ChatStreamReduction {
+  const run = withRunId(previous, chunk);
+  const state = run.state;
+  const markStreaming = !state.streaming;
+
+  switch (chunk.type) {
+    case "unknown":
+    case "text-start":
+    case "text-end":
+      return reduction(state, { runIdChanged: run.changed });
+
+    case "text-delta": {
+      const textSegments = state.textSegments.length > 0 ? [...state.textSegments] : [""];
+      textSegments[textSegments.length - 1] += chunk.payload.text;
+      return reduction(
+        { ...state, textSegments, streaming: true },
+        {
+          runIdChanged: run.changed,
+          assistantChanged: true,
+          statusChanged: markStreaming,
+        },
+      );
+    }
+
+    case "reasoning-delta":
+      return reduction(
+        {
+          ...state,
+          reasoning: state.reasoning + chunk.payload.text,
+          streaming: true,
+        },
+        {
+          runIdChanged: run.changed,
+          assistantChanged: true,
+          statusChanged: markStreaming,
+        },
+      );
+
+    case "tool-call":
+      return reduction(
+        {
+          ...state,
+          streaming: true,
+          toolEvents: [
+            ...state.toolEvents,
+            {
+              id: chunk.payload.toolCallId,
+              toolName: chunk.payload.toolName,
+              status: "running",
+            },
+          ],
+        },
+        {
+          runIdChanged: run.changed,
+          assistantChanged: true,
+          statusChanged: markStreaming,
+          toolEventsChanged: true,
+        },
+      );
+
+    case "tool-call-approval": {
+      const approvalRunId = chunk.runId ?? chunk.payload.runId ?? state.runId;
+      if (!approvalRunId) {
+        return reduction(state, { runIdChanged: run.changed });
+      }
+      const approval: PendingApproval = {
+        toolName: chunk.payload.toolName,
+        toolCallId: chunk.payload.toolCallId,
+        runId: approvalRunId,
+        input: chunk.payload.args,
+      };
+      const exists = state.pendingApprovals.some(
+        (current) => current.toolCallId === approval.toolCallId,
+      );
+      return reduction(
+        {
+          ...state,
+          streaming: true,
+          pendingApprovals: exists ? state.pendingApprovals : [...state.pendingApprovals, approval],
+        },
+        {
+          runIdChanged: run.changed,
+          assistantChanged: true,
+          statusChanged: markStreaming,
+          pendingApprovalsChanged: !exists,
+        },
+      );
+    }
+
+    case "tool-result":
+    case "tool-error":
+      return reduction(
+        {
+          ...state,
+          toolEvents: state.toolEvents.map((event) =>
+            event.id === chunk.payload.toolCallId
+              ? { ...event, status: chunk.type === "tool-result" ? "done" : "error" }
+              : event,
+          ),
+        },
+        {
+          runIdChanged: run.changed,
+          toolEventsChanged: true,
+        },
+      );
+
+    case "tool-output": {
+      const progress = GenieWriterEventSchema.safeParse(chunk.payload.output);
+      if (!progress.success) return reduction(state, { runIdChanged: run.changed });
+      return reduction(
+        {
+          ...state,
+          toolEvents: state.toolEvents.map((event) =>
+            event.id === chunk.payload.toolCallId
+              ? { ...event, progress: [...(event.progress ?? []), progress.data] }
+              : event,
+          ),
+        },
+        {
+          runIdChanged: run.changed,
+          toolEventsChanged: true,
+        },
+      );
+    }
+
+    case "error": {
+      const detail = chunk.payload?.error ?? chunk.payload?.message;
+      return reduction(state, {
+        runIdChanged: run.changed,
+        error:
+          typeof detail === "string" && detail ? detail : "The assistant stream reported an error.",
+      });
+    }
+
+    default: {
+      const exhaustive: never = chunk;
+      return exhaustive;
+    }
+  }
+}

package/src/react/chat-stream.ts

@@ -0,0 +1,117 @@
+import { feedback } from "@dbx-tools/shared-mastra";
+import { useCallback } from "react";
+import type { MastraStreamResponse } from "../support/mastra-stream.ts";
+import type {
+  ThreadMessageWriter,
+  ThreadSessionReader,
+  ThreadSessionUpdater,
+} from "./chat-sessions.ts";
+import {
+  chatStreamAssistantMessage,
+  createChatStreamState,
+  reduceChatStreamChunk,
+} from "./chat-stream-reducer.ts";
+
+/** Read the MLflow trace id captured by the server on a stream response. */
+const readMlflowTraceId = (stream: unknown): string | undefined => {
+  const headers = (stream as { headers?: { get?: (name: string) => string | null } })?.headers;
+  return headers?.get?.(feedback.MLFLOW_TRACE_ID_HEADER)?.trim() || undefined;
+};
+
+class StreamAborted extends Error {}
+
+interface UseChatStreamOptions {
+  getSession: ThreadSessionReader;
+  updateSession: ThreadSessionUpdater;
+  writeMessages: ThreadMessageWriter;
+}
+
+/** Translate one validated Mastra data stream into thread messages and tool state. */
+export function useChatStream({ getSession, updateSession, writeMessages }: UseChatStreamOptions) {
+  return useCallback(
+    async (
+      threadId: string,
+      stream: MastraStreamResponse,
+      assistantId: string,
+      runIdRef: { current: string | null },
+      signal: AbortSignal,
+    ) => {
+      const traceId = readMlflowTraceId(stream);
+      if (traceId) {
+        updateSession(threadId, (session) =>
+          session.feedbackByMessage[assistantId]?.traceId === traceId
+            ? session
+            : {
+                ...session,
+                feedbackByMessage: {
+                  ...session.feedbackByMessage,
+                  [assistantId]: {
+                    ...session.feedbackByMessage[assistantId],
+                    traceId,
+                  },
+                },
+              },
+        );
+      }
+
+      const session = getSession(threadId);
+      const existing = session.messages.find((message) => message.id === assistantId);
+      let state = createChatStreamState({
+        ...(existing ? { existing } : {}),
+        pendingApprovals: session.pendingApprovalsByMessage[assistantId] ?? [],
+        runId: runIdRef.current,
+        toolEvents: session.toolEventsByMessage[assistantId] ?? [],
+      });
+
+      const upsertAssistant = () => {
+        const next = [...getSession(threadId).messages];
+        const index = next.findIndex((message) => message.id === assistantId);
+        const message = chatStreamAssistantMessage(assistantId, state);
+        if (index === -1) next.push(message);
+        else next[index] = message;
+        writeMessages(threadId, next);
+      };
+
+      try {
+        await stream.processDataStream({
+          onChunk: async (chunk) => {
+            if (signal.aborted) throw new StreamAborted();
+            const next = reduceChatStreamChunk(state, chunk);
+            state = next.state;
+            runIdRef.current = state.runId;
+
+            if (Object.values(next.changes).some(Boolean)) {
+              updateSession(threadId, (current) => ({
+                ...current,
+                ...(next.changes.runId ? { runId: state.runId } : {}),
+                ...(next.changes.status ? { status: "streaming" as const } : {}),
+                ...(next.changes.toolEvents
+                  ? {
+                      toolEventsByMessage: {
+                        ...current.toolEventsByMessage,
+                        [assistantId]: state.toolEvents,
+                      },
+                    }
+                  : {}),
+                ...(next.changes.pendingApprovals
+                  ? {
+                      pendingApprovalsByMessage: {
+                        ...current.pendingApprovalsByMessage,
+                        [assistantId]: state.pendingApprovals,
+                      },
+                    }
+                  : {}),
+              }));
+            }
+            if (next.assistantChanged) upsertAssistant();
+            if (next.error) throw new Error(next.error);
+          },
+        });
+      } catch (error) {
+        if (error instanceof StreamAborted || signal.aborted) return;
+        throw error;
+      }
+    },
+    [getSession, updateSession, writeMessages],
+  );
+}

package/src/react/chat-thread-layout.tsx

@@ -0,0 +1,214 @@
+import {
+  Button,
+  Tooltip,
+  TooltipContent,
+  TooltipProvider,
+  TooltipTrigger,
+  cn,
+} from "@dbx-tools/ui-appkit/react";
+import { PanelLeftIcon, PanelRightIcon } from "lucide-react";
+import React, { useEffect, useRef, useState } from "react";
+import { ThreadSidebar, type ThreadSidebarProps } from "./thread-sidebar.tsx";
+import { ThreadTabs } from "./thread-tabs.tsx";
+import type { ChatViewProps } from "./types.ts";
+import { closeThreadTab, nextActiveThreadTab, syncThreadTabs } from "../support/thread-tabs.ts";
+
+const SIDE_PANEL_MIN_WIDTH_PX = 768;
+
+const useIsNarrow = (ref: React.RefObject<HTMLElement | null>): boolean => {
+  const [isNarrow, setIsNarrow] = useState(() =>
+    typeof window === "undefined" ? false : window.innerWidth < SIDE_PANEL_MIN_WIDTH_PX,
+  );
+  useEffect(() => {
+    const element = ref.current;
+    if (!element || typeof ResizeObserver === "undefined") return;
+    const observer = new ResizeObserver((entries) => {
+      const width = entries[0]?.contentRect.width ?? element.clientWidth;
+      if (width > 0) setIsNarrow(width < SIDE_PANEL_MIN_WIDTH_PX);
+    });
+    observer.observe(element);
+    return () => observer.disconnect();
+  }, [ref]);
+  return isNarrow;
+};
+
+type ChatThreadLayoutProps = {
+  children: React.ReactNode;
+  className: ChatViewProps["className"];
+  threads: ChatViewProps["threads"];
+  threadPlacement: NonNullable<ChatViewProps["threadPlacement"]>;
+  activeThreadId: ChatViewProps["activeThreadId"];
+  streamingThreadIds: NonNullable<ChatViewProps["streamingThreadIds"]>;
+  isLoadingThreads: NonNullable<ChatViewProps["isLoadingThreads"]>;
+  onSelectThread: ChatViewProps["onSelectThread"];
+  onNewThread: ChatViewProps["onNewThread"];
+  onDeleteThread: ChatViewProps["onDeleteThread"];
+  onRenameThread: ChatViewProps["onRenameThread"];
+  onCancelThread: ChatViewProps["onCancelThread"];
+  sidebarOpen: ChatViewProps["sidebarOpen"];
+  onToggleSidebar: ChatViewProps["onToggleSidebar"];
+};
+
+type ThreadDrawerProps = {
+  list: Omit<ThreadSidebarProps, "onHide" | "side" | "className">;
+  side: "left" | "right";
+  onClose: () => void;
+  onSelectThread: NonNullable<ChatViewProps["onSelectThread"]>;
+  onNewThread: ChatViewProps["onNewThread"];
+};
+
+const ThreadDrawer = ({ list, side, onClose, onSelectThread, onNewThread }: ThreadDrawerProps) => (
+  <div className={cn("fixed inset-0 z-40 flex", side === "right" && "justify-end")}>
+    <div className="absolute inset-0 bg-black/50" onClick={onClose} aria-hidden="true" />
+    <ThreadSidebar
+      {...list}
+      onHide={onClose}
+      side={side}
+      onSelect={(id) => {
+        onSelectThread(id);
+        onClose();
+      }}
+      {...(onNewThread
+        ? {
+            onNew: () => {
+              onNewThread();
+              onClose();
+            },
+          }
+        : {})}
+      className="relative z-10 w-[85vw] max-w-xs shadow-xl"
+    />
+  </div>
+);
+
+/** Internal responsive shell that owns conversation navigation surfaces. */
+export const ChatThreadLayout = ({
+  children,
+  className,
+  threads,
+  threadPlacement,
+  activeThreadId,
+  streamingThreadIds,
+  isLoadingThreads,
+  onSelectThread,
+  onNewThread,
+  onDeleteThread,
+  onRenameThread,
+  onCancelThread,
+  sidebarOpen: sidebarOpenProp,
+  onToggleSidebar,
+}: ChatThreadLayoutProps) => {
+  const rootRef = useRef<HTMLDivElement>(null);
+  const isNarrow = useIsNarrow(rootRef);
+  const showThreads = Boolean(threads && onSelectThread) && threadPlacement !== "disabled";
+  const placement = threadPlacement === "auto" ? (isNarrow ? "top" : "left") : threadPlacement;
+  const tabbedThreads = showThreads && placement === "top";
+  const dockedSide = placement === "right" ? "right" : "left";
+  const dockedThreads = showThreads && (placement === "left" || placement === "right");
+
+  const [internalSidebarOpen, setInternalSidebarOpen] = useState(true);
+  const inlineSidebarOpen = sidebarOpenProp ?? internalSidebarOpen;
+  const toggleInlineSidebar = () => {
+    if (onToggleSidebar) onToggleSidebar();
+    else setInternalSidebarOpen((open) => !open);
+  };
+
+  const [drawerOpen, setDrawerOpen] = useState(false);
+  useEffect(() => {
+    if (!isNarrow) setDrawerOpen(false);
+  }, [isNarrow]);
+
+  const sidebarOpen = isNarrow ? drawerOpen : inlineSidebarOpen;
+  const toggleSidebar = () => {
+    if (isNarrow) setDrawerOpen((open) => !open);
+    else toggleInlineSidebar();
+  };
+  const showSidebarToggle = dockedThreads && (isNarrow || !inlineSidebarOpen);
+  const SidebarToggleIcon = dockedSide === "right" ? PanelRightIcon : PanelLeftIcon;
+
+  const threadListProps: Omit<ThreadSidebarProps, "onHide" | "side" | "className"> = {
+    threads: threads ?? [],
+    ...(activeThreadId ? { activeThreadId } : {}),
+    streamingThreadIds,
+    isLoading: isLoadingThreads,
+    onSelect: (id) => onSelectThread?.(id),
+    ...(onNewThread ? { onNew: onNewThread } : {}),
+    ...(onDeleteThread ? { onDelete: onDeleteThread } : {}),
+    ...(onRenameThread ? { onRename: onRenameThread } : {}),
+    ...(onCancelThread ? { onCancel: onCancelThread } : {}),
+  };
+
+  const [openTabIds, setOpenTabIds] = useState<string[]>([]);
+  useEffect(() => {
+    if (!tabbedThreads) return;
+    setOpenTabIds((previous) => syncThreadTabs(previous, threads ?? [], activeThreadId));
+  }, [tabbedThreads, threads, activeThreadId]);
+
+  const closeTab = (threadId: string) => {
+    const fallback = nextActiveThreadTab(openTabIds, threadId);
+    setOpenTabIds((previous) => closeThreadTab(previous, threadId));
+    if (threadId !== activeThreadId) return;
+    if (fallback) onSelectThread?.(fallback);
+    else onNewThread?.();
+  };
+
+  return (
+    <TooltipProvider delayDuration={200}>
+      <div
+        ref={rootRef}
+        className={cn(
+          "flex h-full min-h-0",
+          dockedThreads && dockedSide === "right" && "flex-row-reverse",
+          className,
+        )}
+      >
+        {dockedThreads &&
+          (isNarrow
+            ? drawerOpen &&
+              onSelectThread && (
+                <ThreadDrawer
+                  list={threadListProps}
+                  side={dockedSide}
+                  onClose={toggleSidebar}
+                  onSelectThread={onSelectThread}
+                  onNewThread={onNewThread}
+                />
+              )
+            : inlineSidebarOpen && (
+                <ThreadSidebar {...threadListProps} onHide={toggleSidebar} side={dockedSide} />
+              ))}
+        <div className="flex h-full min-w-0 flex-1 flex-col">
+          {tabbedThreads && (
+            <ThreadTabs {...threadListProps} openThreadIds={openTabIds} onCloseTab={closeTab} />
+          )}
+          {showSidebarToggle && (
+            <div
+              className={cn(
+                "mx-auto flex w-full max-w-4xl items-center gap-2 px-3 pb-2 pt-1 text-xs text-muted-foreground md:gap-3 md:px-6",
+                dockedSide === "right" && "justify-end",
+              )}
+            >
+              <Tooltip>
+                <TooltipTrigger asChild>
+                  <Button
+                    type="button"
+                    variant="ghost"
+                    size="icon-sm"
+                    onClick={toggleSidebar}
+                    aria-label={sidebarOpen ? "Hide conversations" : "Show conversations"}
+                  >
+                    <SidebarToggleIcon className="size-4" />
+                  </Button>
+                </TooltipTrigger>
+                <TooltipContent>
+                  {sidebarOpen ? "Hide conversations" : "Show conversations"}
+                </TooltipContent>
+              </Tooltip>
+            </div>
+          )}
+          {children}
+        </div>
+      </div>
+    </TooltipProvider>
+  );
+};