AgentServer/apps/server/src/services/runs.ts

// Run 调度：排队、并发控制、执行、中断与重启恢复（docs/design/02-architecture.md §3.2）
import fsp from 'node:fs/promises'
import path from 'node:path'
import { eq, inArray } from 'drizzle-orm'
import PQueue from 'p-queue'
import { z } from 'zod'
import type { AgentEventType, EventPayloads, ProviderId, RunDto, SendMessageBody, SessionDto } from '@agent-server/shared'
import type { DB } from '../db'
import { runs } from '../db/schema'
import { conflict, notFound } from '../lib/errors'
import { newId, nowIso } from '../lib/ids'
import type { AgentProvider, RunInput } from '../providers/types'
import { ProviderRunError } from '../providers/types'
import { resolveAccess } from './access-policy'
import { parse, type AuditService, type KnowledgeService, type ProfileService } from './catalog'
import type { EventStore } from './event-store'
import type { FileService } from './files'
import type { InteractionService } from './interactions'
import type { CommunitySkillService } from './community-skills'
import type { PluginService } from './plugins'
import type { RunTokenRegistry } from '../mcp-gateway/tokens'
import { gatewayUrl } from '../config'
import type { SessionRow, SessionService } from './sessions'
import type { SettingsService } from './settings'
import { buildSystemPrompt } from './prompt'

type RunRow = typeof runs.$inferSelect

const sendSchema: z.ZodType<SendMessageBody> = z.object({
  text: z.string().trim().min(1).max(100_000),
  attachments: z.array(z.string()).max(50).optional(),
})

const IMAGE_MIMES = new Set(['image/png', 'image/jpeg', 'image/gif', 'image/webp'])
const MAX_INLINE_IMAGE = 5 * 1024 * 1024

interface ActiveRun {
  runId: string
  sessionId: string
  controller: AbortController
  status: 'queued' | 'running' | 'waiting_input'
}

interface Logger { info(o: object, m?: string): void; warn(o: object, m?: string): void; error(o: object, m?: string): void; child(b: object): Logger }

export class RunManager {
  private global: PQueue
  private perProvider: Record<ProviderId, PQueue>
  private active = new Map<string, ActiveRun>()

  constructor(
    private deps: {
      db: DB
      log: Logger
      settings: SettingsService
      sessions: SessionService
      profiles: ProfileService
      knowledge: KnowledgeService
      files: FileService
      events: EventStore
      interactions: InteractionService
      audit: AuditService
      plugins: PluginService
      communitySkills: CommunitySkillService
      gatewayTokens: RunTokenRegistry
      providers: Partial<Record<ProviderId, AgentProvider>>
    },
  ) {
    const c = deps.settings.get().limits.concurrency
    this.global = new PQueue({ concurrency: c.global })
    this.perProvider = { claude: new PQueue({ concurrency: Math.max(c.claude, 1) }), codex: new PQueue({ concurrency: Math.max(c.codex, 1) }) }
    deps.settings.on('changed', s => {
      this.global.concurrency = s.limits.concurrency.global
      this.perProvider.claude.concurrency = Math.max(s.limits.concurrency.claude, 1)
      this.perProvider.codex.concurrency = Math.max(s.limits.concurrency.codex, 1)
    })
    deps.sessions.statusOf = id => this.sessionStatus(id)
    this.recoverAfterRestart()
  }

  sessionStatus(sessionId: string): SessionDto['status'] {
    const a = this.activeForSession(sessionId)
    if (!a) return 'idle'
    return a.status === 'waiting_input' ? 'waiting_input' : 'running'
  }

  activeForSession(sessionId: string): ActiveRun | undefined {
    for (const a of this.active.values()) if (a.sessionId === sessionId) return a
    return undefined
  }

  stats() {
    return { running: this.global.pending, queued: this.global.size + this.perProvider.claude.size + this.perProvider.codex.size, active: this.active.size }
  }

  /** 发送消息：登记 Run 并排队，立即返回 runId */
  send(session: SessionRow, body: SendMessageBody): RunDto {
    const data = parse(sendSchema, body)
    if (this.activeForSession(session.id)) throw conflict('session_busy', '会话正在运行，请等待完成或先中断')
    const provider = this.deps.providers[session.provider]
    if (!provider || !this.deps.settings.get().providers[session.provider].enabled) throw conflict('provider_unavailable', `引擎 ${session.provider} 当前不可用`)
    this.deps.sessions.assertProfileAllowed(session)
    const attachments = this.deps.files.getMany(session.id, data.attachments ?? [])

    const runId = newId('run')
    this.deps.db.insert(runs).values({ id: runId, sessionId: session.id, status: 'queued' }).run()
    const active: ActiveRun = { runId, sessionId: session.id, controller: new AbortController(), status: 'queued' }
    this.active.set(runId, active)

    this.deps.sessions.ensureTitle(session.id, data.text)
    this.deps.events.append(session.id, runId, 'user.message', {
      text: data.text,
      attachments: attachments.map(f => ({ fileId: f.id, name: path.basename(f.relPath), mime: f.mime, size: f.size })),
    })
    const queue = this.perProvider[session.provider]
    this.deps.events.append(session.id, runId, 'run.queued', { position: queue.size + queue.pending })

    void queue.add(() => this.global.add(() => this.execute(runId, session.id, provider, data.text, attachments.map(f => f.id))))
      .catch(err => this.deps.log.error({ err, runId }, 'run task crashed'))
    return this.get(runId)
  }

  interrupt(runId: string, sessionId?: string) {
    const a = this.active.get(runId)
    if (!a || (sessionId && a.sessionId !== sessionId)) {
      const row = this.deps.db.select().from(runs).where(eq(runs.id, runId)).get()
      if (!row || (sessionId && row.sessionId !== sessionId)) throw notFound('Run')
      return // 已结束
    }
    a.controller.abort()
    this.deps.interactions.cancelRun(runId)
    if (a.status === 'queued') this.finish(a, 'interrupted', { reason: 'user' })
  }

  get(runId: string): RunDto {
    const r = this.deps.db.select().from(runs).where(eq(runs.id, runId)).get()
    if (!r) throw notFound('Run')
    return toRunDto(r)
  }

  private async execute(runId: string, sessionId: string, provider: AgentProvider, text: string, attachmentIds: string[]) {
    const a = this.active.get(runId)
    if (!a || a.controller.signal.aborted) return
    const { deps } = this
    const log = deps.log.child({ runId, sessionId })
    const emit = <T extends AgentEventType>(type: T, payload: EventPayloads[T]) => { deps.events.append(sessionId, runId, type, payload) }
    const startedAt = Date.now()

    a.status = 'running'
    deps.db.update(runs).set({ status: 'running', startedAt: nowIso() }).where(eq(runs.id, runId)).run()

    let scanTimer: NodeJS.Timeout | undefined
    let gatewayToken: string | null = null
    let session: SessionRow | undefined
    const scanFiles = async () => {
      if (!session) return
      for (const c of await deps.files.scan(sessionId, session.workspaceDir)) {
        emit('file.changed', { fileId: c.file.id, path: c.file.path, kind: c.kind, size: c.file.size, mime: c.file.mime })
      }
    }

    try {
      session = deps.sessions.getRow(sessionId)
      const settings = deps.settings.get()
      const { profile, permissionMode } = deps.sessions.effectiveProfile(session)
      const knowledgeRoots = deps.knowledge.list()
      const access = resolveAccess({
        policy: profile.fileAccess,
        knowledgeRoots,
        selectedKnowledgeRootIds: profile.knowledgeRootIds,
        workspaceDir: session.workspaceDir,
        settings,
      })
      emit('run.started', { provider: session.provider, model: session.model })
      const sessionRow = session
      gatewayToken = deps.gatewayTokens.issue({
        sessionId, runId, provider: session.provider, workspaceDir: session.workspaceDir, access, permissionMode, settings,
        knowledgeRoots: knowledgeRoots.filter(k => k.enabled && profile.knowledgeRootIds.includes(k.id)),
        audit: (kind, summary, detail) => deps.audit.log({ actor: `agent:${sessionRow.provider}`, sessionId, kind, summary, detail }),
      }, { stablePerSession: session.provider === 'codex' })

      const result = await provider.run({
        sessionId,
        runId,
        model: session.model,
        permissionMode,
        providerSessionId: session.providerSessionId,
        workspaceDir: session.workspaceDir,
        systemPromptAppend: buildSystemPrompt({
          workspaceDir: session.workspaceDir,
          settings,
          knowledgeRoots: knowledgeRoots.filter(k => k.enabled && profile.knowledgeRootIds.includes(k.id)),
          profilePrompt: profile.systemPrompt,
          sessionPrompt: session.systemPromptAppend,
          permissionMode,
        }),
        allowedTools: profile.allowedTools,
        disallowedTools: profile.disallowedTools,
        skills: profile.skills,
        // 插件只在 Claude 中加载；专有许可的插件不能提供给其他引擎（ADR-0013）。第三方技能 Codex 经 extraRoots 加载（ADR-0017）
        plugins: session.provider === 'claude' ? await this.claudePlugins() : [],
        // Codex 插件与应用连接器使用账号绑定的数据，只在后台 Playground（source = admin）中开放
        codexPlugins: session.provider === 'codex' && session.source === 'admin',
        input: await this.buildInput(session, text, attachmentIds),
        access,
        settings,
        signal: a.controller.signal,
        emit: (type, payload) => {
          emit(type, payload)
          if (type === 'tool.completed') {
            clearTimeout(scanTimer)
            scanTimer = setTimeout(() => void scanFiles().catch(() => {}), 300)
          }
        },
        interact: async req => {
          a.status = 'waiting_input'
          deps.db.update(runs).set({ status: 'waiting_input' }).where(eq(runs.id, runId)).run()
          try {
            return await deps.interactions.request(sessionId, runId, req, a.controller.signal)
          } finally {
            if (!deps.interactions.hasPending(runId) && !a.controller.signal.aborted) {
              a.status = 'running'
              deps.db.update(runs).set({ status: 'running' }).where(eq(runs.id, runId)).run()
            }
          }
        },
        setProviderSessionId: id => deps.sessions.setProviderSessionId(sessionId, id),
        audit: (kind, summary, detail) => deps.audit.log({ actor: `agent:${session!.provider}`, sessionId, kind, summary, detail }),
        log,
        gateway: { url: gatewayUrl(), token: gatewayToken },
      })

      clearTimeout(scanTimer)
      await scanFiles().catch(err => log.warn({ err }, 'file scan failed'))
      if (result.usage) {
        emit('usage', result.usage)
        deps.db.update(runs).set({
          inputTokens: result.usage.inputTokens, outputTokens: result.usage.outputTokens,
          cacheReadTokens: result.usage.cacheReadTokens, cacheWriteTokens: result.usage.cacheWriteTokens, costUsd: result.usage.costUsd ?? null,
        }).where(eq(runs.id, runId)).run()
      }
      this.finish(a, 'completed', { durationMs: Date.now() - startedAt })
    } catch (err) {
      clearTimeout(scanTimer)
      await scanFiles().catch(() => {})
      if (a.controller.signal.aborted) {
        this.finish(a, 'interrupted', { reason: 'user' })
      } else {
        const code = err instanceof ProviderRunError ? err.code : 'internal_error'
        const message = err instanceof Error ? err.message : String(err)
        log.error({ err }, 'run failed')
        this.finish(a, 'failed', { code, message })
      }
    } finally {
      if (gatewayToken) deps.gatewayTokens.revoke(gatewayToken)
    }
  }

  private finish(a: ActiveRun, status: 'completed', p: EventPayloads['run.completed']): void
  private finish(a: ActiveRun, status: 'failed', p: EventPayloads['run.failed']): void
  private finish(a: ActiveRun, status: 'interrupted', p: EventPayloads['run.interrupted']): void
  private finish(a: ActiveRun, status: 'completed' | 'failed' | 'interrupted', p: never) {
    if (!this.active.has(a.runId)) return
    this.active.delete(a.runId)
    this.deps.db.update(runs).set({
      status, endedAt: nowIso(), error: status === 'failed' ? (p as EventPayloads['run.failed']).message : null,
    }).where(eq(runs.id, a.runId)).run()
    this.deps.events.append(a.sessionId, a.runId, `run.${status}`, p)
  }

  /** Claude 运行时加载的插件：已启用的插件 + 第三方技能目录 */
  private async claudePlugins() {
    const community = this.deps.communitySkills.pluginRef()
    return [...await this.deps.plugins.loadable().catch(() => []), ...(community ? [community] : [])]
  }

  private async buildInput(session: SessionRow, text: string, attachmentIds: string[]): Promise<RunInput> {
    const input: RunInput = { text, images: [], attachments: [] }
    for (const f of this.deps.files.getMany(session.id, attachmentIds)) {
      const name = path.basename(f.relPath)
      input.attachments.push({ path: f.relPath, name, mime: f.mime, size: f.size })
      if (IMAGE_MIMES.has(f.mime) && f.size <= MAX_INLINE_IMAGE) {
        const buf = await fsp.readFile(path.join(session.workspaceDir, f.relPath))
        input.images.push({ mime: f.mime, base64: buf.toString('base64'), name })
      }
    }
    return input
  }

  private recoverAfterRestart() {
    const stale = this.deps.db.select().from(runs).where(inArray(runs.status, ['queued', 'running', 'waiting_input'])).all()
    for (const r of stale) {
      this.deps.db.update(runs).set({ status: 'interrupted', endedAt: nowIso(), error: '服务重启' }).where(eq(runs.id, r.id)).run()
      try {
        this.deps.events.append(r.sessionId, r.id, 'run.interrupted', { reason: 'restart' })
      } catch {
        // 会话已被删除
      }
    }
    if (stale.length) this.deps.log.warn({ count: stale.length }, '服务重启：未完成的 Run 已标记为 interrupted')
  }
}

function toRunDto(r: RunRow): RunDto {
  return {
    id: r.id, sessionId: r.sessionId, status: r.status, startedAt: r.startedAt, endedAt: r.endedAt, error: r.error,
    usage: r.inputTokens === null ? null : {
      inputTokens: r.inputTokens, outputTokens: r.outputTokens ?? 0, cacheReadTokens: r.cacheReadTokens ?? 0,
      cacheWriteTokens: r.cacheWriteTokens ?? 0, costUsd: r.costUsd ?? undefined,
    },
  }
}

function safe<T>(fn: () => T): T | undefined {
  try {
    return fn()
  } catch {
    return undefined
  }
}