diff --git a/ui/goose2/src/app/AppShell.tsx b/ui/goose2/src/app/AppShell.tsx index 45374f30f2..b3f7160550 100644 --- a/ui/goose2/src/app/AppShell.tsx +++ b/ui/goose2/src/app/AppShell.tsx @@ -16,6 +16,10 @@ import { useAgentStore } from "@/features/agents/stores/agentStore"; import { useProjectStore } from "@/features/projects/stores/projectStore"; import { findExistingDraft } from "@/features/chat/lib/newChat"; import { DEFAULT_CHAT_TITLE } from "@/features/chat/lib/sessionTitle"; +import { + resolveSessionModelPreference, + sanitizeSessionModelPreference, +} from "@/features/chat/lib/sessionModelPreference"; import { useAppStartup } from "./hooks/useAppStartup"; import { useHomeSessionStateSync } from "./hooks/useHomeSessionStateSync"; import { loadStoredHomeSessionId } from "./lib/homeSessionStorage"; @@ -147,13 +151,14 @@ 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, workingDir); + await acpLoadSession(sessionId, gooseSessionId, workingDir); const tFlush = performance.now(); useChatStore.getState().setSessionLoading(sessionId, false); const buffer = getAndDeleteReplayBuffer(sessionId); @@ -236,6 +241,9 @@ export function AppShell({ children }: { children?: React.ReactNode }) { homeSession.id, sessionModelPreference.providerId, workingDir, + { + personaId: homeSession.personaId, + }, ); const shouldClearHomeModel = sessionModelPreference.providerId !== homeSession.providerId || @@ -307,12 +315,6 @@ export function AppShell({ children }: { children?: React.ReactNode }) { ); const providerId = project?.preferredProvider ?? agentStore.selectedProvider ?? "goose"; - const sessionModelPreference = - await resolveSupportedSessionModelPreference( - providerId, - providerInventoryEntries, - project?.preferredModel ?? undefined, - ); const sessionState = useChatSessionStore.getState(); const chatState = useChatStore.getState(); const existingDraft = findExistingDraft({ @@ -336,10 +338,39 @@ export function AppShell({ children }: { children?: React.ReactNode }) { return existingDraft; } - const workingDir = await resolveSessionCwd(project); + if (project) { + const modelPreference = resolveSessionModelPreference({ + providerId, + preferredModel: project.preferredModel ?? undefined, + }); + const sessionModelPreference = sanitizeSessionModelPreference( + modelPreference, + providerInventoryEntries.get(modelPreference.providerId), + ); + const session = sessionStore.createLocalSession({ + title, + projectId: project.id, + providerId: sessionModelPreference.providerId, + modelId: sessionModelPreference.modelId, + modelName: sessionModelPreference.modelName, + }); + sessionStore.setActiveSession(session.id); + setActiveView("chat"); + chatStore.setActiveSession(session.id); + perfLog( + `[perf:newtab] ${session.id.slice(0, 8)} created local project session in ${(performance.now() - tStart).toFixed(1)}ms`, + ); + return session; + } + + const sessionModelPreference = + await resolveSupportedSessionModelPreference( + providerId, + providerInventoryEntries, + ); + const workingDir = await resolveSessionCwd(null); const session = await sessionStore.createSession({ title, - projectId: project?.id, providerId: sessionModelPreference.providerId, workingDir, modelId: sessionModelPreference.modelId, @@ -500,11 +531,21 @@ export function AppShell({ children }: { children?: React.ReactNode }) { if (!workingDir) { return; } - await acpPrepareSession( + const gooseSessionId = await acpPrepareSession( sessionId, session.providerId ?? agentStore.selectedProvider ?? "goose", workingDir, + { + personaId: session.personaId, + ...(projectId ? { projectId } : {}), + ...(!session.acpSessionId ? { knownNew: true } : {}), + }, ); + if (gooseSessionId && session.acpSessionId !== gooseSessionId) { + sessionStore.updateSession(sessionId, { + acpSessionId: gooseSessionId, + }); + } })().catch((error) => { console.error( "Failed to update ACP session project working directory:", 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 5ad52eac54..6d007e04d0 100644 --- a/ui/goose2/src/features/chat/hooks/__tests__/useChatSessionController.test.ts +++ b/ui/goose2/src/features/chat/hooks/__tests__/useChatSessionController.test.ts @@ -95,6 +95,14 @@ vi.mock("../useAgentModelPickerState", () => ({ modelStatusMessage: mockPickerState.modelStatusMessage, handleProviderChange: vi.fn(), handleModelChange: (modelId: string) => { + if (modelId === "gpt-4.1") { + onModelSelected?.({ + id: modelId, + name: modelId, + displayName: "GPT-4.1", + providerId: "openai", + }); + } if (modelId === "claude-sonnet-4") { onModelSelected?.({ id: modelId, @@ -163,6 +171,7 @@ describe("useChatSessionController", () => { sessions: [ { id: "session-1", + acpSessionId: "session-1", title: "Chat", providerId: "openai", modelId: "gpt-4o", @@ -216,6 +225,85 @@ describe("useChatSessionController", () => { modelName: "Claude Sonnet 4", }); }); + + it("binds an optimistic project session before setting its model", async () => { + mockAcpPrepareSession.mockResolvedValueOnce("goose-session-1"); + useProjectStore.setState({ + projects: [ + { + id: "project-1", + name: "Project", + description: "", + prompt: "", + icon: "", + color: "#000000", + preferredProvider: "openai", + preferredModel: null, + workingDirs: ["/tmp/project"], + useWorktrees: false, + order: 0, + archivedAt: null, + createdAt: "2026-04-20T00:00:00.000Z", + updatedAt: "2026-04-20T00:00:00.000Z", + artifactsDir: "/tmp/project/artifacts", + }, + ], + loading: false, + activeProjectId: null, + }); + useChatSessionStore.setState({ + sessions: [ + { + id: "local-session-1", + title: "Chat", + projectId: "project-1", + providerId: "openai", + createdAt: "2026-04-20T00:00:00.000Z", + updatedAt: "2026-04-20T00:00:00.000Z", + messageCount: 0, + }, + ], + activeSessionId: null, + isLoading: false, + hasHydratedSessions: true, + contextPanelOpenBySession: {}, + activeWorkspaceBySession: {}, + }); + + const { result } = renderHook(() => + useChatSessionController({ sessionId: "local-session-1" }), + ); + + act(() => { + result.current.handleModelChange("gpt-4.1"); + }); + + await waitFor(() => { + expect(mockAcpPrepareSession).toHaveBeenCalledWith( + "local-session-1", + "openai", + "/tmp/project", + { + personaId: undefined, + projectId: "project-1", + knownNew: true, + }, + ); + }); + + await waitFor(() => { + expect( + useChatSessionStore.getState().getSession("local-session-1"), + ).toMatchObject({ + acpSessionId: "goose-session-1", + modelId: "gpt-4.1", + modelName: "GPT-4.1", + }); + }); + + expect(mockAcpSetModel).toHaveBeenCalledWith("local-session-1", "gpt-4.1"); + }); + it("restores the previous stored model preference when setting a model fails", async () => { window.localStorage.setItem( "goose:preferredModelsByAgent", @@ -324,6 +412,7 @@ describe("useChatSessionController", () => { sessions: [ { id: "session-2", + acpSessionId: "session-2", title: "Chat", providerId: "openai", createdAt: "2026-04-21T00:00:00.000Z", @@ -383,6 +472,7 @@ describe("useChatSessionController", () => { sessions: [ { id: "session-3", + acpSessionId: "session-3", title: "Chat", providerId: "openai", createdAt: "2026-04-21T00:00:00.000Z", diff --git a/ui/goose2/src/features/chat/hooks/prepareBoundSession.ts b/ui/goose2/src/features/chat/hooks/prepareBoundSession.ts new file mode 100644 index 0000000000..5b9a8e2547 --- /dev/null +++ b/ui/goose2/src/features/chat/hooks/prepareBoundSession.ts @@ -0,0 +1,59 @@ +import { acpPrepareSession, acpSetModel } from "@/shared/api/acp"; +import { useChatSessionStore } from "../stores/chatSessionStore"; +import type { PreferredModelSelection } from "./useResolvedAgentModelPicker"; + +interface PrepareBoundSessionArgs { + sessionId: string; + providerId: string; + workingDir: string; + personaId?: string; + projectId?: string | null; + modelSelection?: PreferredModelSelection | null; +} + +export async function prepareBoundSession({ + sessionId, + providerId, + workingDir, + personaId, + projectId, + modelSelection, +}: PrepareBoundSessionArgs): Promise { + const sessionStore = useChatSessionStore.getState(); + const sessionBeforePrepare = sessionStore.getSession(sessionId); + const gooseSessionId = await acpPrepareSession( + sessionId, + providerId, + workingDir, + { + personaId, + ...(projectId ? { projectId } : {}), + ...(!sessionBeforePrepare?.acpSessionId ? { knownNew: true } : {}), + }, + ); + if ( + gooseSessionId && + sessionStore.getSession(sessionId)?.acpSessionId !== gooseSessionId + ) { + sessionStore.updateSession(sessionId, { acpSessionId: gooseSessionId }); + } + if (!modelSelection?.id) { + return; + } + + const liveSession = sessionStore.getSession(sessionId); + const modelAlreadyApplied = + Boolean(sessionBeforePrepare?.acpSessionId) && + liveSession?.modelId === modelSelection.id && + liveSession?.modelName === modelSelection.name; + + if (modelAlreadyApplied) { + return; + } + + await acpSetModel(sessionId, modelSelection.id); + sessionStore.updateSession(sessionId, { + modelId: modelSelection.id, + modelName: modelSelection.name, + }); +} diff --git a/ui/goose2/src/features/chat/hooks/useChatSessionController.ts b/ui/goose2/src/features/chat/hooks/useChatSessionController.ts index f2acb12871..24d5357f2f 100644 --- a/ui/goose2/src/features/chat/hooks/useChatSessionController.ts +++ b/ui/goose2/src/features/chat/hooks/useChatSessionController.ts @@ -24,11 +24,11 @@ import { supportsContextCompactionControls, } from "../lib/autoCompact"; import { resolveSessionCwd } from "@/features/projects/lib/sessionCwdSelection"; -import { acpPrepareSession, acpSetModel } from "@/shared/api/acp"; import { useResolvedAgentModelPicker, type PreferredModelSelection, } from "./useResolvedAgentModelPicker"; +import { prepareBoundSession } from "./prepareBoundSession"; interface UseChatSessionControllerOptions { sessionId: string | null; @@ -159,7 +159,9 @@ export function useChatSessionController({ providerId: string, nextProject = project, nextWorkspacePath = activeWorkspace?.path, + personaId = selectedPersonaId ?? undefined, modelSelection?: PreferredModelSelection | null, + nextProjectId = nextProject?.id ?? session?.projectId ?? null, ) => { if (!sessionId) { return; @@ -168,28 +170,22 @@ export function useChatSessionController({ nextProject, nextWorkspacePath, ); - await acpPrepareSession(sessionId, providerId, workingDir); - if (!modelSelection?.id) { - return; - } - - const sessionStore = useChatSessionStore.getState(); - const liveSession = sessionStore.getSession(sessionId); - const modelAlreadyApplied = - liveSession?.modelId === modelSelection.id && - liveSession?.modelName === modelSelection.name; - - if (modelAlreadyApplied) { - return; - } - - await acpSetModel(sessionId, modelSelection.id); - sessionStore.updateSession(sessionId, { - modelId: modelSelection.id, - modelName: modelSelection.name, + await prepareBoundSession({ + sessionId, + providerId, + workingDir, + personaId, + projectId: nextProjectId, + modelSelection, }); }, - [activeWorkspace?.path, project, sessionId], + [ + activeWorkspace?.path, + project, + selectedPersonaId, + session?.projectId, + sessionId, + ], ); const prepareSelectedProvider = useCallback( (providerId: string, modelSelection?: PreferredModelSelection | null) => @@ -197,9 +193,11 @@ export function useChatSessionController({ providerId, project, activeWorkspace?.path, + selectedPersonaId ?? undefined, modelSelection, + project?.id ?? null, ), - [activeWorkspace?.path, prepareCurrentSession, project], + [activeWorkspace?.path, prepareCurrentSession, project, selectedPersonaId], ); const prevProjectIdRef = useRef(session?.projectId); @@ -323,7 +321,9 @@ export function useChatSessionController({ selectedProvider, nextProject, activeWorkspace?.path, + selectedPersonaId ?? undefined, effectiveModelSelection, + projectId, ).catch((error) => { console.error("Failed to update ACP session working directory:", error); }); @@ -332,6 +332,7 @@ export function useChatSessionController({ activeWorkspace?.path, effectiveModelSelection, prepareCurrentSession, + selectedPersonaId, selectedProvider, sessionId, ], @@ -421,11 +422,14 @@ export function useChatSessionController({ { onMessageAccepted: sessionId ? onMessageAccepted : undefined, ensurePrepared: selectedProvider - ? () => + ? (personaId?: string) => prepareCurrentSession( selectedProvider, project, activeWorkspace?.path, + personaId, + effectiveModelSelection, + project?.id ?? session?.projectId ?? null, ) : undefined, }, @@ -727,7 +731,9 @@ export function useChatSessionController({ nextProviderId, nextProject, activeWorkspace?.path, + nextPersonaId, pendingModelSelection, + nextProjectId ?? null, ); if (cancelled) { return; diff --git a/ui/goose2/src/features/chat/hooks/useResolvedAgentModelPicker.ts b/ui/goose2/src/features/chat/hooks/useResolvedAgentModelPicker.ts index 81509350a5..d31f02c58c 100644 --- a/ui/goose2/src/features/chat/hooks/useResolvedAgentModelPicker.ts +++ b/ui/goose2/src/features/chat/hooks/useResolvedAgentModelPicker.ts @@ -302,6 +302,7 @@ export function useResolvedAgentModelPicker({ const previousModelName = session.modelName; const providerChanged = Boolean(nextProviderId) && nextProviderId !== session.providerId; + const needsPreparation = providerChanged || !session.acpSessionId; if (providerChanged && nextProviderId) { useChatSessionStore @@ -317,7 +318,7 @@ export function useResolvedAgentModelPicker({ void (async () => { try { - if (providerChanged && nextProviderId) { + if (needsPreparation && nextProviderId) { await prepareSelectedProvider(nextProviderId); } await acpSetModel(sessionId, modelId); 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 735be14993..72506ed64c 100644 --- a/ui/goose2/src/features/chat/stores/__tests__/chatSessionStore.test.ts +++ b/ui/goose2/src/features/chat/stores/__tests__/chatSessionStore.test.ts @@ -4,10 +4,14 @@ import { useChatSessionStore, type ChatSession } from "../chatSessionStore"; const mockAcpCreateSession = vi.fn(); const mockAcpListSessions = vi.fn(); +const mockAcpPrepareSession = vi.fn(); +const mockAcpSetModel = vi.fn(); vi.mock("@/shared/api/acp", () => ({ acpCreateSession: (...args: unknown[]) => mockAcpCreateSession(...args), acpListSessions: (...args: unknown[]) => mockAcpListSessions(...args), + acpPrepareSession: (...args: unknown[]) => mockAcpPrepareSession(...args), + acpSetModel: (...args: unknown[]) => mockAcpSetModel(...args), })); vi.mock("@/shared/api/acpApi", () => ({ @@ -31,6 +35,7 @@ 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", @@ -90,6 +95,30 @@ describe("chatSessionStore", () => { }); }); + describe("createLocalSession", () => { + it("creates an unbacked local draft session without touching ACP", () => { + const session = useChatSessionStore.getState().createLocalSession({ + title: "New Chat", + projectId: "project-1", + providerId: "openai", + modelId: "gpt-4.1", + modelName: "GPT-4.1", + }); + + expect(mockAcpCreateSession).not.toHaveBeenCalled(); + expect(session).toMatchObject({ + title: "New Chat", + projectId: "project-1", + providerId: "openai", + modelId: "gpt-4.1", + modelName: "GPT-4.1", + messageCount: 0, + }); + expect(session.acpSessionId).toBeUndefined(); + expect(useChatSessionStore.getState().sessions[0]).toEqual(session); + }); + }); + describe("loadSessions", () => { it("loads sessions from ACP and maps them correctly", async () => { mockAcpListSessions.mockResolvedValue([ @@ -196,6 +225,36 @@ describe("chatSessionStore", () => { expect(state.activeSessionId).toBeNull(); }); + it("keeps unbacked local draft sessions during ACP refresh", async () => { + const localSession = useChatSessionStore.getState().createLocalSession({ + title: "Local Draft", + projectId: "project-1", + }); + useChatSessionStore.getState().setActiveSession(localSession.id); + + mockAcpListSessions.mockResolvedValue([ + { + sessionId: "acp-1", + title: "ACP Session", + updatedAt: "2026-04-02", + createdAt: "2026-04-02", + archivedAt: null, + userSetName: false, + messageCount: 1, + providerId: null, + modelId: null, + }, + ]); + + await useChatSessionStore.getState().loadSessions(); + + const state = useChatSessionStore.getState(); + expect(state.sessions.map((session) => session.id)).toContain( + localSession.id, + ); + expect(state.activeSessionId).toBe(localSession.id); + }); + it("sets isLoading during fetch", async () => { let resolvePromise: (value: AcpSessionInfo[]) => void = () => {}; mockAcpListSessions.mockReturnValue( diff --git a/ui/goose2/src/features/chat/stores/chatSessionStore.ts b/ui/goose2/src/features/chat/stores/chatSessionStore.ts index a421d6a0df..6af4b4566a 100644 --- a/ui/goose2/src/features/chat/stores/chatSessionStore.ts +++ b/ui/goose2/src/features/chat/stores/chatSessionStore.ts @@ -2,6 +2,8 @@ import { create } from "zustand"; import { acpCreateSession, acpListSessions, + acpPrepareSession, + acpSetModel, type AcpSessionInfo, } from "@/shared/api/acp"; import type { Session } from "@/shared/types/chat"; @@ -18,6 +20,7 @@ import { export interface ChatSession { id: string; + acpSessionId?: string; title: string; projectId?: string | null; providerId?: string; @@ -74,8 +77,26 @@ interface CreateSessionOpts { modelName?: string; } +export interface SessionBindingModel { + id: string; + name?: string; +} + +interface PrepareSessionBindingArgs { + sessionId: string; + providerId: string; + workingDir: string; + personaId?: string; + projectId?: string | null; + model?: SessionBindingModel | null; +} + interface ChatSessionStoreActions { createSession: (opts?: CreateSessionOpts) => Promise; + createLocalSession: ( + opts?: Omit, + ) => ChatSession; + prepareSessionBinding: (args: PrepareSessionBindingArgs) => Promise; loadSessions: () => Promise; updateSession: (id: string, patch: Partial) => void; addSession: (session: ChatSession) => void; @@ -99,6 +120,7 @@ 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, @@ -119,9 +141,45 @@ function sortByUpdatedAtDesc(sessions: ChatSession[]): ChatSession[] { ); } +function mergeAcpSessionsWithLocalState( + currentSessions: ChatSession[], + acpSessions: AcpSessionInfo[], +): ChatSession[] { + const localSessionByAcpId = new Map(); + for (const session of currentSessions) { + if (session.acpSessionId && session.id !== session.acpSessionId) { + localSessionByAcpId.set(session.acpSessionId, session); + } + } + + const localOnlySessions = currentSessions.filter( + (session) => !session.acpSessionId, + ); + const mergedAcpSessions = acpSessions.map((session) => { + const chatSession = acpSessionToChatSession(session); + const localSession = localSessionByAcpId.get(session.sessionId); + + if (!localSession) { + return chatSession; + } + + return { + ...chatSession, + id: localSession.id, + modelName: + localSession.modelId === chatSession.modelId + ? localSession.modelName + : undefined, + }; + }); + + return sortByUpdatedAtDesc([...localOnlySessions, ...mergedAcpSessions]); +} + export function sessionToChatSession(session: Session): ChatSession { return { id: session.id, + acpSessionId: session.id, title: session.title, projectId: session.projectId, providerId: session.providerId, @@ -154,10 +212,11 @@ export const useChatSessionStore = create((set, get) => ({ const { sessionId } = await acpCreateSession(providerId, opts.workingDir, { personaId: opts.personaId, modelId: opts.modelId, - projectId: opts.projectId, + ...(opts.projectId ? { projectId: opts.projectId } : {}), }); const chatSession: ChatSession = { id: sessionId, + acpSessionId: sessionId, title: opts.title ?? DEFAULT_CHAT_TITLE, projectId: opts.projectId, providerId, @@ -173,12 +232,77 @@ export const useChatSessionStore = create((set, get) => ({ return chatSession; }, + createLocalSession: (opts) => { + const now = new Date().toISOString(); + const chatSession: ChatSession = { + id: crypto.randomUUID(), + title: opts?.title ?? DEFAULT_CHAT_TITLE, + projectId: opts?.projectId, + providerId: opts?.providerId ?? "goose", + personaId: opts?.personaId, + modelId: opts?.modelId, + modelName: opts?.modelName, + createdAt: now, + updatedAt: now, + messageCount: 0, + }; + set((state) => ({ sessions: [chatSession, ...state.sessions] })); + return chatSession; + }, + + prepareSessionBinding: async ({ + sessionId, + providerId, + workingDir, + personaId, + projectId, + model, + }) => { + const sessionBeforePrepare = get().getSession(sessionId); + const gooseSessionId = await acpPrepareSession( + sessionId, + providerId, + workingDir, + { + personaId, + ...(projectId ? { projectId } : {}), + ...(!sessionBeforePrepare?.acpSessionId ? { knownNew: true } : {}), + }, + ); + if ( + gooseSessionId && + get().getSession(sessionId)?.acpSessionId !== gooseSessionId + ) { + get().updateSession(sessionId, { acpSessionId: gooseSessionId }); + } + if (!model?.id) { + return; + } + + const liveSession = get().getSession(sessionId); + const modelAlreadyApplied = + Boolean(sessionBeforePrepare?.acpSessionId) && + liveSession?.modelId === model.id && + liveSession?.modelName === model.name; + + if (modelAlreadyApplied) { + return; + } + + await acpSetModel(sessionId, model.id); + get().updateSession(sessionId, { + modelId: model.id, + modelName: model.name, + }); + }, + loadSessions: async () => { set({ isLoading: true }); try { const acpSessions = await acpListSessions(); - const sessions = sortByUpdatedAtDesc( - acpSessions.map(acpSessionToChatSession), + const sessions = mergeAcpSessionsWithLocalState( + get().sessions, + acpSessions, ); const activeSessionId = get().activeSessionId; const activeSessionStillExists = @@ -209,23 +333,24 @@ 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 && - updatedSession && + acpSessionId && patch.title ) { - acpRenameSession(updatedSession.id, patch.title).catch((err: unknown) => + acpRenameSession(acpSessionId, patch.title).catch((err: unknown) => console.error("Failed to rename session in backend:", err), ); } // Persist projectId change to backend - if ("projectId" in patch && updatedSession) { - updateSessionProject(updatedSession.id, patch.projectId ?? null).catch( + if ("projectId" in patch && acpSessionId) { + updateSessionProject(acpSessionId, patch.projectId ?? null).catch( (err: unknown) => console.error("Failed to update session project in backend:", err), ); @@ -233,16 +358,20 @@ 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 === session.id, + (candidate) => candidate.id === normalizedSession.id, ); if (existing >= 0) { const updated = [...state.sessions]; - updated[existing] = { ...updated[existing], ...session }; + updated[existing] = { ...updated[existing], ...normalizedSession }; return { sessions: updated }; } - return { sessions: [session, ...state.sessions] }; + return { sessions: [normalizedSession, ...state.sessions] }; }); }, @@ -257,8 +386,8 @@ export const useChatSessionStore = create((set, get) => ({ state.activeSessionId === id ? null : state.activeSessionId, })); const session = get().sessions.find((candidate) => candidate.id === id); - if (session) { - acpArchiveSession(session.id).catch((err: unknown) => + if (session?.acpSessionId) { + acpArchiveSession(session.acpSessionId).catch((err: unknown) => console.error("Failed to archive session in backend:", err), ); } @@ -271,8 +400,8 @@ export const useChatSessionStore = create((set, get) => ({ ), })); const session = get().sessions.find((candidate) => candidate.id === id); - if (session) { - acpUnarchiveSession(session.id).catch((err: unknown) => + if (session?.acpSessionId) { + acpUnarchiveSession(session.acpSessionId).catch((err: unknown) => console.error("Failed to unarchive session in backend:", err), ); } diff --git a/ui/goose2/src/shared/api/__tests__/acp.test.ts b/ui/goose2/src/shared/api/__tests__/acp.test.ts index e428afd930..9f2143a508 100644 --- a/ui/goose2/src/shared/api/__tests__/acp.test.ts +++ b/ui/goose2/src/shared/api/__tests__/acp.test.ts @@ -101,7 +101,7 @@ describe("acpPrepareSession", () => { await expect( acpPrepareSession("acp-session-1", "openai", "/tmp/project"), - ).resolves.toBeUndefined(); + ).resolves.toBe("acp-session-1"); expect(mockLoadSession).toHaveBeenCalledWith( "acp-session-1", @@ -124,4 +124,30 @@ describe("acpPrepareSession", () => { expect(mockNewSession).not.toHaveBeenCalled(); expect(mockSetProvider).not.toHaveBeenCalled(); }); + + it("creates a known-new session without probing load or resetting provider", async () => { + mockNewSession.mockResolvedValueOnce({ sessionId: "goose-session-1" }); + + const sessionRegistry = await import("../acpSessionRegistry"); + const { acpPrepareSession } = await import("../acp"); + + await expect( + acpPrepareSession("local-session", "openai", "/tmp/project", { + projectId: "project-1", + knownNew: true, + }), + ).resolves.toBe("goose-session-1"); + + expect(mockLoadSession).not.toHaveBeenCalled(); + expect(mockNewSession).toHaveBeenCalledWith( + "/tmp/project", + "openai", + "project-1", + undefined, + ); + expect(mockSetProvider).not.toHaveBeenCalled(); + expect(sessionRegistry.getGooseSessionId("local-session")).toBe( + "goose-session-1", + ); + }); }); diff --git a/ui/goose2/src/shared/api/acp.ts b/ui/goose2/src/shared/api/acp.ts index aeb3a8df5e..367da81a40 100644 --- a/ui/goose2/src/shared/api/acp.ts +++ b/ui/goose2/src/shared/api/acp.ts @@ -27,9 +27,13 @@ export interface AcpSendMessageOptions { images?: [string, string][]; } -export interface AcpCreateSessionOptions { +export interface AcpPrepareSessionOptions { personaId?: string; projectId?: string; + knownNew?: boolean; +} + +export interface AcpCreateSessionOptions extends AcpPrepareSessionOptions { modelId?: string | null; } @@ -83,7 +87,11 @@ export async function acpSendMessage( const sid = sessionId.slice(0, 8); const tStart = performance.now(); - if (!sessionRegistry.isSessionPrepared(sessionId)) { + const gooseSessionId = sessionRegistry.getGooseSessionId( + sessionId, + personaId, + ); + if (!gooseSessionId) { throw new Error("Session not prepared. Call acpPrepareSession first."); } @@ -110,7 +118,7 @@ export async function acpSendMessage( } const messageId = crypto.randomUUID(); - setActiveMessageId(sessionId, messageId); + setActiveMessageId(gooseSessionId, messageId); perfLog( `[perf:send] ${sid} acpSendMessage → prompt(len=${prompt.length}, imgs=${images?.length ?? 0})`, @@ -120,7 +128,7 @@ export async function acpSendMessage( if (personaId) meta.personaId = personaId; try { await directAcp.prompt( - sessionId, + gooseSessionId, content, Object.keys(meta).length > 0 ? meta : undefined, ); @@ -129,7 +137,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(sessionId); + clearActiveMessageId(gooseSessionId); } } @@ -138,16 +146,27 @@ export async function acpPrepareSession( sessionId: string, providerId: string, workingDir: string, -): Promise { + options: AcpPrepareSessionOptions = {}, +): Promise { const sid = sessionId.slice(0, 8); const t0 = performance.now(); perfLog( `[perf:prepare] ${sid} acpPrepareSession start (provider=${providerId})`, ); - await sessionRegistry.prepareSession(sessionId, providerId, workingDir); + const gooseSessionId = await sessionRegistry.prepareSession( + sessionId, + providerId, + workingDir, + { + personaId: options.personaId, + projectId: options.projectId, + knownNew: options.knownNew, + }, + ); perfLog( `[perf:prepare] ${sid} acpPrepareSession done in ${(performance.now() - t0).toFixed(1)}ms`, ); + return gooseSessionId; } export async function acpCreateSession( @@ -174,7 +193,8 @@ export async function acpSetModel( sessionId: string, modelId: string, ): Promise { - return directAcp.setModel(sessionId, modelId); + const gooseSessionId = sessionRegistry.getGooseSessionId(sessionId); + return directAcp.setModel(gooseSessionId ?? sessionId, modelId); } export type { AcpSessionInfo }; @@ -196,7 +216,19 @@ export async function acpSearchSessions( query: string, sessionIds: string[], ): Promise { - return searchSessionsViaExports(query, sessionIds); + const localByGooseSessionId = new Map(); + const gooseSessionIds = sessionIds.map((sessionId) => { + const gooseSessionId = + sessionRegistry.getGooseSessionId(sessionId) ?? sessionId; + localByGooseSessionId.set(gooseSessionId, sessionId); + return gooseSessionId; + }); + + const results = await searchSessionsViaExports(query, gooseSessionIds); + return results.map((result) => ({ + ...result, + sessionId: localByGooseSessionId.get(result.sessionId) ?? result.sessionId, + })); } /** @@ -205,21 +237,37 @@ export async function acpSearchSessions( * This triggers message replay via SessionNotification events that the * notification handler picks up automatically. */ -export async function acpLoadSession( +export function acpLoadSession( sessionId: string, workingDir?: string, +): Promise; +export function acpLoadSession( + sessionId: string, + gooseSessionId: string, + workingDir?: string, +): Promise; +export async function acpLoadSession( + sessionId: string, + gooseSessionIdOrWorkingDir?: string, + workingDir?: string, ): Promise { - const effectiveWorkingDir = workingDir ?? "~"; + const gooseSessionId = + workingDir === undefined + ? (sessionRegistry.getGooseSessionId(sessionId) ?? sessionId) + : (gooseSessionIdOrWorkingDir ?? sessionId); + const effectiveWorkingDir = + workingDir === undefined ? (gooseSessionIdOrWorkingDir ?? "~") : workingDir; const sid = sessionId.slice(0, 8); const t0 = performance.now(); - const rollbackSessionRegistration = sessionRegistry.registerPreparedSession( + const rollbackSessionRegistration = sessionRegistry.registerSession( sessionId, + gooseSessionId, "goose", effectiveWorkingDir, ); try { perfLog(`[perf:load] ${sid} acpLoadSession → client.loadSession`); - await directAcp.loadSession(sessionId, effectiveWorkingDir); + await directAcp.loadSession(gooseSessionId, effectiveWorkingDir); perfLog( `[perf:load] ${sid} client.loadSession resolved in ${(performance.now() - t0).toFixed(1)}ms`, ); @@ -231,7 +279,8 @@ export async function acpLoadSession( /** Export a session as JSON via the goose binary. */ export async function acpExportSession(sessionId: string): Promise { - return directAcp.exportSession(sessionId); + const gooseSessionId = sessionRegistry.getGooseSessionId(sessionId); + return directAcp.exportSession(gooseSessionId ?? sessionId); } /** Import a session from JSON via the goose binary. Returns new session metadata. */ @@ -243,11 +292,19 @@ export async function acpImportSession(json: string): Promise { export async function acpDuplicateSession( sessionId: string, ): Promise { - return directAcp.forkSession(sessionId); + const gooseSessionId = sessionRegistry.getGooseSessionId(sessionId); + return directAcp.forkSession(gooseSessionId ?? sessionId); } /** Cancel an in-progress ACP session so the backend stops streaming. */ -export async function acpCancelSession(sessionId: string): Promise { - await directAcp.cancelSession(sessionId); +export async function acpCancelSession( + sessionId: string, + personaId?: string, +): Promise { + const gooseSessionId = sessionRegistry.getGooseSessionId( + sessionId, + personaId, + ); + await directAcp.cancelSession(gooseSessionId ?? sessionId); return true; } diff --git a/ui/goose2/src/shared/api/acpNotificationHandler.ts b/ui/goose2/src/shared/api/acpNotificationHandler.ts index 2f60a91413..a511f000b1 100644 --- a/ui/goose2/src/shared/api/acpNotificationHandler.ts +++ b/ui/goose2/src/shared/api/acpNotificationHandler.ts @@ -32,6 +32,10 @@ import { getToolCallIdentity, getToolChainSummary, } from "./acpToolCallIdentity"; +import { + getLocalSessionId, + subscribeToSessionRegistration, +} from "./acpSessionRegistry"; import { perfLog } from "@/shared/lib/perfLog"; // Pre-set message ID for the next live stream per session. @@ -50,6 +54,20 @@ interface LivePerf { chunkCount: number; } const livePerf = new Map(); +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); +}); const toolCallStatusFromUpdate = (status: string): ToolCallStatus => status === "failed" ? "error" : "completed"; @@ -129,7 +147,9 @@ export function clearActiveMessageId(sessionId: string): void { export async function handleSessionNotification( notification: SessionNotification, ): Promise { - const sessionId = notification.sessionId; + const gooseSessionId = notification.sessionId; + const localSessionId = getLocalSessionId(gooseSessionId); + const sessionId = localSessionId ?? gooseSessionId; const { update } = notification; const isReplay = useChatStore.getState().loadingSessionIds.has(sessionId); @@ -144,20 +164,20 @@ export async function handleSessionNotification( } perf.lastAt = now; perf.count += 1; - handleReplay(sessionId, update); + handleReplay(sessionId, gooseSessionId, update); } else { - const perf = livePerf.get(sessionId); + const perf = livePerf.get(gooseSessionId); if (perf && update.sessionUpdate === "agent_message_chunk") { perf.chunkCount += 1; if (perf.firstChunkAt === null) { perf.firstChunkAt = performance.now(); - const sid = sessionId.slice(0, 8); + const sid = gooseSessionId.slice(0, 8); perfLog( `[perf:stream] ${sid} first agent_message_chunk at ttft=${(perf.firstChunkAt - perf.sendStartedAt).toFixed(1)}ms`, ); } } - handleLive(sessionId, update); + handleLive(sessionId, gooseSessionId, update); } } @@ -179,7 +199,11 @@ function getChunkMessageId(update: SessionUpdate): string | null { : null; } -function handleReplay(sessionId: string, update: SessionUpdate): void { +function handleReplay( + sessionId: string, + gooseSessionId: string, + update: SessionUpdate, +): void { switch (update.sessionUpdate) { case "agent_message_chunk": { const msg = ensureReplayAssistantMessage( @@ -320,7 +344,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, update); break; default: @@ -328,7 +352,11 @@ function handleReplay(sessionId: string, update: SessionUpdate): void { } } -function handleLive(sessionId: string, update: SessionUpdate): void { +function handleLive( + sessionId: string, + gooseSessionId: string, + update: SessionUpdate, +): void { const store = useChatStore.getState(); switch (update.sessionUpdate) { @@ -336,6 +364,7 @@ function handleLive(sessionId: string, update: SessionUpdate): void { const messageId = ensureLiveAssistantMessage( sessionId, getChunkMessageId(update) ?? undefined, + gooseSessionId, ); if (update.content.type === "text" && "text" in update.content) { @@ -346,7 +375,11 @@ function handleLive(sessionId: string, update: SessionUpdate): void { } case "tool_call": { - const messageId = ensureLiveAssistantMessage(sessionId); + const messageId = ensureLiveAssistantMessage( + sessionId, + undefined, + gooseSessionId, + ); const identity = getToolCallIdentity(update); const chainSummary = getToolChainSummary(update); @@ -377,7 +410,9 @@ function handleLive(sessionId: string, update: SessionUpdate): void { sessionId, update.toolCallId, ); - const messageId = ownerMessageId ?? ensureLiveAssistantMessage(sessionId); + const messageId = + ownerMessageId ?? + ensureLiveAssistantMessage(sessionId, undefined, gooseSessionId); const patch = toolCallUpdatePatch(update); if ( @@ -460,7 +495,7 @@ function handleLive(sessionId: string, update: SessionUpdate): void { case "session_info_update": case "config_option_update": case "usage_update": - handleShared(sessionId, update); + handleShared(sessionId, gooseSessionId, update); break; default: @@ -468,7 +503,11 @@ function handleLive(sessionId: string, update: SessionUpdate): void { } } -function handleShared(sessionId: string, update: SessionUpdate): void { +function handleShared( + sessionId: string, + gooseSessionId: string, + update: SessionUpdate, +): void { switch (update.sessionUpdate) { case "session_info_update": { handleSessionInfoUpdate(sessionId, update); @@ -517,6 +556,13 @@ function handleShared(sessionId: string, update: SessionUpdate): void { case "usage_update": { const usage = update as SessionUpdate & { sessionUpdate: "usage_update" }; + if (!getLocalSessionId(gooseSessionId)) { + pendingUsageUpdates.set(gooseSessionId, { + accumulatedTotal: usage.used, + contextLimit: usage.size, + }); + break; + } useChatStore.getState().updateTokenState(sessionId, { accumulatedTotal: usage.used, @@ -563,6 +609,7 @@ function findLiveMessageIdWithToolCall( function ensureLiveAssistantMessage( sessionId: string, preferredMessageId?: string | null, + gooseSessionId = sessionId, ): string { const store = useChatStore.getState(); const existingStreamingMessageId = findStreamingMessageId(sessionId); @@ -577,7 +624,7 @@ function ensureLiveAssistantMessage( const messageId = preferredMessageId ?? - presetMessageIds.get(sessionId) ?? + presetMessageIds.get(gooseSessionId) ?? existingStreamingMessageId ?? crypto.randomUUID(); @@ -597,13 +644,14 @@ function ensureLiveAssistantMessage( store.setPendingAssistantProvider(sessionId, null); store.setStreamingMessageId(sessionId, messageId); - clearActiveMessageId(sessionId); + clearActiveMessageId(gooseSessionId); 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 index 7c7de01e2b..e409692327 100644 --- a/ui/goose2/src/shared/api/acpSessionRegistry.ts +++ b/ui/goose2/src/shared/api/acpSessionRegistry.ts @@ -2,33 +2,84 @@ import * as acpApi from "./acpApi"; import { perfLog } from "@/shared/lib/perfLog"; interface PreparedSession { + gooseSessionId: string; providerId: string; workingDir: string; } +interface PrepareSessionOptions { + personaId?: string; + projectId?: string; + knownNew?: boolean; +} + +type SessionRegistrationListener = ( + localSessionId: string, + gooseSessionId: string, +) => void; + const prepared = new Map(); +const gooseToLocal = new Map(); +const registrationListeners = new Set(); + +function makeKey(sessionId: string, personaId?: string): string { + return personaId && personaId.length > 0 + ? `${sessionId}__${personaId}` + : sessionId; +} + +function notifySessionRegistered( + localSessionId: string, + gooseSessionId: string, +): void { + for (const listener of registrationListeners) { + listener(localSessionId, gooseSessionId); + } +} + +function restoreGooseRegistration( + gooseSessionId: string, + localSessionId: string | undefined, +): void { + if (localSessionId === undefined) { + gooseToLocal.delete(gooseSessionId); + return; + } + + gooseToLocal.set(gooseSessionId, localSessionId); +} + +export function subscribeToSessionRegistration( + listener: SessionRegistrationListener, +): () => void { + registrationListeners.add(listener); + return () => registrationListeners.delete(listener); +} export async function prepareSession( sessionId: string, providerId: string, workingDir: string, -): Promise { + options: PrepareSessionOptions = {}, +): Promise { const sid = sessionId.slice(0, 8); + const { personaId, projectId, knownNew = false } = options; + const key = makeKey(sessionId, personaId); - const existing = prepared.get(sessionId); + const existing = prepared.get(key) ?? prepared.get(sessionId); if (existing) { const tReuse = performance.now(); let changed = false; if (existing.workingDir !== workingDir) { - await acpApi.updateWorkingDir(sessionId, workingDir); + await acpApi.updateWorkingDir(existing.gooseSessionId, workingDir); existing.workingDir = workingDir; changed = true; } if (existing.providerId !== providerId) { const tProv = performance.now(); - await acpApi.setProvider(sessionId, providerId); + await acpApi.setProvider(existing.gooseSessionId, providerId); perfLog( - `[perf:prepare] ${sid} reuse setProvider(${providerId}) in ${(performance.now() - tProv).toFixed(1)}ms`, + `[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; @@ -36,45 +87,125 @@ export async function prepareSession( perfLog( `[perf:prepare] ${sid} reuse existing session (updates=${changed}) in ${(performance.now() - tReuse).toFixed(1)}ms`, ); - return; + return existing.gooseSessionId; } - const tLoad = performance.now(); - await acpApi.loadSession(sessionId, workingDir); - perfLog( - `[perf:prepare] ${sid} registry loadSession ok in ${(performance.now() - tLoad).toFixed(1)}ms`, - ); + let gooseSessionId: string; + let providerApplied = false; - const tProv = performance.now(); - await acpApi.setProvider(sessionId, providerId); - perfLog( - `[perf:prepare] ${sid} registry setProvider(${providerId}) in ${(performance.now() - tProv).toFixed(1)}ms`, - ); + if (knownNew) { + const tNew = performance.now(); + const response = await acpApi.newSession( + workingDir, + providerId, + projectId, + personaId, + ); + gooseSessionId = response.sessionId; + providerApplied = true; + perfLog( + `[perf:prepare] ${sid} registry newSession done in ${(performance.now() - tNew).toFixed(1)}ms (goose_sid=${gooseSessionId.slice(0, 8)})`, + ); + } else { + const tLoad = performance.now(); + await acpApi.loadSession(sessionId, workingDir); + gooseSessionId = sessionId; + perfLog( + `[perf:prepare] ${sid} registry loadSession ok in ${(performance.now() - tLoad).toFixed(1)}ms`, + ); + } - const entry = { providerId, workingDir }; + if (!providerApplied) { + const tProv = performance.now(); + await acpApi.setProvider(gooseSessionId, providerId); + perfLog( + `[perf:prepare] ${sid} registry setProvider(${providerId}) in ${(performance.now() - tProv).toFixed(1)}ms (goose_sid=${gooseSessionId.slice(0, 8)})`, + ); + } + + 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; + 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 isSessionPrepared(sessionId: string): boolean { - return prepared.has(sessionId); + return getGooseSessionId(sessionId) !== null; } export function registerPreparedSession( sessionId: string, providerId: string, workingDir: string, +): () => void { + return registerSession(sessionId, sessionId, providerId, workingDir); +} + +export function registerSession( + sessionId: string, + gooseSessionId: string, + providerId: string, + workingDir: string, ): () => void { const previousEntry = prepared.get(sessionId); - const entry = { providerId, workingDir }; + const previousGooseEntry = prepared.get(gooseSessionId); + 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); } + if (gooseSessionId !== sessionId) { + prepared.delete(gooseSessionId); + if (previousGooseEntry) { + prepared.set(gooseSessionId, previousGooseEntry); + } + } + + restoreGooseRegistration(gooseSessionId, previousGooseSessionLocal); + if (previousEntry && previousEntry.gooseSessionId !== gooseSessionId) { + restoreGooseRegistration( + previousEntry.gooseSessionId, + previousSessionGooseLocal, + ); + } }; }