Fix Goose2 ACP notification state handling

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

Author: aharvard

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

@@ -40,7 +40,6 @@ import type {
   MessageContent,
   TextContent,
   ImageContent,
-  McpAppContent,
   ToolResponseContent,
   ThinkingContent,
   ReasoningContent as ReasoningContentType,
@@ -226,10 +225,8 @@ function renderContentBlock(
     case "toolResponse":
       // Handled by groupContentSections toolChain rendering
       return null;
-    case "mcpApp": {
-      const mcpApp = content as McpAppContent;
-      return <McpAppView key={`mcp-app-${index}`} payload={mcpApp.payload} />;
-    }
+    case "mcpApp":
+      return <McpAppView key={`mcp-app-${index}`} payload={content.payload} />;
     case "thinking": {
       const th = content as ThinkingContent;
       return (

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/features/chat/ui/__tests__/MessageBubble.test.tsx

@@ -380,59 +380,6 @@ describe("MessageBubble", () => {
     expect(screen.queryByText("Tool result")).not.toBeInTheDocument();
   });
 
-  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>");
-  });
-
   it("renders thinking content as Reasoning block", () => {
     const msg = assistantMessage([{ type: "thinking", text: "deep thoughts" }]);
     render(<MessageBubble message={msg} />);

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

@@ -25,7 +25,10 @@ import {
   ensureReplayAssistantMessage,
   getTrackedReplayAssistantMessageId,
 } from "./acpReplayAssistant";
-import { getLocalSessionId } from "./acpSessionTracker";
+import {
+  getLocalSessionId,
+  subscribeToSessionRegistration,
+} from "./acpSessionTracker";
 import { perfLog } from "@/shared/lib/perfLog";
 
 // Pre-set message ID for the next live stream per goose session
@@ -44,6 +47,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,
@@ -78,7 +95,8 @@ export async function handleSessionNotification(
   notification: SessionNotification,
 ): Promise<void> {
   const gooseSessionId = notification.sessionId;
-  const sessionId = getLocalSessionId(gooseSessionId) ?? gooseSessionId;
+  const localSessionId = getLocalSessionId(gooseSessionId);
+  const sessionId = localSessionId ?? gooseSessionId;
   const { update } = notification;
   const isReplay = useChatStore.getState().loadingSessionIds.has(sessionId);
 
@@ -93,7 +111,7 @@ export async function handleSessionNotification(
     }
     perf.lastAt = now;
     perf.count += 1;
-    handleReplay(sessionId, gooseSessionId, update);
+    handleReplay(sessionId, gooseSessionId, localSessionId, update);
   } else {
     const perf = livePerf.get(gooseSessionId);
     if (perf && update.sessionUpdate === "agent_message_chunk") {
@@ -106,7 +124,7 @@ export async function handleSessionNotification(
         );
       }
     }
-    handleLive(sessionId, gooseSessionId, update);
+    handleLive(sessionId, gooseSessionId, localSessionId, update);
   }
 }
 
@@ -125,6 +143,7 @@ export function clearReplayPerf(sessionId: string): void {
 function handleReplay(
   sessionId: string,
   gooseSessionId: string,
+  localSessionId: string | null,
   update: SessionUpdate,
 ): void {
   switch (update.sessionUpdate) {
@@ -247,7 +266,7 @@ function handleReplay(
     case "session_info_update":
     case "config_option_update":
     case "usage_update":
-      handleShared(sessionId, update);
+      handleShared(sessionId, gooseSessionId, localSessionId, update);
       break;
 
     default:
@@ -258,6 +277,7 @@ function handleReplay(
 function handleLive(
   sessionId: string,
   gooseSessionId: string,
+  localSessionId: string | null,
   update: SessionUpdate,
 ): void {
   const store = useChatStore.getState();
@@ -350,15 +370,20 @@ function handleLive(
     case "session_info_update":
     case "config_option_update":
     case "usage_update":
-      handleShared(sessionId, update);
+      handleShared(sessionId, gooseSessionId, localSessionId, update);
       break;
 
     default:
       break;
   }
 }
 
-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 & {
@@ -410,7 +435,6 @@ function handleShared(sessionId: string, update: SessionUpdate): void {
             currentModelId;
 
           const sessionStore = useChatSessionStore.getState();
-          sessionStore.setSessionModels(sessionId, availableModels);
           sessionStore.updateSession(
             sessionId,
             { modelId: currentModelId, modelName: currentModelName },
@@ -423,7 +447,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,
       });
@@ -485,6 +518,7 @@ function ensureLiveAssistantMessage(
 
 export function clearMessageTracking(): void {
   presetMessageIds.clear();
+  pendingUsageUpdates.clear();
   clearReplayAssistantTracking();
 }