Preserve selected branch across project chats (#9010)

Signed-off-by: morgmart <98432065+morgmart@users.noreply.github.com>
Signed-off-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: Douwe Osinga <douwe@squareup.com>

ui/goose2/src/app/AppShell.tsx

@@ -47,6 +47,7 @@ import { perfLog } from "@/shared/lib/perfLog";
 import { useProviderInventoryStore } from "@/features/providers/stores/providerInventoryStore";
 import type { SkillInfo } from "@/features/skills/api/skills";
 import { toChatSkillDraft } from "@/features/skills/lib/skillChatPrompt";
+import { resolveInheritedProjectWorkspace } from "@/features/chat/lib/workspaceContext";
 import { OnboardingFlow } from "@/features/onboarding/ui/OnboardingFlow";
 import { useOnboardingGate } from "@/features/onboarding/hooks/useOnboardingGate";
 import { Spinner } from "@/shared/ui/spinner";
@@ -166,6 +167,7 @@ export function AppShell({ children }: { children?: React.ReactNode }) {
   const createSession = useChatSessionStore((s) => s.createSession);
   const patchSession = useChatSessionStore((s) => s.patchSession);
   const setActiveSession = useChatSessionStore((s) => s.setActiveSession);
+  const setActiveWorkspace = useChatSessionStore((s) => s.setActiveWorkspace);
   const archiveSession = useChatSessionStore((s) => s.archiveSession);
   const selectedProvider = useAgentStore(selectSelectedProvider);
   const projects = useProjectStore(selectProjects);
@@ -205,7 +207,14 @@ export function AppShell({ children }: { children?: React.ReactNode }) {
             .getState()
             .projects.find((p) => p.id === session.projectId) ?? null)
         : null;
-      const workingDir = await resolveSessionCwd(project);
+      const activeWorkspace =
+        session?.id != null
+          ? useChatSessionStore.getState().activeWorkspaceBySession[session.id]
+          : undefined;
+      const workingDir = await resolveSessionCwd(
+        project,
+        activeWorkspace?.path ?? session?.workingDir,
+      );
       await acpLoadSession(sessionId, workingDir);
       const tFlush = performance.now();
       useChatStore.getState().setSessionLoading(sessionId, false);
@@ -411,6 +420,12 @@ export function AppShell({ children }: { children?: React.ReactNode }) {
         );
       const sessionState = useChatSessionStore.getState();
       const chatState = useChatStore.getState();
+      const inheritedWorkspace = resolveInheritedProjectWorkspace({
+        projectId: project?.id,
+        sessions: sessionState.sessions,
+        activeSessionId: sessionState.activeSessionId,
+        activeWorkspaceBySession: sessionState.activeWorkspaceBySession,
+      });
       const existingDraft = findExistingDraft({
         sessions: sessionState.sessions,
         activeSessionId: sessionState.activeSessionId,
@@ -423,6 +438,12 @@ export function AppShell({ children }: { children?: React.ReactNode }) {
       });
 
       if (existingDraft) {
+        if (inheritedWorkspace) {
+          setActiveWorkspace(existingDraft.id, inheritedWorkspace);
+          patchSession(existingDraft.id, {
+            workingDir: inheritedWorkspace.path,
+          });
+        }
         setActiveSession(existingDraft.id);
         clearSettingsSectionUrl();
         setActiveView("chat");
@@ -433,7 +454,10 @@ export function AppShell({ children }: { children?: React.ReactNode }) {
         return existingDraft;
       }
 
-      const workingDir = await resolveSessionCwd(project);
+      const workingDir = await resolveSessionCwd(
+        project,
+        inheritedWorkspace?.path,
+      );
       const session = await createSession({
         title,
         projectId: project?.id,
@@ -442,6 +466,9 @@ export function AppShell({ children }: { children?: React.ReactNode }) {
         modelId: sessionModelPreference.modelId,
         modelName: sessionModelPreference.modelName,
       });
+      if (inheritedWorkspace) {
+        setActiveWorkspace(session.id, inheritedWorkspace);
+      }
       setActiveSession(session.id);
       clearSettingsSectionUrl();
       setActiveView("chat");
@@ -454,8 +481,10 @@ export function AppShell({ children }: { children?: React.ReactNode }) {
     [
       selectedProvider,
       createSession,
+      patchSession,
       providerInventoryEntries,
       setActiveSession,
+      setActiveWorkspace,
       setChatActiveSession,
     ],
   );
@@ -639,6 +668,7 @@ export function AppShell({ children }: { children?: React.ReactNode }) {
           workingDir,
           modelId: session.modelId,
         });
+        patchSession(sessionId, { workingDir });
       })().catch((error) => {
         console.error(
           "Failed to update ACP session project working directory:",
@@ -646,7 +676,7 @@ export function AppShell({ children }: { children?: React.ReactNode }) {
         );
       });
     },
-    [selectedProvider],
+    [selectedProvider, patchSession],
   );
 
   const handleRenameChat = useCallback(

ui/goose2/src/features/chat/hooks/__tests__/useChat.attachments.test.ts

@@ -31,6 +31,7 @@ describe("useChat attachments", () => {
       sessions: [],
       activeSessionId: null,
       isLoading: false,
+      isContextPanelOpen: false,
       contextPanelOpenBySession: {},
       activeWorkspaceBySession: {},
     });

ui/goose2/src/features/chat/hooks/__tests__/useChat.skillChips.test.ts

@@ -38,6 +38,7 @@ describe("useChat skill chips", () => {
       sessions: [],
       activeSessionId: null,
       isLoading: false,
+      isContextPanelOpen: false,
       contextPanelOpenBySession: {},
       activeWorkspaceBySession: {},
     });

ui/goose2/src/features/chat/hooks/__tests__/useChat.test.ts

@@ -73,6 +73,7 @@ describe("useChat", () => {
       sessions: [],
       activeSessionId: null,
       isLoading: false,
+      isContextPanelOpen: false,
       contextPanelOpenBySession: {},
       activeWorkspaceBySession: {},
     });

ui/goose2/src/features/chat/hooks/__tests__/useChatSessionController.test.ts

@@ -186,6 +186,7 @@ describe("useChatSessionController", () => {
       activeSessionId: null,
       isLoading: false,
       hasHydratedSessions: true,
+      isContextPanelOpen: false,
       contextPanelOpenBySession: {},
       activeWorkspaceBySession: {},
     });

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

@@ -369,11 +369,13 @@ export function useChat(
     setPendingAssistantProvider,
   ]);
 
-  const getWorkingDir = useCallback(
-    () =>
-      useChatSessionStore.getState().activeWorkspaceBySession[sessionId]?.path,
-    [sessionId],
-  );
+  const getWorkingDir = useCallback(() => {
+    const sessionStore = useChatSessionStore.getState();
+    return (
+      sessionStore.activeWorkspaceBySession[sessionId]?.path ??
+      sessionStore.getSession(sessionId)?.workingDir
+    );
+  }, [sessionId]);
 
   const compactConversation = useCallback(
     async (overridePersona?: { id: string; name?: string }) => {

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

@@ -170,6 +170,7 @@ export function useChatSessionController({
         workingDir,
         modelId: modelSelection?.id,
       });
+      useChatSessionStore.getState().patchSession(sessionId, { workingDir });
       if (!result.applied || !modelSelection?.id) {
         return result.applied;
       }

ui/goose2/src/features/chat/lib/workspaceContext.test.ts

@@ -0,0 +1,55 @@
+import { describe, expect, it } from "vitest";
+import type { ChatSession } from "../stores/chatSessionStore";
+import { resolveInheritedProjectWorkspace } from "./workspaceContext";
+
+function makeSession(overrides: Partial<ChatSession> = {}): ChatSession {
+  return {
+    id: "session-1",
+    title: "New chat",
+    projectId: "project-1",
+    workingDir: "/repo/main",
+    createdAt: "2026-01-01T00:00:00.000Z",
+    updatedAt: "2026-01-01T00:00:00.000Z",
+    messageCount: 0,
+    ...overrides,
+  };
+}
+
+describe("resolveInheritedProjectWorkspace", () => {
+  it("inherits the active workspace for another chat in the same project", () => {
+    expect(
+      resolveInheritedProjectWorkspace({
+        projectId: "project-1",
+        sessions: [makeSession()],
+        activeSessionId: "session-1",
+        activeWorkspaceBySession: {
+          "session-1": { path: "/repo/feature", branch: "feature" },
+        },
+      }),
+    ).toEqual({ path: "/repo/feature", branch: "feature" });
+  });
+
+  it("falls back to the active session working directory after reload", () => {
+    expect(
+      resolveInheritedProjectWorkspace({
+        projectId: "project-1",
+        sessions: [makeSession({ workingDir: "/repo/feature" })],
+        activeSessionId: "session-1",
+        activeWorkspaceBySession: {},
+      }),
+    ).toEqual({ path: "/repo/feature", branch: null });
+  });
+
+  it("does not inherit workspace across projects", () => {
+    expect(
+      resolveInheritedProjectWorkspace({
+        projectId: "project-2",
+        sessions: [makeSession()],
+        activeSessionId: "session-1",
+        activeWorkspaceBySession: {
+          "session-1": { path: "/repo/feature", branch: "feature" },
+        },
+      }),
+    ).toBeUndefined();
+  });
+});

ui/goose2/src/features/chat/lib/workspaceContext.ts

@@ -0,0 +1,38 @@
+import type {
+  ActiveWorkspace,
+  ChatSession,
+} from "@/features/chat/stores/chatSessionStore";
+
+interface ResolveInheritedProjectWorkspaceOptions {
+  projectId: string | undefined;
+  sessions: ChatSession[];
+  activeSessionId: string | null;
+  activeWorkspaceBySession: Record<string, ActiveWorkspace>;
+}
+
+export function resolveInheritedProjectWorkspace({
+  projectId,
+  sessions,
+  activeSessionId,
+  activeWorkspaceBySession,
+}: ResolveInheritedProjectWorkspaceOptions): ActiveWorkspace | undefined {
+  if (!projectId || !activeSessionId) {
+    return undefined;
+  }
+
+  const activeSession = sessions.find(
+    (session) => session.id === activeSessionId,
+  );
+  if (activeSession?.projectId !== projectId) {
+    return undefined;
+  }
+
+  const activeWorkspace = activeWorkspaceBySession[activeSessionId];
+  if (activeWorkspace?.path) {
+    return activeWorkspace;
+  }
+
+  return activeSession.workingDir
+    ? { path: activeSession.workingDir, branch: null }
+    : undefined;
+}

ui/goose2/src/features/chat/stores/__tests__/chatSessionStore.test.ts

@@ -23,6 +23,7 @@ function resetStore() {
     activeSessionId: null,
     isLoading: false,
     hasHydratedSessions: false,
+    isContextPanelOpen: false,
     contextPanelOpenBySession: {},
     activeWorkspaceBySession: {},
   });
@@ -49,6 +50,7 @@ function seedSession(overrides: Partial<ChatSession> = {}): ChatSession {
 
 describe("chatSessionStore", () => {
   beforeEach(() => {
+    window.localStorage.removeItem("goose:context-panel-open");
     resetStore();
     vi.clearAllMocks();
   });
@@ -271,6 +273,20 @@ describe("chatSessionStore", () => {
     });
   });
 
+  describe("context panel preference", () => {
+    it("stores context panel open state as a global preference", () => {
+      useChatSessionStore.getState().setContextPanelOpen("session-1", true);
+
+      expect(useChatSessionStore.getState().isContextPanelOpen).toBe(true);
+      expect(window.localStorage.getItem("goose:context-panel-open")).toBe("1");
+
+      useChatSessionStore.getState().setContextPanelOpen("session-2", false);
+
+      expect(useChatSessionStore.getState().isContextPanelOpen).toBe(false);
+      expect(window.localStorage.getItem("goose:context-panel-open")).toBe("0");
+    });
+  });
+
   describe("archiveSession", () => {
     it("sets archivedAt on the session", async () => {
       const session = seedSession();

ui/goose2/src/features/chat/stores/chatSessionStore.ts

@@ -14,6 +14,8 @@ import {
   unarchiveSession as acpUnarchiveSession,
 } from "@/shared/api/acpApi";
 
+const CONTEXT_PANEL_OPEN_STORAGE_KEY = "goose:context-panel-open";
+
 export interface ChatSession {
   id: string;
   title: string;
@@ -58,6 +60,7 @@ interface ChatSessionStoreState {
   activeSessionId: string | null;
   isLoading: boolean;
   hasHydratedSessions: boolean;
+  isContextPanelOpen: boolean;
   contextPanelOpenBySession: Record<string, boolean>;
   activeWorkspaceBySession: Record<string, ActiveWorkspace>;
 }
@@ -116,6 +119,29 @@ function sortByUpdatedAtDesc(sessions: ChatSession[]): ChatSession[] {
   );
 }
 
+function loadContextPanelOpenPreference(): boolean {
+  if (typeof window === "undefined") return false;
+
+  try {
+    return window.localStorage.getItem(CONTEXT_PANEL_OPEN_STORAGE_KEY) === "1";
+  } catch {
+    return false;
+  }
+}
+
+function persistContextPanelOpenPreference(open: boolean): void {
+  if (typeof window === "undefined") return;
+
+  try {
+    window.localStorage.setItem(
+      CONTEXT_PANEL_OPEN_STORAGE_KEY,
+      open ? "1" : "0",
+    );
+  } catch {
+    // localStorage may be unavailable
+  }
+}
+
 export function sessionToChatSession(session: Session): ChatSession {
   return {
     id: session.id,
@@ -138,6 +164,7 @@ export const useChatSessionStore = create<ChatSessionStore>((set, get) => ({
   activeSessionId: null,
   isLoading: false,
   hasHydratedSessions: false,
+  isContextPanelOpen: loadContextPanelOpenPreference(),
   contextPanelOpenBySession: {},
   activeWorkspaceBySession: {},
 
@@ -255,13 +282,9 @@ export const useChatSessionStore = create<ChatSessionStore>((set, get) => ({
     set({ activeSessionId: sessionId });
   },
 
-  setContextPanelOpen: (sessionId, open) => {
-    set((state) => ({
-      contextPanelOpenBySession: {
-        ...state.contextPanelOpenBySession,
-        [sessionId]: open,
-      },
-    }));
+  setContextPanelOpen: (_sessionId, open) => {
+    persistContextPanelOpenPreference(open);
+    set({ isContextPanelOpen: open, contextPanelOpenBySession: {} });
   },
 
   setActiveWorkspace: (sessionId, context) => {

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

@@ -27,6 +27,7 @@ interface ChatContextPanelProps {
     color?: string;
     workingDirs?: string[];
   } | null;
+  sessionWorkingDir?: string | null;
   setOpen: (sessionId: string, open: boolean) => void;
 }
 
@@ -35,6 +36,7 @@ export function ChatContextPanel({
   isOpen,
   label,
   project,
+  sessionWorkingDir,
   setOpen,
 }: ChatContextPanelProps) {
   const shouldReduceMotion = useReducedMotion();
@@ -104,6 +106,7 @@ export function ChatContextPanel({
                   projectName={project?.name}
                   projectColor={project?.color}
                   projectWorkingDirs={project?.workingDirs ?? []}
+                  sessionWorkingDir={sessionWorkingDir}
                 />
               </aside>
             </motion.div>

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

@@ -47,9 +47,7 @@ export function ChatView({
 }: ChatViewProps) {
   const { t } = useTranslation("chat");
   const mountStart = useRef(performance.now());
-  const isContextPanelOpen = useChatSessionStore(
-    (s) => s.contextPanelOpenBySession[sessionId] ?? false,
-  );
+  const isContextPanelOpen = useChatSessionStore((s) => s.isContextPanelOpen);
   const setContextPanelOpen = useChatSessionStore((s) => s.setContextPanelOpen);
   const [isLoadingIndicatorMounted, setIsLoadingIndicatorMounted] =
     useState(false);
@@ -189,6 +187,7 @@ export function ChatView({
           isOpen={isContextPanelOpen}
           label={contextPanelLabel}
           project={controller.project}
+          sessionWorkingDir={controller.session?.workingDir}
           setOpen={setContextPanelOpen}
         />
       </div>

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

@@ -26,6 +26,7 @@ interface ContextPanelProps {
   projectName?: string;
   projectColor?: string;
   projectWorkingDirs?: string[];
+  sessionWorkingDir?: string | null;
 }
 
 type ContextPanelTab = "details" | "files";
@@ -64,6 +65,7 @@ export function ContextPanel({
   projectName,
   projectColor,
   projectWorkingDirs = [],
+  sessionWorkingDir,
 }: ContextPanelProps) {
   const { t } = useTranslation("chat");
   const [activeTab, setActiveTab] = useState<ContextPanelTab>("details");
@@ -77,7 +79,8 @@ export function ContextPanel({
   );
   const setActiveWorkspace = useChatSessionStore((s) => s.setActiveWorkspace);
 
-  const gitTargetPath = activeContext?.path ?? primaryWorkspaceRoot;
+  const gitTargetPath =
+    activeContext?.path ?? sessionWorkingDir ?? primaryWorkspaceRoot;
   const {
     data: gitState,
     error,
@@ -230,6 +233,7 @@ export function ContextPanel({
             projectName={projectName}
             projectColor={projectColor}
             projectWorkingDirs={projectWorkingDirs}
+            sessionWorkingDir={sessionWorkingDir}
             gitState={gitState}
             isLoading={isLoading}
             isFetching={isFetching}

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

@@ -55,6 +55,7 @@ vi.mock("../../hooks/useChatSessionController", () => ({
 vi.mock("../../stores/chatSessionStore", () => ({
   useChatSessionStore: (selector: (state: unknown) => unknown) =>
     selector({
+      isContextPanelOpen: false,
       contextPanelOpenBySession: {},
       setContextPanelOpen: vi.fn(),
     }),

ui/goose2/src/features/chat/ui/widgets/WorkspaceWidget.tsx

@@ -12,6 +12,7 @@ interface WorkspaceWidgetProps {
   projectName?: string;
   projectColor?: string;
   projectWorkingDirs: string[];
+  sessionWorkingDir?: string | null;
   gitState: GitState | undefined;
   isLoading: boolean;
   isFetching: boolean;
@@ -44,6 +45,7 @@ export function WorkspaceWidget({
   projectName,
   projectColor,
   projectWorkingDirs,
+  sessionWorkingDir,
   gitState,
   isLoading,
   isFetching,
@@ -62,7 +64,8 @@ export function WorkspaceWidget({
   onToggleOpen,
 }: WorkspaceWidgetProps) {
   const { t } = useTranslation("chat");
-  const primaryWorkspaceRoot = projectWorkingDirs[0] ?? null;
+  const primaryWorkspaceRoot =
+    activeContext?.path ?? sessionWorkingDir ?? projectWorkingDirs[0] ?? null;
 
   const gitErrorMessage =
     error instanceof Error ? error.message : t("contextPanel.errors.gitRead");

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

@@ -20,6 +20,7 @@ describe("ACP session info updates", () => {
       activeSessionId: null,
       isLoading: false,
       hasHydratedSessions: false,
+      isContextPanelOpen: false,
       contextPanelOpenBySession: {},
       activeWorkspaceBySession: {},
     });