@marimo-team/frontend 0.15.5 → 0.16.0-dev96986

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,16 @@ 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);
+      return lastItem?.sessionUpdate === "current_mode_update" ? (
+        <CurrentModeBlock data={lastItem} />
+      ) : null;
+    }
+
     return (
       <SimpleAccordion title={items[0].sessionUpdate}>
         <JsonOutput data={items} format="tree" className="max-h-64" />
@@ -509,6 +520,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 +653,31 @@ 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>
+    );
+  }
+
+  const noContent = !textContent || textContent.length === 0;
+  const noDiffs = !diffs || diffs.length === 0;
+  if (noContent && noDiffs && 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,11 +47,12 @@ 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
 
   <data_handling>
-  - Use pandas for data manipulation
+  - Use polars for data manipulation
   - Implement proper data validation
   - Handle missing values appropriately
   - Use efficient data structures
@@ -60,7 +62,7 @@ export function getAgentPrompt(filename: string) {
   <visualization>
   - For matplotlib: use plt.gca() as the last expression instead of plt.show()
   - For plotly: return the figure object directly
-  - For altair: return the chart object directly
+  - For altair: return the chart object directly. Add tooltips where appropriate. You can pass polars dataframes directly to altair.
   - Include proper labels, titles, and color schemes
   - Make visualizations interactive where appropriate
   </visualization>
@@ -74,7 +76,7 @@ export function getAgentPrompt(filename: string) {
   </ui_elements>
 
   <sql>
-  - When writing duckdb, prefer using marimo's SQL cells, which start with _df = mo.sql(query)
+  - When writing duckdb, prefer using marimo's SQL cells, which start with df = mo.sql(f"""<your query>""") for DuckDB, or df = mo.sql(f"""<your query>""", engine=engine) for other SQL engines.
   - See the SQL with duckdb example for an example on how to do this
   - Don't add comments in cells that use mo.sql()
   - Consider using \`vega_datasets\` for common example datasets
@@ -115,19 +117,35 @@ 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([
     `
   import marimo as mo
-  import matplotlib.pyplot as plt
+  import altair as alt
+  import polars as pl
   import numpy as np
     `,
     `
@@ -138,12 +156,18 @@ export function getAgentPrompt(filename: string) {
   x = np.random.rand(n_points.value)
   y = np.random.rand(n_points.value)
 
-  plt.figure(figsize=(8, 6))
-  plt.scatter(x, y, alpha=0.7)
-  plt.title(f"Scatter plot with {n_points.value} points")
-  plt.xlabel("X axis")
-  plt.ylabel("Y axis")
-  plt.gca()
+  df = pl.DataFrame({"x": x, "y": y})
+
+  chart = alt.Chart(df).mark_circle(opacity=0.7).encode(
+      x=alt.X('x', title='X axis'),
+      y=alt.Y('y', title='Y axis')
+  ).properties(
+      title=f"Scatter plot with {n_points.value} points",
+      width=400,
+      height=300
+  )
+
+  chart
     `,
   ])}
   </example>
@@ -152,11 +176,11 @@ export function getAgentPrompt(filename: string) {
   ${formatCells([
     `
   import marimo as mo
-  import pandas as pd
+  import polars as pl
   from vega_datasets import data
     `,
     `
-  cars_df = data.cars()
+  cars_df = pl.DataFrame(data.cars())
   mo.ui.data_explorer(cars_df)
     `,
   ])}
@@ -166,43 +190,57 @@ export function getAgentPrompt(filename: string) {
   ${formatCells([
     `
   import marimo as mo
-  import pandas as pd
-  import matplotlib.pyplot as plt
-  import seaborn as sns
+  import polars as pl
+  import altair as alt
     `,
     `
-  iris = sns.load_dataset('iris')
+  iris = pl.read_csv("hf://datasets/scikit-learn/iris/Iris.csv")
     `,
     `
   species_selector = mo.ui.dropdown(
-      options=["All"] + iris["species"].unique().tolist(),
+      options=["All"] + iris["Species"].unique().to_list(),
       value="All",
-      label="Species"
+      label="Species",
   )
   x_feature = mo.ui.dropdown(
-      options=iris.select_dtypes('number').columns.tolist(),
-      value="sepal_length",
-      label="X Feature"
+      options=iris.select(pl.col(pl.Float64, pl.Int64)).columns,
+      value="SepalLengthCm",
+      label="X Feature",
   )
   y_feature = mo.ui.dropdown(
-      options=iris.select_dtypes('number').columns.tolist(),
-      value="sepal_width",
-      label="Y Feature"
+      options=iris.select(pl.col(pl.Float64, pl.Int64)).columns,
+      value="SepalWidthCm",
+      label="Y Feature",
   )
   mo.hstack([species_selector, x_feature, y_feature])
     `,
     `
-  filtered_data = iris if species_selector.value == "All" else iris[iris["species"] == species_selector.value]
-
-  plt.figure(figsize=(10, 6))
-  sns.scatterplot(
-      data=filtered_data,
-      x=x_feature.value,
-      y=y_feature.value,
-      hue="species"
+  filtered_data = iris if species_selector.value == "All" else iris.filter(pl.col("Species") == species_selector.value)
+
+  chart = alt.Chart(filtered_data).mark_circle().encode(
+      x=alt.X(x_feature.value, title=x_feature.value),
+      y=alt.Y(y_feature.value, title=y_feature.value),
+      color='Species'
+  ).properties(
+      title=f"{y_feature.value} vs {x_feature.value}",
+      width=500,
+      height=400
   )
-  plt.title(f"{y_feature.value} vs {x_feature.value}")
-  plt.gca()
+
+  chart
+    `,
+  ])}
+  </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>
@@ -212,14 +250,21 @@ export function getAgentPrompt(filename: string) {
     `
   import marimo as mo
   import altair as alt
-  import pandas as pd
+  import polars as pl
     `,
     `# Load dataset
-  cars_df = pd.read_csv('<https://raw.githubusercontent.com/vega/vega-datasets/master/data/cars.json>')
-  _chart = alt.Chart(cars_df).mark_point().encode(
-      x='Horsepower',
-      y='Miles_per_Gallon',
-      color='Origin',
+  weather = pl.read_csv("https://raw.githubusercontent.com/vega/vega-datasets/refs/heads/main/data/weather.csv")
+  weather_dates = weather.with_columns(
+      pl.col("date").str.strptime(pl.Date, format="%Y-%m-%d")
+  )
+  _chart = (
+      alt.Chart(weather_dates)
+      .mark_point()
+      .encode(
+          x="date:T",
+          y="temp_max",
+          color="location",
+      )
   )
   `,
     "chart = mo.ui.altair_chart(_chart)\nchart",
@@ -248,33 +293,31 @@ export function getAgentPrompt(filename: string) {
 
   <example title="SQL with duckdb">
   ${formatCells([
-    "import marimo as mo\nimport pandas as pd",
-    `cars_df = pd.read_csv('<https://raw.githubusercontent.com/vega/vega-datasets/master/data/cars.json>')`,
-    `_df = mo.sql("SELECT * from cars_df WHERE Miles_per_Gallon > 20")`,
+    "import marimo as mo\n  import polars as pl",
+    `weather = pl.read_csv('https://raw.githubusercontent.com/vega/vega-datasets/refs/heads/main/data/weather.csv')`,
+    `seattle_weather_df = mo.sql(
+      f"""
+      SELECT * FROM weather WHERE location = 'Seattle';
+      """
+  )`,
   ])}
   </example>`;
 }
 
 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<