ai-novel-studio/server/utils/jobs.ts

import { eq } from 'drizzle-orm'
import type { JobInfo, JobType, ProjectEvent } from '#shared/types'

type JobRow = typeof tables.jobs.$inferSelect

export function toJobInfo(row: JobRow): JobInfo {
  const { prompt: _p, projectId: _pid, ...rest } = row
  return rest
}

/** 一次 AI 调用（或一组调用）的记录：持久化到数据库，并实时推送给前端 */
export class Job {
  readonly id = newId()
  private row: JobRow

  constructor(readonly projectId: string, type: JobType, opts: { chapterId?: string | null, agentName?: string, modelLabel?: string } = {}) {
    const db = useDb()
    this.row = db.insert(tables.jobs).values({
      id: this.id,
      projectId,
      type,
      chapterId: opts.chapterId ?? null,
      agentName: opts.agentName ?? '',
      modelLabel: opts.modelLabel ?? ''
    }).returning().get()
    this.emit()
  }

  private update(patch: Partial<JobRow>) {
    this.row = useDb().update(tables.jobs).set(patch).where(eq(tables.jobs.id, this.id)).returning().get()!
    this.emit()
  }

  private emit() {
    emitProject(this.projectId, { type: 'job', job: toJobInfo(this.row) })
  }

  log(message: string) {
    const line = `[${new Date().toLocaleTimeString('zh-CN', { hour12: false })}] ${message}`
    this.update({ log: this.row.log ? `${this.row.log}\n${line}` : line })
  }

  setPrompt(system: string, prompt: string) {
    this.update({ prompt: { system, prompt } })
  }

  addUsage(usage: { inputTokens?: number, outputTokens?: number } | undefined) {
    if (!usage) return
    this.update({
      inputTokens: this.row.inputTokens + (usage.inputTokens ?? 0),
      outputTokens: this.row.outputTokens + (usage.outputTokens ?? 0)
    })
  }

  done() {
    this.update({ status: 'done', finishedAt: new Date() })
  }

  fail(err: unknown) {
    const cancelled = isAbort(err)
    this.update({ status: cancelled ? 'cancelled' : 'error', error: cancelled ? '' : describeError(err), finishedAt: new Date() })
  }
}

export function isAbort(err: unknown) {
  const e = err as { name?: string } | undefined
  return e?.name === 'AbortError' || e?.name === 'ResponseAborted' || (err instanceof Error && /aborted|abort/i.test(err.message) && /signal|user/i.test(err.message))
}

/* ---------------- 项目级互斥运行 ---------------- */

interface ActiveRun {
  label: string
  controller: AbortController
  startedAt: number
}

const activeRuns = new Map<string, ActiveRun>()

export function getActiveRun(projectId: string) {
  const run = activeRuns.get(projectId)
  return run ? { label: run.label, startedAt: run.startedAt } : null
}

/**
 * 每个项目同一时间只允许一个生成任务（保证章节连贯、避免状态竞争）。
 * 任务在后台执行，接口立即返回；进度通过 SSE 推送。
 */
export function runExclusive(projectId: string, label: string, fn: (signal: AbortSignal) => Promise<void>) {
  if (activeRuns.has(projectId)) {
    throw createError({ statusCode: 409, message: `当前项目正在执行「${activeRuns.get(projectId)!.label}」，请等待完成或先停止` })
  }
  const controller = new AbortController()
  activeRuns.set(projectId, { label, controller, startedAt: Date.now() })
  emitProject(projectId, { type: 'pipeline', running: true, message: label })
  fn(controller.signal)
    .catch((err) => {
      if (!isAbort(err)) {
        console.error(`[run:${label}]`, err)
        emitProject(projectId, { type: 'log', level: 'error', message: `${label}失败：${describeError(err)}` })
      }
    })
    .finally(() => {
      activeRuns.delete(projectId)
      emitProject(projectId, { type: 'pipeline', running: false })
    })
}

export function stopRun(projectId: string) {
  const run = activeRuns.get(projectId)
  if (!run) return false
  run.controller.abort()
  return true
}

export function emitAll(projectId: string, ...events: ProjectEvent[]) {
  for (const e of events) emitProject(projectId, e)
}