@marimo-team/islands 0.15.2 → 0.15.3

package/src/core/ai/context/registry.ts

@@ -2,6 +2,10 @@
 
 import type { Completion } from "@codemirror/autocomplete";
 import type { Resource } from "@marimo-team/codemirror-mcp";
+import type { FileUIPart } from "ai";
+import { Memoize } from "typescript-memoize";
+import { Logger } from "@/utils/Logger";
+import { MultiMap } from "@/utils/multi-map";
 import type { TypedString } from "@/utils/typed";
 
 /**
@@ -41,6 +45,12 @@ export abstract class AIContextProvider<
   /** Format completion */
   abstract formatCompletion(item: T): Completion;
 
+  /** Get attachments for context items (optional, async) */
+  async getAttachments(_items: T[]): Promise<FileUIPart[]> {
+    // Default implementation returns no attachments
+    return [];
+  }
+
   asURI(id: string): ContextLocatorId {
     return `${this.contextType}://${id}` as ContextLocatorId;
   }
@@ -118,6 +128,7 @@ export class AIContextRegistry<T extends AIContextItem> {
     );
   }
 
+  @Memoize()
   getAllItems(): T[] {
     return [...this.providers].flatMap((provider) => provider.getItems());
   }
@@ -177,4 +188,72 @@ export class AIContextRegistry<T extends AIContextItem> {
       })
       .join("\n\n");
   }
+
+  /**
+   * Get attachments for mentioned items
+   */
+  async getAttachmentsForContext(
+    contextIds: ContextLocatorId[],
+  ): Promise<FileUIPart[]> {
+    const allItems = new Map<ContextLocatorId, T>(
+      this.getAllItems().map((item) => [item.uri as ContextLocatorId, item]),
+    );
+
+    const contextInfo: T[] = [];
+    for (const contextId of contextIds) {
+      const item = allItems.get(contextId);
+      if (item) {
+        contextInfo.push(item);
+      }
+    }
+
+    if (contextInfo.length === 0) {
+      return [];
+    }
+
+    // Group items by provider type to batch attachment requests
+    const itemsByProvider = new MultiMap<string, T>();
+    for (const item of contextInfo) {
+      const providerType = item.type;
+      itemsByProvider.add(providerType, item);
+    }
+
+    // Collect attachments from all providers
+    const attachmentPromises = [...itemsByProvider.entries()].map(
+      async ([providerType, items]) => {
+        const provider = this.getProvider(providerType);
+        if (!provider) {
+          return [];
+        }
+        try {
+          return await provider.getAttachments(items);
+        } catch (error) {
+          Logger.error("Error getting attachments from provider", error);
+          return [];
+        }
+      },
+    );
+
+    const attachmentResults = await Promise.all(attachmentPromises);
+    const results = attachmentResults.flat();
+
+    // Print attachments to the console with a rich image preview
+    if (import.meta.env.DEV) {
+      for (const attachment of results) {
+        // If it's an image, print a rich preview
+        if (
+          /^data:image\/(png|jpeg|jpg|gif|svg\+xml);base64,/.test(
+            attachment.url,
+          )
+        ) {
+          Logger.debug(
+            "%c ",
+            `font-size:1px;padding:140px 180px;background:url('${attachment.url}') no-repeat;background-size:contain;`,
+          );
+        }
+      }
+    }
+
+    return results;
+  }
 }

package/src/core/ai/state.ts

@@ -1,13 +1,15 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import type { Message as AIMessage } from "@ai-sdk/react";
+import type { UIMessage } from "@ai-sdk/react";
+import type { FileUIPart } from "ai";
 import { atom } from "jotai";
 import { atomWithStorage } from "jotai/utils";
+import { uniqueBy } from "@/utils/arrays";
 import { adaptForLocalStorage } from "@/utils/storage";
 import type { TypedString } from "@/utils/typed";
 import type { CellId } from "../cells/ids";
 
-const KEY = "marimo:ai:chatState:v4";
+const KEY = "marimo:ai:chatState:v5";
 
 export type ChatId = TypedString<"ChatId">;
 
@@ -27,13 +29,14 @@ export interface Message {
   role: "user" | "assistant" | "data" | "system";
   content: string;
   timestamp: number;
-  parts?: AIMessage["parts"];
+  parts?: UIMessage["parts"];
+  attachments?: FileUIPart[];
 }
 
 export interface Chat {
   id: ChatId;
   title: string;
-  messages: Message[];
+  messages: UIMessage[];
   createdAt: number;
   updatedAt: number;
 }
@@ -43,6 +46,20 @@ export interface ChatState {
   activeChatId: ChatId | null;
 }
 
+function removeEmptyChats(chatState: Map<ChatId, Chat>): Map<ChatId, Chat> {
+  const result = new Map<ChatId, Chat>();
+
+  // Dedupe messages with the same id
+  for (const [chatId, chat] of chatState.entries()) {
+    if (chat.messages.length === 0) {
+      continue;
+    }
+    const dedupedMessages = uniqueBy(chat.messages, (message) => message.id);
+    result.set(chatId, { ...chat, messages: dedupedMessages });
+  }
+  return result;
+}
+
 export const chatStateAtom = atomWithStorage<ChatState>(
   KEY,
   {
@@ -51,7 +68,7 @@ export const chatStateAtom = atomWithStorage<ChatState>(
   },
   adaptForLocalStorage({
     toSerializable: (value: ChatState) => ({
-      chats: [...value.chats.entries()],
+      chats: [...removeEmptyChats(value.chats).entries()],
       activeChatId: value.activeChatId,
     }),
     fromSerializable: (value) => ({

package/src/core/alerts/state.ts

@@ -16,6 +16,13 @@ export interface MissingPackageAlert {
 export interface InstallingPackageAlert {
   kind: "installing";
   packages: PackageInstallationStatus;
+  logs?: { [key: string]: string } | null;
+  log_status?: "append" | "start" | "done" | null;
+}
+
+export interface StartupLogsAlert {
+  content: string;
+  status: "append" | "start" | "done";
 }
 
 export function isMissingPackageAlert(
@@ -39,26 +46,87 @@ interface AlertState {
     | Identified<MissingPackageAlert>
     | Identified<InstallingPackageAlert>
     | null;
+  startupLogsAlert: StartupLogsAlert | null;
+  packageLogs: { [packageName: string]: string };
 }
 
 const { valueAtom: alertAtom, useActions } = createReducerAndAtoms(
-  () => ({ packageAlert: null }) as AlertState,
+  () =>
+    ({
+      packageAlert: null,
+      startupLogsAlert: null,
+      packageLogs: {},
+    }) as AlertState,
   {
     addPackageAlert: (
       state,
       alert: MissingPackageAlert | InstallingPackageAlert,
     ) => {
+      const newPackageLogs = { ...state.packageLogs };
+
+      // Handle streaming logs for installing package alerts
+      if (isInstallingPackageAlert(alert) && alert.logs && alert.log_status) {
+        for (const [packageName, newContent] of Object.entries(alert.logs)) {
+          switch (alert.log_status) {
+            case "start":
+              // Start new log for this package
+              newPackageLogs[packageName] = newContent;
+
+              break;
+
+            case "append": {
+              // Append to existing log
+              const prevContent = newPackageLogs[packageName] || "";
+              newPackageLogs[packageName] = prevContent + newContent;
+
+              break;
+            }
+            case "done": {
+              // Append final content and mark as done
+              const prevContent = newPackageLogs[packageName] || "";
+              newPackageLogs[packageName] = prevContent + newContent;
+
+              break;
+            }
+            // No default
+          }
+        }
+      }
+
+      const existingAlert = state.packageAlert;
+      const alertId = existingAlert?.id || generateUUID();
+
       return {
         ...state,
-        packageAlert: { id: generateUUID(), ...alert },
+        packageAlert: { id: alertId, ...alert },
+        packageLogs: newPackageLogs,
       };
     },
 
     clearPackageAlert: (state, id: string) => {
       return state.packageAlert !== null && state.packageAlert.id === id
-        ? { ...state, packageAlert: null }
+        ? { ...state, packageAlert: null, packageLogs: {} }
         : state;
     },
+
+    addStartupLog: (
+      state,
+      logData: { content: string; status: "append" | "start" | "done" },
+    ) => {
+      const prevContent = state.startupLogsAlert?.content || "";
+      return {
+        ...state,
+        startupLogsAlert: {
+          ...state.startupLogsAlert,
+          content: prevContent + logData.content,
+          status: logData.status,
+        },
+      };
+    },
+
+    clearStartupLogsAlert: (state) => {
+      return { ...state, startupLogsAlert: null };
+    },
   },
 );
 

package/src/core/cells/cell.ts

@@ -47,7 +47,7 @@ export function transitionCell(
     case "idle":
       if (cell.runStartTimestamp) {
         nextCell.runElapsedTimeMs = Time.fromSeconds(
-          (message.timestamp - cell.runStartTimestamp) as Seconds,
+          ((message.timestamp ?? 0) - cell.runStartTimestamp) as Seconds,
         ).toMilliseconds();
         nextCell.runStartTimestamp = null;
         nextCell.staleInputs = false;
@@ -210,7 +210,7 @@ export function outputIsStale(
     status === "running" &&
     output !== null &&
     runStartTimestamp !== null &&
-    output.timestamp > runStartTimestamp;
+    (output.timestamp ?? 0) > runStartTimestamp;
 
   // If loading and output has not been received while running
   if (loading && !outputReceivedWhileRunning) {

package/src/core/cells/cells.ts

@@ -1372,7 +1372,7 @@ const {
 
 function isCellCodeHidden(state: NotebookState, cellId: CellId): boolean {
   return (
-    state.cellData[cellId].config.hide_code &&
+    Boolean(state.cellData[cellId].config.hide_code) &&
     !state.untouchedNewCells.has(cellId)
   );
 }

package/src/core/cells/logs.ts

@@ -56,7 +56,7 @@ export function getCellLogsForMessage(cell: CellMessage): CellLog[] {
       CellLogLogger.log({
         level: "stderr",
         cellId: cell.cell_id as CellId,
-        timestamp: cell.timestamp,
+        timestamp: cell.timestamp ?? 0,
         message: JSON.stringify(error),
       });
     });

package/src/core/cells/runs.ts

@@ -45,6 +45,7 @@ const {
     opts: { cellOperation: CellMessage; code: string },
   ): RunsState => {
     const { cellOperation, code } = opts;
+    const timestamp = cellOperation.timestamp ?? 0;
     const runId = cellOperation.run_id as RunId | undefined;
     if (!runId) {
       return state;
@@ -84,11 +85,11 @@ const {
               code: code.slice(0, MAX_CODE_LENGTH),
               elapsedTime: 0,
               status: status,
-              startTime: cellOperation.timestamp,
+              startTime: timestamp,
             },
           ],
         ]),
-        runStartTime: cellOperation.timestamp,
+        runStartTime: timestamp,
       };
 
       // Manage run history size
@@ -129,12 +130,12 @@ const {
 
       const startTime =
         cellOperation.status === "running"
-          ? cellOperation.timestamp
+          ? timestamp
           : existingCellRun.startTime;
 
       const elapsedTime =
         status === "success" || status === "error"
-          ? cellOperation.timestamp - existingCellRun.startTime
+          ? timestamp - existingCellRun.startTime
           : undefined;
 
       nextCellRuns.set(cellOperation.cell_id as CellId, {
@@ -149,7 +150,7 @@ const {
         code: code.slice(0, MAX_CODE_LENGTH),
         elapsedTime: 0,
         status: status,
-        startTime: cellOperation.timestamp,
+        startTime: timestamp,
       });
     }
 

package/src/core/codemirror/ai/resources.ts

@@ -12,30 +12,72 @@ import {
   resourcesField,
   resourceTheme,
 } from "@marimo-team/codemirror-mcp";
-import { getAIContextRegistry } from "@/core/ai/context/context";
+import {
+  getAIContextRegistry,
+  getFileContextProvider,
+} from "@/core/ai/context/context";
 import type { AIContextItem } from "@/core/ai/context/registry";
 import type { JotaiStore } from "@/core/state/jotai";
 import { Logger } from "@/utils/Logger";
+import { contextCallbacks } from "./state";
+
+const NONE_RESOURCE_TYPE = "_none_";
+const NONE_RESOURCE = [
+  {
+    uri: "",
+    name: "No resources",
+    type: NONE_RESOURCE_TYPE,
+    data: {},
+  },
+];
+const NONE_RESOURCE_FORMAT_COMPLETION = {
+  info: "Variables, dataframes, and tables will appear here.",
+  apply: () => {
+    return;
+  },
+};
+
+export function resourceExtension(opts: {
+  language: Language;
+  store: JotaiStore;
+  onAddFiles?: (files: File[]) => void;
+}): Extension[] {
+  const { language, store, onAddFiles } = opts;
 
-export function resourceExtension(
-  language: Language,
-  store: JotaiStore,
-): Extension[] {
   return [
     language.data.of({
+      // Resource completion for static resources (variables, tables, etc.)
       autocomplete: resourceCompletion(
         async (): Promise<Resource[]> => {
           const registry = getAIContextRegistry(store);
           const resources = registry.getAllItems();
+          if (resources.length === 0) {
+            return NONE_RESOURCE;
+          }
           return resources;
         },
         (resource) => {
+          if (resource.type === NONE_RESOURCE_TYPE) {
+            return NONE_RESOURCE_FORMAT_COMPLETION;
+          }
+
           const registry = getAIContextRegistry(store);
           const provider = registry.getProvider(resource.type);
           return provider?.formatCompletion(resource as AIContextItem) || {};
         },
       ),
     }),
+    contextCallbacks.of({
+      addAttachment: (attachment) => onAddFiles?.([attachment]),
+    }),
+    // Dynamic file completion
+    ...(onAddFiles
+      ? [
+          language.data.of({
+            autocomplete: getFileContextProvider().createCompletionSource(),
+          }),
+        ]
+      : []),
     resourceDecorations,
     resourceInputFilter,
     resourcesField.init(() => {

package/src/core/codemirror/ai/state.ts

@@ -0,0 +1,12 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { singleFacet } from "../facet";
+
+interface ContextCallbacks {
+  addAttachment?: (attachment: File) => void;
+}
+
+/**
+ * State for completion callbacks
+ */
+export const contextCallbacks = singleFacet<ContextCallbacks>();

package/src/core/codemirror/language/__tests__/sql.test.ts

@@ -664,6 +664,9 @@ describe("tablesCompletionSource", () => {
                   source: "duckdb",
                   source_type: "local",
                   type: "table",
+                  num_columns: 0,
+                  num_rows: 0,
+                  variable_name: null,
                   columns: [
                     {
                       name: "id",
@@ -690,6 +693,9 @@ describe("tablesCompletionSource", () => {
                   source: "duckdb",
                   source_type: "local",
                   type: "table",
+                  num_columns: 0,
+                  num_rows: 0,
+                  variable_name: null,
                   columns: [
                     {
                       name: "order_id",
@@ -763,6 +769,9 @@ describe("tablesCompletionSource", () => {
                   source: "postgres",
                   source_type: "local",
                   type: "table",
+                  num_columns: 0,
+                  num_rows: 0,
+                  variable_name: null,
                   columns: [
                     {
                       name: "col1",
@@ -788,6 +797,9 @@ describe("tablesCompletionSource", () => {
                   source: "postgres",
                   source_type: "local",
                   type: "table",
+                  num_columns: 0,
+                  num_rows: 0,
+                  variable_name: null,
                   columns: [
                     {
                       name: "col2",
@@ -856,6 +868,9 @@ describe("tablesCompletionSource", () => {
                   source: "postgres",
                   source_type: "local",
                   type: "table",
+                  num_columns: 0,
+                  num_rows: 0,
+                  variable_name: null,
                   columns: [
                     {
                       name: "col1",
@@ -875,6 +890,9 @@ describe("tablesCompletionSource", () => {
                   source: "postgres",
                   source_type: "local",
                   type: "table",
+                  num_columns: 0,
+                  num_rows: 0,
+                  variable_name: null,
                   columns: [
                     {
                       name: "col2",
@@ -900,6 +918,9 @@ describe("tablesCompletionSource", () => {
                   source: "postgres",
                   source_type: "local",
                   type: "table",
+                  num_columns: 0,
+                  num_rows: 0,
+                  variable_name: null,
                   columns: [
                     {
                       name: "col2",
@@ -925,6 +946,9 @@ describe("tablesCompletionSource", () => {
                   source: "postgres",
                   source_type: "local",
                   type: "table",
+                  num_columns: 0,
+                  num_rows: 0,
+                  variable_name: null,
                   columns: [
                     {
                       name: "col2",
@@ -1002,6 +1026,9 @@ describe("tablesCompletionSource", () => {
                   source: "postgres",
                   source_type: "local",
                   type: "table",
+                  num_columns: 0,
+                  num_rows: 0,
+                  variable_name: null,
                   columns: [
                     {
                       name: "id",
@@ -1070,6 +1097,9 @@ describe("tablesCompletionSource", () => {
                   source: "postgres",
                   source_type: "local",
                   type: "table",
+                  num_columns: 0,
+                  num_rows: 0,
+                  variable_name: null,
                   columns: [],
                 },
               ],
@@ -1109,6 +1139,9 @@ describe("tablesCompletionSource", () => {
                   source: "postgres",
                   source_type: "local",
                   type: "table",
+                  num_columns: 0,
+                  num_rows: 0,
+                  variable_name: null,
                   columns: [
                     {
                       name: "id",
@@ -1134,6 +1167,9 @@ describe("tablesCompletionSource", () => {
                   source: "postgres",
                   source_type: "local",
                   type: "table",
+                  num_columns: 0,
+                  num_rows: 0,
+                  variable_name: null,
                   columns: [
                     {
                       name: "order_id",
@@ -1279,6 +1315,9 @@ describe("tablesCompletionSource", () => {
                         source: "postgres",
                         source_type: "local",
                         type: "table",
+                        num_columns: 0,
+                        num_rows: 0,
+                        variable_name: null,
                         columns: [
                           {
                             name: "id",
@@ -1469,6 +1508,9 @@ describe("tablesCompletionSource", () => {
                   source: "duckdb",
                   source_type: "local",
                   type: "table",
+                  num_columns: 0,
+                  num_rows: 0,
+                  variable_name: null,
                   columns: [
                     {
                       name: "col1",
@@ -1613,6 +1655,9 @@ const mockConnection: DataSourceConnection = {
               source: "duckdb",
               source_type: "local",
               type: "table",
+              num_columns: 0,
+              num_rows: 0,
+              variable_name: null,
               columns: [
                 {
                   name: "col1",

package/src/core/codemirror/markdown/__tests__/commands.test.ts

@@ -267,6 +267,7 @@ describe("insertImage", () => {
 
     mockRequestClient.sendCreateFileOrFolder.mockResolvedValueOnce({
       success: true,
+      message: null,
       info: {
         path: "nested/public/hello.png",
         name: "hello.png",

package/src/core/codemirror/theme/dark.ts

@@ -8,7 +8,7 @@ export const darkTheme = [
   createTheme({
     variant: "dark",
     settings: {
-      background: "#282c34",
+      background: "var(--cm-background)",
       foreground: "#abb2bf",
       caret: "#528bff",
       selection: "#3E4451",

package/src/core/config/capabilities.ts

@@ -11,5 +11,5 @@ export const capabilitiesAtom = atom<Capabilities>({
 });
 
 export function hasCapability(key: keyof Capabilities): boolean {
-  return store.get(capabilitiesAtom)[key];
+  return store.get(capabilitiesAtom)[key] ?? false;
 }