From e78628caecb95f34eece5747931520e7900e8e3c Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Mon, 27 Apr 2026 11:11:22 +0000 Subject: [PATCH] feat(ui): implement turn-boundary snapshot refetch and bind progress UI - Added `patchSnapshot` to `CoreStateProvider` for partial global state updates. - Created `useRefetchSnapshotOnTurnEnd` hook with 750ms debounce to refetch authoritative user state. - Wired refetch into `ChatRuntimeProvider` on turn completion. - Refactored `ContextGatheringStep` to read progress from global snapshot. - Added Vitest unit tests for the new hook and its debounce logic. - Updated documentation in `app/src/providers/README.md`. Closes #924 Co-authored-by: oxoxDev <164490987+oxoxDev@users.noreply.github.com> --- .../useRefetchSnapshotOnTurnEnd.test.ts | 104 ++++++++++++++++++ app/src/hooks/useRefetchSnapshotOnTurnEnd.ts | 37 +++++++ .../onboarding/steps/ContextGatheringStep.tsx | 39 +++++-- app/src/providers/ChatRuntimeProvider.tsx | 6 +- app/src/providers/CoreStateProvider.tsx | 13 +++ app/src/providers/README.md | 24 ++++ app/src/types/api.ts | 5 + 7 files changed, 215 insertions(+), 13 deletions(-) create mode 100644 app/src/hooks/__tests__/useRefetchSnapshotOnTurnEnd.test.ts create mode 100644 app/src/hooks/useRefetchSnapshotOnTurnEnd.ts create mode 100644 app/src/providers/README.md diff --git a/app/src/hooks/__tests__/useRefetchSnapshotOnTurnEnd.test.ts b/app/src/hooks/__tests__/useRefetchSnapshotOnTurnEnd.test.ts new file mode 100644 index 0000000000..f330c35efb --- /dev/null +++ b/app/src/hooks/__tests__/useRefetchSnapshotOnTurnEnd.test.ts @@ -0,0 +1,104 @@ +import { act, renderHook } from '@testing-library/react'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import { useCoreState } from '../../providers/CoreStateProvider'; +import { userApi } from '../../services/api/userApi'; +import { useRefetchSnapshotOnTurnEnd } from '../useRefetchSnapshotOnTurnEnd'; + +vi.mock('../../providers/CoreStateProvider', () => ({ + useCoreState: vi.fn(), +})); + +vi.mock('../../services/api/userApi', () => ({ + userApi: { + getMe: vi.fn(), + }, +})); + +describe('useRefetchSnapshotOnTurnEnd', () => { + const mockPatchSnapshot = vi.fn(); + + beforeEach(() => { + vi.useFakeTimers(); + vi.mocked(useCoreState).mockReturnValue({ + patchSnapshot: mockPatchSnapshot, + } as any); + }); + + afterEach(() => { + vi.restoreAllMocks(); + vi.useRealTimers(); + }); + + it('refetches and patches snapshot after 750ms', async () => { + const mockUser1 = { _id: 'user1', firstName: 'Jules' }; + vi.mocked(userApi.getMe).mockResolvedValue(mockUser1 as any); + + const { result } = renderHook(() => useRefetchSnapshotOnTurnEnd()); + + act(() => { + result.current.refetch(); + }); + + expect(userApi.getMe).not.toHaveBeenCalled(); + + await act(async () => { + vi.advanceTimersByTime(750); + }); + + expect(userApi.getMe).toHaveBeenCalledTimes(1); + expect(mockPatchSnapshot).toHaveBeenCalledWith({ currentUser: mockUser1 }); + }); + + it('three rapid finalize events → one getMe call', async () => { + const mockUser1 = { _id: 'user1', firstName: 'Jules' }; + vi.mocked(userApi.getMe).mockResolvedValue(mockUser1 as any); + + const { result } = renderHook(() => useRefetchSnapshotOnTurnEnd()); + + act(() => { + result.current.refetch(); + vi.advanceTimersByTime(300); + result.current.refetch(); + vi.advanceTimersByTime(300); + result.current.refetch(); + }); + + expect(userApi.getMe).not.toHaveBeenCalled(); + + await act(async () => { + vi.advanceTimersByTime(750); + }); + + expect(userApi.getMe).toHaveBeenCalledTimes(1); + }); + + it('mock userApi.getMe() to return two payloads in sequence → UI re-reads second value', async () => { + const mockUser1 = { _id: 'user1', firstName: 'Jules', onboarding_status: { finished: false } }; + const mockUser2 = { _id: 'user1', firstName: 'Jules', onboarding_status: { finished: true } }; + + vi.mocked(userApi.getMe) + .mockResolvedValueOnce(mockUser1 as any) + .mockResolvedValueOnce(mockUser2 as any); + + const { result } = renderHook(() => useRefetchSnapshotOnTurnEnd()); + + // First refetch + act(() => { + result.current.refetch(); + }); + await act(async () => { + vi.advanceTimersByTime(750); + }); + expect(mockPatchSnapshot).toHaveBeenLastCalledWith({ currentUser: mockUser1 }); + + // Second refetch + act(() => { + result.current.refetch(); + }); + await act(async () => { + vi.advanceTimersByTime(750); + }); + expect(mockPatchSnapshot).toHaveBeenLastCalledWith({ currentUser: mockUser2 }); + }); +}); diff --git a/app/src/hooks/useRefetchSnapshotOnTurnEnd.ts b/app/src/hooks/useRefetchSnapshotOnTurnEnd.ts new file mode 100644 index 0000000000..44ecf1ce55 --- /dev/null +++ b/app/src/hooks/useRefetchSnapshotOnTurnEnd.ts @@ -0,0 +1,37 @@ +import { useCallback, useRef } from 'react'; + +import { useCoreState } from '../providers/CoreStateProvider'; +import { userApi } from '../services/api/userApi'; + +/** + * Hook to refetch the authoritative user state from the backend after a chat + * turn finishes. Updates the global snapshot in CoreStateProvider. + * + * Includes a 750ms debounce to collapse multiple rapid turn-finalized events. + */ +export function useRefetchSnapshotOnTurnEnd() { + const { patchSnapshot } = useCoreState(); + const debounceTimerRef = useRef(null); + + const refetch = useCallback(() => { + if (debounceTimerRef.current !== null) { + window.clearTimeout(debounceTimerRef.current); + } + + debounceTimerRef.current = window.setTimeout(() => { + debounceTimerRef.current = null; + + // Fire-and-forget on a microtask + void (async () => { + try { + const user = await userApi.getMe(); + patchSnapshot({ currentUser: user }); + } catch (error) { + console.warn('[useRefetchSnapshotOnTurnEnd] failed to refetch user state:', error); + } + })(); + }, 750); + }, [patchSnapshot]); + + return { refetch }; +} diff --git a/app/src/pages/onboarding/steps/ContextGatheringStep.tsx b/app/src/pages/onboarding/steps/ContextGatheringStep.tsx index c20f9483e2..ce97bbaf65 100644 --- a/app/src/pages/onboarding/steps/ContextGatheringStep.tsx +++ b/app/src/pages/onboarding/steps/ContextGatheringStep.tsx @@ -13,10 +13,11 @@ * External calls still go through core (auth, proxy, billing). Only the * stage-by-stage orchestration lives in the renderer. */ -import { useRef, useState } from 'react'; +import { useMemo, useRef, useState } from 'react'; import Button from '../../../components/ui/Button'; import WhatLeavesLink from '../../../features/privacy/WhatLeavesLink'; +import { useCoreState } from '../../../providers/CoreStateProvider'; import { callCoreRpc } from '../../../services/coreRpcClient'; import OnboardingNextButton from '../components/OnboardingNextButton'; @@ -148,22 +149,36 @@ const ContextGatheringStep = ({ onNext, onBack: _onBack, }: ContextGatheringStepProps) => { - const [stageStatuses, setStageStatuses] = useState>(() => { + const { snapshot } = useCoreState(); + const [localStageStatuses, setLocalStageStatuses] = useState>(() => { const initial: Record = {}; for (const s of STAGES) initial[s.id] = 'pending'; return initial; }); - const [stageDetails, setStageDetails] = useState>({}); - const [finished, setFinished] = useState(false); + const [localStageDetails, setLocalStageDetails] = useState>({}); + const [localFinished, setLocalFinished] = useState(false); const [started, setStarted] = useState(false); const [error, setError] = useState(null); const ranRef = useRef(false); const hasGmail = connectedSources.some(s => s.includes('gmail')); + const stageStatuses = useMemo( + () => snapshot.currentUser?.onboarding_status?.stages ?? localStageStatuses, + [snapshot.currentUser?.onboarding_status?.stages, localStageStatuses] + ); + const stageDetails = useMemo( + () => snapshot.currentUser?.onboarding_status?.details ?? localStageDetails, + [snapshot.currentUser?.onboarding_status?.details, localStageDetails] + ); + const finished = useMemo( + () => snapshot.currentUser?.onboarding_status?.finished ?? localFinished, + [snapshot.currentUser?.onboarding_status?.finished, localFinished] + ); + const setStage = (id: Stage['id'], status: StageStatus, detail?: string) => { - setStageStatuses(prev => ({ ...prev, [id]: status })); - if (detail !== undefined) setStageDetails(prev => ({ ...prev, [id]: detail })); + setLocalStageStatuses(prev => ({ ...prev, [id]: status })); + if (detail !== undefined) setLocalStageDetails(prev => ({ ...prev, [id]: detail })); }; const handleStart = () => { @@ -174,9 +189,9 @@ const ContextGatheringStep = ({ if (!hasGmail) { const skipped: Record = {}; for (const s of STAGES) skipped[s.id] = 'skipped'; - setStageStatuses(skipped); - setStageDetails({ 'gmail-search': 'Gmail not connected' }); - setFinished(true); + setLocalStageStatuses(skipped); + setLocalStageDetails({ 'gmail-search': 'Gmail not connected' }); + setLocalFinished(true); return; } @@ -197,7 +212,7 @@ const ContextGatheringStep = ({ setStage('gmail-search', 'skipped', 'No LinkedIn URL found in mailbox'); setStage('linkedin-scrape', 'skipped'); setStage('build-profile', 'skipped'); - setFinished(true); + setLocalFinished(true); return; } } catch (e) { @@ -205,7 +220,7 @@ const ContextGatheringStep = ({ setStage('gmail-search', 'error', e instanceof Error ? e.message : String(e)); setStage('linkedin-scrape', 'skipped'); setStage('build-profile', 'skipped'); - setFinished(true); + setLocalFinished(true); return; } @@ -238,7 +253,7 @@ const ContextGatheringStep = ({ setStage('build-profile', 'error', e instanceof Error ? e.message : String(e)); } - setFinished(true); + setLocalFinished(true); } const completedCount = STAGES.filter(s => { diff --git a/app/src/providers/ChatRuntimeProvider.tsx b/app/src/providers/ChatRuntimeProvider.tsx index 2f8b7db901..1e197c7817 100644 --- a/app/src/providers/ChatRuntimeProvider.tsx +++ b/app/src/providers/ChatRuntimeProvider.tsx @@ -1,6 +1,7 @@ import debug from 'debug'; import { useCallback, useEffect, useRef } from 'react'; +import { useRefetchSnapshotOnTurnEnd } from '../hooks/useRefetchSnapshotOnTurnEnd'; import { requestUsageRefresh } from '../hooks/usageRefresh'; import { type ChatInferenceStartEvent, @@ -56,6 +57,7 @@ function rtLog(message: string, fields?: Record { const dispatch = useAppDispatch(); + const { refetch: refetchSnapshot } = useRefetchSnapshotOnTurnEnd(); const socketStatus = useAppSelector(selectSocketStatus); const toolTimelineByThread = useAppSelector(state => state.chatRuntime.toolTimelineByThread); const inferenceStatusByThread = useAppSelector( @@ -556,6 +558,7 @@ const ChatRuntimeProvider = ({ children }: { children: React.ReactNode }) => { reason: 'chat_done', }); requestUsageRefresh(); + refetchSnapshot(); dispatch(endInferenceTurn({ threadId: event.thread_id })); dispatch(setActiveThread(null)); })(); @@ -574,6 +577,7 @@ const ChatRuntimeProvider = ({ children }: { children: React.ReactNode }) => { reason: 'chat_done', }); requestUsageRefresh(); + refetchSnapshot(); dispatch(endInferenceTurn({ threadId: event.thread_id })); dispatch(setActiveThread(null)); }, @@ -636,7 +640,7 @@ const ChatRuntimeProvider = ({ children }: { children: React.ReactNode }) => { rtLog('unsubscribe_chat_events'); cleanup(); }; - }, [dispatch, resolveVisibleThreadForProactive, socketStatus]); + }, [dispatch, resolveVisibleThreadForProactive, socketStatus, refetchSnapshot]); return <>{children}; }; diff --git a/app/src/providers/CoreStateProvider.tsx b/app/src/providers/CoreStateProvider.tsx index 7055f0c6ec..bffc1d2afe 100644 --- a/app/src/providers/CoreStateProvider.tsx +++ b/app/src/providers/CoreStateProvider.tsx @@ -64,6 +64,7 @@ interface CoreStateContextValue extends CoreState { setOnboardingCompletedFlag: (value: boolean) => Promise; setEncryptionKey: (value: string | null) => Promise; setPrimaryWalletAddress: (value: string | null) => Promise; + patchSnapshot: (patch: Partial) => void; setOnboardingTasks: (value: CoreOnboardingTasks | null) => Promise; storeSessionToken: (token: string, user?: object) => Promise; clearSession: () => Promise; @@ -424,6 +425,16 @@ export default function CoreStateProvider({ children }: { children: ReactNode }) }); }, [commitState, refresh]); + const patchSnapshot = useCallback( + (patch: Partial) => { + commitState(previous => ({ + ...previous, + snapshot: { ...previous.snapshot, ...patch }, + })); + }, + [commitState] + ); + const value = useMemo( () => ({ ...state, @@ -431,6 +442,7 @@ export default function CoreStateProvider({ children }: { children: ReactNode }) refreshTeams, refreshTeamMembers, refreshTeamInvites, + patchSnapshot, setAnalyticsEnabled, setOnboardingCompletedFlag, setEncryptionKey: value => updateLocalState({ encryptionKey: value }), @@ -445,6 +457,7 @@ export default function CoreStateProvider({ children }: { children: ReactNode }) refreshTeamInvites, refreshTeamMembers, refreshTeams, + patchSnapshot, setAnalyticsEnabled, setOnboardingCompletedFlag, state, diff --git a/app/src/providers/README.md b/app/src/providers/README.md new file mode 100644 index 0000000000..412550591f --- /dev/null +++ b/app/src/providers/README.md @@ -0,0 +1,24 @@ +# App Providers + +This directory contains the React context providers that manage the global state and services of the application. + +## CoreStateProvider + +Manages the authoritative global state of the application, including user authentication, session tokens, and the application snapshot. + +### Turn-Boundary Refetch Contract + +To ensure that the UI stays in sync with the backend state (especially during onboarding and context gathering), the application follows a refetch-on-turn-end contract: + +- **Refetch Timing**: After every agent reply completes (the `chat_done` event in `ChatRuntimeProvider`), the application refetches the authoritative user state via `userApi.getMe()`. +- **Debounce**: Multiple rapid turn-finalized events within 750ms are collapsed into a single refetch call to avoid unnecessary network traffic. +- **Single Source of Truth**: The refetched state is merged into the global snapshot using `patchSnapshot`. Components should bind to this global snapshot to ensure they reflect the latest backend state without requiring a full remount. +- **Fire-and-Forget**: The refetch operation is non-blocking and fires on a microtask after the chat UI has painted the final response. + +## ChatRuntimeProvider + +Manages the live chat state, including message streaming, tool execution timeline, and subagent orchestration. It subscribes to socket events and updates the Redux store. + +## SocketProvider + +Manages the Socket.IO connection to the Rust core, providing the underlying transport for real-time chat events. diff --git a/app/src/types/api.ts b/app/src/types/api.ts index 26134ed669..f3ba6ad479 100644 --- a/app/src/types/api.ts +++ b/app/src/types/api.ts @@ -63,6 +63,11 @@ export interface User { languageCode?: string; waitlist?: string; activeTeamId: string; + onboarding_status?: { + stages: Record; + details: Record; + finished: boolean; + }; } // Billing types