@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 { ModelMessage, UIMessage } from '@tanstack/ai'
import { validateReconstructChatStores } from './types'
import type {
  AIPersistence,
  ChatTranscriptStores,
  MessagePage,
  MessageStore,
} from './types'

const MAX_PAGE_SIZE = 500

/**
 * 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`.
 * `page` is set only when the GET included a valid `limit`. `truncated` is true
 * when older UI messages exist. `cursor` is the opaque `before` token for the
 * next older window.
 */
export interface ReconstructedChat {
  messages: Array<UIMessage>
  activeRun: { runId: string } | null
  interrupts: {
    runId: string
    pending: Array<Record<string, unknown>>
  } | null
  page?: { truncated: false } | { truncated: true; cursor: string }
}

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.
 *
 * Paging is opt-in. A valid `?limit=` (positive integer, capped at 500) returns
 * the newest window of UI messages plus `page`. `?before=` walks to an older
 * window. Invalid `limit` (`0`, negative, NaN) is ignored and the full
 * transcript is returned. `activeRun` and `interrupts` are never paged.
 *
 * 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 requestUrl = new URL(request.url)
  const param = options?.param ?? 'threadId'
  const threadId = requestUrl.searchParams.get(param) ?? ''
  const pageSize = parsePageSize(requestUrl.searchParams.get('limit'))
  const before = parseBefore(requestUrl.searchParams.get('before'))

  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 === ''
      ? []
      : pageSize === undefined
        ? await messageStore.loadThread(threadId)
        : await messageStore.loadThread(threadId, {
            limit: pageSize + 1,
            ...(before === undefined ? {} : { before }),
          })
  // 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 isPaging = pageSize !== undefined && threadId !== ''
  const transcript = !isPaging
    ? {
        messages: modelMessagesToUIMessages(threadMessages(stored)),
      }
    : Array.isArray(stored)
      ? await windowFromArray({
          stored,
          messageStore,
          threadId,
          pageSize,
          before,
        })
      : windowFromMessagePage(stored, pageSize)
  const body: ReconstructedChat = {
    messages: transcript.messages,
    activeRun: active ? { runId: active.runId } : null,
    interrupts: firstPending
      ? {
          runId: firstPending.runId,
          pending: pending.map((record) => record.payload),
        }
      : null,
    ...('page' in transcript ? { page: transcript.page } : {}),
  }
  return new Response(JSON.stringify(body), {
    headers: {
      'content-type': 'application/json',
      'cache-control': 'no-store',
    },
  })
}

function parsePageSize(raw: string | null) {
  if (raw == null) return
  const pageSize = Number(raw)
  const isValidPageSize = Number.isInteger(pageSize) && pageSize > 0
  if (!isValidPageSize) return
  return Math.min(pageSize, MAX_PAGE_SIZE)
}

function parseBefore(raw: string | null) {
  if (raw == null || raw === '') return
  return raw
}

function threadMessages(
  loaded: Array<ModelMessage> | MessagePage,
): Array<ModelMessage> {
  return Array.isArray(loaded) ? loaded : loaded.messages
}

function completePage() {
  return { truncated: false as const }
}

function truncatedPage(cursor: string) {
  return { truncated: true as const, cursor }
}

function pageFromCursor(cursor: string | undefined) {
  if (cursor === undefined || cursor === '') {
    return completePage()
  }
  return truncatedPage(cursor)
}

function newestUiWindow(messages: Array<UIMessage>, pageSize: number) {
  const truncated = messages.length > pageSize
  if (!truncated) {
    return { messages, page: completePage() }
  }
  const uiWindow = messages.slice(messages.length - pageSize)
  return {
    messages: uiWindow,
    page: pageFromCursor(uiWindow[0]?.id),
  }
}

function uiBeforeCursor(messages: Array<UIMessage>, cursor: string) {
  const cut = messages.findIndex((message) => message.id === cursor)
  if (cut === -1) return
  return messages.slice(0, cut)
}

function windowFromMessagePage(page: MessagePage, pageSize: number) {
  const ui = modelMessagesToUIMessages(page.messages)
  if (ui.length > pageSize) {
    // Extra slice uses a library-minted cursor. Keeping the adapter cursor
    // after dropping the oldest row would skip that row on the next GET.
    return newestUiWindow(ui, pageSize)
  }
  if (page.truncated) {
    return {
      messages: ui,
      page: pageFromCursor(page.cursor),
    }
  }
  return { messages: ui, page: completePage() }
}

async function windowFromArray(input: {
  stored: Array<ModelMessage>
  messageStore: MessageStore
  threadId: string
  pageSize: number
  before: string | undefined
}) {
  const { stored, messageStore, threadId, pageSize, before } = input
  if (before === undefined) {
    return newestUiWindow(modelMessagesToUIMessages(stored), pageSize)
  }
  // Array adapters own no cursor. Apply `before` to the full transcript so an
  // adapter that ignored the hint cannot return the same newest page forever.
  const full = threadMessages(await messageStore.loadThread(threadId))
  const older = uiBeforeCursor(modelMessagesToUIMessages(full), before)
  if (older === undefined) {
    return { messages: [], page: truncatedPage(before) }
  }
  return newestUiWindow(older, pageSize)
}