diff --git a/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx b/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx index b19721b91c..939ca44878 100644 --- a/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx +++ b/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx @@ -67,6 +67,8 @@ function renderUseResumeOnLoad({ messages = [], getMessages: getMessagesOverride, submission = null, + isSubmitting = false, + attachedGenerationCreatedAt = null, conversationId = CONVERSATION_ID, messagesLoaded = true, onSubmission, @@ -81,6 +83,8 @@ function renderUseResumeOnLoad({ messages?: TMessage[]; getMessages?: () => TMessage[] | undefined; submission?: TSubmission | null; + isSubmitting?: boolean; + attachedGenerationCreatedAt?: number | null; conversationId?: string; messagesLoaded?: boolean; onSubmission?: (submission: TSubmission | null) => void; @@ -100,6 +104,11 @@ function renderUseResumeOnLoad({ const initializeState = (snapshot: MutableSnapshot) => { snapshot.set(store.conversationByIndex(0), buildConversation(conversationId)); snapshot.set(store.submissionByIndex(0), submission); + snapshot.set(store.isSubmittingFamily(0), isSubmitting); + snapshot.set( + store.activeGenerationCreatedAtByConvoId(conversationId), + attachedGenerationCreatedAt, + ); if (submissionStart != null) { snapshot.set(store.submissionStartFamily(0), submissionStart); } @@ -644,6 +653,105 @@ describe('useResumeOnLoad', () => { expect(mockUseStreamStatus).toHaveBeenCalledWith(CONVERSATION_ID, true); }); + it('attaches to a continuation when the finished run left its submission installed', async () => { + const observedSubmissions: Array = []; + mockUseStreamStatus.mockReturnValue(INACTIVE_STATUS); + + /** + * The FINAL path never clears the submission atom, so a pane that just + * finished a run still holds it. When a background tool dispatch (or any + * server-side continuation) then starts the next generation, that stale + * bookkeeping is the only thing standing between the announcement and the + * attachment — and it is not an attachment: nothing is streaming. + */ + const { rerender } = renderUseResumeOnLoad({ + submission: buildSubmission(CONVERSATION_ID), + isSubmitting: false, + messages: [buildUserMessage(CONVERSATION_ID)], + onSubmission: (currentSubmission) => observedSubmissions.push(currentSubmission), + }); + await act(async () => { + await Promise.resolve(); + }); + + mockUseStreamStatus.mockClear(); + mockUseActiveJobs.mockReturnValue({ + data: { activeJobIds: [CONVERSATION_ID] }, + dataUpdatedAt: 2, + }); + mockUseStreamStatus.mockReturnValue(ACTIVE_STATUS); + rerender(); + await act(async () => { + await Promise.resolve(); + }); + + expect(mockUseStreamStatus).toHaveBeenCalledWith(CONVERSATION_ID, true); + const attached = observedSubmissions[observedSubmissions.length - 1] as + | (TSubmission & { resumeStreamId?: string }) + | null; + expect(attached?.resumeStreamId).toBe(CONVERSATION_ID); + }); + + it('leaves a resumed attachment alone before its stream has opened', async () => { + const observedSubmissions: Array = []; + mockUseStreamStatus.mockReturnValue(INACTIVE_STATUS); + + /** + * A resume installs its submission and only reads as submitting once the + * stream opens, so between those two points a submission exists with + * `isSubmitting` still false. An announcement landing there must not + * mistake it for the leftovers of a finished run and tear down the + * attachment this hook just built. The epoch is stamped first for exactly + * this reason. + */ + const { rerender } = renderUseResumeOnLoad({ + submission: buildSubmission(CONVERSATION_ID), + isSubmitting: false, + attachedGenerationCreatedAt: 4242, + messages: [buildUserMessage(CONVERSATION_ID)], + onSubmission: (currentSubmission) => observedSubmissions.push(currentSubmission), + }); + await act(async () => { + await Promise.resolve(); + }); + + mockUseActiveJobs.mockReturnValue({ + data: { activeJobIds: [CONVERSATION_ID] }, + dataUpdatedAt: 2, + }); + rerender(); + await act(async () => { + await Promise.resolve(); + }); + + expect(observedSubmissions[observedSubmissions.length - 1]).not.toBeNull(); + }); + + it('leaves a streaming attachment alone when its own run is announced', async () => { + mockUseStreamStatus.mockReturnValue(INACTIVE_STATUS); + const { rerender } = renderUseResumeOnLoad({ + submission: buildSubmission(CONVERSATION_ID), + isSubmitting: true, + messages: [buildUserMessage(CONVERSATION_ID)], + }); + await act(async () => { + await Promise.resolve(); + }); + + mockUseStreamStatus.mockClear(); + /** This pane's own run is what the list is reporting. */ + mockUseActiveJobs.mockReturnValue({ + data: { activeJobIds: [CONVERSATION_ID] }, + dataUpdatedAt: 2, + }); + rerender(); + await act(async () => { + await Promise.resolve(); + }); + + expect(mockUseStreamStatus).not.toHaveBeenCalledWith(CONVERSATION_ID, true); + }); + it('does not re-arm for a conversation that is not the one being viewed', async () => { const observedSubmissions: Array = []; mockUseStreamStatus.mockReturnValue(INACTIVE_STATUS); diff --git a/client/src/hooks/SSE/useResumeOnLoad.ts b/client/src/hooks/SSE/useResumeOnLoad.ts index 1727758b24..904833e9d8 100644 --- a/client/src/hooks/SSE/useResumeOnLoad.ts +++ b/client/src/hooks/SSE/useResumeOnLoad.ts @@ -250,6 +250,10 @@ export default function useResumeOnLoad( const setSubmission = useSetRecoilState(store.submissionByIndex(runIndex)); const setSubmissionStart = useSetRecoilState(store.submissionStartFamily(runIndex)); const currentSubmission = useRecoilValue(store.submissionByIndex(runIndex)); + const isSubmitting = useRecoilValue(store.isSubmittingFamily(runIndex)); + const attachedGenerationCreatedAt = useRecoilValue( + store.activeGenerationCreatedAtByConvoId(conversationId ?? ''), + ); const currentConversation = useRecoilValue(store.conversationByIndex(runIndex)); const endpoint = currentConversation?.endpoint; const endpointType = currentConversation?.endpointType; @@ -425,6 +429,23 @@ export default function useResumeOnLoad( !!currentSubmission && (hasExplicitSubmissionMatch || hasHydratedMessageMatch); const hasStaleSubmissionForDifferentConvo = !!currentSubmission && submissionConvoId != null && submissionConvoId !== conversationId; + /** + * A submission only stands in for an attachment while one is actually live. + * The FINAL path never clears the atom, so a pane that just finished a run + * keeps holding the submission it completed — bookkeeping, not a stream. + * + * Two signals, because neither covers the whole lifetime on its own. + * `isSubmitting` is raised by `ask` before a submission is installed and held + * across reconnect backoff, so sends and recovery both read as attached. It + * is not raised for a resumed attachment until that stream opens, though, and + * this effect installs the submission well before then — so an announcement + * landing in that window would tear down the attachment it had just built. + * The generation epoch closes it: this hook stamps it immediately before + * installing a resume submission, `subscribeToStream` stamps it for every + * other attachment, and terminal teardown is what clears it. + */ + const hasLiveSubmissionForThisConvo = + hasActiveSubmissionForThisConvo && (isSubmitting || attachedGenerationCreatedAt != null); /** * A run this pane did not start — another tab, another device, a scheduled @@ -745,7 +766,7 @@ export default function useResumeOnLoad( !resumableEnabled || !hasActiveJobForThisConvo || !conversationId || - hasActiveSubmissionForThisConvo + hasLiveSubmissionForThisConvo ) { if (!hasActiveJobForThisConvo) { answeredActiveJobRef.current = null; @@ -765,6 +786,17 @@ export default function useResumeOnLoad( answeredActiveJobRef.current = { conversationId, at: now }; console.log('[ResumeOnLoad] Active job announced without an attachment', { conversationId }); + /** + * A finished submission still installed here is the one thing standing + * between this announcement and the attachment: it is what the check below + * reads as "already attached", and what disables the status query. The run + * it describes is over, so release it — a server-side continuation such as + * a background tool dispatch finishing is a different generation, and the + * resume path is what owns building the submission for it. + */ + if (hasActiveSubmissionForThisConvo) { + setSubmission(null); + } queryClient.invalidateQueries({ queryKey: streamStatusQueryKey(conversationId) }); queryClient.invalidateQueries({ queryKey: [QueryKeys.messages, conversationId] }); processedConvoRef.current = null; @@ -774,7 +806,10 @@ export default function useResumeOnLoad( resumableEnabled, hasActiveJobForThisConvo, hasActiveSubmissionForThisConvo, + hasLiveSubmissionForThisConvo, + attachedGenerationCreatedAt, activeJobsUpdatedAt, + setSubmission, queryClient, ]); }