feat: add support for Codex agent

frontend/package.json

@@ -166,7 +166,7 @@
     "thememirror": "^2.0.1",
     "timestring": "^7.0.0",
     "typescript-memoize": "^1.1.1",
-    "use-acp": "0.2.4",
+    "use-acp": "0.2.5",
     "use-resize-observer": "^9.1.0",
     "vega-lite": "^5.23.0",
     "vega-loader": "^4.5.3",

frontend/src/components/ai/ai-provider-icon.tsx

@@ -15,14 +15,15 @@ import type { ProviderId } from "@/core/ai/ids/ids";
 import { cn } from "@/utils/cn";
 import marimoIcon from "../../assets/icon-32x32.png?inline";
 
-type Aliases = "claude" | "gemini";
+type Aliases = "claude" | "gemini" | "codex";
 
 const icons: Record<ProviderId | Aliases, string> = {
   openai: OpenAIIcon,
   anthropic: AnthropicIcon,
   claude: AnthropicIcon,
   gemini: GeminiIcon,
   google: GeminiIcon,
+  codex: OpenAIIcon,
   ollama: OllamaIcon,
   azure: AzureIcon,
   bedrock: BedrockIcon,

frontend/src/components/chat/acp/__tests__/state.test.ts

@@ -9,7 +9,6 @@ import {
   type ExternalAgentId,
   getAgentConnectionCommand,
   getAgentDisplayName,
-  getAllAgentIds,
   getSessionsByAgent,
   removeSession,
   type TabId,
@@ -632,18 +631,6 @@ describe("state utility functions", () => {
     });
   });
 
-  describe("getAllAgentIds", () => {
-    it("should return all available agent IDs", () => {
-      const agentIds = getAllAgentIds();
-      expect(agentIds).toMatchInlineSnapshot(`
-        [
-          "claude",
-          "gemini",
-        ]
-      `);
-    });
-  });
-
   describe("getAgentDisplayName", () => {
     it("should capitalize agent names", () => {
       expect({

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

@@ -381,7 +381,12 @@
             </div>
             <div className="flex flex-row">
               <Tooltip content="Add context">
-                <Button variant="text" size="icon" onClick={handleAddContext}>
+                <Button
+                  variant="text"
+                  size="icon"
+                  onClick={handleAddContext}
+                  disabled={isLoading}
+                >
                   <AtSignIcon className="h-3.5 w-3.5" />
                 </Button>
               </Tooltip>
@@ -392,6 +397,7 @@
                   className="cursor-pointer"
                   onClick={() => fileInputRef.current?.click()}
                   title="Attach a file"
+                  disabled={isLoading}
                 >
                   <PaperclipIcon className="h-3.5 w-3.5" />
                 </Button>
@@ -725,7 +731,7 @@
       // We don't want to disconnect so users can switch between different
       // panels without losing their session
     };
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [wsUrl]);
 
   const handleNewSession = useEvent(async () => {
@@ -820,7 +826,7 @@
     };
 
     createOrResumeSession();
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [isConnected, agent, tabLastActiveSessionId, activeSessionId]);
 
   // Handler for prompt submission

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

@@ -42,6 +42,11 @@ const AVAILABLE_AGENTS = [
     displayName: "Gemini",
     iconId: "google",
   },
+  {
+    id: "codex",
+    displayName: "Codex",
+    iconId: "openai",
+  },
 ] as const;
 
 interface AgentMenuItemProps {

frontend/src/components/chat/acp/blocks.tsx

@@ -22,7 +22,8 @@ import {
   XIcon,
 } from "lucide-react";
 import React from "react";
-import { mergeToolCalls } from "use-acp";
+import { JsonRpcError, mergeToolCalls } from "use-acp";
+import { z } from "zod";
 import { ReadonlyDiff } from "@/components/editor/code/readonly-diff";
 import { JsonOutput } from "@/components/editor/output/JsonOutput";
 import { Button } from "@/components/ui/button";
@@ -31,6 +32,7 @@ import {
   PopoverContent,
   PopoverTrigger,
 } from "@/components/ui/popover";
+import { uniqueByTakeLast } from "@/utils/arrays";
 import { logNever } from "@/utils/assertNever";
 import { cn } from "@/utils/cn";
 import { Strings } from "@/utils/strings";
@@ -102,13 +104,20 @@ export const ErrorBlock = (props: {
   onRetry?: () => void;
   onDismiss?: () => void;
 }) => {
-  const { message } = props.data;
+  const error = props.data;
+  let message = props.data.message;
 
   // Don't show WebSocket connection errors
   if (message.includes("WebSocket")) {
     return null;
   }
 
+  if (error instanceof JsonRpcError) {
+    const dataStr =
+      typeof error.data === "string" ? error.data : JSON.stringify(error.data);
+    message = `${dataStr} (${error.code})`;
+  }
+
   return (
     <div
       className="border border-[var(--red-6)] bg-[var(--red-2)] rounded-lg p-4 my-2"
@@ -307,7 +316,9 @@ export const AgentThoughtsBlock = (props: {
 };
 
 export const PlansBlock = (props: { data: PlanNotificationEvent[] }) => {
-  const plans = props.data.flatMap((item) => item.entries);
+  // Dedupe plans by text, take the last one which may have a status update
+  let plans = props.data.flatMap((item) => item.entries);
+  plans = uniqueByTakeLast(plans, (item) => item.content);
 
   return (
     <div className="rounded-lg border bg-background p-2 text-xs">
@@ -529,6 +540,7 @@ export const SessionNotificationsBlock = <
   data: T[];
   startTimestamp: number;
   endTimestamp: number;
+  isLastBlock: boolean;
 }) => {
   if (props.data.length === 0) {
     return null;
@@ -537,7 +549,9 @@ export const SessionNotificationsBlock = <
 
   const renderItems = (items: T[]) => {
     if (isToolCalls(items)) {
-      return <ToolNotificationsBlock data={items} />;
+      return (
+        <ToolNotificationsBlock data={items} isLastBlock={props.isLastBlock} />
+      );
     }
     if (isAgentThoughts(items)) {
       return (
@@ -591,6 +605,7 @@ export const CurrentModeBlock = (props: {
 
 export const ToolNotificationsBlock = (props: {
   data: (ToolCallNotificationEvent | ToolCallUpdateNotificationEvent)[];
+  isLastBlock: boolean;
 }) => {
   const toolCalls = mergeToolCalls(props.data);
 
@@ -604,7 +619,9 @@ export const ToolNotificationsBlock = (props: {
               ? "success"
               : item.status === "failed"
                 ? "error"
-                : item.status === "in_progress" || item.status === "pending"
+                : (item.status === "in_progress" ||
+                      item.status === "pending") &&
+                    !props.isLastBlock
                   ? "loading"
                   : undefined
           }
@@ -631,7 +648,7 @@ export const DiffBlocks = (props: {
         return (
           <div
             key={item.path}
-            className="border rounded-md overflow-hidden bg-[var(--gray-2)] overflow-y-auto scrollbar-thin"
+            className="border rounded-md overflow-hidden bg-[var(--gray-2)] overflow-y-auto scrollbar-thin max-h-64"
           >
             {/* File path header */}
             <div className="px-2 py-1 bg-[var(--gray-2)] border-b text-xs font-medium text-[var(--gray-11)]">
@@ -651,8 +668,12 @@ export const DiffBlocks = (props: {
 function toolTitle(
   item: Pick<ToolCallUpdateNotificationEvent, "title" | "kind" | "locations">,
 ) {
-  const prefix =
-    item.title || Strings.startCase(item.kind || "") || "Tool call";
+  let title = item.title;
+  // Hack: sometimes title comes back: "undefined", so lets undo that
+  if (title === '"undefined"') {
+    title = undefined;
+  }
+  const prefix = title || Strings.startCase(item.kind || "") || "Tool call";
   const firstLocation = item.locations?.[0];
   // Add the first location if it is not in the title already
   if (firstLocation && !prefix.includes(firstLocation.path)) {
@@ -699,6 +720,22 @@ export const ToolBodyBlock = (props: {
 
   // Completely empty
   if (!content && !hasLocations && rawInput) {
+    // HACK: if the raw input is `abs_path`, `old_string`, `new_string` then handle it as if it is a diff
+    const rawDiff = rawDiffSchema.safeParse(rawInput);
+    if (rawDiff.success) {
+      return (
+        <DiffBlocks
+          data={[
+            {
+              type: "diff",
+              oldText: rawDiff.data.old_string,
+              newText: rawDiff.data.new_string,
+              path: rawDiff.data.abs_path,
+            },
+          ]}
+        />
+      );
+    }
     // Show rawInput
     return (
       <pre className="bg-[var(--slate-2)] p-1 text-muted-foreground border border-[var(--slate-4)] rounded text-xs overflow-auto scrollbar-thin max-h-64">
@@ -728,3 +765,9 @@ export const ToolBodyBlock = (props: {
     </div>
   );
 };
+
+const rawDiffSchema = z.object({
+  abs_path: z.string(),
+  old_string: z.string(),
+  new_string: z.string(),
+});

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

@@ -11,7 +11,7 @@ import type { ExternalAgentSessionId, SessionSupportType } from "./types";
 
 // Types
 export type TabId = TypedString<"TabId">;
-export type ExternalAgentId = "claude" | "gemini";
+export type ExternalAgentId = "claude" | "gemini" | "codex";
 
 // No agents support loading sessions, so we limit to 1, otherwise
 // this is confusing to the user when switching between sessions
@@ -220,7 +220,7 @@ export function getSessionsByAgent(
 }
 
 export function getAllAgentIds(): ExternalAgentId[] {
-  return ["claude", "gemini"];
+  return ["claude", "gemini", "codex"];
 }
 
 export function getAgentDisplayName(agentId: ExternalAgentId): string {
@@ -251,6 +251,12 @@ const AGENT_CONFIG: Record<ExternalAgentId, AgentConfig> = {
     webSocketUrl: "ws://localhost:3019/message",
     sessionSupport: "single",
   },
+  codex: {
+    port: 3021,
+    command: "npx @zed-industries/codex-acp",
+    webSocketUrl: "ws://localhost:3021/message",
+    sessionSupport: "single",
+  },
 };
 
 export function getAgentSessionSupport(

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

@@ -50,7 +50,10 @@ export const AgentThread = ({
     return true;
   });
 
-  const renderNotification = (group: NotificationEvent[]) => {
+  const renderNotification = (
+    group: NotificationEvent[],
+    isLastBlock: boolean,
+  ) => {
     if (group.length === 0) {
       return null;
     }
@@ -90,6 +93,7 @@ export const AgentThread = ({
           data={data}
           startTimestamp={startTimestamp}
           endTimestamp={endTimestamp}
+          isLastBlock={isLastBlock}
         />
       );
     }
@@ -98,9 +102,12 @@ export const AgentThread = ({
 
   return (
     <div className="flex flex-col gap-4 px-2 pb-10 flex-1">
-      {combinedNotifications.map((notification) => (
+      {combinedNotifications.map((notification, idx) => (
         <React.Fragment key={notification[0].id}>
-          {renderNotification(notification)}
+          {renderNotification(
+            notification,
+            idx === combinedNotifications.length - 1,
+          )}
         </React.Fragment>
       ))}
       {combinedNotifications.length === 0 && <ReadyToChatBlock />}

frontend/src/components/chat/chat-panel.tsx

@@ -372,7 +372,12 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
           </div>
           <div className="flex flex-row">
             <Tooltip content="Add context">
-              <Button variant="text" size="icon" onClick={onAddContext}>
+              <Button
+                variant="text"
+                size="icon"
+                onClick={onAddContext}
+                disabled={isLoading}
+              >
                 <AtSignIcon className="h-3.5 w-3.5" />
               </Button>
             </Tooltip>
@@ -385,6 +390,7 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
                     className="cursor-pointer"
                     onClick={() => fileInputRef.current?.click()}
                     title="Attach a file"
+                    disabled={isLoading}
                   >
                     <PaperclipIcon className="h-3.5 w-3.5" />
                   </Button>

frontend/src/components/editor/code/readonly-diff.tsx

@@ -15,6 +15,10 @@ export const ReadonlyDiff = memo(
         unifiedMergeView({
           original: props.original,
           mergeControls: false,
+          collapseUnchanged: {
+            margin: 3,
+            minSize: 4,
+          },
         }),
       ];
     }, [props.original, props.modified, theme]);

frontend/src/utils/arrays.ts

@@ -74,6 +74,9 @@ export function arrayToggle<T>(arr: T[], item: T): T[] {
   return result;
 }
 
+/**
+ * Unique by a key function.
+ */
 export function uniqueBy<T>(arr: T[], key: (item: T) => string): T[] {
   const result = [];
   const seen = new Set();
@@ -87,6 +90,19 @@ export function uniqueBy<T>(arr: T[], key: (item: T) => string): T[] {
   return result;
 }
 
+/**
+ * Unique by a key function, taking the last item for each key.
+ */
+export function uniqueByTakeLast<T>(arr: T[], key: (item: T) => string): T[] {
+  // Create a map of keys to items
+  const map = new Map<string, T>();
+  for (const item of arr) {
+    const k = key(item);
+    map.set(k, item);
+  }
+  return Array.from(map.values());
+}
+
 /**
  * Get the next index in the list, wrapping around to the start or end if necessary.
  * @param currentIndex - The current index, or null if there is no current index.