@tanstack/ai-persistence/src/reconstruct.ts

Composable state persistence for TanStack AI messages, runs, interrupts, metadata, and locks.

import { modelMessagesToUIMessages } from '@tanstack/ai'
import type { UIMessage } from '@tanstack/ai'
import { validateReconstructChatStores } from './types'
import type { AIPersistence, ChatTranscriptStores } from './types'

/**
 * The JSON body `reconstructChat` returns and a server-authoritative client
 * hydrates from on mount.
 *
 * `messages` is the stored transcript as UI messages (ready to paint).
 * `activeRun` is a cursor to a run still generating for the thread, or `null` —
 * resolved from the STABLE thread id via `stores.runs.findActiveRun`, so the
 * client learns "there is a live run to tail" without ever handling a run id.
 * `interrupts` is the thread's pending human-in-the-loop interrupts (tool
 * approvals, client-tool/generic waits) and the run they paused, or `null` —
 * so a reload (or another device) re-prompts the approval from the SERVER, not
 * from client storage. Resolved via `stores.interrupts.listPending`.
 */
export interface ReconstructedChat {
  messages: Array<UIMessage>
  activeRun: { runId: string } | null
  interrupts: {
    runId: string
    pending: Array<Record<string, unknown>>
  } | null
}

export interface ReconstructChatOptions {
  /** Query parameter carrying the thread id. Defaults to `threadId`. */
  param?: string
  /**
   * Authorize access to the requested thread before loading history.
   *
   * ⚠️ Without this, any caller who knows or guesses `?threadId=` receives the
   * full transcript. Multi-user / multi-tenant deployments **must** supply
   * an authorization check (session → owned threads) or resolve a validated
   * thread id in the route and pass it via a custom `param` that only your
   * server sets.
   *
   * Return:
   * - `true` to allow the load
   * - `false` for a default `403` response
   * - a `Response` to return as-is (e.g. `401` with a body)
   */
  authorize?: (
    threadId: string,
    request: Request,
  ) => boolean | Response | Promise<boolean | Response>
}

/**
 * Build the JSON `Response` a server-authoritative client hydrates from on load
 * (see the client-persistence guide). Reads the thread id from the request query
 * (`?threadId=` by default) and returns `{ messages, activeRun, interrupts }`
 * ({@link ReconstructedChat}):
 *
 * - `messages` — the stored transcript as UI messages.
 * - `activeRun` — `{ runId }` if a run is still generating for the thread (so the
 *   client tails it via the durability stream), else `null`. Resolved via the
 *   required `stores.runs.findActiveRun`; `null` when the `runs` store is absent.
 * - `interrupts` — `{ runId, pending }` if the thread has pending human-in-the-loop
 *   interrupts (a paused approval / wait) and the run they paused, else `null`, so
 *   a reload re-prompts the decision from the server. Resolved via the optional
 *   `stores.interrupts.listPending`; `null` when that store is absent.
 *
 * Requires `stores.messages`. Returns an empty transcript with no active run
 * and no interrupts when the thread id is missing or the thread is unknown, so
 * the caller never has to special-case a first load.
 *
 * This helper does **not** enforce tenancy by itself. Pass
 * {@link ReconstructChatOptions.authorize} (or wrap the call in your own
 * session gate) before exposing it on a public route.
 *
 * ```ts
 * export async function GET(request: Request) {
 *   return reconstructChat(persistence, request, {
 *     authorize: async (threadId, req) => {
 *       const userId = await getSessionUserId(req)
 *       return userId != null && (await userOwnsThread(userId, threadId))
 *     },
 *   })
 * }
 * ```
 */
export async function reconstructChat(
  persistence: AIPersistence<ChatTranscriptStores>,
  request: Request,
  options?: ReconstructChatOptions,
): Promise<Response> {
  validateReconstructChatStores(persistence)
  const messageStore = persistence.stores.messages
  if (!messageStore) {
    // validateReconstructChatStores already throws; this narrows for TypeScript.
    throw new Error('reconstructChat requires stores.messages.')
  }

  const param = options?.param ?? 'threadId'
  const threadId = new URL(request.url).searchParams.get(param) ?? ''

  if (threadId && options?.authorize) {
    const decision = await options.authorize(threadId, request)
    if (decision instanceof Response) {
      return decision
    }
    if (!decision) {
      return new Response(JSON.stringify({ error: 'Forbidden' }), {
        status: 403,
        headers: {
          'content-type': 'application/json',
          'cache-control': 'no-store',
        },
      })
    }
  }

  // Resolve the active run BEFORE reading the transcript. `withPersistence`
  // persists the final transcript BEFORE marking a run complete, so observing
  // "no active run" here guarantees the transcript read below is the FINAL one.
  // Reading them in the other order opens a finish-window race: a fast run that
  // completes between the two reads would return a stale streaming snapshot with
  // `activeRun: null`, leaving the client stuck on the partial (no run to tail).
  const active = threadId
    ? await persistence.stores.runs?.findActiveRun(threadId)
    : null
  const stored = threadId ? await messageStore.loadThread(threadId) : []
  // Pending interrupts for the thread, so a reload re-prompts the approval from
  // the server. Each stored `payload` is the full interrupt descriptor the
  // client hydrates; they share the run they paused.
  const pending = threadId
    ? ((await persistence.stores.interrupts?.listPending(threadId)) ?? [])
    : []
  const firstPending = pending[0]
  const body: ReconstructedChat = {
    messages: modelMessagesToUIMessages(stored),
    activeRun: active ? { runId: active.runId } : null,
    interrupts: firstPending
      ? {
          runId: firstPending.runId,
          pending: pending.map((record) => record.payload),
        }
      : null,
  }
  return new Response(JSON.stringify(body), {
    headers: {
      'content-type': 'application/json',
      'cache-control': 'no-store',
    },
  })
}