@dbx-tools/ui-mastra 0.9.99 → 0.9.105

package/README.md

@@ -2,11 +2,21 @@
 
 React chat UI for the AppKit-Mastra plugin.
 
-Import this package when a Databricks App needs a production-ready chat surface
-for [`@dbx-tools/appkit-mastra`](../../node/appkit-mastra): streaming
-assistant responses, model selection, Genie progress events, inline charts/data
-tables, tool approvals, conversation history, thread management, export, and
-MLflow feedback.
+Add a complete chat surface for
+[`@dbx-tools/appkit-mastra`](../../node/appkit-mastra), including streaming,
+model selection, Genie progress, charts, tables, approvals, history, threads,
+export, and MLflow feedback.
+
+## Quick Start
+
+```tsx
+import { MastraChat } from "@dbx-tools/ui-mastra/react";
+import "@dbx-tools/ui-mastra/styles.css";
+
+export function Assistant() {
+  return <MastraChat showModelPicker enableExport />;
+}
+```
 
 Key features:
 
@@ -47,7 +57,7 @@ Key features:
 - Export menu for PDF and Markdown, resolving charts and tables so
   exported conversations remain useful offline.
 
-## Why Not Just AppKit UI?
+## Choose This UI Or AppKit UI
 
 Use native `@databricks/appkit-ui` when you need its general primitives, Genie
 chat component, or Model Serving hooks directly against native AppKit plugins.
@@ -69,7 +79,7 @@ understand Mastra-specific behavior:
   has no steering.
 - Conversation export that resolves those embeds into Markdown or PDF.
 
-## Add The Styles
+## Stylesheet Setup
 
 ```css
 @import "@databricks/appkit-ui/styles.css";

package/index.ts

@@ -3,7 +3,7 @@
 // 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.99";
+export const PACKAGE_VERSION = "0.9.105";
 export * as reactBubbles from "./src/react/bubbles.tsx";
 export * as reactChatApprovals from "./src/react/chat-approvals.ts";
 export * as reactChatComposer from "./src/react/chat-composer.tsx";
@@ -71,7 +71,7 @@ export type { ChartChrome } from "./src/support/chart-theme.ts";
 export type { ExportFormat, ExportBrand, EmbedResolver, ExportChatOptions } from "./src/support/export.ts";
 export { MastraPluginClient, useMastraConfig, useMastraClient, useMastraMlflowExperiment, useMastraModels, useMastraDefaultModel, useMastraSuggestions, useMastraThreads, useChartFetch, useStatementFetch } from "./src/support/mastra-client.ts";
 export type { MastraStreamResponse, MastraMemoryThread, ByIdFetchState } from "./src/support/mastra-client.ts";
-export { modelStorageKey, readStoredModel, storeSelectedModel } from "./src/support/model-selection.ts";
+export { defaultReasoningEffort, modelStorageKey, readStoredModel, storeSelectedModel } from "./src/support/model-selection.ts";
 export type { MastraRequestContextInput, MastraRequestContextSource, MastraRequestContextSnapshot } from "./src/support/request-context.ts";
 export { DEFAULT_THREAD_SESSION_KEY } from "./src/support/thread-sessions.ts";
 export type { SessionQueuedSteer, SessionRun, ThreadSession } from "./src/support/thread-sessions.ts";

package/package.json

@@ -29,15 +29,15 @@
     "react-dom": "~19.2.4"
   },
   "dependencies": {
-    "@dbx-tools/shared-core": "^0.9.99",
-    "@dbx-tools/shared-genie": "^0.9.99",
-    "@dbx-tools/shared-mastra": "^0.9.99",
-    "@dbx-tools/shared-model": "^0.9.99",
-    "@dbx-tools/ui": "^0.9.99",
+    "@dbx-tools/shared-core": "^0.9.105",
+    "@dbx-tools/shared-genie": "^0.9.105",
+    "@dbx-tools/shared-mastra": "^0.9.105",
+    "@dbx-tools/shared-model": "^0.9.105",
+    "@dbx-tools/ui": "^0.9.105",
     "@mastra/client-js": "1.50.0",
     "@mastra/core": "1.71.0",
     "@tanstack/react-table": "^8.21.3",
-    "ai": "^5.0.0",
+    "ai": "^7.0.0",
     "echarts": "^6.0.0",
     "echarts-for-react": "^3.0.2",
     "lucide-react": "^0.554.0",
@@ -50,7 +50,7 @@
   "publishConfig": {
     "access": "public"
   },
-  "version": "0.9.99",
+  "version": "0.9.105",
   "dbxToolsConfig": {
     "tags": [
       "ui"

package/src/react/bubbles.tsx

@@ -349,11 +349,17 @@ export const AssistantBubble = ({
   feedbackValue,
   mlflowExperimentUrl,
 }: AssistantBubbleProps) => {
-  const toolEvents = mergeToolEvents(
-    toolEventsFromParts(message.parts),
-    events,
-    isLast && status === "error",
+  const pendingApprovals = mergePendingApprovals(
+    collectPendingApprovals(message.parts),
+    externalApprovals,
   );
+  const terminalToolState =
+    isLast && status === "error"
+      ? "error"
+      : pendingApprovals.length === 0 && (!isLast || status === "ready")
+        ? "done"
+        : undefined;
+  const toolEvents = mergeToolEvents(toolEventsFromParts(message.parts), events, terminalToolState);
   const reasoning = [getReasoningText(message.parts), genieReasoningText(toolEvents)]
     .filter(Boolean)
     .join("\n\n");
@@ -397,13 +403,8 @@ export const AssistantBubble = ({
   // {@link DataSlot} inline; unknown / expired ids resolve as
   // nothing so the prose flows unaffected. Suggested questions
   // stay gated on settle to avoid pop-in mid-stream.
-  const pendingApprovals = mergePendingApprovals(
-    collectPendingApprovals(message.parts),
-    externalApprovals,
-  );
-
   return (
-    <Item className="items-start gap-3 border-none bg-transparent p-0">
+    <Item className="items-start gap-2.5 border-none bg-transparent p-0">
       <ItemMedia>
         <RoleAvatar role="assistant" />
       </ItemMedia>
@@ -413,7 +414,7 @@ export const AssistantBubble = ({
        * bubble past the chat's `max-w-4xl`. Flex children default to
        * `min-width: auto`, which sizes to content and overflows.
        */}
-      <ItemContent className="min-w-0 gap-2">
+      <ItemContent className="min-w-0 gap-1.5">
         {/*
          * Tool session pill leads the message. Charts and tables are
          * embedded inline in the prose (via `[chart:<id>]` / `[data:<id>]`
@@ -535,7 +536,7 @@ export const UserBubble = ({ message }: { message: UIMessage }) => {
     .map((p) => p.text)
     .join("");
   return (
-    <Item className="items-start gap-3 border-none bg-transparent p-0">
+    <Item className="items-start gap-2.5 border-none bg-transparent p-0">
       <ItemMedia>
         <RoleAvatar role="user" />
       </ItemMedia>

package/src/react/chat-composer.tsx

@@ -353,7 +353,7 @@ export const ChatComposer = ({
             }}
             placeholder="Send a message..."
             rows={1}
-            className="field-sizing-fixed min-h-10 max-h-64 w-full flex-none overflow-y-auto px-4 pb-2 pt-4 text-xs"
+            className="field-sizing-fixed min-h-10 max-h-64 w-full flex-none overflow-y-auto px-4 pb-2 pt-4 text-base md:text-xs"
           />
           <InputGroupAddon
             align="block-end"

package/src/react/chat-stream.ts

@@ -1,7 +1,7 @@
 import { errorUtils, log, stringUtils } from "@dbx-tools/shared-core";
 import { feedback } from "@dbx-tools/shared-mastra";
 import { isStaleMastraResumeError } from "@dbx-tools/shared-mastra/resume";
-import { readUIMessageStream, type UIMessage } from "ai";
+import { readUIMessageStream, type UIMessage, type UIMessageChunk } from "ai";
 import { useCallback } from "react";
 import type {
   ThreadMessageWriter,
@@ -20,6 +20,43 @@ const readMlflowTraceId = (stream: unknown): string | 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>,
+): ReadableStream<UIMessageChunk> {
+  let reader: ReadableStreamDefaultReader<UIMessageChunk> | undefined;
+  let cancelled = false;
+  return new ReadableStream<UIMessageChunk>({
+    async start(controller) {
+      reader = stream.getReader();
+      try {
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) {
+            controller.close();
+            return;
+          }
+          controller.enqueue(value);
+          if (value.type === "finish" || value.type === "abort") {
+            controller.close();
+            await reader.cancel("terminal AI SDK chunk received").catch(() => undefined);
+            return;
+          }
+        }
+      } catch (error) {
+        if (!cancelled) controller.error(error);
+      } finally {
+        reader.releaseLock();
+        reader = undefined;
+      }
+    },
+    cancel(reason) {
+      cancelled = true;
+      return reader?.cancel(reason);
+    },
+  });
+}
+
 /** Replace or append one assistant message without disturbing the transcript. */
 function upsertAssistant(messages: UIMessage[], assistant: UIMessage): UIMessage[] {
   const next = [...messages];
@@ -69,7 +106,7 @@ export function useChatStream({ getSession, updateSession, writeMessages }: UseC
         );
         for await (const message of readUIMessageStream({
           ...(existing ? { message: existing } : {}),
-          stream: stream.stream,
+          stream: closeOnTerminalChunk(stream.stream),
           terminateOnError: true,
         })) {
           if (signal.aborted) throw new StreamAborted();

package/src/react/chat-transcript.tsx

@@ -219,7 +219,7 @@ export const ChatTranscript = ({
         ) : (
           <div
             ref={controller.contentRef}
-            className="mx-auto flex w-full max-w-4xl flex-col gap-4 px-4 py-4 md:px-6"
+            className="mx-auto flex w-full max-w-4xl flex-col gap-3 px-4 py-3 md:px-6"
           >
             {(isLoadingMore || isLoadingHistory) && (
               <div className="flex items-center justify-center gap-2 py-1 text-xs text-muted-foreground">

package/src/react/markdown.tsx

@@ -258,6 +258,8 @@ export const AssistantMarkdown = ({
     isAnimating={animate}
     className={cn(
       "prose prose-sm dark:prose-invert max-w-none break-words",
+      "leading-relaxed prose-p:my-1.5 prose-p:leading-relaxed",
+      "prose-ul:my-1.5 prose-ol:my-1.5 prose-li:my-0.5",
       "prose-headings:font-semibold prose-headings:tracking-tight",
       "prose-h1:text-lg prose-h1:mt-4 prose-h1:mb-2",
       "prose-h2:text-base prose-h2:mt-4 prose-h2:mb-2",
@@ -288,13 +290,13 @@ export const ToolMarkdown = ({ children }: { children: string }) => (
     controls={false}
     className={cn(
       "prose prose-sm dark:prose-invert max-w-none break-words",
-      "text-[11px] leading-snug text-muted-foreground",
+      "text-xs leading-snug text-muted-foreground",
       "prose-p:my-0.5 prose-p:leading-snug",
       "prose-ul:my-0.5 prose-ul:pl-4 prose-ol:my-0.5 prose-ol:pl-4",
       "prose-li:my-0 prose-li:leading-snug prose-li:marker:text-muted-foreground/60",
       "prose-headings:my-1 prose-headings:text-xs prose-headings:font-semibold",
       "prose-strong:text-foreground/90 prose-strong:font-medium",
-      "prose-code:text-[10px] prose-code:font-medium",
+      "prose-code:text-xs prose-code:font-medium",
     )}
   >
     {children}
@@ -384,11 +386,11 @@ export const SqlBlock = ({ sql }: { sql: string }) => {
         value={formatted}
         className="absolute right-1.5 top-1.5 z-10 bg-background/70 opacity-0 backdrop-blur transition-opacity group-hover:opacity-100 focus-visible:opacity-100"
       />
-      <pre className="max-w-full overflow-x-auto rounded-md border border-border bg-background p-3 font-mono text-[11px] leading-relaxed">
+      <pre className="max-w-full overflow-x-auto rounded-md border border-border bg-background p-3 font-mono text-xs leading-relaxed">
         {html === null ? (
           <code>{formatted}</code>
         ) : (
-          <code dangerouslySetInnerHTML={{ __html: html }} />
+          <code data-dbx-highlighted-code dangerouslySetInnerHTML={{ __html: html }} />
         )}
       </pre>
     </div>
@@ -404,11 +406,15 @@ export const JsonBlock = ({ json, className }: { json: string; className?: strin
   return (
     <pre
       className={cn(
-        "max-w-full overflow-auto whitespace-pre-wrap break-words font-mono text-[11px] leading-relaxed",
+        "max-w-full overflow-auto whitespace-pre-wrap break-words font-mono text-xs leading-relaxed",
         className,
       )}
     >
-      {html === null ? <code>{json}</code> : <code dangerouslySetInnerHTML={{ __html: html }} />}
+      {html === null ? (
+        <code>{json}</code>
+      ) : (
+        <code data-dbx-highlighted-code dangerouslySetInnerHTML={{ __html: html }} />
+      )}
     </pre>
   );
 };

package/src/react/mastra-chat.tsx

@@ -24,6 +24,7 @@ import {
   type MastraStreamResponse,
 } from "../support/mastra-client.ts";
 import {
+  defaultReasoningEffort,
   modelStorageKey,
   readStoredModel,
   storeSelectedModel,
@@ -368,11 +369,14 @@ export const useMastraChat = <TValues extends Record<string, unknown> = Record<s
   const activeModel = model || defaultModelId || undefined;
   const activeModelOption = models.find((option) => option.name === activeModel);
   const reasoningEfforts = activeModelOption?.reasoningEfforts ?? EMPTY_REASONING_EFFORTS;
+  const reasoningDefault = defaultReasoningEffort(reasoningEfforts);
+  const reasoningModelKey = `${activeModel ?? ""}:${reasoningEfforts.join(",")}`;
+  const initializedReasoningModelRef = useRef<string | undefined>(undefined);
   useEffect(() => {
-    if (reasoningEffort && !reasoningEfforts.includes(reasoningEffort)) {
-      setReasoningEffort(undefined);
-    }
-  }, [reasoningEffort, reasoningEfforts]);
+    if (initializedReasoningModelRef.current === reasoningModelKey) return;
+    initializedReasoningModelRef.current = reasoningModelKey;
+    setReasoningEffort(reasoningDefault);
+  }, [reasoningDefault, reasoningModelKey]);
   // Starter suggestions: an explicit `options.suggestions` always
   // wins (including `[]` to force none) and is rendered verbatim;
   // otherwise auto-source the agent's Genie space sample questions.
@@ -450,8 +454,42 @@ export const useMastraChat = <TValues extends Record<string, unknown> = Record<s
       });
       const runIdRef = { current: getSession(threadId).runId };
       try {
-        const stream = await open(controller.signal);
-        await processStream(threadId, stream, assistantId, runIdRef, controller.signal);
+        try {
+          const stream = await open(controller.signal);
+          await processStream(threadId, stream, assistantId, runIdRef, controller.signal);
+        } catch (initialError) {
+          const runId = runIdRef.current;
+          if (
+            controller.signal.aborted ||
+            getSession(threadId).runToken !== token ||
+            !mastraClient.backgroundTurns ||
+            !runId
+          ) {
+            throw initialError;
+          }
+          logger.warn("chat stream disconnected; reconnecting to background run", {
+            runId,
+            error: errorUtils.errorMessage(initialError),
+          });
+          const messagesBeforeReplay = getSession(threadId).messages;
+          writeMessages(
+            threadId,
+            messagesBeforeReplay.filter((message) => message.id !== assistantId),
+          );
+          const streamThreadId = threadId === DEFAULT_THREAD_SESSION_KEY ? undefined : threadId;
+          try {
+            const stream = await mastraClient.observeAgentStream({
+              agentId,
+              runId,
+              threadId: streamThreadId,
+              signal: controller.signal,
+            });
+            await processStream(threadId, stream, assistantId, runIdRef, controller.signal);
+          } catch (reconnectError) {
+            writeMessages(threadId, messagesBeforeReplay);
+            throw reconnectError;
+          }
+        }
         updateSession(threadId, (session) => {
           if (session.runToken !== token) return session;
           return {
@@ -500,7 +538,15 @@ export const useMastraChat = <TValues extends Record<string, unknown> = Record<s
         }));
       }
     },
-    [getSession, processStream, refreshThreadsSoon, updateSession],
+    [
+      agentId,
+      getSession,
+      mastraClient,
+      processStream,
+      refreshThreadsSoon,
+      updateSession,
+      writeMessages,
+    ],
   );
 
   const runStream = useCallback(
@@ -567,9 +613,26 @@ export const useMastraChat = <TValues extends Record<string, unknown> = Record<s
   const stop = useCallback(
     (threadId?: string) => {
       const key = threadId ?? activeKey;
+      const current = getSession(key);
+      if (!isSessionRunning(current)) return;
+      const streamThreadId = key === DEFAULT_THREAD_SESSION_KEY ? undefined : key;
+      if (mastraClient.backgroundTurns && current.runId) {
+        void mastraClient
+          .abortAgentRun({
+            agentId,
+            runId: current.runId,
+            threadId: streamThreadId,
+          })
+          .catch((error: unknown) => {
+            logger.warn("failed to abort background run", {
+              runId: current.runId,
+              error: errorUtils.errorMessage(error),
+            });
+          });
+      }
+      current.abortController?.abort();
       updateSession(key, (session) => {
         if (!isSessionRunning(session)) return session;
-        session.abortController?.abort();
         return {
           ...session,
           abortController: null,
@@ -579,7 +642,7 @@ export const useMastraChat = <TValues extends Record<string, unknown> = Record<s
         };
       });
     },
-    [activeKey, updateSession],
+    [activeKey, agentId, getSession, mastraClient, updateSession],
   );
 
   const handleApproval = useChatApprovals({

package/src/react/tool-pill.tsx

@@ -277,7 +277,7 @@ const MessageGroupBody = ({
   return (
     <div className="flex flex-col gap-1.5">
       {!omitQuestion && group.question && (
-        <div className="rounded border-l-2 border-primary/40 bg-background/40 px-2 py-1 text-[11px] italic leading-snug text-muted-foreground">
+        <div className="rounded border-l-2 border-primary/40 bg-background/40 px-2 py-1 text-xs italic leading-snug text-muted-foreground">
           {group.question}
         </div>
       )}
@@ -290,13 +290,11 @@ const MessageGroupBody = ({
         // for the same reason: code is the heaviest content here,
         // and most readers only want a glance.
         <Collapsible key={bucket.key} className="rounded border border-border/60 bg-background/40">
-          <CollapsibleTrigger className="group flex w-full items-center gap-1.5 px-2 py-1 text-left text-[11px] uppercase tracking-wide text-muted-foreground hover:text-foreground">
+          <CollapsibleTrigger className="group flex w-full items-center gap-1.5 px-2 py-1 text-left text-xs uppercase tracking-wide text-muted-foreground hover:text-foreground">
             <ChevronDownIcon className="size-3 shrink-0 transition-transform group-data-[state=closed]:-rotate-90" />
             <span>{queryBuckets.length > 1 ? `Query ${i + 1}` : "Query"}</span>
             {bucket.query?.title ? (
-              <span className="min-w-0 flex-1 truncate normal-case text-muted-foreground/70">
-                {bucket.query.title}
-              </span>
+              <span className="min-w-0 flex-1 truncate normal-case">{bucket.query.title}</span>
             ) : null}
           </CollapsibleTrigger>
           <CollapsibleContent>
@@ -387,7 +385,7 @@ export function formatRawToolPayload(value: unknown): string {
 
 const RawToolPayload = ({ label, value }: { label: "Request" | "Response"; value: unknown }) => (
   <Collapsible className="rounded border border-border/60 bg-background/40">
-    <CollapsibleTrigger className="group flex w-full items-center gap-1.5 px-2 py-1 text-left text-[11px] uppercase tracking-wide text-muted-foreground hover:text-foreground">
+    <CollapsibleTrigger className="group flex w-full items-center gap-1.5 px-2 py-1 text-left text-xs uppercase tracking-wide text-muted-foreground hover:text-foreground">
       <ChevronDownIcon className="size-3 shrink-0 transition-transform group-data-[state=closed]:-rotate-90" />
       <span>{label}</span>
     </CollapsibleTrigger>
@@ -458,7 +456,7 @@ const WebSearchProgressDetails = ({ event }: { event: ToolEvent }) => {
         return (
           <div
             key={group.key}
-            className="flex min-w-0 items-center gap-1.5 text-[11px] leading-4"
+            className="flex min-w-0 items-center gap-1.5 text-xs leading-4"
             aria-label={done ? [query, results].filter(Boolean).join(", ") : `Searching ${query}`}
           >
             {done ? (
@@ -605,7 +603,7 @@ export const ToolSessionPill = ({ events }: { events: ToolEvent[] }) => {
   return (
     <div
       className={cn(
-        "rounded-md border bg-background/30 px-2 py-1.5 transition-colors",
+        "rounded-md border bg-background/30 px-2 py-1 transition-colors",
         // Loud-but-not-jarring "in flight" treatment when any row
         // is running: primary border + soft ring. Failed sessions
         // (no longer running) pick up the destructive border;
@@ -634,7 +632,7 @@ export const ToolSessionPill = ({ events }: { events: ToolEvent[] }) => {
             )}
           >
             {verb}
-            {countSuffix && <span className="text-muted-foreground/70">{countSuffix}</span>}
+            {countSuffix}
           </span>
           {/*
            * Trailing "live" pip on the right edge while any tool is
@@ -650,7 +648,7 @@ export const ToolSessionPill = ({ events }: { events: ToolEvent[] }) => {
           )}
         </CollapsibleTrigger>
         <CollapsibleContent>
-          <div className="mt-2 flex flex-col gap-1.5">
+          <div className="mt-1.5 flex flex-col gap-1">
             {events.map((event) => (
               <ToolCallRow key={event.id} event={event} />
             ))}

package/src/support/_shiki-runtime.ts

@@ -14,12 +14,13 @@ import python from "shiki/langs/python.mjs";
 import sql from "shiki/langs/sql.mjs";
 import typescript from "shiki/langs/typescript.mjs";
 import yaml from "shiki/langs/yaml.mjs";
-import githubLight from "shiki/themes/github-light.mjs";
+import githubDarkHighContrast from "shiki/themes/github-dark-high-contrast.mjs";
+import githubLightHighContrast from "shiki/themes/github-light-high-contrast.mjs";
 
 /** Create the shared highlighter with the exact configured language set. */
 export function createHighlighter() {
   return createHighlighterCore({
-    themes: [githubLight],
+    themes: [githubLightHighContrast, githubDarkHighContrast],
     langs: [sql, python, typescript, javascript, json, bash, yaml, markdown],
     engine: createOnigurumaEngine(import("shiki/wasm")),
   });

package/src/support/mastra-client.ts

@@ -72,6 +72,8 @@ export class MastraPluginClient extends MastraClient {
    * mode) keeps the OBO-gated behavior.
    */
   readonly chatAlwaysAvailable: boolean;
+  /** Whether turns continue server-side after the initiating stream disconnects. */
+  readonly backgroundTurns: boolean;
 
   constructor(config: MastraClientConfig) {
     super({
@@ -83,8 +85,9 @@ export class MastraPluginClient extends MastraClient {
     this.basePath = config.basePath;
     this.defaultAgent = config.defaultAgent;
     this.agents = config.agents;
-    this.feedbackEnabled = config.feedbackEnabled;
-    this.chatAlwaysAvailable = config.chatAlwaysAvailable;
+    this.feedbackEnabled = config.feedbackEnabled ?? false;
+    this.chatAlwaysAvailable = config.chatAlwaysAvailable ?? false;
+    this.backgroundTurns = config.backgroundTurns ?? false;
   }
 
   /**
@@ -137,6 +140,33 @@ export class MastraPluginClient extends MastraClient {
     return { headers, stream };
   }
 
+  async #uiReconnect(options: {
+    path: string;
+    chatId: string;
+    routing: { threadId?: string; model?: string };
+    signal?: AbortSignal;
+  }): Promise<MastraStreamResponse> {
+    let headers = new Headers();
+    const origin = typeof window !== "undefined" ? window.location.origin : "http://localhost";
+    const transport = new DefaultChatTransport<UIMessage>({
+      api: new URL(options.path, origin).toString(),
+      credentials: "include",
+      headers: this.#routingHeaders(options.routing),
+      prepareReconnectToStreamRequest: ({ api }) => ({ api }),
+      fetch: async (input, init) => {
+        const response = await fetch(input, init);
+        headers = new Headers(response.headers);
+        return response;
+      },
+    });
+    const stream = await transport.reconnectToStream({
+      chatId: options.chatId,
+      abortSignal: options.signal,
+    });
+    if (!stream) throw new Error(`Run ${options.chatId} is no longer available`);
+    return { headers, stream };
+  }
+
   #requestClient(
     routing: { threadId?: string; model?: string },
     signal?: AbortSignal,
@@ -188,6 +218,45 @@ export class MastraPluginClient extends MastraClient {
     });
   }
 
+  /** Reconnect to the cached event stream for one background agent run. */
+  async observeAgentStream(params: {
+    agentId: string;
+    runId: string;
+    threadId?: string;
+    model?: string;
+    signal?: AbortSignal;
+  }): Promise<MastraStreamResponse> {
+    return this.#uiReconnect({
+      path: `${this.basePath}${routes.MASTRA_ROUTES.chat}/${encodeURIComponent(params.agentId)}${routes.MASTRA_ROUTES.runs}/${encodeURIComponent(params.runId)}`,
+      chatId: params.runId,
+      routing: params,
+      signal: params.signal,
+    });
+  }
+
+  /** Explicitly stop one background agent run on the server. */
+  async abortAgentRun(params: {
+    agentId: string;
+    runId: string;
+    threadId?: string;
+    signal?: AbortSignal;
+  }): Promise<boolean> {
+    const origin = typeof window !== "undefined" ? window.location.origin : "http://localhost";
+    const response = await this.#mutateJson(
+      new URL(
+        `${this.basePath}${routes.MASTRA_ROUTES.chat}/${encodeURIComponent(params.agentId)}${routes.MASTRA_ROUTES.runs}/${encodeURIComponent(params.runId)}${routes.MASTRA_ROUTES.abort}`,
+        origin,
+      ).toString(),
+      "POST",
+      wire.MastraRunAbortResponseSchema,
+      {
+        headers: this.#routingHeaders({ threadId: params.threadId }),
+        signal: params.signal,
+      },
+    );
+    return response.aborted;
+  }
+
   /**
    * Resume a suspended `requireApproval` tool through the same official chat
    * route using Mastra's native `resumeData` contract.

package/src/support/model-selection.ts

@@ -1,5 +1,18 @@
+import {
+  ReasoningEffort as ReasoningEffortValue,
+  type ReasoningEffort,
+} from "@dbx-tools/shared-model";
+
 type StorageLike = Pick<Storage, "getItem" | "setItem">;
 
+/** Default reasoning on for capable models while preserving an explicit off choice. */
+export const defaultReasoningEffort = (
+  efforts: readonly ReasoningEffort[],
+): ReasoningEffort | undefined =>
+  efforts.includes(ReasoningEffortValue.Medium)
+    ? ReasoningEffortValue.Medium
+    : efforts.find((effort) => effort !== ReasoningEffortValue.None);
+
 /** Namespace the last-selected model by plugin mount and agent. */
 export const modelStorageKey = (basePath: string, agentId: string): string =>
   `dbx-mastra-model:${basePath}:${agentId}`;

package/src/support/shiki-plugin.ts

@@ -27,13 +27,11 @@ const LANGUAGES = [
   "markdown",
 ] as const;
 
-/**
- * Single light theme. The demo has no dark-mode toggle, and a single
- * theme guarantees every token carries a concrete `color` (shiki's
- * dual-theme mode emits CSS-variable-only tokens that need extra CSS
- * wiring to paint).
- */
-const THEME: BundledTheme = "github-light";
+/** High-contrast theme pair used for every syntax surface. */
+const THEMES = [
+  "github-light-high-contrast",
+  "github-dark-high-contrast",
+] as const satisfies readonly [BundledTheme, BundledTheme];
 
 /** Languages we can tokenize, as a set for O(1) support checks. */
 const SUPPORTED = new Set<string>(LANGUAGES);
@@ -52,13 +50,40 @@ function loadHighlighter(): Promise<HighlighterCore> {
   return _loading;
 }
 
-/** Tokenize `code` with the active theme into Streamdown's result shape. */
+function splitThemeColor(
+  value: string | undefined,
+  darkVariable: "--shiki-dark" | "--shiki-dark-bg",
+): { light?: string; dark?: string } {
+  if (!value) return {};
+  const marker = `;${darkVariable}:`;
+  const index = value.indexOf(marker);
+  if (index < 0) return { light: value };
+  return {
+    light: value.slice(0, index),
+    dark: value.slice(index + marker.length),
+  };
+}
+
+/** Tokenize `code` with the active themes into Streamdown's result shape. */
 function highlightTokens(h: HighlighterCore, code: string, language: string) {
   const { tokens, fg, bg, rootStyle } = h.codeToTokens(code, {
     lang: language as BundledLanguage,
-    theme: THEME,
+    themes: { light: THEMES[0], dark: THEMES[1] },
   });
-  return { tokens, fg, bg, rootStyle };
+  const foreground = splitThemeColor(fg, "--shiki-dark");
+  const background = splitThemeColor(bg, "--shiki-dark-bg");
+  const darkStyles = [
+    foreground.dark ? `--shiki-dark:${foreground.dark}` : "",
+    background.dark ? `--shiki-dark-bg:${background.dark}` : "",
+  ]
+    .filter(Boolean)
+    .join(";");
+  return {
+    tokens,
+    fg: foreground.light,
+    bg: background.light,
+    rootStyle: [rootStyle, darkStyles].filter(Boolean).join(";"),
+  };
 }
 
 /** Escape HTML-significant characters (from the shared string utils). */
@@ -76,17 +101,16 @@ export async function highlightToHtml(code: string, language: string): Promise<s
   if (!SUPPORTED.has(language)) return escapeHtml(code);
   const h = await loadHighlighter();
   try {
-    const { tokens } = h.codeToTokens(code, {
-      lang: language as BundledLanguage,
-      theme: THEME,
-    });
+    const { tokens } = highlightTokens(h, code, language);
     return tokens
       .map((line) =>
         line
-          .map(
-            (token) =>
-              `<span style="color:${token.color ?? "inherit"}">${escapeHtml(token.content)}</span>`,
-          )
+          .map((token) => {
+            const light = token.htmlStyle?.color ?? token.color ?? "inherit";
+            const dark = token.htmlStyle?.["--shiki-dark"] ?? light;
+            const style = `--sdm-c:${light};--shiki-dark:${dark}`;
+            return `<span style="${escapeHtml(style)}">${escapeHtml(token.content)}</span>`;
+          })
           .join(""),
       )
       .join("\n");
@@ -109,7 +133,7 @@ export function createShikiPlugin(): CodeHighlighterPlugin {
     name: "shiki",
     type: "code-highlighter",
     getSupportedLanguages: () => [...LANGUAGES],
-    getThemes: () => [THEME, THEME],
+    getThemes: () => [...THEMES],
     supportsLanguage: (language) => isSupported(language),
     highlight: (options, callback) => {
       if (!isSupported(options.language)) return null;

package/src/support/tool-events.ts

@@ -4,7 +4,7 @@ import {
   TOOL_PROGRESS_PART_TYPE,
   ToolProgressPartDataSchema,
 } from "@dbx-tools/shared-mastra/wire";
-import { getToolOrDynamicToolName, isToolOrDynamicToolUIPart, type UIMessage } from "ai";
+import { getToolOrDynamicToolName, isToolUIPart, type UIMessage } from "ai";
 import type { ToolEvent } from "../react/types.ts";
 
 /** Collapse one reasoning event into a compact paragraph. */
@@ -32,7 +32,7 @@ export function genieReasoningText(events: ToolEvent[]): string {
       if (text) sections.push(text);
     }
   }
-  return sections.join("\n\n");
+  return sections.join("\n");
 }
 
 /** Drop markdown table rows so query samples never land in Thoughts. */
@@ -46,7 +46,7 @@ function stripMarkdownTables(text: string): string {
 
 /** Project native AI SDK tool parts from persisted messages onto pill state. */
 export function toolEventsFromParts(parts: UIMessage["parts"]): ToolEvent[] {
-  const events = parts.filter(isToolOrDynamicToolUIPart).map<ToolEvent>((part) => {
+  const events = parts.filter(isToolUIPart).map<ToolEvent>((part) => {
     const error = part.state === "output-error";
     const done = part.state === "output-available";
     return {
@@ -78,21 +78,33 @@ export function toolEventsFromParts(parts: UIMessage["parts"]): ToolEvent[] {
 export function mergeToolEvents(
   persisted: ToolEvent[],
   live: ToolEvent[] | undefined,
-  terminalError = false,
+  terminalState?: "done" | "error",
 ): ToolEvent[] {
   const merged = new Map(persisted.map((event) => [event.id, event]));
   for (const event of live ?? []) {
     const prior = merged.get(event.id);
-    merged.set(event.id, prior ? { ...prior, ...event } : event);
+    merged.set(
+      event.id,
+      prior
+        ? {
+            ...prior,
+            ...event,
+            status: prior.status === "running" ? event.status : prior.status,
+            ...(prior.output !== undefined ? { output: prior.output } : {}),
+          }
+        : event,
+    );
   }
   const events = [...merged.values()];
-  if (!terminalError) return events;
+  if (!terminalState) return events;
   return events.map((event) =>
     event.status === "running"
       ? {
           ...event,
-          status: "error",
-          output: { error: "The assistant stream ended before this tool returned." },
+          status: terminalState,
+          ...(terminalState === "error"
+            ? { output: { error: "The assistant stream ended before this tool returned." } }
+            : {}),
         }
       : event,
   );