InfoRSI/eval/lib/replay.ts

import { createServer, type ServerResponse } from "node:http";
import type { AddressInfo } from "node:net";
import { mkdirSync, writeFileSync } from "node:fs";
import { join } from "node:path";

// 回放用的假模型服务（任务书 12 第四部分第 3 条）：本机起一个兼容 OpenAI 的流式接口，
// 用 Pi 自己的 PI_CODING_AGENT_DIR 指一份 models.json 过来，正式代码不知道对面是假模型。
// 每个模型名一份按顺序交出的提交清单：每次被派任务（会话的第一次请求），交出下一份；之后的请求回一句话结束。
// 交完了就不再交，只回一句话：不为凑结果改内容。被退回时模型没有别的可交，会话就以"没交"结束。

/** 一份要交的提交：research_submit 的参数原样（payload，以及可能的 replaces），from 写它从哪来，只用于报告。 */
export interface ReplayItem {
  from: string;
  arguments: { payload: unknown; replaces?: string };
}

/** 回放脚本：模型名 → 按顺序交出的提交。没有要交的模型（实验员、核验员）写空数组。 */
export interface ReplayScript {
  source: string;
  models: Record<string, ReplayItem[]>;
}

export interface Replay {
  /** 起研究的子进程要带的环境变量。 */
  env: Record<string, string>;
  /** 每个模型名已经交出的提交（按交出先后）。 */
  delivered: Record<string, string[]>;
  /** 每个模型名还没交出的提交。 */
  remaining: () => Record<string, string[]>;
  close: () => Promise<void>;
}

export const REPLAY_PROVIDER = "replay";

/**
 * 起服务并写好 Pi 的 models.json、auth.json。delayMs 是每次回复前等多久（杀进程检验要一段"正在写"的窗口，0 就是不等）。
 * 一份提交只有在回复完整写出后才算交出：回复前对面断开（控制进程被杀），这一份放回清单最前面，下次被派任务时还交它。
 */
export async function startReplay(script: ReplayScript, agentDir: string, delayMs: number): Promise<Replay> {
  const queues = new Map(Object.entries(script.models).map(([model, items]) => [model, [...items]]));
  const delivered: Record<string, string[]> = Object.fromEntries([...queues.keys()].map((m) => [m, []]));

  const server = createServer((req, res) => {
    const chunks: Buffer[] = [];
    req.on("data", (c: Buffer) => chunks.push(c));
    req.on("end", () => void answer(JSON.parse(Buffer.concat(chunks).toString()), res));
  });

  async function answer(body: { model: string; messages: Array<{ role: string }> }, res: ServerResponse): Promise<void> {
    const queue = queues.get(body.model);
    if (queue === undefined) throw new Error(`回放脚本里没有模型 ${body.model}`);
    // 会话的第一次请求里还没有助手的回复：这是新派的任务。
    const item = body.messages.some((m) => m.role === "assistant") ? undefined : queue.shift();
    res.on("close", () => {
      if (!res.writableFinished && item !== undefined) queue.unshift(item);
    });
    await new Promise((ok) => setTimeout(ok, delayMs));
    if (res.destroyed) return;
    const send = (delta: unknown, finish: string | null) =>
      res.write(`data: ${JSON.stringify({ id: "r", object: "chat.completion.chunk", created: 1, model: body.model, choices: [{ index: 0, delta, finish_reason: finish }] })}\n\n`);
    res.writeHead(200, { "content-type": "text/event-stream" });
    if (item === undefined) {
      send({ role: "assistant", content: "好的，我这次没有别的要交了。" }, null);
      send({}, "stop");
    } else {
      const args = JSON.stringify(item.arguments);
      send({ role: "assistant", tool_calls: [{ index: 0, id: "call-1", type: "function", function: { name: "research_submit", arguments: args } }] }, null);
      send({}, "tool_calls");
    }
    res.write(`data: ${JSON.stringify({ id: "r", object: "chat.completion.chunk", created: 1, model: body.model, choices: [], usage: { prompt_tokens: 100, completion_tokens: 10 } })}\n\n`);
    res.end("data: [DONE]\n\n");
    if (item !== undefined) delivered[body.model]!.push(item.from);
  }

  await new Promise<void>((ok) => server.listen(0, "127.0.0.1", ok));
  const baseUrl = `http://127.0.0.1:${(server.address() as AddressInfo).port}/v1`;
  const models = [...queues.keys()].map((id) => ({
    id,
    reasoning: false,
    input: ["text"],
    contextWindow: 1_000_000,
    maxTokens: 32_000,
    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
  }));
  mkdirSync(agentDir, { recursive: true });
  writeFileSync(join(agentDir, "models.json"), JSON.stringify({ providers: { [REPLAY_PROVIDER]: { baseUrl, api: "openai-completions", apiKey: "k", models } } }, null, 2));
  writeFileSync(join(agentDir, "auth.json"), "{}");
  return {
    env: { PI_CODING_AGENT_DIR: agentDir },
    delivered,
    remaining: () => Object.fromEntries([...queues].map(([m, q]) => [m, q.map((i) => i.from)])),
    close: () => new Promise<void>((ok) => server.close(() => ok())),
  };
}

/** 回放用的各角色模型：研究者按脚本里的模型名，实验员、核验员各一个空清单的名字。 */
export function replayModels(researchers: string[], jev: string): unknown {
  const m = (model: string) => ({ provider: REPLAY_PROVIDER, model, thinking: "off" });
  return { researchers: researchers.map(m), experimenter: m("experimenter"), verifier: m("verifier"), jev };
}