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..e30c18e8e --- /dev/null +++ b/lib/messages/__tests__/getChatSnapshot.test.ts @@ -0,0 +1,88 @@ +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("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 () => ({}), + }); + + await expect(getChatSnapshot("sess-1", "chat-1", "tok")).rejects.toThrow( + "Chat not found", + ); + }); + + 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 () => { + 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..7e3867c17 --- /dev/null +++ b/lib/messages/getChatSnapshot.ts @@ -0,0 +1,45 @@ +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) { + 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[]; + isStreaming?: boolean; + }; + + return { + messages: data.messages ?? [], + isStreaming: data.isStreaming ?? false, + }; +}