UHD_Media_Tool/packages/agent/src/pi/runtime.ts

import {
  AgentHarness,
  type AgentHarness as AgentHarnessInstance,
  type AgentHarnessTool,
  type AgentLane,
  type AgentMessage,
  BACKGROUND_CONTEXT,
  type Context,
  type JsonlSessionRepo,
  type RunResult,
} from "@earendil-works/pi-agent-core";
import type { Usage } from "@earendil-works/pi-ai";
import type { AgentBudgetKind, AgentRunEvent, AgentRunEventBody, AgentUsageTotals } from "@umt/contracts";
import type { TSchema } from "typebox";
import { isCapabilityError } from "../capability/errors.ts";
import type { CapabilityRegistry, RegisteredCapability } from "../capability/registry.ts";
import { toolNameToCapabilityId } from "../capability/tool-name.ts";
import { redactForModel, redactString, summarizeForLog } from "../redact.ts";
import { EventChannel } from "../runtime/event-channel.ts";
import {
  type AgentRunHandle,
  type AgentRuntimePort,
  type CreateRunSpec,
  emptyUsage,
  type RunOutcome,
} from "../runtime/port.ts";
import { type ResolvedAgentSettings, resolveAgentSettings } from "../settings.ts";
import { compactionFromSettings, type PiModelResolver, resolveModelFromSettings } from "./model.ts";
import { openRunSession } from "./session.ts";

export interface PiAgentRuntimeOptions {
  registry: CapabilityRegistry;
  /** 默认按设置构造真实上游;测试用 FakeModelProvider 覆盖 */
  resolveModel?: PiModelResolver;
  /**
   * 每个事件产生时回调(server 用它写 agent_run_event 再推 SSE)。
   * 同一 run 内按 seq 顺序串行调用(前一个 Promise 结束才调下一个);回调抛错只记入 onEventError,不影响运行。
   * close() 会等待已排队的回调全部完成。
   */
  onEvent?: (event: AgentRunEvent) => void | Promise<void>;
  onEventError?: (error: unknown, event: AgentRunEvent) => void;
  now?: () => Date;
}

export const DEFAULT_COMPACTION_INSTRUCTIONS =
  "保留:本次目标、约束与授权范围、证据引用(对象 ID/记录)、已完成的动作及其结果(已提交≠已入库)、剩余事项。不要保留任何凭证。";

const LANE = "main";
const SUMMARY_LIMIT = 600;

/** Pi(pi-agent-core AgentHarness)实现的 AgentRuntimePort */
export class PiAgentRuntime implements AgentRuntimePort {
  constructor(private readonly options: PiAgentRuntimeOptions) {}

  async createRun(spec: CreateRunSpec): Promise<AgentRunHandle> {
    const ctx = BACKGROUND_CONTEXT;
    const settings = resolveAgentSettings(spec.settings);
    const capabilities = this.options.registry.resolve(spec.capabilities);
    const { models, model } = (this.options.resolveModel ?? resolveModelFromSettings)(settings);
    const opened = await openRunSession(spec.sessionDir, spec.runId, ctx);

    const run = new PiAgentRun(spec, settings, capabilities, this.options, opened.path, opened.repo);
    try {
      const { harness, open } = await AgentHarness.create(
        {
          session: opened.session,
          models,
          model,
          systemPrompt: spec.systemPrompt,
          tools: capabilities.map((c) => run.toTool(c)),
          toolContext: {},
          thinkingLevel: settings.thinkingLevel,
          compaction: compactionFromSettings(settings, model.contextWindow),
          retry: { enabled: settings.maxRetries > 0, maxRetries: settings.maxRetries, baseDelayMs: 1000 },
          streamOptions: { timeoutMs: settings.requestTimeoutMs, maxRetries: settings.maxRetries },
          toolExecution: "sequential",
        },
        ctx,
      );
      const lane = await harness.lane(LANE, ctx);
      run.attach(harness, lane, open.filter((o) => o.lane === LANE).length);
    } catch (error) {
      await opened.session.close(ctx).catch(() => undefined);
      await opened.repo.close(ctx).catch(() => undefined);
      throw error;
    }
    return run;
  }
}

type BudgetState = { kind: AgentBudgetKind; message: string };

class PiAgentRun implements AgentRunHandle {
  readonly runId: string;
  readonly sessionPath: string;
  pendingOperations = 0;

  private readonly channel = new EventChannel<AgentRunEvent>();
  private readonly allowed = new Map<string, string>();
  private readonly rejected = new Set<string>();
  private readonly abortController = new AbortController();
  private harness!: AgentHarnessInstance<object>;
  private lane!: AgentLane;
  private readonly firstSeq: number;
  private seq: number;
  private delivery: Promise<void> = Promise.resolve();
  private toolCalls = 0;
  private usage: AgentUsageTotals = emptyUsage();
  private finalText = "";
  private budget: BudgetState | undefined;
  private userAbort: string | undefined;
  private startedAt: number | undefined;
  private deadlineTimer: NodeJS.Timeout | undefined;
  private closed = false;

  constructor(
    private readonly spec: CreateRunSpec,
    private readonly settings: ResolvedAgentSettings,
    capabilities: RegisteredCapability[],
    private readonly options: PiAgentRuntimeOptions,
    sessionPath: string,
    private readonly repo: JsonlSessionRepo,
  ) {
    this.runId = spec.runId;
    this.sessionPath = sessionPath;
    this.firstSeq = spec.firstSeq ?? 0;
    this.seq = this.firstSeq;
    for (const c of capabilities) this.allowed.set(c.toolName, c.descriptor.id);
  }

  /** 能力 → Pi 工具。执行时统一走注册中心 invoke(schema/作用域/授权/幂等/审计) */
  toTool(entry: RegisteredCapability): AgentHarnessTool<object> {
    const { descriptor } = entry;
    const registry = this.options.registry;
    const actor = this.spec.actor;
    return {
      name: entry.toolName,
      label: descriptor.title,
      description: `${descriptor.description}\n[能力 ${descriptor.id} v${descriptor.version};类别 ${descriptor.kind};风险 ${descriptor.riskLevel}]`,
      parameters: registry.inputJsonSchema(descriptor.id) as unknown as TSchema,
      replay: descriptor.replay,
      ...(descriptor.kind === "execute" ? { executionMode: "sequential" as const } : {}),
      execute: async (toolCallId, params, _onUpdate, _toolCtx, invocation, context) => {
        const signal = context.abortSignal ?? this.abortController.signal;
        try {
          const output = await registry.invoke(descriptor.id, params, {
            actorId: actor.actorId,
            taskId: actor.taskId,
            scopeIds: actor.scopeIds,
            ...(actor.allowedKinds ? { allowedKinds: actor.allowedKinds } : {}),
            signal,
            source: "agent",
            invocationId: invocation.invocationId ?? toolCallId,
            agentRunId: this.runId,
          });
          return {
            content: [{ type: "text", text: JSON.stringify(redactForModel(output) ?? null) }],
            details: undefined,
          };
        } catch (error) {
          if (isCapabilityError(error)) {
            throw new Error(JSON.stringify({ error: { code: error.code, message: redactString(error.message) } }));
          }
          throw new Error(redactString(error instanceof Error ? error.message : String(error)));
        }
      },
    };
  }

  attach(harness: AgentHarnessInstance<object>, lane: AgentLane, pendingOperations: number): void {
    this.harness = harness;
    this.lane = lane;
    this.pendingOperations = pendingOperations;
    const { events, hooks } = harness;

    // 白名单二次校验 + 工具调用次数预算
    hooks.on("before_tool", (event) => {
      const decision = checkToolCall(event.toolName, this.allowed);
      if (!decision.allowed) {
        this.reject(event.toolCallId, event.toolName, decision.reason);
        return { block: { reason: decision.reason } };
      }
      this.toolCalls += 1;
      if (this.toolCalls > this.settings.maxToolCalls) {
        const reason = `工具调用次数超出预算(${this.settings.maxToolCalls})`;
        this.exceed("tool_calls", reason);
        this.rejected.add(event.toolCallId);
        this.emit({
          type: "tool_call_rejected",
          toolCallId: event.toolCallId,
          toolName: event.toolName,
          capabilityId: decision.capabilityId,
          reason,
        });
        return { block: { reason, terminate: true } };
      }
      return undefined;
    });
    hooks.on("before_request", () => {
      this.checkDeadline();
      if (this.budget) this.scheduleAbort();
      return undefined;
    });

    events.on("message_update", (e) => {
      if (e.event.type === "text_delta") this.emit({ type: "text_delta", delta: e.event.delta });
      else if (e.event.type === "thinking_delta") this.emit({ type: "thinking_delta", delta: e.event.delta });
    });
    events.on("message_end", (e) => {
      const text = assistantText(e.message);
      if (text?.trim()) this.finalText = text;
    });
    events.on("tool_start", (e) => {
      if (this.rejected.has(e.toolCallId)) return;
      const capabilityId = this.allowed.get(e.toolName);
      if (!capabilityId) {
        // 模型调用了未注册的工具(例如 bash):harness 已合成 unavailable 错误,这里只产生拒绝事件
        this.reject(e.toolCallId, e.toolName, `工具 ${e.toolName} 不在本次白名单内`);
        return;
      }
      this.emit({
        type: "tool_call_start",
        toolCallId: e.toolCallId,
        capabilityId,
        inputSummary: summarizeForLog(e.args, SUMMARY_LIMIT),
        recovery: e.recovery === true,
      });
    });
    events.on("tool_end", (e) => {
      if (this.rejected.has(e.toolCallId)) return;
      const capabilityId = this.allowed.get(e.toolName) ?? e.toolName;
      const text = e.result.content
        .map((c) => (c.type === "text" ? c.text : `[${c.type}]`))
        .join("")
        .trim();
      this.emit({
        type: "tool_call_end",
        toolCallId: e.toolCallId,
        capabilityId,
        isError: e.isError,
        outputSummary: summarizeForLog(text, SUMMARY_LIMIT),
        recovery: e.recovery === true,
      });
    });
    events.on("usage", (e) => {
      this.usage = toAgentUsage(e.totals);
      const cost = this.settings.costBudget;
      if (cost?.maxTotalTokens !== undefined && this.usage.totalTokens > cost.maxTotalTokens) {
        this.exceed("tokens", `token 用量 ${this.usage.totalTokens} 超出预算 ${cost.maxTotalTokens}`);
      } else if (cost?.maxCostUsd !== undefined && this.usage.costUsd > cost.maxCostUsd) {
        this.exceed("cost", `费用 $${this.usage.costUsd.toFixed(4)} 超出预算 $${cost.maxCostUsd}`);
      }
    });
    events.on("compaction_start", (e) => this.emit({ type: "compaction", phase: "start", reason: e.reason }));
    events.on("compaction_end", (e) =>
      this.emit({ type: "compaction", phase: "end", reason: e.reason, status: e.status }),
    );
  }

  events(fromSeq = this.firstSeq): AsyncIterable<AgentRunEvent> {
    return this.channel.iterate(fromSeq - this.firstSeq);
  }

  history(): AgentRunEvent[] {
    return this.channel.snapshot();
  }

  async prompt(text?: string): Promise<RunOutcome> {
    const message = text ?? this.spec.goal;
    return this.drive(false, () => this.lane.prompt(message, undefined, BACKGROUND_CONTEXT));
  }

  async resume(): Promise<RunOutcome> {
    return this.drive(true, async () => {
      const r = await this.lane.resume(BACKGROUND_CONTEXT);
      if (!r.ok && r.error._tag === "NothingToResume") return undefined;
      return r as RunResult;
    });
  }

  async abort(reason = "用户取消"): Promise<void> {
    this.userAbort = reason;
    this.abortController.abort(reason);
    await this.lane.abort(BACKGROUND_CONTEXT).catch(() => undefined);
  }

  async compact(customInstructions = DEFAULT_COMPACTION_INSTRUCTIONS): Promise<void> {
    const r = await this.lane.compact({ customInstructions }, BACKGROUND_CONTEXT);
    if (!r.ok && r.error._tag !== "NothingToCompact") throw new Error(`压缩失败: ${r.error._tag}`);
  }

  async close(): Promise<void> {
    if (this.closed) return;
    this.closed = true;
    this.clearDeadline();
    const ctx: Context = BACKGROUND_CONTEXT;
    await this.harness?.close(ctx).catch(() => undefined);
    await this.repo.close(ctx).catch(() => undefined);
    this.channel.close();
    await this.delivery;
  }

  // ---- internals ----

  private async drive(resumed: boolean, start: () => Promise<RunResult | undefined>): Promise<RunOutcome> {
    if (this.closed) throw new Error(`run ${this.runId} 已关闭`);
    this.startedAt ??= Date.now();
    this.emit({ type: "run_started", goal: this.spec.goal, resumed });
    this.checkDeadline();
    if (this.budget) return this.finish(undefined);
    this.armDeadline();
    let result: RunResult | undefined;
    try {
      result = await start();
    } finally {
      this.clearDeadline();
    }
    if (result && !result.ok) throw new Error(`Pi 拒绝了本次操作: ${result.error._tag} ${result.error.message}`);
    this.pendingOperations = 0;
    return this.finish(result);
  }

  private finish(result: RunResult | undefined): RunOutcome {
    const record = result?.ok ? result.value : undefined;
    const status = record && "status" in record ? record.status : undefined;
    let outcome: RunOutcome;
    if (this.budget) {
      outcome = this.failed("budget", this.budget.message, this.budget.kind);
    } else if (status === "aborted" || this.userAbort) {
      outcome = this.failed("aborted", this.userAbort ?? "已中止");
    } else if (status === "failed") {
      const err = record && "error" in record ? record.error : undefined;
      outcome = this.failed("error", redactString(err ? `${err.code}: ${err.message}` : "运行失败"));
    } else if (status === "suspended") {
      outcome = this.failed("error", "上游进入延迟模式(deferred),v0.1 不支持");
    } else {
      outcome = { status: "completed", finalText: this.finalText, usage: this.usage };
      this.emit({ type: "run_completed", finalText: outcome.finalText, usage: this.usage });
      return outcome;
    }
    this.emit({
      type: "run_failed",
      reason: outcome.reason ?? "error",
      message: outcome.message ?? "",
      usage: this.usage,
      ...(outcome.budget ? { budget: outcome.budget } : {}),
    });
    return outcome;
  }

  private failed(reason: RunOutcome["reason"] & string, message: string, budget?: AgentBudgetKind): RunOutcome {
    return {
      status: "failed",
      reason,
      message,
      finalText: this.finalText,
      usage: this.usage,
      ...(budget ? { budget } : {}),
    };
  }

  private reject(toolCallId: string, toolName: string, reason: string): void {
    if (this.rejected.has(toolCallId)) return;
    this.rejected.add(toolCallId);
    const capabilityId = toolNameToCapabilityId(toolName);
    this.emit({ type: "tool_call_rejected", toolCallId, toolName, ...(capabilityId ? { capabilityId } : {}), reason });
  }

  private exceed(kind: AgentBudgetKind, message: string): void {
    if (this.budget) return;
    this.budget = { kind, message };
    this.scheduleAbort();
  }

  /** 不能在事件/钩子回调里同步等待 abort(会与事件投递互相等待),放到下一个宏任务 */
  private scheduleAbort(): void {
    setTimeout(() => {
      this.abortController.abort("budget");
      void this.lane.abort(BACKGROUND_CONTEXT).catch(() => undefined);
    }, 0);
  }

  private checkDeadline(): void {
    if (this.startedAt !== undefined && Date.now() - this.startedAt > this.settings.maxDurationMs) {
      this.exceed("duration", `运行时长超出预算(${this.settings.maxDurationMs}ms)`);
    }
  }

  private armDeadline(): void {
    const remaining = (this.startedAt ?? Date.now()) + this.settings.maxDurationMs - Date.now();
    this.deadlineTimer = setTimeout(
      () => this.exceed("duration", `运行时长超出预算(${this.settings.maxDurationMs}ms)`),
      Math.max(0, remaining),
    );
  }

  private clearDeadline(): void {
    if (this.deadlineTimer) clearTimeout(this.deadlineTimer);
    this.deadlineTimer = undefined;
  }

  private emit(body: AgentRunEventBody): void {
    const event = { ...body, runId: this.runId, seq: this.seq++, at: new Date().toISOString() } as AgentRunEvent;
    this.channel.push(event);
    const onEvent = this.options.onEvent;
    if (onEvent) {
      // 串行投递:落库顺序与 seq 一致
      this.delivery = this.delivery
        .then(() => onEvent(event))
        .catch((error) => this.options.onEventError?.(error, event));
    }
  }
}

/** before_tool 白名单判定(纯函数,便于单测) */
export function checkToolCall(
  toolName: string,
  allowed: ReadonlyMap<string, string>,
): { allowed: true; capabilityId: string } | { allowed: false; reason: string; capabilityId?: string } {
  const capabilityId = allowed.get(toolName);
  if (capabilityId) return { allowed: true, capabilityId };
  const mapped = toolNameToCapabilityId(toolName);
  return {
    allowed: false,
    reason: `工具 ${toolName} 不在本次白名单内`,
    ...(mapped ? { capabilityId: mapped } : {}),
  };
}

function assistantText(message: AgentMessage): string | undefined {
  const m = message as { role?: string; content?: unknown };
  if (m.role !== "assistant" || !Array.isArray(m.content)) return undefined;
  return (m.content as Array<{ type: string; text?: string }>)
    .filter((c) => c.type === "text")
    .map((c) => c.text ?? "")
    .join("");
}

function toAgentUsage(u: Usage): AgentUsageTotals {
  return {
    inputTokens: u.input,
    outputTokens: u.output,
    cacheReadTokens: u.cacheRead,
    cacheWriteTokens: u.cacheWrite,
    totalTokens: u.totalTokens,
    costUsd: u.cost.total,
  };
}