From 2d9c4cd04463c96efbb32294efc0f7ab1250fcaa Mon Sep 17 00:00:00 2001 From: Hasan Kumar <769950+mintuhouse@users.noreply.github.com> Date: Fri, 23 Oct 2026 14:58:05 +0530 Subject: [PATCH 1/3] Adopt immutable TanStack client snapshots in the SDK --- .../src/routes/docs/-content/sdk/headless.md | 2 + sdk/src/core/messages.ts | 7 +- sdk/src/core/sandbox.ts | 2 +- sdk/src/core/session.test.ts | 27 ++- sdk/src/core/session.ts | 212 +++++++++++++----- sdk/src/core/widget-renders.test.ts | 126 ++++++++++- sdk/src/widget/components/chat-transcript.tsx | 4 +- sdk/src/widget/lib/utils.ts | 2 +- 8 files changed, 311 insertions(+), 71 deletions(-) diff --git a/platform/src/routes/docs/-content/sdk/headless.md b/platform/src/routes/docs/-content/sdk/headless.md index ac15f2eee..8e6c95f6e 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 77914751b..be2234e8a 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 e8608a494..4206c0386 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.test.ts b/sdk/src/core/session.test.ts index 0e0f822d5..82ed13b56 100644 --- a/sdk/src/core/session.test.ts +++ b/sdk/src/core/session.test.ts @@ -171,17 +171,40 @@ 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.find((snapshot) => + snapshot.messages.some((message) => + message.parts.some((part) => part.type === "text" && part.content === "I wrote it."), + ), + )! + expect(firstDelta.messages.at(-1)?.parts).toContainEqual({ type: "text", content: "I" }) + expect(reply.messages.at(-1)).not.toBe(firstDelta.messages.at(-1)) + 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( + snapshots.some((snapshot, index) => { + const previous = snapshots[index - 1] + return previous?.status !== snapshot.status && previous?.messages === snapshot.messages + }), + ).toBe(true) }) // 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 a6f94c179..b440b9ce4 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 @@ -927,6 +983,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 +998,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 +1020,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 +1042,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,8 +1128,44 @@ export function createAstralBeamChat( } let selectedId: string | undefined let client = createClient() + 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 = () => { const previous = client + unsubscribeSnapshot?.() + unsubscribeSnapshot = undefined + snapshotMessages = undefined // Swap first, so the old client's reports on stop and detach are dropped. client = createClient() clientStatus = "ready" @@ -1097,7 +1174,7 @@ export function createAstralBeamChat( liveToolCalls.clear() liveToolMessageIds.clear() update({ activeToolCallIds: [] }) - if (started) client.attach() + if (started) attachClient() previous.dispose() } @@ -1106,7 +1183,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 +1203,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, + } : {}), } }) @@ -1151,6 +1235,9 @@ export function createAstralBeamChat( toolResults = threadToolResults.get(threadId ?? "") ?? new Map() historicalToolIds.clear() threadRecords = [] + messagePresentation.clear() + projectedMessages = new WeakMap() + presentationVersion++ disposeRenders() update({ thread: undefined, @@ -1176,7 +1263,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) @@ -1229,7 +1316,7 @@ export function createAstralBeamChat( update({ olderMessagesLoading: false }) try { client.detach() - client.attach() + attachClient() await hydration if (explicit && generation === selectionGeneration) update({ error: undefined }) } catch (error) { @@ -1251,6 +1338,7 @@ export function createAstralBeamChat( update({ olderMessagesLoading: true }) try { await client.loadOlderMessages() + if (generation === selectionGeneration && unsubscribeSnapshot) observeSnapshot() } catch (error) { if (generation === selectionGeneration) reportError(error) } finally { @@ -1462,7 +1550,7 @@ export function createAstralBeamChat( if (started) return started = true client.updateOptions({ tools: declareTools() }) - client.attach() + attachClient() if (requestController.signal.aborted) requestController = new AbortController() unsubscribeAuthentication = subscribeAuthentication(authentication, observeAuthentication) observeAuthentication() @@ -1592,6 +1680,8 @@ export function createAstralBeamChat( }, dispose: () => { releaseSubmission("The chat closed before this was sent.") + unsubscribeSnapshot?.() + unsubscribeSnapshot = undefined client.detach() disposeRenders() client.stop() @@ -1607,7 +1697,11 @@ 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() + 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 2b6d72a22..9a9cd7f0d 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,66 @@ 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 () => { + const cleanup = vi.fn() + const onRenderWidget = vi.fn((_request: WidgetRenderRequest) => cleanup) + const chat = createTestSession({ + threadId: "new", + widgets: { card: { description: "A host card" } }, + onRenderWidget, + }) + await vi.waitFor(() => expect(chat.getState().auth.status).toBe("ready")) + const native = mocked.client! + const readSnapshot = vi.spyOn(native, "getSnapshot") + chat.dispose() + const disposed = chat.getState() + readSnapshot.mockClear() + native.setMessagesManually([{ id: "closed", role: "assistant", parts: [] }]) + expect(readSnapshot).not.toHaveBeenCalled() + expect(chat.getState()).toBe(disposed) + + chat.start() + expect(mocked.client).toBe(native) + readSnapshot.mockClear() + native.setMessagesManually([ + { id: "restarted", role: "assistant", parts: [{ type: "text", content: "Reply" }] }, + ]) + expect(readSnapshot).toHaveBeenCalledOnce() + expect(chat.getState().messages[0]?.id).toBe("restarted") + + native.setMessagesManually([ + { + id: "retained-widget", + role: "assistant", + parts: [ + { + type: "tool-call", + id: "widget-call", + name: RENDER_WIDGET_TOOL, + arguments: '{"widget":"card","props":{}}', + input: { widget: "card", props: {} }, + state: "complete", + output: { widget: "card", rendered: true }, + }, + ], + }, + ]) + await vi.waitFor(() => expect(onRenderWidget).toHaveBeenCalledOnce()) + const retained = native.getSnapshot().messages + chat.dispose() + expect(cleanup).toHaveBeenCalledOnce() + chat.start() + await vi.waitFor(() => expect(onRenderWidget).toHaveBeenCalledTimes(2)) + expect(native.getSnapshot().messages).toBe(retained) + expect(onRenderWidget.mock.calls[1]?.[0].toolCallId).toBe("widget-call") }) function savedWidget(widget: string, id: string) { @@ -239,8 +300,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 +388,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 +454,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,6 +498,17 @@ 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 unsubscribeNative = native.subscribeSnapshot(() => + snapshots.push({ native: native.getSnapshot(), session: chat.getState() }), + ) const renderIds: unknown[] = [] const unsubscribe = chat.subscribe(() => { const part = chat.getState().messages.find((message) => message.id === "assistant-live") @@ -422,14 +518,36 @@ test("later stream snapshots preserve a live widget without rendering it again", }) await chat.sendMessage("Continue") unsubscribe() + unsubscribeNative() expect(chat.getState().error).toBeUndefined() expect(renderedPartDuringSend).toMatchObject({ 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 8bb7d2764..1bc05ccde 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 ab3680304..24b32ac35 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 [] From 87da9532e4ad5052410ebab5d5bf96b3894483df Mon Sep 17 00:00:00 2001 From: Hasan Kumar <769950+mintuhouse@users.noreply.github.com> Date: Fri, 23 Oct 2026 19:39:46 +0530 Subject: [PATCH 2/3] Preserve saved sessions across immediate restarts --- sdk/src/core/session-persistence.test.ts | 133 ++++++++++++++++------- sdk/src/core/session.test.ts | 17 +-- sdk/src/core/session.ts | 27 ++++- sdk/src/core/widget-renders.test.ts | 67 +++++------- 4 files changed, 148 insertions(+), 96 deletions(-) diff --git a/sdk/src/core/session-persistence.test.ts b/sdk/src/core/session-persistence.test.ts index 83ffd685d..9d07ca3c9 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 82ed13b56..6af5125d0 100644 --- a/sdk/src/core/session.test.ts +++ b/sdk/src/core/session.test.ts @@ -188,23 +188,16 @@ test("the reply after a server tool round keeps its first text delta", async () message.parts.some((part) => part.type === "text" && part.content === "I"), ), )! - const reply = snapshots.find((snapshot) => - snapshot.messages.some((message) => - message.parts.some((part) => part.type === "text" && part.content === "I wrote it."), - ), - )! + const reply = snapshots.findLast((snapshot) => snapshot.messages.at(-1)?.id === "reply")! expect(firstDelta.messages.at(-1)?.parts).toContainEqual({ type: "text", content: "I" }) - expect(reply.messages.at(-1)).not.toBe(firstDelta.messages.at(-1)) 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( - snapshots.some((snapshot, index) => { - const previous = snapshots[index - 1] - return previous?.status !== snapshot.status && previous?.messages === snapshot.messages - }), - ).toBe(true) + 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 b440b9ce4..44b5c1570 100644 --- a/sdk/src/core/session.ts +++ b/sdk/src/core/session.ts @@ -1549,12 +1549,35 @@ export function createAstralBeamChat( const start = () => { if (started) return started = true + const retained = client + const generation = selectionGeneration + const interruptedHydration = + requestController.signal.aborted && client.getIsHydrating() ? hydration : undefined client.updateOptions({ tools: declareTools() }) - attachClient() if (requestController.signal.aborted) requestController = new AbortController() unsubscribeAuthentication = subscribeAuthentication(authentication, observeAuthentication) observeAuthentication() stopAuthentication = startAuthentication(authentication) + attachClient() + // A reattach can reuse the history read cancelled by disposal. Retry only that read, + // after TanStack releases it, while this restarted selection still owns the client. + if (interruptedHydration) { + hydration = interruptedHydration.catch((error: unknown) => { + if ( + started && + client === retained && + generation === selectionGeneration && + error instanceof Error && + error.name === "AbortError" + ) { + client.detach() + attachClient() + return hydration + } + throw error + }) + void hydration.catch(() => undefined) + } } if (!deferStart) start() @@ -1687,7 +1710,6 @@ export function createAstralBeamChat( client.stop() requestController.abort() selectionGeneration++ - threadKey = undefined capabilitiesGeneration++ capabilityIdentity = undefined stopAuthentication?.() @@ -1699,6 +1721,7 @@ export function createAstralBeamChat( queueMicrotask(() => { if (!started) { client.dispose() + threadKey = undefined messagePresentation.clear() projectedMessages = new WeakMap() } diff --git a/sdk/src/core/widget-renders.test.ts b/sdk/src/core/widget-renders.test.ts index 9a9cd7f0d..de5705e32 100644 --- a/sdk/src/core/widget-renders.test.ts +++ b/sdk/src/core/widget-renders.test.ts @@ -102,57 +102,45 @@ test("API scope changes clear private MCP presentation and ignore old client eve }) 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: "new", + threadId: thread.id, widgets: { card: { description: "A host card" } }, onRenderWidget, }) - await vi.waitFor(() => expect(chat.getState().auth.status).toBe("ready")) + await vi.waitFor(() => expect(onRenderWidget).toHaveBeenCalledOnce()) const native = mocked.client! - const readSnapshot = vi.spyOn(native, "getSnapshot") + const retained = native.getSnapshot().messages chat.dispose() const disposed = chat.getState() - readSnapshot.mockClear() - native.setMessagesManually([{ id: "closed", role: "assistant", parts: [] }]) - expect(readSnapshot).not.toHaveBeenCalled() + 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) - readSnapshot.mockClear() - native.setMessagesManually([ - { id: "restarted", role: "assistant", parts: [{ type: "text", content: "Reply" }] }, - ]) - expect(readSnapshot).toHaveBeenCalledOnce() - expect(chat.getState().messages[0]?.id).toBe("restarted") + expect(onRenderWidget.mock.calls[1]?.[0].toolCallId).toBe( + "saved:assistant-retained:part-retained", + ) + native.setMessagesManually(retained) + expect(chat.getState().messages).toHaveLength(1) - native.setMessagesManually([ - { - id: "retained-widget", - role: "assistant", - parts: [ - { - type: "tool-call", - id: "widget-call", - name: RENDER_WIDGET_TOOL, - arguments: '{"widget":"card","props":{}}', - input: { widget: "card", props: {} }, - state: "complete", - output: { widget: "card", rendered: true }, - }, - ], - }, - ]) - await vi.waitFor(() => expect(onRenderWidget).toHaveBeenCalledOnce()) - const retained = native.getSnapshot().messages + // Restart also restores widgets when the native snapshot and presentation cache have not changed. chat.dispose() - expect(cleanup).toHaveBeenCalledOnce() + expect(cleanup).toHaveBeenCalledTimes(2) chat.start() - await vi.waitFor(() => expect(onRenderWidget).toHaveBeenCalledTimes(2)) - expect(native.getSnapshot().messages).toBe(retained) - expect(onRenderWidget.mock.calls[1]?.[0].toolCallId).toBe("widget-call") + await vi.waitFor(() => expect(onRenderWidget).toHaveBeenCalledTimes(3)) }) function savedWidget(widget: string, id: string) { @@ -506,19 +494,16 @@ test("later stream snapshots preserve attachment IDs and a live widget without r native: ReturnType session: ReturnType }> = [] - const unsubscribeNative = native.subscribeSnapshot(() => - snapshots.push({ native: native.getSnapshot(), session: chat.getState() }), - ) 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) }) await chat.sendMessage("Continue") unsubscribe() - unsubscribeNative() expect(chat.getState().error).toBeUndefined() expect(renderedPartDuringSend).toMatchObject({ id: "saved:assistant-live:part-live", From 2c4b828fb87c3a90efcc0b56dd915b5aedaae62e Mon Sep 17 00:00:00 2001 From: Hasan Kumar <769950+mintuhouse@users.noreply.github.com> Date: Fri, 23 Oct 2026 19:53:12 +0530 Subject: [PATCH 3/3] Reuse client replacement for interrupted history loads --- sdk/src/core/session.ts | 34 +++++++--------------------------- 1 file changed, 7 insertions(+), 27 deletions(-) diff --git a/sdk/src/core/session.ts b/sdk/src/core/session.ts index 44b5c1570..1b925f3c6 100644 --- a/sdk/src/core/session.ts +++ b/sdk/src/core/session.ts @@ -946,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" @@ -1161,18 +1162,18 @@ export function createAstralBeamChat( } client.attach() } - const replaceClient = () => { + 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) attachClient() previous.dispose() @@ -1234,6 +1235,7 @@ export function createAstralBeamChat( requestController = new AbortController() toolResults = threadToolResults.get(threadId ?? "") ?? new Map() historicalToolIds.clear() + liveToolMessageIds.clear() threadRecords = [] messagePresentation.clear() projectedMessages = new WeakMap() @@ -1313,7 +1315,6 @@ export function createAstralBeamChat( const refreshThread = async (explicit = true) => { if (!state.thread) return const generation = selectionGeneration - update({ olderMessagesLoading: false }) try { client.detach() attachClient() @@ -1549,35 +1550,14 @@ export function createAstralBeamChat( const start = () => { if (started) return started = true - const retained = client - const generation = selectionGeneration - const interruptedHydration = - requestController.signal.aborted && client.getIsHydrating() ? hydration : undefined + const interruptedHydration = requestController.signal.aborted && client.getIsHydrating() client.updateOptions({ tools: declareTools() }) if (requestController.signal.aborted) requestController = new AbortController() unsubscribeAuthentication = subscribeAuthentication(authentication, observeAuthentication) observeAuthentication() stopAuthentication = startAuthentication(authentication) - attachClient() - // A reattach can reuse the history read cancelled by disposal. Retry only that read, - // after TanStack releases it, while this restarted selection still owns the client. - if (interruptedHydration) { - hydration = interruptedHydration.catch((error: unknown) => { - if ( - started && - client === retained && - generation === selectionGeneration && - error instanceof Error && - error.name === "AbortError" - ) { - client.detach() - attachClient() - return hydration - } - throw error - }) - void hydration.catch(() => undefined) - } + if (interruptedHydration) replaceClient(client.getSnapshot().messages) + else attachClient() } if (!deferStart) start()