@marimo-team/frontend 0.15.5 → 0.16.0-dev96986

package/src/components/chat/acp/agent-panel.tsx

@@ -1,23 +1,32 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import { useAtom } from "jotai";
+import { useAtom, useAtomValue } from "jotai";
 import { capitalize } from "lodash-es";
 import {
+  AtSignIcon,
   BotMessageSquareIcon,
+  PaperclipIcon,
   RefreshCwIcon,
+  SendIcon,
+  SquareIcon,
   StopCircleIcon,
 } from "lucide-react";
-import React, { memo, useEffect, useRef, useState } from "react";
+import React, { memo, useEffect, useMemo, useRef, useState } from "react";
 import useEvent from "react-use-event-hook";
 import { useAcpClient } from "use-acp";
 import {
   ConnectionStatus,
   PermissionRequest,
 } from "@/components/chat/acp/common";
-import { PromptInput } from "@/components/editor/ai/add-cell-with-ai";
+import {
+  type AdditionalCompletions,
+  PromptInput,
+} from "@/components/editor/ai/add-cell-with-ai";
 import { PanelEmptyState } from "@/components/editor/chrome/panels/empty-state";
 import { Spinner } from "@/components/icons/spinner";
 import { Button } from "@/components/ui/button";
+import { Input } from "@/components/ui/input";
+import { Tooltip, TooltipProvider } from "@/components/ui/tooltip";
 import { cn } from "@/utils/cn";
 import { Logger } from "@/utils/Logger";
 import { AgentDocs } from "./agent-docs";
@@ -34,27 +43,50 @@ import {
 } from "./state";
 import { AgentThread } from "./thread";
 import "./agent-panel.css";
+import type { Completion } from "@codemirror/autocomplete";
+import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
 import type {
-  ReadTextFileResponse,
+  ContentBlock,
   RequestPermissionResponse,
-  WriteTextFileResponse,
 } from "@zed-industries/agent-client-protocol";
+import {
+  addContextCompletion,
+  CONTEXT_TRIGGER,
+} from "@/components/editor/ai/completion-utils";
+import {
+  Select,
+  SelectContent,
+  SelectGroup,
+  SelectItem,
+  SelectLabel,
+  SelectTrigger,
+} from "@/components/ui/select";
 import { toast } from "@/components/ui/use-toast";
 import { useRequestClient } from "@/core/network/requests";
 import { filenameAtom } from "@/core/saving/file-state";
 import { store } from "@/core/state/jotai";
 import { Functions } from "@/utils/functions";
 import { Paths } from "@/utils/paths";
+import { FileAttachmentPill } from "../chat-components";
+import {
+  convertFilesToResourceLinks,
+  parseContextFromPrompt,
+} from "./context-utils";
 import { getAgentPrompt } from "./prompt";
 import type {
   AgentConnectionState,
   AgentPendingPermission,
+  AvailableCommands,
   ExternalAgentSessionId,
   NotificationEvent,
+  SessionMode,
 } from "./types";
 
 const logger = Logger.get("agents");
 
+// File attachment constants
+const SUPPORTED_ATTACHMENT_TYPES = ["image/*", "text/*"];
+
 interface AgentTitleProps {
   currentAgentId?: ExternalAgentId;
 }
@@ -171,35 +203,45 @@ interface EmptyStateProps {
 }
 
 const EmptyState = memo<EmptyStateProps>(
-  ({ currentAgentId, connectionState, onConnect, onDisconnect }) => (
-    <div className="flex flex-col h-full">
-      <AgentPanelHeader
-        connectionState={connectionState}
-        currentAgentId={currentAgentId}
-        onConnect={onConnect}
-        onDisconnect={onDisconnect}
-        hasActiveSession={false}
-      />
-      <SessionTabs />
-      <div className="flex-1 flex items-center justify-center p-6">
-        <div className="max-w-md w-full space-y-6">
-          <PanelEmptyState
-            title="No Agent Sessions"
-            description="Create a new session to start a conversation"
-            action={<AgentSelector className="border-y-1 rounded" />}
-            icon={<BotMessageSquareIcon />}
-          />
-          {connectionState.status === "disconnected" && (
-            <AgentDocs
-              className="border-t pt-6"
-              title="Connect to an agent"
-              description="Start agents by running these commands in your terminal:"
+  ({ currentAgentId, connectionState, onConnect, onDisconnect }) => {
+    const filename = useAtomValue(filenameAtom);
+    return (
+      <div className="flex flex-col h-full">
+        <AgentPanelHeader
+          connectionState={connectionState}
+          currentAgentId={currentAgentId}
+          onConnect={onConnect}
+          onDisconnect={onDisconnect}
+          hasActiveSession={false}
+        />
+        <SessionTabs />
+        <div className="flex-1 flex items-center justify-center p-6">
+          <div className="max-w-md w-full space-y-6">
+            <PanelEmptyState
+              title="No Agent Sessions"
+              description="Create a new session to start a conversation"
+              action={<AgentSelector className="border-y-1 rounded" />}
+              icon={<BotMessageSquareIcon />}
             />
-          )}
+            {connectionState.status === "disconnected" && (
+              <AgentDocs
+                className="border-t pt-6"
+                title="Connect to an agent"
+                description={
+                  <>
+                    Start agents by running these commands in your terminal:
+                    <br />
+                    Note: This must be in the directory{" "}
+                    {Paths.dirname(filename ?? "")}
+                  </>
+                }
+              />
+            )}
+          </div>
         </div>
       </div>
-    </div>
-  ),
+    );
+  },
 );
 EmptyState.displayName = "EmptyState";
 
@@ -248,8 +290,14 @@ interface PromptAreaProps {
   isLoading: boolean;
   activeSessionId: ExternalAgentSessionId | null;
   promptValue: string;
+  commands: AvailableCommands | undefined;
   onPromptValueChange: (value: string) => void;
   onPromptSubmit: (e: KeyboardEvent | undefined, prompt: string) => void;
+  onAddFiles: (files: File[]) => void;
+  onStop: () => void;
+  fileInputRef: React.RefObject<HTMLInputElement | null>;
+  sessionMode?: SessionMode;
+  onModeChange?: (mode: string) => void;
 }
 
 const PromptArea = memo<PromptAreaProps>(
@@ -257,36 +305,191 @@ const PromptArea = memo<PromptAreaProps>(
     isLoading,
     activeSessionId,
     promptValue,
+    commands,
     onPromptValueChange,
     onPromptSubmit,
-  }) => (
-    <div
-      className={cn(
-        "px-3 py-2 border-t bg-background flex-shrink-0 min-h-[80px]",
-        (isLoading || !activeSessionId) && "opacity-50 pointer-events-none",
-      )}
-    >
-      <PromptInput
-        value={promptValue}
-        onChange={isLoading ? Functions.NOOP : onPromptValueChange}
-        onSubmit={onPromptSubmit}
-        onClose={Functions.NOOP}
-        placeholder={isLoading ? "Processing..." : "Ask your AI agent..."}
-        className={isLoading ? "opacity-50 pointer-events-none" : ""}
-        maxHeight="120px"
-      />
-    </div>
-  ),
+    onAddFiles,
+    onStop,
+    fileInputRef,
+    sessionMode,
+    onModeChange,
+  }) => {
+    const inputRef = useRef<ReactCodeMirrorRef | null>(null);
+    const promptCompletions: AdditionalCompletions | undefined = useMemo(() => {
+      if (!commands) {
+        return undefined;
+      }
+      // sentence has to begin with '/' to trigger autocomplete
+      return {
+        triggerCompletionRegex: /^\/(\w+)?/,
+        completions: commands.map(
+          (prompt): Completion => ({
+            label: `/${prompt.name}`,
+            info: prompt.description,
+          }),
+        ),
+      };
+    }, [commands]);
+
+    const handleSendClick = useEvent(() => {
+      if (promptValue.trim()) {
+        onPromptSubmit(undefined, promptValue);
+      }
+    });
+
+    const handleAddContext = useEvent(() => {
+      // For now, just append @ to the current value
+      addContextCompletion(inputRef);
+    });
+
+    return (
+      <div className="border-t bg-background flex-shrink-0">
+        <div
+          className={cn(
+            "px-3 py-2 min-h-[80px]",
+            (isLoading || !activeSessionId) && "opacity-50 pointer-events-none",
+          )}
+        >
+          <PromptInput
+            inputRef={inputRef}
+            value={promptValue}
+            onChange={isLoading ? Functions.NOOP : onPromptValueChange}
+            onSubmit={onPromptSubmit}
+            additionalCompletions={promptCompletions}
+            onClose={Functions.NOOP}
+            onAddFiles={onAddFiles}
+            placeholder={
+              isLoading
+                ? "Processing..."
+                : `Ask anything, ${CONTEXT_TRIGGER} to include context about tables or dataframes`
+            }
+            className={isLoading ? "opacity-50 pointer-events-none" : ""}
+            maxHeight="120px"
+          />
+        </div>
+        <TooltipProvider>
+          <div className="px-3 py-2 border-t border-border/20 flex flex-row items-center justify-between">
+            <div className="flex items-center gap-2">
+              {sessionMode && onModeChange && (
+                <ModeSelector
+                  sessionMode={sessionMode}
+                  onModeChange={onModeChange}
+                />
+              )}
+            </div>
+            <div className="flex flex-row">
+              <Tooltip content="Add context">
+                <Button variant="text" size="icon" onClick={handleAddContext}>
+                  <AtSignIcon className="h-3.5 w-3.5" />
+                </Button>
+              </Tooltip>
+              <Tooltip content="Attach a file">
+                <Button
+                  variant="text"
+                  size="icon"
+                  className="cursor-pointer"
+                  onClick={() => fileInputRef.current?.click()}
+                  title="Attach a file"
+                >
+                  <PaperclipIcon className="h-3.5 w-3.5" />
+                </Button>
+              </Tooltip>
+              <Input
+                ref={fileInputRef}
+                type="file"
+                multiple={true}
+                hidden={true}
+                onChange={(event) => {
+                  if (event.target.files) {
+                    onAddFiles([...event.target.files]);
+                  }
+                }}
+                accept={SUPPORTED_ATTACHMENT_TYPES.join(",")}
+              />
+              <Tooltip content={isLoading ? "Stop" : "Submit"}>
+                <Button
+                  variant="text"
+                  size="sm"
+                  className="h-6 w-6 p-0 hover:bg-muted/30 cursor-pointer"
+                  onClick={isLoading ? onStop : handleSendClick}
+                  disabled={isLoading ? false : !promptValue.trim()}
+                >
+                  {isLoading ? (
+                    <SquareIcon className="h-3 w-3 fill-current" />
+                  ) : (
+                    <SendIcon className="h-3 w-3" />
+                  )}
+                </Button>
+              </Tooltip>
+            </div>
+          </div>
+        </TooltipProvider>
+      </div>
+    );
+  },
 );
 PromptArea.displayName = "PromptArea";
 
+interface ModeSelectorProps {
+  sessionMode: SessionMode;
+  onModeChange: (mode: string) => void;
+}
+
+const ModeSelector = memo<ModeSelectorProps>(
+  ({ sessionMode, onModeChange }) => {
+    const availableModes = sessionMode?.availableModes || [];
+    const currentModeId = sessionMode?.currentModeId;
+    if (availableModes.length === 0) {
+      return null;
+    }
+
+    const modeOptions = availableModes.map((mode) => ({
+      value: mode.id,
+      label: mode.name,
+      subtitle: mode.description ?? "",
+    }));
+    const currentMode = modeOptions.find((opt) => opt.value === currentModeId);
+
+    return (
+      <Select value={currentModeId} onValueChange={onModeChange}>
+        <SelectTrigger className="h-6 text-xs border-border shadow-none! ring-0! bg-muted hover:bg-muted/30 py-0 px-2 gap-1 capitalize">
+          {currentMode?.label ?? currentModeId}
+        </SelectTrigger>
+        <SelectContent>
+          <SelectGroup>
+            <SelectLabel>Agent Mode</SelectLabel>
+            {modeOptions.map((option) => (
+              <SelectItem
+                key={option.value}
+                value={option.value}
+                className="text-xs"
+              >
+                <div className="flex flex-col">
+                  {option.label}
+                  {option.subtitle && (
+                    <div className="text-muted-foreground text-xs pt-1 block">
+                      {option.subtitle}
+                    </div>
+                  )}
+                </div>
+              </SelectItem>
+            ))}
+          </SelectGroup>
+        </SelectContent>
+      </Select>
+    );
+  },
+);
+ModeSelector.displayName = "ModeSelector";
+
 interface ChatContentProps {
   hasNotifications: boolean;
+  agentId: ExternalAgentId | undefined;
   connectionState: AgentConnectionState;
   sessionId: ExternalAgentSessionId | null;
   notifications: NotificationEvent[];
   pendingPermission: AgentPendingPermission;
-  onResolvePermission: (option: unknown) => void;
+  onResolvePermission: (option: RequestPermissionResponse) => void;
   onRetryConnection?: () => void;
   onRetryLastAction?: () => void;
   onDismissError?: (errorId: string) => void;
@@ -295,6 +498,7 @@ interface ChatContentProps {
 const ChatContent = memo<ChatContentProps>(
   ({
     hasNotifications,
+    agentId,
     connectionState,
     notifications,
     pendingPermission,
@@ -306,6 +510,7 @@ const ChatContent = memo<ChatContentProps>(
   }) => {
     const [isScrolledToBottom, setIsScrolledToBottom] = useState(true);
     const scrollContainerRef = useRef<HTMLDivElement>(null);
+    const isDisconnected = connectionState.status === "disconnected";
 
     // Scroll handler to determine if we're at the bottom of the chat
     const handleScroll = useEvent(() => {
@@ -375,12 +580,30 @@ const ChatContent = memo<ChatContentProps>(
               />
             </div>
           ) : (
-            <div className="flex items-center justify-center h-full min-h-[200px]">
+            <div className="flex items-center justify-center h-full min-h-[200px] flex-col">
               <PanelEmptyState
                 title="Waiting for agent"
                 description="Your AI agent will appear here when active"
                 icon={<BotMessageSquareIcon />}
               />
+              {isDisconnected && agentId && (
+                <AgentDocs
+                  className="border-t pt-6 px-5"
+                  title="Make sure you're connected to an agent"
+                  description="Run this command in your terminal:"
+                  agents={[agentId]}
+                />
+              )}
+              {isDisconnected && (
+                <Button
+                  variant="outline"
+                  onClick={onRetryConnection}
+                  type="button"
+                  className="mt-4"
+                >
+                  Retry
+                </Button>
+              )}
             </div>
           )}
         </div>
@@ -408,6 +631,8 @@ function getCwd() {
 const AgentPanel: React.FC = () => {
   const [isLoading, setIsLoading] = useState(false);
   const [promptValue, setPromptValue] = useState("");
+  const [files, setFiles] = useState<File[]>();
+  const fileInputRef = useRef<HTMLInputElement>(null);
 
   const [selectedTab] = useAtom(selectedTabAtom);
   const [sessionState, setSessionState] = useAtom(agentSessionStateAtom);
@@ -421,7 +646,7 @@ const AgentPanel: React.FC = () => {
   const acpClient = useAcpClient({
     wsUrl,
     clientOptions: {
-      readTextFile: (request): Promise<ReadTextFileResponse> => {
+      readTextFile: (request) => {
         logger.debug("Agent requesting file read", {
           path: request.path,
         });
@@ -429,7 +654,7 @@ const AgentPanel: React.FC = () => {
           content: response.contents || "",
         }));
       },
-      writeTextFile: (request): Promise<WriteTextFileResponse> => {
+      writeTextFile: (request) => {
         logger.debug("Agent requesting file write", {
           path: request.path,
           contentLength: request.content.length,
@@ -437,7 +662,7 @@ const AgentPanel: React.FC = () => {
         return sendUpdateFile({
           path: request.path,
           contents: request.content,
-        }).then(() => null);
+        }).then(() => ({}));
       },
     },
     autoConnect: false, // We'll manage connection manually based on active session
@@ -449,11 +674,25 @@ const AgentPanel: React.FC = () => {
     connectionState,
     notifications,
     pendingPermission,
+    availableCommands,
     resolvePermission,
+    sessionMode,
     activeSessionId,
     agent,
   } = acpClient;
 
+  useEffect(() => {
+    agent?.initialize({
+      protocolVersion: 1,
+      clientCapabilities: {
+        fs: {
+          readTextFile: true,
+          writeTextFile: true,
+        },
+      },
+    });
+  }, [agent]);
+
   // Auto-connect to agent when we have an active session, but only once per session
   useEffect(() => {
     if (wsUrl === NO_WS_SET) {
@@ -581,6 +820,7 @@ const AgentPanel: React.FC = () => {
       });
       setIsLoading(true);
       setPromptValue("");
+      setFiles(undefined);
 
       // Update session title with first message if it's still the default
       if (selectedTab?.title.startsWith("New ")) {
@@ -597,26 +837,47 @@ const AgentPanel: React.FC = () => {
         return;
       }
 
+      const promptBlocks: ContentBlock[] = [{ type: "text", text: prompt }];
+
+      // Parse context from the prompt
+      const { contextBlocks, attachmentBlocks } =
+        await parseContextFromPrompt(prompt);
+      promptBlocks.push(...contextBlocks, ...attachmentBlocks);
+
+      // Add manually uploaded files as resource links
+      if (files && files.length > 0) {
+        const fileResourceLinks = await convertFilesToResourceLinks(files);
+        promptBlocks.push(...fileResourceLinks);
+      }
+
+      const hasGivenRules = notifications.some(
+        (notification) =>
+          notification.type === "session_notification" &&
+          notification.data.update.sessionUpdate === "user_message_chunk",
+      );
+      if (!hasGivenRules) {
+        promptBlocks.push(
+          {
+            type: "resource_link",
+            uri: filename,
+            mimeType: "text/x-python",
+            name: filename,
+          },
+          {
+            type: "resource",
+            resource: {
+              uri: "marimo_rules.md",
+              mimeType: "text/markdown",
+              text: getAgentPrompt(filename),
+            },
+          },
+        );
+      }
+
       try {
         await agent.prompt({
           sessionId: activeSessionId,
-          prompt: [
-            { type: "text", text: prompt },
-            {
-              type: "resource_link",
-              uri: filename,
-              mimeType: "text/x-python",
-              name: filename,
-            },
-            {
-              type: "resource",
-              resource: {
-                uri: "marimo_rules.md",
-                mimeType: "text/markdown",
-                text: getAgentPrompt(filename),
-              },
-            },
-          ],
+          prompt: promptBlocks,
         });
       } catch (error) {
         logger.error("Failed to send prompt", { error });
@@ -635,6 +896,21 @@ const AgentPanel: React.FC = () => {
     setIsLoading(false);
   });
 
+  // Handler for adding files
+  const handleAddFiles = useEvent((newFiles: File[]) => {
+    if (newFiles.length === 0) {
+      return;
+    }
+    setFiles((prev) => [...(prev ?? []), ...newFiles]);
+  });
+
+  // Handler for removing files
+  const handleRemoveFile = useEvent((fileToRemove: File) => {
+    if (files) {
+      setFiles(files.filter((f) => f !== fileToRemove));
+    }
+  });
+
   // Handler for manual connect
   const handleManualConnect = useEvent(() => {
     logger.debug("Manual connect requested", {
@@ -652,6 +928,33 @@ const AgentPanel: React.FC = () => {
     disconnect();
   });
 
+  const handleModeChange = useEvent((mode: string) => {
+    logger.debug("Mode change requested", {
+      sessionId: activeSessionId,
+      mode,
+    });
+    if (!agent) {
+      toast({
+        title: "Agent not connected",
+        description: "Please connect to an agent to change the mode",
+        variant: "danger",
+      });
+      return;
+    }
+    if (!agent.setSessionMode) {
+      toast({
+        title: "Mode change not supported",
+        description: "The agent does not support mode changes",
+        variant: "danger",
+      });
+      return;
+    }
+    void agent.setSessionMode?.({
+      sessionId: activeSessionId as string,
+      modeId: mode,
+    });
+  });
+
   const hasNotifications = notifications.length > 0;
   const hasActiveSessions = sessionState.sessions.length > 0;
 
@@ -680,6 +983,7 @@ const AgentPanel: React.FC = () => {
       <SessionTabs />
 
       <ChatContent
+        agentId={selectedTab?.agentId}
         sessionId={activeSessionId}
         hasNotifications={hasNotifications}
         connectionState={connectionState}
@@ -690,7 +994,7 @@ const AgentPanel: React.FC = () => {
             sessionId: activeSessionId,
             option,
           });
-          resolvePermission(option as RequestPermissionResponse);
+          resolvePermission(option);
         }}
         onRetryConnection={handleManualConnect}
       />
@@ -701,12 +1005,30 @@ const AgentPanel: React.FC = () => {
         onStop={handleStop}
       />
 
+      {files && files.length > 0 && (
+        <div className="flex flex-row gap-1 flex-wrap p-3 border-t">
+          {files.map((file) => (
+            <FileAttachmentPill
+              file={file}
+              key={file.name}
+              onRemove={() => handleRemoveFile(file)}
+            />
+          ))}
+        </div>
+      )}
+
       <PromptArea
         isLoading={isLoading}
         activeSessionId={activeSessionId}
         promptValue={promptValue}
         onPromptValueChange={setPromptValue}
         onPromptSubmit={handlePromptSubmit}
+        onAddFiles={handleAddFiles}
+        onStop={handleStop}
+        fileInputRef={fileInputRef}
+        commands={availableCommands}
+        sessionMode={sessionMode}
+        onModeChange={handleModeChange}
       />
     </div>
   );

package/src/components/chat/acp/agent-selector.tsx

@@ -13,7 +13,9 @@ import {
   DropdownMenuSeparator,
   DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu";
+import { useFilename } from "@/core/saving/filename";
 import { cn } from "@/utils/cn";
+import { Paths } from "@/utils/paths";
 import { AgentDocs } from "./agent-docs";
 import {
   type AgentSession,
@@ -79,6 +81,7 @@ AgentMenuItem.displayName = "AgentMenuItem";
 
 export const AgentSelector: React.FC<AgentSelectorProps> = memo(
   ({ onSessionCreated, className }) => {
+    const filename = useFilename();
     const [sessionState, setSessionState] = useAtom(agentSessionStateAtom);
     const setActiveTab = useSetAtom(selectedTabAtom);
     const [isOpen, setIsOpen] = useState(false);
@@ -123,7 +126,10 @@ export const AgentSelector: React.FC<AgentSelectorProps> = memo(
           <div className="px-2 py-2">
             <div className="text-xs font-medium text-muted-foreground mb-3">
               To start an external agent, run the following command in your
-              terminal:
+              terminal.
+              <br />
+              Note: This must be in the directory{" "}
+              {Paths.dirname(filename ?? "")}
             </div>
             <AgentDocs
               agents={AVAILABLE_AGENTS.map((agent) => agent.id)}