call goose serve from tauri frontend via goose-acp client (#8549)

Co-authored-by: Jack Amadeo <jackamadeo@squareup.com>

ui/goose2/src/app/AppShell.tsx

@@ -98,6 +98,7 @@ export function AppShell({ children }: { children?: React.ReactNode }) {
           ? resolveEffectiveWorkingDir(null, await getHomeDir())
           : undefined);
       await acpLoadSession(sessionId, gooseSessionId, workingDir);
+      useChatStore.getState().setSessionLoading(sessionId, false);
       const t3 = performance.now();
       console.log(
         `[perf:load] ${sessionId.slice(0, 8)} acpLoadSession resolved in ${(t3 - t2).toFixed(1)}ms (total ${(t3 - t0).toFixed(1)}ms)`,

ui/goose2/src/app/hooks/useAppStartup.ts

@@ -1,10 +1,22 @@
 import { useEffect } from "react";
 import { useAgentStore } from "@/features/agents/stores/agentStore";
 import { useChatSessionStore } from "@/features/chat/stores/chatSessionStore";
+import { USE_DIRECT_ACP } from "@/shared/api/acpFeatureFlag";
+import { setNotificationHandler, getClient } from "@/shared/api/acpConnection";
+import notificationHandler from "@/shared/api/acpNotificationHandler";
 
 export function useAppStartup() {
   useEffect(() => {
     (async () => {
+      if (USE_DIRECT_ACP) {
+        try {
+          setNotificationHandler(notificationHandler);
+          await getClient();
+        } catch (err) {
+          console.error("Failed to initialize ACP connection:", err);
+        }
+      }
+
       const store = useAgentStore.getState();
       const loadPersonas = async () => {
         store.setPersonasLoading(true);

ui/goose2/src/features/chat/hooks/useChat.ts

@@ -246,6 +246,9 @@ export function useChat(
           ),
         });
 
+        store.setChatState(sessionId, "idle");
+        store.setStreamingMessageId(sessionId, null);
+
         if (wasDraft) {
           const promoted = sessionStore.getSession(sessionId);
           if (promoted) {

ui/goose2/src/shared/api/acp.ts

@@ -1,4 +1,12 @@
 import { invoke } from "@tauri-apps/api/core";
+import { USE_DIRECT_ACP } from "./acpFeatureFlag";
+import * as directAcp from "./acpApi";
+import * as sessionTracker from "./acpSessionTracker";
+import {
+  setActiveMessageId,
+  clearActiveMessageId,
+} from "./acpNotificationHandler";
+import { searchSessionsViaExports } from "./sessionSearch";
 
 export interface AcpProvider {
   id: string;
@@ -21,6 +29,9 @@ export interface AcpPrepareSessionOptions {
 
 /** Discover ACP providers installed on the system. */
 export async function discoverAcpProviders(): Promise<AcpProvider[]> {
+  if (USE_DIRECT_ACP) {
+    return directAcp.listProviders();
+  }
   return invoke("discover_acp_providers");
 }
 
@@ -31,6 +42,39 @@ export async function acpSendMessage(
   prompt: string,
   options: AcpSendMessageOptions = {},
 ): Promise<void> {
+  if (USE_DIRECT_ACP) {
+    const { systemPrompt, personaId, images } = options;
+
+    const gooseSessionId = sessionTracker.getGooseSessionId(
+      sessionId,
+      personaId,
+    );
+    if (!gooseSessionId) {
+      throw new Error("Session not prepared. Call acpPrepareSession first.");
+    }
+
+    const hasSystem = systemPrompt && systemPrompt.trim().length > 0;
+    const effectivePrompt = hasSystem
+      ? `<persona-instructions>\n${systemPrompt}\n</persona-instructions>\n\n<user-message>\n${prompt}\n</user-message>`
+      : prompt;
+
+    const content: import("@agentclientprotocol/sdk").ContentBlock[] = [
+      { type: "text", text: effectivePrompt },
+    ];
+    if (images) {
+      for (const [data, mimeType] of images) {
+        content.push({ type: "image", data, mimeType } as any);
+      }
+    }
+
+    const messageId = crypto.randomUUID();
+    setActiveMessageId(gooseSessionId, messageId);
+
+    await directAcp.prompt(gooseSessionId, content);
+
+    clearActiveMessageId(gooseSessionId);
+    return;
+  }
   const { systemPrompt, workingDir, personaId, personaName, images } = options;
   return invoke("acp_send_message", {
     sessionId,
@@ -50,6 +94,16 @@ export async function acpPrepareSession(
   providerId: string,
   options: AcpPrepareSessionOptions = {},
 ): Promise<void> {
+  if (USE_DIRECT_ACP) {
+    const workingDir = options.workingDir ?? "~/.goose/artifacts";
+    await sessionTracker.prepareSession(
+      sessionId,
+      providerId,
+      workingDir,
+      options.personaId,
+    );
+    return;
+  }
   const { workingDir, personaId } = options;
   return invoke("acp_prepare_session", {
     sessionId,
@@ -63,6 +117,10 @@ export async function acpSetModel(
   sessionId: string,
   modelId: string,
 ): Promise<void> {
+  if (USE_DIRECT_ACP) {
+    const gooseSessionId = sessionTracker.getGooseSessionId(sessionId);
+    return directAcp.setModel(gooseSessionId ?? sessionId, modelId);
+  }
   return invoke("acp_set_model", {
     sessionId,
     modelId,
@@ -87,6 +145,9 @@ export interface AcpSessionSearchResult {
 
 /** List all sessions known to the goose binary. */
 export async function acpListSessions(): Promise<AcpSessionInfo[]> {
+  if (USE_DIRECT_ACP) {
+    return directAcp.listSessions();
+  }
   return invoke("acp_list_sessions");
 }
 
@@ -94,6 +155,9 @@ export async function acpSearchSessions(
   query: string,
   sessionIds: string[],
 ): Promise<AcpSessionSearchResult[]> {
+  if (USE_DIRECT_ACP) {
+    return searchSessionsViaExports(query, sessionIds);
+  }
   return invoke("acp_search_sessions", { query, sessionIds });
 }
 
@@ -108,6 +172,17 @@ export async function acpLoadSession(
   gooseSessionId: string,
   workingDir?: string,
 ): Promise<void> {
+  if (USE_DIRECT_ACP) {
+    const effectiveWorkingDir = workingDir ?? "~/.goose/artifacts";
+    await directAcp.loadSession(gooseSessionId, effectiveWorkingDir);
+    sessionTracker.registerSession(
+      sessionId,
+      gooseSessionId,
+      "goose",
+      effectiveWorkingDir,
+    );
+    return;
+  }
   return invoke("acp_load_session", {
     sessionId,
     gooseSessionId,
@@ -117,11 +192,17 @@ export async function acpLoadSession(
 
 /** Export a session as JSON via the goose binary. */
 export async function acpExportSession(sessionId: string): Promise<string> {
+  if (USE_DIRECT_ACP) {
+    return directAcp.exportSession(sessionId);
+  }
   return invoke("acp_export_session", { sessionId });
 }
 
 /** Import a session from JSON via the goose binary. Returns new session metadata. */
 export async function acpImportSession(json: string): Promise<AcpSessionInfo> {
+  if (USE_DIRECT_ACP) {
+    return directAcp.importSession(json);
+  }
   return invoke("acp_import_session", { json });
 }
 
@@ -129,6 +210,11 @@ export async function acpImportSession(json: string): Promise<AcpSessionInfo> {
 export async function acpDuplicateSession(
   sessionId: string,
 ): Promise<AcpSessionInfo> {
+  if (USE_DIRECT_ACP) {
+    const gooseSessionId =
+      sessionTracker.getGooseSessionId(sessionId) ?? sessionId;
+    return directAcp.forkSession(gooseSessionId);
+  }
   return invoke("acp_duplicate_session", { sessionId });
 }
 
@@ -137,6 +223,14 @@ export async function acpCancelSession(
   sessionId: string,
   personaId?: string,
 ): Promise<boolean> {
+  if (USE_DIRECT_ACP) {
+    const gooseSessionId = sessionTracker.getGooseSessionId(
+      sessionId,
+      personaId,
+    );
+    await directAcp.cancelSession(gooseSessionId ?? sessionId);
+    return true;
+  }
   return invoke("acp_cancel_session", {
     sessionId,
     personaId: personaId ?? null,

ui/goose2/src/shared/api/acpApi.ts

@@ -0,0 +1,129 @@
+import type {
+  ContentBlock,
+  NewSessionResponse,
+  LoadSessionResponse,
+  PromptResponse,
+} from "@agentclientprotocol/sdk";
+import { getClient } from "./acpConnection";
+
+export interface AcpProvider {
+  id: string;
+  label: string;
+}
+
+export interface AcpSessionInfo {
+  sessionId: string;
+  title: string | null;
+  updatedAt: string | null;
+  messageCount: number;
+}
+
+const DEPRECATED_PROVIDER_IDS = new Set(["claude-code", "codex", "gemini-cli"]);
+
+export async function listProviders(): Promise<AcpProvider[]> {
+  const client = await getClient();
+  const result = await client.goose.GooseProvidersList({});
+  return (result as any).providers
+    .filter((p: any) => !DEPRECATED_PROVIDER_IDS.has(p.id))
+    .map((p: any) => ({ id: p.id, label: p.label }));
+}
+
+export async function listSessions(): Promise<AcpSessionInfo[]> {
+  const client = await getClient();
+  // GooseClient.unstable_listSessions doesn't work with SDK 0.19 (renamed to listSessions).
+  // Bypass GooseClient and call the connection directly. Fix when ui/acp is updated.
+  const conn = (client as any).conn;
+  const response = await conn.listSessions({});
+  return response.sessions.map((info: any) => ({
+    sessionId: info.sessionId,
+    title: info.title ?? null,
+    updatedAt: info.updatedAt ?? null,
+    messageCount: (info._meta?.messageCount as number) ?? 0,
+  }));
+}
+
+export async function exportSession(sessionId: string): Promise<string> {
+  const client = await getClient();
+  const result = await client.goose.GooseSessionExport({ sessionId });
+  return (result as any).data;
+}
+
+export async function importSession(json: string): Promise<AcpSessionInfo> {
+  const client = await getClient();
+  const result = await client.goose.GooseSessionImport({ data: json });
+  return result as unknown as AcpSessionInfo;
+}
+
+export async function forkSession(sessionId: string): Promise<AcpSessionInfo> {
+  const client = await getClient();
+  const response = await client.unstable_forkSession({
+    sessionId,
+    cwd: "~/.goose/artifacts",
+  });
+  return {
+    sessionId: response.sessionId,
+    title: (response._meta?.title as string) ?? null,
+    updatedAt: null,
+    messageCount: (response._meta?.messageCount as number) ?? 0,
+  };
+}
+
+export async function setModel(
+  sessionId: string,
+  modelId: string,
+): Promise<void> {
+  const client = await getClient();
+  await client.setSessionConfigOption({
+    sessionId,
+    configId: "model",
+    value: modelId,
+  });
+}
+
+export async function setProvider(
+  sessionId: string,
+  providerId: string,
+): Promise<void> {
+  const client = await getClient();
+  await client.setSessionConfigOption({
+    sessionId,
+    configId: "provider",
+    value: providerId,
+  });
+}
+
+export async function updateWorkingDir(
+  sessionId: string,
+  workingDir: string,
+): Promise<void> {
+  const client = await getClient();
+  await client.extMethod("goose/working_dir/update", { sessionId, workingDir });
+}
+
+export async function cancelSession(sessionId: string): Promise<void> {
+  const client = await getClient();
+  await client.cancel({ sessionId });
+}
+
+export async function newSession(
+  workingDir: string,
+): Promise<NewSessionResponse> {
+  const client = await getClient();
+  return client.newSession({ cwd: workingDir, mcpServers: [] });
+}
+
+export async function loadSession(
+  sessionId: string,
+  workingDir: string,
+): Promise<LoadSessionResponse> {
+  const client = await getClient();
+  return client.loadSession({ sessionId, cwd: workingDir, mcpServers: [] });
+}
+
+export async function prompt(
+  sessionId: string,
+  content: ContentBlock[],
+): Promise<PromptResponse> {
+  const client = await getClient();
+  return client.prompt({ sessionId, prompt: content });
+}

ui/goose2/src/shared/api/acpConnection.ts

@@ -0,0 +1,112 @@
+import { invoke } from "@tauri-apps/api/core";
+import { GooseClient } from "@aaif/goose-acp";
+import {
+  PROTOCOL_VERSION,
+  type Client,
+  type SessionNotification,
+  type RequestPermissionRequest,
+  type RequestPermissionResponse,
+} from "@agentclientprotocol/sdk";
+import { createWebSocketStream } from "./createWebSocketStream";
+
+let notificationHandler: AcpNotificationHandler | null = null;
+
+export interface AcpNotificationHandler {
+  handleSessionNotification(notification: SessionNotification): Promise<void>;
+}
+
+export function setNotificationHandler(handler: AcpNotificationHandler): void {
+  notificationHandler = handler;
+}
+
+let clientPromise: Promise<GooseClient> | null = null;
+let resolvedClient: GooseClient | null = null;
+
+function createClientCallbacks(): () => Client {
+  return () => ({
+    requestPermission: async (
+      args: RequestPermissionRequest,
+    ): Promise<RequestPermissionResponse> => {
+      const optionId = args.options?.[0]?.optionId ?? "approve";
+      return {
+        outcome: {
+          outcome: "selected",
+          optionId,
+        },
+      };
+    },
+
+    sessionUpdate: async (notification: SessionNotification): Promise<void> => {
+      if (notificationHandler) {
+        await notificationHandler.handleSessionNotification(notification);
+      }
+    },
+  });
+}
+
+function monitorConnection(client: GooseClient): void {
+  client.closed
+    .then(() => {
+      console.warn(
+        "[acp] Connection closed. Will reconnect on next getClient().",
+      );
+      resolvedClient = null;
+      clientPromise = null;
+    })
+    .catch(() => {
+      console.warn(
+        "[acp] Connection error. Will reconnect on next getClient().",
+      );
+      resolvedClient = null;
+      clientPromise = null;
+    });
+}
+
+async function initializeConnection(): Promise<GooseClient> {
+  const wsUrl: string = await invoke("get_goose_serve_url");
+
+  const stream = createWebSocketStream(wsUrl);
+
+  const client = new GooseClient(createClientCallbacks(), stream);
+
+  await client.initialize({
+    protocolVersion: PROTOCOL_VERSION,
+    clientCapabilities: {},
+    clientInfo: {
+      name: "goose2",
+      version: "0.1.0",
+    },
+  });
+
+  monitorConnection(client);
+
+  return client;
+}
+
+export async function getClient(): Promise<GooseClient> {
+  if (resolvedClient) {
+    return resolvedClient;
+  }
+
+  if (!clientPromise) {
+    clientPromise = initializeConnection()
+      .then((client) => {
+        resolvedClient = client;
+        return client;
+      })
+      .catch((error) => {
+        clientPromise = null;
+        throw error;
+      });
+  }
+
+  return clientPromise;
+}
+
+export function isClientReady(): boolean {
+  return resolvedClient !== null;
+}
+
+export function getClientSync(): GooseClient | null {
+  return resolvedClient;
+}

ui/goose2/src/shared/api/acpFeatureFlag.ts

@@ -0,0 +1 @@
+export const USE_DIRECT_ACP = true;

ui/goose2/src/shared/api/acpNotificationHandler.ts

@@ -0,0 +1,420 @@
+import type {
+  SessionNotification,
+  SessionUpdate,
+} from "@agentclientprotocol/sdk";
+import { useChatStore } from "@/features/chat/stores/chatStore";
+import { useChatSessionStore } from "@/features/chat/stores/chatSessionStore";
+import {
+  ensureReplayBuffer,
+  getBufferedMessage,
+  findLatestUnpairedToolRequest,
+} from "@/features/chat/hooks/replayBuffer";
+import type {
+  ToolRequestContent,
+  ToolResponseContent,
+} from "@/shared/types/messages";
+import type { AcpNotificationHandler } from "./acpConnection";
+import { getLocalSessionId } from "./acpSessionTracker";
+
+// Pre-set message ID for the next live stream per goose session
+const presetMessageIds = new Map<string, string>();
+
+export function setActiveMessageId(
+  gooseSessionId: string,
+  messageId: string,
+): void {
+  presetMessageIds.set(gooseSessionId, messageId);
+}
+
+export function clearActiveMessageId(gooseSessionId: string): void {
+  presetMessageIds.delete(gooseSessionId);
+}
+
+export async function handleSessionNotification(
+  notification: SessionNotification,
+): Promise<void> {
+  const gooseSessionId = notification.sessionId;
+  const sessionId = getLocalSessionId(gooseSessionId) ?? gooseSessionId;
+  const { update } = notification;
+  const isReplay = useChatStore.getState().loadingSessionIds.has(sessionId);
+
+  if (isReplay) {
+    handleReplay(sessionId, update);
+  } else {
+    handleLive(sessionId, gooseSessionId, update);
+  }
+}
+
+function handleReplay(sessionId: string, update: SessionUpdate): void {
+  switch (update.sessionUpdate) {
+    case "agent_message_chunk": {
+      const messageId = update.messageId ?? crypto.randomUUID();
+      const buffer = ensureReplayBuffer(sessionId);
+      if (!getBufferedMessage(sessionId, messageId)) {
+        buffer.push({
+          id: messageId,
+          role: "assistant",
+          created: Date.now(),
+          content: [],
+          metadata: {
+            userVisible: true,
+            agentVisible: true,
+            completionStatus: "inProgress",
+          },
+        });
+      }
+      const msg = getBufferedMessage(sessionId, messageId);
+      if (msg && update.content.type === "text" && "text" in update.content) {
+        const last = msg.content[msg.content.length - 1];
+        if (last?.type === "text") {
+          (last as { type: "text"; text: string }).text += update.content.text;
+        } else {
+          msg.content.push({ type: "text", text: update.content.text });
+        }
+      }
+      break;
+    }
+
+    case "user_message_chunk": {
+      const messageId = update.messageId ?? crypto.randomUUID();
+      const buffer = ensureReplayBuffer(sessionId);
+      const existing = getBufferedMessage(sessionId, messageId);
+      if (
+        !existing &&
+        update.content.type === "text" &&
+        "text" in update.content
+      ) {
+        buffer.push({
+          id: messageId,
+          role: "user",
+          created: Date.now(),
+          content: [{ type: "text", text: update.content.text }],
+          metadata: { userVisible: true, agentVisible: true },
+        });
+      } else if (
+        existing &&
+        update.content.type === "text" &&
+        "text" in update.content
+      ) {
+        const last = existing.content[existing.content.length - 1];
+        if (last?.type === "text") {
+          (last as { type: "text"; text: string }).text += update.content.text;
+        } else {
+          existing.content.push({ type: "text", text: update.content.text });
+        }
+      }
+      break;
+    }
+
+    case "tool_call": {
+      const msg = findMessageInBuffer(sessionId, update.toolCallId);
+      if (msg) {
+        msg.content.push({
+          type: "toolRequest",
+          id: update.toolCallId,
+          name: update.title,
+          arguments: {},
+          status: "executing",
+          startedAt: Date.now(),
+        });
+      }
+      break;
+    }
+
+    case "tool_call_update": {
+      const msg = findMessageWithToolCall(sessionId, update.toolCallId);
+      if (msg) {
+        if (update.title) {
+          const tc = msg.content.find(
+            (c) => c.type === "toolRequest" && c.id === update.toolCallId,
+          );
+          if (tc && tc.type === "toolRequest") {
+            (tc as ToolRequestContent).name = update.title;
+          }
+        }
+        if (update.status === "completed" || update.status === "failed") {
+          const tc = msg.content.find(
+            (c) => c.type === "toolRequest" && c.id === update.toolCallId,
+          );
+          if (tc && tc.type === "toolRequest") {
+            const idx = msg.content.indexOf(tc);
+            if (idx >= 0) {
+              msg.content[idx] = {
+                ...tc,
+                status: "completed",
+              } as ToolRequestContent;
+            }
+          }
+          const resultText = extractToolResultText(update);
+          msg.content.push({
+            type: "toolResponse",
+            id: update.toolCallId,
+            name: (tc as ToolRequestContent)?.name ?? "",
+            result: resultText,
+            isError: update.status === "failed",
+          });
+        }
+      }
+      break;
+    }
+
+    case "session_info_update":
+    case "config_option_update":
+    case "usage_update":
+      handleShared(sessionId, update);
+      break;
+
+    default:
+      break;
+  }
+}
+
+function handleLive(
+  sessionId: string,
+  gooseSessionId: string,
+  update: SessionUpdate,
+): void {
+  const store = useChatStore.getState();
+
+  switch (update.sessionUpdate) {
+    case "agent_message_chunk": {
+      const messageId =
+        update.messageId ??
+        presetMessageIds.get(gooseSessionId) ??
+        crypto.randomUUID();
+      const existing = store.messagesBySession[sessionId]?.find(
+        (m) => m.id === messageId,
+      );
+
+      if (!existing) {
+        store.addMessage(sessionId, {
+          id: messageId,
+          role: "assistant",
+          created: Date.now(),
+          content: [],
+          metadata: {
+            userVisible: true,
+            agentVisible: true,
+            completionStatus: "inProgress",
+          },
+        });
+        store.setPendingAssistantProvider(sessionId, null);
+        store.setStreamingMessageId(sessionId, messageId);
+      }
+
+      if (update.content.type === "text" && "text" in update.content) {
+        store.setStreamingMessageId(sessionId, messageId);
+        store.updateStreamingText(sessionId, update.content.text);
+      }
+      break;
+    }
+
+    case "tool_call": {
+      const messageId = findStreamingMessageId(sessionId);
+      if (!messageId) break;
+
+      const toolRequest: ToolRequestContent = {
+        type: "toolRequest",
+        id: update.toolCallId,
+        name: update.title,
+        arguments: {},
+        status: "executing",
+        startedAt: Date.now(),
+      };
+      store.setStreamingMessageId(sessionId, messageId);
+      store.appendToStreamingMessage(sessionId, toolRequest);
+      break;
+    }
+
+    case "tool_call_update": {
+      const messageId = findStreamingMessageId(sessionId);
+      if (!messageId) break;
+
+      if (update.title) {
+        store.updateMessage(sessionId, messageId, (msg) => ({
+          ...msg,
+          content: msg.content.map((c) =>
+            c.type === "toolRequest" && c.id === update.toolCallId
+              ? { ...c, name: update.title! }
+              : c,
+          ),
+        }));
+      }
+
+      if (update.status === "completed" || update.status === "failed") {
+        const streamingMessage = store.messagesBySession[sessionId]?.find(
+          (m) => m.id === messageId,
+        );
+        const toolRequest = streamingMessage
+          ? findLatestUnpairedToolRequest(streamingMessage.content)
+          : null;
+
+        store.updateMessage(sessionId, messageId, (msg) => ({
+          ...msg,
+          content: msg.content.map((block) =>
+            block.type === "toolRequest" && block.id === update.toolCallId
+              ? { ...block, status: "completed" }
+              : block,
+          ),
+        }));
+
+        const resultText = extractToolResultText(update);
+        const toolResponse: ToolResponseContent = {
+          type: "toolResponse",
+          id: update.toolCallId,
+          name: toolRequest?.name ?? "",
+          result: resultText,
+          isError: update.status === "failed",
+        };
+        store.setStreamingMessageId(sessionId, messageId);
+        store.appendToStreamingMessage(sessionId, toolResponse);
+      }
+      break;
+    }
+
+    case "session_info_update":
+    case "config_option_update":
+    case "usage_update":
+      handleShared(sessionId, update);
+      break;
+
+    default:
+      break;
+  }
+}
+
+function handleShared(sessionId: string, update: SessionUpdate): void {
+  switch (update.sessionUpdate) {
+    case "session_info_update": {
+      const info = update as SessionUpdate & {
+        sessionUpdate: "session_info_update";
+      };
+      if ("title" in info && info.title) {
+        const session = useChatSessionStore.getState().getSession(sessionId);
+        if (session && !session.userSetName) {
+          useChatSessionStore
+            .getState()
+            .updateSession(
+              sessionId,
+              { title: info.title as string },
+              { persistOverlay: false },
+            );
+        }
+      }
+      break;
+    }
+
+    case "config_option_update": {
+      const configUpdate = update as SessionUpdate & {
+        sessionUpdate: "config_option_update";
+      };
+      if ("options" in configUpdate && Array.isArray(configUpdate.options)) {
+        const modelOption = configUpdate.options.find(
+          (opt: any) => opt.category === "model",
+        );
+        if (modelOption?.kind?.type === "select") {
+          const select = modelOption.kind;
+          const currentModelId = select.currentValue;
+          const availableModels: Array<{ id: string; name: string }> = [];
+
+          if (select.options?.type === "ungrouped") {
+            for (const v of select.options.values) {
+              availableModels.push({ id: v.value, name: v.name });
+            }
+          } else if (select.options?.type === "grouped") {
+            for (const group of select.options.groups) {
+              for (const v of group.options) {
+                availableModels.push({ id: v.value, name: v.name });
+              }
+            }
+          }
+
+          const currentModelName =
+            availableModels.find((m) => m.id === currentModelId)?.name ??
+            currentModelId;
+
+          const sessionStore = useChatSessionStore.getState();
+          sessionStore.setSessionModels(sessionId, availableModels);
+          sessionStore.updateSession(
+            sessionId,
+            { modelId: currentModelId, modelName: currentModelName },
+            { persistOverlay: false },
+          );
+        }
+      }
+      break;
+    }
+
+    case "usage_update": {
+      const usage = update as SessionUpdate & { sessionUpdate: "usage_update" };
+      useChatStore.getState().updateTokenState(sessionId, {
+        accumulatedTotal: usage.used,
+        contextLimit: usage.size,
+      });
+      break;
+    }
+
+    default:
+      break;
+  }
+}
+
+// Helpers
+
+function findStreamingMessageId(sessionId: string): string | null {
+  return useChatStore.getState().getSessionRuntime(sessionId)
+    .streamingMessageId;
+}
+
+function findMessageInBuffer(
+  sessionId: string,
+  _toolCallId: string,
+): ReturnType<typeof getBufferedMessage> {
+  const buffer = ensureReplayBuffer(sessionId);
+  return buffer[buffer.length - 1];
+}
+
+function findMessageWithToolCall(
+  sessionId: string,
+  toolCallId: string,
+): ReturnType<typeof getBufferedMessage> {
+  const buffer = ensureReplayBuffer(sessionId);
+  for (let i = buffer.length - 1; i >= 0; i--) {
+    const msg = buffer[i];
+    if (
+      msg.content.some((c) => c.type === "toolRequest" && c.id === toolCallId)
+    ) {
+      return msg;
+    }
+  }
+  return buffer[buffer.length - 1];
+}
+
+function extractToolResultText(update: {
+  content?: Array<any> | null;
+  rawOutput?: unknown;
+}): string {
+  if (update.content && update.content.length > 0) {
+    for (const item of update.content) {
+      if (item.type === "content" && item.content?.type === "text") {
+        return item.content.text;
+      }
+    }
+  }
+  if (update.rawOutput !== undefined && update.rawOutput !== null) {
+    return typeof update.rawOutput === "string"
+      ? update.rawOutput
+      : JSON.stringify(update.rawOutput);
+  }
+  return "";
+}
+
+export function clearMessageTracking(): void {
+  presetMessageIds.clear();
+}
+
+const handler: AcpNotificationHandler = {
+  handleSessionNotification,
+};
+
+export default handler;

ui/goose2/src/shared/api/acpSessionTracker.ts

@@ -0,0 +1,86 @@
+import * as acpApi from "./acpApi";
+
+interface PreparedSession {
+  gooseSessionId: string;
+  providerId: string;
+  workingDir: string;
+}
+
+const prepared = new Map<string, PreparedSession>();
+const gooseToLocal = new Map<string, string>();
+
+function makeKey(sessionId: string, personaId?: string): string {
+  if (personaId && personaId.length > 0) {
+    return `${sessionId}__${personaId}`;
+  }
+  return sessionId;
+}
+
+export async function prepareSession(
+  sessionId: string,
+  providerId: string,
+  workingDir: string,
+  personaId?: string,
+): Promise<string> {
+  const key = makeKey(sessionId, personaId);
+
+  const existing = prepared.get(key) ?? prepared.get(sessionId);
+  if (existing) {
+    if (existing.workingDir !== workingDir) {
+      await acpApi.updateWorkingDir(existing.gooseSessionId, workingDir);
+      existing.workingDir = workingDir;
+    }
+    if (existing.providerId !== providerId) {
+      await acpApi.setProvider(existing.gooseSessionId, providerId);
+      existing.providerId = providerId;
+    }
+    return existing.gooseSessionId;
+  }
+
+  let gooseSessionId: string | null = null;
+
+  try {
+    await acpApi.loadSession(sessionId, workingDir);
+    gooseSessionId = sessionId;
+  } catch {}
+
+  if (!gooseSessionId) {
+    const response = await acpApi.newSession(workingDir);
+    gooseSessionId = response.sessionId;
+  }
+
+  await acpApi.setProvider(gooseSessionId, providerId);
+
+  prepared.set(key, { gooseSessionId, providerId, workingDir });
+  prepared.set(sessionId, { gooseSessionId, providerId, workingDir });
+  gooseToLocal.set(gooseSessionId, sessionId);
+
+  return gooseSessionId;
+}
+
+export function getGooseSessionId(
+  sessionId: string,
+  personaId?: string,
+): string | null {
+  const key = makeKey(sessionId, personaId);
+  return (
+    prepared.get(key)?.gooseSessionId ??
+    prepared.get(sessionId)?.gooseSessionId ??
+    null
+  );
+}
+
+export function getLocalSessionId(gooseSessionId: string): string | null {
+  return gooseToLocal.get(gooseSessionId) ?? null;
+}
+
+export function registerSession(
+  sessionId: string,
+  gooseSessionId: string,
+  providerId: string,
+  workingDir: string,
+): void {
+  const entry = { gooseSessionId, providerId, workingDir };
+  prepared.set(sessionId, entry);
+  gooseToLocal.set(gooseSessionId, sessionId);
+}

ui/goose2/src/shared/api/createWebSocketStream.ts

@@ -0,0 +1,80 @@
+import type { AnyMessage, Stream } from "@agentclientprotocol/sdk";
+
+export function createWebSocketStream(wsUrl: string): Stream {
+  const ws = new WebSocket(wsUrl);
+
+  const incoming: AnyMessage[] = [];
+  const waiters: Array<() => void> = [];
+  let closed = false;
+
+  function pushMessage(msg: AnyMessage): void {
+    incoming.push(msg);
+    const waiter = waiters.shift();
+    if (waiter) waiter();
+  }
+
+  function waitForMessage(): Promise<void> {
+    if (incoming.length > 0 || closed) return Promise.resolve();
+    return new Promise<void>((resolve) => waiters.push(resolve));
+  }
+
+  const openPromise = new Promise<void>((resolve, reject) => {
+    ws.addEventListener("open", () => resolve(), { once: true });
+    ws.addEventListener(
+      "error",
+      (event) => {
+        reject(new Error(`WebSocket connection failed: ${event}`));
+      },
+      { once: true },
+    );
+  });
+
+  ws.addEventListener("message", (event) => {
+    if (typeof event.data !== "string") return;
+    try {
+      const msg = JSON.parse(event.data) as AnyMessage;
+      pushMessage(msg);
+    } catch {
+      // ignore malformed JSON
+    }
+  });
+
+  ws.addEventListener("close", () => {
+    closed = true;
+    for (const waiter of waiters) waiter();
+    waiters.length = 0;
+  });
+
+  ws.addEventListener("error", () => {
+    closed = true;
+    for (const waiter of waiters) waiter();
+    waiters.length = 0;
+  });
+
+  const readable = new ReadableStream<AnyMessage>({
+    async pull(controller) {
+      await waitForMessage();
+      while (incoming.length > 0) {
+        controller.enqueue(incoming.shift()!);
+      }
+      if (closed && incoming.length === 0) {
+        controller.close();
+      }
+    },
+  });
+
+  const writable = new WritableStream<AnyMessage>({
+    async write(msg) {
+      await openPromise;
+      ws.send(JSON.stringify(msg));
+    },
+    close() {
+      ws.close();
+    },
+    abort() {
+      ws.close();
+    },
+  });
+
+  return { readable, writable };
+}

ui/goose2/src/shared/api/sessionSearch.ts

@@ -0,0 +1,206 @@
+import { exportSession } from "./acpApi";
+
+const SNIPPET_PREFIX = 40;
+const SNIPPET_SUFFIX = 60;
+
+type MessageRole = "user" | "assistant" | "system";
+
+const SEARCHABLE_ROLES = new Set<MessageRole>(["user", "assistant", "system"]);
+
+const SEARCHABLE_BLOCK_TYPES = new Set([
+  "text",
+  "input_text",
+  "output_text",
+  "systemNotification",
+  "system_notification",
+]);
+
+const SKIPPED_BLOCK_TYPES = new Set([
+  "toolRequest",
+  "toolResponse",
+  "thinking",
+  "redactedThinking",
+  "reasoning",
+  "image",
+]);
+
+export interface SessionSearchResult {
+  sessionId: string;
+  snippet: string;
+  messageId: string;
+  messageRole?: MessageRole;
+  matchCount: number;
+}
+
+interface ParsedMessage {
+  id: string;
+  role: MessageRole | null;
+  texts: string[];
+}
+
+export async function searchSessionsViaExports(
+  query: string,
+  sessionIds: string[],
+): Promise<SessionSearchResult[]> {
+  const trimmed = query.trim();
+  if (!trimmed) return [];
+
+  const unique = [...new Set(sessionIds)];
+  const results: SessionSearchResult[] = [];
+
+  for (const sessionId of unique) {
+    try {
+      const exported = await exportSession(sessionId);
+      const result = searchSession(sessionId, exported, trimmed);
+      if (result) results.push(result);
+    } catch {
+      // skip sessions that fail to export
+    }
+  }
+
+  return results;
+}
+
+function searchSession(
+  sessionId: string,
+  json: string,
+  query: string,
+): SessionSearchResult | null {
+  const root = safeParse(json);
+  if (!root) return null;
+
+  const conversation = root.conversation ?? root.messages;
+  if (!conversation) return null;
+
+  const messages = flattenMessages(conversation);
+  if (!messages.length) return null;
+
+  let firstMatch: {
+    messageId: string;
+    role: MessageRole | null;
+    snippet: string;
+  } | null = null;
+  let matchCount = 0;
+
+  for (const msg of messages) {
+    for (const text of msg.texts) {
+      const count = countMatches(text, query);
+      if (!count) continue;
+      matchCount += count;
+      firstMatch ??= {
+        messageId: msg.id,
+        role: msg.role,
+        snippet: buildSnippet(text, query),
+      };
+    }
+  }
+
+  if (!firstMatch) return null;
+
+  return {
+    sessionId,
+    snippet: firstMatch.snippet,
+    messageId: firstMatch.messageId,
+    messageRole: firstMatch.role ?? undefined,
+    matchCount,
+  };
+}
+
+function safeParse(json: string): Record<string, any> | null {
+  try {
+    return JSON.parse(json);
+  } catch {
+    return null;
+  }
+}
+
+function flattenMessages(value: unknown): ParsedMessage[] {
+  if (Array.isArray(value)) return value.flatMap(flattenMessages);
+  if (!isObject(value)) return [];
+
+  if ("message" in value) return flattenMessages(value.message);
+  if ("messages" in value) return flattenMessages(value.messages);
+
+  const msg = tryParseMessage(value);
+  return msg ? [msg] : [];
+}
+
+function tryParseMessage(obj: Record<string, unknown>): ParsedMessage | null {
+  if (!("role" in obj) || !("content" in obj || "text" in obj)) return null;
+
+  const role = toRole(obj.role);
+  const texts =
+    obj.content !== undefined
+      ? getSearchableTexts(obj.content, role)
+      : typeof obj.text === "string" && role && obj.text.trim()
+        ? [obj.text.trim()]
+        : [];
+
+  if (!texts.length) return null;
+
+  return {
+    id: typeof obj.id === "string" ? obj.id : crypto.randomUUID(),
+    role,
+    texts,
+  };
+}
+
+function getSearchableTexts(
+  value: unknown,
+  role: MessageRole | null,
+): string[] {
+  if (typeof value === "string") {
+    return role && SEARCHABLE_ROLES.has(role) && value.trim()
+      ? [value.trim()]
+      : [];
+  }
+  if (Array.isArray(value)) return value.flatMap((v) => getBlockText(v, role));
+  if (isObject(value)) return getBlockText(value, role);
+  return [];
+}
+
+function getBlockText(value: unknown, role: MessageRole | null): string[] {
+  if (!isObject(value)) return [];
+  const type = value.type as string | undefined;
+  const text = (value.text as string | undefined)?.trim();
+  if (!text) return [];
+
+  if (SKIPPED_BLOCK_TYPES.has(type ?? "")) return [];
+  if (SEARCHABLE_BLOCK_TYPES.has(type ?? "")) return [text];
+  return role && SEARCHABLE_ROLES.has(role) ? [text] : [];
+}
+
+function toRole(value: unknown): MessageRole | null {
+  if (typeof value !== "string") return null;
+  const r = value.trim().toLowerCase();
+  return SEARCHABLE_ROLES.has(r as MessageRole) ? (r as MessageRole) : null;
+}
+
+function isObject(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function countMatches(text: string, query: string): number {
+  const hay = text.toLowerCase();
+  const needle = query.toLowerCase();
+  if (!needle) return 0;
+
+  let count = 0;
+  let pos = hay.indexOf(needle);
+  while (pos !== -1) {
+    count++;
+    pos = hay.indexOf(needle, pos + needle.length);
+  }
+  return count;
+}
+
+function buildSnippet(text: string, query: string): string {
+  const idx = text.toLowerCase().indexOf(query.toLowerCase());
+  const at = idx >= 0 ? idx : 0;
+  const start = Math.max(0, at - SNIPPET_PREFIX);
+  const end = Math.min(text.length, at + query.length + SNIPPET_SUFFIX);
+
+  const prefix = start > 0 ? "..." : "";
+  const suffix = end < text.length ? "..." : "";
+  return `${prefix}${text.substring(start, end).trim()}${suffix}`;
+}