diff --git a/platform/src/routes/docs/-content/sdk/headless.md b/platform/src/routes/docs/-content/sdk/headless.md index ac15f2ee..8e6c95f6 100644 --- a/platform/src/routes/docs/-content/sdk/headless.md +++ b/platform/src/routes/docs/-content/sdk/headless.md @@ -56,6 +56,8 @@ chat.dispose() ## Reading the transcript +`messages` is a readonly snapshot. Treat its messages, parts, and metadata as immutable, and use session actions to change the conversation. Derive display values with methods such as `map`, `filter`, or `toSorted` instead of mutating the transcript. + The core exports the part helpers the widget itself renders with. - `isSandboxTool`, `readSandboxFileWrite`, `readSandboxCommandRun`, `readSandboxArtifact`, `collectSandboxActivity`. diff --git a/sdk/src/core/messages.ts b/sdk/src/core/messages.ts index 77914751..be2234e8 100644 --- a/sdk/src/core/messages.ts +++ b/sdk/src/core/messages.ts @@ -10,7 +10,7 @@ export function isSettledToolCall(part: { state: string; output?: unknown }): bo // A busy stream can be silent for a while (server-side reasoning, follow-ups after tool // results), so a "Thinking…" affordance shows until the newest assistant part visibly // makes progress. -export function lastPartInProgress(messages: UIMessage[]): boolean { +export function lastPartInProgress(messages: readonly UIMessage[]): boolean { const lastMessage = messages.at(-1) const lastPart = lastMessage?.role === "assistant" ? lastMessage.parts.at(-1) : undefined return ( @@ -25,7 +25,10 @@ export function lastPartInProgress(messages: UIMessage[]): boolean { // the message goes unanswered, and the redelivered call can re-execute a side-effecting // tool — so those windows count as busy too. Questionnaires and calls to tools this mount // never implemented stay interactive: the pre-send settle resolves them. -export function hasPendingToolRun(messages: UIMessage[], toolNames: ReadonlySet): boolean { +export function hasPendingToolRun( + messages: readonly UIMessage[], + toolNames: ReadonlySet, +): boolean { return messages.some((message) => message.parts.some( (part) => diff --git a/sdk/src/core/sandbox.ts b/sdk/src/core/sandbox.ts index e8608a49..4206c038 100644 --- a/sdk/src/core/sandbox.ts +++ b/sdk/src/core/sandbox.ts @@ -123,7 +123,7 @@ export function readSandboxCommandRun(part: ToolCallPart): SandboxCommandRun | u * The whole conversation's sandbox work, for the panel. A path written more than once appears * once, holding its latest content, because the panel answers "what is in the sandbox now". */ -export function collectSandboxActivity(messages: UIMessage[]): SandboxActivity { +export function collectSandboxActivity(messages: readonly UIMessage[]): SandboxActivity { const files = new Map() const commands: SandboxCommandRun[] = [] for (const message of messages) { diff --git a/sdk/src/core/session-persistence.test.ts b/sdk/src/core/session-persistence.test.ts index 83ffd685..9d07ca3c 100644 --- a/sdk/src/core/session-persistence.test.ts +++ b/sdk/src/core/session-persistence.test.ts @@ -1714,31 +1714,58 @@ test.each(["before", "after"] as const)( }, ) -test("a session restarted like React Strict Mode keeps loading history and sending", async () => { - let sends = 0 - stubSessionFetch((input: string | URL, init?: RequestInit) => { - const path = new URL(input).pathname - if (path.endsWith("/config")) return sessionResponse({ attachments: true }) - if (path.endsWith("/threads") && init?.method === "POST") - return sessionResponse({ ...thread, title: "", version: 0 }) - if (path.endsWith("/chat")) return Promise.resolve(acceptedStream(String(++sends))) - return sessionResponse(page()) - }) - const chat = createAstralBeamChat({ fetchAstralBeamToken: chatTestToken, threadId: "new" }, true) - try { - chat.start() - chat.dispose() - chat.start() - await vi.waitFor(() => expect(chat.getState().auth.status).toBe("ready")) - for (const text of ["first", "second"]) { - expect(await chat.sendMessage(text)).toMatchObject({ phase: "accepted" }) - expect(chat.getState().history).toBe("idle") +test.each(["new", "pending"])( + "a %s session restarted like React Strict Mode keeps loading history and sending", + async (initial) => { + let sends = 0 + let reads = 0 + let cancelled = false + stubSessionFetch((input: string | URL, init?: RequestInit) => { + const path = new URL(input).pathname + if (path.endsWith("/config")) return sessionResponse({ attachments: true }) + if (path.endsWith("/threads") && init?.method === "POST") + return sessionResponse({ ...thread, title: "", version: 0 }) + if (path.endsWith("/chat")) return Promise.resolve(acceptedStream(String(++sends))) + if (path.endsWith("/messages") && ++reads === 1 && initial === "pending") { + const signal = init!.signal! + return new Promise((_resolve, reject) => + signal.addEventListener( + "abort", + () => { + cancelled = true + reject(new DOMException("History request cancelled", "AbortError")) + }, + { once: true }, + ), + ) + } + return sessionResponse(page()) + }) + const chat = createAstralBeamChat( + { fetchAstralBeamToken: chatTestToken, threadId: initial === "new" ? "new" : thread.id }, + true, + ) + try { + chat.start() + if (initial !== "new") await vi.waitUntil(() => reads === 1) + chat.dispose() + expect(cancelled).toBe(initial === "pending") + chat.start() + await vi.waitFor(() => { + expect(chat.getState().auth.status).toBe("ready") + expect(chat.getState().history).toBe("idle") + expect(chat.getState().thread?.id).toBe(initial === "new" ? undefined : thread.id) + }) + for (const text of ["first", "second"]) { + expect(await chat.sendMessage(text)).toMatchObject({ phase: "accepted" }) + expect(chat.getState().history).toBe("idle") + } + expect(sends).toBe(2) + } finally { + chat.dispose() } - expect(sends).toBe(2) - } finally { - chat.dispose() - } -}) + }, +) test("a send before the configured conversation opens returns to its caller", async () => { const sends = vi.fn() @@ -1756,14 +1783,23 @@ test("a send before the configured conversation opens returns to its caller", as expect(chat.getState().submission).toBeUndefined() }) -test.each(["history", "thread"] as const)( - "Stop returns a send still waiting for its %s without dispatching it", +test.each(["history", "thread", "restarted history"] as const)( + "a send waiting for %s remains undispatched until its wait completes or Stop refuses it", async (wait) => { + const restarting = wait === "restarted history" + const waitingForHistory = wait !== "thread" const sends = vi.fn() - let release: (() => void) | undefined - const hold = (response: Response) => - new Promise((resolve) => { - release = () => resolve(response) + const release: Array<() => void> = [] + const hold = (response: Response, signal?: AbortSignal | null) => + new Promise((resolve, reject) => { + release.push(() => resolve(response)) + signal?.addEventListener( + "abort", + () => reject(new DOMException("Cancelled", "AbortError")), + { + once: true, + }, + ) }) let reads = 0 stubSessionFetch((input: string | URL, init?: RequestInit) => { @@ -1776,23 +1812,38 @@ test.each(["history", "thread"] as const)( if (path.endsWith("/threads") && init?.method === "POST") return hold(Response.json({ ...thread, title: "", version: 0 })) const response = Response.json(page()) - return wait === "history" && ++reads === 2 ? hold(response) : Promise.resolve(response) + return waitingForHistory && (++reads === 2 || (restarting && reads === 3)) + ? hold(response, restarting ? init?.signal : undefined) + : Promise.resolve(response) }) - const chat = createTestSession(wait === "history" ? { threadId: thread.id } : {}) + const chat = createTestSession(waitingForHistory ? { threadId: thread.id } : {}) await vi.waitFor(() => expect(chat.getState().auth.status).toBe("ready")) - if (wait === "history") { + if (waitingForHistory) { await vi.waitUntil(() => chat.getState().thread) void chat.refreshThread() } - expect(chat.getState().history).toBe(wait === "history" ? "refreshing" : "idle") - const outcome = chat.sendMessage("Stop me") - await vi.waitFor(() => expect(release).toBeDefined()) - chat.stop() - expect(chat.getState()).toMatchObject({ status: "ready", submission: { phase: "refused" } }) - release!() - expect(await outcome).toEqual(expect.objectContaining({ phase: "refused", content: "Stop me" })) - expect((await outcome).error).toBeUndefined() + if (restarting) { + await vi.waitUntil(() => release.length === 1) + chat.dispose() + chat.start() + } + expect(chat.getState().history).toBe(waitingForHistory ? "refreshing" : "idle") + const outcome = chat.sendMessage("Wait for me") + await vi.waitFor(() => expect(release).toHaveLength(restarting ? 2 : 1)) + await new Promise((resolve) => setTimeout(resolve, 0)) expect(sends).not.toHaveBeenCalled() + if (!restarting) chat.stop() + expect(chat.getState()).toMatchObject({ + status: restarting ? "submitted" : "ready", + submission: { phase: restarting ? "submitting" : "refused" }, + }) + release.at(-1)!() + expect(await outcome).toMatchObject({ + phase: restarting ? "accepted" : "refused", + content: "Wait for me", + }) + expect((await outcome).error).toBeUndefined() + expect(sends).toHaveBeenCalledTimes(restarting ? 1 : 0) }, ) diff --git a/sdk/src/core/session.test.ts b/sdk/src/core/session.test.ts index 0e0f822d..6af5125d 100644 --- a/sdk/src/core/session.test.ts +++ b/sdk/src/core/session.test.ts @@ -171,17 +171,33 @@ test("the reply after a server tool round keeps its first text delta", async () const chat = createTestSession({ streamCallbacks: { onFinish }, }) - const statuses: string[] = [] - chat.subscribe(() => statuses.push(chat.getState().status)) + const snapshots: ReturnType[] = [] + chat.subscribe(() => snapshots.push(chat.getState())) await chat.sendMessage("Write a file") expect(onFinish.mock.calls.at(-1)?.[0].parts.at(-1)).toEqual({ type: "text", content: "I wrote it.", }) // The tool round's intermediate RUN_FINISHED must not make the widget look idle mid-turn. + const statuses = snapshots.map((snapshot) => snapshot.status) expect( statuses.slice(statuses.indexOf("streaming"), statuses.lastIndexOf("streaming")), ).not.toContain("ready") + const firstDelta = snapshots.find((snapshot) => + snapshot.messages.some((message) => + message.parts.some((part) => part.type === "text" && part.content === "I"), + ), + )! + const reply = snapshots.findLast((snapshot) => snapshot.messages.at(-1)?.id === "reply")! + expect(firstDelta.messages.at(-1)?.parts).toContainEqual({ type: "text", content: "I" }) + expect(firstDelta.messages.find((message) => message.id === "call-turn")).toBeDefined() + expect(reply.messages.find((message) => message.id === "call-turn")).toBe( + firstDelta.messages.find((message) => message.id === "call-turn"), + ) + expect(reply.status).toBe("ready") + expect(reply.messages).toBe( + snapshots.findLast((snapshot) => snapshot.status === "streaming")?.messages, + ) }) // A React host rebuilds its tool objects every render, so publishing a fresh `agentTools` for an diff --git a/sdk/src/core/session.ts b/sdk/src/core/session.ts index a6f94c17..1b925f3c 100644 --- a/sdk/src/core/session.ts +++ b/sdk/src/core/session.ts @@ -1,5 +1,6 @@ import { ChatClient, + type ChatClientSnapshot, type ChatClientState, type ConnectionAdapter, fetchServerSentEvents, @@ -168,7 +169,7 @@ export interface ChatSubmission { } export interface AstralBeamChatState { - messages: UIMessage[] + messages: ChatClientSnapshot["messages"] /** The run status: "ready", "submitted", "streaming", or "error". */ status: ChatClientState /** The latest failure outside a submission, such as history loading or a failed reply. */ @@ -372,6 +373,34 @@ export function createAstralBeamChat( let threadKey: string | undefined let startWithNewThread = false let threadRecords: Awaited>["messages"] = [] + type MessagePresentation = { + attachmentParts?: UIMessage["parts"] + renderIds?: ReadonlyMap + } + const messagePresentation = new Map() + let presentationVersion = 0 + let projectedMessages = new WeakMap< + UIMessage, + { presentation: MessagePresentation | undefined; message: UIMessage } + >() + const rememberAttachments = (messages: readonly UIMessage[]) => { + for (const message of messages) { + // Accepted user content is immutable. Keep its saved download parts across model snapshots. + if ( + message.role !== "user" || + messagePresentation.get(message.id)?.attachmentParts || + !message.parts.some((part) => "fileId" in part) + ) + continue + const parts = message.parts.map((part) => Object.freeze({ ...part })) + Object.freeze(parts) + messagePresentation.set(message.id, { + ...messagePresentation.get(message.id), + attachmentParts: parts, + }) + presentationVersion++ + } + } const clientId = newUuid() const acceptSubmission = () => { if (!submission || !settle(submission, "accepted")) return @@ -379,11 +408,13 @@ export function createAstralBeamChat( submission.callbacks?.onAccepted?.() } const replaceAcceptedMessage = (record: ChatHistoryPageEncodedMessagesItem) => { - const messages = [...client.getMessages()] + const messages = [...client.getSnapshot().messages] const index = messages.findLastIndex((message) => message.role === "user") if (index !== -1) { messages[index] = projectThreadMessages([record])[0]! + rememberAttachments([messages[index]]) client.setMessagesManually(messages) + if (unsubscribeSnapshot) observeSnapshot() } } const liveToolCalls = new Set() @@ -525,22 +556,61 @@ export function createAstralBeamChat( return { widget: input.widget, rendered: true } } + const projectMessages = (messages: ChatClientSnapshot["messages"]) => + Object.freeze( + messages.map((message) => { + let presentation = messagePresentation.get(message.id) + const cached = projectedMessages.get(message) + if (cached && cached.presentation === presentation) return cached.message + let parts = presentation?.attachmentParts ?? message.parts + let renderIds = presentation?.renderIds + const nextParts = parts.map((part) => { + if ( + part.type !== "tool-call" || + part.name !== RENDER_WIDGET_TOOL || + part.state !== "complete" || + !part.input + ) + return part + const stored: ChatToolCallPart = part + const alias = savedToolCallId(message.id, part.id) + const canonical = stored.applicationPartId + ? savedToolCallId(message.id, stored.applicationPartId, stored.responseTargetId) + : alias + const upstream = stored.upstreamToolCallId + const renderId = + stored.widgetRenderId ?? + renderIds?.get(alias) ?? + renderIds?.get(canonical) ?? + (upstream && + liveToolMessageIds.get(upstream) === message.id && + restoredWidgets.has(upstream) + ? upstream + : part.id) + if (renderIds?.get(alias) !== renderId || renderIds.get(canonical) !== renderId) { + renderIds = new Map(renderIds).set(alias, renderId).set(canonical, renderId) + } + return stored.widgetRenderId === renderId + ? part + : Object.freeze({ ...part, widgetRenderId: renderId }) + }) + if (nextParts.some((part, index) => part !== parts[index])) { + Object.freeze(nextParts) + parts = nextParts + } + if (renderIds && renderIds !== presentation?.renderIds) { + presentation = { ...presentation, renderIds } + messagePresentation.set(message.id, presentation) + presentationVersion++ + } + const projected = parts === message.parts ? message : Object.freeze({ ...message, parts }) + projectedMessages.set(message, { presentation, message: projected }) + return projected + }), + ) + const restoreCompletedWidgets = (messages: readonly UIMessage[]) => { if (state.thread?.role === "viewer") return - const renderIds = new Map() - for (const message of state.messages) { - for (const part of message.parts) { - if (part.type !== "tool-call" || part.name !== RENDER_WIDGET_TOOL) continue - const stored = part as ChatToolCallPart - if (!stored.widgetRenderId) continue - renderIds.set(savedToolCallId(message.id, part.id), stored.widgetRenderId) - if (stored.applicationPartId) - renderIds.set( - savedToolCallId(message.id, stored.applicationPartId, stored.responseTargetId), - stored.widgetRenderId, - ) - } - } for (const message of messages) { for (const part of message.parts) { if ( @@ -550,22 +620,8 @@ export function createAstralBeamChat( !part.input ) continue - const stored = part as ChatToolCallPart - const upstream = stored.upstreamToolCallId - const renderId = - stored.widgetRenderId ?? - renderIds.get(savedToolCallId(message.id, part.id)) ?? - (stored.applicationPartId - ? renderIds.get( - savedToolCallId(message.id, stored.applicationPartId, stored.responseTargetId), - ) - : undefined) ?? - (upstream && - liveToolMessageIds.get(upstream) === message.id && - restoredWidgets.has(upstream) - ? upstream - : part.id) - stored.widgetRenderId = renderId + const stored: ChatToolCallPart = part + const renderId = stored.widgetRenderId ?? part.id if (restoredWidgets.has(renderId) || restoringWidgets.has(renderId)) continue restoringWidgets.add(renderId) const generation = selectionGeneration @@ -890,6 +946,7 @@ export function createAstralBeamChat( // Like TanStack Query's isLoading and isRefetching, loading means nothing is shown yet. if (!options?.before) update({ + olderMessagesLoading: false, history: state.messages.length > 0 || state.thread?.hasMessages === false ? "refreshing" @@ -927,6 +984,7 @@ export function createAstralBeamChat( older ? state.pendingInteractions : history.pendingInteractions, liveToolMessageIds, ).filter((message) => pageIds.has(message.id)) + rememberAttachments(messages) if (older) update({ messagesCursor: history.page_after ?? undefined }) else { if (history.thread.role === "viewer") disposeRenders() @@ -941,7 +999,6 @@ export function createAstralBeamChat( // An empty saved thread must remove rejected optimistic input. if (messages.length === 0) next.setMessagesManually([]) } - restoreCompletedWidgets(messages) // Foreground execution has no replay stream to join. Pending interactions remain // application-owned because several turns can wait independently in one thread. return { @@ -964,7 +1021,10 @@ export function createAstralBeamChat( (history) => new Promise((resolve) => queueMicrotask(() => { - if (client === next && readGeneration === historyGeneration) update({ history }) + if (client === next && readGeneration === historyGeneration) { + if (unsubscribeSnapshot) observeSnapshot() + update({ history }) + } resolve() }), ), @@ -983,24 +1043,6 @@ export function createAstralBeamChat( queue: "drop", tools: declareTools(), forwardedProps: forwardedProps(), - onMessagesChange: (messages) => { - if (client !== next) return - // Model snapshots omit client render IDs. Retain the existing slot before publishing them. - restoreCompletedWidgets(messages) - // User messages are immutable. Keep saved download identities across model snapshots. - const savedMessages = new Map(state.messages.map((message) => [message.id, message])) - for (const message of messages) { - if (message.role !== "user") continue - const saved = savedMessages.get(message.id) - if (saved?.parts.some((part) => "fileId" in part)) message.parts = saved.parts - } - update({ messages, sandbox: collectSandboxActivity(messages) }) - }, - onStatusChange: (status) => { - if (client !== next) return - clientStatus = status - update(status === "submitted" ? { usageLimitReached: undefined } : {}) - }, onInterruptStateChange: (interrupts) => { if (client === next) debug?.("tool", "interrupt state changed", interrupts) }, @@ -1087,17 +1129,53 @@ export function createAstralBeamChat( } let selectedId: string | undefined let client = createClient() - const replaceClient = () => { + let unsubscribeSnapshot: (() => void) | undefined + let snapshotMessages: ChatClientSnapshot["messages"] | undefined + let projectedPresentationVersion = -1 + const observeSnapshot = () => { + const snapshot = client.getSnapshot() + const messagesChanged = + snapshotMessages !== snapshot.messages || projectedPresentationVersion !== presentationVersion + if (!messagesChanged && clientStatus === snapshot.status) return + const next: Partial = {} + if (clientStatus !== snapshot.status && snapshot.status === "submitted") + next.usageLimitReached = undefined + clientStatus = snapshot.status + if (messagesChanged) { + next.messages = projectMessages(snapshot.messages) + next.sandbox = collectSandboxActivity(next.messages) + snapshotMessages = snapshot.messages + projectedPresentationVersion = presentationVersion + } + update(next) + if (messagesChanged) restoreCompletedWidgets(state.messages) + } + const attachClient = () => { + if (!unsubscribeSnapshot) { + const observed = client + const generation = selectionGeneration + unsubscribeSnapshot = observed.subscribeSnapshot(() => { + if (client === observed && generation === selectionGeneration) observeSnapshot() + }) + observeSnapshot() + restoreCompletedWidgets(state.messages) + } + client.attach() + } + const replaceClient = (messages?: ChatClientSnapshot["messages"]) => { const previous = client + unsubscribeSnapshot?.() + unsubscribeSnapshot = undefined + snapshotMessages = undefined // Swap first, so the old client's reports on stop and detach are dropped. client = createClient() + if (messages) client.setMessagesManually(messages) clientStatus = "ready" previous.stop() previous.detach() liveToolCalls.clear() - liveToolMessageIds.clear() update({ activeToolCallIds: [] }) - if (started) client.attach() + if (started) attachClient() previous.dispose() } @@ -1106,7 +1184,7 @@ export function createAstralBeamChat( const scopeCompletedToolCalls = () => { historicalToolIds.clear() client.setMessagesManually( - state.messages.map((message) => { + client.getSnapshot().messages.map((message) => { const aliases = new Map() const parts = message.parts.map((part) => { if (part.type !== "tool-call" || webPartMetadata(part).providerExecuted) return part @@ -1126,7 +1204,14 @@ export function createAstralBeamChat( id, ...(part.name === RENDER_WIDGET_TOOL && liveToolMessageIds.get(upstreamToolCallId) === message.id - ? { widgetRenderId: stored.widgetRenderId ?? part.id } + ? { + widgetRenderId: + stored.widgetRenderId ?? + messagePresentation + .get(message.id) + ?.renderIds?.get(savedToolCallId(message.id, part.id)) ?? + part.id, + } : {}), } }) @@ -1150,7 +1235,11 @@ export function createAstralBeamChat( requestController = new AbortController() toolResults = threadToolResults.get(threadId ?? "") ?? new Map() historicalToolIds.clear() + liveToolMessageIds.clear() threadRecords = [] + messagePresentation.clear() + projectedMessages = new WeakMap() + presentationVersion++ disposeRenders() update({ thread: undefined, @@ -1176,7 +1265,7 @@ export function createAstralBeamChat( submission = submissions.get(id) if (submission) update({ submission: project(submission) }) try { - if (!started) client.attach() + if (!started) attachClient() await hydration if (generation !== selectionGeneration) return if (threadKey) selectedThread(threadKey, id) @@ -1226,10 +1315,9 @@ export function createAstralBeamChat( const refreshThread = async (explicit = true) => { if (!state.thread) return const generation = selectionGeneration - update({ olderMessagesLoading: false }) try { client.detach() - client.attach() + attachClient() await hydration if (explicit && generation === selectionGeneration) update({ error: undefined }) } catch (error) { @@ -1251,6 +1339,7 @@ export function createAstralBeamChat( update({ olderMessagesLoading: true }) try { await client.loadOlderMessages() + if (generation === selectionGeneration && unsubscribeSnapshot) observeSnapshot() } catch (error) { if (generation === selectionGeneration) reportError(error) } finally { @@ -1461,12 +1550,14 @@ export function createAstralBeamChat( const start = () => { if (started) return started = true + const interruptedHydration = requestController.signal.aborted && client.getIsHydrating() client.updateOptions({ tools: declareTools() }) - client.attach() if (requestController.signal.aborted) requestController = new AbortController() unsubscribeAuthentication = subscribeAuthentication(authentication, observeAuthentication) observeAuthentication() stopAuthentication = startAuthentication(authentication) + if (interruptedHydration) replaceClient(client.getSnapshot().messages) + else attachClient() } if (!deferStart) start() @@ -1592,12 +1683,13 @@ export function createAstralBeamChat( }, dispose: () => { releaseSubmission("The chat closed before this was sent.") + unsubscribeSnapshot?.() + unsubscribeSnapshot = undefined client.detach() disposeRenders() client.stop() requestController.abort() selectionGeneration++ - threadKey = undefined capabilitiesGeneration++ capabilityIdentity = undefined stopAuthentication?.() @@ -1607,7 +1699,12 @@ export function createAstralBeamChat( // React Strict Mode immediately restarts a committed session. Dispose the client only if // that restart did not happen: https://react.dev/reference/react/useEffect#caveats queueMicrotask(() => { - if (!started) client.dispose() + if (!started) { + client.dispose() + threadKey = undefined + messagePresentation.clear() + projectedMessages = new WeakMap() + } }) listeners.clear() }, diff --git a/sdk/src/core/widget-renders.test.ts b/sdk/src/core/widget-renders.test.ts index 2b6d72a2..de5705e3 100644 --- a/sdk/src/core/widget-renders.test.ts +++ b/sdk/src/core/widget-renders.test.ts @@ -12,6 +12,7 @@ interface ClientTool { const mocked = vi.hoisted(() => ({ tools: [] as readonly ClientTool[], onCustomEvent: undefined as ((name: string, data: unknown) => void) | undefined, + client: undefined as import("@tanstack/ai-client").ChatClient | undefined, })) // Capture the declared render callback while retaining the real client's stream and history behavior. @@ -22,6 +23,7 @@ vi.mock("@tanstack/ai-client", async (importOriginal) => { ChatClient: class extends actual.ChatClient { constructor(options: ConstructorParameters[0]) { super(options) + mocked.client = this mocked.tools = (options.tools ?? []) as readonly ClientTool[] mocked.onCustomEvent = (name, data) => options.onCustomEvent?.(name, data, {}) } @@ -81,6 +83,7 @@ test("API scope changes clear private MCP presentation and ignore old client eve const chat = createTestSession({ threadId: "new" }) await vi.waitFor(() => expect(chat.getState().auth.status).toBe("ready")) const oldEvent = mocked.onCustomEvent! + const oldClient = mocked.client! const event = { tools: [{ name: "mcp-tool", title: "Private account" }], unavailable: ["private-id"], @@ -90,8 +93,54 @@ test("API scope changes clear private MCP presentation and ignore old client eve expect(chat.getState().unavailableMcpConnections).toEqual(["private-id"]) chat.updateOptions({ apiUrl: "https://replacement.example/api" }) oldEvent("astralbeam_mcp_tools", event) + oldClient.setMessagesManually([ + { id: "stale", role: "assistant", parts: [{ type: "text", content: "Private account" }] }, + ]) expect(chat.getState().serverToolTitles).toEqual({}) expect(chat.getState().unavailableMcpConnections).toEqual([]) + expect(chat.getState().messages).toEqual([]) +}) + +test("snapshot observation stops on disposal and resumes on immediate restart", async () => { + savedMessages = savedWidget("card", "retained") + const cleanup = vi.fn() + const onRenderWidget = vi.fn((_request: WidgetRenderRequest) => cleanup) + const chat = createTestSession({ + threadId: thread.id, + widgets: { card: { description: "A host card" } }, + onRenderWidget, + }) + await vi.waitFor(() => expect(onRenderWidget).toHaveBeenCalledOnce()) + const native = mocked.client! + const retained = native.getSnapshot().messages + chat.dispose() + const disposed = chat.getState() + expect(cleanup).toHaveBeenCalledOnce() + native.setMessagesManually([ + ...retained, + { id: "detached", role: "assistant", parts: [{ type: "text", content: "Reply" }] }, + ]) + expect(chat.getState()).toBe(disposed) + + chat.start() + expect(chat.getState().messages.at(-1)?.id).toBe("detached") + await vi.waitFor(() => { + expect(chat.getState().auth.status).toBe("ready") + expect(chat.getState().history).toBe("idle") + expect(onRenderWidget).toHaveBeenCalledTimes(2) + }) + expect(mocked.client).toBe(native) + expect(onRenderWidget.mock.calls[1]?.[0].toolCallId).toBe( + "saved:assistant-retained:part-retained", + ) + native.setMessagesManually(retained) + expect(chat.getState().messages).toHaveLength(1) + + // Restart also restores widgets when the native snapshot and presentation cache have not changed. + chat.dispose() + expect(cleanup).toHaveBeenCalledTimes(2) + chat.start() + await vi.waitFor(() => expect(onRenderWidget).toHaveBeenCalledTimes(3)) }) function savedWidget(widget: string, id: string) { @@ -239,8 +288,15 @@ test.each(["definition", "renderer", "failed-render"])( : onRenderWidget, }) await vi.waitFor(() => expect(chat.getState().messages).toHaveLength(1)) + const nativeMessages = mocked.client!.getSnapshot().messages chat.updateOptions({ widgets, onRenderWidget }) await vi.waitFor(() => expect(onRenderWidget).toHaveBeenCalledTimes(1)) + expect(mocked.client!.getSnapshot().messages).toBe(nativeMessages) + expect(nativeMessages[0]?.parts[0]).not.toHaveProperty("widgetRenderId") + expect(chat.getState().messages[0]?.parts[0]).toHaveProperty( + "widgetRenderId", + "saved:assistant-late:part-late", + ) chat.updateOptions({ widgets: { ...widgets } }) await chat.refreshThread() expect(onRenderWidget).toHaveBeenCalledTimes(1) @@ -320,18 +376,34 @@ test("an old async restoration cannot release a newly selected thread's pending await vi.waitFor(() => expect(onRenderWidget).toHaveBeenCalledTimes(1)) }) -test("later stream snapshots preserve a live widget without rendering it again", async () => { +test("later stream snapshots preserve attachment IDs and a live widget without rendering again", async () => { const fetchHistory = fetch let sends = 0 let renderedPartDuringSend: unknown const onRenderWidget = vi.fn<(request: WidgetRenderRequest) => void>() + const attachment = { + id: "uploaded", + role: "user", + created_at: thread.created_at, + parts: [ + { + type: "document", + mediaType: "application/pdf", + source: { type: "attachment", id: "file" }, + }, + ], + } + savedMessages = [attachment] vi.stubGlobal("fetch", (input: RequestInfo | URL, init?: RequestInit) => { const path = input instanceof Request ? input.url : String(input) if (!path.endsWith("/chat")) return fetchHistory(input, init) sends++ const first = sends === 1 - if (!first) renderedPartDuringSend = chat.getState().messages[0]?.parts[0] - savedMessages = savedWidget("card", "live") + if (!first) + renderedPartDuringSend = chat + .getState() + .messages.find((message) => message.id === "assistant-live")?.parts[0] + savedMessages = [attachment, ...savedWidget("card", "live")] const run = { threadId: thread.id, runId: `run-${sends}` } const chunks = [ { type: "RUN_STARTED", ...run }, @@ -370,6 +442,7 @@ test("later stream snapshots preserve a live widget without rendering it again", { type: "MESSAGES_SNAPSHOT", messages: [ + { id: "uploaded", role: "user", content: "Provider attachment projection" }, { id: "assistant-live", role: "assistant", @@ -413,10 +486,19 @@ test("later stream snapshots preserve a live widget without rendering it again", expect(chat.getState().error).toBeUndefined() await vi.waitFor(() => expect(onRenderWidget).toHaveBeenCalledTimes(1)) expect(onRenderWidget.mock.calls[0]?.[0].toolCallId).toBe("reused-provider-id") + const attachmentParts = chat + .getState() + .messages.find((message) => message.id === "uploaded")!.parts + const native = mocked.client! + const snapshots: Array<{ + native: ReturnType + session: ReturnType + }> = [] const renderIds: unknown[] = [] const unsubscribe = chat.subscribe(() => { - const part = chat.getState().messages.find((message) => message.id === "assistant-live") - ?.parts[0] + const session = chat.getState() + snapshots.push({ native: native.getSnapshot(), session }) + const part = session.messages.find((message) => message.id === "assistant-live")?.parts[0] if (part?.type === "tool-call") renderIds.push("widgetRenderId" in part ? part.widgetRenderId : undefined) }) @@ -427,9 +509,30 @@ test("later stream snapshots preserve a live widget without rendering it again", id: "saved:assistant-live:part-live", widgetRenderId: "reused-provider-id", }) - expect(chat.getState().messages[0]?.parts[0]).toMatchObject({ + expect( + chat.getState().messages.find((message) => message.id === "assistant-live")?.parts[0], + ).toMatchObject({ widgetRenderId: "reused-provider-id", }) + expect( + chat.getState().messages.find((message) => message.id === "uploaded")?.parts[0], + ).toMatchObject({ type: "document", fileId: "file" }) + const providerSnapshot = snapshots.find((snapshot) => + snapshot.native.messages.some((message) => + message.parts.some( + (part) => part.type === "text" && part.content === "Provider attachment projection", + ), + ), + )! + expect( + providerSnapshot.session.messages.find((message) => message.id === "uploaded")?.parts, + ).toBe(attachmentParts) + expect( + providerSnapshot.native.messages.find((message) => message.id === "uploaded")?.parts, + ).toEqual([{ type: "text", content: "Provider attachment projection" }]) + expect( + providerSnapshot.native.messages.find((message) => message.id === "assistant-live")?.parts[0], + ).not.toHaveProperty("widgetRenderId") expect(onRenderWidget).toHaveBeenCalledTimes(1) expect(renderIds.length).toBeGreaterThan(0) expect(new Set(renderIds)).toEqual(new Set(["reused-provider-id"])) diff --git a/sdk/src/widget/components/chat-transcript.tsx b/sdk/src/widget/components/chat-transcript.tsx index 8bb7d276..1bc05ccd 100644 --- a/sdk/src/widget/components/chat-transcript.tsx +++ b/sdk/src/widget/components/chat-transcript.tsx @@ -23,7 +23,7 @@ import { UserMessageBody } from "./user-message-body.tsx" interface ChatTranscriptProps { readOnly?: boolean | undefined - messages: UIMessage[] + messages: readonly UIMessage[] /** Name of the host's empty-state slot; when set, it replaces the default empty state. */ emptySlot?: string | undefined /** Headline of the empty transcript; defaults to `DEFAULT_EMPTY_HEADLINE`. */ @@ -173,7 +173,7 @@ export function ChatTranscript({ }: ChatTranscriptProps) { const loadingEarlier = useRef(false) const viewport = useRef(null) - const automaticPage = useRef(undefined) + const automaticPage = useRef(undefined) useEffect(() => { const element = viewport.current if (!element || !hasOlder || loadingOlder || loadingHistory) return diff --git a/sdk/src/widget/lib/utils.ts b/sdk/src/widget/lib/utils.ts index ab368030..24b32ac3 100644 --- a/sdk/src/widget/lib/utils.ts +++ b/sdk/src/widget/lib/utils.ts @@ -60,7 +60,7 @@ export function getMessageText(message: UIMessage): string { } /** The conversation's visible text as Markdown, one section per user or assistant message. */ -export function transcriptMarkdown(title: string, messages: UIMessage[]): string { +export function transcriptMarkdown(title: string, messages: readonly UIMessage[]): string { const sections = messages.flatMap((message) => { const text = getMessageText(message).trim() if (!text || (message.role !== "user" && message.role !== "assistant")) return []