[goose2] MCP Apps: hydrate and replay app payloads in Goose2 (#8632)

Signed-off-by: Andrew Harvard <aharvard@squareup.com>

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

@@ -35,6 +35,10 @@ export function getBufferedMessage(
   return replayBuffers.get(sessionId)?.find((m) => m.id === messageId);
 }
 
+export function getReplayBuffer(sessionId: string): Message[] | undefined {
+  return replayBuffers.get(sessionId);
+}
+
 export function getAndDeleteReplayBuffer(
   sessionId: string,
 ): Message[] | undefined {

ui/goose2/src/features/chat/ui/McpAppView.tsx

@@ -0,0 +1,22 @@
+import { CodeBlock } from "@/shared/ui/ai-elements/code-block";
+import { useTranslation } from "react-i18next";
+import type { McpAppPayload } from "@/shared/types/messages";
+
+interface McpAppViewProps {
+  payload: McpAppPayload;
+}
+
+export function McpAppView({ payload }: McpAppViewProps) {
+  const { t } = useTranslation("chat");
+
+  // Currently we just render the MCP App payload as JSON.
+  // Up next, we'll replace this with actual HTML rendering and host bridging.
+  return (
+    <div className="my-3" data-testid="mcp-app-view">
+      <div className="mb-2 text-muted-foreground text-xs uppercase tracking-wide">
+        {t("message.mcpAppUnderConstruction")}
+      </div>
+      <CodeBlock code={JSON.stringify(payload, null, 2)} language="json" />
+    </div>
+  );
+}

ui/goose2/src/features/chat/ui/MessageBubble.tsx

@@ -32,6 +32,7 @@ import {
 } from "@/shared/ui/ai-elements/reasoning";
 import { ToolChainCards, type ToolChainItem } from "./ToolChainCards";
 import { ClickableImage } from "./ClickableImage";
+import { McpAppView } from "./McpAppView";
 import { useArtifactLinkHandler } from "@/features/chat/hooks/useArtifactLinkHandler";
 import type {
   Message,
@@ -97,10 +98,9 @@ interface ContentSection {
   items: MessageContent[] | ToolChainItem[];
 }
 
-/** Keep only content blocks whose audience includes "user" (or has no audience). */
 function filterUserVisibleContent(content: MessageContent[]): MessageContent[] {
   return content.filter((b) => {
-    const aud = b.annotations?.audience;
+    const aud = "annotations" in b ? b.annotations?.audience : undefined;
     return !aud || aud.length === 0 || aud.includes("user");
   });
 }
@@ -232,6 +232,8 @@ function renderContentBlock(
     case "toolResponse":
       // Handled by groupContentSections toolChain rendering
       return null;
+    case "mcpApp":
+      return <McpAppView key={`mcp-app-${index}`} payload={content.payload} />;
     case "thinking":
     case "reasoning": {
       const text = (content as ThinkingContent | ReasoningContentType).text;

ui/goose2/src/features/chat/ui/__tests__/MessageBubble.mcpApp.test.tsx

@@ -0,0 +1,81 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { render, screen } from "@testing-library/react";
+import { MessageBubble } from "../MessageBubble";
+import { useAgentStore } from "@/features/agents/stores/agentStore";
+import type { Message } from "@/shared/types/messages";
+
+vi.mock("@tauri-apps/plugin-opener", () => ({
+  openPath: vi.fn(),
+}));
+
+function assistantMessage(
+  content: Message["content"],
+  overrides: Partial<Message> = {},
+): Message {
+  return {
+    id: "a1",
+    role: "assistant",
+    created: Date.now(),
+    content,
+    ...overrides,
+  };
+}
+
+describe("MessageBubble MCP app rendering", () => {
+  beforeEach(() => {
+    useAgentStore.setState({ personas: [] });
+  });
+
+  it("renders MCP App blocks", () => {
+    const msg = assistantMessage([
+      {
+        type: "toolRequest",
+        id: "tool-1",
+        name: "weather: open app",
+        arguments: {},
+        status: "completed",
+      },
+      {
+        type: "toolResponse",
+        id: "tool-1",
+        name: "weather: open app",
+        result: "done",
+        isError: false,
+      },
+      {
+        type: "mcpApp",
+        id: "tool-1",
+        payload: {
+          sessionId: "local-session",
+          gooseSessionId: "goose-session",
+          toolCallId: "tool-1",
+          toolCallTitle: "weather: open app",
+          source: "toolCallUpdateMeta",
+          tool: {
+            name: "weather__open_app",
+            extensionName: "weather",
+            resourceUri: "ui://weather/app",
+          },
+          resource: {
+            result: {
+              contents: [
+                {
+                  uri: "ui://weather/app",
+                  mimeType: "text/html",
+                  text: "<div>Hello</div>",
+                },
+              ],
+            },
+          },
+        },
+      },
+    ]);
+
+    render(<MessageBubble message={msg} />);
+
+    const mcpAppView = screen.getByTestId("mcp-app-view");
+    expect(mcpAppView).toBeInTheDocument();
+    expect(mcpAppView).toHaveTextContent("ui://weather/app");
+    expect(mcpAppView).toHaveTextContent("<div>Hello</div>");
+  });
+});

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

@@ -0,0 +1,293 @@
+import { waitFor } from "@testing-library/react";
+import { beforeEach, describe, expect, it } from "vitest";
+import {
+  clearReplayBuffer,
+  getReplayBuffer,
+} from "@/features/chat/hooks/replayBuffer";
+import { useChatStore } from "@/features/chat/stores/chatStore";
+import type { McpAppPayload } from "@/shared/types/messages";
+import {
+  clearMessageTracking,
+  handleSessionNotification,
+  setActiveMessageId,
+} from "../acpNotificationHandler";
+import { registerSession } from "../acpSessionTracker";
+
+function createMcpAppPayload(): McpAppPayload {
+  return {
+    sessionId: "local-session",
+    gooseSessionId: "goose-session",
+    toolCallId: "tool-1",
+    toolCallTitle: "mcp_app_bench__inspect_host_info",
+    source: "toolCallUpdateMeta",
+    tool: {
+      name: "mcp_app_bench__inspect_host_info",
+      extensionName: "mcp_app_bench",
+      resourceUri: "ui://inspect-host-info",
+    },
+    resource: {
+      result: null,
+    },
+  };
+}
+
+describe("acpNotificationHandler", () => {
+  beforeEach(() => {
+    clearMessageTracking();
+    clearReplayBuffer("local-session");
+    clearReplayBuffer("goose-session");
+    useChatStore.setState({
+      messagesBySession: {},
+      sessionStateById: {},
+      queuedMessageBySession: {},
+      draftsBySession: {},
+      activeSessionId: null,
+      isConnected: false,
+      loadingSessionIds: new Set<string>(),
+      scrollTargetMessageBySession: {},
+    });
+  });
+
+  it("keeps tool calls that arrive before the first text chunk on the pending assistant message", async () => {
+    registerSession(
+      "local-session",
+      "goose-session",
+      "goose",
+      "/Users/aharvard/.goose/artifacts",
+    );
+    setActiveMessageId("goose-session", "assistant-1");
+
+    await handleSessionNotification({
+      sessionId: "goose-session",
+      update: {
+        sessionUpdate: "tool_call",
+        toolCallId: "tool-1",
+        title: "mcp_app_bench__inspect_host_info",
+      },
+    } as never);
+
+    await handleSessionNotification({
+      sessionId: "goose-session",
+      update: {
+        sessionUpdate: "tool_call_update",
+        toolCallId: "tool-1",
+        status: "completed",
+        content: [
+          {
+            type: "content",
+            content: {
+              type: "text",
+              text: "Opened the Host Info inspector.",
+            },
+          },
+        ],
+        _meta: {
+          goose: {
+            mcpApp: {
+              toolName: "mcp_app_bench__inspect_host_info",
+              extensionName: "mcp_app_bench",
+              resourceUri: "ui://inspect-host-info",
+            },
+          },
+        },
+      },
+    } as never);
+
+    await handleSessionNotification({
+      sessionId: "goose-session",
+      update: {
+        sessionUpdate: "agent_message_chunk",
+        content: {
+          type: "text",
+          text: "The Host Info inspector is now open.",
+        },
+      },
+    } as never);
+
+    await waitFor(() => {
+      const message =
+        useChatStore.getState().messagesBySession["local-session"]?.[0];
+      expect(message?.content.some((block) => block.type === "mcpApp")).toBe(
+        true,
+      );
+    });
+
+    const [message] =
+      useChatStore.getState().messagesBySession["local-session"];
+    expect(message.id).toBe("assistant-1");
+    expect(message.content.map((block) => block.type)).toEqual([
+      "toolRequest",
+      "toolResponse",
+      "mcpApp",
+      "text",
+    ]);
+    expect(message.content[0]).toMatchObject({
+      type: "toolRequest",
+      id: "tool-1",
+      name: "mcp_app_bench__inspect_host_info",
+      status: "completed",
+    });
+    expect(message.content[1]).toMatchObject({
+      type: "toolResponse",
+      id: "tool-1",
+      name: "mcp_app_bench__inspect_host_info",
+      result: "Opened the Host Info inspector.",
+      isError: false,
+    });
+    expect(message.content[2]).toMatchObject({
+      type: "mcpApp",
+      id: "tool-1",
+      payload: createMcpAppPayload(),
+    });
+    expect(message.content[3]).toMatchObject({
+      type: "text",
+      text: "The Host Info inspector is now open.",
+    });
+    expect(
+      useChatStore.getState().getSessionRuntime("local-session")
+        .streamingMessageId,
+    ).toBe("assistant-1");
+  });
+
+  it("replay keeps tool and MCP app content on an assistant message when tool events arrive before text", async () => {
+    const replaySessionId = "replay-goose-session";
+    useChatStore.setState({
+      loadingSessionIds: new Set<string>([replaySessionId]),
+    });
+
+    await handleSessionNotification({
+      sessionId: replaySessionId,
+      update: {
+        sessionUpdate: "user_message_chunk",
+        messageId: "user-1",
+        content: {
+          type: "text",
+          text: "run the app bench",
+        },
+      },
+    } as never);
+
+    await handleSessionNotification({
+      sessionId: replaySessionId,
+      update: {
+        sessionUpdate: "tool_call",
+        toolCallId: "tool-1",
+        title: "mcp_app_bench__inspect_host_info",
+      },
+    } as never);
+
+    await handleSessionNotification({
+      sessionId: replaySessionId,
+      update: {
+        sessionUpdate: "tool_call_update",
+        toolCallId: "tool-1",
+        status: "completed",
+        content: [
+          {
+            type: "content",
+            content: {
+              type: "text",
+              text: "Opened the Host Info inspector.",
+            },
+          },
+        ],
+        _meta: {
+          goose: {
+            mcpApp: {
+              toolName: "mcp_app_bench__inspect_host_info",
+              extensionName: "mcp_app_bench",
+              resourceUri: "ui://inspect-host-info",
+            },
+          },
+        },
+      },
+    } as never);
+
+    await handleSessionNotification({
+      sessionId: replaySessionId,
+      update: {
+        sessionUpdate: "agent_message_chunk",
+        messageId: "assistant-1",
+        content: {
+          type: "text",
+          text: "The Host Info inspector is now open.",
+        },
+      },
+    } as never);
+
+    const buffer = getReplayBuffer(replaySessionId);
+    expect(buffer).toHaveLength(2);
+    expect(buffer?.[0]).toMatchObject({
+      id: "user-1",
+      role: "user",
+      content: [{ type: "text", text: "run the app bench" }],
+    });
+    expect(
+      buffer?.[0]?.content.some((block) => block.type === "toolRequest"),
+    ).toBe(false);
+
+    expect(buffer?.[1]?.id).toBe("assistant-1");
+    expect(buffer?.[1]?.role).toBe("assistant");
+    expect(buffer?.[1]?.content.map((block) => block.type)).toEqual([
+      "toolRequest",
+      "toolResponse",
+      "mcpApp",
+      "text",
+    ]);
+    expect(buffer?.[1]?.content[2]).toMatchObject({
+      type: "mcpApp",
+      id: "tool-1",
+      payload: {
+        ...createMcpAppPayload(),
+        sessionId: replaySessionId,
+        gooseSessionId: replaySessionId,
+      },
+    });
+  });
+
+  it("replay preserves gooseSessionId in MCP app payloads before tracker registration", async () => {
+    const replaySessionId = "replay-goose-session-2";
+    useChatStore.setState({
+      loadingSessionIds: new Set<string>([replaySessionId]),
+    });
+
+    await handleSessionNotification({
+      sessionId: replaySessionId,
+      update: {
+        sessionUpdate: "tool_call",
+        toolCallId: "tool-1",
+        title: "mcp_app_bench__inspect_host_info",
+      },
+    } as never);
+
+    await handleSessionNotification({
+      sessionId: replaySessionId,
+      update: {
+        sessionUpdate: "tool_call_update",
+        toolCallId: "tool-1",
+        status: "completed",
+        _meta: {
+          goose: {
+            mcpApp: {
+              toolName: "mcp_app_bench__inspect_host_info",
+              extensionName: "mcp_app_bench",
+              resourceUri: "ui://inspect-host-info",
+            },
+          },
+        },
+      },
+    } as never);
+
+    const buffer = getReplayBuffer(replaySessionId);
+    const assistant = buffer?.[0];
+    const mcpAppBlock = assistant?.content.find(
+      (block) => block.type === "mcpApp",
+    );
+    expect(mcpAppBlock).toMatchObject({
+      type: "mcpApp",
+      payload: expect.objectContaining({
+        gooseSessionId: replaySessionId,
+      }),
+    });
+  });
+});

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

@@ -97,17 +97,19 @@ export async function acpSendMessage(
   const tPrompt = performance.now();
   const meta: Record<string, unknown> = {};
   if (personaId) meta.personaId = personaId;
-  await directAcp.prompt(
-    gooseSessionId,
-    content,
-    Object.keys(meta).length > 0 ? meta : undefined,
-  );
-  const tDone = performance.now();
-  perfLog(
-    `[perf:send] ${sid} prompt() resolved in ${(tDone - tPrompt).toFixed(1)}ms (total acpSendMessage ${(tDone - tStart).toFixed(1)}ms)`,
-  );
-
-  clearActiveMessageId(gooseSessionId);
+  try {
+    await directAcp.prompt(
+      gooseSessionId,
+      content,
+      Object.keys(meta).length > 0 ? meta : undefined,
+    );
+    const tDone = performance.now();
+    perfLog(
+      `[perf:send] ${sid} prompt() resolved in ${(tDone - tPrompt).toFixed(1)}ms (total acpSendMessage ${(tDone - tStart).toFixed(1)}ms)`,
+    );
+  } finally {
+    clearActiveMessageId(gooseSessionId);
+  }
 }
 
 /** Prepare or warm an ACP session ahead of the first prompt. */

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

@@ -6,10 +6,14 @@ import {
   getAndDeleteReplayBuffer,
 } from "@/features/chat/hooks/replayBuffer";
 import { registerSession } from "./acpSessionTracker";
-import { handleSessionNotification } from "./acpNotificationHandler";
+import {
+  clearMessageTracking,
+  handleSessionNotification,
+} from "./acpNotificationHandler";
 
 describe("acpNotificationHandler", () => {
   beforeEach(() => {
+    clearMessageTracking();
     clearReplayBuffer("draft-session-1");
     clearReplayBuffer("draft-session-2");
     useChatStore.setState({

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

@@ -15,6 +15,17 @@ import type {
   ToolResponseContent,
 } from "@/shared/types/messages";
 import type { AcpNotificationHandler } from "./acpConnection";
+import {
+  attachMcpAppPayload,
+  extractToolResultText,
+  findReplayMessageWithToolCall,
+} from "./acpToolCallContent";
+import {
+  clearReplayAssistantMessage,
+  clearReplayAssistantTracking,
+  ensureReplayAssistantMessage,
+  getTrackedReplayAssistantMessageId,
+} from "./acpReplayAssistant";
 import {
   getLocalSessionId,
   subscribeToSessionRegistration,
@@ -23,47 +34,6 @@ import { perfLog } from "@/shared/lib/perfLog";
 
 // Pre-set message ID for the next live stream per goose session
 const presetMessageIds = new Map<string, string>();
-const pendingUsageUpdates = new Map<string, SessionUpdate[]>();
-
-function shouldBufferPendingUpdate(update: SessionUpdate): boolean {
-  return update.sessionUpdate === "usage_update";
-}
-
-function queuePendingUsageUpdate(
-  gooseSessionId: string,
-  update: SessionUpdate,
-): void {
-  const pending = pendingUsageUpdates.get(gooseSessionId);
-  if (pending) {
-    pending.push(update);
-    return;
-  }
-  pendingUsageUpdates.set(gooseSessionId, [update]);
-}
-
-function flushPendingUsageUpdates(
-  localSessionId: string,
-  gooseSessionId: string,
-): void {
-  const pending = pendingUsageUpdates.get(gooseSessionId);
-  if (!pending?.length) {
-    return;
-  }
-
-  pendingUsageUpdates.delete(gooseSessionId);
-
-  for (const update of pending) {
-    if (useChatStore.getState().loadingSessionIds.has(localSessionId)) {
-      handleReplay(localSessionId, update);
-    } else {
-      handleLive(localSessionId, gooseSessionId, update);
-    }
-  }
-}
-
-subscribeToSessionRegistration((localSessionId, gooseSessionId) => {
-  flushPendingUsageUpdates(localSessionId, gooseSessionId);
-});
 
 // Per-session perf counters for replay/live streaming.
 interface ReplayPerf {
@@ -78,6 +48,20 @@ interface LivePerf {
   chunkCount: number;
 }
 const livePerf = new Map<string, LivePerf>();
+const pendingUsageUpdates = new Map<
+  string,
+  { accumulatedTotal: number; contextLimit: number }
+>();
+
+subscribeToSessionRegistration((localSessionId, gooseSessionId) => {
+  const pendingUsage = pendingUsageUpdates.get(gooseSessionId);
+  if (!pendingUsage) {
+    return;
+  }
+
+  useChatStore.getState().updateTokenState(localSessionId, pendingUsage);
+  pendingUsageUpdates.delete(gooseSessionId);
+});
 
 export function setActiveMessageId(
   gooseSessionId: string,
@@ -112,32 +96,23 @@ export async function handleSessionNotification(
   notification: SessionNotification,
 ): Promise<void> {
   const gooseSessionId = notification.sessionId;
-  const { update } = notification;
   const localSessionId = getLocalSessionId(gooseSessionId);
-
-  if (!localSessionId) {
-    if (shouldBufferPendingUpdate(update)) {
-      queuePendingUsageUpdate(gooseSessionId, update);
-    }
-    return;
-  }
-
-  const isReplay = useChatStore
-    .getState()
-    .loadingSessionIds.has(localSessionId);
+  const sessionId = localSessionId ?? gooseSessionId;
+  const { update } = notification;
+  const isReplay = useChatStore.getState().loadingSessionIds.has(sessionId);
 
   if (isReplay) {
-    const sid = localSessionId.slice(0, 8);
-    let perf = replayPerf.get(localSessionId);
+    const sid = sessionId.slice(0, 8);
+    let perf = replayPerf.get(sessionId);
     const now = performance.now();
     if (!perf) {
       perf = { firstAt: now, lastAt: now, count: 0 };
-      replayPerf.set(localSessionId, perf);
+      replayPerf.set(sessionId, perf);
       perfLog(`[perf:replay] ${sid} first notification received`);
     }
     perf.lastAt = now;
     perf.count += 1;
-    handleReplay(localSessionId, update);
+    handleReplay(sessionId, gooseSessionId, localSessionId, update);
   } else {
     const perf = livePerf.get(gooseSessionId);
     if (perf && update.sessionUpdate === "agent_message_chunk") {
@@ -150,7 +125,7 @@ export async function handleSessionNotification(
         );
       }
     }
-    handleLive(localSessionId, gooseSessionId, update);
+    handleLive(sessionId, gooseSessionId, localSessionId, update);
   }
 }
 
@@ -166,25 +141,18 @@ export function clearReplayPerf(sessionId: string): void {
   replayPerf.delete(sessionId);
 }
 
-function handleReplay(sessionId: string, update: SessionUpdate): void {
+function handleReplay(
+  sessionId: string,
+  gooseSessionId: string,
+  localSessionId: string | null,
+  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);
+      const msg = ensureReplayAssistantMessage(
+        sessionId,
+        update.messageId ?? null,
+      );
       if (msg && update.content.type === "text" && "text" in update.content) {
         const last = msg.content[msg.content.length - 1];
         if (last?.type === "text") {
@@ -197,6 +165,7 @@ function handleReplay(sessionId: string, update: SessionUpdate): void {
     }
 
     case "user_message_chunk": {
+      clearReplayAssistantMessage(sessionId);
       if (update.content.type !== "text" || !("text" in update.content)) break;
       const messageId = update.messageId ?? crypto.randomUUID();
       const buffer = ensureReplayBuffer(sessionId);
@@ -228,22 +197,25 @@ function handleReplay(sessionId: string, update: SessionUpdate): void {
     }
 
     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(),
-        });
-      }
+      const msg = ensureReplayAssistantMessage(sessionId);
+      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);
+      const replayMessageId = getTrackedReplayAssistantMessageId(sessionId);
+      const msg =
+        findReplayMessageWithToolCall(sessionId, update.toolCallId) ??
+        (replayMessageId
+          ? getBufferedMessage(sessionId, replayMessageId)
+          : undefined);
       if (msg) {
         if (update.title) {
           const tc = msg.content.find(
@@ -274,6 +246,19 @@ function handleReplay(sessionId: string, update: SessionUpdate): void {
             result: resultText,
             isError: update.status === "failed",
           });
+          if (update.status === "completed") {
+            attachMcpAppPayload(
+              sessionId,
+              update.toolCallId,
+              (tc as ToolRequestContent)?.name ?? update.title ?? "",
+              update,
+              true,
+              {
+                gooseSessionId,
+                replayMessageId,
+              },
+            );
+          }
         }
       }
       break;
@@ -282,7 +267,7 @@ function handleReplay(sessionId: string, update: SessionUpdate): void {
     case "session_info_update":
     case "config_option_update":
     case "usage_update":
-      handleShared(sessionId, update);
+      handleShared(sessionId, gooseSessionId, localSessionId, update);
       break;
 
     default:
@@ -293,36 +278,19 @@ function handleReplay(sessionId: string, update: SessionUpdate): void {
 function handleLive(
   sessionId: string,
   gooseSessionId: string,
+  localSessionId: string | null,
   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,
+      const messageId = ensureLiveAssistantMessage(
+        sessionId,
+        gooseSessionId,
+        update.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);
@@ -331,8 +299,7 @@ function handleLive(
     }
 
     case "tool_call": {
-      const messageId = findStreamingMessageId(sessionId);
-      if (!messageId) break;
+      const messageId = ensureLiveAssistantMessage(sessionId, gooseSessionId);
 
       const toolRequest: ToolRequestContent = {
         type: "toolRequest",
@@ -348,8 +315,7 @@ function handleLive(
     }
 
     case "tool_call_update": {
-      const messageId = findStreamingMessageId(sessionId);
-      if (!messageId) break;
+      const messageId = ensureLiveAssistantMessage(sessionId, gooseSessionId);
 
       if (update.title) {
         store.updateMessage(sessionId, messageId, (msg) => ({
@@ -389,6 +355,15 @@ function handleLive(
         };
         store.setStreamingMessageId(sessionId, messageId);
         store.appendToStreamingMessage(sessionId, toolResponse);
+        if (update.status === "completed") {
+          attachMcpAppPayload(
+            sessionId,
+            update.toolCallId,
+            toolRequest?.name ?? update.title ?? "",
+            update,
+            false,
+          );
+        }
       }
       break;
     }
@@ -396,7 +371,7 @@ function handleLive(
     case "session_info_update":
     case "config_option_update":
     case "usage_update":
-      handleShared(sessionId, update);
+      handleShared(sessionId, gooseSessionId, localSessionId, update);
       break;
 
     default:
@@ -404,7 +379,12 @@ function handleLive(
   }
 }
 
-function handleShared(sessionId: string, update: SessionUpdate): void {
+function handleShared(
+  sessionId: string,
+  gooseSessionId: string,
+  localSessionId: string | null,
+  update: SessionUpdate,
+): void {
   switch (update.sessionUpdate) {
     case "session_info_update": {
       const info = update as SessionUpdate & {
@@ -463,7 +443,16 @@ function handleShared(sessionId: string, update: SessionUpdate): void {
 
     case "usage_update": {
       const usage = update as SessionUpdate & { sessionUpdate: "usage_update" };
-      useChatStore.getState().updateTokenState(sessionId, {
+
+      if (!localSessionId) {
+        pendingUsageUpdates.set(gooseSessionId, {
+          accumulatedTotal: usage.used,
+          contextLimit: usage.size,
+        });
+        break;
+      }
+
+      useChatStore.getState().updateTokenState(localSessionId, {
         accumulatedTotal: usage.used,
         contextLimit: usage.size,
       });
@@ -475,8 +464,6 @@ function handleShared(sessionId: string, update: SessionUpdate): void {
   }
 }
 
-// Helpers
-
 function findStreamingMessageId(sessionId: string): string | null {
   return useChatStore.getState().getSessionRuntime(sessionId)
     .streamingMessageId;
@@ -489,52 +476,53 @@ function makeTextBlock(
   return { type: "text", text, ...(ann ? { annotations: ann } : {}) };
 }
 
-function findMessageInBuffer(
+function ensureLiveAssistantMessage(
   sessionId: string,
-  _toolCallId: string,
-): ReturnType<typeof getBufferedMessage> {
-  const buffer = ensureReplayBuffer(sessionId);
-  return buffer[buffer.length - 1];
-}
+  gooseSessionId: string,
+  preferredMessageId?: string | null,
+): string {
+  const store = useChatStore.getState();
+  const existingStreamingMessageId = findStreamingMessageId(sessionId);
+  const messages = store.messagesBySession[sessionId] ?? [];
 
-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;
-    }
+  if (
+    existingStreamingMessageId &&
+    messages.some((message) => message.id === existingStreamingMessageId)
+  ) {
+    return existingStreamingMessageId;
   }
-  return buffer[buffer.length - 1];
-}
 
-function extractToolResultText(update: {
-  // biome-ignore lint/suspicious/noExplicitAny: ACP SDK ToolCallContent type is complex
-  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;
-      }
-    }
+  const messageId =
+    preferredMessageId ??
+    presetMessageIds.get(gooseSessionId) ??
+    existingStreamingMessageId ??
+    crypto.randomUUID();
+
+  if (!messages.some((message) => message.id === messageId)) {
+    store.addMessage(sessionId, {
+      id: messageId,
+      role: "assistant",
+      created: Date.now(),
+      content: [],
+      metadata: {
+        userVisible: true,
+        agentVisible: true,
+        completionStatus: "inProgress",
+      },
+    });
   }
-  if (update.rawOutput !== undefined && update.rawOutput !== null) {
-    return typeof update.rawOutput === "string"
-      ? update.rawOutput
-      : JSON.stringify(update.rawOutput);
-  }
-  return "";
+
+  store.setPendingAssistantProvider(sessionId, null);
+  store.setStreamingMessageId(sessionId, messageId);
+  clearActiveMessageId(gooseSessionId);
+
+  return messageId;
 }
 
 export function clearMessageTracking(): void {
   presetMessageIds.clear();
+  pendingUsageUpdates.clear();
+  clearReplayAssistantTracking();
 }
 
 const handler: AcpNotificationHandler = {

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

@@ -0,0 +1,64 @@
+import {
+  ensureReplayBuffer,
+  getBufferedMessage,
+} from "@/features/chat/hooks/replayBuffer";
+import type { Message } from "@/shared/types/messages";
+
+const replayAssistantMessageIds = new Map<string, string>();
+
+export function getTrackedReplayAssistantMessageId(
+  sessionId: string,
+): string | null {
+  return replayAssistantMessageIds.get(sessionId) ?? null;
+}
+
+export function ensureReplayAssistantMessage(
+  sessionId: string,
+  preferredMessageId?: string | null,
+): Message {
+  const trackedMessageId = replayAssistantMessageIds.get(sessionId);
+
+  if (preferredMessageId) {
+    const preferredMessage = getBufferedMessage(sessionId, preferredMessageId);
+    if (preferredMessage?.role === "assistant") {
+      replayAssistantMessageIds.set(sessionId, preferredMessageId);
+      return preferredMessage;
+    }
+  }
+
+  if (trackedMessageId) {
+    const trackedMessage = getBufferedMessage(sessionId, trackedMessageId);
+    if (trackedMessage?.role === "assistant") {
+      if (preferredMessageId && trackedMessage.id !== preferredMessageId) {
+        trackedMessage.id = preferredMessageId;
+        replayAssistantMessageIds.set(sessionId, preferredMessageId);
+      }
+      return trackedMessage;
+    }
+  }
+
+  const messageId = preferredMessageId ?? crypto.randomUUID();
+  const buffer = ensureReplayBuffer(sessionId);
+  const message: Message = {
+    id: messageId,
+    role: "assistant",
+    created: Date.now(),
+    content: [],
+    metadata: {
+      userVisible: true,
+      agentVisible: true,
+      completionStatus: "inProgress",
+    },
+  };
+  buffer.push(message);
+  replayAssistantMessageIds.set(sessionId, messageId);
+  return message;
+}
+
+export function clearReplayAssistantMessage(sessionId: string): void {
+  replayAssistantMessageIds.delete(sessionId);
+}
+
+export function clearReplayAssistantTracking(): void {
+  replayAssistantMessageIds.clear();
+}

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

@@ -0,0 +1,150 @@
+import type { SessionUpdate } from "@agentclientprotocol/sdk";
+import { useChatStore } from "@/features/chat/stores/chatStore";
+import {
+  getReplayBuffer,
+  getBufferedMessage,
+} from "@/features/chat/hooks/replayBuffer";
+import type { McpAppContent, MessageContent } from "@/shared/types/messages";
+import { buildMcpAppPayloadFromToolUpdate } from "./mcpAppToolUpdate";
+
+export function findReplayMessageWithToolCall(
+  sessionId: string,
+  toolCallId: string,
+): ReturnType<typeof getBufferedMessage> {
+  const buffer = getReplayBuffer(sessionId);
+  if (!buffer) {
+    return undefined;
+  }
+  for (let index = buffer.length - 1; index >= 0; index -= 1) {
+    const message = buffer[index];
+    if (
+      message.content.some(
+        (content) =>
+          content.type === "toolRequest" && content.id === toolCallId,
+      )
+    ) {
+      return message;
+    }
+  }
+  return undefined;
+}
+
+export function extractToolResultText(update: {
+  // biome-ignore lint/suspicious/noExplicitAny: ACP SDK ToolCallContent type is complex
+  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 attachMcpAppPayload(
+  sessionId: string,
+  toolCallId: string,
+  toolCallTitle: string,
+  update: SessionUpdate,
+  isReplay: boolean,
+  options?: {
+    gooseSessionId?: string | null;
+    replayMessageId?: string | null;
+  },
+): void {
+  const payload = buildMcpAppPayloadFromToolUpdate(
+    sessionId,
+    toolCallId,
+    toolCallTitle,
+    update,
+    options?.gooseSessionId,
+  );
+  if (!payload) {
+    return;
+  }
+
+  const block: McpAppContent = {
+    type: "mcpApp",
+    id: toolCallId,
+    payload,
+  };
+
+  if (isReplay) {
+    const message =
+      findReplayMessageWithToolCall(sessionId, toolCallId) ??
+      (options?.replayMessageId
+        ? getBufferedMessage(sessionId, options.replayMessageId)
+        : undefined);
+    if (message) {
+      message.content = insertMcpAppContent(message.content, block);
+      return;
+    }
+  }
+
+  const store = useChatStore.getState();
+  const message = [...(store.messagesBySession[sessionId] ?? [])]
+    .reverse()
+    .find((candidate) =>
+      candidate.content.some(
+        (content) =>
+          content.type === "toolRequest" && content.id === toolCallId,
+      ),
+    );
+  if (!message) {
+    return;
+  }
+
+  store.updateMessage(sessionId, message.id, (current) => ({
+    ...current,
+    content: insertMcpAppContent(current.content, block),
+  }));
+}
+
+function insertMcpAppContent(
+  content: MessageContent[],
+  block: McpAppContent,
+): MessageContent[] {
+  if (content.some((item) => item.type === "mcpApp" && item.id === block.id)) {
+    return content;
+  }
+
+  const insertAfterIndex = findMcpAppAnchorIndex(content, block.id);
+  if (insertAfterIndex === -1) {
+    return [...content, block];
+  }
+
+  return [
+    ...content.slice(0, insertAfterIndex + 1),
+    block,
+    ...content.slice(insertAfterIndex + 1),
+  ];
+}
+
+function findMcpAppAnchorIndex(
+  content: MessageContent[],
+  toolCallId: string,
+): number {
+  for (let index = content.length - 1; index >= 0; index -= 1) {
+    const block = content[index];
+    if (block.type === "toolResponse" && block.id === toolCallId) {
+      return index;
+    }
+  }
+
+  for (let index = content.length - 1; index >= 0; index -= 1) {
+    const block = content[index];
+    if (block.type === "toolRequest" && block.id === toolCallId) {
+      return index;
+    }
+  }
+
+  return -1;
+}

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

@@ -0,0 +1,62 @@
+import type {
+  GooseMcpAppToolPayload,
+  GooseReadResourceResult,
+  GooseToolCallUpdateMeta,
+  GooseToolMetadata,
+} from "@aaif/goose-sdk";
+import type { SessionUpdate } from "@agentclientprotocol/sdk";
+import type { McpAppPayload } from "@/shared/types/messages";
+import { getGooseSessionId } from "./acpSessionTracker";
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null;
+}
+
+function extractMcpAppPayload(
+  update: SessionUpdate,
+): GooseMcpAppToolPayload | null {
+  if (update.sessionUpdate !== "tool_call_update" || !isRecord(update._meta)) {
+    return null;
+  }
+
+  const meta = update._meta as GooseToolCallUpdateMeta;
+  const payload = meta.goose?.mcpApp;
+  return isRecord(payload) ? (payload as GooseMcpAppToolPayload) : null;
+}
+
+export function buildMcpAppPayloadFromToolUpdate(
+  sessionId: string,
+  toolCallId: string,
+  toolCallTitle: string,
+  update: SessionUpdate,
+  gooseSessionIdOverride?: string | null,
+): McpAppPayload | null {
+  const payload = extractMcpAppPayload(update);
+  if (!payload) {
+    return null;
+  }
+
+  return {
+    sessionId,
+    gooseSessionId:
+      gooseSessionIdOverride ?? getGooseSessionId(sessionId) ?? null,
+    toolCallId,
+    toolCallTitle,
+    source: "toolCallUpdateMeta",
+    tool: {
+      name: payload.toolName,
+      extensionName: payload.extensionName,
+      resourceUri: payload.resourceUri,
+      meta: isRecord(payload.toolMeta)
+        ? (payload.toolMeta as GooseToolMetadata)
+        : undefined,
+    },
+    resource: {
+      result:
+        (payload.resourceResult as GooseReadResourceResult | null) ?? null,
+      ...(typeof payload.readError === "string"
+        ? { readError: payload.readError }
+        : {}),
+    },
+  };
+}

ui/goose2/src/shared/i18n/locales/en/chat.json

@@ -126,6 +126,7 @@
   "message": {
     "copied": "Copied",
     "defaultImageAlt": "Attached",
+    "mcpAppUnderConstruction": "🚧 MCP App Rendering is under construction",
     "redactedThinking": "(thinking redacted)"
   },
   "persona": {

ui/goose2/src/shared/i18n/locales/es/chat.json

@@ -126,6 +126,7 @@
   "message": {
     "copied": "Copiado",
     "defaultImageAlt": "Adjunto",
+    "mcpAppUnderConstruction": "🚧 La renderización de MCP App está en construcción",
     "redactedThinking": "(pensamiento redactado)"
   },
   "persona": {

ui/goose2/src/shared/types/messages.ts

@@ -1,3 +1,8 @@
+import type {
+  GooseReadResourceResult,
+  GooseToolMetadata,
+} from "@aaif/goose-sdk";
+
 export type ChatAttachmentKind = "image" | "file" | "directory";
 
 export interface ChatImageAttachmentDraft {
@@ -89,6 +94,30 @@ export interface ToolResponseContent {
   annotations?: ContentAnnotations;
 }
 
+export interface McpAppPayload {
+  sessionId: string;
+  gooseSessionId: string | null;
+  toolCallId: string;
+  toolCallTitle: string;
+  source: "toolCallUpdateMeta";
+  tool: {
+    name: string;
+    extensionName: string;
+    resourceUri: string;
+    meta?: GooseToolMetadata;
+  };
+  resource: {
+    result: GooseReadResourceResult | null;
+    readError?: string;
+  };
+}
+
+export interface McpAppContent {
+  type: "mcpApp";
+  id: string;
+  payload: McpAppPayload;
+}
+
 export interface ThinkingContent {
   type: "thinking";
   text: string;
@@ -129,6 +158,7 @@ export type MessageContent =
   | ImageContent
   | ToolRequestContent
   | ToolResponseContent
+  | McpAppContent
   | ThinkingContent
   | RedactedThinkingContent
   | ReasoningContent
@@ -181,6 +211,9 @@ export function isToolRequest(c: MessageContent): c is ToolRequestContent {
 export function isToolResponse(c: MessageContent): c is ToolResponseContent {
   return c.type === "toolResponse";
 }
+export function isMcpApp(c: MessageContent): c is McpAppContent {
+  return c.type === "mcpApp";
+}
 export function isThinking(c: MessageContent): c is ThinkingContent {
   return c.type === "thinking";
 }

ui/pnpm-lock.yaml

@@ -651,6 +651,9 @@ importers:
       '@modelcontextprotocol/ext-apps':
         specifier: ^0.3.1
         version: 0.3.1(@modelcontextprotocol/sdk@1.27.1(zod@3.25.76))(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(zod@3.25.76)
+      '@modelcontextprotocol/sdk':
+        specifier: ^1.27.0
+        version: 1.27.1(zod@3.25.76)
       zod:
         specifier: ^3.25.76
         version: 3.25.76

ui/sdk/package.json

@@ -41,6 +41,7 @@
   },
   "dependencies": {
     "@modelcontextprotocol/ext-apps": "^0.3.1",
+    "@modelcontextprotocol/sdk": "^1.27.0",
     "zod": "^3.25.76"
   },
   "peerDependencies": {

ui/sdk/src/mcp-apps.ts

@@ -2,7 +2,17 @@ import type {
   Implementation,
   InitializeRequest,
 } from "@agentclientprotocol/sdk";
-import { RESOURCE_MIME_TYPE } from "@modelcontextprotocol/ext-apps";
+import { RESOURCE_MIME_TYPE } from "@modelcontextprotocol/ext-apps/app-bridge";
+import type {
+  McpUiAppResourceConfig,
+  McpUiAppToolConfig,
+} from "@modelcontextprotocol/ext-apps/server";
+import type {
+  BlobResourceContents,
+  ReadResourceResult,
+  TextResourceContents,
+  Tool,
+} from "@modelcontextprotocol/sdk/types.js";
 
 export const GOOSE_MCP_UI_EXTENSION_ID = "io.modelcontextprotocol/ui" as const;
 
@@ -14,6 +24,50 @@ export interface GooseMcpHostCapabilities {
   extensions: Record<string, GooseMcpUiExtensionSettings>;
 }
 
+export type GooseToolUiMetadata = Extract<
+  McpUiAppToolConfig["_meta"],
+  { ui: unknown }
+>["ui"];
+
+export type GooseToolMetadata = NonNullable<Tool["_meta"]> & {
+  ui?: GooseToolUiMetadata;
+  goose_extension?: string;
+};
+
+export type GooseSessionTool = Tool & {
+  meta?: GooseToolMetadata;
+  _meta?: GooseToolMetadata;
+};
+
+export type GooseTextResourceContents = TextResourceContents;
+
+export type GooseBlobResourceContents = BlobResourceContents;
+
+export type GooseResourceContents = TextResourceContents | BlobResourceContents;
+
+export type GooseReadResourceResult = ReadResourceResult;
+
+export type GooseResourceMetadata = NonNullable<
+  Extract<NonNullable<McpUiAppResourceConfig["_meta"]>, { ui?: unknown }>["ui"]
+>;
+
+export interface GooseMcpAppToolPayload {
+  toolName: string;
+  extensionName: string;
+  resourceUri: string;
+  toolMeta?: GooseToolMetadata;
+  resourceResult?: GooseReadResourceResult | null;
+  readError?: string;
+}
+
+export interface GooseToolCallUpdateMeta {
+  goose?: {
+    mcpApp?: GooseMcpAppToolPayload;
+    [key: string]: unknown;
+  };
+  [key: string]: unknown;
+}
+
 export interface GooseClientMeta {
   goose: {
     mcpHostCapabilities: GooseMcpHostCapabilities;