@marimo-team/islands 0.16.4 → 0.16.5

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/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/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

@@ -8,7 +8,7 @@ import type {
 } 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 { Boosts, Sections } from "../common";
 import { DatasourceContextProvider } from "../datasource";
 
 // Mock data for testing
@@ -268,8 +268,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 +310,7 @@ describe("DatasourceContextProvider", () => {
         detail: "DuckDB",
         boost: Boosts.MEDIUM,
         type: "datasource",
-        section: "Data Sources",
+        section: Sections.DATA_SOURCES,
       });
 
       expect(completion.info).toBeDefined();
@@ -355,7 +354,7 @@ describe("DatasourceContextProvider", () => {
         detail: "PostgreSQL",
         boost: Boosts.MEDIUM,
         type: "datasource",
-        section: "Data Sources",
+        section: Sections.DATA_SOURCES,
       });
     });
 
@@ -380,7 +379,7 @@ describe("DatasourceContextProvider", () => {
         detail: "DuckDB",
         boost: Boosts.MEDIUM,
         type: "datasource",
-        section: "Data Sources",
+        section: Sections.DATA_SOURCES,
       });
     });
   });

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>;

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

@@ -17,7 +17,7 @@ import type { DataSourceConnection, DataTable } from "@/core/kernel/messages";
 import type { AIContextItem } from "../registry";
 import { AIContextProvider } from "../registry";
 import { contextToXml } from "../utils";
-import { Boosts } from "./common";
+import { Boosts, Sections } from "./common";
 
 type NamedDatasource = Omit<
   DataSourceConnection,
@@ -53,24 +53,35 @@ export class DatasourceContextProvider extends AIContextProvider<DatasourceConte
   }
 
   getItems(): DatasourceContextItem[] {
-    return [...this.connectionsMap.values()].map((connection) => {
-      let description = "Database schema.";
-      const data: DatasourceContextItem["data"] = {
-        connection: connection,
-      };
-      if (INTERNAL_SQL_ENGINES.has(connection.name)) {
-        data.tables = this.dataframes;
-        description = "Database schema and the dataframes that can be queried";
-      }
-
-      return {
-        uri: this.asURI(connection.name),
-        name: connection.name,
-        description: description,
-        type: this.contextType,
-        data: data,
-      };
-    });
+    return [...this.connectionsMap.values()]
+      .map((connection): DatasourceContextItem | null => {
+        let description = "Database schema.";
+        const data: DatasourceContextItem["data"] = {
+          connection: connection,
+        };
+
+        if (INTERNAL_SQL_ENGINES.has(connection.name)) {
+          data.tables = this.dataframes;
+          description =
+            "Database schema and the dataframes that can be queried";
+        }
+
+        // Hide empty datasources
+        const hasNoTables =
+          connection.databases.length === 0 && (data.tables?.length ?? 0) === 0;
+        if (hasNoTables) {
+          return null;
+        }
+
+        return {
+          uri: this.asURI(connection.name),
+          name: connection.name,
+          description: description,
+          type: this.contextType,
+          data: data,
+        };
+      })
+      .filter(Boolean);
   }
 
   formatContext(item: DatasourceContextItem): string {
@@ -115,7 +126,7 @@ export class DatasourceContextProvider extends AIContextProvider<DatasourceConte
       detail: dbDisplayName(dataConnection.dialect),
       boost: Boosts.MEDIUM,
       type: this.contextType,
-      section: "Data Sources",
+      section: Sections.DATA_SOURCES,
       info: () => {
         const infoContainer = document.createElement("div");
         infoContainer.classList.add("mo-cm-tooltip", "docs-documentation");

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

@@ -9,7 +9,7 @@ import { logNever } from "@/utils/assertNever";
 import { PluralWord } from "@/utils/pluralize";
 import { type AIContextItem, AIContextProvider } from "../registry";
 import { contextToXml } from "../utils";
-import { Boosts } from "./common";
+import { Sections } from "./common";
 
 export interface ErrorContextItem extends AIContextItem {
   type: "error";
@@ -111,10 +111,9 @@ export class ErrorContextProvider extends AIContextProvider<ErrorContextItem> {
         label: "@Errors",
         displayLabel: "Errors",
         detail: `${item.data.errors.length} ${errorsTxt.pluralize(item.data.errors.length)}`,
-        boost: Boosts.ERROR,
         type: "error",
         apply: "@Errors",
-        section: "Errors",
+        section: Sections.ERROR,
         info: () => {
           const infoContainer = document.createElement("div");
           infoContainer.classList.add(
@@ -150,7 +149,7 @@ export class ErrorContextProvider extends AIContextProvider<ErrorContextItem> {
     return {
       label: "Error",
       displayLabel: "Error",
-      boost: Boosts.ERROR,
+      section: Sections.ERROR,
     };
   }
 

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

@@ -15,7 +15,7 @@ import { type Base64String, base64ToDataURL } from "@/utils/json/base64";
 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 FileContextItem extends AIContextItem {
   type: "file";
   data: {
@@ -190,7 +190,7 @@ export class FileContextProvider extends AIContextProvider<FileContextItem> {
     return {
       ...this.createBasicCompletion(item),
       type: "file",
-      section: "File",
+      section: Sections.FILE,
       boost: data.isDirectory ? Boosts.MEDIUM : Boosts.LOW,
       detail: data.path,
       displayLabel: `${icon} ${name}`,

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

@@ -9,7 +9,7 @@ import type { DataTable } from "@/core/kernel/messages";
 import type { AIContextItem } from "../registry";
 import { AIContextProvider } from "../registry";
 import { contextToXml } from "../utils";
-import { Boosts } from "./common";
+import { Boosts, Sections } from "./common";
 
 export interface TableContextItem extends AIContextItem {
   type: "data";
@@ -38,21 +38,46 @@ export class TableContextProvider extends AIContextProvider<TableContextItem> {
 
   formatContext(item: TableContextItem): string {
     const { data } = item;
-    const { columns, source, num_rows, num_columns, name } = data;
+    const { columns, source, num_rows, num_columns, name, variable_name } =
+      data;
+
+    // Build shape information
     const shape = [
-      num_rows == null ? undefined : `${num_rows} rows`,
-      num_columns == null ? undefined : `${num_columns} columns`,
+      num_rows != null ? `${num_rows} rows` : undefined,
+      num_columns != null ? `${num_columns} columns` : undefined,
     ]
       .filter(Boolean)
       .join(", ");
 
     let details = "";
+
+    // Add shape information
     if (shape) {
-      details += `\nShape: ${shape}`;
+      details += `Shape: ${shape}\n`;
+    }
+
+    // Add variable name if available
+    if (variable_name) {
+      details += `Variable: ${variable_name}\n`;
     }
 
+    // Add column information with sample values
     if (columns && columns.length > 0) {
-      details += `\nColumns:\n${columns.map((col) => `  - ${col.name}: ${col.type}`).join("\n")}`;
+      details += "Columns:\n";
+      for (const col of columns) {
+        let columnInfo = `  ${col.name} (${col.type})`;
+
+        // Add sample values if available
+        if (col.sample_values && col.sample_values.length > 0) {
+          const samples = col.sample_values
+            .slice(0, 3) // Limit to first 3 samples
+            .map((val) => (val === null ? "null" : String(val)))
+            .join(", ");
+          columnInfo += ` - samples: [${samples}]`;
+        }
+
+        details += `${columnInfo}\n`;
+      }
     }
 
     return contextToXml({
@@ -61,7 +86,7 @@ export class TableContextProvider extends AIContextProvider<TableContextItem> {
         name: name,
         source: source ?? "unknown",
       },
-      details: details,
+      details: details.trim(),
     });
   }
 
@@ -78,7 +103,10 @@ export class TableContextProvider extends AIContextProvider<TableContextItem> {
           : Boosts.REMOTE_TABLE,
       type: getTableType(table),
       apply: `@${tableName}`,
-      section: getTableType(table) === "dataframe" ? "Dataframe" : "Table",
+      section: {
+        name: getTableType(table) === "dataframe" ? "Dataframe" : "Table",
+        rank: Sections.TABLE.rank,
+      },
       info: () => this.createTableInfoElement(tableName, table),
     };
   }

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

@@ -6,7 +6,7 @@ import type { DatasetTablesMap } from "@/core/datasets/data-source-connections";
 import type { Variable, Variables } from "@/core/variables/types";
 import { type AIContextItem, AIContextProvider } from "../registry";
 import { contextToXml } from "../utils";
-import { Boosts } from "./common";
+import { Sections } from "./common";
 
 export interface VariableContextItem extends AIContextItem {
   type: "variable";
@@ -58,9 +58,8 @@ export class VariableContextProvider extends AIContextProvider<VariableContextIt
       label: `@${variable.name}`,
       displayLabel: variable.name,
       detail: variable.dataType ?? "",
-      boost: Boosts.VARIABLE,
       type: this.contextType,
-      section: "Variable",
+      section: Sections.VARIABLE,
       info: () => {
         return createVariableInfoElement(variable);
       },

package/src/core/cells/__tests__/cells.test.ts

@@ -12,7 +12,7 @@ import {
   it,
   vi,
 } from "vitest";
-import type { CellHandle } from "@/components/editor/Cell";
+import type { CellHandle } from "@/components/editor/notebook-cell";
 import { CellId } from "@/core/cells/ids";
 import { foldAllBulk, unfoldAllBulk } from "@/core/codemirror/editing/commands";
 import { adaptiveLanguageConfiguration } from "@/core/codemirror/language/extension";
@@ -2094,9 +2094,9 @@ describe("cell reducer", () => {
     `);
   });
 
-  it("can create and update a setup cell", () => {
+  it("can create and noop-update a setup cell", () => {
     // Create the setup cell
-    actions.upsertSetupCell({ code: "# Setup code" });
+    actions.addSetupCellIfDoesntExist({ code: "# Setup code" });
 
     // Check that setup cell was created
     expect(state.cellData[SETUP_CELL_ID].id).toBe(SETUP_CELL_ID);
@@ -2106,17 +2106,17 @@ describe("cell reducer", () => {
     expect(state.cellIds.inOrderIds).toContain(SETUP_CELL_ID);
 
     // Update the setup cell
-    actions.upsertSetupCell({ code: "# Updated setup code" });
+    actions.addSetupCellIfDoesntExist({ code: "# Updated setup code" });
 
     // Check that the same setup cell was updated, not duplicated
-    expect(state.cellData[SETUP_CELL_ID].code).toBe("# Updated setup code");
+    expect(state.cellData[SETUP_CELL_ID].code).toBe("# Setup code");
     expect(state.cellData[SETUP_CELL_ID].edited).toBe(true);
     expect(state.cellIds.inOrderIds).toContain(SETUP_CELL_ID);
   });
 
   it("can delete and undelete the setup cell", () => {
     // Create the setup cell
-    actions.upsertSetupCell({ code: "# Setup code" });
+    actions.addSetupCellIfDoesntExist({ code: "# Setup code" });
 
     // Check that setup cell was created
     expect(state.cellData[SETUP_CELL_ID].id).toBe(SETUP_CELL_ID);