@mastra/playground-ui 3.0.0-alpha.2 → 3.0.0-alpha.4

package/dist/components/assistant-ui/tool-fallback.d.ts

@@ -0,0 +1,3 @@
+import { ToolCallContentPartComponent } from '@assistant-ui/react';
+
+export declare const ToolFallback: ToolCallContentPartComponent;

package/dist/index.es.js

@@ -1,6 +1,6 @@
 import { jsx, jsxs, Fragment } from 'react/jsx-runtime';
 import { MessagePrimitive, ActionBarPrimitive, BranchPickerPrimitive, ThreadPrimitive, ComposerPrimitive, useExternalStoreRuntime, AssistantRuntimeProvider } from '@assistant-ui/react';
-import { CheckIcon, CopyIcon, ChevronLeftIcon, ChevronRightIcon, ArrowUp, Copy, Search, RefreshCcwIcon, ChevronRight, SortAsc, SortDesc, Braces, Clock1, ChevronDown, XIcon, Check, LoaderCircle, ChevronUpIcon, ChevronDownIcon, ExternalLinkIcon, X, Footprints, CircleCheck, CircleX, AlertCircleIcon, Plus, CalendarIcon, Loader2 } from 'lucide-react';
+import { CheckIcon, CopyIcon, ChevronUpIcon, ChevronDownIcon, ChevronLeftIcon, ChevronRightIcon, ArrowUp, Copy, Search, RefreshCcwIcon, ChevronRight, SortAsc, SortDesc, Braces, Clock1, ChevronDown, XIcon, Check, LoaderCircle, ExternalLinkIcon, X, Footprints, CircleCheck, CircleX, AlertCircleIcon, Plus, CalendarIcon, Loader2 } from 'lucide-react';
 import * as React from 'react';
 import React__default, { forwardRef, memo, useState, useRef, useEffect, createContext, useContext, useMemo, useCallback, Suspense, Fragment as Fragment$1 } from 'react';
 import { Slot } from '@radix-ui/react-slot';
@@ -3120,9 +3120,38 @@ const defaultComponents = unstable_memoizeMarkdownComponents({
   SyntaxHighlighter: SyntaxHighlighter$1
 });
 
-const AssistantMessage = ({ ToolFallback }) => {
-  return /* @__PURE__ */ jsxs(MessagePrimitive.Root, { className: "grid group grid-cols-[auto_auto_1fr] grid-rows-[auto_1fr] max-w-[var(--thread-max-width)] relative w-full py-4", children: [
-    /* @__PURE__ */ jsx("div", { className: "text-foreground max-w-[calc(var(--thread-max-width)*0.8)] sm:max-w-[70%] break-words leading-7 col-span-2 py-2 col-start-2 row-start-1 my-1.5", children: /* @__PURE__ */ jsx(MessagePrimitive.Content, { components: { Text: MarkdownText, tools: { Fallback: ToolFallback } } }) }),
+const ToolFallback$1 = ({ toolName, argsText, result }) => {
+  const [isCollapsed, setIsCollapsed] = useState(true);
+  return /* @__PURE__ */ jsxs("div", { className: "mb-2 flex w-full flex-col gap-3 rounded-lg border py-3 text-sm", children: [
+    /* @__PURE__ */ jsxs("div", { className: "flex items-center gap-2 px-4", children: [
+      /* @__PURE__ */ jsx(CheckIcon, { className: "size-4" }),
+      /* @__PURE__ */ jsxs("p", { className: "", children: [
+        "Used tool: ",
+        /* @__PURE__ */ jsx("b", { children: toolName })
+      ] }),
+      /* @__PURE__ */ jsx("div", { className: "flex-grow" }),
+      /* @__PURE__ */ jsx(Button, { onClick: () => setIsCollapsed(!isCollapsed), children: isCollapsed ? /* @__PURE__ */ jsx(ChevronUpIcon, {}) : /* @__PURE__ */ jsx(ChevronDownIcon, {}) })
+    ] }),
+    !isCollapsed && /* @__PURE__ */ jsxs("div", { className: "flex flex-col gap-2 border-t pt-2", children: [
+      /* @__PURE__ */ jsx("div", { className: "px-4", children: /* @__PURE__ */ jsx("pre", { className: "whitespace-pre-wrap", children: argsText }) }),
+      result !== void 0 && /* @__PURE__ */ jsxs("div", { className: "border-t border-dashed px-4 pt-2", children: [
+        /* @__PURE__ */ jsx("p", { className: "font-semibold", children: "Result:" }),
+        /* @__PURE__ */ jsx("pre", { className: "whitespace-pre-wrap", children: typeof result === "string" ? result : JSON.stringify(result, null, 2) })
+      ] })
+    ] })
+  ] });
+};
+
+const AssistantMessage = ({
+  ToolFallback: ToolFallbackCustom
+}) => {
+  return /* @__PURE__ */ jsxs(MessagePrimitive.Root, { className: "grid group grid-cols-[auto_auto_1fr] grid-rows-[auto_1fr] max-w-[var(--thread-max-width)] relative w-full", children: [
+    /* @__PURE__ */ jsx("div", { className: "text-foreground max-w-[calc(var(--thread-max-width)*0.8)] sm:max-w-[70%] break-words leading-7 col-span-2 py-2 col-start-2 row-start-1", children: /* @__PURE__ */ jsx(
+      MessagePrimitive.Content,
+      {
+        components: { Text: MarkdownText, tools: { Fallback: ToolFallbackCustom || ToolFallback$1 } }
+      }
+    ) }),
     /* @__PURE__ */ jsx(AssistantActionBar, {}),
     /* @__PURE__ */ jsx(BranchPicker, { className: "col-start-2 row-start-2 -ml-2 mr-2" })
   ] });
@@ -3169,7 +3198,7 @@ const UserMessage = () => {
       style: {
         placeItems: "end"
       },
-      className: "grid w-full py-4",
+      className: "grid w-full",
       children: /* @__PURE__ */ jsx(
         "div",
         {
@@ -3477,21 +3506,49 @@ function MastraRuntimeProvider({
     const hasNewInitialMessages = initialMessages && initialMessages?.length > messages?.length;
     if (messages.length === 0 || currentThreadId !== threadId || hasNewInitialMessages && currentThreadId === threadId) {
       if (initialMessages && threadId && memory) {
-        setMessages(initialMessages);
+        const convertedMessages = initialMessages?.map((message) => {
+          if (message?.toolInvocations?.length > 0) {
+            return {
+              ...message,
+              content: message.toolInvocations.map((toolInvocation) => ({
+                type: "tool-call",
+                toolCallId: toolInvocation?.toolCallId,
+                toolName: toolInvocation?.toolName,
+                args: toolInvocation?.args,
+                result: toolInvocation?.result
+              }))
+            };
+          }
+          return message;
+        }).filter(Boolean);
+        setMessages(convertedMessages);
         setCurrentThreadId(threadId);
       }
     }
-  }, [initialMessages, threadId, memory, messages]);
+  }, [initialMessages, threadId, memory]);
   const mastra = new MastraClient({
     baseUrl: baseUrl || ""
   });
+  const agent = mastra.getAgent(agentId);
   const onNew = async (message) => {
     if (message.content[0]?.type !== "text") throw new Error("Only text messages are supported");
     const input = message.content[0].text;
     setMessages((currentConversation) => [...currentConversation, { role: "user", content: input }]);
     setIsRunning(true);
     try {
-      const agent = mastra.getAgent(agentId);
+      let updater = function() {
+        setMessages((currentConversation) => {
+          const message2 = {
+            role: "assistant",
+            content: [{ type: "text", text: content }]
+          };
+          if (!assistantMessageAdded) {
+            assistantMessageAdded = true;
+            return [...currentConversation, message2];
+          }
+          return [...currentConversation.slice(0, -1), message2];
+        });
+      };
       const response = await agent.stream({
         messages: [
           {
@@ -3502,65 +3559,98 @@ function MastraRuntimeProvider({
         runId: agentId,
         ...memory ? { threadId, resourceId: agentId } : {}
       });
-      const reader = response.body?.getReader();
-      const decoder = new TextDecoder();
-      let buffer = "";
-      let assistantMessage = "";
-      let assistantMessageAdded = false;
-      let errorMessage = "";
-      if (!reader) {
-        throw new Error("No reader found");
+      if (!response.body) {
+        throw new Error("No response body");
       }
-      try {
-        while (true) {
-          const { done, value } = await reader.read();
-          if (done) break;
-          const chunk = decoder.decode(value);
-          buffer += chunk;
-          const matches = buffer.matchAll(/0:"((?:\\.|(?!").)*?)"/g);
-          const errorMatches = buffer.matchAll(/3:"((?:\\.|(?!").)*?)"/g);
-          if (errorMatches) {
-            for (const match of errorMatches) {
-              const content = match[1];
-              errorMessage += content;
-              setMessages((currentConversation) => [
-                ...currentConversation.slice(0, -1),
+      const parts = [];
+      let content = "";
+      let currentTextPart = null;
+      let assistantMessageAdded = false;
+      await response.processDataStream({
+        onTextPart(value) {
+          if (currentTextPart == null) {
+            currentTextPart = {
+              type: "text",
+              text: value
+            };
+            parts.push(currentTextPart);
+          } else {
+            currentTextPart.text += value;
+          }
+          content += value;
+          updater();
+        },
+        async onToolCallPart(value) {
+          setMessages((currentConversation) => {
+            const lastMessage = currentConversation[currentConversation.length - 1];
+            if (lastMessage && lastMessage.role === "assistant") {
+              const updatedMessage = {
+                ...lastMessage,
+                content: Array.isArray(lastMessage.content) ? [
+                  ...lastMessage.content,
+                  {
+                    type: "tool-call",
+                    toolCallId: value.toolCallId,
+                    toolName: value.toolName,
+                    args: value.args
+                  }
+                ] : [
+                  ...typeof lastMessage.content === "string" ? [{ type: "text", text: lastMessage.content }] : [],
+                  {
+                    type: "tool-call",
+                    toolCallId: value.toolCallId,
+                    toolName: value.toolName,
+                    args: value.args
+                  }
+                ]
+              };
+              return [...currentConversation.slice(0, -1), updatedMessage];
+            }
+            const newMessage = {
+              role: "assistant",
+              content: [
+                { type: "text", text: content },
                 {
-                  role: "assistant",
-                  content: [{ type: "text", text: errorMessage }],
-                  isError: true
+                  type: "tool-call",
+                  toolCallId: value.toolCallId,
+                  toolName: value.toolName,
+                  args: value.args
                 }
-              ]);
-            }
-          }
-          for (const match of matches) {
-            const content = match[1].replace(/\\"/g, '"').replace(/\\n/g, "\n");
-            assistantMessage += content;
-            setMessages((currentConversation) => {
-              const message2 = {
-                role: "assistant",
-                content: [{ type: "text", text: assistantMessage }]
+              ]
+            };
+            return [...currentConversation, newMessage];
+          });
+        },
+        async onToolResultPart(value) {
+          setMessages((currentConversation) => {
+            const lastMessage = currentConversation[currentConversation.length - 1];
+            if (lastMessage && lastMessage.role === "assistant" && Array.isArray(lastMessage.content)) {
+              const updatedContent = lastMessage.content.map((part) => {
+                if (typeof part === "object" && part.type === "tool-call" && part.toolCallId === value.toolCallId) {
+                  return {
+                    ...part,
+                    result: value.result
+                  };
+                }
+                return part;
+              });
+              const updatedMessage = {
+                ...lastMessage,
+                content: updatedContent
               };
-              const lastMessage = currentConversation[currentConversation.length - 1];
-              if (lastMessage.id) {
-                return currentConversation;
-              }
-              if (!assistantMessageAdded) {
-                assistantMessageAdded = true;
-                return [...currentConversation, message2];
-              }
-              return [...currentConversation.slice(0, -1), message2];
-            });
-          }
-          buffer = "";
+              return [...currentConversation.slice(0, -1), updatedMessage];
+            }
+            return currentConversation;
+          });
+        },
+        onErrorPart(error) {
+          throw new Error(error);
         }
-      } finally {
-        reader.releaseLock();
-        setIsRunning(false);
-        setTimeout(() => {
-          refreshThreadList?.();
-        }, 500);
-      }
+      });
+      setIsRunning(false);
+      setTimeout(() => {
+        refreshThreadList?.();
+      }, 500);
     } catch (error) {
       console.error("Error occurred in MastraRuntimeProvider", error);
       setIsRunning(false);
@@ -4862,6 +4952,7 @@ function usePolling({
 
 const useTraces = (componentName, baseUrl, isWorkflow = false) => {
   const [traces, setTraces] = useState([]);
+  const { setTraces: setTraceContextTraces } = useContext(TraceContext);
   const client = useMemo(
     () => new MastraClient({
       baseUrl: baseUrl || ""
@@ -4884,11 +4975,15 @@ const useTraces = (componentName, baseUrl, isWorkflow = false) => {
       throw error2;
     }
   }, [client, componentName, isWorkflow]);
-  const onSuccess = useCallback((newTraces) => {
-    if (newTraces.length > 0) {
-      setTraces(() => newTraces);
-    }
-  }, []);
+  const onSuccess = useCallback(
+    (newTraces) => {
+      if (newTraces.length > 0) {
+        setTraces(() => newTraces);
+        setTraceContextTraces(() => newTraces);
+      }
+    },
+    [setTraceContextTraces]
+  );
   const onError = useCallback((error2) => {
     toast.error(error2.message);
   }, []);