fix: make project new chat optimistic

Signed-off-by: Matt Toohey <contact@matttoohey.com>
2026-07-03 14:10:03 +02:00 · 2026-05-01 16:43:30 +10:00
parent b17fd4f287
commit bf8fc7162b
11 changed files with 747 additions and 100 deletions
@@ -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:",
@@ -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",
@@ -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<void> {
+  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,
+  });
+}
@@ -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;
@@ -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);
@@ -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> = {}): 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(
@@ -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<ChatSession>;
+  createLocalSession: (
+    opts?: Omit<CreateSessionOpts, "workingDir">,
+  ) => ChatSession;
+  prepareSessionBinding: (args: PrepareSessionBindingArgs) => Promise<void>;
  loadSessions: () => Promise<void>;
  updateSession: (id: string, patch: Partial<ChatSession>) => 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<string, ChatSession>();
+  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<ChatSessionStore>((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<ChatSessionStore>((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<ChatSessionStore>((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<ChatSessionStore>((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<ChatSessionStore>((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<ChatSessionStore>((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),
      );
    }
@@ -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",
+    );
+  });
 });
@@ -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<void> {
+  options: AcpPrepareSessionOptions = {},
+): Promise<string> {
  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<void> {
-  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<AcpSessionSearchResult[]> {
-  return searchSessionsViaExports(query, sessionIds);
+  const localByGooseSessionId = new Map<string, string>();
+  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<void>;
+export function acpLoadSession(
+  sessionId: string,
+  gooseSessionId: string,
+  workingDir?: string,
+): Promise<void>;
+export async function acpLoadSession(
+  sessionId: string,
+  gooseSessionIdOrWorkingDir?: string,
+  workingDir?: string,
 ): Promise<void> {
-  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<string> {
-  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<AcpSessionInfo> {
 export async function acpDuplicateSession(
  sessionId: string,
 ): Promise<AcpSessionInfo> {
-  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<boolean> {
-  await directAcp.cancelSession(sessionId);
+export async function acpCancelSession(
+  sessionId: string,
+  personaId?: string,
+): Promise<boolean> {
+  const gooseSessionId = sessionRegistry.getGooseSessionId(
+    sessionId,
+    personaId,
+  );
+  await directAcp.cancelSession(gooseSessionId ?? sessionId);
  return true;
 }
@@ -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<string, LivePerf>();
+const pendingUsageUpdates = new Map<
+  string,
+  { accumulatedTotal: number; contextLimit: number }
+>();
+
+subscribeToSessionRegistration((localSessionId, gooseSessionId) => {
+  const pendingUsage = pendingUsageUpdates.get(gooseSessionId);
+  if (!pendingUsage) {
+    return;
+  }
+
+  useChatStore.getState().updateTokenState(localSessionId, pendingUsage);
+  pendingUsageUpdates.delete(gooseSessionId);
+});

 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<void> {
-  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();
 }

@@ -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<string, PreparedSession>();
+const gooseToLocal = new Map<string, string>();
+const registrationListeners = new Set<SessionRegistrationListener>();
+
+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<void> {
+  options: PrepareSessionOptions = {},
+): Promise<string> {
  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,
+      );
+    }
  };
 }