@marimo-team/islands 0.15.2 → 0.15.4

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

@@ -0,0 +1,263 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { atom } from "jotai";
+import { atomWithStorage } from "jotai/utils";
+import { capitalize } from "lodash-es";
+import type { TypedString } from "@/utils/typed";
+import { generateUUID } from "@/utils/uuid";
+import type { ExternalAgentSessionId, SessionSupportType } from "./types";
+
+// Types
+export type TabId = TypedString<"TabId">;
+export type ExternalAgentId = "claude" | "gemini";
+
+// No agents support loading sessions, so we limit to 1, otherwise
+// this is confusing to the user when switching between sessions
+const MAX_SESSIONS = 1;
+
+export interface AgentSession {
+  tabId: TabId;
+  agentId: ExternalAgentId;
+  title: string;
+  createdAt: number;
+  lastUsedAt: number;
+  // Store the actual agent session ID for resumption
+  externalAgentSessionId: ExternalAgentSessionId | null;
+}
+
+export interface AgentSessionState {
+  sessions: AgentSession[];
+  activeTabId: TabId | null;
+}
+
+// Constants
+const STORAGE_KEY = "marimo:acp:sessions:v1";
+
+// Atoms
+export const agentSessionStateAtom = atomWithStorage<AgentSessionState>(
+  STORAGE_KEY,
+  {
+    sessions: [],
+    activeTabId: null,
+  },
+);
+
+export const selectedTabAtom = atom(
+  (get) => {
+    const state = get(agentSessionStateAtom);
+    if (!state.activeTabId) {
+      return null;
+    }
+    return (
+      state.sessions.find((session) => session.tabId === state.activeTabId) ||
+      null
+    );
+  },
+  (_get, set, activeTabId: TabId | null) => {
+    set(agentSessionStateAtom, (prev) => ({
+      ...prev,
+      activeTabId: activeTabId,
+    }));
+  },
+);
+
+// Utilities
+function generateTabId(): TabId {
+  // Our tab ID for internal session management
+  return `tab_${generateUUID()}` as TabId;
+}
+
+export function truncateTitle(title: string, maxLength = 20): string {
+  if (title.length <= maxLength) {
+    return title;
+  }
+  return `${title.slice(0, maxLength - 3)}...`;
+}
+
+export function addSession(
+  state: AgentSessionState,
+  session: {
+    agentId: ExternalAgentId;
+    firstMessage?: string;
+  },
+): AgentSessionState {
+  const sessionSupport = getAgentSessionSupport(session.agentId);
+
+  const now = Date.now();
+  const title = session.firstMessage
+    ? truncateTitle(session.firstMessage.trim())
+    : `New ${session.agentId} session`;
+  const tabId = generateTabId();
+
+  if (sessionSupport === "single") {
+    // For single session agents, replace any existing session for this agent
+    const existingSessions = state.sessions.filter(
+      (s) => s.agentId === session.agentId,
+    );
+    const otherSessions = state.sessions.filter(
+      (s) => s.agentId !== session.agentId,
+    );
+
+    if (existingSessions.length > 0) {
+      // Replace the existing session (overwrite it)
+      const existingSession = existingSessions[0];
+      const updatedSession: AgentSession = {
+        agentId: session.agentId,
+        title,
+        createdAt: now,
+        lastUsedAt: now,
+        tabId: existingSession.tabId, // Keep the same ID to maintain tab reference
+        externalAgentSessionId: null, // Clear the external session ID
+      };
+
+      return {
+        ...state,
+        sessions: [...otherSessions.slice(0, MAX_SESSIONS - 1), updatedSession],
+        activeTabId: updatedSession.tabId,
+      };
+    }
+  }
+
+  // For multiple session agents or when no existing session exists
+  return {
+    ...state,
+    sessions: [
+      ...state.sessions.slice(0, MAX_SESSIONS - 1),
+      {
+        agentId: session.agentId,
+        tabId,
+        title,
+        createdAt: now,
+        lastUsedAt: now,
+        externalAgentSessionId: null,
+      },
+    ],
+    activeTabId: tabId,
+  };
+}
+
+export function removeSession(
+  state: AgentSessionState,
+  sessionId: TabId,
+): AgentSessionState {
+  const filteredSessions = state.sessions.filter((s) => s.tabId !== sessionId);
+  const newActiveTabId =
+    state.activeTabId === sessionId
+      ? filteredSessions.length > 0
+        ? filteredSessions[filteredSessions.length - 1].tabId
+        : null
+      : state.activeTabId;
+
+  return {
+    sessions: filteredSessions,
+    activeTabId: newActiveTabId,
+  };
+}
+
+export function updateSessionTitle(
+  state: AgentSessionState,
+  title: string,
+): AgentSessionState {
+  const selectedTab = state.activeTabId;
+  if (!selectedTab) {
+    return state;
+  }
+  return {
+    ...state,
+    sessions: state.sessions.map((session) =>
+      session.tabId === selectedTab
+        ? { ...session, title: truncateTitle(title) }
+        : session,
+    ),
+  };
+}
+
+export function updateSessionLastUsed(
+  state: AgentSessionState,
+  sessionId: TabId,
+): AgentSessionState {
+  return {
+    ...state,
+    sessions: state.sessions.map((session) =>
+      session.tabId === sessionId
+        ? { ...session, lastUsedAt: Date.now() }
+        : session,
+    ),
+  };
+}
+
+/**
+ * Update the sessionId for the current;y selected tab
+ */
+export function updateSessionExternalAgentSessionId(
+  state: AgentSessionState,
+  externalAgentSessionId: ExternalAgentSessionId,
+): AgentSessionState {
+  const selectedTab = state.activeTabId;
+  if (!selectedTab) {
+    return state;
+  }
+  return {
+    ...state,
+    sessions: state.sessions.map((session) =>
+      session.tabId === selectedTab
+        ? { ...session, externalAgentSessionId, lastUsedAt: Date.now() }
+        : session,
+    ),
+  };
+}
+
+export function getSessionsByAgent(
+  sessions: AgentSession[],
+  agentId: ExternalAgentId,
+): AgentSession[] {
+  return sessions
+    .filter((session) => session.agentId === agentId)
+    .sort((a, b) => b.lastUsedAt - a.lastUsedAt);
+}
+
+export function getAllAgentIds(): ExternalAgentId[] {
+  return ["claude", "gemini"];
+}
+
+export function getAgentDisplayName(agentId: ExternalAgentId): string {
+  return capitalize(agentId);
+}
+
+export function getAgentWebSocketUrl(agentId: ExternalAgentId): string {
+  return AGENT_CONFIG[agentId].webSocketUrl;
+}
+
+interface AgentConfig {
+  port: number;
+  command: string;
+  webSocketUrl: string;
+  sessionSupport: SessionSupportType;
+}
+
+const AGENT_CONFIG: Record<ExternalAgentId, AgentConfig> = {
+  claude: {
+    port: 3017,
+    command: "npx @zed-industries/claude-code-acp",
+    webSocketUrl: "ws://localhost:3017/message",
+    sessionSupport: "single",
+  },
+  gemini: {
+    port: 3019,
+    command: "npx @google/gemini-cli --experimental-acp",
+    webSocketUrl: "ws://localhost:3019/message",
+    sessionSupport: "single",
+  },
+};
+
+export function getAgentSessionSupport(
+  agentId: ExternalAgentId,
+): SessionSupportType {
+  return AGENT_CONFIG[agentId].sessionSupport;
+}
+
+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} `;
+}

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

@@ -0,0 +1,121 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import React from "react";
+import { groupNotifications } from "use-acp";
+import {
+  ConnectionChangeBlock,
+  ErrorBlock,
+  SessionNotificationsBlock,
+} from "./blocks";
+
+type NotificationEvent = Awaited<
+  ReturnType<typeof groupNotifications>
+>[number][number];
+
+interface AgentThreadProps {
+  notifications: NotificationEvent[];
+  isConnected: boolean;
+  onRetryConnection?: () => void;
+  onRetryLastAction?: () => void;
+  onDismissError?: (errorId: string) => void;
+}
+
+export const AgentThread = ({
+  notifications,
+  isConnected,
+  onRetryConnection,
+  onRetryLastAction,
+  onDismissError,
+}: AgentThreadProps) => {
+  let combinedNotifications = groupNotifications(notifications);
+
+  // Filter out all connection changes unless it is the last one
+  combinedNotifications = combinedNotifications.filter((group, index) => {
+    const isLast = index === combinedNotifications.length - 1;
+    if (isLast) {
+      return true;
+    }
+    if (isConnectionChangeGroup(group)) {
+      return false;
+    }
+    return true;
+  });
+
+  const renderNotification = (group: NotificationEvent[]) => {
+    if (group.length === 0) {
+      return null;
+    }
+
+    if (isErrorGroup(group)) {
+      const lastError = group[group.length - 1];
+      return (
+        <ErrorBlock
+          key={lastError.id}
+          data={lastError.data}
+          onRetry={onRetryLastAction}
+          onDismiss={
+            onDismissError ? () => onDismissError(lastError.id) : undefined
+          }
+        />
+      );
+    }
+    if (isConnectionChangeGroup(group)) {
+      const lastConnectionChange = group[group.length - 1];
+      return (
+        <ConnectionChangeBlock
+          key={lastConnectionChange.id}
+          data={lastConnectionChange.data}
+          isConnected={isConnected}
+          onRetry={onRetryConnection}
+          timestamp={lastConnectionChange.timestamp}
+        />
+      );
+    }
+    if (isSessionNotificationGroup(group)) {
+      const startTimestamp = group[0].timestamp;
+      const endTimestamp = group[group.length - 1].timestamp;
+      const data = group.map((item) => item.data.update);
+      return (
+        <SessionNotificationsBlock
+          data={data}
+          startTimestamp={startTimestamp}
+          endTimestamp={endTimestamp}
+        />
+      );
+    }
+    return "Unknown notification type";
+  };
+
+  return (
+    <div className="flex flex-col gap-4 px-2 pb-10">
+      {combinedNotifications.map((notification) => (
+        <React.Fragment key={notification[0].id}>
+          {renderNotification(notification)}
+        </React.Fragment>
+      ))}
+    </div>
+  );
+};
+
+function isErrorGroup(
+  group: NotificationEvent[],
+): group is Array<Extract<NotificationEvent, { type: "error" }>> {
+  // We only check the first since we know the group is the same type
+  return group[0].type === "error";
+}
+
+function isConnectionChangeGroup(
+  group: NotificationEvent[],
+): group is Array<Extract<NotificationEvent, { type: "connection_change" }>> {
+  // We only check the first since we know the group is the same type
+  return group[0].type === "connection_change";
+}
+
+function isSessionNotificationGroup(
+  group: NotificationEvent[],
+): group is Array<
+  Extract<NotificationEvent, { type: "session_notification" }>
+> {
+  // We only check the first since we know the group is the same type
+  return group[0].type === "session_notification";
+}

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

@@ -0,0 +1,63 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import type { ContentBlock } from "@zed-industries/agent-client-protocol";
+import type { groupNotifications, useAcpClient } from "use-acp";
+
+export type NotificationEvent = Awaited<
+  ReturnType<typeof groupNotifications>
+>[number][number];
+
+export type AgentConnectionState = ReturnType<
+  typeof useAcpClient
+>["connectionState"];
+
+export type AgentPendingPermission = ReturnType<
+  typeof useAcpClient
+>["pendingPermission"];
+
+// Notification events
+
+export type ErrorNotificationEvent = Extract<
+  NotificationEvent,
+  { type: "error" }
+>;
+export type ConnectionChangeNotificationEvent = Extract<
+  NotificationEvent,
+  { type: "connection_change" }
+>;
+export type SessionNotificationEvent = Extract<
+  NotificationEvent,
+  { type: "session_notification" }
+>;
+
+// Session notification events
+
+export type SessionNotificationEventData =
+  SessionNotificationEvent["data"]["update"];
+export type SessionNotificationEventType =
+  SessionNotificationEventData["sessionUpdate"];
+
+export type NotificationDataOf<T extends SessionNotificationEventType> =
+  Extract<SessionNotificationEventData, { sessionUpdate: T }>;
+
+// Session notification event data
+
+export type UserNotificationEvent = NotificationDataOf<"user_message_chunk">;
+export type AgentNotificationEvent = NotificationDataOf<"agent_message_chunk">;
+export type AgentThoughtNotificationEvent =
+  NotificationDataOf<"agent_thought_chunk">;
+export type ToolCallNotificationEvent = NotificationDataOf<"tool_call">;
+export type ToolCallUpdateNotificationEvent =
+  NotificationDataOf<"tool_call_update">;
+export type PlanNotificationEvent = NotificationDataOf<"plan">;
+
+export type ContentBlockType = ContentBlock["type"];
+export type ContentBlockOf<T extends ContentBlockType> = Extract<
+  ContentBlock,
+  { type: T }
+>;
+
+// Agent session support types
+export type SessionSupportType = "single" | "multiple";
+export type ExternalAgentSessionId = NonNullable<
+  ReturnType<typeof useAcpClient>["activeSessionId"]
+>;

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

@@ -0,0 +1,45 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import type { NotificationDataOf, SessionNotificationEventData } from "./types";
+
+export function isToolCalls(
+  group: SessionNotificationEventData[],
+): group is Array<NotificationDataOf<"tool_call" | "tool_call_update">> {
+  // We only check the first since we know the group is the same type
+  const first = group[0];
+  return (
+    first.sessionUpdate === "tool_call" ||
+    first.sessionUpdate === "tool_call_update"
+  );
+}
+
+export function isAgentThoughts(
+  group: SessionNotificationEventData[],
+): group is Array<NotificationDataOf<"agent_thought_chunk">> {
+  // We only check the first since we know the group is the same type
+  const first = group[0];
+  return first.sessionUpdate === "agent_thought_chunk";
+}
+
+export function isUserMessages(
+  group: SessionNotificationEventData[],
+): group is Array<NotificationDataOf<"user_message_chunk">> {
+  // We only check the first since we know the group is the same type
+  const first = group[0];
+  return first.sessionUpdate === "user_message_chunk";
+}
+
+export function isAgentMessages(
+  group: SessionNotificationEventData[],
+): group is Array<NotificationDataOf<"agent_message_chunk">> {
+  // We only check the first since we know the group is the same type
+  const first = group[0];
+  return first.sessionUpdate === "agent_message_chunk";
+}
+
+export function isPlans(
+  group: SessionNotificationEventData[],
+): group is Array<NotificationDataOf<"plan">> {
+  // We only check the first since we know the group is the same type
+  const first = group[0];
+  return first.sessionUpdate === "plan";
+}

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

@@ -0,0 +1,71 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { FileUIPart } from "ai";
+import { FileIcon, FileTextIcon, ImageIcon, XIcon } from "lucide-react";
+import { useState } from "react";
+import { cn } from "@/utils/cn";
+
+export const AttachmentRenderer = ({
+  attachment,
+}: {
+  attachment: FileUIPart;
+}) => {
+  if (attachment.mediaType?.startsWith("image/")) {
+    return (
+      <img
+        src={attachment.url}
+        alt={attachment.filename}
+        className="max-h-[100px] max-w-[100px] object-contain mb-1.5"
+      />
+    );
+  }
+
+  return (
+    <div className="flex flex-row gap-1 items-center text-xs">
+      <FileIcon className="h-3 w-3 mt-0.5" />
+      {attachment.filename}
+    </div>
+  );
+};
+
+export const FileAttachmentPill = ({
+  file,
+  className,
+  onRemove,
+}: {
+  file: File;
+  className?: string;
+  onRemove: () => void;
+}) => {
+  const [isHovered, setIsHovered] = useState(false);
+
+  return (
+    <div
+      className={cn(
+        "py-1 px-1.5 bg-muted rounded-md cursor-pointer flex flex-row gap-1 items-center text-xs",
+        className,
+      )}
+      onMouseEnter={() => setIsHovered(true)}
+      onMouseLeave={() => setIsHovered(false)}
+    >
+      {isHovered ? (
+        <XIcon className="h-3 w-3 mt-0.5" onClick={onRemove} />
+      ) : (
+        renderFileIcon(file)
+      )}
+      {file.name}
+    </div>
+  );
+};
+
+function renderFileIcon(file: File): React.ReactNode {
+  const classNames = "h-3 w-3 mt-0.5";
+
+  if (file.type.startsWith("image/")) {
+    return <ImageIcon className={classNames} />;
+  } else if (file.type.startsWith("text/")) {
+    return <FileTextIcon className={classNames} />;
+  }
+
+  return <FileIcon className={classNames} />;
+}