spb/autoscholar

1import { chmod, mkdir, readFile, writeFile } from "node:fs/promises";2import { hostname } from "node:os";3import { join, resolve } from "node:path";4import { ControlPlaneError, WorkerClient } from "@autoscholar/maclustr-adapter";5import { createLogger, redact } from "@autoscholar/observability";6import { type RunnerHost, runResearch, simulateResearch } from "@autoscholar/research-runner";7import type { ArtifactKind, RunEventInput, StageKey } from "@autoscholar/schemas";8import { collectNodeMetadata, loadAverage, sdkVersion, workerVersion } from "./node-meta.js";9import { runSelfTest } from "./selftest.js";1011const log = createLogger("worker");1213interface Config {14  controlPlaneUrl: string;15  workerName: string;16  workerRoot: string;17  token?: string;18  enrollmentSecret?: string;19  mode: "real" | "simulated";20  maxConcurrent: number;21  profiles: string[];22  apiKey?: string;23  extraDomains: string[];24  tectonicBin?: string;25  latexTemplateDir: string;26  pollMs: number;27}2829function loadConfig(): Config {30  const e = process.env;31  const workerRoot = e.WORKER_ROOT ?? "/var/autoscholar";32  return {33    controlPlaneUrl: e.CONTROL_PLANE_URL ?? "http://127.0.0.1:8470",34    workerName: e.WORKER_NAME ?? hostname(),35    workerRoot,36    token: e.WORKER_TOKEN || undefined,37    enrollmentSecret: e.WORKER_ENROLLMENT_SECRET || undefined,38    mode: e.WORKER_MODE === "simulated" ? "simulated" : "real",39    maxConcurrent: Number(e.WORKER_MAX_CONCURRENT ?? 1),40    profiles: (e.WORKER_PROFILES ?? "standard")41      .split(",")42      .map((s) => s.trim())43      .filter(Boolean),44    apiKey: e.ANTHROPIC_API_KEY || undefined,45    extraDomains: (e.RESEARCH_ALLOWED_DOMAINS ?? "")46      .split(",")47      .map((s) => s.trim())48      .filter(Boolean),49    tectonicBin: e.TECTONIC_BIN || undefined,50    latexTemplateDir: e.LATEX_TEMPLATE_DIR ?? resolve(workerRoot, "templates/latex"),51    pollMs: Number(e.WORKER_POLL_MS ?? 4000),52  };53}5455class Worker {56  private readonly client: WorkerClient;57  private readonly running = new Map<58    string,59    { fencingToken: number; cancel: boolean; events: RunEventInput[] }60  >();61  private stopping = false;62  private draining = false;63  private tokenPath: string;6465  constructor(private readonly cfg: Config) {66    this.client = new WorkerClient({ baseUrl: cfg.controlPlaneUrl, token: cfg.token });67    this.tokenPath = join(cfg.workerRoot, "state/worker-token");68  }6970  async start() {71    await mkdir(join(this.cfg.workerRoot, "runs"), { recursive: true });72    await mkdir(join(this.cfg.workerRoot, "state"), { recursive: true });73    if (!this.cfg.token) {74      try {75        const saved = (await readFile(this.tokenPath, "utf8")).trim();76        if (saved) this.client.setToken(saved);77      } catch {78        /* none */79      }80    }81    const node = collectNodeMetadata({82      simulated: this.cfg.mode === "simulated",83      profiles: this.cfg.profiles,84    });85    if (this.cfg.mode === "real" && !node.sandboxAvailable)86      log.error(87        "sandbox-exec not available: this worker will never be dispatched research (sandboxAvailable=false)",88      );89    if (this.cfg.mode === "real" && !this.cfg.apiKey)90      log.error("ANTHROPIC_API_KEY missing: research sessions cannot start");91    // enroll if needed92    let enrolled = false;93    for (let i = 0; i < 30 && !enrolled && !this.stopping; i++) {94      try {95        await this.client.heartbeat({96          status: "online",97          capacity: this.cfg.maxConcurrent,98          running: [],99          load: loadAverage(),100          node,101          leases: [],102        });103        enrolled = true;104      } catch (e) {105        if (e instanceof ControlPlaneError && e.status === 401 && this.cfg.enrollmentSecret) {106          try {107            const r = await this.client.enroll(this.cfg.enrollmentSecret, this.cfg.workerName, node);108            this.client.setToken(r.token);109            await writeFile(this.tokenPath, r.token, { mode: 0o600 });110            await chmod(this.tokenPath, 0o600);111            log.info({ workerId: r.workerId }, "enrolled with control plane");112            enrolled = true;113          } catch (e2) {114            log.error({ err: (e2 as Error).message }, "enrollment failed");115          }116        } else log.warn({ err: (e as Error).message, attempt: i }, "control plane unreachable; retrying");117        if (!enrolled) await new Promise((r) => setTimeout(r, 5000));118      }119    }120    if (!enrolled) throw new Error("could not reach/enroll with the control plane");121    log.info(122      {123        name: this.cfg.workerName,124        mode: this.cfg.mode,125        sdk: sdkVersion(),126        worker: workerVersion(),127        controlPlane: this.cfg.controlPlaneUrl,128        sandbox: node.sandboxAvailable,129      },130      "worker online",131    );132    void this.heartbeatLoop(node);133    void this.leaseLoop();134  }135136  private async heartbeatLoop(node: ReturnType<typeof collectNodeMetadata>) {137    while (!this.stopping) {138      try {139        const leases = [...this.running.entries()].map(([runId, r]) => ({140          runId,141          fencingToken: r.fencingToken,142        }));143        const hb = await this.client.heartbeat({144          status: this.draining ? "draining" : "online",145          capacity: this.cfg.maxConcurrent,146          running: [...this.running.keys()],147          load: loadAverage(),148          leases,149          node: { tools: node.tools, sandboxAvailable: node.sandboxAvailable },150        });151        this.draining = hb.drain;152        for (const id of hb.cancel) {153          const r = this.running.get(id);154          if (r) r.cancel = true;155        }156        for (const id of hb.revokedLeases) {157          const r = this.running.get(id);158          if (r) {159            log.warn({ runId: id }, "lease revoked by control plane (fencing) — aborting local work");160            r.cancel = true;161          }162        }163        await this.flushEvents();164      } catch (e) {165        log.warn({ err: (e as Error).message }, "heartbeat failed");166      }167      await new Promise((r) => setTimeout(r, 20_000));168    }169  }170171  private async flushEvents() {172    for (const [runId, r] of this.running) {173      if (r.events.length === 0) continue;174      const batch = r.events.splice(0, 200);175      try {176        await this.client.events(177          runId,178          r.fencingToken,179          batch.map((e) => ({ ...e, level: e.level ?? "info" })),180        );181      } catch (e) {182        if (e instanceof ControlPlaneError && e.status === 409) {183          r.cancel = true;184          log.warn({ runId }, "lease lost while posting events; cancelling local work");185        } else r.events.unshift(...batch.slice(0, 50));186      }187    }188  }189190  private async leaseLoop() {191    while (!this.stopping) {192      if (!this.draining && this.running.size < this.cfg.maxConcurrent) {193        try {194          const { job } = await this.client.lease(this.cfg.profiles);195          if (job) {196            log.info({ runId: job.runId, attempt: job.attemptNo, resume: job.resumeFromStage }, "leased job");197            void this.execute(job);198          }199        } catch (e) {200          log.warn({ err: (e as Error).message }, "lease request failed");201        }202      }203      await new Promise((r) => setTimeout(r, this.cfg.pollMs));204    }205  }206207  private async execute(job: Awaited<ReturnType<WorkerClient["lease"]>>["job"] & object) {208    const state = {209      fencingToken: job.fencingToken,210      cancel: job.cancelRequested,211      events: [] as RunEventInput[],212    };213    this.running.set(job.runId, state);214    const renew = setInterval(async () => {215      try {216        const r = await this.client.renew(job.runId, job.fencingToken);217        if (r.cancelRequested) state.cancel = true;218        await this.flushEvents();219      } catch (e) {220        if (e instanceof ControlPlaneError && e.status === 409) {221          state.cancel = true;222          log.warn({ runId: job.runId }, "lease lost (stale fencing token); aborting");223        }224      }225    }, 30_000);226    const host: RunnerHost = {227      workspaceRoot: this.cfg.workerRoot,228      apiKey: this.cfg.apiKey ?? "",229      extraDomains: this.cfg.extraDomains,230      tectonicBin: this.cfg.tectonicBin,231      tectonicCache: join(this.cfg.workerRoot, "cache/tectonic"),232      latexTemplateDir: this.cfg.latexTemplateDir,233      workerName: this.cfg.workerName,234      sdkVersion: sdkVersion(),235      toolchain: collectNodeMetadata({236        simulated: this.cfg.mode === "simulated",237        profiles: this.cfg.profiles,238      }).tools,239      emit: (ev) => {240        state.events.push(redact(ev));241        if (state.events.length >= 50) void this.flushEvents();242      },243      stage: async (report) => {244        await this.flushEvents();245        await this.client.stage(job.runId, { fencingToken: job.fencingToken, ...report });246      },247      upload: async (path: string, kind: ArtifactKind, contentType: string, buf: Buffer) => {248        await this.client.uploadArtifact(249          job.runId,250          job.fencingToken,251          {252            path,253            kind,254            contentType,255            stage: (job.resumeFromStage ?? "preflight") as StageKey,256            synthetic: this.cfg.mode === "simulated",257          },258          buf,259        );260      },261      isCancelled: () => state.cancel || this.stopping,262    };263    let outcome: Awaited<ReturnType<typeof runResearch>>;264    try {265      const runnerJob = {266        runId: job.runId,267        attemptNo: job.attemptNo,268        manifest: job.manifest,269        instructionText: job.instructionText,270        resumeFromStage: job.resumeFromStage,271        checkpoint: job.checkpoint,272      };273      if (this.cfg.mode === "simulated") {274        const sim = process.env.WORKER_SIM_BEHAVIOUR ?? "";275        outcome = await simulateResearch(runnerJob, host, {276          failAtStage: sim.startsWith("fail:") ? sim.slice(5) : undefined,277          hang: sim === "hang",278          delayMs: Number(process.env.WORKER_SIM_DELAY_MS ?? 0),279        });280      } else {281        if (!this.cfg.apiKey) throw new Error("ANTHROPIC_API_KEY not configured on this worker");282        outcome = await runResearch(runnerJob, host);283      }284      if (state.restart && outcome.outcome === "cancelled") {285        outcome = {286          ...outcome,287          outcome: "failed",288          error: "WORKER_RESTART: worker daemon restarted; resume from checkpoint",289        };290      }291    } catch (e) {292      outcome = {293        outcome: "failed",294        error: `runner crashed: ${(e as Error).message}`.slice(0, 2000),295        costUsd: 0,296        turns: 0,297        stagesCompleted: [],298        resourceUsage: { wallSeconds: 0, diskBytes: 0, peakRssBytes: 0 },299      };300      log.error(301        { runId: job.runId, err: (e as Error).message, stack: (e as Error).stack?.split("\n").slice(0, 5) },302        "runner crashed",303      );304    }305    clearInterval(renew);306    await this.flushEvents();307    try {308      await this.client.complete(job.runId, {309        fencingToken: job.fencingToken,310        outcome: outcome.outcome,311        error: outcome.error,312        costUsd: outcome.costUsd,313        summary: { stagesCompleted: outcome.stagesCompleted, turns: outcome.turns },314        resourceUsage: outcome.resourceUsage,315      });316      log.info({ runId: job.runId, outcome: outcome.outcome, costUsd: outcome.costUsd }, "attempt completed");317    } catch (e) {318      log.error({ runId: job.runId, err: (e as Error).message }, "could not report completion (lease lost?)");319    }320    this.running.delete(job.runId);321  }322323  async stop() {324    this.stopping = true;325    // graceful restart: abort sessions (isCancelled() → true via `stopping`) but report a retryable failure, not a cancel326    for (const r of this.running.values()) r.restart = true;327    const t0 = Date.now();328    while (this.running.size > 0 && Date.now() - t0 < 60_000) await new Promise((r) => setTimeout(r, 500));329  }330}331332const cfg = loadConfig();333if (process.argv[2] === "selftest") {334  runSelfTest(cfg.workerRoot, cfg.apiKey, cfg.latexTemplateDir, cfg.tectonicBin)335    .then((r) => {336      console.log(JSON.stringify(r, null, 2));337      process.exit(r.ok ? 0 : 1);338    })339    .catch((e) => {340      console.error(e);341      process.exit(1);342    });343} else {344  const w = new Worker(cfg);345  process.on("SIGTERM", () => void w.stop().then(() => process.exit(0)));346  process.on("SIGINT", () => void w.stop().then(() => process.exit(0)));347  w.start().catch((e) => {348    log.fatal({ err: (e as Error).message }, "worker failed to start");349    process.exit(1);350  });351}352