InfoRSI/src/flow/loop.ts

import { mkdirSync } from "node:fs";
import { resolve } from "node:path";
import type { ResearchConfig } from "../config/index.ts";
import { cellNotRunEvent, diagnosisNotRunEvent, forwardNotRunEvent } from "../materials/index.ts";
import { Ledger } from "../ledger/index.ts";
import { Runner } from "../runner/index.ts";
import { fullConfig, readFrozen } from "./frozen.ts";
import { EVENTS, nextStep, type Next } from "./next-step.ts";
import { ledgerFile } from "./paths.ts";
import { recover } from "./recover.ts";
import { baselineRun, saveBaseline } from "./steps/baseline.ts";
import { ablationRun, selectCells, settle, unsealChosen } from "./steps/ablation.ts";
import { diagnosisDedupe, diagnosisRun, diagnosisSettle, selectDiagnosisExperiments } from "./steps/diagnosis.ts";
import { forwardRun, adoptBatch } from "./steps/forward.ts";
import { dispatchStage } from "./steps/dispatch.ts";
import { dedupe, freezeProtocols, selectBatch } from "./steps/round.ts";
import { assignEntries } from "./steps/round-zero.ts";
import { freezeConfig, halt, openStudy, USER_STOPPED, type Study, type StudyEnding } from "./study.ts";

// 主循环：读账本 → 问下一步 → 做这一步 → 回到开头（任务书 12 第 2.4 节）。每做完一步都重新读账本。
// 每次启动走同一条路：新研究先冻结配置；之后把研究接上（核对镜像和提示词）、核对未完成的动作，再进入主循环。

export type StudyInput = {
  /** 研究目录：账本、代码仓库、每次运行和每次模型调用的目录都在这里（任务书 12 第 2.1 节）。 */
  studyDir: string;
  /** 新研究交 checkConfig 检查过的配置；接着跑一份已有的研究写 null，以账本里冻结的那份为准。 */
  config: ResearchConfig | null;
  /** 用户停下研究：正在跑的模型会话立即停，正在跑的容器不管它，下次启动时由恢复入口接着等。 */
  stop: AbortSignal;
};

export async function runStudy(input: StudyInput): Promise<StudyEnding> {
  const dir = resolve(input.studyDir);
  const config = input.config === null ? null : fullConfig(input.config);
  mkdirSync(dir, { recursive: true });
  const ledger = Ledger.open(ledgerFile(dir));
  try {
    const frozen = readFrozen(ledger.records()) !== null;
    if (!frozen && config === null) throw new Error(`研究目录 ${dir} 里还没有研究，配置不能写 null`);
    if (frozen && config !== null) throw new Error(`研究目录 ${dir} 里已经有冻结的配置：接着跑时配置写 null，以账本里冻结的那份为准`);
    // Docker 守护进程的地址交给 docker 命令按自己的环境变量和 context 决定。
    const runner = new Runner(null);
    // 恢复入口：把研究接上，核对未完成的动作，之后才问下一步（任务书 12 第 2.5 节）。新研究冻结配置之后走同一条路。
    const begin = async (): Promise<Study | StudyEnding> => {
      const opened = await openStudy(dir, ledger, runner, input.stop);
      if ("halt" in opened) return halt(ledger, opened.halt, []);
      const recovered = await recover(opened);
      if (recovered === "stopped") return USER_STOPPED;
      if (recovered !== null) return { status: "halted", reason: recovered.halt };
      return opened;
    };
    let study: Study | null = null;
    if (frozen) {
      const begun = await begin();
      if ("status" in begun) return begun;
      study = begun;
    }
    for (;;) {
      if (input.stop.aborted) return USER_STOPPED;
      const next = nextStep(ledger.records());
      if (next.do === "freezeConfig") {
        await freezeConfig(dir, ledger, runner, config!);
        const begun = await begin();
        if ("status" in begun) return begun;
        study = begun;
        continue;
      }
      const ended = await doStep(study!, next);
      if (ended !== null) return ended;
    }
  } finally {
    ledger.close();
  }
}

/** 做一步。交回 null 表示接着走下一步。 */
async function doStep(s: Study, next: Exclude<Next, { do: "freezeConfig" }>): Promise<StudyEnding | null> {
  switch (next.do) {
    case "saveBaseline":
      saveBaseline(s);
      return null;
    case "baselineRun":
    case "forwardRun":
    case "ablationRun":
    case "diagnosisRun": {
      const ran =
        next.do === "baselineRun" ? await baselineRun(s, next.run)
        : next.do === "forwardRun" ? await forwardRun(s, next.round, next.candidate, next.run)
        : next.do === "diagnosisRun" ? await diagnosisRun(s, next.round, next.run)
        : await ablationRun(s, next.round, next.cell, next.run);
      if (ran === "stopped") return USER_STOPPED;
      if (ran === "done") return null;
      return halt(s.ledger, `运行 ${next.run} 没有启动，名额已退还；本进程不重跑，下次启动时起新编号重跑。Docker 的报错：${ran.notStarted}`, [next.run]);
    }
    case "skip":
      s.ledger.append("event", { event: EVENTS.skipped, reason: next.reason, objects: [next.step] });
      return null;
    case "openStage":
      s.ledger.append("stageOpen", { stageId: next.stageId, stage: next.stage, round: next.round, roster: next.roster });
      return null;
    case "closeStage":
      s.ledger.append("stageClose", { stageId: next.stageId });
      return null;
    case "dispatch":
      return dispatchStage(s, next);
    case "assignEntries":
      assignEntries(s);
      return null;
    case "diagnosisDedupe":
      await diagnosisDedupe(s, next.round);
      return null;
    case "selectDiagnosis":
      selectDiagnosisExperiments(s, next.round);
      return null;
    case "diagnosisNotRun":
      s.ledger.append("event", diagnosisNotRunEvent(next.round, next.reason));
      return null;
    case "diagnosisSettle":
      await diagnosisSettle(s, next.round);
      return null;
    case "dedupe":
      await dedupe(s, next.round);
      return null;
    case "freezeProtocols":
      freezeProtocols(s, next.round);
      return null;
    case "selectBatch":
      selectBatch(s, next.round);
      return null;
    case "forwardNotRun":
      s.ledger.append("event", forwardNotRunEvent(next.candidate, next.reason));
      return null;
    case "adopt":
      adoptBatch(s, next.round);
      return null;
    case "selectCells":
      selectCells(s, next.round);
      return null;
    case "cellNotRun":
      s.ledger.append("event", cellNotRunEvent(next.cell, next.reason));
      return null;
    case "settle":
      settle(s, next.round);
      return null;
    case "unsealChosen":
      unsealChosen(s, next.round);
      return null;
    case "halt":
      return halt(s.ledger, next.reason, next.objects);
  }
}