@marimo-team/frontend 0.15.5 → 0.16.0

package/src/components/chat/acp/blocks.tsx

@@ -20,7 +20,6 @@ import {
   XIcon,
 } from "lucide-react";
 import React from "react";
-import { Streamdown } from "streamdown";
 import { mergeToolCalls } from "use-acp";
 import { JsonOutput } from "@/components/editor/output/JsonOutput";
 import { Button } from "@/components/ui/button";
@@ -33,12 +32,14 @@ import { logNever } from "@/utils/assertNever";
 import { cn } from "@/utils/cn";
 import { type Base64String, base64ToDataURL } from "@/utils/json/base64";
 import { Strings } from "@/utils/strings";
+import { MarkdownRenderer } from "../markdown-renderer";
 import { SimpleAccordion } from "./common";
 import type {
   AgentNotificationEvent,
   AgentThoughtNotificationEvent,
   ConnectionChangeNotificationEvent,
   ContentBlockOf,
+  CurrentModeUpdateNotificationEvent,
   ErrorNotificationEvent,
   PlanNotificationEvent,
   SessionNotificationEventData,
@@ -307,7 +308,7 @@ export const AgentMessagesBlock = (props: {
 export const ContentBlocks = (props: { data: ContentBlock[] }) => {
   const renderBlock = (block: ContentBlock) => {
     if (block.type === "text") {
-      return <Streamdown>{block.text}</Streamdown>;
+      return <MarkdownRenderer content={block.text} />;
     }
     if (block.type === "image") {
       return <ImageBlock data={block} />;
@@ -371,7 +372,7 @@ export const ResourceBlock = (props: { data: ContentBlockOf<"resource"> }) => {
           </span>
         </PopoverTrigger>
         <PopoverContent className="max-h-96 overflow-y-auto scrollbar-thin">
-          <Streamdown>{props.data.resource.text}</Streamdown>
+          <MarkdownRenderer content={props.data.resource.text} />
         </PopoverContent>
       </Popover>
     );
@@ -495,6 +496,18 @@ export const SessionNotificationsBlock = <
       return <PlansBlock data={items} />;
     }
 
+    if (kind === "available_commands_update") {
+      return null; // nothing to show
+    }
+    if (kind === "current_mode_update") {
+      const lastItem = items.at(-1);
+      if (lastItem?.sessionUpdate !== "current_mode_update") {
+        return null;
+      } else {
+        return <CurrentModeBlock data={lastItem} />;
+      }
+    }
+
     return (
       <SimpleAccordion title={items[0].sessionUpdate}>
         <JsonOutput data={items} format="tree" className="max-h-64" />
@@ -509,6 +522,13 @@ export const SessionNotificationsBlock = <
   );
 };
 
+export const CurrentModeBlock = (props: {
+  data: CurrentModeUpdateNotificationEvent;
+}) => {
+  const { currentModeId } = props.data;
+  return <div>Mode: {currentModeId}</div>;
+};
+
 export const ToolNotificationsBlock = (props: {
   data: Array<ToolCallNotificationEvent | ToolCallUpdateNotificationEvent>;
 }) => {
@@ -635,19 +655,29 @@ export const ToolBodyBlock = (props: {
     | Omit<ToolCallNotificationEvent, "sessionUpdate">
     | Omit<ToolCallUpdateNotificationEvent, "sessionUpdate">;
 }) => {
-  const content = props.data.content
+  const { content, locations, status, kind, rawInput } = props.data;
+  const textContent = content
     ?.filter((item) => item.type === "content")
     .map((item) => item.content);
-  const diffs = props.data.content?.filter((item) => item.type === "diff");
-  const locations = props.data.locations;
-  const isFailed = props.data.status === "failed";
+  const diffs = content?.filter((item) => item.type === "diff");
+  const isFailed = status === "failed";
   const hasLocations = locations && locations.length > 0;
 
-  if (content?.length === 0 && diffs?.length === 0 && hasLocations) {
+  // Completely empty
+  if (!content && !hasLocations && rawInput) {
+    // Show rawInput
+    return (
+      <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">
+        <JsonOutput data={rawInput} format="tree" />
+      </pre>
+    );
+  }
+
+  if (content?.length && hasLocations) {
     return (
       <div className="flex flex-col gap-2 pr-2">
         <span className="text-xs text-muted-foreground">
-          {capitalize(props.data.kind || "")}{" "}
+          {capitalize(kind || "")}{" "}
           {locations?.map((item) => item.path).join(", ")}
         </span>
       </div>
@@ -657,7 +687,7 @@ export const ToolBodyBlock = (props: {
   return (
     <div className="flex flex-col gap-2 pr-2">
       {locations && <LocationsBlock data={locations} />}
-      {content && <ContentBlocks data={content} />}
+      {textContent && <ContentBlocks data={textContent} />}
       {diffs && !isFailed && <DiffBlocks data={diffs} />}
     </div>
   );

package/src/components/chat/acp/common.tsx

@@ -1,5 +1,6 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
+import type { RequestPermissionResponse } from "@zed-industries/agent-client-protocol";
 import {
   CheckCircleIcon,
   Loader2,
@@ -151,7 +152,7 @@ ConnectionStatus.displayName = "ConnectionStatus";
 
 interface PermissionRequestProps {
   permission: NonNullable<AgentPendingPermission>;
-  onResolve: (option: unknown) => void;
+  onResolve: (option: RequestPermissionResponse) => void;
 }
 
 export const PermissionRequest: React.FC<PermissionRequestProps> = memo(
@@ -179,7 +180,14 @@ export const PermissionRequest: React.FC<PermissionRequestProps> = memo(
                   ? "text-[var(--blue-10)]"
                   : "text-[var(--red-10)]"
               }
-              onClick={() => onResolve(option)}
+              onClick={() =>
+                onResolve({
+                  outcome: {
+                    outcome: "selected",
+                    optionId: option.optionId,
+                  },
+                })
+              }
             >
               {option.kind.startsWith("allow") && (
                 <CheckCircleIcon className="h-3 w-3 mr-1" />

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

@@ -0,0 +1,127 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { ContentBlock } from "@zed-industries/agent-client-protocol";
+import type { FileUIPart } from "ai";
+import { getAIContextRegistry } from "@/core/ai/context/context";
+import { store } from "@/core/state/jotai";
+import { blobToString } from "@/utils/fileToBase64";
+import { Logger } from "@/utils/Logger";
+
+export interface ContextParseResult {
+  contextBlocks: ContentBlock[];
+  attachmentBlocks: ContentBlock[];
+}
+
+/**
+ * Converts File objects to agent protocol resource_link content blocks
+ */
+export async function convertFilesToResourceLinks(
+  files: File[],
+): Promise<ContentBlock[]> {
+  const resourceLinks: ContentBlock[] = [];
+
+  for (const file of files) {
+    try {
+      const dataUrl = await blobToString(file, "dataUrl");
+      resourceLinks.push({
+        type: "resource_link",
+        uri: dataUrl,
+        mimeType: file.type,
+        name: file.name,
+      });
+    } catch (error) {
+      Logger.error("Error converting file to resource link", {
+        fileName: file.name,
+        error,
+      });
+    }
+  }
+
+  return resourceLinks;
+}
+
+/**
+ * Converts AI context registry attachments to agent protocol resource_link content blocks
+ */
+async function convertAiAttachmentsToResourceLinks(
+  attachments: FileUIPart[],
+): Promise<ContentBlock[]> {
+  const resourceLinks: ContentBlock[] = [];
+
+  for (const attachment of attachments) {
+    resourceLinks.push({
+      type: "resource_link",
+      uri: attachment.url,
+      mimeType: attachment.mediaType,
+      name: attachment.filename ?? attachment.url,
+    });
+  }
+
+  return resourceLinks;
+}
+
+/**
+ * Parses context from the prompt value and returns content blocks for agent prompts.
+ * Extracts context references using @ notation and converts them to resource blocks.
+ * Also handles attachments from the AI context registry.
+ */
+export async function parseContextFromPrompt(
+  promptValue: string,
+): Promise<ContextParseResult> {
+  const contextBlocks: ContentBlock[] = [];
+  const attachmentBlocks: ContentBlock[] = [];
+
+  // Skip if no '@' in the input
+  if (!promptValue.includes("@")) {
+    return { contextBlocks, attachmentBlocks };
+  }
+
+  try {
+    const registry = getAIContextRegistry(store);
+    const contextIds = registry.parseAllContextIds(promptValue);
+
+    if (contextIds.length === 0) {
+      return { contextBlocks, attachmentBlocks };
+    }
+
+    // Get context string for the registry
+    const contextString = registry.formatContextForAI(contextIds);
+
+    if (contextString.trim()) {
+      // Create a resource block with the context information
+      contextBlocks.push({
+        type: "resource",
+        resource: {
+          uri: "context.md",
+          mimeType: "text/markdown",
+          text: contextString,
+        },
+      });
+    }
+
+    // Get attachments from the AI context registry
+    try {
+      const aiAttachments = await registry.getAttachmentsForContext(contextIds);
+      if (aiAttachments.length > 0) {
+        const resourceLinks =
+          await convertAiAttachmentsToResourceLinks(aiAttachments);
+        attachmentBlocks.push(...resourceLinks);
+        Logger.debug("Added AI context attachments", {
+          count: aiAttachments.length,
+        });
+      }
+    } catch (error) {
+      Logger.error("Error getting AI context attachments", { error });
+    }
+
+    Logger.debug("Parsed context for agent", {
+      contextIds,
+      contextLength: contextString.length,
+      attachmentCount: attachmentBlocks.length,
+    });
+  } catch (error) {
+    Logger.error("Error parsing context for agent", { error });
+  }
+
+  return { contextBlocks, attachmentBlocks };
+}

package/src/components/chat/acp/prompt.ts

@@ -3,7 +3,7 @@
 export function getAgentPrompt(filename: string) {
   return `
   I am currently editing a marimo notebook.
-  You can read or write to the notebook at @${filename}
+  You can read or write to the notebook at ${filename}
 
   If you make edits to the notebook, only edit the contents inside the function decorator with @app.cell.
   marimo will automatically handle adding the parameters and return statement of the function. For example,
@@ -11,7 +11,7 @@ export function getAgentPrompt(filename: string) {
 
   \`\`\`
   @app.cell
-  def __():
+  def _():
     <your code here>
     return
   \`\`\`
@@ -37,6 +37,7 @@ export function getAgentPrompt(filename: string) {
   7. The last expression in a cell is automatically displayed, just like in Jupyter notebooks.
   8. Don't include comments in markdown cells
   9. Don't include comments in SQL cells
+  10. Never define anything using \`global\`.
 
   ## Reactivity
 
@@ -46,6 +47,7 @@ export function getAgentPrompt(filename: string) {
   - UI elements trigger updates when their values change without explicit callbacks
   - UI element values are accessed through \`.value\` attribute
   - You cannot access a UI element's value in the same cell where it's defined
+  - Cells prefixed with an underscore (e.g. _my_var) are local to the cell and cannot be accessed by other cells
 
   ## Best Practices
 
@@ -115,14 +117,29 @@ export function getAgentPrompt(filename: string) {
 
   - \`mo.md(text)\` - display markdown
   - \`mo.stop(predicate, output=None)\` - stop execution conditionally
+  - \`mo.output.append(value)\` - append to the output when it is not the last expression
+  - \`mo.output.replace(value)\` - replace the output when it is not the last expression
   - \`mo.Html(html)\` - display HTML
   - \`mo.image(image)\` - display an image
   - \`mo.hstack(elements)\` - stack elements horizontally
   - \`mo.vstack(elements)\` - stack elements vertically
   - \`mo.tabs(elements)\` - create a tabbed interface
 
+
+
   ## Examples
 
+  <example title="Markdown ccell">
+  ${formatCells([
+    `
+  mo.md("""
+  # Hello world
+  This is a _markdown_ **cell**.
+  """)
+    `,
+  ])}
+  </example>
+
   <example title="Basic UI with reactivity">
   ${formatCells([
     `
@@ -207,6 +224,19 @@ export function getAgentPrompt(filename: string) {
   ])}
   </example>
 
+  <example title="Conditional Outputs">
+  ${formatCells([
+    `
+  mo.stop(not data.value, mo.md("No data to display"))
+
+  if mode.value == "scatter":
+    mo.output.replace(render_scatter(data.value))
+  else:
+    mo.output.replace(render_bar_chart(data.value))
+    `,
+  ])}
+  </example>
+
   <example title="Interactive chart with Altair">
   ${formatCells([
     `
@@ -256,25 +286,19 @@ export function getAgentPrompt(filename: string) {
 }
 
 function formatCells(cells: string[]) {
-  // Option 1:
-  // return cells.map((cell) => {
-  //   return `# Cell ${cell}`;
-  // });
-
   const indent = "    ";
   const indentCode = (code: string) => {
     return code
       .trim()
       .split("\n")
-      .map((line) => indent + line)
+      .map((line) => (indent + line).trimEnd())
       .join("\n");
   };
 
-  // Option 2:
   const formatCell = (cell: string) => {
     return `
   @app.cell
-  def __():
+  def _():
   ${indentCode(cell)}
       return
   `;

package/src/components/chat/acp/state.ts

@@ -259,5 +259,5 @@ export function getAgentSessionSupport(
 export function getAgentConnectionCommand(agentId: ExternalAgentId): string {
   const port = AGENT_CONFIG[agentId].port;
   const command = AGENT_CONFIG[agentId].command;
-  return `npx supergateway --stdio\\\n  "${command}" \\\n   --outputTransport ws --port ${port} `;
+  return `npx stdio-to-ws "${command}" --port ${port}`;
 }

package/src/components/chat/acp/types.ts

@@ -2,10 +2,16 @@
 import type { ContentBlock } from "@zed-industries/agent-client-protocol";
 import type { groupNotifications, useAcpClient } from "use-acp";
 
+export type SessionMode = ReturnType<typeof useAcpClient>["sessionMode"];
+
 export type NotificationEvent = Awaited<
   ReturnType<typeof groupNotifications>
 >[number][number];
 
+export type AvailableCommands = ReturnType<
+  typeof useAcpClient
+>["availableCommands"];
+
 export type AgentConnectionState = ReturnType<
   typeof useAcpClient
 >["connectionState"];
@@ -49,6 +55,8 @@ export type ToolCallNotificationEvent = NotificationDataOf<"tool_call">;
 export type ToolCallUpdateNotificationEvent =
   NotificationDataOf<"tool_call_update">;
 export type PlanNotificationEvent = NotificationDataOf<"plan">;
+export type CurrentModeUpdateNotificationEvent =
+  NotificationDataOf<"current_mode_update">;
 
 export type ContentBlockType = ContentBlock["type"];
 export type ContentBlockOf<T extends ContentBlockType> = Extract<

package/src/components/chat/chat-panel.tsx

@@ -48,7 +48,6 @@ import { DEFAULT_AI_MODEL } from "@/core/config/config-schema";
 import { FeatureFlagged } from "@/core/config/feature-flag";
 import { useRequestClient } from "@/core/network/requests";
 import { useRuntimeManager } from "@/core/runtime/config";
-import type { ChatMessage } from "@/plugins/impl/chat/types";
 import { ErrorBanner } from "@/plugins/impl/common/error-banner";
 import { cn } from "@/utils/cn";
 import { timeAgo } from "@/utils/dates";
@@ -59,7 +58,6 @@ import { PromptInput } from "../editor/ai/add-cell-with-ai";
 import {
   addContextCompletion,
   CONTEXT_TRIGGER,
-  getAICompletionBodyWithAttachments,
 } from "../editor/ai/completion-utils";
 import { PanelEmptyState } from "../editor/chrome/panels/empty-state";
 import { CopyClipboardIcon } from "../icons/copy-icon";
@@ -68,8 +66,11 @@ import { Tooltip, TooltipProvider } from "../ui/tooltip";
 import { toast } from "../ui/use-toast";
 import { AttachmentRenderer, FileAttachmentPill } from "./chat-components";
 import {
+  buildCompletionRequestBody,
   convertToFileUIPart,
   generateChatTitle,
+  handleToolCall,
+  hasPendingToolCalls,
   isLastMessageReasoning,
 } from "./chat-utils";
 import { MarkdownRenderer } from "./markdown-renderer";
@@ -233,6 +234,7 @@ const ChatMessageDisplay: React.FC<ChatMessageProps> = memo(
                   index={i}
                   toolName={part.type}
                   result={part.output}
+                  className="my-2"
                   state={part.state}
                 />
               );
@@ -265,6 +267,7 @@ const ChatMessageDisplay: React.FC<ChatMessageProps> = memo(
                     toolName={part.type}
                     result={part.output}
                     state={part.state}
+                    className="my-2"
                   />
                 );
 
@@ -282,10 +285,11 @@ const ChatMessageDisplay: React.FC<ChatMessageProps> = memo(
                 Logger.error("Unhandled part type:", part.type);
                 try {
                   return (
-                    <MarkdownRenderer
-                      key={i}
-                      content={JSON.stringify(part, null, 2)}
-                    />
+                    <div className="text-xs text-muted-foreground" key={i}>
+                      <MarkdownRenderer
+                        content={JSON.stringify(part, null, 2)}
+                      />
+                    </div>
                   );
                 } catch (error) {
                   Logger.error("Error rendering part:", part.type, error);
@@ -429,7 +433,7 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
               </>
             )}
 
-            <Tooltip content="Submit">
+            <Tooltip content={isLoading ? "Stop" : "Submit"}>
               <Button
                 variant="text"
                 size="sm"
@@ -565,80 +569,20 @@ const ChatPanelBody = () => {
     id: chatId,
   } = useChat({
     id: activeChatId,
-    // Only automatically submit if we have tool calls but no text response yet
-    sendAutomaticallyWhen: ({ messages }) => {
-      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;
-    },
+    sendAutomaticallyWhen: ({ messages }) => hasPendingToolCalls(messages),
     messages: activeChat?.messages || [], // initial messages
     transport: new DefaultChatTransport({
       api: runtimeManager.getAiURL("chat").toString(),
       headers: runtimeManager.headers(),
       prepareSendMessagesRequest: async (options) => {
-        // Map from parts to a single string
-        function toContent(parts: UIMessage["parts"]): string {
-          return parts
-            .map((part) => (part.type === "text" ? part.text : ""))
-            .join("\n");
-        }
-
-        const input = toContent(options.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 mapMessage(m: UIMessage, isLastMessage: boolean): ChatMessage {
-          const parts = m.parts;
-          if (isLastMessage) {
-            parts.push(...completionBody.attachments);
-          }
-          return {
-            role: m.role,
-            content: toContent(m.parts),
-            parts: parts,
-          };
-        }
+        const completionBody = await buildCompletionRequestBody(
+          options.messages,
+        );
 
         return {
           body: {
             ...options,
-            ...completionBody.body,
-            messages: options.messages.map((m, idx) => ({
-              ...m,
-              ...mapMessage(m, idx === options.messages.length - 1),
-            })),
+            ...completionBody,
           },
         };
       },
@@ -653,24 +597,15 @@ const ChatPanelBody = () => {
       });
     },
     onToolCall: async ({ toolCall }) => {
-      try {
-        const response = await invokeAiTool({
+      await handleToolCall({
+        invokeAiTool,
+        addToolResult,
+        toolCall: {
           toolName: toolCall.toolName,
-          arguments: toolCall.input as Record<string, never>,
-        });
-        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)}`,
-        });
-      }
+          input: toolCall.input as Record<string, never>,
+        },
+      });
     },
     onError: (error) => {
       Logger.error("An error occurred:", error);