Files
aigen/server/utils/generationLog.ts
T

160 lines
5.2 KiB
TypeScript

import { existsSync, mkdirSync, readFileSync, renameSync, writeFileSync } from 'node:fs'
import { join } from 'node:path'
import { ownerRoot } from '~/server/utils/library'
import type { StudioJobKind, StudioJobPayload } from '~/server/utils/studioQueue'
export type GenerationLogStatus = 'queued' | 'running' | 'complete' | 'error' | 'cancelled'
export type GenerationLogEntry = {
id: string
studioJobId?: string
createdAt: number
updatedAt: number
kind: StudioJobKind
name: string
prompt: string
status: GenerationLogStatus
lastError?: string
payload: StudioJobPayload
}
const MAX_ENTRIES = 80
const writeChains = new Map<string, Promise<unknown>>()
function logPath(owner: string) {
return join(ownerRoot(owner), 'generation-log.json')
}
function ensureOwner(owner: string) {
mkdirSync(ownerRoot(owner), { recursive: true })
}
function readLog(owner: string): GenerationLogEntry[] {
ensureOwner(owner)
const path = logPath(owner)
if (!existsSync(path)) return []
try {
const parsed = JSON.parse(readFileSync(path, 'utf8'))
return Array.isArray(parsed) ? parsed.filter(item => item && typeof item === 'object' && item.id && item.payload) : []
} catch {
return []
}
}
function writeLog(owner: string, entries: GenerationLogEntry[]) {
ensureOwner(owner)
const path = logPath(owner)
const tmp = `${path}.tmp`
writeFileSync(tmp, JSON.stringify(entries, null, 2))
renameSync(tmp, path)
}
function mutate<T>(owner: string, fn: (entries: GenerationLogEntry[]) => T): Promise<T> {
const prev = writeChains.get(owner) || Promise.resolve()
const run = prev.then(() => {
const entries = readLog(owner)
const result = fn(entries)
writeLog(owner, entries)
return result
})
writeChains.set(owner, run.then(() => undefined, () => undefined))
return run
}
function snippetName(payload: StudioJobPayload, kind: StudioJobKind) {
const named = String(payload.name || '').trim()
if (named) return named.slice(0, 80)
if (kind === 'music') return String(payload.prompt || 'Track').slice(0, 80)
return String(payload.prompt || 'Untitled').slice(0, 80)
}
export function listGenerationLog(owner: string, limit = 40) {
return readLog(owner).slice(0, Math.max(1, Math.min(MAX_ENTRIES, limit)))
}
export function getGenerationLogEntry(owner: string, id: string) {
const entry = readLog(owner).find(item => item.id === id)
if (!entry) throw createError({ statusCode: 404, statusMessage: 'Generation log entry not found' })
return entry
}
export async function recordGenerationQueued(params: {
ownerKey: string
studioJobId: string
kind: StudioJobKind
payload: StudioJobPayload
}) {
const now = Date.now()
return mutate(params.ownerKey, (entries) => {
const existing = entries.find(item => item.studioJobId === params.studioJobId)
if (existing) {
existing.updatedAt = now
existing.status = 'queued'
existing.payload = structuredClone(params.payload)
existing.name = snippetName(params.payload, params.kind)
existing.prompt = params.payload.prompt || existing.prompt
existing.kind = params.kind
existing.lastError = undefined
entries.splice(entries.indexOf(existing), 1)
entries.unshift(existing)
return existing
}
const entry: GenerationLogEntry = {
id: crypto.randomUUID(),
studioJobId: params.studioJobId,
createdAt: now,
updatedAt: now,
kind: params.kind,
name: snippetName(params.payload, params.kind),
prompt: params.payload.prompt || '',
status: 'queued',
payload: structuredClone(params.payload)
}
entries.unshift(entry)
entries.splice(MAX_ENTRIES)
return entry
})
}
export async function updateGenerationLogByStudioJob(
owner: string,
studioJobId: string,
patch: { status?: GenerationLogStatus; lastError?: string; payload?: StudioJobPayload }
) {
return mutate(owner, (entries) => {
const entry = entries.find(item => item.studioJobId === studioJobId)
if (!entry) return null
entry.updatedAt = Date.now()
if (patch.status) entry.status = patch.status
if (patch.lastError !== undefined) entry.lastError = patch.lastError || undefined
if (patch.payload) {
entry.payload = structuredClone(patch.payload)
entry.name = snippetName(patch.payload, entry.kind)
entry.prompt = patch.payload.prompt || entry.prompt
}
return structuredClone(entry)
})
}
export async function requeueGenerationLog(owner: string, id: string) {
const entry = getGenerationLogEntry(owner, id)
const { addStudioJob, kickStudioQueue } = await import('~/server/utils/studioQueue')
const studio = await addStudioJob({
ownerKey: owner,
familyId: crypto.randomUUID(),
kind: entry.kind,
payload: structuredClone(entry.payload)
})
await kickStudioQueue()
return { studioJobId: studio.id, entry: listGenerationLog(owner, 1)[0] || getGenerationLogEntry(owner, id) }
}
export async function deleteGenerationLogEntry(owner: string, id: string) {
return mutate(owner, (entries) => {
const index = entries.findIndex(item => item.id === id)
if (index < 0) throw createError({ statusCode: 404, statusMessage: 'Generation log entry not found' })
const [removed] = entries.splice(index, 1)
return removed
})
}