@dbx-tools/ui-mastra 0.9.131 → 0.9.136

package/src/react/chat-feedback.ts

@@ -1,23 +1,58 @@
 import { errorUtils, log } from "@dbx-tools/shared-core";
+import { genieModel } from "@dbx-tools/shared-genie";
+import type { FeedbackMode } from "@dbx-tools/shared-mastra/feedback";
 import type { UIMessage } from "ai";
 import { useCallback, useRef } from "react";
 import type { ThreadSessionUpdater } from "./chat-sessions.ts";
 import type { FeedbackSubmission, MessageFeedback } from "./types.ts";
 import type { MastraPluginClient } from "../support/mastra-client.ts";
+import { toolEventsFromParts } from "../support/tool-events.ts";
 
 const logger = log.logger("ui-mastra/chat");
 
 interface UseChatFeedbackOptions {
   activeKey: string;
   feedbackByMessage: Record<string, MessageFeedback>;
+  feedbackMode: FeedbackMode;
   mastraClient: MastraPluginClient;
   updateSession: ThreadSessionUpdater;
 }
 
-/** Submit trace-scoped MLflow feedback with optimistic thumbs state. */
+/** Genie message target recovered from one completed ask_genie tool result. */
+export interface GenieFeedbackTarget {
+  spaceId: string;
+  conversationId: string;
+  messageId: string;
+}
+
+/** Read the latest Genie target from an assistant message's tool results. */
+export function genieFeedbackTarget(message: UIMessage): GenieFeedbackTarget | undefined {
+  for (const event of [...toolEventsFromParts(message.parts)].reverse()) {
+    if (!event.toolName.startsWith("ask_genie")) continue;
+    const output =
+      typeof event.output === "object" && event.output !== null && "message" in event.output
+        ? event.output.message
+        : undefined;
+    const parsed = genieModel.GenieMessageSchema.safeParse(output);
+    if (!parsed.success) continue;
+    const genieMessage = parsed.data;
+    const messageId = genieMessage.message_id ?? genieMessage.id;
+    if (genieMessage.space_id && genieMessage.conversation_id && messageId) {
+      return {
+        spaceId: genieMessage.space_id,
+        conversationId: genieMessage.conversation_id,
+        messageId,
+      };
+    }
+  }
+  return undefined;
+}
+
+/** Submit feedback to the single destination selected by the server. */
 export function useChatFeedback({
   activeKey,
   feedbackByMessage,
+  feedbackMode,
   mastraClient,
   updateSession,
 }: UseChatFeedbackOptions) {
@@ -27,34 +62,55 @@ export function useChatFeedback({
   return useCallback(
     async (message: UIMessage, submission: FeedbackSubmission) => {
       const traceId = feedbackByMessageRef.current[message.id]?.traceId;
-      if (!traceId) return;
+      const target = feedbackMode === "genie" ? genieFeedbackTarget(message) : undefined;
+      if (feedbackMode === "none") return;
+
       if (submission.value) {
         updateSession(activeKey, (session) => ({
           ...session,
           feedbackByMessage: {
             ...session.feedbackByMessage,
-            [message.id]: { traceId, value: submission.value },
+            [message.id]: {
+              ...(traceId ? { traceId } : {}),
+              value: submission.value,
+            },
           },
         }));
       }
       try {
-        const result = await mastraClient.feedback({
-          traceId,
-          ...(submission.value !== undefined ? { value: submission.value === "up" } : {}),
-          ...(submission.comment ? { comment: submission.comment } : {}),
-        });
-        if (!result.ok) {
-          logger.warn("feedback not recorded (trace may still be exporting)", {
+        let result;
+        if (feedbackMode === "genie") {
+          if (!target || !submission.value || !submission.request) return;
+          result = await mastraClient.feedback({
+            mode: "genie",
+            ...target,
+            rating: submission.value === "up" ? "POSITIVE" : "NEGATIVE",
+            request: submission.request,
+            ...(submission.comment ? { comment: submission.comment } : {}),
+          });
+        } else {
+          if (!traceId) return;
+          result = await mastraClient.feedback({
+            mode: "mlflow",
             traceId,
+            ...(submission.value !== undefined ? { value: submission.value === "up" } : {}),
+            ...(submission.comment ? { comment: submission.comment } : {}),
+          });
+        }
+        if (!result.ok) {
+          logger.warn("feedback not recorded", {
+            feedbackMode,
+            ...(traceId ? { traceId } : {}),
           });
         }
       } catch (error) {
         logger.error("feedback error", {
-          traceId,
+          feedbackMode,
+          ...(traceId ? { traceId } : {}),
           error: errorUtils.toError(error),
         });
       }
     },
-    [activeKey, mastraClient, updateSession],
+    [activeKey, feedbackMode, mastraClient, updateSession],
   );
 }

package/src/react/chat-transcript.tsx

@@ -15,7 +15,9 @@ import { isToolUIPart } from "ai";
 import { ArrowDownIcon, MessageSquareIcon, RefreshCwIcon, TriangleAlertIcon } from "lucide-react";
 import React, { useCallback, useEffect, useLayoutEffect, useRef, useState } from "react";
 import { AssistantBubble, UserBubble } from "./bubbles.tsx";
+import { genieFeedbackTarget } from "./chat-feedback.ts";
 import type { ChatViewProps } from "./types.ts";
+import { messageText } from "../support/export.ts";
 
 const BOTTOM_THRESHOLD_PX = 24;
 const TOP_LOAD_MORE_THRESHOLD_PX = 120;
@@ -156,6 +158,7 @@ type ChatTranscriptProps = {
   pendingApprovalsByMessage: NonNullable<ChatViewProps["pendingApprovalsByMessage"]>;
   onExportMessage: ChatViewProps["onExportMessage"];
   feedbackByMessage: NonNullable<ChatViewProps["feedbackByMessage"]>;
+  feedbackMode: NonNullable<ChatViewProps["feedbackMode"]>;
   onFeedback: ChatViewProps["onFeedback"];
   mlflowExperimentUrl: ChatViewProps["mlflowExperimentUrl"];
 };
@@ -176,6 +179,7 @@ export const ChatTranscript = ({
   pendingApprovalsByMessage,
   onExportMessage,
   feedbackByMessage,
+  feedbackMode,
   onFeedback,
   mlflowExperimentUrl,
 }: ChatTranscriptProps) => {
@@ -229,6 +233,17 @@ export const ChatTranscript = ({
               const isLast = index === messages.length - 1;
               if (message.role === "assistant") {
                 const messageFeedback = feedbackByMessage[message.id];
+                const genieTarget =
+                  feedbackMode === "genie" ? genieFeedbackTarget(message) : undefined;
+                const canSubmitFeedback =
+                  feedbackMode === "mlflow"
+                    ? Boolean(messageFeedback?.traceId)
+                    : Boolean(genieTarget);
+                const requestMessage = [...messages.slice(0, index)]
+                  .reverse()
+                  .find((candidate) => candidate.role === "user");
+                const feedbackRequest = requestMessage ? messageText(requestMessage) : "";
+                const bubbleFeedbackMode = feedbackMode === "genie" ? "genie" : "mlflow";
                 return (
                   <AssistantBubble
                     key={message.id}
@@ -243,10 +258,12 @@ export const ChatTranscript = ({
                     {...(onExportMessage
                       ? { onExport: (format) => onExportMessage(message, format) }
                       : {})}
-                    {...(onFeedback && messageFeedback
+                    {...(onFeedback && canSubmitFeedback
                       ? {
                           onFeedback: (submission) => onFeedback(message, submission),
-                          ...(messageFeedback.value
+                          feedbackMode: bubbleFeedbackMode,
+                          ...(feedbackRequest ? { feedbackRequest } : {}),
+                          ...(messageFeedback?.value
                             ? { feedbackValue: messageFeedback.value }
                             : {}),
                         }

package/src/react/chat-view.tsx

@@ -64,6 +64,7 @@ export const ChatView = ({
   onExportConversation,
   onExportMessage,
   feedbackByMessage = {},
+  feedbackMode = "none",
   onFeedback,
   mlflowExperimentUrl,
 }: ChatViewProps) => {
@@ -113,6 +114,7 @@ export const ChatView = ({
         pendingApprovalsByMessage={pendingApprovalsByMessage}
         onExportMessage={onExportMessage}
         feedbackByMessage={feedbackByMessage}
+        feedbackMode={feedbackMode}
         onFeedback={onFeedback}
         mlflowExperimentUrl={mlflowExperimentUrl}
       />

package/src/react/feedback-controls.tsx

@@ -25,15 +25,22 @@ import type { FeedbackSubmission, FeedbackValue } from "./types.ts";
 
 /** Render per-message rating and comment controls for an MLflow-backed trace. */
 export const FeedbackControls = ({
+  mode = "mlflow",
+  request,
   value,
   onSubmit,
 }: {
+  /** Server-selected feedback destination. */
+  mode?: "mlflow" | "genie";
+  /** Original user request shown in Genie feedback. */
+  request?: string;
   /** Last thumbs the user chose, so the active button stays highlighted. */
   value?: FeedbackValue;
   onSubmit: (submission: FeedbackSubmission) => void | Promise<void>;
 }) => {
   const [open, setOpen] = useState(false);
   const [comment, setComment] = useState("");
+  const [rating, setRating] = useState<FeedbackValue | undefined>(value);
   const [sending, setSending] = useState(false);
 
   const submitComment = async () => {
@@ -49,6 +56,103 @@ export const FeedbackControls = ({
     }
   };
 
+  const submitGenieFeedback = async () => {
+    if (!rating || !request?.trim() || sending) return;
+    setSending(true);
+    try {
+      await onSubmit({
+        value: rating,
+        request: request.trim(),
+        ...(comment.trim() ? { comment: comment.trim() } : {}),
+      });
+      setOpen(false);
+    } finally {
+      setSending(false);
+    }
+  };
+
+  if (mode === "genie") {
+    return (
+      <Dialog open={open} onOpenChange={setOpen}>
+        <Tooltip>
+          <TooltipTrigger asChild>
+            <DialogTrigger asChild>
+              <Button
+                type="button"
+                size="icon"
+                variant="ghost"
+                className="size-7"
+                aria-label="Give feedback"
+              >
+                <MessageSquareTextIcon className="size-3" />
+              </Button>
+            </DialogTrigger>
+          </TooltipTrigger>
+          <TooltipContent>Give feedback</TooltipContent>
+        </Tooltip>
+        <DialogContent className="sm:max-w-md">
+          <DialogHeader>
+            <DialogTitle>Feedback</DialogTitle>
+            <DialogDescription>Send feedback directly to this Genie response.</DialogDescription>
+          </DialogHeader>
+          <div className="space-y-4">
+            <fieldset className="space-y-2">
+              <legend className="text-sm font-medium">Rating</legend>
+              <div className="flex gap-2">
+                <Button
+                  type="button"
+                  variant={rating === "up" ? "default" : "outline"}
+                  aria-pressed={rating === "up"}
+                  onClick={() => setRating("up")}
+                >
+                  <ThumbsUpIcon className="size-4" />
+                  Good
+                </Button>
+                <Button
+                  type="button"
+                  variant={rating === "down" ? "destructive" : "outline"}
+                  aria-pressed={rating === "down"}
+                  onClick={() => setRating("down")}
+                >
+                  <ThumbsDownIcon className="size-4" />
+                  Bad
+                </Button>
+              </div>
+            </fieldset>
+            <div className="space-y-1.5">
+              <p className="text-sm font-medium">Request</p>
+              <p className="rounded-md border bg-muted/40 px-3 py-2 text-sm text-foreground">
+                {request}
+              </p>
+            </div>
+            <label className="block space-y-1.5 text-sm font-medium">
+              <span>Comment</span>
+              <Textarea
+                value={comment}
+                maxLength={5_000}
+                onChange={(event) => setComment(event.target.value)}
+                placeholder="Optional details"
+                rows={4}
+              />
+            </label>
+          </div>
+          <DialogFooter>
+            <Button type="button" variant="ghost" onClick={() => setOpen(false)}>
+              Cancel
+            </Button>
+            <Button
+              type="button"
+              disabled={!rating || !request?.trim() || sending}
+              onClick={() => void submitGenieFeedback()}
+            >
+              {sending ? "Sending..." : "Send feedback"}
+            </Button>
+          </DialogFooter>
+        </DialogContent>
+      </Dialog>
+    );
+  }
+
   return (
     <>
       <Tooltip>

package/src/react/mastra-chat.tsx

@@ -3,6 +3,7 @@ import { isStaleMastraResumeError } from "@dbx-tools/shared-mastra/resume";
 import type { ReasoningEffort } from "@dbx-tools/shared-model";
 import { getAuthStatus } from "@dbx-tools/ui/auth/react";
 import { useBrand } from "@dbx-tools/ui/branding/react";
+import type { RequestContext } from "@mastra/client-js";
 import type { FileUIPart, UIMessage } from "ai";
 import { type ReactNode, useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { toChronologicalUiMessages } from "./_history-messages.ts";
@@ -42,7 +43,6 @@ import {
 } from "../support/model-selection.ts";
 import {
   snapshotRequestContext,
-  type MastraRequestContextSnapshot,
   type MastraRequestContextSource,
 } from "../support/request-context.ts";
 import { threadStorageKey } from "../support/thread-persistence.ts";
@@ -238,14 +238,9 @@ export interface UseMastraChatOptions<
    */
   enableClear?: boolean;
   /**
-   * Surface per-message feedback controls (thumbs up/down + a comment
-   * modal) that log to MLflow as trace assessments.
-   *
-   * Defaults to auto: controls render when the server reports MLflow
-   * feedback is available. Pass `false` to hide them or `true` to request
-   * them explicitly. Regardless of this flag, controls only render when the
-   * server reports MLflow logging is enabled (`clientConfig.feedbackEnabled`)
-   * and the turn produced a trace id.
+   * Surface per-message feedback controls for the destination selected by the
+   * server. Genie Agent Mode posts Rating and Comment to Genie; traced agent
+   * mode posts thumbs/comment assessments to MLflow.
    */
   enableFeedback?: boolean;
   /**
@@ -329,8 +324,9 @@ export const useMastraChat = <TValues extends Record<string, unknown> = Record<s
   // Feedback defaults to automatic server capability detection. An explicit
   // false hides it; true keeps it requested but cannot bypass the server gate.
   const enableFeedback = options.enableFeedback !== false;
-  const feedbackAvailable = enableFeedback && mastraClient.feedbackEnabled;
-  const { url: mlflowExperimentUrl } = useMastraMlflowExperiment(feedbackAvailable);
+  const feedbackMode = enableFeedback ? mastraClient.feedbackMode : "none";
+  const feedbackAvailable = feedbackMode !== "none";
+  const { url: mlflowExperimentUrl } = useMastraMlflowExperiment(feedbackMode === "mlflow");
   const [resourceScope, setResourceScope] = useState<string | null>(null);
   useEffect(() => {
     let cancelled = false;
@@ -577,6 +573,7 @@ export const useMastraChat = <TValues extends Record<string, unknown> = Record<s
   const submitFeedback = useChatFeedback({
     activeKey,
     feedbackByMessage: activeSession.feedbackByMessage,
+    feedbackMode,
     mastraClient,
     updateSession,
   });
@@ -758,7 +755,7 @@ export const useMastraChat = <TValues extends Record<string, unknown> = Record<s
     (
       threadId: string,
       messages: UIMessage | UIMessage[],
-      requestContext?: MastraRequestContextSnapshot,
+      requestContext?: RequestContext,
     ) => {
       const turnMessages = Array.isArray(messages) ? messages : [messages];
       const userMessage = lastUserMessage(turnMessages);
@@ -1493,12 +1490,11 @@ export const useMastraChat = <TValues extends Record<string, unknown> = Record<s
           onExportMessage: exportMessage,
         }
       : {}),
-    // Feedback: only expose the state + handler (which light up the
-    // per-bubble thumbs / comment controls in ChatView) when feedback
-    // is enabled AND the server can log to MLflow. `feedbackByMessage`
-    // still gates per-message on a captured trace id.
+    // Feedback: expose one server-selected destination. MLflow messages still
+    // require trace state; Genie messages derive their target from ask_genie.
     ...(feedbackAvailable
       ? {
+          feedbackMode,
           feedbackByMessage: activeSession.feedbackByMessage,
           onFeedback: submitFeedback,
           ...(mlflowExperimentUrl ? { mlflowExperimentUrl } : {}),

package/src/react/thread-sidebar.tsx

@@ -152,7 +152,9 @@ export const ThreadSidebar = ({
               New chat
             </Button>
           )}
-          {headerActions}
+          {headerActions ? (
+            <div className={cn("shrink-0", !onNew && "ml-auto")}>{headerActions}</div>
+          ) : null}
           {onHide && (
             <Tooltip>
               <TooltipTrigger asChild>
@@ -162,7 +164,7 @@ export const ThreadSidebar = ({
                   size="icon"
                   onClick={onHide}
                   aria-label="Hide conversations"
-                  className={cn("size-8 shrink-0", !onNew && "ml-auto")}
+                  className={cn("size-8 shrink-0", !onNew && !headerActions && "ml-auto")}
                 >
                   <HideIcon className="size-4" />
                 </Button>

package/src/react/types.ts

@@ -70,6 +70,7 @@ export type FeedbackValue = "up" | "down";
  */
 export type FeedbackSubmission = {
   value?: FeedbackValue;
+  request?: string;
   comment?: string;
 };
 
@@ -81,7 +82,7 @@ export type FeedbackSubmission = {
  * active thumb stays highlighted.
  */
 export type MessageFeedback = {
-  traceId: string;
+  traceId?: string;
   value?: FeedbackValue;
 };
 
@@ -364,6 +365,8 @@ export type ChatViewProps = {
    * response header when MLflow logging is enabled.
    */
   feedbackByMessage?: Record<string, MessageFeedback>;
+  /** Resolved feedback destination selected by the server. */
+  feedbackMode?: "none" | "mlflow" | "genie";
   /**
    * Submit thumbs / comment feedback for an assistant message. When
    * provided (and the message has a {@link feedbackByMessage} entry),

package/src/support/export.ts

@@ -217,7 +217,7 @@ function splitSegments(text: string): Segment[] {
 }
 
 /** Concatenate a message's `text` parts (matches the on-screen bubbles). */
-function messageText(message: UIMessage): string {
+export function messageText(message: UIMessage): string {
   return message.parts
     .filter((p): p is { type: "text"; text: string } => p.type === "text")
     .map((p) => p.text)

package/src/support/genie-client.ts

@@ -3,17 +3,19 @@ import type {
   ListAgentSuspendedRunsResponse,
   ListMemoryThreadMessagesResponse,
   ListMemoryThreadsResponse,
+  RequestContext,
 } from "@mastra/client-js";
 import { validateUIMessages, type UIMessage, type UIMessageChunk } from "ai";
 import { MastraPluginClient, type MastraStreamResponse } from "./mastra-client.ts";
 
 /**
  * {@link MastraPluginClient} that talks to one Genie space directly through
- * the Mastra plugin's `/genie/:alias` routes instead of an agent and Mastra
- * Memory. Genie owns all state: its conversations are the threads (the
- * conversation id is the thread id), its messages are the history, and each
- * turn streams the same AI SDK chunks as an agent's `ask_genie` call, so the
- * Mastra chat renders it unchanged.
+ * the Mastra plugin's `/genie/:alias` routes instead of an outer agent. Genie
+ * owns conversation state while the server may add Mastra request context and
+ * resource-scoped memory. Its conversations are the threads (the conversation
+ * id is the thread id), its messages are the history, and each turn streams
+ * the same AI SDK chunks as an agent's `ask_genie` call, so the Mastra chat
+ * renders it unchanged.
  *
  * A brand-new thread has no Genie conversation until its first turn; the
  * stream announces the conversation id, which {@link MastraStreamResponse.threadId}
@@ -23,7 +25,7 @@ export class GenieChatClient extends MastraPluginClient {
   /** Genie space alias every route is scoped to. */
   readonly alias: string;
   override readonly backgroundTurns = false;
-  override readonly feedbackEnabled = false;
+  override readonly feedbackMode = "none";
   override readonly threadRename = false;
   override readonly attachments = false;
   /** Thread ids known to be Genie conversations (listed, loaded, or announced). */
@@ -54,6 +56,7 @@ export class GenieChatClient extends MastraPluginClient {
     messages: UIMessage[];
     runId: string;
     threadId?: string;
+    requestContext?: RequestContext;
     signal?: AbortSignal;
   }): Promise<MastraStreamResponse> {
     const userMessage = [...params.messages].reverse().find((message) => message.role === "user");
@@ -67,7 +70,11 @@ export class GenieChatClient extends MastraPluginClient {
       path: this.#url(routes.MASTRA_ROUTES.chat),
       chatId: params.threadId ?? params.runId,
       messages: userMessage ? [userMessage] : [],
-      body: { question, ...(conversationId ? { conversationId } : {}) },
+      body: {
+        question,
+        ...(conversationId ? { conversationId } : {}),
+        ...(params.requestContext ? { requestContext: params.requestContext } : {}),
+      },
       routing: {},
       signal: params.signal,
     });

package/src/support/index.ts

@@ -26,7 +26,7 @@ export { GenieChatClient } from "./genie-client.ts";
 export { MastraChatRequestError, MastraPluginClient, useMastraConfig, useMastraClient, useMastraMlflowExperiment, useMastraModels, useMastraDefaultModel, useMastraSuggestions, useMastraThreads, useChartFetch, useStatementFetch, useVisualizationFetch } from "./mastra-client.ts";
 export type { MastraStreamResponse, MastraMemoryThread, ByIdFetchState } from "./mastra-client.ts";
 export { defaultReasoningEffort, modelStorageKey, readStoredModel, storeSelectedModel } from "./model-selection.ts";
-export type { MastraRequestContextInput, MastraRequestContextSource, MastraRequestContextSnapshot } from "./request-context.ts";
+export type { MastraRequestContextInput, MastraRequestContextSource } from "./request-context.ts";
 export { DEFAULT_THREAD_SESSION_KEY } from "./thread-sessions.ts";
 export type { SessionQueuedSteer, SessionRun, ThreadSession } from "./thread-sessions.ts";
 export { THREAD_TAB_SEED_MAX } from "./thread-tabs.ts";

package/src/support/mastra-client.ts

@@ -20,10 +20,10 @@ import {
   type ListAgentSuspendedRunsResponse,
   type ListMemoryThreadMessagesResponse,
   type ListMemoryThreadsResponse,
+  type RequestContext,
 } from "@mastra/client-js";
 import { DefaultChatTransport, type UIMessage, type UIMessageChunk } from "ai";
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
-import type { MastraRequestContextSnapshot } from "./request-context.ts";
 
 type NativeAgentClient = ReturnType<MastraClient["getAgent"]>;
 
@@ -79,13 +79,8 @@ export class MastraPluginClient extends MastraClient {
   readonly defaultAgent: string;
   /** Registered agent ids, surfaced for pickers. */
   readonly agents: readonly string[];
-  /**
-   * Whether the server can log feedback to MLflow. When `false`, the
-   * chat UI hides thumbs / comment controls and {@link feedback} would
-   * be rejected server-side, so callers gate on this before offering
-   * feedback affordances.
-   */
-  readonly feedbackEnabled: boolean;
+  /** Resolved feedback destination published by the server. */
+  readonly feedbackMode: feedback.FeedbackMode;
   /**
    * Whether the agents run their Databricks calls as the app service principal
    * rather than on-behalf-of the signed-in user. When `true` the chat works for
@@ -111,7 +106,7 @@ export class MastraPluginClient extends MastraClient {
     this.basePath = config.basePath;
     this.defaultAgent = config.defaultAgent;
     this.agents = config.agents;
-    this.feedbackEnabled = config.feedbackEnabled ?? false;
+    this.feedbackMode = config.feedbackMode ?? "none";
     this.chatAlwaysAvailable = config.chatAlwaysAvailable ?? false;
     this.backgroundTurns = config.backgroundTurns ?? false;
   }
@@ -227,7 +222,7 @@ export class MastraPluginClient extends MastraClient {
     threadId?: string;
     model?: string;
     reasoningEffort?: ReasoningEffort;
-    requestContext?: MastraRequestContextSnapshot;
+    requestContext?: RequestContext;
     signal?: AbortSignal;
   }): Promise<MastraStreamResponse> {
     return this.uiStream({
@@ -295,7 +290,7 @@ export class MastraPluginClient extends MastraClient {
       runId: string;
       toolCallId: string;
       threadId?: string;
-      requestContext?: MastraRequestContextSnapshot;
+      requestContext?: RequestContext;
       signal?: AbortSignal;
     },
   ): Promise<MastraStreamResponse> {
@@ -312,7 +307,7 @@ export class MastraPluginClient extends MastraClient {
       toolCallId: string;
       threadId?: string;
       reason?: string;
-      requestContext?: MastraRequestContextSnapshot;
+      requestContext?: RequestContext;
       signal?: AbortSignal;
     },
   ): Promise<MastraStreamResponse> {
@@ -327,7 +322,7 @@ export class MastraPluginClient extends MastraClient {
       toolCallId: string;
       threadId?: string;
       reason?: string;
-      requestContext?: MastraRequestContextSnapshot;
+      requestContext?: RequestContext;
       signal?: AbortSignal;
     },
   ): Promise<MastraStreamResponse> {

package/src/support/request-context.ts

@@ -1,5 +1,4 @@
-import { object, type SerializableValue } from "@dbx-tools/shared-core";
-import type { RequestContext } from "@mastra/client-js";
+import { RequestContext } from "@mastra/client-js";
 
 /** Typed plain values or Mastra's native typed request-context container. */
 export type MastraRequestContextInput<TValues extends Record<string, unknown>> =
@@ -9,22 +8,18 @@ export type MastraRequestContextInput<TValues extends Record<string, unknown>> =
 export type MastraRequestContextSource<TValues extends Record<string, unknown>> =
   MastraRequestContextInput<TValues> | (() => MastraRequestContextInput<TValues> | undefined);
 
-/** JSON-safe immutable snapshot retained for one run and its continuation. */
-export type MastraRequestContextSnapshot = Record<string, SerializableValue>;
-
 /**
- * Materialize and validate one run's application context. Infrastructure keys
- * remain server-owned; this only snapshots caller-provided JSON values.
+ * Materialize one run's application context through Mastra's native container
+ * and serialization policy.
  */
 export function snapshotRequestContext<TValues extends Record<string, unknown>>(
   source: MastraRequestContextSource<TValues> | undefined,
-): MastraRequestContextSnapshot | undefined {
+): RequestContext | undefined {
   const input = typeof source === "function" ? source() : source;
   if (input === undefined) return undefined;
-  const candidate =
-    "toJSON" in input && typeof input.toJSON === "function" ? input.toJSON() : input;
-  if (!object.isRecord(candidate) || !object.isSerializableValue(candidate)) {
-    throw new TypeError("Mastra request context must be a JSON-serializable record");
-  }
-  return structuredClone(candidate) as MastraRequestContextSnapshot;
+  const context =
+    input instanceof RequestContext
+      ? input
+      : new RequestContext<unknown>(Object.entries(input));
+  return new RequestContext<unknown>(Object.entries(structuredClone(context.toJSON())));
 }