@dbx-tools/ui-mastra 0.1.9

package/src/support/mastra-stream.ts

@@ -0,0 +1,66 @@
+/** One chunk from a Mastra agent SSE stream (`data: { type, payload, ... }`). */
+export interface MastraStreamChunk {
+  type: string;
+  payload?: unknown;
+  runId?: string;
+}
+
+/** Response from agent streaming endpoints with {@link processMastraStream}. */
+export type MastraStreamResponse = Response & {
+  processDataStream: (options: {
+    onChunk: (chunk: MastraStreamChunk) => void | Promise<void>;
+  }) => Promise<void>;
+};
+
+/**
+ * Parse Mastra agent SSE (`data: …` lines) and invoke `onChunk` per event.
+ * Mirrors `@mastra/client-js`'s internal reader without the
+ * `processChatResponse_vNext` side channel that throws when a resumed
+ * `approve-tool-call` stream emits `tool-result` without a new `tool-call`.
+ */
+export async function processMastraStream(options: {
+  stream: ReadableStream<Uint8Array>;
+  onChunk: (chunk: MastraStreamChunk) => void | Promise<void>;
+  signal?: AbortSignal;
+}): Promise<void> {
+  const reader = options.stream.getReader();
+  const decoder = new TextDecoder();
+  let buffer = "";
+  const abort = () => void reader.cancel();
+  if (options.signal?.aborted) abort();
+  else options.signal?.addEventListener("abort", abort, { once: true });
+  try {
+    for (;;) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      buffer += decoder.decode(value, { stream: true });
+      const lines = buffer.split("\n\n");
+      buffer = lines.pop() ?? "";
+      for (const line of lines) {
+        if (!line.startsWith("data: ")) continue;
+        const data = line.slice(6);
+        if (data === "[DONE]") return;
+        let json: MastraStreamChunk;
+        try {
+          json = JSON.parse(data) as MastraStreamChunk;
+        } catch {
+          continue;
+        }
+        if (json) await options.onChunk(json);
+      }
+    }
+  } finally {
+    options.signal?.removeEventListener("abort", abort);
+    reader.releaseLock();
+  }
+}
+
+/** Attach {@link processMastraStream} to a fetch `Response`. */
+export function asMastraStreamResponse(response: Response): MastraStreamResponse {
+  const streamResponse = response as MastraStreamResponse;
+  streamResponse.processDataStream = async ({ onChunk }) => {
+    if (!response.body) throw new Error("No response body");
+    await processMastraStream({ stream: response.body, onChunk });
+  };
+  return streamResponse;
+}

package/src/support/shiki-plugin.ts

@@ -0,0 +1,134 @@
+import { string } from "@dbx-tools/shared-core";
+import {
+  createHighlighter,
+  type BundledLanguage,
+  type BundledTheme,
+  type Highlighter,
+} from "shiki";
+import type { CodeHighlighterPlugin } from "streamdown";
+
+// Streamdown 2.x ships syntax highlighting as an opt-in plugin and has
+// no shiki dependency of its own; without a `code` plugin every fenced
+// block renders as uncolored plaintext. This module provides a small
+// shiki-backed highlighter to wire in via `plugins={{ code }}`.
+
+/**
+ * Languages we highlight in the chat. SQL is the primary one (Genie
+ * query previews); the rest cover code the assistant might emit. Kept
+ * to a curated set so shiki only bundles these grammars.
+ */
+const LANGUAGES: BundledLanguage[] = [
+  "sql",
+  "python",
+  "typescript",
+  "javascript",
+  "json",
+  "bash",
+  "yaml",
+  "markdown",
+];
+
+/**
+ * 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";
+
+/** Languages we can tokenize, as a set for O(1) support checks. */
+const SUPPORTED = new Set<BundledLanguage>(LANGUAGES);
+
+let _highlighter: Highlighter | null = null;
+let _loading: Promise<Highlighter> | null = null;
+
+/** Lazily create the shared shiki highlighter (singleton, loaded once). */
+function loadHighlighter(): Promise<Highlighter> {
+  _loading ??= createHighlighter({ themes: [THEME], langs: LANGUAGES }).then((h) => {
+    _highlighter = h;
+    return h;
+  });
+  return _loading;
+}
+
+/** Tokenize `code` with the active theme into Streamdown's result shape. */
+function highlightTokens(h: Highlighter, code: string, language: BundledLanguage) {
+  const { tokens, fg, bg, rootStyle } = h.codeToTokens(code, {
+    lang: language,
+    theme: THEME,
+  });
+  return { tokens, fg, bg, rootStyle };
+}
+
+/** Escape HTML-significant characters (from the shared string utils). */
+const escapeHtml = string.escapeHtml;
+
+/**
+ * Highlight `code` into minimal inline HTML: one colored `<span>` per
+ * token, lines joined by real newlines, with no line-number gutter or
+ * per-line wrapper elements. Meant to drop straight into a
+ * `<pre><code>` so the rendered text stays cleanly selectable and
+ * copyable. Falls back to plain escaped text when the language isn't
+ * supported or shiki fails to parse the snippet.
+ */
+export async function highlightToHtml(code: string, language: string): Promise<string> {
+  if (!SUPPORTED.has(language as BundledLanguage)) return escapeHtml(code);
+  const h = await loadHighlighter();
+  try {
+    const { tokens } = h.codeToTokens(code, {
+      lang: language as BundledLanguage,
+      theme: THEME,
+    });
+    return tokens
+      .map((line) =>
+        line
+          .map(
+            (token) =>
+              `<span style="color:${token.color ?? "inherit"}">${escapeHtml(token.content)}</span>`,
+          )
+          .join(""),
+      )
+      .join("\n");
+  } catch {
+    return escapeHtml(code);
+  }
+}
+
+/**
+ * shiki-backed code highlighter plugin for Streamdown. The highlighter
+ * loads asynchronously: the first call for any block returns `null`
+ * and resolves through the `callback` once shiki is ready, after which
+ * results are synchronous. Unsupported languages return `null` so the
+ * block stays plaintext rather than throwing.
+ */
+export function createShikiPlugin(): CodeHighlighterPlugin {
+  const isSupported = (language: string): language is BundledLanguage =>
+    SUPPORTED.has(language as BundledLanguage);
+
+  return {
+    name: "shiki",
+    type: "code-highlighter",
+    getSupportedLanguages: () => LANGUAGES,
+    getThemes: () => [THEME, THEME],
+    supportsLanguage: (language) => isSupported(language),
+    highlight: (options, callback) => {
+      if (!isSupported(options.language)) return null;
+      const language = options.language;
+      if (_highlighter) {
+        try {
+          return highlightTokens(_highlighter, options.code, language);
+        } catch {
+          return null;
+        }
+      }
+      void loadHighlighter().then((h) => {
+        try {
+          callback?.(highlightTokens(h, options.code, language));
+        } catch {
+          // Parse failure for this block - leave it as plaintext.
+        }
+      });
+      return null;
+    },
+  };
+}

package/src/support/thread-sessions.ts

@@ -0,0 +1,55 @@
+import type { UIMessage } from "ai";
+import type {
+  ChatStatus,
+  MessageFeedback,
+  PendingApproval,
+  ToolEvent,
+} from "../react/types";
+
+/** Session-scoped transcript + stream state for one conversation thread. */
+export type ThreadSession = {
+  messages: UIMessage[];
+  status: ChatStatus;
+  error: Error | null;
+  toolEventsByMessage: Record<string, ToolEvent[]>;
+  pendingApprovalsByMessage: Record<string, PendingApproval[]>;
+  feedbackByMessage: Record<string, MessageFeedback>;
+  abortController: AbortController | null;
+  runToken: number;
+  assistantId: string | null;
+  runId: string | null;
+  historyLoaded: boolean;
+  hasMoreHistory: boolean;
+  historyPage: number;
+  lastUserText: string | null;
+};
+
+/** Map key for the classic single-thread chat (no explicit thread id). */
+export const DEFAULT_THREAD_SESSION_KEY = "__session__";
+
+export function createThreadSession(): ThreadSession {
+  return {
+    messages: [],
+    status: "ready",
+    error: null,
+    toolEventsByMessage: {},
+    pendingApprovalsByMessage: {},
+    feedbackByMessage: {},
+    abortController: null,
+    runToken: 0,
+    assistantId: null,
+    runId: null,
+    historyLoaded: false,
+    hasMoreHistory: false,
+    historyPage: 0,
+    lastUserText: null,
+  };
+}
+
+export function isSessionRunning(session: ThreadSession): boolean {
+  return session.status === "submitted" || session.status === "streaming";
+}
+
+export function sessionKey(activeThreadId: string | undefined): string {
+  return activeThreadId ?? DEFAULT_THREAD_SESSION_KEY;
+}

package/tsconfig.json

@@ -0,0 +1,43 @@
+// ~~ Generated by projen. To modify, edit .projenrc.js and run "pnpm exec projen".
+{
+  "compilerOptions": {
+    "rootDir": "src",
+    "outDir": "lib",
+    "alwaysStrict": true,
+    "declaration": true,
+    "esModuleInterop": true,
+    "experimentalDecorators": true,
+    "inlineSourceMap": true,
+    "inlineSources": true,
+    "lib": [
+      "ES2022",
+      "DOM",
+      "DOM.Iterable"
+    ],
+    "module": "ESNext",
+    "noEmitOnError": false,
+    "noFallthroughCasesInSwitch": true,
+    "noImplicitAny": true,
+    "noImplicitReturns": true,
+    "noImplicitThis": true,
+    "noUnusedLocals": true,
+    "noUnusedParameters": true,
+    "resolveJsonModule": true,
+    "strict": true,
+    "strictNullChecks": true,
+    "strictPropertyInitialization": true,
+    "stripInternal": true,
+    "target": "ES2022",
+    "types": [],
+    "moduleResolution": "bundler",
+    "skipLibCheck": true,
+    "jsx": "react-jsx"
+  },
+  "include": [
+    "src/**/*.ts",
+    "src/**/*.tsx"
+  ],
+  "exclude": [
+    "node_modules"
+  ]
+}