lukasa1993/tanstack-skills-extracted/tanstack-ai-persistence-build-prisma-adapter

model ChatThread {
  threadId     String @id @map("thread_id")
  messagesJson String @map("messages_json")
  updatedAt    BigInt @map("updated_at")
 
  @@map("chat_threads")
}
 
model ChatRun {
  runId           String  @id @map("run_id")
  threadId        String  @map("thread_id")
  status          String
  startedAt       BigInt  @map("started_at")
  finishedAt      BigInt? @map("finished_at")
  error           String?
  errorCode       String? @map("error_code")
  usageJson       String? @map("usage_json")
  sandboxKey      String? @map("sandbox_key")
  detachedSince   BigInt? @map("detached_since")
  cancelRequested Boolean? @map("cancel_requested")
  driverEpoch     Int?     @map("driver_epoch")
  parentRunId     String?  @map("parent_run_id")
  subagentRunId   String?  @map("subagent_run_id")
  name            String?
 
  @@index([threadId, status])
  @@index([threadId, startedAt])
  // Powers listByParentRun: children of one parent, oldest startedAt first.
  @@index([parentRunId, startedAt])
  // Powers listReclaimable: status = 'running' AND detachedSince <= cutoff.
  @@index([status, detachedSince])
  @@map("chat_runs")
}
 
model ChatInterrupt {
  interruptId  String  @id @map("interrupt_id")
  runId        String  @map("run_id")
  threadId     String  @map("thread_id")
  status       String
  requestedAt  BigInt  @map("requested_at")
  resolvedAt   BigInt? @map("resolved_at")
  payloadJson  String  @map("payload_json")
  responseJson String? @map("response_json")
 
  @@index([threadId, requestedAt])
  @@map("chat_interrupts")
}
 
model ChatMetadata {
  namespace String
  key       String
  valueJson String @map("value_json")
 
  @@id([namespace, key])
  @@map("chat_metadata")
}

import { defineAIPersistence } from '@tanstack/ai-persistence'
import type {
  ChatInterrupt,
  ChatRun,
  Prisma,
  PrismaClient,
} from '@prisma/client'
import type { ModelMessage, TokenUsage } from '@tanstack/ai'
import type {
  ChatPersistence,
  InterruptRecord,
  InterruptStatus,
  InterruptStore,
  MessageStore,
  MetadataStore,
  RunRecord,
  RunStatus,
  RunStore,
} from '@tanstack/ai-persistence'
import { prisma } from '@/lib/prisma'
// Trusts the shape the stores themselves wrote — nothing else writes these
// columns.
function parseJson<T>(raw: string): T {
  return JSON.parse(raw)
}
const RUN_STATUSES: ReadonlyArray<RunStatus> = [
  'running',
  'interrupted',
  'completed',
  'failed',
  'aborted',
]
const INTERRUPT_STATUSES: ReadonlyArray<InterruptStatus> = [
  'pending',
  'resolved',
  'cancelled',
]
// The column is a plain String, so narrow instead of trusting it.
function toRunStatus(value: string): RunStatus {
  const status = RUN_STATUSES.find((candidate) => candidate === value)
  if (!status) throw new Error(`Unknown run status: ${value}`)
  return status
}
function toInterruptStatus(value: string): InterruptStatus {
  const status = INTERRUPT_STATUSES.find((candidate) => candidate === value)
  if (!status) throw new Error(`Unknown interrupt status: ${value}`)
  return status
}
// Records omit absent optionals so they compare cleanly against the reference
// in-memory backend.
function mapRun(row: ChatRun): RunRecord {
  return {
    runId: row.runId,
    threadId: row.threadId,
    status: toRunStatus(row.status),
    startedAt: Number(row.startedAt),
    ...(row.finishedAt != null ? { finishedAt: Number(row.finishedAt) } : {}),
    ...(row.error != null
      ? {
          error: {
            message: row.error,
            ...(row.errorCode != null ? { code: row.errorCode } : {}),
          },
        }
      : {}),
    ...(row.usageJson != null
      ? { usage: parseJson<TokenUsage>(row.usageJson) }
      : {}),
    ...(row.sandboxKey != null ? { sandboxKey: row.sandboxKey } : {}),
    ...(row.detachedSince != null
      ? { detachedSince: Number(row.detachedSince) }
      : {}),
    ...(row.cancelRequested != null
      ? { cancelRequested: row.cancelRequested }
      : {}),
    ...(row.driverEpoch != null ? { driverEpoch: row.driverEpoch } : {}),
    ...(row.parentRunId != null ? { parentRunId: row.parentRunId } : {}),
    ...(row.subagentRunId != null ? { subagentRunId: row.subagentRunId } : {}),
    ...(row.name != null ? { name: row.name } : {}),
  }
}
function mapInterrupt(row: ChatInterrupt): InterruptRecord {
  return {
    interruptId: row.interruptId,
    runId: row.runId,
    threadId: row.threadId,
    status: toInterruptStatus(row.status),
    requestedAt: Number(row.requestedAt),
    payload: parseJson<Record<string, unknown>>(row.payloadJson),
    ...(row.resolvedAt != null ? { resolvedAt: Number(row.resolvedAt) } : {}),
    ...(row.responseJson != null
      ? { response: parseJson<unknown>(row.responseJson) }
      : {}),
  }
}
function createMessageStore(db: PrismaClient): MessageStore {
  return {
    async loadThread(threadId) {
      const row = await db.chatThread.findUnique({ where: { threadId } })
      // Unknown thread is [], never null.
      return row ? parseJson<Array<ModelMessage>>(row.messagesJson) : []
    },
    // Full overwrite — `messages` is the complete authoritative transcript.
    async saveThread(threadId, messages) {
      const messagesJson = JSON.stringify(messages)
      const updatedAt = BigInt(Date.now())
      await db.chatThread.upsert({
        where: { threadId },
        create: { threadId, messagesJson, updatedAt },
        update: { messagesJson, updatedAt },
      })
    },
  }
}