@marimo-team/frontend 0.15.5 → 0.16.0

package/src/components/chat/chat-utils.ts

@@ -1,7 +1,15 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import type { FileUIPart, UIMessage } from "ai";
+import type { components } from "@marimo-team/marimo-api";
+import type { FileUIPart, ToolUIPart, UIMessage } from "ai";
+import type {
+  InvokeAiToolRequest,
+  InvokeAiToolResponse,
+} from "@/core/network/types";
+import type { ChatMessage } from "@/plugins/impl/chat/types";
 import { blobToString } from "@/utils/fileToBase64";
+import { Logger } from "@/utils/Logger";
+import { getAICompletionBodyWithAttachments } from "../editor/ai/completion-utils";
 
 export function generateChatTitle(message: string): string {
   return message.length > 50 ? `${message.slice(0, 50)}...` : message;
@@ -48,3 +56,121 @@ export function isLastMessageReasoning(messages: UIMessage[]): boolean {
   const lastPart = parts[parts.length - 1];
   return lastPart.type === "reasoning";
 }
+
+function stringifyTextParts(parts: UIMessage["parts"]): string {
+  return parts
+    .map((part) => (part.type === "text" ? part.text : ""))
+    .join("\n");
+}
+
+export async function buildCompletionRequestBody(
+  messages: UIMessage[],
+): Promise<{
+  messages: ChatMessage[];
+  context?: (null | components["schemas"]["AiCompletionContext"]) | undefined;
+  includeOtherCode: string;
+  selectedText?: string | null | undefined;
+}> {
+  const input = stringifyTextParts(messages.flatMap((m) => m.parts));
+  const completionBody = await getAICompletionBodyWithAttachments({ input });
+
+  // Map from UIMessage to our ChatMessage type
+  // If it's the last message, add the attachments from the completion body
+  function toChatMessage(message: UIMessage, isLast: boolean): ChatMessage {
+    // Clone parts to avoid mutating the original message
+    const parts = [...message.parts];
+    if (isLast) {
+      parts.push(...completionBody.attachments);
+    }
+    return {
+      role: message.role,
+      content: stringifyTextParts(message.parts), // This is no longer used in the backend
+      parts,
+    };
+  }
+
+  return {
+    ...completionBody.body,
+    messages: messages.map((m, idx) =>
+      toChatMessage(m, idx === messages.length - 1),
+    ),
+  };
+}
+
+interface AddToolResult {
+  tool: string;
+  toolCallId: string;
+  output: unknown;
+}
+
+export async function handleToolCall({
+  invokeAiTool,
+  addToolResult,
+  toolCall,
+}: {
+  invokeAiTool: (request: InvokeAiToolRequest) => Promise<InvokeAiToolResponse>;
+  addToolResult: (result: AddToolResult) => Promise<void>;
+  toolCall: {
+    toolName: string;
+    toolCallId: string;
+    input: Record<string, never>;
+  };
+}) {
+  try {
+    const response = await invokeAiTool({
+      toolName: toolCall.toolName,
+      arguments: toolCall.input,
+    });
+    addToolResult({
+      tool: toolCall.toolName,
+      toolCallId: toolCall.toolCallId,
+      output: response.result || response.error,
+    });
+  } catch (error) {
+    Logger.error("Tool call failed:", error);
+    addToolResult({
+      tool: toolCall.toolName,
+      toolCallId: toolCall.toolCallId,
+      output: `Error: ${error instanceof Error ? error.message : String(error)}`,
+    });
+  }
+}
+
+/**
+ * Checks if we should send a message automatically based on the messages.
+ * We only want to send a message if we have completed tool calls and there is no reply yet.
+ */
+export function hasPendingToolCalls(messages: UIMessage[]): boolean {
+  if (messages.length === 0) {
+    return false;
+  }
+
+  const lastMessage = messages[messages.length - 1];
+  const parts = lastMessage.parts;
+
+  if (parts.length === 0) {
+    return false;
+  }
+
+  // Only auto-send if the last message is an assistant message
+  // Because assistant messages are the ones that can have tool calls
+  if (lastMessage.role !== "assistant") {
+    return false;
+  }
+
+  const toolParts = parts.filter((part) =>
+    part.type.startsWith("tool-"),
+  ) as ToolUIPart[];
+
+  const hasCompletedToolCalls = toolParts.some(
+    (part) => part.state === "output-available",
+  );
+
+  // Check if the last part has any text content
+  const lastPart = parts[parts.length - 1];
+  const hasTextContent =
+    lastPart.type === "text" && lastPart.text?.trim().length > 0;
+
+  // Only auto-send if we have completed tool calls and there is no reply yet
+  return hasCompletedToolCalls && !hasTextContent;
+}

package/src/components/chat/markdown-renderer.css

@@ -1,4 +1,10 @@
+@source "../../../node_modules/streamdown/dist/index.js";
+
 .mo-markdown-renderer {
+  /* Make headers a bit smaller */
+  --text-2xl: 1rem;
+  --text-3xl: 1.2rem;
+
   h1,
   h2,
   h3,
@@ -12,4 +18,37 @@
     margin-top: 8px;
     margin-bottom: 12px;
   }
+
+  pre {
+    width: 100%;
+  }
+
+  [data-code-block-container="true"] {
+    margin-top: 0 !important;
+  }
+
+  [data-code-block-header="true"] {
+    padding: 4px !important;
+  }
+
+  code {
+    font-size: inherit !important;
+  }
+
+  li.task-list-item {
+    list-style-type: none;
+  }
+
+  li {
+    padding-left: 6px;
+
+    > input {
+      margin-right: 6px;
+    }
+  }
+
+  ul,
+  ol {
+    padding-left: 10px;
+  }
 }

package/src/components/chat/markdown-renderer.tsx

@@ -3,10 +3,8 @@
 import { EditorView } from "@codemirror/view";
 import { useAtomValue } from "jotai";
 import { BetweenHorizontalStartIcon } from "lucide-react";
-import { marked } from "marked";
-import { memo, Suspense, useEffect, useMemo, useState } from "react";
-import Markdown, { type Components } from "react-markdown";
-import remarkGfm from "remark-gfm";
+import { memo, Suspense, useEffect, useState } from "react";
+import { Streamdown, type StreamdownProps } from "streamdown";
 import { Button, type ButtonProps } from "@/components/ui/button";
 import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
 import { useCellActions } from "@/core/cells/cells";
@@ -171,6 +169,8 @@ const CopyButton: React.FC<ButtonProps> = ({ onClick, ...props }) => {
   );
 };
 
+type Components = StreamdownProps["components"];
+
 const COMPONENTS: Components = {
   code: ({ children, className }) => {
     const language = className?.replace("language-", "");
@@ -186,42 +186,11 @@ const COMPONENTS: Components = {
   },
 };
 
-function parseMarkdownIntoBlocks(markdown: string): string[] {
-  const tokens = marked.lexer(markdown);
-  return tokens.map((token) => token.raw);
-}
-
-const PLUGINS = [remarkGfm];
-
-const MemoizedMarkdownBlock = memo(
-  ({ content }: { content: string }) => {
-    return (
-      <Markdown
-        components={COMPONENTS}
-        remarkPlugins={PLUGINS}
-        className="mo-markdown-renderer prose dark:prose-invert max-w-none prose-pre:pl-0"
-      >
-        {content}
-      </Markdown>
-    );
-  },
-  (prevProps, nextProps) => prevProps.content === nextProps.content,
-);
-
-MemoizedMarkdownBlock.displayName = "MemoizedMarkdownBlock";
-
 export const MarkdownRenderer = memo(({ content }: { content: string }) => {
-  const blocks = useMemo(() => parseMarkdownIntoBlocks(content), [content]);
-
   return (
-    <>
-      {blocks.map((block, index) => (
-        <MemoizedMarkdownBlock
-          content={block}
-          key={`markdown-block-${index}`}
-        />
-      ))}
-    </>
+    <Streamdown components={COMPONENTS} className="mo-markdown-renderer">
+      {content}
+    </Streamdown>
   );
 });
 MarkdownRenderer.displayName = "MarkdownRenderer";

package/src/components/chat/tool-call-accordion.tsx

@@ -1,13 +1,16 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import type { ToolUIPart } from "ai";
+import { isEmpty } from "lodash-es";
 import {
   CheckCircleIcon,
+  InfoIcon,
   Loader2,
   WrenchIcon,
   XCircleIcon,
 } from "lucide-react";
 import React from "react";
+import { z } from "zod";
 import {
   Accordion,
   AccordionContent,
@@ -16,12 +19,101 @@ import {
 } from "@/components/ui/accordion";
 import { cn } from "@/utils/cn";
 
+// Zod schema matching the Python SuccessResult dataclass
+const SuccessResultSchema = z
+  .object({
+    status: z.string().default("success"),
+    auth_required: z.boolean().default(false),
+    action_url: z.any(),
+    next_steps: z.any(),
+    meta: z.any(),
+    message: z.string().nullish(),
+  })
+  .passthrough();
+
+type SuccessResult = z.infer<typeof SuccessResultSchema>;
+
+const PrettySuccessResult: React.FC<{ data: SuccessResult }> = ({ data }) => {
+  const {
+    status,
+    auth_required,
+    action_url: _action_url,
+    meta: _meta,
+    next_steps: _next_steps,
+    message,
+    ...rest
+  } = data;
+
+  return (
+    <div className="py-1 flex flex-col gap-1">
+      {/* Status */}
+      <div className="flex items-center gap-2">
+        <span className="text-xs font-medium text-[var(--grass-11)] capitalize">
+          {status}
+        </span>
+        {auth_required && (
+          <span className="text-xs px-2 py-0.5 bg-[var(--amber-2)] text-[var(--amber-11)] rounded-full">
+            Auth Required
+          </span>
+        )}
+      </div>
+
+      {/* Message */}
+      {message && (
+        <div className="flex items-start gap-2">
+          <InfoIcon className="h-3 w-3 text-[var(--blue-11)] mt-0.5 flex-shrink-0" />
+          <div className="text-xs text-foreground">{message}</div>
+        </div>
+      )}
+
+      {/* Data */}
+      {rest && (
+        <div className="flex flex-col gap-1">
+          {Object.entries(rest).map(([key, value]) => {
+            if (isEmpty(value)) {
+              return null;
+            }
+            return (
+              <div key={key}>
+                <div className="text-xs font-medium text-muted-foreground mb-1 capitalize">
+                  {key}:
+                </div>
+                <pre className="bg-[var(--slate-2)] p-1 text-muted-foreground border border-[var(--slate-4)] rounded text-xs overflow-auto scrollbar-thin max-h-64">
+                  {JSON.stringify(value, null, 2)}
+                </pre>
+              </div>
+            );
+          })}
+        </div>
+      )}
+    </div>
+  );
+};
+
+const ResultRenderer: React.FC<{ result: unknown }> = ({ result }) => {
+  // Try to parse the result with our Zod schema
+  const parseResult = SuccessResultSchema.safeParse(result);
+
+  if (parseResult.success) {
+    // If it matches the SuccessResult schema, show the pretty UI
+    return <PrettySuccessResult data={parseResult.data} />;
+  }
+
+  // Otherwise, fall back to the current JSON viewer
+  return (
+    <div className="text-xs font-medium text-muted-foreground mb-1">
+      {typeof result === "string" ? result : JSON.stringify(result, null, 2)}
+    </div>
+  );
+};
+
 interface ToolCallAccordionProps {
   toolName: string;
   result: unknown;
   error?: string;
   index?: number;
   state?: ToolUIPart["state"];
+  className?: string;
 }
 
 export const ToolCallAccordion: React.FC<ToolCallAccordionProps> = ({
@@ -30,6 +122,7 @@ export const ToolCallAccordion: React.FC<ToolCallAccordionProps> = ({
   error,
   index = 0,
   state,
+  className,
 }) => {
   const hasResult = state === "output-available" && (result || error);
   const status = error ? "error" : hasResult ? "success" : "loading";
@@ -39,9 +132,9 @@ export const ToolCallAccordion: React.FC<ToolCallAccordionProps> = ({
       case "loading":
         return <Loader2 className="h-3 w-3 animate-spin" />;
       case "error":
-        return <XCircleIcon className="h-3 w-3 text-destructive" />;
+        return <XCircleIcon className="h-3 w-3 text-[var(--red-11)]" />;
       case "success":
-        return <CheckCircleIcon className="h-3 w-3 text-green-600" />;
+        return <CheckCircleIcon className="h-3 w-3 text-[var(--grass-11)]" />;
       default:
         return <WrenchIcon className="h-3 w-3" />;
     }
@@ -49,13 +142,13 @@ export const ToolCallAccordion: React.FC<ToolCallAccordionProps> = ({
 
   const getStatusText = () => {
     if (status === "loading") {
-      return "Running: ";
+      return "Running";
     }
     if (error) {
-      return "Failed: ";
+      return "Failed";
     }
     if (hasResult) {
-      return "Done: ";
+      return "Done";
     }
     return "Tool call";
   };
@@ -65,20 +158,22 @@ export const ToolCallAccordion: React.FC<ToolCallAccordionProps> = ({
       key={`tool-${index}`}
       type="single"
       collapsible={true}
-      className="w-full"
+      className={cn("w-full", className)}
     >
       <AccordionItem value="tool-call" className="border-0">
         <AccordionTrigger
           className={cn(
-            "h-6 text-xs border-border shadow-none! ring-0! bg-muted hover:bg-muted/30 py-0 px-2 gap-1 rounded-sm [&[data-state=open]>svg]:rotate-180",
-            status === "error" && "text-destructive/80",
-            status === "success" && "text-green-600/80",
+            "h-6 text-xs border-border shadow-none! ring-0! bg-muted/60 hover:bg-muted py-0 px-2 gap-1 rounded-sm [&[data-state=open]>svg]:rotate-180 hover:no-underline",
+            status === "error" && "text-[var(--red-11)]/80",
+            status === "success" && "text-[var(--grass-11)]/80",
           )}
         >
           <span className="flex items-center gap-1">
             {getStatusIcon()}
-            {getStatusText()}:{" "}
-            <code className="font-mono text-xs">{toolName}</code>
+            {getStatusText()}:
+            <code className="font-mono text-xs">
+              {formatToolName(toolName)}
+            </code>
           </span>
         </AccordionTrigger>
         <AccordionContent className="pb-2 px-2">
@@ -86,25 +181,16 @@ export const ToolCallAccordion: React.FC<ToolCallAccordionProps> = ({
           {hasResult && (
             <div className="space-y-3">
               {result !== undefined && result !== null && (
-                <div>
-                  <div className="text-xs font-medium text-muted-foreground mt-2 mb-1">
-                    Result:
-                  </div>
-                  <div className="text-xs font-medium text-muted-foreground mb-1">
-                    {typeof result === "string"
-                      ? result
-                      : JSON.stringify(result, null, 2)}
-                  </div>
-                </div>
+                <ResultRenderer result={result} />
               )}
 
               {/* Error */}
               {error && (
                 <div>
-                  <div className="text-xs font-medium text-destructive mb-1">
+                  <div className="text-xs font-medium text-[var(--red-11)] mb-1">
                     Error:
                   </div>
-                  <div className="bg-destructive/10 border border-destructive/20 rounded-md p-3 text-sm text-destructive">
+                  <div className="bg-[var(--red-2)] border border-[var(--red-6)] rounded-md p-3 text-sm text-[var(--red-11)]">
                     {error}
                   </div>
                 </div>
@@ -116,3 +202,7 @@ export const ToolCallAccordion: React.FC<ToolCallAccordionProps> = ({
     </Accordion>
   );
 };
+
+function formatToolName(toolName: string) {
+  return toolName.replace("tool-", "");
+}

package/src/components/editor/Cell.tsx

@@ -72,6 +72,10 @@ import { CollapsedCellBanner, CollapseToggle } from "./cell/collapse";
 import { DeleteButton } from "./cell/DeleteButton";
 import { PendingDeleteConfirmation } from "./cell/PendingDeleteConfirmation";
 import { RunButton } from "./cell/RunButton";
+import {
+  StagedAICellBackground,
+  StagedAICellFooter,
+} from "./cell/StagedAICell";
 import { useDeleteCellCallback } from "./cell/useDeleteCell";
 import { useRunCell } from "./cell/useRunCells";
 import { HideCodeButton } from "./code/readonly-python-code";
@@ -571,6 +575,7 @@ const EditableCellComponent = ({
           >
             {cellOutput === "above" && outputArea}
             <div className={cn("tray")} data-hidden={isMarkdownCodeHidden}>
+              <StagedAICellBackground cellId={cellId} />
               <div className="absolute right-2 -top-4 z-10">
                 <CellToolbar
                   edited={cellData.edited}
@@ -729,6 +734,7 @@ const EditableCellComponent = ({
             />
             <PendingDeleteConfirmation cellId={cellId} />
           </div>
+          <StagedAICellFooter cellId={cellId} />
           {isCollapsed && (
             <CollapsedCellBanner
               onClick={() => actions.expandCell({ cellId })}

package/src/components/editor/actions/name-cell-input.tsx

@@ -19,12 +19,14 @@ interface Props
   value: string;
   onChange: (newName: string) => void;
   placeholder?: string;
+  onEnterKey?: () => void;
 }
 
 export const NameCellInput: React.FC<Props> = ({
   value,
   onChange,
   placeholder,
+  onEnterKey,
   ...props
 }) => {
   const ref = useRef<HTMLInputElement>(null);
@@ -53,7 +55,10 @@ export const NameCellInput: React.FC<Props> = ({
       ref={ref}
       placeholder={placeholder}
       className="shadow-none! hover:shadow-none focus:shadow-none focus-visible:shadow-none"
-      onKeyDown={Events.onEnter(Events.stopPropagation())}
+      onKeyDown={Events.onEnter((e) => {
+        Events.stopPropagation()(e);
+        onEnterKey?.();
+      })}
       {...props}
     />
   );

package/src/components/editor/actions/useCellActionButton.tsx

@@ -73,9 +73,10 @@ export interface CellActionButtonProps
 
 interface Props {
   cell: CellActionButtonProps | null;
+  closePopover?: () => void;
 }
 
-export function useCellActionButtons({ cell }: Props) {
+export function useCellActionButtons({ cell, closePopover }: Props) {
   const {
     createNewCell: createCell,
     updateCellConfig,
@@ -177,6 +178,7 @@ export function useCellActionButtons({ cell }: Props) {
             placeholder={"cell name"}
             value={name}
             onChange={(newName) => updateCellName({ cellId, name: newName })}
+            onEnterKey={() => closePopover?.()}
           />
         ),
       },