diff --git a/apps/web/src/components/ChatView.tsx b/apps/web/src/components/ChatView.tsx index 7fc7669fe60..94ed9830303 100644 --- a/apps/web/src/components/ChatView.tsx +++ b/apps/web/src/components/ChatView.tsx @@ -209,118 +209,131 @@ type EnvironmentUnavailableState = { }; type ThreadPlanCatalogEntry = Pick; +type ThreadPlanCatalogStoreState = ReturnType; -function useThreadPlanCatalog(threadIds: readonly ThreadId[]): ThreadPlanCatalogEntry[] { - return useStore( - useMemo(() => { - let previousThreadIds: readonly ThreadId[] = []; - let previousResult: ThreadPlanCatalogEntry[] = []; - let previousEntries = new Map< - ThreadId, - { - shell: object | null; - proposedPlanIds: readonly string[] | undefined; - proposedPlansById: Record | undefined; - entry: ThreadPlanCatalogEntry; +interface ThreadPlanCatalogCacheEntry { + shell: object | null; + proposedPlanIds: readonly string[] | undefined; + proposedPlansById: Record | undefined; + entry: ThreadPlanCatalogEntry; +} + +function createThreadPlanCatalogSelector(threadIds: readonly ThreadId[]) { + let previousThreadIds: readonly ThreadId[] = []; + let previousResult: ThreadPlanCatalogEntry[] = []; + let previousEntries = new Map(); + + return (state: ThreadPlanCatalogStoreState): ThreadPlanCatalogEntry[] => { + const sameThreadIds = + previousThreadIds.length === threadIds.length && + previousThreadIds.every((id, index) => id === threadIds[index]); + const nextEntries = new Map(); + const nextResult: ThreadPlanCatalogEntry[] = []; + let changed = !sameThreadIds; + + for (const threadId of threadIds) { + let shell: object | undefined; + let proposedPlanIds: readonly string[] | undefined; + let proposedPlansById: Record | undefined; + + for (const environmentState of Object.values(state.environmentStateById)) { + const matchedShell = environmentState.threadShellById[threadId]; + if (!matchedShell) { + continue; } - >(); - - return (state) => { - const sameThreadIds = - previousThreadIds.length === threadIds.length && - previousThreadIds.every((id, index) => id === threadIds[index]); - const nextEntries = new Map< - ThreadId, - { - shell: object | null; - proposedPlanIds: readonly string[] | undefined; - proposedPlansById: Record | undefined; - entry: ThreadPlanCatalogEntry; - } - >(); - const nextResult: ThreadPlanCatalogEntry[] = []; - let changed = !sameThreadIds; - - for (const threadId of threadIds) { - let shell: object | undefined; - let proposedPlanIds: readonly string[] | undefined; - let proposedPlansById: Record | undefined; - - for (const environmentState of Object.values(state.environmentStateById)) { - const matchedShell = environmentState.threadShellById[threadId]; - if (!matchedShell) { - continue; - } - shell = matchedShell; - proposedPlanIds = environmentState.proposedPlanIdsByThreadId[threadId]; - proposedPlansById = environmentState.proposedPlanByThreadId[threadId] as - | Record - | undefined; - break; - } + shell = matchedShell; + proposedPlanIds = environmentState.proposedPlanIdsByThreadId[threadId]; + proposedPlansById = environmentState.proposedPlanByThreadId[threadId] as + | Record + | undefined; + break; + } - if (!shell) { - const previous = previousEntries.get(threadId); - if ( - previous && - previous.shell === null && - previous.proposedPlanIds === undefined && - previous.proposedPlansById === undefined - ) { - nextEntries.set(threadId, previous); - continue; - } - changed = true; - nextEntries.set(threadId, { - shell: null, - proposedPlanIds: undefined, - proposedPlansById: undefined, - entry: { id: threadId, proposedPlans: EMPTY_PROPOSED_PLANS }, - }); - continue; - } + if (!shell) { + const previous = previousEntries.get(threadId); + if ( + previous && + previous.shell === null && + previous.proposedPlanIds === undefined && + previous.proposedPlansById === undefined + ) { + nextEntries.set(threadId, previous); + continue; + } + changed = true; + nextEntries.set(threadId, { + shell: null, + proposedPlanIds: undefined, + proposedPlansById: undefined, + entry: { id: threadId, proposedPlans: EMPTY_PROPOSED_PLANS }, + }); + continue; + } - const previous = previousEntries.get(threadId); - if ( - previous && - previous.shell === shell && - previous.proposedPlanIds === proposedPlanIds && - previous.proposedPlansById === proposedPlansById - ) { - nextEntries.set(threadId, previous); - nextResult.push(previous.entry); - continue; - } + const previous = previousEntries.get(threadId); + if ( + previous && + previous.shell === shell && + previous.proposedPlanIds === proposedPlanIds && + previous.proposedPlansById === proposedPlansById + ) { + nextEntries.set(threadId, previous); + nextResult.push(previous.entry); + continue; + } - changed = true; - const proposedPlans = - proposedPlanIds && proposedPlanIds.length > 0 && proposedPlansById - ? proposedPlanIds.flatMap((planId) => { - const proposedPlan = proposedPlansById?.[planId]; - return proposedPlan ? [proposedPlan] : []; - }) - : EMPTY_PROPOSED_PLANS; - const entry = { id: threadId, proposedPlans }; - nextEntries.set(threadId, { - shell, - proposedPlanIds, - proposedPlansById, - entry, - }); - nextResult.push(entry); - } + changed = true; + const proposedPlans = + proposedPlanIds && proposedPlanIds.length > 0 && proposedPlansById + ? proposedPlanIds.flatMap((planId) => { + const proposedPlan = proposedPlansById?.[planId]; + return proposedPlan ? [proposedPlan] : []; + }) + : EMPTY_PROPOSED_PLANS; + const entry = { id: threadId, proposedPlans }; + nextEntries.set(threadId, { + shell, + proposedPlanIds, + proposedPlansById, + entry, + }); + nextResult.push(entry); + } - if (!changed && previousResult.length === nextResult.length) { - return previousResult; - } + if (!changed && previousResult.length === nextResult.length) { + return previousResult; + } - previousThreadIds = threadIds; - previousEntries = nextEntries; - previousResult = nextResult; - return nextResult; - }; - }, [threadIds]), + previousThreadIds = threadIds; + previousEntries = nextEntries; + previousResult = nextResult; + return nextResult; + }; +} + +function deriveThreadPlanCatalogThreadIds(input: { + activeThreadId: ThreadId | null; + sourceProposedPlanThreadId: ThreadId | null; +}): readonly ThreadId[] { + const threadIds: ThreadId[] = []; + if (input.activeThreadId) { + threadIds.push(input.activeThreadId); + } + if ( + input.sourceProposedPlanThreadId && + input.sourceProposedPlanThreadId !== input.activeThreadId + ) { + threadIds.push(input.sourceProposedPlanThreadId); + } + return threadIds; +} + +function useThreadPlanCatalog(threadIds: readonly ThreadId[]): ThreadPlanCatalogEntry[] { + const selectThreadPlanCatalog = useMemo( + () => createThreadPlanCatalogSelector(threadIds), + [threadIds], ); + return useStore(selectThreadPlanCatalog); } function formatOutgoingPrompt(params: { @@ -1028,19 +1041,16 @@ export default function ChatView(props: ChatViewProps) { return openTerminalThreadKeys.filter((nextThreadKey) => existingThreadKeys.has(nextThreadKey)); }, [draftThreadKeys, openTerminalThreadKeys, serverThreadKeys]); const activeLatestTurn = activeThread?.latestTurn ?? null; - const threadPlanCatalog = useThreadPlanCatalog( - useMemo(() => { - const threadIds: ThreadId[] = []; - if (activeThread?.id) { - threadIds.push(activeThread.id); - } - const sourceThreadId = activeLatestTurn?.sourceProposedPlan?.threadId; - if (sourceThreadId && sourceThreadId !== activeThread?.id) { - threadIds.push(sourceThreadId); - } - return threadIds; - }, [activeLatestTurn?.sourceProposedPlan?.threadId, activeThread?.id]), + const sourceProposedPlanThreadId = activeLatestTurn?.sourceProposedPlan?.threadId ?? null; + const threadPlanCatalogThreadIds = useMemo( + () => + deriveThreadPlanCatalogThreadIds({ + activeThreadId, + sourceProposedPlanThreadId, + }), + [activeThreadId, sourceProposedPlanThreadId], ); + const threadPlanCatalog = useThreadPlanCatalog(threadPlanCatalogThreadIds); useEffect(() => { setMountedTerminalThreadKeys((currentThreadIds) => { const nextThreadIds = reconcileMountedTerminalThreadIds({