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'
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',
]
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
}
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 } : {}),
}
}
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 } })
return row ? parseJson<Array<ModelMessage>>(row.messagesJson) : []
},
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 },
})
},
}
}
function createRunStore(db: PrismaClient): RunStore {
return {
async get(runId) {
const row = await db.chatRun.findUnique({ where: { runId } })
return row ? mapRun(row) : null
},
async createOrResume({ runId, threadId, startedAt, status }) {
const row = await db.chatRun.upsert({
where: { runId },
create: {
runId,
threadId,
status: status ?? 'running',
startedAt: BigInt(startedAt),
},
update: {},
})
return mapRun(row)
},
async update(runId, patch) {
const data: Prisma.ChatRunUpdateManyMutationInput = {}
if (patch.status !== undefined) data.status = patch.status
if (patch.finishedAt !== undefined) {
data.finishedAt = BigInt(patch.finishedAt)
}
if (patch.error !== undefined) {
data.error = patch.error.message
data.errorCode = patch.error.code ?? null
}
if (patch.usage !== undefined)
data.usageJson = JSON.stringify(patch.usage)
if ('sandboxKey' in patch) data.sandboxKey = patch.sandboxKey ?? null
if ('detachedSince' in patch) {
data.detachedSince =
patch.detachedSince === undefined ? null : BigInt(patch.detachedSince)
}
if ('cancelRequested' in patch)
data.cancelRequested = patch.cancelRequested ?? null
if ('driverEpoch' in patch) data.driverEpoch = patch.driverEpoch ?? null
if (Object.keys(data).length === 0) return
await db.chatRun.updateMany({ where: { runId }, data })
},
async findActiveRun(threadId) {
const row = await db.chatRun.findFirst({
where: { threadId, status: 'running' },
orderBy: { startedAt: 'desc' },
})
return row ? mapRun(row) : null
},
async listByThread(threadId) {
const rows = await db.chatRun.findMany({
where: { threadId },
orderBy: { startedAt: 'asc' },
})
return rows.map(mapRun)
},
async listReclaimable({ now, ttlMs }) {
const cutoff = BigInt(now - ttlMs)
const rows = await db.chatRun.findMany({
where: {
status: 'running',
detachedSince: { not: null, lte: cutoff },
},
})
return rows.map(mapRun)
},
}
}
function createInterruptStore(db: PrismaClient): InterruptStore {
const listWhere = async (where: Prisma.ChatInterruptWhereInput) => {
const rows = await db.chatInterrupt.findMany({
where,
orderBy: { requestedAt: 'asc' },
})
return rows.map(mapInterrupt)
}
return {
async create(record) {
await db.chatInterrupt.upsert({
where: { interruptId: record.interruptId },
create: {
interruptId: record.interruptId,
runId: record.runId,
threadId: record.threadId,
status: 'pending',
requestedAt: BigInt(record.requestedAt),
payloadJson: JSON.stringify(record.payload),
...(record.response !== undefined
? { responseJson: JSON.stringify(record.response) }
: {}),
},
update: {},
})
},
async resolve(interruptId, response) {
await db.chatInterrupt.updateMany({
where: { interruptId },
data: {
status: 'resolved',
resolvedAt: BigInt(Date.now()),
...(response !== undefined
? { responseJson: JSON.stringify(response) }
: {}),
},
})
},
async cancel(interruptId) {
await db.chatInterrupt.updateMany({
where: { interruptId },
data: { status: 'cancelled', resolvedAt: BigInt(Date.now()) },
})
},
async get(interruptId) {
const row = await db.chatInterrupt.findUnique({ where: { interruptId } })
return row ? mapInterrupt(row) : null
},
list: (threadId) => listWhere({ threadId }),
listPending: (threadId) => listWhere({ threadId, status: 'pending' }),
listByRun: (runId) => listWhere({ runId }),
listPendingByRun: (runId) => listWhere({ runId, status: 'pending' }),
}
}
function createMetadataStore(db: PrismaClient): MetadataStore {
return {
async get(namespace, key) {
const row = await db.chatMetadata.findUnique({
where: { namespace_key: { namespace, key } },
})
return row ? parseJson<unknown>(row.valueJson) : null
},
async set(namespace, key, value) {
if (value == null) {
throw new TypeError(
`Cannot store ${value} for (${namespace}, ${key}) — use delete() to clear metadata.`,
)
}
const valueJson = JSON.stringify(value)
await db.chatMetadata.upsert({
where: { namespace_key: { namespace, key } },
create: { namespace, key, valueJson },
update: { valueJson },
})
},
async delete(namespace, key) {
await db.chatMetadata.deleteMany({ where: { namespace, key } })
},
}
}
export const chatPersistence: ChatPersistence = defineAIPersistence({
stores: {
messages: createMessageStore(prisma),
runs: createRunStore(prisma),
interrupts: createInterruptStore(prisma),
metadata: createMetadataStore(prisma),
},
})
