@marimo-team/islands 0.16.4 → 0.17.0

package/src/components/icons/copy-icon.tsx

@@ -9,10 +9,11 @@ import { Tooltip } from "../ui/tooltip";
 import { toast } from "../ui/use-toast";
 
 interface Props {
-  value: string | (() => string);
+  value: string | ((event: React.MouseEvent) => string);
   className?: string;
-  tooltip?: string | false;
+  tooltip?: React.ReactNode | false;
   toastTitle?: string;
+  ariaLabel?: string;
 }
 
 export const CopyClipboardIcon: React.FC<Props> = ({
@@ -20,11 +21,12 @@ export const CopyClipboardIcon: React.FC<Props> = ({
   className,
   tooltip,
   toastTitle,
+  ariaLabel,
 }) => {
   const [isCopied, setIsCopied] = useState(false);
 
-  const handleCopy = Events.stopPropagation(async () => {
-    const valueToCopy = typeof value === "function" ? value() : value;
+  const handleCopy = Events.stopPropagation(async (event: React.MouseEvent) => {
+    const valueToCopy = typeof value === "function" ? value(event) : value;
     await copyToClipboard(valueToCopy).then(() => {
       setIsCopied(true);
       setTimeout(() => setIsCopied(false), 2000);
@@ -35,7 +37,11 @@ export const CopyClipboardIcon: React.FC<Props> = ({
   });
 
   const button = (
-    <button type="button" onClick={handleCopy} aria-label="Copy to clipboard">
+    <button
+      type="button"
+      onClick={handleCopy}
+      aria-label={ariaLabel ?? "Copy to clipboard"}
+    >
       {isCopied ? (
         <CheckIcon className={cn(className, "text-(--grass-11)")} />
       ) : (

package/src/components/mcp/hooks.ts

@@ -0,0 +1,48 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { components } from "@marimo-team/marimo-api";
+import { useState } from "react";
+import { API } from "@/core/network/api";
+import { useAsyncData } from "@/hooks/useAsyncData";
+import { toast } from "../ui/use-toast";
+
+export type MCPStatus = components["schemas"]["MCPStatusResponse"];
+export type MCPRefreshResponse = components["schemas"]["MCPRefreshResponse"];
+
+/**
+ * Hook to fetch MCP server status
+ */
+export function useMCPStatus() {
+  return useAsyncData<MCPStatus>(async () => {
+    return API.get<MCPStatus>("/ai/mcp/status");
+  }, []);
+}
+
+/**
+ * Hook to refresh MCP server configuration
+ */
+export function useMCPRefresh() {
+  const [isRefreshing, setIsRefreshing] = useState(false);
+
+  const refresh = async () => {
+    setIsRefreshing(true);
+    try {
+      await API.post<object, MCPRefreshResponse>("/ai/mcp/refresh", {});
+      toast({
+        title: "MCP refreshed",
+        description: "MCP server configuration has been refreshed successfully",
+      });
+    } catch (error) {
+      toast({
+        title: "Refresh failed",
+        description:
+          error instanceof Error ? error.message : "Failed to refresh MCP",
+        variant: "danger",
+      });
+    } finally {
+      setIsRefreshing(false);
+    }
+  };
+
+  return { refresh, isRefreshing };
+}

package/src/components/mcp/mcp-status-indicator.tsx

@@ -0,0 +1,144 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { Loader2, PlugIcon, RefreshCwIcon } from "lucide-react";
+import { API } from "@/core/network/api";
+import { cn } from "@/utils/cn";
+import { Button } from "../ui/button";
+import { Popover, PopoverContent, PopoverTrigger } from "../ui/popover";
+import { Tooltip } from "../ui/tooltip";
+import { toast } from "../ui/use-toast";
+import { useMCPStatus } from "./hooks";
+
+/**
+ * MCP Status indicator component
+ * Shows a small icon with status color and a popover with detailed information
+ */
+export const MCPStatusIndicator: React.FC = () => {
+  const { data: status, refetch, isFetching } = useMCPStatus();
+
+  const handleRefresh = async () => {
+    try {
+      await API.post<object, { success: boolean }>("/ai/mcp/refresh", {});
+      toast({
+        title: "MCP refreshed",
+        description: "MCP server configuration has been refreshed",
+      });
+      refetch();
+    } catch (error) {
+      toast({
+        title: "Refresh failed",
+        description:
+          error instanceof Error ? error.message : "Failed to refresh MCP",
+        variant: "danger",
+      });
+    }
+  };
+
+  const servers = status?.servers || {};
+  const hasServers = Object.keys(servers).length > 0;
+
+  return (
+    <Popover>
+      <Tooltip content="MCP Status">
+        <PopoverTrigger asChild={true}>
+          <Button variant="text" size="icon">
+            <PlugIcon
+              className={cn(
+                "h-4 w-4",
+                status?.status === "ok" && "text-green-500",
+                status?.status === "partial" && "text-yellow-500",
+                status?.status === "error" && hasServers && "text-red-500",
+              )}
+            />
+          </Button>
+        </PopoverTrigger>
+      </Tooltip>
+      <PopoverContent className="w-[320px]" align="start" side="right">
+        <div className="space-y-3">
+          <div className="flex items-center justify-between">
+            <h4 className="font-medium text-sm">MCP Server Status</h4>
+            <Button
+              variant="ghost"
+              size="xs"
+              onClick={handleRefresh}
+              disabled={isFetching}
+            >
+              {isFetching ? (
+                <Loader2 className="h-3 w-3 animate-spin" />
+              ) : (
+                <RefreshCwIcon className="h-3 w-3" />
+              )}
+            </Button>
+          </div>
+          {status && (
+            <div className="text-xs space-y-2">
+              {hasServers && (
+                <div className="flex justify-between items-center">
+                  <span className="text-muted-foreground">Overall:</span>
+                  <McpStatusText status={status.status} />
+                </div>
+              )}
+              {status.error && (
+                <div className="text-xs text-red-500 bg-red-50 dark:bg-red-950/20 p-2 rounded">
+                  {status.error}
+                </div>
+              )}
+              {hasServers && (
+                <div className="space-y-1">
+                  <div className="text-muted-foreground font-medium">
+                    Servers:
+                  </div>
+                  {Object.entries(servers).map(([name, serverStatus]) => (
+                    <div
+                      key={name}
+                      className="flex justify-between items-center pl-2"
+                    >
+                      <span className="text-muted-foreground truncate max-w-[180px]">
+                        {name}
+                      </span>
+                      <McpStatusText status={serverStatus} />
+                    </div>
+                  ))}
+                </div>
+              )}
+              {!hasServers && (
+                <div className="text-muted-foreground text-center py-2">
+                  No MCP servers configured. <br /> Configure under{" "}
+                  <b>Settings &gt; AI &gt; MCP</b>
+                </div>
+              )}
+            </div>
+          )}
+        </div>
+      </PopoverContent>
+    </Popover>
+  );
+};
+
+export const McpStatusText: React.FC<{
+  status:
+    | "ok"
+    | "partial"
+    | "error"
+    | "failed"
+    | "disconnected"
+    | "pending"
+    | "connected";
+}> = ({ status }) => {
+  return (
+    <span
+      className={cn(
+        "text-xs font-medium",
+        status === "ok" && "text-green-500",
+        status === "partial" && "text-yellow-500",
+        status === "error" && "text-red-500",
+        status === "failed" && "text-red-500",
+        status === "disconnected" && "text-gray-500",
+        status === "pending" && "text-yellow-500",
+        status === "connected" && "text-green-500",
+      )}
+    >
+      {status}
+    </span>
+  );
+};

package/src/components/shortcuts/renderShortcut.tsx

@@ -102,7 +102,9 @@ export const MinimalHotkeys: React.FC<{
   );
 };
 
-function prettyPrintHotkey(key: string): [label: string, symbol?: string] {
+export function prettyPrintHotkey(
+  key: string,
+): [label: string, symbol?: string] {
   const platform = isPlatformMac() ? "mac" : "default";
 
   const lowerKey = key.toLowerCase();

package/src/components/ui/button.tsx

@@ -70,6 +70,8 @@ const buttonVariants = cva(
         link: "underline-offset-4 hover:underline text-link",
         linkDestructive:
           "underline-offset-4 hover:underline text-destructive underline-destructive",
+        outlineDestructive:
+          "border border-destructive text-destructive hover:bg-destructive/10",
       },
       size: {
         default: "h-10 py-2 px-4",

package/src/components/ui/confirmation-button.tsx

@@ -0,0 +1,85 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import React from "react";
+import {
+  AlertDialog,
+  AlertDialogAction,
+  AlertDialogCancel,
+  AlertDialogContent,
+  AlertDialogDescription,
+  AlertDialogDestructiveAction,
+  AlertDialogFooter,
+  AlertDialogHeader,
+  AlertDialogTitle,
+  AlertDialogTrigger,
+} from "./alert-dialog";
+
+interface ConfirmationButtonProps {
+  /**
+   * The button element to trigger the confirmation dialog
+   */
+  children: React.ReactElement;
+  /**
+   * Title of the confirmation dialog
+   */
+  title: string;
+  /**
+   * Description/message of the confirmation dialog
+   */
+  description: string;
+  /**
+   * Callback when the user confirms the action
+   */
+  onConfirm: () => void;
+  /**
+   * Text for the confirm button (default: "Continue")
+   */
+  confirmText?: string;
+  /**
+   * Text for the cancel button (default: "Cancel")
+   */
+  cancelText?: string;
+  /**
+   * Whether to use destructive styling for the confirm button
+   */
+  destructive?: boolean;
+}
+
+export const ConfirmationButton: React.FC<ConfirmationButtonProps> = ({
+  children,
+  title,
+  description,
+  onConfirm,
+  confirmText = "Continue",
+  cancelText = "Cancel",
+  destructive = false,
+}) => {
+  const [open, setOpen] = React.useState(false);
+
+  const handleConfirm = () => {
+    onConfirm();
+    setOpen(false);
+  };
+
+  const ActionComponent = destructive
+    ? AlertDialogDestructiveAction
+    : AlertDialogAction;
+
+  return (
+    <AlertDialog open={open} onOpenChange={setOpen}>
+      <AlertDialogTrigger asChild>{children}</AlertDialogTrigger>
+      <AlertDialogContent>
+        <AlertDialogHeader>
+          <AlertDialogTitle>{title}</AlertDialogTitle>
+          <AlertDialogDescription>{description}</AlertDialogDescription>
+        </AlertDialogHeader>
+        <AlertDialogFooter>
+          <AlertDialogCancel>{cancelText}</AlertDialogCancel>
+          <ActionComponent onClick={handleConfirm}>
+            {confirmText}
+          </ActionComponent>
+        </AlertDialogFooter>
+      </AlertDialogContent>
+    </AlertDialog>
+  );
+};

package/src/components/ui/number-field.tsx

@@ -7,8 +7,10 @@ import {
   Button,
   type ButtonProps,
   Input as RACInput,
+  useLocale,
 } from "react-aria-components";
 import { cn } from "@/utils/cn";
+import { maxFractionalDigits } from "@/utils/numbers";
 
 export interface NumberFieldProps extends AriaNumberFieldProps {
   placeholder?: string;
@@ -17,12 +19,13 @@ export interface NumberFieldProps extends AriaNumberFieldProps {
 
 export const NumberField = React.forwardRef<HTMLInputElement, NumberFieldProps>(
   ({ placeholder, variant = "default", ...props }, ref) => {
+    const { locale } = useLocale();
     return (
       <AriaNumberField
         {...props}
         formatOptions={{
           minimumFractionDigits: 0,
-          maximumFractionDigits: 100,
+          maximumFractionDigits: maxFractionalDigits(locale),
         }}
       >
         <div

package/src/components/ui/switch.tsx

@@ -12,6 +12,7 @@ const switchVariants = cva(
       size: {
         default: "h-6 w-11 mb-1",
         sm: "h-4.5 w-8.5",
+        xs: "h-4 w-7",
       },
     },
     defaultVariants: {
@@ -28,6 +29,7 @@ const switchThumbVariants = cva(
         default:
           "h-5 w-5 data-[state=checked]:translate-x-5 data-[state=unchecked]:translate-x-0",
         sm: "h-3.5 w-3.5 data-[state=checked]:translate-x-4 data-[state=unchecked]:translate-x-0",
+        xs: "h-3 w-3 data-[state=checked]:translate-x-3 data-[state=unchecked]:translate-x-0",
       },
     },
     defaultVariants: {

package/src/core/ai/context/providers/__tests__/__snapshots__/tables.test.ts.snap

@@ -1,34 +1,28 @@
 // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
 
 exports[`TableContextProvider > formatContext > should format context for basic table > basic-table-context 1`] = `
-"<data name="products" source="memory">
-Shape: 100 rows, 3 columns
+"<data name="products" source="memory">Shape: 100 rows, 3 columns
 Columns:
-  - id: integer
-  - name: string
-  - active: boolean</data>"
+  id (integer) - samples: [sample_id_1, sample_id_2]
+  name (string) - samples: [sample_name_1, sample_name_2]
+  active (boolean) - samples: [sample_active_1, sample_active_2]</data>"
 `;
 
 exports[`TableContextProvider > formatContext > should format context for remote database table > remote-table-context 1`] = `
-"<data name="remote_table" source="postgresql://localhost:5432/mydb">
-Shape: 100 rows, 3 columns
+"<data name="remote_table" source="postgresql://localhost:5432/mydb">Shape: 100 rows, 3 columns
 Columns:
-  - uuid: string
-  - created_at: string
-  - metadata: string</data>"
+  uuid (string) - samples: [sample_uuid_1, sample_uuid_2]
+  created_at (string) - samples: [sample_created_at_1, sample_created_at_2]
+  metadata (string) - samples: [sample_metadata_1, sample_metadata_2]</data>"
 `;
 
-exports[`TableContextProvider > formatContext > should format context for table without columns > no-columns-table-context 1`] = `
-"<data name="no_columns" source="memory">
-Shape: 100 rows, 3 columns</data>"
-`;
+exports[`TableContextProvider > formatContext > should format context for table without columns > no-columns-table-context 1`] = `"<data name="no_columns" source="memory">Shape: 100 rows, 3 columns</data>"`;
 
 exports[`TableContextProvider > formatContext > should format context for table without shape info > no-shape-table-context 1`] = `
-"<data name="no_shape" source="memory">
-Columns:
-  - id: integer
-  - name: string
-  - active: boolean</data>"
+"<data name="no_shape" source="memory">Columns:
+  id (integer) - samples: [sample_id_1, sample_id_2]
+  name (string) - samples: [sample_name_1, sample_name_2]
+  active (boolean) - samples: [sample_active_1, sample_active_2]</data>"
 `;
 
 exports[`TableContextProvider > getItems > should handle dataframe tables with variable names > dataframe-table 1`] = `

package/src/core/ai/context/providers/__tests__/cell-output.test.ts

@@ -187,7 +187,6 @@ describe("CellOutputContextProvider", () => {
       expect(completion.displayLabel).toBe(item.data.cellName);
       expect(completion.detail).toContain("output");
       expect(completion.type).toBe("cell-output");
-      expect(completion.section).toBe("Cell Output");
       expect(typeof completion.info).toBe("function");
     });
   });

package/src/core/ai/context/providers/__tests__/datasource.test.ts

@@ -2,14 +2,15 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
 import { beforeEach, describe, expect, it } from "vitest";
+import type { CellId } from "@/core/cells/ids";
 import type {
   ConnectionsMap,
   DatasetTablesMap,
 } from "@/core/datasets/data-source-connections";
 import { DUCKDB_ENGINE } from "@/core/datasets/engines";
 import type { DataSourceConnection, DataTable } from "@/core/kernel/messages";
-import { Boosts } from "../common";
-import { DatasourceContextProvider } from "../datasource";
+import { Boosts, Sections } from "../common";
+import { DatasourceContextProvider, getDatasourceContext } from "../datasource";
 
 // Mock data for testing
 const createMockDataSourceConnection = (
@@ -268,8 +269,7 @@ describe("DatasourceContextProvider", () => {
       );
 
       const items = providerWithEmpty.getItems();
-      expect(items).toHaveLength(1);
-      expect(items[0].data.connection.databases).toEqual([]);
+      expect(items).toHaveLength(0);
     });
 
     it("should handle connections with databases but no schemas", () => {
@@ -311,7 +311,7 @@ describe("DatasourceContextProvider", () => {
         detail: "DuckDB",
         boost: Boosts.MEDIUM,
         type: "datasource",
-        section: "Data Sources",
+        section: Sections.DATA_SOURCES,
       });
 
       expect(completion.info).toBeDefined();
@@ -355,7 +355,7 @@ describe("DatasourceContextProvider", () => {
         detail: "PostgreSQL",
         boost: Boosts.MEDIUM,
         type: "datasource",
-        section: "Data Sources",
+        section: Sections.DATA_SOURCES,
       });
     });
 
@@ -380,7 +380,7 @@ describe("DatasourceContextProvider", () => {
         detail: "DuckDB",
         boost: Boosts.MEDIUM,
         type: "datasource",
-        section: "Data Sources",
+        section: Sections.DATA_SOURCES,
       });
     });
   });
@@ -631,3 +631,10 @@ describe("DatasourceContextProvider", () => {
     });
   });
 });
+
+describe("getDatasourceContext", () => {
+  it("should return null if no cell ID is found", () => {
+    const context = getDatasourceContext("1" as CellId);
+    expect(context).toBeNull();
+  });
+});

package/src/core/ai/context/providers/__tests__/error.test.ts

@@ -6,7 +6,6 @@ import { beforeEach, describe, expect, it } from "vitest";
 import { MockNotebook } from "@/__mocks__/notebook";
 import { notebookAtom } from "@/core/cells/cells";
 import { type CellId, CellId as CellIdClass } from "@/core/cells/ids";
-import { Boosts } from "../common";
 import { ErrorContextProvider } from "../error";
 
 describe("ErrorContextProvider", () => {
@@ -118,15 +117,20 @@ describe("ErrorContextProvider", () => {
       const items = provider.getItems();
       const completion = provider.formatCompletion(items[0]);
 
-      expect(completion).toMatchObject({
-        label: "@Errors",
-        displayLabel: "Errors",
-        detail: "2 errors",
-        boost: Boosts.ERROR,
-        type: "error",
-        apply: "@Errors",
-        section: "Errors",
-      });
+      expect(completion).toMatchInlineSnapshot(`
+        {
+          "apply": "@Errors",
+          "detail": "2 errors",
+          "displayLabel": "Errors",
+          "info": [Function],
+          "label": "@Errors",
+          "section": {
+            "name": "Error",
+            "rank": 1,
+          },
+          "type": "error",
+        }
+      `);
 
       // Test the info function
       expect(completion.info).toBeDefined();
@@ -179,11 +183,16 @@ describe("ErrorContextProvider", () => {
       };
 
       const completion = provider.formatCompletion(item);
-      expect(completion).toMatchObject({
-        label: "Error",
-        displayLabel: "Error",
-        boost: Boosts.ERROR,
-      });
+      expect(completion).toMatchInlineSnapshot(`
+        {
+          "displayLabel": "Error",
+          "label": "Error",
+          "section": {
+            "name": "Error",
+            "rank": 1,
+          },
+        }
+      `);
     });
   });
 

package/src/core/ai/context/providers/cell-output.ts

@@ -13,7 +13,7 @@ import { parseHtmlContent } from "@/utils/dom";
 import { Logger } from "@/utils/Logger";
 import { type AIContextItem, AIContextProvider } from "../registry";
 import { contextToXml } from "../utils";
-import { Boosts } from "./common";
+import { Boosts, Sections } from "./common";
 
 export interface CellOutputContextItem extends AIContextItem {
   type: "cell-output";
@@ -150,13 +150,14 @@ export class CellOutputContextProvider extends AIContextProvider<CellOutputConte
 
   formatCompletion(item: CellOutputContextItem): Completion {
     const { data } = item;
+
     return {
       label: `@${data.cellName}`,
       displayLabel: data.cellName,
       detail: `${data.outputType} output`,
-      boost: Boosts.CELL_OUTPUT,
+      boost: data.outputType === "media" ? Boosts.HIGH : Boosts.MEDIUM,
       type: this.contextType,
-      section: "Cell Output",
+      section: Sections.CELL_OUTPUT,
       apply: `@${data.cellName}`,
       info: () => {
         const infoContainer = document.createElement("div");
@@ -247,8 +248,7 @@ export class CellOutputContextProvider extends AIContextProvider<CellOutputConte
             "italic",
             "mb-2",
           );
-          mediaDiv.textContent =
-            "Contains media content (image, SVG, or canvas)";
+          mediaDiv.textContent = "A screenshot of the output will be attached";
           infoContainer.append(mediaDiv);
         }
 

package/src/core/ai/context/providers/common.ts

@@ -1,13 +1,22 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
+import type { CompletionSection } from "@codemirror/autocomplete";
+
 /** Number from -99 to 99. Higher numbers are prioritized when surfacing completions. */
 export const Boosts = {
-  LOCAL_TABLE: 5,
-  REMOTE_TABLE: 4,
+  LOCAL_TABLE: 7,
+  REMOTE_TABLE: 5,
   HIGH: 4,
-  VARIABLE: 3,
   MEDIUM: 3,
   CELL_OUTPUT: 2,
   LOW: 2,
-  ERROR: 1,
 } as const;
+
+export const Sections = {
+  ERROR: { name: "Error", rank: 1 },
+  TABLE: { name: "Table", rank: 2 },
+  DATA_SOURCES: { name: "Data Sources", rank: 3 },
+  VARIABLE: { name: "Variable", rank: 4 },
+  CELL_OUTPUT: { name: "Cell Output", rank: 5 },
+  FILE: { name: "File", rank: 6 },
+} satisfies Record<string, CompletionSection>;