From c3db6d522d7d118f1d14c3470353f5480d8b307d Mon Sep 17 00:00:00 2001 From: Arpit Gupta Date: Tue, 2 Jun 2026 22:57:59 +0530 Subject: [PATCH 1/2] feat(chat): resume in-flight workflow stream on chat mount MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reopening a chat mid-generation showed a frozen partial turn — the client fetched persisted messages but never reattached to the live run. The GET chat snapshot already reports isStreaming; surface it and reconnect via the AI SDK resume path to GET /api/chat/{chatId}/stream. - getChatSnapshot replaces getChatMessages, returning { messages, isStreaming } - useMessageLoader fires onActiveStream once when the loaded chat is streaming - workflow transport gains prepareReconnectToStreamRequest → GET stream route - useVercelChat resumes via useChat's resumeStream() on active-stream signal Co-Authored-By: Claude Opus 4.8 --- hooks/useChatTransport.ts | 6 ++ hooks/useMessageLoader.ts | 17 +++-- hooks/useVercelChat.ts | 10 ++- .../__tests__/getChatSnapshot.test.ts | 71 +++++++++++++++++++ lib/messages/getChatMessages.ts | 26 ------- lib/messages/getChatSnapshot.ts | 40 +++++++++++ 6 files changed, 138 insertions(+), 32 deletions(-) create mode 100644 lib/messages/__tests__/getChatSnapshot.test.ts delete mode 100644 lib/messages/getChatMessages.ts create mode 100644 lib/messages/getChatSnapshot.ts diff --git a/hooks/useChatTransport.ts b/hooks/useChatTransport.ts index 5f68b0fca..d2cb67ade 100644 --- a/hooks/useChatTransport.ts +++ b/hooks/useChatTransport.ts @@ -74,6 +74,12 @@ export function useChatTransport({ if (recoupAccessToken) body.recoupAccessToken = recoupAccessToken; return body; }, + // Resume/reconnect targets recoup-api's dedicated GET stream route + // (POST /api/chat/workflow never resumes). The transport `headers` + // above carry the Privy bearer on this request too. + prepareReconnectToStreamRequest: ({ id }) => ({ + api: `${baseUrl}/api/chat/${encodeURIComponent(id)}/stream`, + }), }); }, [baseUrl, chatId, sessionId, getAccessToken]); diff --git a/hooks/useMessageLoader.ts b/hooks/useMessageLoader.ts index 1bd0236e8..dc7a2db0e 100644 --- a/hooks/useMessageLoader.ts +++ b/hooks/useMessageLoader.ts @@ -1,7 +1,7 @@ import { useState, useEffect } from "react"; import { UIMessage } from "ai"; import { usePrivy } from "@privy-io/react-auth"; -import { getChatMessages } from "@/lib/messages/getChatMessages"; +import { getChatSnapshot } from "@/lib/messages/getChatSnapshot"; /** * Loads the persisted UI message stream for a session-scoped chat from @@ -9,12 +9,16 @@ import { getChatMessages } from "@/lib/messages/getChatMessages"; * bootstrap path mounts `` before a session has been minted, and * the in-transition legacy `/chat/{roomId}` route lacks a sessionId * until track I redirects it. + * + * When the loaded chat still has a workflow run in flight, `onActiveStream` + * fires once so the caller can reconnect to the live stream. */ export function useMessageLoader( sessionId: string | undefined, chatId: string | undefined, userId: string | undefined, setMessages: (messages: UIMessage[]) => void, + onActiveStream?: () => void, ) { const { getAccessToken } = usePrivy(); const [isLoading, setIsLoading] = useState(!!(sessionId && chatId)); @@ -39,13 +43,16 @@ export function useMessageLoader( const accessToken = await getAccessToken(); if (!accessToken) return; - const initialMessages = await getChatMessages( + const { messages, isStreaming } = await getChatSnapshot( sessionId, chatId, accessToken, ); - if (initialMessages.length > 0) { - setMessages(initialMessages); + if (messages.length > 0) { + setMessages(messages); + } + if (isStreaming) { + onActiveStream?.(); } } catch (err) { console.error("Error loading messages:", err); @@ -58,7 +65,7 @@ export function useMessageLoader( }; loadMessages(); - }, [userId, sessionId, chatId, getAccessToken, setMessages]); + }, [userId, sessionId, chatId, getAccessToken, setMessages, onActiveStream]); return { isLoading, diff --git a/hooks/useVercelChat.ts b/hooks/useVercelChat.ts index 20d5d5713..0beab0207 100644 --- a/hooks/useVercelChat.ts +++ b/hooks/useVercelChat.ts @@ -184,7 +184,7 @@ export function useVercelChat({ [id, artistId, organizationId, accountIdOverride, model], ); - const { messages, status, stop, sendMessage, setMessages, regenerate } = + const { messages, status, stop, sendMessage, setMessages, regenerate, resumeStream } = useChat({ id, transport, @@ -282,11 +282,19 @@ export function useVercelChat({ // Keep messagesRef in sync with messages messagesLengthRef.current = messages.length; + // Reconnect to an in-flight workflow run when the loaded chat is still + // streaming (e.g. reopened mid-generation). Routes through the + // transport's prepareReconnectToStreamRequest → GET /api/chat/{id}/stream. + const resumeActiveStream = useCallback(() => { + void resumeStream(); + }, [resumeStream]); + const { isLoading: isMessagesLoading, hasError } = useMessageLoader( sessionId, messages.length === 0 ? id : undefined, userId, setMessages, + resumeActiveStream, ); // Only show loading state if: diff --git a/lib/messages/__tests__/getChatSnapshot.test.ts b/lib/messages/__tests__/getChatSnapshot.test.ts new file mode 100644 index 000000000..12a3ea81a --- /dev/null +++ b/lib/messages/__tests__/getChatSnapshot.test.ts @@ -0,0 +1,71 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { getChatSnapshot } from "../getChatSnapshot"; +import { NEW_API_BASE_URL } from "../../consts"; + +const mockFetch = vi.fn(); +global.fetch = mockFetch; + +describe("getChatSnapshot", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("fetches the session-scoped chat with a bearer token", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => ({ messages: [], isStreaming: false }), + }); + + await getChatSnapshot("sess-1", "chat-1", "tok-123"); + + expect(mockFetch).toHaveBeenCalledWith( + `${NEW_API_BASE_URL}/api/sessions/sess-1/chats/chat-1`, + { headers: { Authorization: "Bearer tok-123" } }, + ); + }); + + it("returns messages and isStreaming from the response", async () => { + const messages = [{ id: "m1", role: "user", parts: [] }]; + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => ({ messages, isStreaming: true }), + }); + + const result = await getChatSnapshot("sess-1", "chat-1", "tok"); + + expect(result).toEqual({ messages, isStreaming: true }); + }); + + it("defaults isStreaming to false when absent", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => ({ messages: [] }), + }); + + const result = await getChatSnapshot("sess-1", "chat-1", "tok"); + + expect(result).toEqual({ messages: [], isStreaming: false }); + }); + + it("returns an empty, non-streaming snapshot on a non-ok response", async () => { + mockFetch.mockResolvedValueOnce({ ok: false, json: async () => ({}) }); + + const result = await getChatSnapshot("sess-1", "chat-1", "tok"); + + expect(result).toEqual({ messages: [], isStreaming: false }); + }); + + it("url-encodes both ids", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => ({ messages: [], isStreaming: false }), + }); + + await getChatSnapshot("a/b", "c d", "tok"); + + expect(mockFetch).toHaveBeenCalledWith( + `${NEW_API_BASE_URL}/api/sessions/a%2Fb/chats/c%20d`, + expect.any(Object), + ); + }); +}); diff --git a/lib/messages/getChatMessages.ts b/lib/messages/getChatMessages.ts deleted file mode 100644 index 8b0775599..000000000 --- a/lib/messages/getChatMessages.ts +++ /dev/null @@ -1,26 +0,0 @@ -import { UIMessage } from "ai"; -import { getClientApiBaseUrl } from "@/lib/api/getClientApiBaseUrl"; - -/** - * Fetch the persisted UI message stream for a session-scoped chat from - * recoup-api `GET /api/sessions/{sessionId}/chats/{chatId}`. Each row's - * `parts` blob is the full UIMessage, so the response can be handed - * straight to the Vercel AI SDK as initial messages. - */ -export async function getChatMessages( - sessionId: string, - chatId: string, - accessToken: string, -): Promise { - const response = await fetch( - `${getClientApiBaseUrl()}/api/sessions/${encodeURIComponent(sessionId)}/chats/${encodeURIComponent(chatId)}`, - { - headers: { Authorization: `Bearer ${accessToken}` }, - }, - ); - - if (!response.ok) return []; - - const data = (await response.json()) as { messages?: UIMessage[] }; - return data.messages ?? []; -} diff --git a/lib/messages/getChatSnapshot.ts b/lib/messages/getChatSnapshot.ts new file mode 100644 index 000000000..d4aa71d8b --- /dev/null +++ b/lib/messages/getChatSnapshot.ts @@ -0,0 +1,40 @@ +import { UIMessage } from "ai"; +import { getClientApiBaseUrl } from "@/lib/api/getClientApiBaseUrl"; + +export interface ChatSnapshot { + /** Persisted UI messages; each row's `parts` blob is a full UIMessage. */ + messages: UIMessage[]; + /** True when a workflow run is still in flight and can be resumed. */ + isStreaming: boolean; +} + +/** + * Fetch a session-scoped chat snapshot from recoup-api + * `GET /api/sessions/{sessionId}/chats/{chatId}`: the persisted messages + * plus whether a workflow run is still streaming (so the caller can + * reconnect via `GET /api/chat/{chatId}/stream`). + */ +export async function getChatSnapshot( + sessionId: string, + chatId: string, + accessToken: string, +): Promise { + const response = await fetch( + `${getClientApiBaseUrl()}/api/sessions/${encodeURIComponent(sessionId)}/chats/${encodeURIComponent(chatId)}`, + { + headers: { Authorization: `Bearer ${accessToken}` }, + }, + ); + + if (!response.ok) return { messages: [], isStreaming: false }; + + const data = (await response.json()) as { + messages?: UIMessage[]; + isStreaming?: boolean; + }; + + return { + messages: data.messages ?? [], + isStreaming: data.isStreaming ?? false, + }; +} From 6aa30a8c4be392314ed01f9a26e2fa27fbeddc0d Mon Sep 17 00:00:00 2001 From: Arpit Gupta Date: Tue, 2 Jun 2026 23:35:14 +0530 Subject: [PATCH 2/2] fix(chat): throw on !ok in getChatSnapshot to match lib/ pattern Inherited silent-empty-on-error behavior from the deleted getChatMessages helper, but every other fetch helper under lib/ throws with the server's error text. Silent swallowing hid 404/5xx behind a blank-chat render with no UI signal. useMessageLoader already wraps the call in try/catch and sets hasError, so a thrown error is what the caller expects. Co-Authored-By: Claude Opus 4.8 --- .../__tests__/getChatSnapshot.test.ts | 25 ++++++++++++++++--- lib/messages/getChatSnapshot.ts | 7 +++++- 2 files changed, 27 insertions(+), 5 deletions(-) diff --git a/lib/messages/__tests__/getChatSnapshot.test.ts b/lib/messages/__tests__/getChatSnapshot.test.ts index 12a3ea81a..e30c18e8e 100644 --- a/lib/messages/__tests__/getChatSnapshot.test.ts +++ b/lib/messages/__tests__/getChatSnapshot.test.ts @@ -47,12 +47,29 @@ describe("getChatSnapshot", () => { expect(result).toEqual({ messages: [], isStreaming: false }); }); - it("returns an empty, non-streaming snapshot on a non-ok response", async () => { - mockFetch.mockResolvedValueOnce({ ok: false, json: async () => ({}) }); + it("throws with the server's error text on a non-ok response", async () => { + mockFetch.mockResolvedValueOnce({ + ok: false, + status: 404, + text: async () => "Chat not found", + json: async () => ({}), + }); - const result = await getChatSnapshot("sess-1", "chat-1", "tok"); + await expect(getChatSnapshot("sess-1", "chat-1", "tok")).rejects.toThrow( + "Chat not found", + ); + }); - expect(result).toEqual({ messages: [], isStreaming: false }); + it("throws a generic HTTP error when the body has no text", async () => { + mockFetch.mockResolvedValueOnce({ + ok: false, + status: 500, + text: async () => "", + }); + + await expect(getChatSnapshot("sess-1", "chat-1", "tok")).rejects.toThrow( + "HTTP 500", + ); }); it("url-encodes both ids", async () => { diff --git a/lib/messages/getChatSnapshot.ts b/lib/messages/getChatSnapshot.ts index d4aa71d8b..7e3867c17 100644 --- a/lib/messages/getChatSnapshot.ts +++ b/lib/messages/getChatSnapshot.ts @@ -26,7 +26,12 @@ export async function getChatSnapshot( }, ); - if (!response.ok) return { messages: [], isStreaming: false }; + if (!response.ok) { + const errorText = await response.text().catch(() => ""); + throw new Error( + errorText || `Failed to fetch chat snapshot: HTTP ${response.status}`, + ); + } const data = (await response.json()) as { messages?: UIMessage[];