From de471bc2b2a34de831542ae50c461826fa61460d Mon Sep 17 00:00:00 2001 From: Lifei Zhou Date: Tue, 5 May 2026 12:40:34 +1000 Subject: [PATCH] refactor: goose 2 ui used acp session id (#8985) --- ui/goose2/src/app/AppShell.tsx | 9 +- .../__tests__/useChat.compaction.test.ts | 49 ++--- .../useChat.personaPreparation.test.ts | 7 - .../__tests__/useChat.skillChips.test.ts | 7 - .../chat/hooks/__tests__/useChat.test.ts | 96 -------- .../useChatSessionController.test.ts | 2 - ui/goose2/src/features/chat/hooks/useChat.ts | 67 ++---- .../chat/hooks/useChatSessionController.ts | 14 +- .../stores/__tests__/chatSessionStore.test.ts | 40 +--- .../features/chat/stores/chatSessionStore.ts | 31 +-- ui/goose2/src/features/chat/ui/McpAppView.tsx | 11 +- .../chat/ui/__tests__/McpAppView.test.tsx | 1 - .../__tests__/MessageBubble.mcpApp.test.tsx | 3 +- .../chat/ui/__tests__/mcpAppPayload.test.ts | 1 - .../hooks/__tests__/useSessionSearch.test.ts | 3 +- .../sessions/hooks/useSessionSearch.ts | 8 +- .../lib/buildSessionSearchResults.test.ts | 5 +- .../sessions/lib/buildSessionSearchResults.ts | 6 +- .../src/shared/api/__tests__/acp.test.ts | 98 +++++++-- .../__tests__/acpNotificationHandler.test.ts | 64 +++--- .../__tests__/acpSessionInfoUpdate.test.ts | 2 - .../api/__tests__/acpToolCallStatus.test.ts | 18 +- ui/goose2/src/shared/api/acp.ts | 69 ++---- .../shared/api/acpNotificationHandler.test.ts | 49 +---- .../src/shared/api/acpNotificationHandler.ts | 101 ++------- .../src/shared/api/acpSessionRegistry.ts | 80 +++++++ ui/goose2/src/shared/api/acpSessionTracker.ts | 206 ------------------ .../src/shared/api/acpToolCallContent.ts | 2 - ui/goose2/src/shared/api/mcpAppToolUpdate.ts | 4 - ui/goose2/src/shared/types/messages.ts | 1 - 30 files changed, 319 insertions(+), 735 deletions(-) create mode 100644 ui/goose2/src/shared/api/acpSessionRegistry.ts delete mode 100644 ui/goose2/src/shared/api/acpSessionTracker.ts diff --git a/ui/goose2/src/app/AppShell.tsx b/ui/goose2/src/app/AppShell.tsx index 88d4a99156..8c9c323979 100644 --- a/ui/goose2/src/app/AppShell.tsx +++ b/ui/goose2/src/app/AppShell.tsx @@ -107,14 +107,13 @@ export function AppShell({ children }: { children?: React.ReactNode }) { const t1 = performance.now(); perfLog(`[perf:load] ${sid} import in ${(t1 - t0).toFixed(1)}ms`); const session = useChatSessionStore.getState().getSession(sessionId); - const gooseSessionId = session?.acpSessionId ?? sessionId; const project = session?.projectId ? (useProjectStore .getState() .projects.find((p) => p.id === session.projectId) ?? null) : null; const workingDir = await resolveSessionCwd(project); - await acpLoadSession(sessionId, gooseSessionId, workingDir); + await acpLoadSession(sessionId, workingDir); const tFlush = performance.now(); useChatStore.getState().setSessionLoading(sessionId, false); const buffer = getAndDeleteReplayBuffer(sessionId); @@ -197,9 +196,6 @@ export function AppShell({ children }: { children?: React.ReactNode }) { homeSession.id, sessionModelPreference.providerId, workingDir, - { - personaId: homeSession.personaId, - }, ); const shouldClearHomeModel = sessionModelPreference.providerId !== homeSession.providerId || @@ -468,9 +464,6 @@ export function AppShell({ children }: { children?: React.ReactNode }) { sessionId, session.providerId ?? agentStore.selectedProvider ?? "goose", workingDir, - { - personaId: session.personaId, - }, ); })().catch((error) => { console.error( diff --git a/ui/goose2/src/features/chat/hooks/__tests__/useChat.compaction.test.ts b/ui/goose2/src/features/chat/hooks/__tests__/useChat.compaction.test.ts index 3d6ef86be7..295c0fd69e 100644 --- a/ui/goose2/src/features/chat/hooks/__tests__/useChat.compaction.test.ts +++ b/ui/goose2/src/features/chat/hooks/__tests__/useChat.compaction.test.ts @@ -6,7 +6,6 @@ import { clearReplayBuffer, ensureReplayBuffer } from "../replayBuffer"; const mockAcpSendMessage = vi.fn(); const mockAcpLoadSession = vi.fn(); -const mockGetGooseSessionId = vi.fn(); vi.mock("@/shared/api/acp", () => ({ acpSendMessage: (...args: unknown[]) => mockAcpSendMessage(...args), @@ -16,10 +15,6 @@ vi.mock("@/shared/api/acp", () => ({ acpSetModel: vi.fn(), })); -vi.mock("@/shared/api/acpSessionTracker", () => ({ - getGooseSessionId: (...args: unknown[]) => mockGetGooseSessionId(...args), -})); - import { useChat } from "../useChat"; function createDeferredPromise() { @@ -51,7 +46,6 @@ describe("useChat compaction", () => { beforeEach(() => { mockAcpSendMessage.mockReset(); mockAcpLoadSession.mockReset(); - mockGetGooseSessionId.mockReset(); clearReplayBuffer("session-1"); useChatStore.setState({ messagesBySession: {}, @@ -62,11 +56,9 @@ describe("useChat compaction", () => { }); mockAcpSendMessage.mockResolvedValue(undefined); mockAcpLoadSession.mockResolvedValue(undefined); - mockGetGooseSessionId.mockReturnValue(null); }); it("reloads compacted history after sending the compact command", async () => { - mockGetGooseSessionId.mockReturnValue("goose-session-1"); mockAcpLoadSession.mockImplementation(async (sessionId: string) => { const buffer = ensureReplayBuffer(sessionId); buffer.push(createTextMessage("user-1", "user", "Before compact")); @@ -93,11 +85,7 @@ describe("useChat compaction", () => { "/compact", undefined, ); - expect(mockAcpLoadSession).toHaveBeenCalledWith( - "session-1", - "goose-session-1", - undefined, - ); + expect(mockAcpLoadSession).toHaveBeenCalledWith("session-1", undefined); const messages = useChatStore.getState().messagesBySession["session-1"]; const runtime = useChatStore.getState().getSessionRuntime("session-1"); @@ -135,12 +123,6 @@ describe("useChat compaction", () => { const ensurePrepared = vi.fn(async (personaId?: string) => { preparedPersonaId = personaId; }); - mockGetGooseSessionId.mockImplementation( - (_sessionId: string, personaId?: string) => - personaId === "persona-a" && preparedPersonaId === "persona-a" - ? "goose-session-a" - : null, - ); const { result } = renderHook(() => useChat( @@ -160,15 +142,11 @@ describe("useChat compaction", () => { expect(mockAcpSendMessage).toHaveBeenCalledWith("session-1", "/compact", { personaId: "persona-a", }); - expect(mockAcpLoadSession).toHaveBeenCalledWith( - "session-1", - "goose-session-a", - undefined, - ); + expect(mockAcpLoadSession).toHaveBeenCalledWith("session-1", undefined); + expect(preparedPersonaId).toBe("persona-a"); }); it("blocks new sends while compaction is in flight", async () => { - mockGetGooseSessionId.mockReturnValue("goose-session-1"); const compactDeferred = createDeferredPromise(); mockAcpSendMessage.mockImplementation( (_sessionId: string, prompt: string) => @@ -215,7 +193,6 @@ describe("useChat compaction", () => { }); it("ignores a second compact request while the first one is still in flight", async () => { - mockGetGooseSessionId.mockReturnValue("goose-session-1"); const compactDeferred = createDeferredPromise(); mockAcpSendMessage.mockImplementation( (_sessionId: string, prompt: string) => @@ -254,13 +231,22 @@ describe("useChat compaction", () => { ).toBe("idle"); }); - it("surfaces an error when compacting before the session is prepared", async () => { - const { result } = renderHook(() => useChat("session-1")); + it("surfaces an error when preparing for compaction fails", async () => { + const ensurePrepared = vi + .fn() + .mockRejectedValue(new Error("prepare failed")); + + const { result } = renderHook(() => + useChat("session-1", undefined, undefined, undefined, { + ensurePrepared, + }), + ); await act(async () => { await result.current.compactConversation(); }); + expect(ensurePrepared).toHaveBeenCalledWith(undefined); expect(mockAcpSendMessage).not.toHaveBeenCalled(); expect(mockAcpLoadSession).not.toHaveBeenCalled(); @@ -272,11 +258,10 @@ describe("useChat compaction", () => { { type: "systemNotification", notificationType: "error", - text: "Session not prepared. Send a message before compacting.", + text: "prepare failed", }, ]); - expect(runtime.error).toBe( - "Session not prepared. Send a message before compacting.", - ); + expect(runtime.error).toBe("prepare failed"); + expect(runtime.chatState).toBe("idle"); }); }); diff --git a/ui/goose2/src/features/chat/hooks/__tests__/useChat.personaPreparation.test.ts b/ui/goose2/src/features/chat/hooks/__tests__/useChat.personaPreparation.test.ts index a4e19cfc0a..70edd53968 100644 --- a/ui/goose2/src/features/chat/hooks/__tests__/useChat.personaPreparation.test.ts +++ b/ui/goose2/src/features/chat/hooks/__tests__/useChat.personaPreparation.test.ts @@ -8,7 +8,6 @@ import { clearReplayBuffer } from "../replayBuffer"; const mockAcpSendMessage = vi.fn(); const mockAcpCancelSession = vi.fn(); const mockAcpLoadSession = vi.fn(); -const mockGetGooseSessionId = vi.fn(); vi.mock("@/shared/api/acp", () => ({ acpSendMessage: (...args: unknown[]) => mockAcpSendMessage(...args), @@ -16,10 +15,6 @@ vi.mock("@/shared/api/acp", () => ({ acpLoadSession: (...args: unknown[]) => mockAcpLoadSession(...args), })); -vi.mock("@/shared/api/acpSessionTracker", () => ({ - getGooseSessionId: (...args: unknown[]) => mockGetGooseSessionId(...args), -})); - import { useChat } from "../useChat"; describe("useChat persona preparation", () => { @@ -27,7 +22,6 @@ describe("useChat persona preparation", () => { mockAcpSendMessage.mockReset(); mockAcpCancelSession.mockReset(); mockAcpLoadSession.mockReset(); - mockGetGooseSessionId.mockReset(); clearReplayBuffer("session-1"); useChatStore.setState({ messagesBySession: {}, @@ -72,7 +66,6 @@ describe("useChat persona preparation", () => { mockAcpSendMessage.mockResolvedValue(undefined); mockAcpCancelSession.mockResolvedValue(true); mockAcpLoadSession.mockResolvedValue(undefined); - mockGetGooseSessionId.mockReturnValue(null); }); it("prepares the override persona before prompting", async () => { diff --git a/ui/goose2/src/features/chat/hooks/__tests__/useChat.skillChips.test.ts b/ui/goose2/src/features/chat/hooks/__tests__/useChat.skillChips.test.ts index f73b955049..c359585f01 100644 --- a/ui/goose2/src/features/chat/hooks/__tests__/useChat.skillChips.test.ts +++ b/ui/goose2/src/features/chat/hooks/__tests__/useChat.skillChips.test.ts @@ -9,7 +9,6 @@ const mockAcpCancelSession = vi.fn(); const mockAcpLoadSession = vi.fn(); const mockAcpPrepareSession = vi.fn(); const mockAcpSetModel = vi.fn(); -const mockGetGooseSessionId = vi.fn(); vi.mock("@/shared/api/acp", () => ({ acpSendMessage: (...args: unknown[]) => mockAcpSendMessage(...args), @@ -19,10 +18,6 @@ vi.mock("@/shared/api/acp", () => ({ acpSetModel: (...args: unknown[]) => mockAcpSetModel(...args), })); -vi.mock("@/shared/api/acpSessionTracker", () => ({ - getGooseSessionId: (...args: unknown[]) => mockGetGooseSessionId(...args), -})); - import { useChat } from "../useChat"; describe("useChat skill chips", () => { @@ -32,9 +27,7 @@ describe("useChat skill chips", () => { mockAcpLoadSession.mockReset(); mockAcpPrepareSession.mockReset(); mockAcpSetModel.mockReset(); - mockGetGooseSessionId.mockReset(); mockAcpSendMessage.mockResolvedValue(undefined); - mockGetGooseSessionId.mockReturnValue(null); useChatStore.setState({ messagesBySession: {}, sessionStateById: {}, diff --git a/ui/goose2/src/features/chat/hooks/__tests__/useChat.test.ts b/ui/goose2/src/features/chat/hooks/__tests__/useChat.test.ts index 039fda785a..bd2d91a521 100644 --- a/ui/goose2/src/features/chat/hooks/__tests__/useChat.test.ts +++ b/ui/goose2/src/features/chat/hooks/__tests__/useChat.test.ts @@ -11,7 +11,6 @@ const mockAcpCancelSession = vi.fn(); const mockAcpLoadSession = vi.fn(); const mockAcpPrepareSession = vi.fn(); const mockAcpSetModel = vi.fn(); -const mockGetGooseSessionId = vi.fn(); vi.mock("@/shared/api/acp", () => ({ acpSendMessage: (...args: unknown[]) => mockAcpSendMessage(...args), @@ -21,10 +20,6 @@ vi.mock("@/shared/api/acp", () => ({ acpSetModel: (...args: unknown[]) => mockAcpSetModel(...args), })); -vi.mock("@/shared/api/acpSessionTracker", () => ({ - getGooseSessionId: (...args: unknown[]) => mockGetGooseSessionId(...args), -})); - import { useChat } from "../useChat"; function addStreamingAssistantMessage( @@ -66,7 +61,6 @@ describe("useChat", () => { mockAcpLoadSession.mockReset(); mockAcpPrepareSession.mockReset(); mockAcpSetModel.mockReset(); - mockGetGooseSessionId.mockReset(); clearReplayBuffer("session-1"); clearReplayBuffer("session-2"); useChatStore.setState({ @@ -115,96 +109,6 @@ describe("useChat", () => { mockAcpLoadSession.mockResolvedValue(undefined); mockAcpPrepareSession.mockResolvedValue(undefined); mockAcpSetModel.mockResolvedValue(undefined); - mockGetGooseSessionId.mockReturnValue(null); - }); - - it("cancels the active override persona instead of the hook default persona", async () => { - const deferred = createDeferredPromise(); - mockAcpSendMessage.mockReturnValue(deferred.promise); - - const { result } = renderHook(() => - useChat("session-1", undefined, undefined, { - id: "persona-a", - name: "Persona A", - }), - ); - - let sendPromise!: Promise; - await act(async () => { - sendPromise = result.current.sendMessage("Hello", { - id: "persona-b", - name: "Persona B", - }); - await Promise.resolve(); - }); - - act(() => { - result.current.stopGeneration(); - }); - - expect(mockAcpSendMessage).toHaveBeenCalledWith("session-1", "Hello", { - systemPrompt: undefined, - personaId: "persona-b", - personaName: "Persona B", - images: undefined, - }); - expect(mockAcpCancelSession).toHaveBeenCalledWith("session-1", "persona-b"); - - deferred.resolve(); - await act(async () => { - await sendPromise; - }); - }); - - it("keeps persona-aware cancellation working after remount", async () => { - const deferred = createDeferredPromise(); - mockAcpSendMessage.mockReturnValue(deferred.promise); - - const firstMount = renderHook(() => - useChat("session-1", undefined, undefined, { - id: "persona-a", - name: "Persona A", - }), - ); - - let sendPromise!: Promise; - await act(async () => { - sendPromise = firstMount.result.current.sendMessage("Hello", { - id: "persona-b", - name: "Persona B", - }); - await Promise.resolve(); - }); - act(() => { - addStreamingAssistantMessage( - "session-1", - "assistant-1", - "persona-b", - "Persona B", - ); - }); - - act(() => { - firstMount.unmount(); - }); - - const secondMount = renderHook(() => - useChat("session-1", undefined, undefined, { - id: "persona-a", - name: "Persona A", - }), - ); - - act(() => { - secondMount.result.current.stopGeneration(); - }); - - expect(mockAcpCancelSession).toHaveBeenCalledWith("session-1", "persona-b"); - - deferred.resolve(); - await act(async () => { - await sendPromise; - }); }); it("marks the streaming message stopped only after cancellation succeeds", async () => { diff --git a/ui/goose2/src/features/chat/hooks/__tests__/useChatSessionController.test.ts b/ui/goose2/src/features/chat/hooks/__tests__/useChatSessionController.test.ts index 2928ad1437..5ad52eac54 100644 --- a/ui/goose2/src/features/chat/hooks/__tests__/useChatSessionController.test.ts +++ b/ui/goose2/src/features/chat/hooks/__tests__/useChatSessionController.test.ts @@ -194,7 +194,6 @@ describe("useChatSessionController", () => { "session-1", "anthropic", "/tmp/project", - { personaId: undefined }, ); }); @@ -342,7 +341,6 @@ describe("useChatSessionController", () => { "session-2", "anthropic", "/tmp/project", - { personaId: undefined }, ); }); diff --git a/ui/goose2/src/features/chat/hooks/useChat.ts b/ui/goose2/src/features/chat/hooks/useChat.ts index 3dd2b0589e..409c2c526f 100644 --- a/ui/goose2/src/features/chat/hooks/useChat.ts +++ b/ui/goose2/src/features/chat/hooks/useChat.ts @@ -13,7 +13,6 @@ import { acpCancelSession, acpLoadSession, } from "@/shared/api/acp"; -import { getGooseSessionId } from "@/shared/api/acpSessionTracker"; import { useAgentStore } from "@/features/agents/stores/agentStore"; import { getSessionTitleFromDraft, @@ -107,24 +106,12 @@ export function useChat( ) { const store = useChatStore(); const abortRef = useRef(null); - const streamingPersonaIdRef = useRef(null); const messages = store.messagesBySession[sessionId] ?? []; const { chatState, tokenState, error, streamingMessageId } = store.getSessionRuntime(sessionId); const isStreaming = chatState === "streaming" || streamingMessageId !== null; - const getStreamingPersonaId = useCallback(() => { - if (!streamingMessageId) { - return null; - } - - return ( - messages.find((message) => message.id === streamingMessageId)?.metadata - ?.personaId ?? null - ); - }, [messages, streamingMessageId]); - const resolvePersonaInfo = useCallback( (overridePersonaId?: string, overridePersonaName?: string) => { if (overridePersonaId) { @@ -237,7 +224,6 @@ export function useChat( const abort = new AbortController(); abortRef.current = abort; - streamingPersonaIdRef.current = effectivePersonaInfo?.id ?? null; try { await options?.ensurePrepared?.(effectivePersonaInfo?.id); @@ -302,7 +288,6 @@ export function useChat( store.setPendingAssistantProvider(sessionId, null); } finally { abortRef.current = null; - streamingPersonaIdRef.current = null; } }, [ @@ -317,8 +302,6 @@ export function useChat( const stopGeneration = useCallback(() => { abortRef.current?.abort(); - const activePersonaId = - streamingPersonaIdRef.current ?? getStreamingPersonaId(); const activeStreamingMessageId = useChatStore .getState() .getSessionRuntime(sessionId).streamingMessageId; @@ -327,7 +310,7 @@ export function useChat( store.setStreamingMessageId(sessionId, null); store.setPendingAssistantProvider(sessionId, null); // Cancel the backend ACP session to stop orphaned streaming events - acpCancelSession(sessionId, activePersonaId ?? undefined) + acpCancelSession(sessionId) .then((wasCancelled) => { if (wasCancelled && activeStreamingMessageId) { markMessageStopped(sessionId, activeStreamingMessageId); @@ -336,7 +319,7 @@ export function useChat( .catch(() => { // Best-effort cancellation — ignore errors }); - }, [getStreamingPersonaId, store, sessionId]); + }, [store, sessionId]); const retryLastMessage = useCallback(async () => { const sessionMessages = store.messagesBySession[sessionId] ?? []; @@ -400,41 +383,25 @@ export function useChat( overridePersona?.id, overridePersona?.name, ); - let gooseSessionId = getGooseSessionId( - sessionId, - effectivePersonaInfo?.id, - ); - - if (!gooseSessionId) { - try { - await options?.ensurePrepared?.(effectivePersonaInfo?.id); - } catch (err) { - const errorMessage = getErrorMessage(err); - store.addMessage( - sessionId, - createSystemNotificationMessage(errorMessage, "error"), - ); - store.setError(sessionId, errorMessage); - return "failed" as CompactConversationResult; - } - gooseSessionId = getGooseSessionId(sessionId, effectivePersonaInfo?.id); - } - - if (!gooseSessionId) { - const errorMessage = - "Session not prepared. Send a message before compacting."; - store.addMessage( - sessionId, - createSystemNotificationMessage(errorMessage, "error"), - ); - store.setError(sessionId, errorMessage); - return "failed" as CompactConversationResult; - } store.setActiveSession(sessionId); store.setChatState(sessionId, "compacting"); store.setStreamingMessageId(sessionId, null); store.setError(sessionId, null); + + try { + await options?.ensurePrepared?.(effectivePersonaInfo?.id); + } catch (err) { + const errorMessage = getErrorMessage(err); + store.addMessage( + sessionId, + createSystemNotificationMessage(errorMessage, "error"), + ); + store.setError(sessionId, errorMessage); + store.setChatState(sessionId, "idle"); + return "failed" as CompactConversationResult; + } + store.setSessionLoading(sessionId, true); clearReplayBuffer(sessionId); @@ -449,7 +416,7 @@ export function useChat( // transient chunks and refresh the session from replay instead. clearReplayBuffer(sessionId); const workingDir = getWorkingDir(); - await acpLoadSession(sessionId, gooseSessionId, workingDir); + await acpLoadSession(sessionId, workingDir); store.setSessionLoading(sessionId, false); diff --git a/ui/goose2/src/features/chat/hooks/useChatSessionController.ts b/ui/goose2/src/features/chat/hooks/useChatSessionController.ts index 4488e9aa00..fa54bdb0cd 100644 --- a/ui/goose2/src/features/chat/hooks/useChatSessionController.ts +++ b/ui/goose2/src/features/chat/hooks/useChatSessionController.ts @@ -157,7 +157,6 @@ export function useChatSessionController({ providerId: string, nextProject = project, nextWorkspacePath = activeWorkspace?.path, - personaId = selectedPersonaId ?? undefined, modelSelection?: PreferredModelSelection | null, ) => { if (!sessionId) { @@ -167,7 +166,7 @@ export function useChatSessionController({ nextProject, nextWorkspacePath, ); - await acpPrepareSession(sessionId, providerId, workingDir, { personaId }); + await acpPrepareSession(sessionId, providerId, workingDir); if (!modelSelection?.id) { return; } @@ -188,7 +187,7 @@ export function useChatSessionController({ modelName: modelSelection.name, }); }, - [activeWorkspace?.path, project, selectedPersonaId, sessionId], + [activeWorkspace?.path, project, sessionId], ); const prepareSelectedProvider = useCallback( (providerId: string, modelSelection?: PreferredModelSelection | null) => @@ -196,10 +195,9 @@ export function useChatSessionController({ providerId, project, activeWorkspace?.path, - selectedPersonaId ?? undefined, modelSelection, ), - [activeWorkspace?.path, prepareCurrentSession, project, selectedPersonaId], + [activeWorkspace?.path, prepareCurrentSession, project], ); const prevProjectIdRef = useRef(session?.projectId); @@ -323,7 +321,6 @@ export function useChatSessionController({ selectedProvider, nextProject, activeWorkspace?.path, - selectedPersonaId ?? undefined, effectiveModelSelection, ).catch((error) => { console.error("Failed to update ACP session working directory:", error); @@ -333,7 +330,6 @@ export function useChatSessionController({ activeWorkspace?.path, effectiveModelSelection, prepareCurrentSession, - selectedPersonaId, selectedProvider, sessionId, ], @@ -423,12 +419,11 @@ export function useChatSessionController({ { onMessageAccepted: sessionId ? onMessageAccepted : undefined, ensurePrepared: selectedProvider - ? (personaId?: string) => + ? () => prepareCurrentSession( selectedProvider, project, activeWorkspace?.path, - personaId, ) : undefined, }, @@ -719,7 +714,6 @@ export function useChatSessionController({ nextProviderId, nextProject, activeWorkspace?.path, - nextPersonaId, pendingModelSelection, ); if (cancelled) { diff --git a/ui/goose2/src/features/chat/stores/__tests__/chatSessionStore.test.ts b/ui/goose2/src/features/chat/stores/__tests__/chatSessionStore.test.ts index 0a22302d74..735be14993 100644 --- a/ui/goose2/src/features/chat/stores/__tests__/chatSessionStore.test.ts +++ b/ui/goose2/src/features/chat/stores/__tests__/chatSessionStore.test.ts @@ -31,7 +31,6 @@ function resetStore() { function makeSession(overrides: Partial = {}): ChatSession { return { id: "session-1", - acpSessionId: "session-1", title: "Test Session", createdAt: "2026-04-01T00:00:00.000Z", updatedAt: "2026-04-01T00:00:00.000Z", @@ -42,7 +41,9 @@ function makeSession(overrides: Partial = {}): ChatSession { function seedSession(overrides: Partial = {}): ChatSession { const session = makeSession(overrides); - useChatSessionStore.getState().addSession(session); + useChatSessionStore.setState((state) => ({ + sessions: [session, ...state.sessions], + })); return session; } @@ -59,6 +60,7 @@ describe("chatSessionStore", () => { const session = await useChatSessionStore.getState().createSession({ title: "New Chat", providerId: "openai", + projectId: "project-1", personaId: "persona-1", modelId: "gpt-4.1", modelName: "GPT-4.1", @@ -69,14 +71,15 @@ describe("chatSessionStore", () => { "openai", "/tmp/project", { + projectId: "project-1", personaId: "persona-1", modelId: "gpt-4.1", }, ); expect(session).toMatchObject({ id: "acp-1", - acpSessionId: "acp-1", title: "New Chat", + projectId: "project-1", providerId: "openai", personaId: "persona-1", modelId: "gpt-4.1", @@ -312,35 +315,4 @@ describe("chatSessionStore", () => { expect(useChatSessionStore.getState().activeSessionId).toBeNull(); }); }); - - describe("addSession", () => { - it("prepends a new session to the list", () => { - const { addSession } = useChatSessionStore.getState(); - addSession( - makeSession({ - id: "imported-1", - title: "Imported Session", - messageCount: 5, - }), - ); - - const sessions = useChatSessionStore.getState().sessions; - expect(sessions[0].id).toBe("imported-1"); - expect(sessions[0].title).toBe("Imported Session"); - expect(sessions[0].messageCount).toBe(5); - }); - - it("does not create a duplicate if session ID already exists", () => { - const { addSession } = useChatSessionStore.getState(); - addSession(makeSession({ id: "dup-1", title: "First", messageCount: 1 })); - addSession( - makeSession({ id: "dup-1", title: "Second", messageCount: 2 }), - ); - - const sessions = useChatSessionStore.getState().sessions; - const matches = sessions.filter((session) => session.id === "dup-1"); - expect(matches).toHaveLength(1); - expect(matches[0].title).toBe("Second"); - }); - }); }); diff --git a/ui/goose2/src/features/chat/stores/chatSessionStore.ts b/ui/goose2/src/features/chat/stores/chatSessionStore.ts index 7a60095333..a421d6a0df 100644 --- a/ui/goose2/src/features/chat/stores/chatSessionStore.ts +++ b/ui/goose2/src/features/chat/stores/chatSessionStore.ts @@ -18,7 +18,6 @@ import { export interface ChatSession { id: string; - acpSessionId?: string; title: string; projectId?: string | null; providerId?: string; @@ -100,7 +99,6 @@ function acpSessionToChatSession(session: AcpSessionInfo): ChatSession { const now = new Date().toISOString(); return { id: session.sessionId, - acpSessionId: session.sessionId, title: normalizeAcpTitle(session.title) ?? "Untitled", projectId: session.projectId ?? undefined, providerId: session.providerId ?? undefined, @@ -124,7 +122,6 @@ function sortByUpdatedAtDesc(sessions: ChatSession[]): ChatSession[] { export function sessionToChatSession(session: Session): ChatSession { return { id: session.id, - acpSessionId: session.id, title: session.title, projectId: session.projectId, providerId: session.providerId, @@ -161,7 +158,6 @@ export const useChatSessionStore = create((set, get) => ({ }); const chatSession: ChatSession = { id: sessionId, - acpSessionId: sessionId, title: opts.title ?? DEFAULT_CHAT_TITLE, projectId: opts.projectId, providerId, @@ -213,24 +209,23 @@ export const useChatSessionStore = create((set, get) => ({ })); const updatedSession = get().sessions.find((session) => session.id === id); - const acpSessionId = updatedSession?.acpSessionId; // Persist title rename to backend if ( "title" in patch && "userSetName" in patch && patch.userSetName && - acpSessionId && + updatedSession && patch.title ) { - acpRenameSession(acpSessionId, patch.title).catch((err: unknown) => + acpRenameSession(updatedSession.id, patch.title).catch((err: unknown) => console.error("Failed to rename session in backend:", err), ); } // Persist projectId change to backend - if ("projectId" in patch && acpSessionId) { - updateSessionProject(acpSessionId, patch.projectId ?? null).catch( + if ("projectId" in patch && updatedSession) { + updateSessionProject(updatedSession.id, patch.projectId ?? null).catch( (err: unknown) => console.error("Failed to update session project in backend:", err), ); @@ -238,20 +233,16 @@ export const useChatSessionStore = create((set, get) => ({ }, addSession: (session) => { - const normalizedSession = { - ...session, - acpSessionId: session.acpSessionId ?? session.id, - }; set((state) => { const existing = state.sessions.findIndex( - (candidate) => candidate.id === normalizedSession.id, + (candidate) => candidate.id === session.id, ); if (existing >= 0) { const updated = [...state.sessions]; - updated[existing] = { ...updated[existing], ...normalizedSession }; + updated[existing] = { ...updated[existing], ...session }; return { sessions: updated }; } - return { sessions: [normalizedSession, ...state.sessions] }; + return { sessions: [session, ...state.sessions] }; }); }, @@ -266,8 +257,8 @@ export const useChatSessionStore = create((set, get) => ({ state.activeSessionId === id ? null : state.activeSessionId, })); const session = get().sessions.find((candidate) => candidate.id === id); - if (session?.acpSessionId) { - acpArchiveSession(session.acpSessionId).catch((err: unknown) => + if (session) { + acpArchiveSession(session.id).catch((err: unknown) => console.error("Failed to archive session in backend:", err), ); } @@ -280,8 +271,8 @@ export const useChatSessionStore = create((set, get) => ({ ), })); const session = get().sessions.find((candidate) => candidate.id === id); - if (session?.acpSessionId) { - acpUnarchiveSession(session.acpSessionId).catch((err: unknown) => + if (session) { + acpUnarchiveSession(session.id).catch((err: unknown) => console.error("Failed to unarchive session in backend:", err), ); } diff --git a/ui/goose2/src/features/chat/ui/McpAppView.tsx b/ui/goose2/src/features/chat/ui/McpAppView.tsx index fd53c248e9..bea9a71837 100644 --- a/ui/goose2/src/features/chat/ui/McpAppView.tsx +++ b/ui/goose2/src/features/chat/ui/McpAppView.tsx @@ -327,13 +327,11 @@ export function McpAppView({ name: string; arguments?: Record; }) => { - const acpSessionId = payload.gooseSessionId ?? payload.sessionId; - setActiveToolInput(args ?? {}); const client = await getClient(); const response = (await client.extMethod("_goose/tool/call", { - sessionId: acpSessionId, + sessionId: payload.sessionId, name: `${payload.tool.extensionName}__${name}`, arguments: args ?? {}, })) as GooseToolCallResponse; @@ -348,22 +346,21 @@ export function McpAppView({ return toolResult; }, - [payload.gooseSessionId, payload.sessionId, payload.tool.extensionName], + [payload.sessionId, payload.tool.extensionName], ); const handleReadResource = useCallback( async ({ uri }: { uri: string }) => { - const acpSessionId = payload.gooseSessionId ?? payload.sessionId; const client = await getClient(); const response = await client.goose.GooseResourceRead({ - sessionId: acpSessionId, + sessionId: payload.sessionId, uri, extensionName: payload.tool.extensionName, }); return (response.result ?? { contents: [] }) as ReadResourceResult; }, - [payload.gooseSessionId, payload.sessionId, payload.tool.extensionName], + [payload.sessionId, payload.tool.extensionName], ); const handleSizeChanged = useCallback( diff --git a/ui/goose2/src/features/chat/ui/__tests__/McpAppView.test.tsx b/ui/goose2/src/features/chat/ui/__tests__/McpAppView.test.tsx index 105848aa3a..6ebaa27a8c 100644 --- a/ui/goose2/src/features/chat/ui/__tests__/McpAppView.test.tsx +++ b/ui/goose2/src/features/chat/ui/__tests__/McpAppView.test.tsx @@ -71,7 +71,6 @@ function createPayload({ } = {}): McpAppPayload { return { sessionId: "local-session", - gooseSessionId: null, toolCallId: "tool-1", toolCallTitle: "inspect messaging", source: "toolCallUpdateMeta", diff --git a/ui/goose2/src/features/chat/ui/__tests__/MessageBubble.mcpApp.test.tsx b/ui/goose2/src/features/chat/ui/__tests__/MessageBubble.mcpApp.test.tsx index 640f618a0e..3c8d2431b2 100644 --- a/ui/goose2/src/features/chat/ui/__tests__/MessageBubble.mcpApp.test.tsx +++ b/ui/goose2/src/features/chat/ui/__tests__/MessageBubble.mcpApp.test.tsx @@ -67,8 +67,7 @@ describe("MessageBubble MCP app rendering", () => { type: "mcpApp", id: "tool-1", payload: { - sessionId: "local-session", - gooseSessionId: "goose-session", + sessionId: "acp-session", toolCallId: "tool-1", toolCallTitle: "weather: open app", source: "toolCallUpdateMeta", diff --git a/ui/goose2/src/features/chat/ui/__tests__/mcpAppPayload.test.ts b/ui/goose2/src/features/chat/ui/__tests__/mcpAppPayload.test.ts index f082f7b0f1..8045a90660 100644 --- a/ui/goose2/src/features/chat/ui/__tests__/mcpAppPayload.test.ts +++ b/ui/goose2/src/features/chat/ui/__tests__/mcpAppPayload.test.ts @@ -5,7 +5,6 @@ import type { McpAppPayload } from "@/shared/types/messages"; function createPayload(csp: unknown): McpAppPayload { return { sessionId: "session-1", - gooseSessionId: null, toolCallId: "tool-1", toolCallTitle: "inspect app", source: "toolCallUpdateMeta", diff --git a/ui/goose2/src/features/sessions/hooks/__tests__/useSessionSearch.test.ts b/ui/goose2/src/features/sessions/hooks/__tests__/useSessionSearch.test.ts index 29b9c383df..f44e36c800 100644 --- a/ui/goose2/src/features/sessions/hooks/__tests__/useSessionSearch.test.ts +++ b/ui/goose2/src/features/sessions/hooks/__tests__/useSessionSearch.test.ts @@ -20,8 +20,7 @@ function createDeferredPromise() { const sessions: ChatSession[] = [ { - id: "session-1", - acpSessionId: "acp-1", + id: "acp-1", title: "Needle notes", createdAt: "2026-04-10T12:00:00Z", updatedAt: "2026-04-10T12:00:00Z", diff --git a/ui/goose2/src/features/sessions/hooks/useSessionSearch.ts b/ui/goose2/src/features/sessions/hooks/useSessionSearch.ts index 225d25f8b7..592ebc903e 100644 --- a/ui/goose2/src/features/sessions/hooks/useSessionSearch.ts +++ b/ui/goose2/src/features/sessions/hooks/useSessionSearch.ts @@ -72,10 +72,8 @@ export function useSessionSearch({ setError(null); setResults(metadataResults); - const acpSessionIds = sessions.map( - (session) => session.acpSessionId ?? session.id, - ); - if (trimmed.length < 2 || acpSessionIds.length === 0) { + const sessionIds = sessions.map((session) => session.id); + if (trimmed.length < 2 || sessionIds.length === 0) { setIsSearching(false); return; } @@ -83,7 +81,7 @@ export function useSessionSearch({ setIsSearching(true); try { - const messageResults = await acpSearchSessions(trimmed, acpSessionIds); + const messageResults = await acpSearchSessions(trimmed, sessionIds); if (requestIdRef.current !== requestId) { return; } diff --git a/ui/goose2/src/features/sessions/lib/buildSessionSearchResults.test.ts b/ui/goose2/src/features/sessions/lib/buildSessionSearchResults.test.ts index 9e3bdc97e9..1a4e4d14f6 100644 --- a/ui/goose2/src/features/sessions/lib/buildSessionSearchResults.test.ts +++ b/ui/goose2/src/features/sessions/lib/buildSessionSearchResults.test.ts @@ -23,8 +23,7 @@ describe("buildSessionSearchResults", () => { it("merges metadata and message matches, preferring message details", () => { const sessions = [ makeSession({ - id: "session-1", - acpSessionId: "acp-1", + id: "acp-1", title: "Needle session", updatedAt: "2026-04-10T12:00:00Z", }), @@ -59,7 +58,7 @@ describe("buildSessionSearchResults", () => { messageId: "message-1", matchCount: 2, }); - expect(results[0].session.id).toBe("session-1"); + expect(results[0].session.id).toBe("acp-1"); }); it("includes metadata-only matches and sorts by updatedAt descending", () => { diff --git a/ui/goose2/src/features/sessions/lib/buildSessionSearchResults.ts b/ui/goose2/src/features/sessions/lib/buildSessionSearchResults.ts index 7770d7e2a9..72967777f0 100644 --- a/ui/goose2/src/features/sessions/lib/buildSessionSearchResults.ts +++ b/ui/goose2/src/features/sessions/lib/buildSessionSearchResults.ts @@ -40,15 +40,13 @@ export function buildSessionSearchResults( return sortByUpdatedAtDesc(sessions) .filter((session) => { - const acpSessionId = session.acpSessionId ?? session.id; return ( metadataMatchIds.has(session.id) || - messageMatchesBySessionId.has(acpSessionId) + messageMatchesBySessionId.has(session.id) ); }) .map((session) => { - const acpSessionId = session.acpSessionId ?? session.id; - const messageMatch = messageMatchesBySessionId.get(acpSessionId); + const messageMatch = messageMatchesBySessionId.get(session.id); if (!messageMatch) { return { session, diff --git a/ui/goose2/src/shared/api/__tests__/acp.test.ts b/ui/goose2/src/shared/api/__tests__/acp.test.ts index 880a9fb4fd..e428afd930 100644 --- a/ui/goose2/src/shared/api/__tests__/acp.test.ts +++ b/ui/goose2/src/shared/api/__tests__/acp.test.ts @@ -1,13 +1,18 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; const mockLoadSession = vi.fn(); +const mockNewSession = vi.fn(); +const mockSetProvider = vi.fn(); +const mockSetModel = vi.fn(); vi.mock("../acpApi", () => ({ listProviders: vi.fn(), prompt: vi.fn(), - setModel: vi.fn(), + setModel: (...args: unknown[]) => mockSetModel(...args), + setProvider: (...args: unknown[]) => mockSetProvider(...args), listSessions: vi.fn(), loadSession: (...args: unknown[]) => mockLoadSession(...args), + newSession: (...args: unknown[]) => mockNewSession(...args), exportSession: vi.fn(), importSession: vi.fn(), forkSession: vi.fn(), @@ -29,29 +34,94 @@ describe("acpLoadSession", () => { vi.resetModules(); }); - it("restores the prior session mapping when replay loading fails", async () => { + it("restores the prior prepared session registration when replay loading fails", async () => { mockLoadSession.mockRejectedValueOnce(new Error("load failed")); - const sessionTracker = await import("../acpSessionTracker"); + const sessionRegistry = await import("../acpSessionRegistry"); const { acpLoadSession } = await import("../acp"); - sessionTracker.registerSession( - "local-session", - "goose-session-1", + sessionRegistry.registerPreparedSession( + "acp-session-1", "goose", "/tmp/original", ); await expect( - acpLoadSession("local-session", "goose-session-2", "/tmp/replay"), + acpLoadSession("acp-session-1", "/tmp/replay"), ).rejects.toThrow("load failed"); - expect(sessionTracker.getGooseSessionId("local-session")).toBe( - "goose-session-1", - ); - expect(sessionTracker.getLocalSessionId("goose-session-1")).toBe( - "local-session", - ); - expect(sessionTracker.getLocalSessionId("goose-session-2")).toBeNull(); + expect(sessionRegistry.isSessionPrepared("acp-session-1")).toBe(true); + }); +}); + +describe("acpCreateSession", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.resetModules(); + }); + + it("uses the ACP session id as the UI session id", async () => { + mockNewSession.mockResolvedValue({ sessionId: "acp-session-1" }); + + const sessionRegistry = await import("../acpSessionRegistry"); + const { acpCreateSession } = await import("../acp"); + + await expect( + acpCreateSession("openai", "/tmp/project", { + projectId: "project-1", + personaId: "persona-1", + modelId: "gpt-4.1", + }), + ).resolves.toEqual({ sessionId: "acp-session-1" }); + + expect(mockNewSession).toHaveBeenCalledWith( + "/tmp/project", + "openai", + "project-1", + "persona-1", + ); + expect(mockLoadSession).not.toHaveBeenCalled(); + expect(mockSetProvider).toHaveBeenCalledWith("acp-session-1", "openai"); + expect(mockSetModel).toHaveBeenCalledWith("acp-session-1", "gpt-4.1"); + expect(sessionRegistry.isSessionPrepared("acp-session-1")).toBe(true); + }); +}); + +describe("acpPrepareSession", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.resetModules(); + }); + + it("loads the existing ACP session instead of creating a replacement", async () => { + mockLoadSession.mockResolvedValue(undefined); + + const sessionRegistry = await import("../acpSessionRegistry"); + const { acpPrepareSession } = await import("../acp"); + + await expect( + acpPrepareSession("acp-session-1", "openai", "/tmp/project"), + ).resolves.toBeUndefined(); + + expect(mockLoadSession).toHaveBeenCalledWith( + "acp-session-1", + "/tmp/project", + ); + expect(mockNewSession).not.toHaveBeenCalled(); + expect(mockSetProvider).toHaveBeenCalledWith("acp-session-1", "openai"); + expect(sessionRegistry.isSessionPrepared("acp-session-1")).toBe(true); + }); + + it("surfaces load failures instead of creating a new ACP session", async () => { + mockLoadSession.mockRejectedValueOnce(new Error("missing session")); + + const { acpPrepareSession } = await import("../acp"); + + await expect( + acpPrepareSession("acp-session-1", "openai", "/tmp/project"), + ).rejects.toThrow("missing session"); + + expect(mockNewSession).not.toHaveBeenCalled(); + expect(mockSetProvider).not.toHaveBeenCalled(); }); }); diff --git a/ui/goose2/src/shared/api/__tests__/acpNotificationHandler.test.ts b/ui/goose2/src/shared/api/__tests__/acpNotificationHandler.test.ts index 41360f4e17..9eff09532c 100644 --- a/ui/goose2/src/shared/api/__tests__/acpNotificationHandler.test.ts +++ b/ui/goose2/src/shared/api/__tests__/acpNotificationHandler.test.ts @@ -11,12 +11,11 @@ import { handleSessionNotification, setActiveMessageId, } from "../acpNotificationHandler"; -import { registerSession } from "../acpSessionTracker"; +import { registerPreparedSession } from "../acpSessionRegistry"; function createMcpAppPayload(): McpAppPayload { return { - sessionId: "local-session", - gooseSessionId: "goose-session", + sessionId: "acp-session", toolCallId: "tool-1", toolCallTitle: "mcp_app_bench__inspect_host_info", source: "toolCallUpdateMeta", @@ -34,8 +33,7 @@ function createMcpAppPayload(): McpAppPayload { describe("acpNotificationHandler", () => { beforeEach(() => { clearMessageTracking(); - clearReplayBuffer("local-session"); - clearReplayBuffer("goose-session"); + clearReplayBuffer("acp-session"); useChatStore.setState({ messagesBySession: {}, sessionStateById: {}, @@ -49,16 +47,11 @@ describe("acpNotificationHandler", () => { }); 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", - ); - setActiveMessageId("goose-session", "assistant-1"); + registerPreparedSession("acp-session", "goose", "/Users/aharvard"); + setActiveMessageId("acp-session", "assistant-1"); await handleSessionNotification({ - sessionId: "goose-session", + sessionId: "acp-session", update: { sessionUpdate: "tool_call", toolCallId: "tool-1", @@ -67,7 +60,7 @@ describe("acpNotificationHandler", () => { } as never); await handleSessionNotification({ - sessionId: "goose-session", + sessionId: "acp-session", update: { sessionUpdate: "tool_call_update", toolCallId: "tool-1", @@ -94,7 +87,7 @@ describe("acpNotificationHandler", () => { } as never); await handleSessionNotification({ - sessionId: "goose-session", + sessionId: "acp-session", update: { sessionUpdate: "agent_message_chunk", content: { @@ -106,14 +99,13 @@ describe("acpNotificationHandler", () => { await waitFor(() => { const message = - useChatStore.getState().messagesBySession["local-session"]?.[0]; + useChatStore.getState().messagesBySession["acp-session"]?.[0]; expect(message?.content.some((block) => block.type === "mcpApp")).toBe( true, ); }); - const [message] = - useChatStore.getState().messagesBySession["local-session"]; + const [message] = useChatStore.getState().messagesBySession["acp-session"]; expect(message.id).toBe("assistant-1"); expect(message.content.map((block) => block.type)).toEqual([ "toolRequest", @@ -146,17 +138,17 @@ describe("acpNotificationHandler", () => { text: "The Host Info inspector is now open.", }); expect( - useChatStore.getState().getSessionRuntime("local-session") + useChatStore.getState().getSessionRuntime("acp-session") .streamingMessageId, ).toBe("assistant-1"); }); it("preserves ACP tool kind and locations on tool requests", async () => { - registerSession("local-session", "goose-session", "goose", "/Users/test"); - setActiveMessageId("goose-session", "assistant-1"); + registerPreparedSession("acp-session", "goose", "/Users/test"); + setActiveMessageId("acp-session", "assistant-1"); await handleSessionNotification({ - sessionId: "goose-session", + sessionId: "acp-session", update: { sessionUpdate: "tool_call", toolCallId: "tool-1", @@ -168,7 +160,7 @@ describe("acpNotificationHandler", () => { } as never); await handleSessionNotification({ - sessionId: "goose-session", + sessionId: "acp-session", update: { sessionUpdate: "tool_call_update", toolCallId: "tool-1", @@ -177,8 +169,7 @@ describe("acpNotificationHandler", () => { }, } as never); - const [message] = - useChatStore.getState().messagesBySession["local-session"]; + const [message] = useChatStore.getState().messagesBySession["acp-session"]; expect(message.content[0]).toMatchObject({ type: "toolRequest", id: "tool-1", @@ -190,16 +181,15 @@ describe("acpNotificationHandler", () => { }); it("preserves structured tool output when ACP provides rawOutput", async () => { - registerSession( - "local-session", - "goose-session", + registerPreparedSession( + "acp-session", "goose", "/Users/aharvard/.goose/artifacts", ); - setActiveMessageId("goose-session", "assistant-1"); + setActiveMessageId("acp-session", "assistant-1"); await handleSessionNotification({ - sessionId: "goose-session", + sessionId: "acp-session", update: { sessionUpdate: "tool_call", toolCallId: "tool-1", @@ -208,7 +198,7 @@ describe("acpNotificationHandler", () => { } as never); await handleSessionNotification({ - sessionId: "goose-session", + sessionId: "acp-session", update: { sessionUpdate: "tool_call_update", toolCallId: "tool-1", @@ -229,8 +219,7 @@ describe("acpNotificationHandler", () => { }, } as never); - const [message] = - useChatStore.getState().messagesBySession["local-session"]; + const [message] = useChatStore.getState().messagesBySession["acp-session"]; expect(message.content[1]).toMatchObject({ type: "toolResponse", id: "tool-1", @@ -244,7 +233,7 @@ describe("acpNotificationHandler", () => { }); it("replay keeps tool and MCP app content on an assistant message when tool events arrive before text", async () => { - const replaySessionId = "replay-goose-session"; + const replaySessionId = "replay-acp-session"; useChatStore.setState({ loadingSessionIds: new Set([replaySessionId]), }); @@ -339,7 +328,6 @@ describe("acpNotificationHandler", () => { payload: { ...createMcpAppPayload(), sessionId: replaySessionId, - gooseSessionId: replaySessionId, }, }); }); @@ -442,8 +430,8 @@ describe("acpNotificationHandler", () => { }); }); - it("replay preserves gooseSessionId in MCP app payloads before tracker registration", async () => { - const replaySessionId = "replay-goose-session-2"; + it("replay attaches MCP app payloads to tool-only assistant messages", async () => { + const replaySessionId = "replay-acp-session-2"; const replayCreated = 1_700_000_240; useChatStore.setState({ loadingSessionIds: new Set([replaySessionId]), @@ -496,7 +484,7 @@ describe("acpNotificationHandler", () => { expect(mcpAppBlock).toMatchObject({ type: "mcpApp", payload: expect.objectContaining({ - gooseSessionId: replaySessionId, + sessionId: replaySessionId, }), }); }); diff --git a/ui/goose2/src/shared/api/__tests__/acpSessionInfoUpdate.test.ts b/ui/goose2/src/shared/api/__tests__/acpSessionInfoUpdate.test.ts index 51d899d122..7f219a6b3e 100644 --- a/ui/goose2/src/shared/api/__tests__/acpSessionInfoUpdate.test.ts +++ b/ui/goose2/src/shared/api/__tests__/acpSessionInfoUpdate.test.ts @@ -28,7 +28,6 @@ describe("ACP session info updates", () => { it("applies generated session info updates to non-user-named sessions", async () => { useChatSessionStore.getState().addSession({ id: "goose-session-title", - acpSessionId: "goose-session-title", title: "New Chat", createdAt: "2026-01-01T00:00:00.000Z", updatedAt: "2026-01-01T00:00:00.000Z", @@ -62,7 +61,6 @@ describe("ACP session info updates", () => { it("ignores generated titles for user-named sessions", async () => { useChatSessionStore.getState().addSession({ id: "goose-session-user-title", - acpSessionId: "goose-session-user-title", title: "My Custom Title", createdAt: "2026-01-01T00:00:00.000Z", updatedAt: "2026-01-01T00:00:00.000Z", diff --git a/ui/goose2/src/shared/api/__tests__/acpToolCallStatus.test.ts b/ui/goose2/src/shared/api/__tests__/acpToolCallStatus.test.ts index dcbc2ba627..d2cb678423 100644 --- a/ui/goose2/src/shared/api/__tests__/acpToolCallStatus.test.ts +++ b/ui/goose2/src/shared/api/__tests__/acpToolCallStatus.test.ts @@ -9,13 +9,13 @@ import { handleSessionNotification, setActiveMessageId, } from "../acpNotificationHandler"; -import { registerSession } from "../acpSessionTracker"; +import { registerPreparedSession } from "../acpSessionRegistry"; describe("ACP tool call status handling", () => { beforeEach(() => { clearMessageTracking(); clearReplayBuffer("replay-failed-tool-session"); - clearReplayBuffer("goose-session"); + clearReplayBuffer("acp-session"); useChatStore.setState({ messagesBySession: {}, sessionStateById: {}, @@ -76,16 +76,15 @@ describe("ACP tool call status handling", () => { }); it("marks failed live tool updates as errors", async () => { - registerSession( - "local-session", - "goose-session", + registerPreparedSession( + "acp-session", "goose", "/Users/aharvard/.goose/artifacts", ); - setActiveMessageId("goose-session", "assistant-1"); + setActiveMessageId("acp-session", "assistant-1"); await handleSessionNotification({ - sessionId: "goose-session", + sessionId: "acp-session", update: { sessionUpdate: "tool_call", toolCallId: "tool-1", @@ -94,7 +93,7 @@ describe("ACP tool call status handling", () => { } as never); await handleSessionNotification({ - sessionId: "goose-session", + sessionId: "acp-session", update: { sessionUpdate: "tool_call_update", toolCallId: "tool-1", @@ -111,8 +110,7 @@ describe("ACP tool call status handling", () => { }, } as never); - const [message] = - useChatStore.getState().messagesBySession["local-session"]; + const [message] = useChatStore.getState().messagesBySession["acp-session"]; expect(message.content[0]).toMatchObject({ type: "toolRequest", id: "tool-1", diff --git a/ui/goose2/src/shared/api/acp.ts b/ui/goose2/src/shared/api/acp.ts index 9a07f7d723..c7940eaf4f 100644 --- a/ui/goose2/src/shared/api/acp.ts +++ b/ui/goose2/src/shared/api/acp.ts @@ -1,7 +1,7 @@ import type { ContentBlock } from "@agentclientprotocol/sdk"; import * as directAcp from "./acpApi"; import type { AcpSessionInfo } from "./acpApi"; -import * as sessionTracker from "./acpSessionTracker"; +import * as sessionRegistry from "./acpSessionRegistry"; import { getCatalogEntry, resolveAgentProviderCatalogId, @@ -27,12 +27,9 @@ export interface AcpSendMessageOptions { images?: [string, string][]; } -export interface AcpPrepareSessionOptions { +export interface AcpCreateSessionOptions { personaId?: string; projectId?: string; -} - -export interface AcpCreateSessionOptions extends AcpPrepareSessionOptions { modelId?: string | null; } @@ -85,8 +82,7 @@ export async function acpSendMessage( const sid = sessionId.slice(0, 8); const tStart = performance.now(); - const gooseSessionId = sessionTracker.getGooseSessionId(sessionId, personaId); - if (!gooseSessionId) { + if (!sessionRegistry.isSessionPrepared(sessionId)) { throw new Error("Session not prepared. Call acpPrepareSession first."); } @@ -113,7 +109,7 @@ export async function acpSendMessage( } const messageId = crypto.randomUUID(); - setActiveMessageId(gooseSessionId, messageId); + setActiveMessageId(sessionId, messageId); perfLog( `[perf:send] ${sid} acpSendMessage → prompt(len=${prompt.length}, imgs=${images?.length ?? 0})`, @@ -123,7 +119,7 @@ export async function acpSendMessage( if (personaId) meta.personaId = personaId; try { await directAcp.prompt( - gooseSessionId, + sessionId, content, Object.keys(meta).length > 0 ? meta : undefined, ); @@ -132,7 +128,7 @@ export async function acpSendMessage( `[perf:send] ${sid} prompt() resolved in ${(tDone - tPrompt).toFixed(1)}ms (total acpSendMessage ${(tDone - tStart).toFixed(1)}ms)`, ); } finally { - clearActiveMessageId(gooseSessionId); + clearActiveMessageId(sessionId); } } @@ -141,24 +137,16 @@ export async function acpPrepareSession( sessionId: string, providerId: string, workingDir: string, - options: AcpPrepareSessionOptions = {}, -): Promise { +): Promise { const sid = sessionId.slice(0, 8); const t0 = performance.now(); perfLog( `[perf:prepare] ${sid} acpPrepareSession start (provider=${providerId})`, ); - const gooseSessionId = await sessionTracker.prepareSession( - sessionId, - providerId, - workingDir, - options.personaId, - options.projectId, - ); + await sessionRegistry.prepareSession(sessionId, providerId, workingDir); perfLog( `[perf:prepare] ${sid} acpPrepareSession done in ${(performance.now() - t0).toFixed(1)}ms`, ); - return gooseSessionId; } export async function acpCreateSession( @@ -166,31 +154,26 @@ export async function acpCreateSession( workingDir: string, options: AcpCreateSessionOptions = {}, ): Promise<{ sessionId: string }> { - const localSessionId = crypto.randomUUID(); - const gooseSessionId = await acpPrepareSession( - localSessionId, - providerId, + const response = await directAcp.newSession( workingDir, - options, - ); - sessionTracker.registerSession( - gooseSessionId, - gooseSessionId, providerId, - workingDir, + options.projectId, + options.personaId, ); + const sessionId = response.sessionId; + await directAcp.setProvider(sessionId, providerId); + sessionRegistry.registerPreparedSession(sessionId, providerId, workingDir); if (options.modelId) { - await directAcp.setModel(gooseSessionId, options.modelId); + await directAcp.setModel(sessionId, options.modelId); } - return { sessionId: gooseSessionId }; + return { sessionId }; } export async function acpSetModel( sessionId: string, modelId: string, ): Promise { - const gooseSessionId = sessionTracker.getGooseSessionId(sessionId); - return directAcp.setModel(gooseSessionId ?? sessionId, modelId); + return directAcp.setModel(sessionId, modelId); } export type { AcpSessionInfo }; @@ -223,21 +206,19 @@ export async function acpSearchSessions( */ export async function acpLoadSession( sessionId: string, - gooseSessionId: string, workingDir?: string, ): Promise { const effectiveWorkingDir = workingDir ?? "~"; const sid = sessionId.slice(0, 8); const t0 = performance.now(); - const rollbackSessionRegistration = sessionTracker.registerSession( + const rollbackSessionRegistration = sessionRegistry.registerPreparedSession( sessionId, - gooseSessionId, "goose", effectiveWorkingDir, ); try { perfLog(`[perf:load] ${sid} acpLoadSession → client.loadSession`); - await directAcp.loadSession(gooseSessionId, effectiveWorkingDir); + await directAcp.loadSession(sessionId, effectiveWorkingDir); perfLog( `[perf:load] ${sid} client.loadSession resolved in ${(performance.now() - t0).toFixed(1)}ms`, ); @@ -261,17 +242,11 @@ export async function acpImportSession(json: string): Promise { export async function acpDuplicateSession( sessionId: string, ): Promise { - const gooseSessionId = - sessionTracker.getGooseSessionId(sessionId) ?? sessionId; - return directAcp.forkSession(gooseSessionId); + return directAcp.forkSession(sessionId); } /** Cancel an in-progress ACP session so the backend stops streaming. */ -export async function acpCancelSession( - sessionId: string, - personaId?: string, -): Promise { - const gooseSessionId = sessionTracker.getGooseSessionId(sessionId, personaId); - await directAcp.cancelSession(gooseSessionId ?? sessionId); +export async function acpCancelSession(sessionId: string): Promise { + await directAcp.cancelSession(sessionId); return true; } diff --git a/ui/goose2/src/shared/api/acpNotificationHandler.test.ts b/ui/goose2/src/shared/api/acpNotificationHandler.test.ts index ee6442d3bc..f2e3c8cd45 100644 --- a/ui/goose2/src/shared/api/acpNotificationHandler.test.ts +++ b/ui/goose2/src/shared/api/acpNotificationHandler.test.ts @@ -1,11 +1,7 @@ import { beforeEach, describe, expect, it } from "vitest"; import type { SessionNotification } from "@agentclientprotocol/sdk"; import { useChatStore } from "@/features/chat/stores/chatStore"; -import { - clearReplayBuffer, - getAndDeleteReplayBuffer, -} from "@/features/chat/hooks/replayBuffer"; -import { registerSession } from "./acpSessionTracker"; +import { clearReplayBuffer } from "@/features/chat/hooks/replayBuffer"; import { clearMessageTracking, handleSessionNotification, @@ -14,8 +10,8 @@ import { describe("acpNotificationHandler", () => { beforeEach(() => { clearMessageTracking(); - clearReplayBuffer("draft-session-1"); - clearReplayBuffer("draft-session-2"); + clearReplayBuffer("acp-session-1"); + clearReplayBuffer("acp-session-2"); useChatStore.setState({ messagesBySession: {}, sessionStateById: {}, @@ -28,9 +24,9 @@ describe("acpNotificationHandler", () => { }); }); - it("buffers usage updates until the local session mapping is registered", async () => { + it("applies usage updates to the ACP session id", async () => { const notification = { - sessionId: "goose-session-1", + sessionId: "acp-session-1", update: { sessionUpdate: "usage_update", used: 512, @@ -40,42 +36,9 @@ describe("acpNotificationHandler", () => { await handleSessionNotification(notification); - expect( - useChatStore.getState().sessionStateById["draft-session-1"], - ).toBeUndefined(); - expect( - useChatStore.getState().sessionStateById["goose-session-1"], - ).toBeUndefined(); - - registerSession("draft-session-1", "goose-session-1", "goose", "/tmp"); - - const runtime = useChatStore - .getState() - .getSessionRuntime("draft-session-1"); + const runtime = useChatStore.getState().getSessionRuntime("acp-session-1"); expect(runtime.tokenState.accumulatedTotal).toBe(512); expect(runtime.tokenState.contextLimit).toBe(8192); expect(runtime.hasUsageSnapshot).toBe(true); }); - - it("does not buffer non-usage updates before the local session mapping exists", async () => { - const notification = { - sessionId: "goose-session-2", - update: { - sessionUpdate: "agent_message_chunk", - messageId: "message-1", - content: { - type: "text", - text: "hello from replay", - }, - }, - } as SessionNotification; - - await handleSessionNotification(notification); - registerSession("draft-session-2", "goose-session-2", "goose", "/tmp"); - - expect(getAndDeleteReplayBuffer("draft-session-2")).toBeUndefined(); - expect( - useChatStore.getState().messagesBySession["draft-session-2"], - ).toBeUndefined(); - }); }); diff --git a/ui/goose2/src/shared/api/acpNotificationHandler.ts b/ui/goose2/src/shared/api/acpNotificationHandler.ts index 749b5d3f56..42c350b256 100644 --- a/ui/goose2/src/shared/api/acpNotificationHandler.ts +++ b/ui/goose2/src/shared/api/acpNotificationHandler.ts @@ -31,14 +31,10 @@ import { } from "./acpReplayAssistant"; import { getReplayCreated, getReplayMessageId } from "./acpReplayMetadata"; import { handleSessionInfoUpdate } from "./acpSessionInfoUpdate"; -import { - getLocalSessionId, - subscribeToSessionRegistration, -} from "./acpSessionTracker"; import { getToolCallIdentity } from "./acpToolCallIdentity"; import { perfLog } from "@/shared/lib/perfLog"; -// Pre-set message ID for the next live stream per goose session +// Pre-set message ID for the next live stream per session. const presetMessageIds = new Map(); // Per-session perf counters for replay/live streaming. @@ -54,10 +50,6 @@ interface LivePerf { chunkCount: number; } const livePerf = new Map(); -const pendingUsageUpdates = new Map< - string, - { accumulatedTotal: number; contextLimit: number } ->(); const toolCallStatusFromUpdate = (status: string): ToolCallStatus => status === "failed" ? "error" : "completed"; @@ -108,33 +100,20 @@ function toolCallUpdatePatch( }; } -subscribeToSessionRegistration((localSessionId, gooseSessionId) => { - const pendingUsage = pendingUsageUpdates.get(gooseSessionId); - if (!pendingUsage) { - return; - } - - useChatStore.getState().updateTokenState(localSessionId, pendingUsage); - pendingUsageUpdates.delete(gooseSessionId); -}); - -export function setActiveMessageId( - gooseSessionId: string, - messageId: string, -): void { - presetMessageIds.set(gooseSessionId, messageId); - livePerf.set(gooseSessionId, { +export function setActiveMessageId(sessionId: string, messageId: string): void { + presetMessageIds.set(sessionId, messageId); + livePerf.set(sessionId, { sendStartedAt: performance.now(), firstChunkAt: null, chunkCount: 0, }); } -export function clearActiveMessageId(gooseSessionId: string): void { - presetMessageIds.delete(gooseSessionId); - const perf = livePerf.get(gooseSessionId); +export function clearActiveMessageId(sessionId: string): void { + presetMessageIds.delete(sessionId); + const perf = livePerf.get(sessionId); if (perf) { - const sid = gooseSessionId.slice(0, 8); + const sid = sessionId.slice(0, 8); const total = performance.now() - perf.sendStartedAt; const ttft = perf.firstChunkAt !== null @@ -143,16 +122,14 @@ export function clearActiveMessageId(gooseSessionId: string): void { perfLog( `[perf:stream] ${sid} stream ended — ttft=${ttft}ms total=${total.toFixed(1)}ms chunks=${perf.chunkCount}`, ); - livePerf.delete(gooseSessionId); + livePerf.delete(sessionId); } } export async function handleSessionNotification( notification: SessionNotification, ): Promise { - const gooseSessionId = notification.sessionId; - const localSessionId = getLocalSessionId(gooseSessionId); - const sessionId = localSessionId ?? gooseSessionId; + const sessionId = notification.sessionId; const { update } = notification; const isReplay = useChatStore.getState().loadingSessionIds.has(sessionId); @@ -167,20 +144,20 @@ export async function handleSessionNotification( } perf.lastAt = now; perf.count += 1; - handleReplay(sessionId, gooseSessionId, localSessionId, update); + handleReplay(sessionId, update); } else { - const perf = livePerf.get(gooseSessionId); + const perf = livePerf.get(sessionId); if (perf && update.sessionUpdate === "agent_message_chunk") { perf.chunkCount += 1; if (perf.firstChunkAt === null) { perf.firstChunkAt = performance.now(); - const sid = gooseSessionId.slice(0, 8); + const sid = sessionId.slice(0, 8); perfLog( `[perf:stream] ${sid} first agent_message_chunk at ttft=${(perf.firstChunkAt - perf.sendStartedAt).toFixed(1)}ms`, ); } } - handleLive(sessionId, gooseSessionId, localSessionId, update); + handleLive(sessionId, update); } } @@ -202,12 +179,7 @@ function getChunkMessageId(update: SessionUpdate): string | null { : null; } -function handleReplay( - sessionId: string, - gooseSessionId: string, - localSessionId: string | null, - update: SessionUpdate, -): void { +function handleReplay(sessionId: string, update: SessionUpdate): void { switch (update.sessionUpdate) { case "agent_message_chunk": { const msg = ensureReplayAssistantMessage( @@ -331,7 +303,6 @@ function handleReplay( update, true, { - gooseSessionId, replayMessageId, }, ); @@ -344,7 +315,7 @@ function handleReplay( case "session_info_update": case "config_option_update": case "usage_update": - handleShared(sessionId, gooseSessionId, localSessionId, update); + handleShared(sessionId, update); break; default: @@ -352,19 +323,13 @@ function handleReplay( } } -function handleLive( - sessionId: string, - gooseSessionId: string, - localSessionId: string | null, - update: SessionUpdate, -): void { +function handleLive(sessionId: string, update: SessionUpdate): void { const store = useChatStore.getState(); switch (update.sessionUpdate) { case "agent_message_chunk": { const messageId = ensureLiveAssistantMessage( sessionId, - gooseSessionId, getChunkMessageId(update) ?? undefined, ); @@ -376,7 +341,7 @@ function handleLive( } case "tool_call": { - const messageId = ensureLiveAssistantMessage(sessionId, gooseSessionId); + const messageId = ensureLiveAssistantMessage(sessionId); const identity = getToolCallIdentity(update); const toolRequest: ToolRequestContent = { @@ -395,7 +360,7 @@ function handleLive( } case "tool_call_update": { - const messageId = ensureLiveAssistantMessage(sessionId, gooseSessionId); + const messageId = ensureLiveAssistantMessage(sessionId); const identity = getToolCallIdentity(update); const patch = toolCallUpdatePatch(update); @@ -460,9 +425,6 @@ function handleLive( toolRequest?.name ?? update.title ?? "", update, false, - { - gooseSessionId, - }, ); } } @@ -472,7 +434,7 @@ function handleLive( case "session_info_update": case "config_option_update": case "usage_update": - handleShared(sessionId, gooseSessionId, localSessionId, update); + handleShared(sessionId, update); break; default: @@ -480,12 +442,7 @@ function handleLive( } } -function handleShared( - sessionId: string, - gooseSessionId: string, - localSessionId: string | null, - update: SessionUpdate, -): void { +function handleShared(sessionId: string, update: SessionUpdate): void { switch (update.sessionUpdate) { case "session_info_update": { handleSessionInfoUpdate(sessionId, update); @@ -535,15 +492,7 @@ function handleShared( case "usage_update": { const usage = update as SessionUpdate & { sessionUpdate: "usage_update" }; - if (!localSessionId) { - pendingUsageUpdates.set(gooseSessionId, { - accumulatedTotal: usage.used, - contextLimit: usage.size, - }); - break; - } - - useChatStore.getState().updateTokenState(localSessionId, { + useChatStore.getState().updateTokenState(sessionId, { accumulatedTotal: usage.used, contextLimit: usage.size, }); @@ -562,7 +511,6 @@ function findStreamingMessageId(sessionId: string): string | null { function ensureLiveAssistantMessage( sessionId: string, - gooseSessionId: string, preferredMessageId?: string | null, ): string { const store = useChatStore.getState(); @@ -578,7 +526,7 @@ function ensureLiveAssistantMessage( const messageId = preferredMessageId ?? - presetMessageIds.get(gooseSessionId) ?? + presetMessageIds.get(sessionId) ?? existingStreamingMessageId ?? crypto.randomUUID(); @@ -598,14 +546,13 @@ function ensureLiveAssistantMessage( store.setPendingAssistantProvider(sessionId, null); store.setStreamingMessageId(sessionId, messageId); - clearActiveMessageId(gooseSessionId); + clearActiveMessageId(sessionId); return messageId; } export function clearMessageTracking(): void { presetMessageIds.clear(); - pendingUsageUpdates.clear(); clearReplayAssistantTracking(); } diff --git a/ui/goose2/src/shared/api/acpSessionRegistry.ts b/ui/goose2/src/shared/api/acpSessionRegistry.ts new file mode 100644 index 0000000000..7c7de01e2b --- /dev/null +++ b/ui/goose2/src/shared/api/acpSessionRegistry.ts @@ -0,0 +1,80 @@ +import * as acpApi from "./acpApi"; +import { perfLog } from "@/shared/lib/perfLog"; + +interface PreparedSession { + providerId: string; + workingDir: string; +} + +const prepared = new Map(); + +export async function prepareSession( + sessionId: string, + providerId: string, + workingDir: string, +): Promise { + const sid = sessionId.slice(0, 8); + + const existing = prepared.get(sessionId); + if (existing) { + const tReuse = performance.now(); + let changed = false; + if (existing.workingDir !== workingDir) { + await acpApi.updateWorkingDir(sessionId, workingDir); + existing.workingDir = workingDir; + changed = true; + } + if (existing.providerId !== providerId) { + const tProv = performance.now(); + await acpApi.setProvider(sessionId, providerId); + perfLog( + `[perf:prepare] ${sid} reuse setProvider(${providerId}) in ${(performance.now() - tProv).toFixed(1)}ms`, + ); + existing.providerId = providerId; + changed = true; + } + perfLog( + `[perf:prepare] ${sid} reuse existing session (updates=${changed}) in ${(performance.now() - tReuse).toFixed(1)}ms`, + ); + return; + } + + const tLoad = performance.now(); + await acpApi.loadSession(sessionId, workingDir); + perfLog( + `[perf:prepare] ${sid} registry loadSession ok in ${(performance.now() - tLoad).toFixed(1)}ms`, + ); + + const tProv = performance.now(); + await acpApi.setProvider(sessionId, providerId); + perfLog( + `[perf:prepare] ${sid} registry setProvider(${providerId}) in ${(performance.now() - tProv).toFixed(1)}ms`, + ); + + const entry = { providerId, workingDir }; + prepared.set(sessionId, entry); + + return; +} + +export function isSessionPrepared(sessionId: string): boolean { + return prepared.has(sessionId); +} + +export function registerPreparedSession( + sessionId: string, + providerId: string, + workingDir: string, +): () => void { + const previousEntry = prepared.get(sessionId); + const entry = { providerId, workingDir }; + + prepared.set(sessionId, entry); + + return () => { + prepared.delete(sessionId); + if (previousEntry) { + prepared.set(sessionId, previousEntry); + } + }; +} diff --git a/ui/goose2/src/shared/api/acpSessionTracker.ts b/ui/goose2/src/shared/api/acpSessionTracker.ts deleted file mode 100644 index dc0c3d493c..0000000000 --- a/ui/goose2/src/shared/api/acpSessionTracker.ts +++ /dev/null @@ -1,206 +0,0 @@ -import * as acpApi from "./acpApi"; -import { perfLog } from "@/shared/lib/perfLog"; - -interface PreparedSession { - gooseSessionId: string; - providerId: string; - workingDir: string; -} - -type SessionRegistrationListener = ( - localSessionId: string, - gooseSessionId: string, -) => void; - -const prepared = new Map(); -const gooseToLocal = new Map(); -const registrationListeners = new Set(); - -function restoreGooseRegistration( - gooseSessionId: string, - localSessionId: string | undefined, -): void { - if (localSessionId === undefined) { - gooseToLocal.delete(gooseSessionId); - return; - } - - gooseToLocal.set(gooseSessionId, localSessionId); -} - -function makeKey(sessionId: string, personaId?: string): string { - if (personaId && personaId.length > 0) { - return `${sessionId}__${personaId}`; - } - return sessionId; -} - -function notifySessionRegistered( - localSessionId: string, - gooseSessionId: string, -): void { - for (const listener of registrationListeners) { - listener(localSessionId, gooseSessionId); - } -} - -export function subscribeToSessionRegistration( - listener: SessionRegistrationListener, -): () => void { - registrationListeners.add(listener); - return () => registrationListeners.delete(listener); -} - -export async function prepareSession( - sessionId: string, - providerId: string, - workingDir: string, - personaId?: string, - projectId?: string, -): Promise { - const sid = sessionId.slice(0, 8); - const key = makeKey(sessionId, personaId); - - const existing = prepared.get(key) ?? prepared.get(sessionId); - if (existing) { - const tReuse = performance.now(); - let changed = false; - if (existing.workingDir !== workingDir) { - await acpApi.updateWorkingDir(existing.gooseSessionId, workingDir); - existing.workingDir = workingDir; - changed = true; - } - if (existing.providerId !== providerId) { - const tProv = performance.now(); - await acpApi.setProvider(existing.gooseSessionId, providerId); - perfLog( - `[perf:prepare] ${sid} reuse setProvider(${providerId}) in ${(performance.now() - tProv).toFixed(1)}ms (goose_sid=${existing.gooseSessionId.slice(0, 8)})`, - ); - existing.providerId = providerId; - changed = true; - } - perfLog( - `[perf:prepare] ${sid} reuse existing session (updates=${changed}) in ${(performance.now() - tReuse).toFixed(1)}ms`, - ); - return existing.gooseSessionId; - } - - let gooseSessionId: string | null = null; - - const tLoad = performance.now(); - try { - await acpApi.loadSession(sessionId, workingDir); - gooseSessionId = sessionId; - perfLog( - `[perf:prepare] ${sid} tracker loadSession ok in ${(performance.now() - tLoad).toFixed(1)}ms`, - ); - } catch { - perfLog( - `[perf:prepare] ${sid} tracker loadSession failed in ${(performance.now() - tLoad).toFixed(1)}ms → newSession`, - ); - } - - if (!gooseSessionId) { - const tNew = performance.now(); - const response = await acpApi.newSession( - workingDir, - providerId, - projectId, - personaId, - ); - gooseSessionId = response.sessionId; - perfLog( - `[perf:prepare] ${sid} tracker newSession done in ${(performance.now() - tNew).toFixed(1)}ms (goose_sid=${gooseSessionId.slice(0, 8)})`, - ); - } - - const gooseSid = gooseSessionId.slice(0, 8); - const tProv = performance.now(); - await acpApi.setProvider(gooseSessionId, providerId); - perfLog( - `[perf:prepare] ${sid} tracker setProvider(${providerId}) in ${(performance.now() - tProv).toFixed(1)}ms (goose_sid=${gooseSid})`, - ); - - const entry = { gooseSessionId, providerId, workingDir }; - prepared.set(key, entry); - prepared.set(sessionId, entry); - prepared.set(gooseSessionId, entry); - gooseToLocal.set(gooseSessionId, sessionId); - notifySessionRegistered(sessionId, gooseSessionId); - - 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 previousEntry = prepared.get(sessionId); - const previousGooseSessionLocal = gooseToLocal.get(gooseSessionId); - const previousSessionGooseLocal = previousEntry - ? gooseToLocal.get(previousEntry.gooseSessionId) - : undefined; - const entry = { gooseSessionId, providerId, workingDir }; - - if ( - previousEntry && - previousEntry.gooseSessionId !== gooseSessionId && - gooseToLocal.get(previousEntry.gooseSessionId) === sessionId - ) { - gooseToLocal.delete(previousEntry.gooseSessionId); - } - - prepared.set(sessionId, entry); - prepared.set(gooseSessionId, entry); - gooseToLocal.set(gooseSessionId, sessionId); - notifySessionRegistered(sessionId, gooseSessionId); - - return () => { - prepared.delete(sessionId); - if (previousEntry) { - prepared.set(sessionId, previousEntry); - } - - restoreGooseRegistration(gooseSessionId, previousGooseSessionLocal); - if (previousEntry && previousEntry.gooseSessionId !== gooseSessionId) { - restoreGooseRegistration( - previousEntry.gooseSessionId, - previousSessionGooseLocal, - ); - } - }; -} - -export function unregisterSession( - sessionId: string, - gooseSessionId?: string, -): void { - const entry = prepared.get(sessionId); - prepared.delete(sessionId); - - const resolvedGooseSessionId = gooseSessionId ?? entry?.gooseSessionId; - if ( - resolvedGooseSessionId && - gooseToLocal.get(resolvedGooseSessionId) === sessionId - ) { - gooseToLocal.delete(resolvedGooseSessionId); - } -} diff --git a/ui/goose2/src/shared/api/acpToolCallContent.ts b/ui/goose2/src/shared/api/acpToolCallContent.ts index ffa8745aad..970ad356a5 100644 --- a/ui/goose2/src/shared/api/acpToolCallContent.ts +++ b/ui/goose2/src/shared/api/acpToolCallContent.ts @@ -66,7 +66,6 @@ export function attachMcpAppPayload( update: SessionUpdate, isReplay: boolean, options?: { - gooseSessionId?: string | null; replayMessageId?: string | null; }, ): void { @@ -75,7 +74,6 @@ export function attachMcpAppPayload( toolCallId, toolCallTitle, update, - options?.gooseSessionId, ); if (!payload) { return; diff --git a/ui/goose2/src/shared/api/mcpAppToolUpdate.ts b/ui/goose2/src/shared/api/mcpAppToolUpdate.ts index 125cc0a853..edd18a23ff 100644 --- a/ui/goose2/src/shared/api/mcpAppToolUpdate.ts +++ b/ui/goose2/src/shared/api/mcpAppToolUpdate.ts @@ -6,7 +6,6 @@ import type { } 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 { return typeof value === "object" && value !== null; @@ -29,7 +28,6 @@ export function buildMcpAppPayloadFromToolUpdate( toolCallId: string, toolCallTitle: string, update: SessionUpdate, - gooseSessionIdOverride?: string | null, ): McpAppPayload | null { const payload = extractMcpAppPayload(update); if (!payload) { @@ -38,8 +36,6 @@ export function buildMcpAppPayloadFromToolUpdate( return { sessionId, - gooseSessionId: - gooseSessionIdOverride ?? getGooseSessionId(sessionId) ?? null, toolCallId, toolCallTitle, source: "toolCallUpdateMeta", diff --git a/ui/goose2/src/shared/types/messages.ts b/ui/goose2/src/shared/types/messages.ts index 8d3ea67fbb..a859d1c6c9 100644 --- a/ui/goose2/src/shared/types/messages.ts +++ b/ui/goose2/src/shared/types/messages.ts @@ -118,7 +118,6 @@ export interface ToolResponseContent { export interface McpAppPayload { sessionId: string; - gooseSessionId: string | null; toolCallId: string; toolCallTitle: string; source: "toolCallUpdateMeta";