@marimo-team/frontend 0.15.3 → 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/tool-call-accordion.tsx

@@ -65,7 +65,7 @@ export const ToolCallAccordion: React.FC<ToolCallAccordionProps> = ({
       key={`tool-${index}`}
       type="single"
       collapsible={true}
-      className="w-full my-4"
+      className="w-full"
     >
       <AccordionItem value="tool-call" className="border-0">
         <AccordionTrigger

package/src/components/editor/chrome/types.ts

@@ -8,6 +8,7 @@ import {
   FileTextIcon,
   FolderTreeIcon,
   FunctionSquareIcon,
+  HatGlassesIcon,
   KeyRoundIcon,
   type LucideIcon,
   NetworkIcon,
@@ -17,6 +18,7 @@ import {
   TextSearchIcon,
   XCircleIcon,
 } from "lucide-react";
+import { getFeatureFlag } from "@/core/config/feature-flag";
 import { isWasm } from "@/core/wasm/utils";
 
 export type PanelType =
@@ -32,6 +34,7 @@ export type PanelType =
   | "datasources"
   | "scratchpad"
   | "chat"
+  | "agents"
   | "secrets"
   | "logs";
 
@@ -86,6 +89,13 @@ export const PANELS: PanelDescriptor[] = [
     tooltip: "Chat with AI",
     position: "sidebar",
   },
+  {
+    type: "agents",
+    Icon: HatGlassesIcon,
+    tooltip: "Agents",
+    position: "sidebar",
+    hidden: !getFeatureFlag("external_agents"),
+  },
   {
     type: "documentation",
     Icon: TextSearchIcon,

package/src/components/editor/chrome/wrapper/app-chrome.tsx

@@ -1,5 +1,10 @@
 /* Copyright 2024 Marimo. All rights reserved. */
-import React, { type PropsWithChildren, Suspense, useEffect } from "react";
+import React, {
+  type PropsWithChildren,
+  Suspense,
+  useEffect,
+  useId,
+} from "react";
 import {
   type ImperativePanelHandle,
   Panel,
@@ -25,6 +30,9 @@ import { handleDragging } from "./utils";
 
 const LazyTerminal = React.lazy(() => import("@/components/terminal/terminal"));
 const LazyChatPanel = React.lazy(() => import("@/components/chat/chat-panel"));
+const LazyAgentPanel = React.lazy(
+  () => import("@/components/chat/acp/agent-panel"),
+);
 const LazyDependencyGraphPanel = React.lazy(
   () => import("@/components/editor/chrome/panels/dependency-graph-panel"),
 );
@@ -55,6 +63,9 @@ export const AppChrome: React.FC<PropsWithChildren> = ({ children }) => {
   const sidebarRef = React.useRef<ImperativePanelHandle>(null);
   const terminalRef = React.useRef<ImperativePanelHandle>(null);
 
+  const helperPanelId = useId();
+  const terminalPanelId = useId();
+
   // sync sidebar
   useEffect(() => {
     if (!sidebarRef.current) {
@@ -161,6 +172,7 @@ export const AppChrome: React.FC<PropsWithChildren> = ({ children }) => {
             {selectedPanel === "snippets" && <LazySnippetsPanel />}
             {selectedPanel === "scratchpad" && <LazyScratchpadPanel />}
             {selectedPanel === "chat" && <LazyChatPanel />}
+            {selectedPanel === "agents" && <LazyAgentPanel />}
             {selectedPanel === "logs" && <LazyLogsPanel />}
             {selectedPanel === "tracing" && <LazyTracingPanel />}
             {selectedPanel === "secrets" && <LazySecretsPanel />}
@@ -173,7 +185,8 @@ export const AppChrome: React.FC<PropsWithChildren> = ({ children }) => {
   const helperPanel = (
     <Panel
       ref={sidebarRef}
-      id="helper"
+      id={`helper-${helperPanelId}`}
+      data-testid="helper"
       key={"helper"}
       collapsedSize={0}
       collapsible={true}
@@ -203,7 +216,8 @@ export const AppChrome: React.FC<PropsWithChildren> = ({ children }) => {
   const terminalPanel = (
     <Panel
       ref={terminalRef}
-      id="terminal"
+      id={`terminal-${terminalPanelId}`}
+      data-testid="terminal"
       key={"terminal"}
       collapsedSize={0}
       collapsible={true}

package/src/core/config/feature-flag.tsx

@@ -13,6 +13,7 @@ export interface ExperimentalFeatures {
   performant_table_charts: boolean;
   mcp_docs: boolean;
   sql_linter: boolean;
+  external_agents: boolean;
   // Add new feature flags here
 }
 
@@ -24,6 +25,7 @@ const defaultValues: ExperimentalFeatures = {
   performant_table_charts: false,
   mcp_docs: false,
   sql_linter: false,
+  external_agents: import.meta.env.DEV,
 };
 
 export function getFeatureFlag<T extends keyof ExperimentalFeatures>(