UHD_Media_Tool/apps/server/src/modules/agent/runner.ts

import { join, resolve } from "node:path";
import {
  type AgentRunHandle,
  type AgentSettings,
  buildSystemPrompt,
  type CapabilityRegistry,
  createFakeRuntime,
  type FakeScript,
  PiAgentRuntime,
  type RunOutcome,
} from "@umt/agent";
import {
  type AgentBudget,
  type AgentGoalReport,
  type AgentRunEvent,
  type AgentRunStatus,
  agentRunEventToBusEvent,
  agentRunEventToRow,
  CAPABILITY_CATALOG,
  type CapabilityKind,
  type GlobalSettings,
  KnownCapabilityId,
} from "@umt/contracts";
import { type AgentRunRow, sql } from "@umt/db";
import type { ServerContext } from "../types.ts";
import { DEMO_SCRIPTS, parseFakeScripts } from "./fake-scripts.ts";
import { buildGoalReport } from "./report.ts";

// runAgent 任务的执行体:从设置装配 Pi(或 fake)运行时 → 事件先写 agent_run_event 再推 SSE →
// 崩溃恢复(pendingOperations>0 时先按幂等键对账未知结果再 resume)→ 消费追加消息 → 汇报以业务表为准。

export interface AgentModuleOptions {
  /** Pi JSONL 会话目录(崩溃恢复点;新 worker 需能读到同一目录) */
  sessionDir?: string;
  /** fake 模式脚本(测试注入);缺省读 settings.agent.fakeScript,再缺省用内置演示脚本 */
  fakeScripts?: FakeScript[];
  /** 覆盖 settings.agent.goalPollSec(测试用小值) */
  goalPollSec?: number;
}

export interface AgentEnv {
  ctx: ServerContext;
  opts: AgentModuleOptions;
}

/** "先给方案"与后台调查只允许这些类别 */
export const PLAN_ONLY_KINDS: CapabilityKind[] = ["read", "preview", "propose"];
/** 后台自动调查的能力白名单:只读/预览 + 身份提案 */
export const INVESTIGATE_CAPABILITIES = [
  "catalog.lookup",
  "release.search",
  "release.inspect",
  "inventory.query",
  "task.inspect",
  "identity.resolve",
] as const;

export interface RunScope {
  capabilities: string[];
  allowedKinds: CapabilityKind[] | null;
  decisionId?: string | null;
}

/** 按模式确定能力白名单与类别上限;plan_only 下给了 execute 类能力直接报错 */
export function scopeFor(mode: "execute" | "plan_only", allowlist?: readonly string[]): RunScope {
  const all = Object.values(CAPABILITY_CATALOG);
  if (mode === "plan_only") {
    const bad = (allowlist ?? []).filter((id) => {
      const d = CAPABILITY_CATALOG[id as KnownCapabilityId];
      return !d || !PLAN_ONLY_KINDS.includes(d.kind);
    });
    if (bad.length) throw new Error(`plan_only 模式只允许 read/preview/propose 类能力,不允许:${bad.join(", ")}`);
    return {
      capabilities: allowlist?.length
        ? [...allowlist]
        : all.filter((d) => PLAN_ONLY_KINDS.includes(d.kind)).map((d) => d.id),
      allowedKinds: PLAN_ONLY_KINDS,
    };
  }
  for (const id of allowlist ?? []) KnownCapabilityId.parse(id);
  return { capabilities: allowlist?.length ? [...allowlist] : all.map((d) => d.id), allowedKinds: null };
}

/**
 * 会话目录优先级:模块装配参数(测试)> settings.agent.sessionDir > AGENT_SESSION_DIR(main 按 DATA_DIR 算好注入)
 * > ./data/agent-sessions
 */
export function sessionDirOf(env: AgentEnv, configured?: string | null): string {
  return resolve(
    env.opts.sessionDir ??
      configured ??
      Deno.env.get("AGENT_SESSION_DIR") ??
      join(Deno.cwd(), "data", "agent-sessions"),
  );
}

export interface DriveResult {
  status: "completed" | "failed" | "cancelled";
  reason: "budget" | "aborted" | "error" | null;
  budget: string | null;
  message: string | null;
  finalText: string;
}

/** 设置 → @umt/agent 的运行参数(apiKey 现场经 secrets 解密,不落日志) */
async function agentSettings(
  env: AgentEnv,
  s: GlobalSettings,
  budget: AgentBudget,
  fake: boolean,
): Promise<AgentSettings> {
  const pi = s.pi;
  let apiKey = "fake-key-not-used";
  if (!fake) {
    if (!pi.apiKeySecretRef) throw new Error("未配置 Pi 上游密钥(设置 → pi.apiKey)");
    apiKey = await env.ctx.secrets.resolve(pi.apiKeySecretRef);
    if (!pi.model) throw new Error("未配置模型(设置 → pi.model)");
  }
  const maxTokens = budget.maxTokens ?? s.agent.maxTokensPerRun ?? undefined;
  const maxCost = budget.maxCostUsd ?? s.agent.maxCostUsdPerRun ?? undefined;
  return {
    provider: pi.provider === "openai-compatible" ? "openai-compatible" : "anthropic",
    ...(pi.baseUrl && !fake ? { baseUrl: pi.baseUrl } : {}),
    apiKey,
    model: pi.model ?? "fake-model",
    thinkingLevel: pi.thinkingLevel,
    contextBudgetTokens: pi.contextBudget,
    compaction: {
      mode: pi.compaction.enabled ? "threshold" : "manual",
      threshold: Math.min(0.95, Math.max(0.5, pi.compaction.thresholdRatio)),
    },
    subagentsEnabled: pi.subagentsEnabled,
    maxToolCalls: budget.maxToolCalls ?? s.agent.maxToolCalls,
    maxDurationMs: (budget.maxDurationSec ?? s.agent.maxDurationSec) * 1000,
    ...(maxTokens !== undefined || maxCost !== undefined
      ? {
          costBudget: {
            ...(maxTokens !== undefined ? { maxTotalTokens: maxTokens } : {}),
            ...(maxCost !== undefined && maxCost > 0 ? { maxCostUsd: maxCost } : {}),
          },
        }
      : {}),
  };
}

/** 事件:先落库(run_id, seq 唯一)再推全局 SSE;未授权的工具失败在总线上标为 denied */
function eventSink(env: AgentEnv) {
  const { ctx } = env;
  return async (e: AgentRunEvent) => {
    const row = agentRunEventToRow(e);
    await ctx.db
      .insertInto("agent_run_event")
      .values({
        run_id: row.runId,
        seq: row.seq,
        type: row.type,
        at: new Date(row.at),
        payload: JSON.stringify(row.payload),
      })
      .onConflict((oc) => oc.columns(["run_id", "seq"]).doNothing())
      .execute();
    await ctx.db
      .updateTable("agent_run")
      .set({
        last_event_seq: sql`greatest(last_event_seq, ${e.seq})`,
        ...(e.type === "tool_call_start" && !e.recovery ? { tool_call_count: sql`tool_call_count + 1` } : {}),
        ...(e.type === "run_completed" || e.type === "run_failed" ? { usage: JSON.stringify(e.usage) } : {}),
      })
      .where("id", "=", e.runId)
      .execute();
    const bus = agentRunEventToBusEvent(e);
    if (!bus) return;
    // run_completed/run_failed 只代表模型本轮结束;agent_run 的最终状态由 finalize 按业务结果发布
    if (bus.type === "agent.run_status" && (e.type === "run_completed" || e.type === "run_failed")) return;
    if (
      bus.type === "agent.tool_call" &&
      e.type === "tool_call_end" &&
      e.isError &&
      /"code":"unauthorized"/.test(e.outputSummary)
    ) {
      bus.payload.phase = "denied";
    }
    void ctx.events.publish(bus.type, bus.payload);
  };
}

export async function publishStatus(ctx: ServerContext, runId: string, status: AgentRunStatus, summary: string | null) {
  void ctx.events.publish("agent.run_status", { runId, status, summary });
}

/** 取下一条待消费的追加消息(原子认领,两个驱动者不会拿到同一条) */
async function claimMessage(ctx: ServerContext, runId: string): Promise<string | null> {
  const r = await sql<{ text: string }>`
    update agent_run_message set status = 'consumed', consumed_at = now()
     where id = (select id from agent_run_message where run_id = ${runId} and status = 'pending'
                  order by created_at, id limit 1 for update skip locked)
    returning text`.execute(ctx.db);
  return r.rows[0]?.text ?? null;
}

/**
 * 崩溃恢复对账:本 run 在账本里停在 running 的有副作用调用,按同一幂等键重跑处理器
 * (处理器自身幂等:下载按计划状态 + qB hash,整理按记录状态),让账本落到真实结果。
 * 之后 resume 时模型拿到"结果未知",再次调用同参数能力会直接得到账本里的结果。
 */
export async function reconcilePending(
  env: AgentEnv,
  registry: CapabilityRegistry,
  run: AgentRunRow,
  signal: AbortSignal,
) {
  const rows = await env.ctx.db
    .selectFrom("capability_invocation")
    .selectAll()
    .where("agent_run_id", "=", run.id)
    .where("status", "=", "running")
    .orderBy("created_at")
    .execute();
  for (const r of rows) {
    try {
      await registry.invoke(r.capability_id, r.input, {
        actorId: r.actor_id,
        taskId: run.task_id ?? run.id,
        scopeIds: [],
        signal,
        source: "task",
        agentRunId: run.id,
        invocationId: `reconcile:${r.id}`,
      });
    } catch (err) {
      env.ctx.logger.warn(
        { runId: run.id, capability: r.capability_id, err: { message: (err as Error).message } },
        "agent 崩溃恢复对账失败",
      );
    }
  }
}

function systemPromptFor(run: AgentRunRow, scope: RunScope): string {
  const constraints: string[] = [
    `可用能力:${scope.capabilities.join(", ")}`,
    "有副作用的能力(提交下载、执行整理、发布策略、重试/取消任务)需要服务端授权;被拒绝时如实说明缺少哪项授权。",
  ];
  if (run.mode === "plan_only")
    constraints.push("本次为“先给方案”:只调查与规划,不提交下载、不做文件变更、不发布策略。");
  if (run.trigger === "event")
    constraints.push("本任务由后台事件触发,是有界的只读调查:收集证据、给出结论与建议,不要反复重试同一操作。");
  return buildSystemPrompt({
    role: run.trigger === "event" ? "媒体库异常调查员" : "媒体库自动化助手",
    goal: run.goal,
    constraints,
  });
}

/** 读 agent_run_event 里最后一轮的结束事件(崩溃发生在"模型已结束、业务状态未写"之间时复用) */
async function lastOutcomeEvent(env: AgentEnv, runId: string) {
  const r = await env.ctx.db
    .selectFrom("agent_run_event")
    .select(["type", "payload", "seq"])
    .where("run_id", "=", runId)
    .where("type", "in", ["run_started", "run_completed", "run_failed"])
    .orderBy("seq", "desc")
    .executeTakeFirst();
  return r ?? null;
}

async function dailyBudgetBlock(env: AgentEnv, s: GlobalSettings, run: AgentRunRow): Promise<string | null> {
  if (run.started_at) return null; // 续跑不重复计数
  const day = new Date();
  day.setHours(0, 0, 0, 0);
  const r = await sql<{ n: number; cost: number | null }>`
    select count(*)::int as n, sum(coalesce((usage->>'costUsd')::float8, 0))::float8 as cost
      from agent_run where started_at >= ${day} and id <> ${run.id}`.execute(env.ctx.db);
  const n = r.rows[0]?.n ?? 0;
  const cost = r.rows[0]?.cost ?? 0;
  if (n >= s.budget.agentMaxRunsPerDay) return `今日 Agent 运行次数已达上限 ${s.budget.agentMaxRunsPerDay}`;
  if (cost >= s.budget.agentMaxCostUsdPerDay) return `今日 Agent 费用已达上限 $${s.budget.agentMaxCostUsdPerDay}`;
  return null;
}

/**
 * 驱动一次模型会话直到本轮结束(含追加消息)。整个函数在一个持久步骤里执行:worker 崩溃后新 worker 从头重跑本函数,
 * 打开同一 JSONL 会话,pendingOperations>0 时对账 + resume,已完成的工具调用不会重放。
 */
export async function drive(
  env: AgentEnv,
  registry: CapabilityRegistry,
  runId: string,
  signal: AbortSignal,
): Promise<DriveResult> {
  const { ctx } = env;
  const run = await ctx.db.selectFrom("agent_run").selectAll().where("id", "=", runId).executeTakeFirstOrThrow();
  if (run.cancel_requested_at || signal.aborted)
    return { status: "cancelled", reason: "aborted", budget: null, message: "运行已被取消", finalText: "" };
  const s = await ctx.settings.all();
  const blocked = await dailyBudgetBlock(env, s, run);
  if (blocked) return { status: "failed", reason: "budget", budget: "daily", message: blocked, finalText: "" };

  const fake = s.pi.provider === "fake";
  const scope = (run.scope ?? {}) as unknown as RunScope;
  let settings: AgentSettings;
  try {
    settings = await agentSettings(env, s, (run.budget ?? {}) as AgentBudget, fake);
  } catch (err) {
    return { status: "failed", reason: "error", budget: null, message: (err as Error).message, finalText: "" };
  }
  await ctx.db
    .updateTable("agent_run")
    .set({
      status: "running",
      started_at: run.started_at ?? new Date(),
      provider: fake ? "fake" : s.pi.provider,
      model: fake ? "fake-model" : s.pi.model,
      thinking_level: s.pi.thinkingLevel,
    })
    .where("id", "=", runId)
    .execute();
  void publishStatus(ctx, runId, "running", null);

  const onEvent = eventSink(env);
  const onEventError = (err: unknown, e: AgentRunEvent) =>
    ctx.logger.warn({ runId, seq: e.seq, err: { message: (err as Error).message } }, "agent 事件落库失败");
  let scripts: FakeScript[] = DEMO_SCRIPTS;
  if (fake) {
    try {
      scripts = env.opts.fakeScripts ?? (s.agent.fakeScript ? parseFakeScripts(s.agent.fakeScript) : DEMO_SCRIPTS);
    } catch (err) {
      return {
        status: "failed",
        reason: "error",
        budget: null,
        message: `fakeScript 不合法:${(err as Error).message}`,
        finalText: "",
      };
    }
  }
  const runtime = fake
    ? createFakeRuntime({ registry, scripts, onEvent, onEventError }).runtime
    : new PiAgentRuntime({ registry, onEvent, onEventError });
  const maxSeq = await ctx.db
    .selectFrom("agent_run_event")
    .select(sql<number | null>`max(seq)`.as("m"))
    .where("run_id", "=", runId)
    .executeTakeFirst();
  const actor = {
    actorId: run.created_by ?? `agent:${runId}`,
    taskId: run.task_id ?? runId,
    scopeIds: [] as string[],
    ...(scope.allowedKinds ? { allowedKinds: scope.allowedKinds } : {}),
  };
  let handle: AgentRunHandle;
  try {
    handle = await runtime.createRun({
      runId,
      goal: run.goal,
      systemPrompt: systemPromptFor(run, scope),
      capabilities: scope.capabilities?.length ? scope.capabilities : Object.keys(CAPABILITY_CATALOG),
      settings,
      sessionDir: sessionDirOf(env, s.agent.sessionDir),
      actor,
      firstSeq: maxSeq?.m === null || maxSeq?.m === undefined ? 0 : Number(maxSeq.m) + 1,
    });
  } catch (err) {
    return { status: "failed", reason: "error", budget: null, message: (err as Error).message, finalText: "" };
  }

  // 取消传播:任务取消信号或 agent_run.cancel_requested_at → 中止会话
  let aborting = false;
  const abortOnce = () => {
    if (aborting) return;
    aborting = true;
    void handle.abort("用户取消");
  };
  signal.addEventListener("abort", abortOnce, { once: true });
  const watch = setInterval(async () => {
    const r = await ctx.db
      .selectFrom("agent_run")
      .select("cancel_requested_at")
      .where("id", "=", runId)
      .executeTakeFirst()
      .catch(() => null);
    if (r?.cancel_requested_at) abortOnce();
  }, 1000);

  try {
    let outcome: RunOutcome | null = null;
    if (handle.pendingOperations > 0) {
      await reconcilePending(env, registry, run, signal);
      outcome = await handle.resume();
    } else {
      const last = await lastOutcomeEvent(env, runId);
      if (!last) outcome = await handle.prompt();
      else if (last.type === "run_started") outcome = await handle.resume();
      else {
        const p = last.payload as { finalText?: string; reason?: string; message?: string; budget?: string };
        outcome =
          last.type === "run_completed"
            ? { status: "completed", finalText: p.finalText ?? "", usage: {} as never }
            : {
                status: "failed",
                reason: (p.reason ?? "error") as RunOutcome["reason"] & string,
                message: p.message ?? "",
                finalText: "",
                usage: {} as never,
                ...(p.budget ? { budget: p.budget as never } : {}),
              };
      }
    }
    while (outcome.status === "completed" && !aborting) {
      const text = await claimMessage(ctx, runId);
      if (text === null) break;
      outcome = await handle.prompt(text);
    }
    if (outcome.status === "completed")
      return { status: "completed", reason: null, budget: null, message: null, finalText: outcome.finalText };
    return {
      status: outcome.reason === "aborted" ? "cancelled" : "failed",
      reason: outcome.reason ?? "error",
      budget: outcome.budget ?? null,
      message: outcome.message ?? null,
      finalText: outcome.finalText,
    };
  } finally {
    clearInterval(watch);
    signal.removeEventListener("abort", abortOnce);
    await handle.close();
  }
}

/** 等待目标关联的下载/入库落定:返回报告(state≠waiting)或 null(继续等) */
export async function goalSettled(env: AgentEnv, runId: string, modelText: string): Promise<AgentGoalReport | null> {
  const r = await buildGoalReport(env.ctx.db, runId, modelText);
  if (r.state !== "waiting") return r;
  await env.ctx.db
    .updateTable("agent_run")
    .set({ status: "waiting", summary: r.summary })
    .where("id", "=", runId)
    .where("status", "in", ["running", "waiting"])
    .execute();
  return null;
}