InfoRSI/tests/agents.test.ts

import assert from "node:assert/strict";
import { mkdirSync, mkdtempSync, realpathSync, rmSync, symlinkSync, writeFileSync } from "node:fs";
import { createServer } from "node:http";
import type { AddressInfo } from "node:net";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { test, type TestContext } from "node:test";
import { parseStreamingJson, validateToolArguments, type AssistantMessageEvent } from "@earendil-works/pi-ai";
import type { ToolDefinition } from "@earendil-works/pi-coding-agent";
import { buildTools, newSubmitState, readSource, toolNames, watchSubmitArguments, type Host, type SubmitState } from "../src/agents/tools.ts";
import { runSession, type KindOf, type SessionTask, type SubmitAttempt } from "../src/agents/index.ts";
import type { ModelUsage } from "../src/vocab/index.ts";

// agents 不花钱的测试（任务书 07 第六部分）：提交参数的两处特殊处理、读代码的目录限制、
// 实验员 shell 的环境变量、三种角色的工具表。真实模型链路由 npm run test:pi 证明。
// 这里直接用 tools.ts 搭工具，回调是测试自己造的任务对象，不经 Pi 会话。

function tempDir(t: TestContext, name: string): string {
  const dir = mkdtempSync(join(tmpdir(), `inforsi-agents-${name}-`));
  t.after(() => rmSync(dir, { recursive: true, force: true }));
  return dir;
}

/** 回调出错直接抛，测试里不应发生。 */
const HOST: Host = {
  guard: (run) => run(),
  fail: (error) => {
    throw error;
  },
};

const signal = new AbortController().signal;
const common = {
  model: { provider: "p", model: "m", thinking: "high" },
  systemPrompt: "角色提示词",
  firstMessage: "第一条消息",
  sessionDir: "/nonexistent-sessions",
  query: () => null,
  usage: () => {},
  stop: signal,
} as const;

/** 造一个任务，记下每次提交回调收到了什么；回调一律收下，编号 s1、s2…… */
function setup(task: SessionTask): { tools: ToolDefinition[]; state: SubmitState; attempts: Array<SubmitAttempt<never>> } {
  const attempts: Array<SubmitAttempt<never>> = [];
  const recorded = {
    ...task,
    submit: async (attempt: SubmitAttempt<never>) => {
      attempts.push(attempt);
      return { accepted: true as const, id: `s${attempts.length}` };
    },
  } as SessionTask;
  const state = newSubmitState();
  return { tools: buildTools(recorded, state, HOST), state, attempts };
}

function researcher(kind: "wrapUp" | "firstPrediction", workDir: string, readableDirs: string[]): SessionTask {
  return { ...common, role: "researcher", kind, workDir, readableDirs, writeNote: () => {}, submit: async () => ({ accepted: true, id: "x" }) };
}

function tool(tools: ToolDefinition[], name: string): ToolDefinition {
  const found = tools.find((t) => t.name === name);
  assert.ok(found, `没有工具 ${name}`);
  return found;
}

/** 按 Pi 流式输出的样子把一段参数原文喂给监听：原文 40 字一片，结尾的参数对象用 Pi 自己的解析函数得到。 */
function stream(state: SubmitState, raw: string): Record<string, unknown> {
  const args = parseStreamingJson(raw) as Record<string, unknown>;
  const events: unknown[] = [{ type: "toolcall_start", contentIndex: 0 }];
  for (let i = 0; i < raw.length; i += 40) events.push({ type: "toolcall_delta", contentIndex: 0, delta: raw.slice(i, i + 40) });
  events.push({ type: "toolcall_end", contentIndex: 0, toolCall: { type: "toolCall", id: "c1", name: "research_submit", arguments: args } });
  for (const event of events) watchSubmitArguments(state, event as AssistantMessageEvent);
  return args;
}

// ---------------------------------------------------------------------------
// 提交参数原文不是合法 JSON（任务书第四部分第 1 条，参考项目 15-real-1、提交 bcc422e）
// ---------------------------------------------------------------------------

test("提交参数原文里半角双引号没转义：退回并说明原因，不收前半截；replaces 在前、在后都一样", async (t) => {
  const head = "## 什么被支持了\n因子 p1 的效果被确认。";
  const rest = '下一步想试 "更紧的排列" 这个说法，然后继续。';
  // 模型正确转义的部分按 JSON 规则写；后半段的半角双引号没转义。
  const esc = (s: string) => JSON.stringify(s).slice(1, -1);
  const raws = {
    replacesFirst: `{"replaces": "s7", "payload": {"body": "${esc(head)}${rest}"}}`,
    payloadFirst: `{"payload": {"body": "${esc(head)}${rest}"}, "replaces": "s7"}`,
  };
  for (const [name, raw] of Object.entries(raws)) {
    assert.throws(() => JSON.parse(raw), SyntaxError, `${name}：原文本身就不是合法 JSON`);
    const { tools, state, attempts } = setup(researcher("wrapUp", tempDir(t, "work"), []));
    const submit = tool(tools, "research_submit");
    const args = stream(state, raw);
    if (name === "replacesFirst") {
      // 不拦的话：Pi 解析出的前半截形状照样合格，正文停在那个引号前。
      const passed = validateToolArguments(
        { name: "research_submit", description: "", parameters: submit.parameters },
        { type: "toolCall", id: "", name: "research_submit", arguments: args },
      ) as { payload: { body: string } };
      assert.equal(passed.payload.body, `${head}下一步想试 `);
    }
    assert.throws(() => submit.prepareArguments!(args), /不是合法的 JSON[\s\S]*半角双引号/, name);
    await Promise.all(state.pending);
    assert.equal(attempts.length, 1, `${name}：交给提交回调记一条退回`);
    const attempt = attempts[0]!;
    assert.equal(attempt.shape, "invalid");
    assert.ok(attempt.shape === "invalid" && /不是合法的 JSON/.test(attempt.reason));
  }
});

/**
 * 本机假模型服务：OpenAI 兼容的流式接口，第一次回复一个提交调用（参数原文按 7 字一片发出），之后回一句话。
 * Pi 经测试自己的 PI_CODING_AGENT_DIR（里面的 models.json 指向这个服务）连过来，不碰本机的 ~/.pi/agent。
 */
async function fakeModel(t: TestContext, rawArguments: string): Promise<void> {
  let replies = 0;
  const server = createServer((req, res) => {
    req.resume();
    req.on("end", () => {
      const send = (delta: unknown, finish: string | null) =>
        res.write(`data: ${JSON.stringify({ id: "x", object: "chat.completion.chunk", created: 1, model: "m", choices: [{ index: 0, delta, finish_reason: finish }] })}\n\n`);
      res.writeHead(200, { "content-type": "text/event-stream" });
      if (replies++ === 0) {
        send({ role: "assistant", tool_calls: [{ index: 0, id: "call-1", type: "function", function: { name: "research_submit", arguments: "" } }] }, null);
        for (let i = 0; i < rawArguments.length; i += 7) send({ tool_calls: [{ index: 0, function: { arguments: rawArguments.slice(i, i + 7) } }] }, null);
        send({}, "tool_calls");
      } else {
        send({ role: "assistant", content: "知道了。" }, null);
        send({}, "stop");
      }
      res.write(`data: ${JSON.stringify({ id: "x", object: "chat.completion.chunk", created: 1, model: "m", choices: [], usage: { prompt_tokens: 100, completion_tokens: 5 } })}\n\n`);
      res.end("data: [DONE]\n\n");
    });
  });
  await new Promise<void>((done) => server.listen(0, "127.0.0.1", done));
  t.after(() => server.close());
  const agentDir = tempDir(t, "pi-agent");
  const model = { id: "m", reasoning: true, input: ["text"], contextWindow: 100000, maxTokens: 1000, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 } };
  const baseUrl = `http://127.0.0.1:${(server.address() as AddressInfo).port}/v1`;
  writeFileSync(join(agentDir, "models.json"), JSON.stringify({ providers: { fake: { baseUrl, api: "openai-completions", apiKey: "k", models: [model] } } }));
  writeFileSync(join(agentDir, "auth.json"), "{}");
  process.env.PI_CODING_AGENT_DIR = agentDir;
  t.after(() => delete process.env.PI_CODING_AGENT_DIR);
}

test("经过真实的 Pi 会话：参数原文不是合法 JSON 的提交交给提交回调记退回，原因回给模型", async (t) => {
  await fakeModel(t, '{"payload": {"body": "前半段，然后"引号"后半段"}}');
  const work = tempDir(t, "session");
  mkdirSync(join(work, "sessions"));
  const attempts: Array<SubmitAttempt<KindOf<"researcher">>> = [];
  const usages: ModelUsage[] = [];
  const outcome = await runSession({
    ...common,
    model: { provider: "fake", model: "m", thinking: "medium" },
    role: "researcher",
    kind: "wrapUp",
    workDir: work,
    sessionDir: join(work, "sessions"),
    readableDirs: [],
    writeNote: () => {},
    usage: (u) => usages.push(u),
    submit: async (attempt) => {
      attempts.push(attempt);
      return { accepted: true, id: "s1" };
    },
  });
  assert.equal(outcome.status, "notSubmitted", JSON.stringify(outcome));
  assert.equal(attempts.length, 1);
  assert.ok(attempts[0]!.shape === "invalid" && /不是合法的 JSON/.test(attempts[0]!.reason), JSON.stringify(attempts));
  assert.deepEqual(attempts[0]!.shape === "invalid" && attempts[0]!.raw, { payload: { body: "前半段，然后" } });
  assert.equal(usages.length, 2, "两次回复各一份用量");
  assert.ok(usages.every((u) => u.inputTokens.status === "measured" && u.inputTokens.value === 100));
});

test("参数原文是合法 JSON 时照常往下走，不登记", (t) => {
  const { tools, state, attempts } = setup(researcher("wrapUp", tempDir(t, "work"), []));
  const args = stream(state, JSON.stringify({ payload: { body: '正文里有 "引号"，正确转义了' } }));
  assert.equal(tool(tools, "research_submit").prepareArguments!(args), args);
  assert.equal(attempts.length, 0);
});

// ---------------------------------------------------------------------------
// 必填字段交 null（任务书第四部分第 2 条，参考项目 14-real-1、提交 fd73b78）
// ---------------------------------------------------------------------------

test("必填的数字、文字交 null：退回并报出字段位置，不让 Pi 转成 0 或 \"null\"", async (t) => {
  const { tools, state, attempts } = setup(researcher("firstPrediction", tempDir(t, "work"), []));
  const submit = tool(tools, "research_submit");
  const args = { payload: { range: { q1: 0.1, q5: null, q9: 0.3 }, body: null } };
  assert.throws(() => submit.prepareArguments!(args), /交的是 null：payload\.range\.q5、payload\.body。/);
  await Promise.all(state.pending);
  assert.equal(attempts.length, 1);
  assert.ok(attempts[0]!.shape === "invalid" && attempts[0]!.reason.includes("payload.range.q5、payload.body"));
});

test("二选一里本来收 null 的字段交 null 照常往下走；Pi 的形状校验不通过也交给提交回调", async (t) => {
  const work = tempDir(t, "work");
  const task: SessionTask = { ...common, role: "experimenter", kind: "ablationImplementation", workDir: work, lightCheck: async () => "", submit: async () => ({ accepted: true, id: "x" }) };
  const { tools, state, attempts } = setup(task);
  const submit = tool(tools, "research_submit");
  const ok = { payload: { alsoOff: [], cannotRun: null, concerns: "无" } };
  assert.equal(submit.prepareArguments!(ok), ok);
  assert.throws(() => submit.prepareArguments!({ payload: { alsoOff: ["x1"], cannotRun: null, concerns: "无" } }), /alsoOff/);
  // 实验员的提交不能替换：参数里根本没有 replaces。
  assert.throws(() => submit.prepareArguments!({ ...ok, replaces: "s1" }), /replaces|additional/i);
  await Promise.all(state.pending);
  assert.deepEqual(attempts.map((a) => a.shape), ["invalid", "invalid"]);
});

// ---------------------------------------------------------------------------
// 读代码：只能读给定目录里的东西
// ---------------------------------------------------------------------------

test("读代码：目录里的文件和目录读得到；目录外的读不到，用 .. 和符号链接绕出去的也读不到", (t) => {
  const base = realpathSync(tempDir(t, "read"));
  const root = join(base, "task");
  const secret = join(base, "other-researcher");
  mkdirSync(join(root, "src"), { recursive: true });
  mkdirSync(secret);
  writeFileSync(join(root, "src", "a.py"), "line1\nline2\n");
  writeFileSync(join(secret, "proposal.md"), "还没公开的提案");
  symlinkSync(secret, join(root, "link-dir"));
  symlinkSync(join(secret, "proposal.md"), join(root, "link-file"));

  assert.equal(readSource([root], join(root, "src", "a.py")), "line1\nline2\n");
  assert.equal(readSource([root], root), "link-dir\nlink-file\nsrc/");
  for (const path of [base, secret, join(secret, "proposal.md"), join(root, "..", "other-researcher", "proposal.md"), join(root, "link-dir"), join(root, "link-file"), join(root, "link-dir", "proposal.md")]) {
    assert.throws(() => readSource([root], path), /不在可读的目录里/, path);
  }
  assert.throws(() => readSource([root], "src/a.py"), /只接受绝对路径/);
  // 前缀相同但不是子目录的也不算在里面。
  mkdirSync(`${root}-copy`);
  assert.throws(() => readSource([root], `${root}-copy`), /不在可读的目录里/);
});

test("读代码：超过单次上限时只显示前 2000 行，并写明截断", (t) => {
  const root = realpathSync(tempDir(t, "page"));
  const lines = Array.from({ length: 2500 }, (_, i) => `第 ${i + 1} 行`);
  writeFileSync(join(root, "big.txt"), lines.join("\n"));
  const shown = readSource([root], join(root, "big.txt"));
  assert.ok(shown.startsWith("第 1 行\n"));
  assert.ok(shown.includes("第 2000 行\n"));
  assert.ok(!shown.includes("第 2001 行"));
  assert.match(shown, /截断：只显示了前 2000 行[\s\S]*共 2500 行/);
  assert.ok(!shown.includes("offset"));
});

// ---------------------------------------------------------------------------
// 实验员的 bash：环境变量只放 PATH、HOME、LANG、LC_ALL、TMPDIR
// ---------------------------------------------------------------------------

test("实验员 shell 里看不到插件进程环境里的密钥，也没有 Pi 的会话信息", async (t) => {
  const work = tempDir(t, "bash");
  process.env.INFORSI_FAKE_API_KEY = "sk-inforsi-fake-0123456789";
  t.after(() => delete process.env.INFORSI_FAKE_API_KEY);
  const task: SessionTask = { ...common, role: "experimenter", kind: "implementation", workDir: work, lightCheck: async () => "", submit: async () => ({ accepted: true, id: "x" }) };
  const bash = tool(setup(task).tools, "bash");
  const result = await bash.execute("call-1", { command: "env; pwd" }, signal, () => {}, undefined as never);
  const output = result.content.map((c) => (c.type === "text" ? c.text : "")).join("");
  assert.ok(!output.includes("sk-inforsi-fake-0123456789"), output);
  const names = output.split("\n").filter((l) => /^[A-Za-z_][A-Za-z0-9_]*=/.test(l)).map((l) => l.slice(0, l.indexOf("=")));
  const allowed = new Set(["PATH", "HOME", "LANG", "LC_ALL", "TMPDIR", "PWD", "SHLVL", "_", "OLDPWD"]);
  assert.deepEqual(names.filter((n) => !allowed.has(n)), [], `shell 里多出的变量：${names.join(", ")}`);
  assert.ok(names.includes("PATH"));
  assert.ok(output.includes(realpathSync(work)), "命令在工作目录里跑");
});

// ---------------------------------------------------------------------------
// 工具表
// ---------------------------------------------------------------------------

test("三种角色各自拿到的工具名正好是任务书第三部分那几个", (t) => {
  const work = tempDir(t, "tools");
  const expected = {
    researcher: ["read_source", "research_query", "research_submit", "write_notes"],
    verifier: ["read_source", "research_query", "research_submit"],
    experimenter: ["bash", "edit", "find", "grep", "ls", "read", "research_check", "research_query", "research_submit", "write"],
  };
  const tasks: SessionTask[] = [
    researcher("wrapUp", work, [work]),
    { ...common, role: "verifier", kind: "verification", workDir: work, readableDirs: [work], submit: async () => ({ accepted: true, id: "x" }) },
    { ...common, role: "experimenter", kind: "implementation", workDir: work, lightCheck: async () => "", submit: async () => ({ accepted: true, id: "x" }) },
  ];
  for (const task of tasks) {
    assert.deepEqual([...toolNames(task.role)].sort(), expected[task.role], task.role);
    // 自定义工具都在名单里；名单里其余的是 Pi 自带、由 Pi 按白名单装的那几个，只有实验员有。
    const custom = setup(task).tools.map((t) => t.name);
    assert.ok(custom.every((name) => toolNames(task.role).includes(name)), task.role);
    const builtins = toolNames(task.role).filter((name) => !custom.includes(name)).sort();
    assert.deepEqual(builtins, task.role === "experimenter" ? ["edit", "find", "grep", "ls", "read", "write"] : [], task.role);
  }
});