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

import {
  type AssistantMessage,
  createModels,
  type FauxContentBlock,
  type FauxResponseFactory,
  fauxAssistantMessage,
  fauxProvider,
  fauxText,
  fauxThinking,
  fauxToolCall,
  type JsonObject,
  type Message,
  type TranscriptContext,
} from "@earendil-works/pi-ai";
import type { CapabilityRegistry } from "../capability/registry.ts";
import { capabilityIdToToolName, toolNameToCapabilityId } from "../capability/tool-name.ts";
import type { ResolvedPiModel } from "./model.ts";
import { PiAgentRuntime, type PiAgentRuntimeOptions } from "./runtime.ts";

/** 前面各步工具调用的结果(按发生顺序),供脚本根据结果决定下一步 */
export interface FakeToolResult {
  toolName: string;
  capabilityId?: string;
  isError: boolean;
  text: string;
  /** text 能解析成 JSON 时的值 */
  json?: unknown;
  /** 工具结果里的图片内容块数(视觉自动化:截图真的交给了模型) */
  images: number;
}

export interface FakeToolCall {
  /** 能力 ID(自动映射为工具名) */
  capability?: string;
  /** 原始工具名(用于模拟模型调用白名单外的工具,如 bash) */
  toolName?: string;
  input: JsonObject | ((results: FakeToolResult[]) => JsonObject);
}

export interface FakeStep {
  thinking?: string;
  text?: string;
  toolCalls?: FakeToolCall[];
}

export interface FakeScript {
  name?: string;
  /** 匹配用户消息(通常就是 goal) */
  match: string | RegExp | ((goal: string) => boolean);
  steps: FakeStep[];
  finalText: string | ((results: FakeToolResult[]) => string);
}

export interface FakeModelProviderOptions {
  contextWindow?: number;
  maxTokens?: number;
  /** 压缩请求时返回的摘要文本 */
  summaryText?: string;
  /** 模拟支持图片输入的模型(model.input 含 image);缺省 false */
  vision?: boolean;
}

const SUMMARY_MARKERS = ["context summarization assistant", "structured context checkpoint summary"];

/**
 * 脚本化假模型:无状态地从 transcript 推导当前步骤(最近一条用户消息之后已有几条 assistant 消息),
 * 因此进程崩溃后在新进程里 resume 也能从正确的步骤继续。不联网、不需要 API Key。
 */
export class FakeModelProvider {
  private readonly faux;
  private readonly scripts: FakeScript[];
  private readonly summaryText: string;
  /** 已收到的模型请求数(含压缩摘要请求) */
  calls = 0;
  /** 最近一次请求看到的系统提示与工具名(测试断言用) */
  lastRequest: { transcript: string } | undefined;
  /** 历次请求里收到的图片内容块(测试断言截图确实交给了模型) */
  images: { mimeType: string; bytes: number }[] = [];

  constructor(scripts: FakeScript[], options: FakeModelProviderOptions = {}) {
    this.scripts = scripts;
    this.summaryText = options.summaryText ?? "## 摘要\n目标:见首条消息。已完成动作:见工具结果。剩余事项:继续执行。";
    this.faux = fauxProvider({
      provider: "umt-fake",
      models: [
        {
          id: "fake-model",
          contextWindow: options.contextWindow ?? 64_000,
          maxTokens: options.maxTokens ?? 4096,
          reasoning: true,
          input: options.vision ? ["text", "image"] : ["text"],
        },
      ],
      tokensPerSecond: 0,
    });
    const factory: FauxResponseFactory = (context) => {
      this.faux.appendResponses([factory]);
      return this.respond(context);
    };
    this.faux.setResponses([factory]);
  }

  /** @internal 供 PiAgentRuntime 使用的模型装配 */
  resolve(): ResolvedPiModel {
    const models = createModels();
    models.setProvider(this.faux.provider);
    return { models, model: this.faux.getModel() };
  }

  private respond(context: TranscriptContext): AssistantMessage {
    this.calls += 1;
    const messages = context.messages as Message[];
    const transcript = messages.map(messageText).join("\n");
    this.lastRequest = { transcript };
    const last = messages[messages.length - 1];
    if (last?.role === "toolResult")
      for (const c of last.content as Array<{ type: string; data?: string; mimeType?: string }>)
        if (c.type === "image") this.images.push({ mimeType: c.mimeType ?? "", bytes: (c.data ?? "").length });
    const head = messages[0] ? messageText(messages[0]) : "";
    if (SUMMARY_MARKERS.some((m) => head.includes(m) || lastUserText(messages).includes(m))) {
      return fauxAssistantMessage(this.summaryText);
    }

    const lastUser = lastUserIndex(messages);
    const goal = lastUser >= 0 ? messageText(messages[lastUser] as Message) : "";
    const script = this.scripts.find((s) => matches(s.match, goal));
    if (!script) {
      return fauxAssistantMessage(`[fake] 没有脚本匹配目标:${goal.slice(0, 80)}`, {
        stopReason: "error",
        errorMessage: "no fake script matched",
      });
    }
    const after = messages.slice(lastUser + 1);
    const stepIndex = after.filter((m) => m.role === "assistant").length;
    const results = toolResults(after);
    const step = script.steps[stepIndex];
    if (!step) {
      const text = typeof script.finalText === "function" ? script.finalText(results) : script.finalText;
      return fauxAssistantMessage(text);
    }
    const blocks: FauxContentBlock[] = [];
    if (step.thinking) blocks.push(fauxThinking(step.thinking));
    if (step.text) blocks.push(fauxText(step.text));
    for (const call of step.toolCalls ?? []) {
      const name = call.toolName ?? capabilityIdToToolName(call.capability ?? "");
      blocks.push(fauxToolCall(name, typeof call.input === "function" ? call.input(results) : call.input));
    }
    const hasTools = (step.toolCalls?.length ?? 0) > 0;
    return fauxAssistantMessage(blocks, hasTools ? { stopReason: "toolUse" } : {});
  }
}

export interface CreateFakeRuntimeOptions extends Omit<PiAgentRuntimeOptions, "resolveModel"> {
  registry: CapabilityRegistry;
  scripts: FakeScript[];
  fake?: FakeModelProviderOptions;
}

/** e2e/测试便捷函数:脚本化假模型 + 真实 PiAgentRuntime(注册中心、白名单、预算、JSONL 全部走真实路径) */
export function createFakeRuntime(options: CreateFakeRuntimeOptions): {
  runtime: PiAgentRuntime;
  fake: FakeModelProvider;
} {
  const { scripts, fake: fakeOptions, ...rest } = options;
  const fake = new FakeModelProvider(scripts, fakeOptions);
  return { runtime: new PiAgentRuntime({ ...rest, resolveModel: () => fake.resolve() }), fake };
}

/** 供 settings 使用的假设置(provider 字段不会真正联网,模型由 FakeModelProvider 覆盖) */
export const FAKE_AGENT_SETTINGS = {
  provider: "anthropic",
  apiKey: "fake-key-not-used",
  model: "fake-model",
} as const;

function matches(match: FakeScript["match"], goal: string): boolean {
  if (typeof match === "string") return goal.includes(match);
  if (match instanceof RegExp) return match.test(goal);
  return match(goal);
}

function messageText(m: Message): string {
  if (typeof m.content === "string") return m.content;
  return (m.content as Array<{ type: string; text?: string }>)
    .map((c) => (c.type === "text" ? (c.text ?? "") : ""))
    .join("");
}

function lastUserIndex(messages: Message[]): number {
  for (let i = messages.length - 1; i >= 0; i--) if (messages[i]?.role === "user") return i;
  return -1;
}

function lastUserText(messages: Message[]): string {
  const i = lastUserIndex(messages);
  return i >= 0 ? messageText(messages[i] as Message) : "";
}

function toolResults(messages: Message[]): FakeToolResult[] {
  const out: FakeToolResult[] = [];
  for (const m of messages) {
    if (m.role !== "toolResult") continue;
    const text = messageText(m);
    let json: unknown;
    // 第一块文本是能力输出;运行时可能在后面追加【预算提醒】等附注块
    const first = (m.content as Array<{ type: string; text?: string }>).find((c) => c.type === "text")?.text ?? text;
    try {
      json = JSON.parse(first);
    } catch {
      json = undefined;
    }
    const capabilityId = toolNameToCapabilityId(m.toolName);
    out.push({
      toolName: m.toolName,
      ...(capabilityId ? { capabilityId } : {}),
      isError: m.isError,
      text,
      ...(json !== undefined ? { json } : {}),
      images: (m.content as Array<{ type: string }>).filter((c) => c.type === "image").length,
    });
  }
  return out;
}