98 lines
9.1 KiB
TypeScript
98 lines
9.1 KiB
TypeScript
import { randomUUID, createHash } from 'node:crypto'
|
|
import { prisma } from '@/lib/prisma'
|
|
import { desktopAssert, desktopId } from './contract.mjs'
|
|
import { taskUuid, taskBusy, validateCreationTask, validateTaskPatch } from './task-contract.mjs'
|
|
import { validateStoredAiModelProtocols } from '@/lib/ai-model-protocol'
|
|
import { desktopExecutionModels } from './models'
|
|
|
|
type Auth = { id: string; account: { id: string } }
|
|
type Row = Record<string, any>
|
|
const hash = (value: string) => createHash('sha256').update(value).digest('hex')
|
|
const columns = 'id,payload,revision,lease_session_id AS "sessionId",lease_until AS "leaseUntil",lease_mode AS "leaseMode"'
|
|
const json = (value: unknown) => JSON.stringify(value)
|
|
async function locked<T>(auth: Auth, projectId: string, action: (tx: any) => Promise<T>) {
|
|
desktopId(projectId)
|
|
return prisma.$transaction(async tx => {
|
|
// Serialize by owner: a second device cannot race a new paid request.
|
|
await tx.$queryRawUnsafe('SELECT 1 FROM (SELECT pg_advisory_xact_lock(hashtext($1))) l', 'desktop-tasks:' + auth.account.id)
|
|
const projects = await tx.$queryRawUnsafe<Row[]>('SELECT document FROM desktop_workspace_projects WHERE id=$1 AND owner_id=$2 AND deleted_at IS NULL', projectId, auth.account.id)
|
|
desktopAssert(projects[0], '项目不存在或无权访问', 404)
|
|
await tx.$executeRawUnsafe(`UPDATE desktop_creation_tasks SET status=CASE WHEN dispatched THEN 'unknown' ELSE 'cancelled' END,
|
|
payload=payload || jsonb_build_object('status',CASE WHEN dispatched THEN 'unknown' ELSE 'cancelled' END,'message',CASE WHEN dispatched THEN '原设备已离线，结果待核查；只能取回原结果' ELSE '原设备未发送请求，任务已取消' END,'updatedAt',NOW()),
|
|
lease_hash=NULL,lease_session_id=NULL,lease_until=NULL,revision=revision+1,updated_at=NOW()
|
|
WHERE owner_id=$1 AND status IN ('queued','running') AND lease_until<NOW()`, auth.account.id)
|
|
return action(tx)
|
|
})
|
|
}
|
|
const dto = (row: Row, auth: Auth) => ({ task: row.payload, revision: row.revision, leaseUntil: row.leaseUntil,
|
|
leasedHere: row.sessionId === auth.id, leaseMode: row.leaseMode })
|
|
export async function listCreationTasks(auth: Auth, projectId: string) {
|
|
return locked(auth,projectId,async tx => {
|
|
const rows: Row[] = await tx.$queryRawUnsafe(`SELECT id,revision,lease_session_id AS "sessionId",lease_until AS "leaseUntil",lease_mode AS "leaseMode",
|
|
jsonb_build_object('id',id,'stage',payload->'stage','status',status,'screenId',payload->'screenId','dispatched',dispatched,
|
|
'requestId',payload->'requestId','remoteTaskId',payload->'remoteTaskId','artifact',payload->'artifact','applied',payload->'applied','generationPlanHash',payload->'input'->'planHash',
|
|
'createdAt',payload->'createdAt','updatedAt',payload->'updatedAt','message',payload->'message',
|
|
'model',jsonb_build_object('id',payload->'model'->'id','key',payload->'model'->'key'),
|
|
'remoteHasSaved',COALESCE(payload->>'text','')<>'' OR payload->'artifact'<>'null'::jsonb,'remoteStub',true) AS payload
|
|
FROM desktop_creation_tasks WHERE owner_id=$1 AND project_id=$2 ORDER BY created_at LIMIT 80`,auth.account.id,projectId)
|
|
return { items: rows.map(row => dto(row,auth)) }
|
|
})
|
|
}
|
|
export async function getCreationTask(auth:Auth,projectId:string,id:string){
|
|
desktopAssert(taskUuid(id),'任务编号无效')
|
|
return locked(auth,projectId,async tx=>{
|
|
const rows:Row[]=await tx.$queryRawUnsafe(`SELECT ${columns} FROM desktop_creation_tasks WHERE id=$1 AND project_id=$2 AND owner_id=$3`,id,projectId,auth.account.id)
|
|
desktopAssert(rows[0],'任务不存在',404);return dto(rows[0],auth)
|
|
})
|
|
}
|
|
export async function createCreationTask(auth: Auth, projectId: string, value: any) {
|
|
const task = validateCreationTask(value,projectId)
|
|
desktopAssert(validateStoredAiModelProtocols([task.model.protocol]).length === 1, '模型协议无效')
|
|
const models=await desktopExecutionModels(),purpose=task.stage==='facts'?'analysis':task.stage==='image'?'image':'copy'
|
|
desktopAssert(models.models.some(model=>model.stage===purpose&&model.key===task.model.key&&model.fingerprint===task.model.fingerprint),'模型配置已变化，请重新读取管理员配置',409,'model_changed')
|
|
return locked(auth,projectId,async tx => {
|
|
const existing: Row[] = await tx.$queryRawUnsafe(`SELECT ${columns} FROM desktop_creation_tasks WHERE id=$1`,task.id)
|
|
desktopAssert(!existing.length,'任务已登记，请读取原任务，不能重复提交',409,'task_exists')
|
|
const blocked: Row[] = await tx.$queryRawUnsafe(`SELECT id FROM desktop_creation_tasks WHERE owner_id=$1 AND (project_id=$2 AND status IN ('queued','running','received','unknown') OR lease_until>NOW()) LIMIT 1`,auth.account.id,projectId)
|
|
desktopAssert(!blocked.length,'此账户已有运行任务或本项目有未知请求，请先读取原结果',409,'task_busy')
|
|
const count: Row[] = await tx.$queryRawUnsafe('SELECT COUNT(*)::int AS count FROM desktop_creation_tasks WHERE owner_id=$1 AND project_id=$2',auth.account.id,projectId)
|
|
desktopAssert(count[0].count<80,'本项目已达到80条任务记录上限')
|
|
const project: Row[] = await tx.$queryRawUnsafe('SELECT document FROM desktop_workspace_projects WHERE id=$1 AND owner_id=$2',projectId,auth.account.id)
|
|
// JSONB key order differs from JSON; compare structurally in PostgreSQL.
|
|
const same: Row[] = await tx.$queryRawUnsafe('SELECT $1::jsonb=$2::jsonb AS same',json(project[0].document),json(task.document))
|
|
desktopAssert(same[0].same,'先同步当前项目及资料，再登记任务',409,'revision_conflict')
|
|
const leaseToken=randomUUID()
|
|
const rows: Row[]=await tx.$queryRawUnsafe(`INSERT INTO desktop_creation_tasks(id,project_id,owner_id,payload,status,lease_session_id,lease_hash,lease_mode,lease_until)
|
|
VALUES($1,$2,$3,$4::jsonb,'queued',$5,$6,'execute',NOW()+INTERVAL '90 seconds') RETURNING ${columns}`,task.id,projectId,auth.account.id,json(task),auth.id,hash(leaseToken))
|
|
return {...dto(rows[0],auth),leaseToken}
|
|
})
|
|
}
|
|
export async function claimCreationTask(auth:Auth,projectId:string,id:string) {
|
|
desktopAssert(taskUuid(id),'任务编号无效')
|
|
return locked(auth,projectId,async tx=>{
|
|
const rows:Row[]=await tx.$queryRawUnsafe(`SELECT ${columns} FROM desktop_creation_tasks WHERE id=$1 AND project_id=$2 AND owner_id=$3`,id,projectId,auth.account.id)
|
|
const row=rows[0];desktopAssert(row,'任务不存在',404)
|
|
desktopAssert(!row.leaseUntil||new Date(row.leaseUntil).getTime()<=Date.now(),'原设备仍持有任务，请等待它结束或租约到期',409,'task_leased')
|
|
desktopAssert(['unknown','received','succeeded','failed'].includes(row.payload.status),'此任务没有可恢复的原结果')
|
|
desktopAssert(row.payload.dispatched || row.payload.text || row.payload.artifact,'此任务尚未发送，无需恢复')
|
|
const busy:Row[]=await tx.$queryRawUnsafe('SELECT id FROM desktop_creation_tasks WHERE owner_id=$1 AND lease_until>NOW() LIMIT 1',auth.account.id)
|
|
desktopAssert(!busy.length,'另一设备仍在执行任务',409,'task_leased')
|
|
const leaseToken=randomUUID()
|
|
const updated:Row[]=await tx.$queryRawUnsafe(`UPDATE desktop_creation_tasks SET lease_session_id=$4,lease_hash=$5,lease_mode='recover',lease_until=NOW()+INTERVAL '90 seconds',revision=revision+1 WHERE id=$1 AND project_id=$2 AND owner_id=$3 RETURNING ${columns}`,id,projectId,auth.account.id,auth.id,hash(leaseToken))
|
|
return {...dto(updated[0],auth),leaseToken}
|
|
})
|
|
}
|
|
export async function checkpointCreationTask(auth:Auth,projectId:string,id:string,input:any) {
|
|
desktopAssert(taskUuid(id)&&taskUuid(input?.leaseToken)&&Number.isInteger(input.revision),'任务租约无效')
|
|
return locked(auth,projectId,async tx=>{
|
|
const rows:Row[]=await tx.$queryRawUnsafe(`SELECT ${columns},lease_hash AS "leaseHash" FROM desktop_creation_tasks WHERE id=$1 AND project_id=$2 AND owner_id=$3 AND lease_until>NOW()`,id,projectId,auth.account.id)
|
|
const row=rows[0]
|
|
desktopAssert(row&&row.sessionId===auth.id&&row.leaseHash===hash(input.leaseToken)&&row.revision===input.revision,'任务租约已变化，请只读刷新；不会重新生成',409,'lease_lost')
|
|
const payload=validateTaskPatch(row.payload,input.patch,row.leaseMode==='recover')
|
|
if(payload.artifact){const found:Row[]=await tx.$queryRawUnsafe('SELECT sha256 FROM desktop_workspace_assets WHERE project_id=$1 AND owner_id=$2 AND source=$3',projectId,auth.account.id,payload.artifact.source);desktopAssert(found[0]?.sha256===payload.artifact.sha256,'先上传原图片检查点，再结束任务',409)}
|
|
const released=Object.hasOwn(input.patch,'status')&&['succeeded','failed','unknown','cancelled'].includes(payload.status)
|
|
const updated:Row[]=await tx.$queryRawUnsafe(`UPDATE desktop_creation_tasks SET payload=$4::jsonb,status=$5,dispatched=$6,revision=revision+1,updated_at=NOW(),lease_until=CASE WHEN $7 THEN NULL ELSE NOW()+INTERVAL '90 seconds' END,lease_hash=CASE WHEN $7 THEN NULL ELSE lease_hash END,lease_session_id=CASE WHEN $7 THEN NULL ELSE lease_session_id END WHERE id=$1 AND project_id=$2 AND owner_id=$3 RETURNING ${columns}`,id,projectId,auth.account.id,json({...payload,updatedAt:new Date().toISOString()}),payload.status,payload.dispatched,released)
|
|
return dto(updated[0],auth)
|
|
})
|
|
}
|