From 570ae9dbc853d5f3ba27b6d1e53f598c8cc2b12c Mon Sep 17 00:00:00 2001 From: Bohdan Triapitsyn Date: Wed, 3 Jun 2026 15:16:36 +0300 Subject: [PATCH] fix: keep worktree sessions in the right group Prevents stale worktree lists from overwriting newly created worktrees Uses the created worktree path when selecting linked worktree sessions --- .../src/components/session/SessionSidebar.tsx | 2 +- .../src/lib/worktrees/worktreeManager.test.ts | 106 +++++++++++++++++ .../ui/src/lib/worktrees/worktreeManager.ts | 107 +++++++++++------- packages/ui/src/sync/session-actions.test.ts | 13 ++- packages/ui/src/sync/session-actions.ts | 20 ++-- 5 files changed, 191 insertions(+), 57 deletions(-) create mode 100644 packages/ui/src/lib/worktrees/worktreeManager.test.ts diff --git a/packages/ui/src/components/session/SessionSidebar.tsx b/packages/ui/src/components/session/SessionSidebar.tsx index a6e29125..2479df12 100644 --- a/packages/ui/src/components/session/SessionSidebar.tsx +++ b/packages/ui/src/components/session/SessionSidebar.tsx @@ -1679,7 +1679,7 @@ export const SessionSidebar: React.FC = ({ setSessionSwitcherOpen(false); } if (options?.sessionId) { - setCurrentSession(options.sessionId); + setCurrentSession(options.sessionId, worktreePath); return; } openNewSessionDraft({ directoryOverride: worktreePath }); diff --git a/packages/ui/src/lib/worktrees/worktreeManager.test.ts b/packages/ui/src/lib/worktrees/worktreeManager.test.ts new file mode 100644 index 00000000..7344a7ad --- /dev/null +++ b/packages/ui/src/lib/worktrees/worktreeManager.test.ts @@ -0,0 +1,106 @@ +import { beforeEach, describe, expect, mock, test } from 'bun:test'; +import type { WorktreeMetadata } from '@/types/worktree'; + +type WorktreeListEntry = { + path?: string; + branch?: string; + head?: string; + name?: string; +}; + +const listCalls: string[] = []; +const listResolvers: Array<(value: WorktreeListEntry[]) => void> = []; +const createdWorktree = { + name: 'feature', + branch: 'feature', + path: '/repo-feature', +}; + +const sessionState = { + availableWorktreesByProject: new Map(), + availableWorktrees: [] as WorktreeMetadata[], +}; + +mock.module('@/lib/openchamberConfig', () => ({ + substituteCommandVariables: (command: string) => command, +})); + +mock.module('@/lib/worktrees/worktreeBootstrap', () => ({ + clearWorktreeBootstrapState: mock(), + markWorktreeBootstrapPending: mock(), +})); + +mock.module('@/lib/worktrees/worktreeStatus', () => ({ + invalidateResolvedProjectRootCache: mock(), + resolveProjectRoot: (directory: string) => Promise.resolve(directory), +})); + +mock.module('@/sync/session-ui-store', () => ({ + useSessionUIStore: { + getState: () => sessionState, + setState: (patch: Partial | ((state: typeof sessionState) => Partial)) => { + const next = typeof patch === 'function' ? patch(sessionState) : patch; + Object.assign(sessionState, next); + }, + }, +})); + +mock.module('@/lib/gitApi', () => ({ + deleteRemoteBranch: mock(), + git: { + worktree: { + list: (directory: string) => { + listCalls.push(directory); + return new Promise((resolve) => { + listResolvers.push(resolve); + }); + }, + create: mock(() => Promise.resolve(createdWorktree)), + remove: mock(() => Promise.resolve({ success: true })), + }, + }, +})); + +const { createWorktree, listProjectWorktrees } = await import('./worktreeManager'); + +const waitForListCallCount = async (count: number): Promise => { + for (let attempt = 0; attempt < 10; attempt += 1) { + if (listCalls.length >= count) { + return; + } + await Promise.resolve(); + } + throw new Error(`Expected ${count} worktree list calls, got ${listCalls.length}`); +}; + +describe('worktreeManager list invalidation', () => { + beforeEach(() => { + listCalls.length = 0; + listResolvers.length = 0; + sessionState.availableWorktreesByProject = new Map(); + sessionState.availableWorktrees = []; + }); + + test('retries an in-flight list when a worktree is created before it resolves', async () => { + const project = { id: 'project-1', path: '/repo' }; + const listing = listProjectWorktrees(project); + + await waitForListCallCount(1); + + await createWorktree(project, { + preferredName: 'feature', + mode: 'new', + branchName: 'feature', + worktreeName: 'feature', + }); + + listResolvers[0]([]); + await waitForListCallCount(2); + listResolvers[1]([createdWorktree]); + + const result = await listing; + + expect(listCalls).toEqual(['/repo', '/repo']); + expect(result.map((entry) => entry.path)).toEqual(['/repo-feature']); + }); +}); diff --git a/packages/ui/src/lib/worktrees/worktreeManager.ts b/packages/ui/src/lib/worktrees/worktreeManager.ts index 2a09698e..f4f9a666 100644 --- a/packages/ui/src/lib/worktrees/worktreeManager.ts +++ b/packages/ui/src/lib/worktrees/worktreeManager.ts @@ -150,8 +150,67 @@ const toCreatePayload = (args: { // Cache worktree listings to avoid repeated git worktree list + rev-parse calls const _worktreeListCache = new Map(); const _worktreeListInflight = new Map>(); +const _worktreeListGeneration = new Map(); const WORKTREE_LIST_CACHE_TTL = 30_000; // 30 seconds +const getWorktreeListGeneration = (projectDirectory: string): number => { + return _worktreeListGeneration.get(projectDirectory) ?? 0; +}; + +const invalidateWorktreeList = (projectDirectory: string): void => { + _worktreeListGeneration.set(projectDirectory, getWorktreeListGeneration(projectDirectory) + 1); + _worktreeListCache.delete(projectDirectory); +}; + +const readProjectWorktrees = async (projectDirectory: string): Promise => { + const metadataProjectDirectory = await resolveProjectRoot(projectDirectory).catch(() => projectDirectory); + const normalizedProjectDirectory = normalizePath(projectDirectory); + + const worktrees = await git.worktree.list(projectDirectory).catch(() => []); + const results: WorktreeMetadata[] = worktrees + .filter((entry) => typeof entry.path === 'string' && entry.path.trim().length > 0) + .map((entry) => { + const worktreePath = normalizePath(entry.path); + const branch = (entry.branch || '').replace(/^refs\/heads\//, '').trim(); + const name = (entry.name || '').trim(); + + // Derive canonical worktree metadata from worktree list entry + const canonical = deriveCanonicalWorktreeFields(entry, worktreePath); + + return { + source: 'sdk' as const, + name: name || deriveSdkWorktreeNameFromDirectory(worktreePath), + path: worktreePath, + projectDirectory: metadataProjectDirectory, + branch: branch, + label: branch || name || deriveSdkWorktreeNameFromDirectory(worktreePath), + worktreeRoot: canonical.worktreeRoot, + worktreeStatus: canonical.worktreeStatus, + headState: canonical.headState, + worktreeSource: canonical.worktreeSource, + }; + }) + .filter((entry) => normalizePath(entry.path) !== normalizedProjectDirectory); + + return results.sort((a, b) => { + const aLabel = (a.label || a.branch || a.path).toLowerCase(); + const bLabel = (b.label || b.branch || b.path).toLowerCase(); + return aLabel.localeCompare(bLabel); + }); +}; + +const readStableProjectWorktrees = async (projectDirectory: string): Promise => { + while (true) { + const generation = getWorktreeListGeneration(projectDirectory); + const worktrees = await readProjectWorktrees(projectDirectory); + + if (generation === getWorktreeListGeneration(projectDirectory)) { + _worktreeListCache.set(projectDirectory, { value: worktrees, at: Date.now() }); + return worktrees; + } + } +}; + export async function listProjectWorktrees(project: ProjectRef): Promise { const projectDirectory = normalizePath(project.path); @@ -165,46 +224,10 @@ export async function listProjectWorktrees(project: ProjectRef): Promise => { - const metadataProjectDirectory = await resolveProjectRoot(projectDirectory).catch(() => projectDirectory); - const normalizedProjectDirectory = normalizePath(projectDirectory); - - const worktrees = await git.worktree.list(projectDirectory).catch(() => []); - const results: WorktreeMetadata[] = worktrees - .filter((entry) => typeof entry.path === 'string' && entry.path.trim().length > 0) - .map((entry) => { - const worktreePath = normalizePath(entry.path); - const branch = (entry.branch || '').replace(/^refs\/heads\//, '').trim(); - const name = (entry.name || '').trim(); - - // Derive canonical worktree metadata from worktree list entry - const canonical = deriveCanonicalWorktreeFields(entry, worktreePath); - - return { - source: 'sdk' as const, - name: name || deriveSdkWorktreeNameFromDirectory(worktreePath), - path: worktreePath, - projectDirectory: metadataProjectDirectory, - branch: branch, - label: branch || name || deriveSdkWorktreeNameFromDirectory(worktreePath), - worktreeRoot: canonical.worktreeRoot, - worktreeStatus: canonical.worktreeStatus, - headState: canonical.headState, - worktreeSource: canonical.worktreeSource, - }; - }) - .filter((entry) => normalizePath(entry.path) !== normalizedProjectDirectory); - - const sorted = results.sort((a, b) => { - const aLabel = (a.label || a.branch || a.path).toLowerCase(); - const bLabel = (b.label || b.branch || b.path).toLowerCase(); - return aLabel.localeCompare(bLabel); - }); - - _worktreeListCache.set(projectDirectory, { value: sorted, at: Date.now() }); - return sorted; - })().finally(() => { - _worktreeListInflight.delete(projectDirectory); + const promise = readStableProjectWorktrees(projectDirectory).finally(() => { + if (_worktreeListInflight.get(projectDirectory) === promise) { + _worktreeListInflight.delete(projectDirectory); + } }); _worktreeListInflight.set(projectDirectory, promise); @@ -255,7 +278,7 @@ export async function createWorktree(project: ProjectRef, args: CreateWorktreeAr markWorktreeBootstrapPending(metadata.path); - _worktreeListCache.delete(projectDirectory); + invalidateWorktreeList(projectDirectory); // The new worktree changes the repo's worktree topology; drop cached root // resolutions so root-branch lookups re-resolve against the new layout. invalidateResolvedProjectRootCache(); @@ -300,7 +323,7 @@ export async function removeProjectWorktree(project: ProjectRef, worktree: Workt clearWorktreeBootstrapState(worktree.path); - _worktreeListCache.delete(normalizePath(project.path)); + invalidateWorktreeList(normalizePath(project.path)); // Removing a worktree changes the repo's worktree topology; drop cached root // resolutions so root-branch lookups re-resolve against the new layout. invalidateResolvedProjectRootCache(); diff --git a/packages/ui/src/sync/session-actions.test.ts b/packages/ui/src/sync/session-actions.test.ts index d8d78b29..4e8608de 100644 --- a/packages/ui/src/sync/session-actions.test.ts +++ b/packages/ui/src/sync/session-actions.test.ts @@ -5,6 +5,7 @@ import type { QuestionRequest } from "@/types/question" // Mock SDK client that records permission.reply / question.reply calls const replyCalls: Array<{ method: string; params: Record }> = [] const scopedClientDirectories: string[] = [] +const registeredSessionDirectories: Array<{ sessionID: string; directory: string }> = [] let sessionRevertResult: { data?: unknown; error?: unknown; response?: { status?: number } } = {} let questionReplyError: unknown | null = null @@ -139,14 +140,18 @@ mock.module("./input-store", () => ({ }, })) -// Mock useGlobalSessionsStore (imported but not used in permission functions) mock.module("@/stores/useGlobalSessionsStore", () => ({ - useGlobalSessionsStore: {}, + useGlobalSessionsStore: { + getState: () => ({ + upsertSession: () => {}, + }), + }, })) -// Mock sync-refs (imported but not used in permission functions) mock.module("./sync-refs", () => ({ - registerSessionDirectory: () => {}, + registerSessionDirectory: (sessionID: string, directory: string) => { + registeredSessionDirectories.push({ sessionID, directory }) + }, })) import { create, type StoreApi } from "zustand" diff --git a/packages/ui/src/sync/session-actions.ts b/packages/ui/src/sync/session-actions.ts index 786d1deb..2c297344 100644 --- a/packages/ui/src/sync/session-actions.ts +++ b/packages/ui/src/sync/session-actions.ts @@ -335,16 +335,16 @@ export async function createSession( parentID: parentID ?? undefined, }, directoryOverride ?? dir()) - const sessionDirectory = (session as { directory?: string }).directory ?? directoryOverride ?? null - // Pre-populate routing index so SSE events arriving before session.created - // can be routed to the correct child store - if (sessionDirectory) { - registerSessionDirectory(session.id, sessionDirectory) - } - useSessionUIStore.getState().setCurrentSession(session.id, sessionDirectory) - useSessionUIStore.getState().markSessionAsOpenChamberCreated(session.id) - useGlobalSessionsStore.getState().upsertSession(session) - return session + const sessionDirectory = (session as { directory?: string | null }).directory ?? null + // Pre-populate routing index so SSE events arriving before session.created + // can be routed to the correct child store + if (sessionDirectory) { + registerSessionDirectory(session.id, sessionDirectory) + } + useSessionUIStore.getState().setCurrentSession(session.id, sessionDirectory) + useSessionUIStore.getState().markSessionAsOpenChamberCreated(session.id) + useGlobalSessionsStore.getState().upsertSession(session) + return session } catch (error) { console.error("[session-actions] createSession failed", error) return null