spb/autoscholar

1import { constantTimeEqual, hashToken, newWorkerToken } from "@autoscholar/auth";2import {3  artifacts,4  checkpoints,5  computeWorkers,6  type Db,7  jobLeases,8  projectInstructionVersions,9  providerUsage,10  researchRuns,11  resourceUsage,12  runAttempts,13  stageExecutions,14  type Tx,15  workerAssignments,16  workerHeartbeats,17} from "@autoscholar/db";18import { redact } from "@autoscholar/observability";19import {20  type ArtifactDescriptor,21  type NodeMetadata,22  type RunEvent,23  type RunManifest,24  STAGES,25  type StageKey,26  WORKER_PROTOCOL_VERSION,27  type WorkerStatus,28} from "@autoscholar/schemas";29import { type FilesystemStorage, runArtifactKey } from "@autoscholar/storage";30import { and, asc, desc, desc as descOrder, eq, inArray, isNull, lt, sql } from "drizzle-orm";31import { audit } from "./audit.js";32import { getEnv } from "./env.js";33import { AppError, conflict, forbidden, invalid, notFound, unauthorized } from "./errors.js";34import { appendEvent, latestCheckpoint, transition } from "./runs.js";35import { getSetting } from "./settings.js";3637export type WorkerRow = typeof computeWorkers.$inferSelect;3839export async function enrollWorker(40  db: Db,41  enrollmentSecret: string,42  name: string,43  node: NodeMetadata,44  lanIp: string | undefined,45  ip: string | null,46) {47  const env = getEnv();48  if (!env.WORKER_ENROLLMENT_SECRET || !constantTimeEqual(enrollmentSecret, env.WORKER_ENROLLMENT_SECRET)) {49    throw unauthorized("invalid enrollment secret");50  }51  if (node.protocolVersion !== WORKER_PROTOCOL_VERSION)52    throw new AppError(53      "protocol_mismatch",54      `worker protocol ${node.protocolVersion} != ${WORKER_PROTOCOL_VERSION}`,55      426,56    );57  const { token, tokenHash } = newWorkerToken();58  const [w] = await db59    .insert(computeWorkers)60    .values({61      name,62      tokenHash,63      status: "online",64      version: node.workerVersion,65      protocolVersion: node.protocolVersion,66      lanIp: lanIp ?? ip,67      nodeMeta: node,68      profiles: node.profiles,69      lastHeartbeatAt: new Date(),70      tokenRotatedAt: new Date(),71    })72    .onConflictDoUpdate({73      target: computeWorkers.name,74      set: {75        tokenHash,76        status: "online",77        version: node.workerVersion,78        protocolVersion: node.protocolVersion,79        lanIp: lanIp ?? ip,80        nodeMeta: node,81        profiles: node.profiles,82        lastHeartbeatAt: new Date(),83        tokenRotatedAt: new Date(),84        disabledAt: null,85      },86    })87    .returning();88  await audit(89    db,90    { type: "worker", id: w!.id, ip },91    "worker.enroll",92    { type: "worker", id: w!.id },93    { name, machineId: node.machineId, sandboxAvailable: node.sandboxAvailable },94  );95  return { workerId: w!.id, token, protocolVersion: WORKER_PROTOCOL_VERSION };96}9798export async function authenticateWorker(db: Db, bearer: string | null | undefined): Promise<WorkerRow> {99  if (!bearer?.startsWith("asw_")) throw unauthorized("worker token required");100  const [w] = await db101    .select()102    .from(computeWorkers)103    .where(eq(computeWorkers.tokenHash, hashToken(bearer)))104    .limit(1);105  if (!w) throw unauthorized("unknown worker token");106  if (w.disabledAt || w.status === "disabled") throw forbidden("worker disabled");107  if (w.status === "quarantined") throw forbidden(`worker quarantined: ${w.quarantineReason ?? ""}`);108  return w;109}110111export async function heartbeat(112  db: Db,113  worker: WorkerRow,114  body: {115    status: WorkerStatus;116    capacity: number;117    running: string[];118    load: Record<string, unknown>;119    node?: Partial<NodeMetadata>;120    leases: Array<{ runId: string; fencingToken: number }>;121  },122) {123  const env = getEnv();124  const status: WorkerStatus = worker.drainRequested ? "draining" : body.status;125  await db126    .update(computeWorkers)127    .set({128      status,129      capacity: body.capacity,130      running: body.running.length,131      lastHeartbeatAt: new Date(),132      nodeMeta: body.node ? { ...(worker.nodeMeta as object), ...body.node } : worker.nodeMeta,133    })134    .where(eq(computeWorkers.id, worker.id));135  await db136    .insert(workerHeartbeats)137    .values({ workerId: worker.id, load: body.load, running: body.running.length });138  // leases the worker thinks it holds but which are no longer its own (fencing) → tell it to stop139  const revokedLeases: string[] = [];140  const cancel: string[] = [];141  for (const l of body.leases) {142    const [lease] = await db.select().from(jobLeases).where(eq(jobLeases.runId, l.runId)).limit(1);143    if (!lease || lease.workerId !== worker.id || lease.fencingToken !== l.fencingToken)144      revokedLeases.push(l.runId);145    else {146      await db147        .update(jobLeases)148        .set({ lastHeartbeatAt: new Date(), expiresAt: new Date(Date.now() + env.LEASE_TTL_SECONDS * 1000) })149        .where(eq(jobLeases.runId, l.runId));150      const [run] = await db151        .select({ cancelRequested: researchRuns.cancelRequested })152        .from(researchRuns)153        .where(eq(researchRuns.id, l.runId))154        .limit(1);155      if (run?.cancelRequested) cancel.push(l.runId);156    }157  }158  return {159    ok: true as const,160    drain: !!worker.drainRequested,161    cancel,162    revokedLeases,163    leaseTtlSeconds: env.LEASE_TTL_SECONDS,164  };165}166167/** Lease the oldest queued run matching the worker's profiles. SKIP LOCKED makes concurrent workers safe. */168export async function leaseJob(db: Db, worker: WorkerRow, profiles: string[]) {169  const env = getEnv();170  if (worker.drainRequested || worker.status === "draining") return { job: null };171  if (env.DISPATCH_ENABLED !== 1 || (await getSetting(db, "dispatch.paused", false))) return { job: null };172  const node = worker.nodeMeta as Partial<NodeMetadata>;173  if (!node.sandboxAvailable) return { job: null }; // never dispatch to a worker that cannot enforce isolation174  return db.transaction(async (tx) => {175    const [run] = await tx176      .select()177      .from(researchRuns)178      .where(179        and(180          inArray(researchRuns.status, ["queued", "retrying"]),181          inArray(researchRuns.computeProfile, profiles.length ? profiles : ["standard"]),182        ),183      )184      .orderBy(asc(researchRuns.queuedAt))185      .limit(1)186      .for("update", { skipLocked: true });187    if (!run) return { job: null };188    const [existingLease] = await tx.select().from(jobLeases).where(eq(jobLeases.runId, run.id)).limit(1);189    if (existingLease && existingLease.expiresAt.getTime() > Date.now()) return { job: null };190    if (run.currentAttemptNo >= run.maxAttempts) {191      await transition(tx, run.id, "failed", {192        reason: `attempt limit (${run.maxAttempts}) reached`,193        actor: "scheduler",194      });195      return { job: null };196    }197    const attemptNo = run.currentAttemptNo + 1;198    const fencingToken = run.fencingToken + 1;199    const cp = await latestCheckpoint(tx, run.id);200    const resumeFromStage = (cp?.stage as StageKey | undefined) ?? null;201    const [attempt] = await tx202      .insert(runAttempts)203      .values({204        runId: run.id,205        attemptNo,206        workerId: worker.id,207        status: "leased",208        fencingToken,209        resumeFromStage,210      })211      .returning();212    const expiresAt = new Date(Date.now() + env.LEASE_TTL_SECONDS * 1000);213    await tx.delete(jobLeases).where(eq(jobLeases.runId, run.id));214    await tx215      .insert(jobLeases)216      .values({ runId: run.id, attemptId: attempt!.id, workerId: worker.id, fencingToken, expiresAt });217    await tx.insert(workerAssignments).values({ runId: run.id, attemptId: attempt!.id, workerId: worker.id });218    await tx219      .update(researchRuns)220      .set({ currentAttemptNo: attemptNo, fencingToken, updatedAt: new Date() })221      .where(eq(researchRuns.id, run.id));222    await transition(tx, run.id, "provisioning", {223      actor: `worker:${worker.name}`,224      attemptId: attempt!.id,225      workerId: worker.id,226      data: { attemptNo, fencingToken, resumeFromStage },227    });228    await tx229      .update(computeWorkers)230      .set({ running: sql`${computeWorkers.running} + 1` })231      .where(eq(computeWorkers.id, worker.id));232    const [iv] = await tx233      .select({ contentText: projectInstructionVersions.contentText })234      .from(projectInstructionVersions)235      .where(eq(projectInstructionVersions.id, run.instructionVersionId))236      .limit(1);237    return {238      job: {239        runId: run.id,240        attemptId: attempt!.id,241        attemptNo,242        fencingToken,243        leaseExpiresAt: expiresAt.toISOString(),244        manifest: run.manifest as RunManifest,245        instructionText: iv!.contentText,246        resumeFromStage,247        checkpoint: (cp?.data as Record<string, unknown> | undefined) ?? null,248        cancelRequested: run.cancelRequested,249      },250    };251  });252}253254/** Validates that the worker holds the current lease with the exact fencing token. */255export async function assertLease(tx: Tx | Db, runId: string, worker: WorkerRow, fencingToken: number) {256  const [lease] = await tx.select().from(jobLeases).where(eq(jobLeases.runId, runId)).limit(1);257  if (!lease || lease.workerId !== worker.id || lease.fencingToken !== fencingToken) {258    throw new AppError("lease_lost", "lease is not held by this worker (stale fencing token)", 409);259  }260  return lease;261}262263export async function renewLease(db: Db, worker: WorkerRow, runId: string, fencingToken: number) {264  const env = getEnv();265  const lease = await assertLease(db, runId, worker, fencingToken);266  const expiresAt = new Date(Date.now() + env.LEASE_TTL_SECONDS * 1000);267  await db268    .update(jobLeases)269    .set({ expiresAt, lastHeartbeatAt: new Date() })270    .where(eq(jobLeases.runId, lease.runId));271  const [run] = await db272    .select({ cancelRequested: researchRuns.cancelRequested, status: researchRuns.status })273    .from(researchRuns)274    .where(eq(researchRuns.id, runId))275    .limit(1);276  return {277    ok: true,278    cancelRequested: !!run?.cancelRequested || run?.status === "cancelling",279    leaseExpiresAt: expiresAt.toISOString(),280  };281}282283export async function ingestEvents(284  db: Db,285  worker: WorkerRow,286  runId: string,287  fencingToken: number,288  events: RunEvent[],289) {290  return db.transaction(async (tx) => {291    const lease = await assertLease(tx, runId, worker, fencingToken);292    let last = 0;293    for (const ev of events)294      last = await appendEvent(tx, runId, { ...redact(ev), attemptId: lease.attemptId, workerId: worker.id });295    return { ok: true as const, lastSeq: last };296  });297}298299export async function reportStage(300  db: Db,301  worker: WorkerRow,302  runId: string,303  report: {304    fencingToken: number;305    stage: StageKey;306    status: "running" | "completed" | "failed" | "skipped";307    costUsd: number;308    turns: number;309    usage?: Record<string, unknown>;310    error?: string;311    checkpoint?: Record<string, unknown>;312  },313) {314  return db.transaction(async (tx) => {315    const lease = await assertLease(tx, runId, worker, report.fencingToken);316    const [run] = await tx.select().from(researchRuns).where(eq(researchRuns.id, runId)).for("update");317    if (!run) throw notFound("run");318    if (report.status === "running") {319      await tx320        .insert(stageExecutions)321        .values({ runId, attemptId: lease.attemptId, stage: report.stage, status: "running" })322        .onConflictDoUpdate({323          target: [stageExecutions.attemptId, stageExecutions.stage],324          set: { status: "running", startedAt: new Date(), endedAt: null, error: null },325        });326      await tx327        .update(researchRuns)328        .set({ currentStage: report.stage, updatedAt: new Date() })329        .where(eq(researchRuns.id, runId));330      if (run.status === "provisioning")331        await transition(tx, runId, "running", {332          actor: `worker:${worker.name}`,333          attemptId: lease.attemptId,334          workerId: worker.id,335        });336      await tx.update(runAttempts).set({ status: "running" }).where(eq(runAttempts.id, lease.attemptId));337    } else {338      await tx339        .insert(stageExecutions)340        .values({341          runId,342          attemptId: lease.attemptId,343          stage: report.stage,344          status: report.status,345          endedAt: new Date(),346          costUsd: String(report.costUsd),347          turns: report.turns,348          usage: report.usage ?? null,349          error: report.error ?? null,350        })351        .onConflictDoUpdate({352          target: [stageExecutions.attemptId, stageExecutions.stage],353          set: {354            status: report.status,355            endedAt: new Date(),356            costUsd: String(report.costUsd),357            turns: report.turns,358            usage: report.usage ?? null,359            error: report.error ?? null,360          },361        });362      if (report.costUsd > 0) {363        const u = (report.usage ?? {}) as Record<string, number | string>;364        await tx.insert(providerUsage).values({365          runId,366          attemptId: lease.attemptId,367          stage: report.stage,368          model: String(u.model ?? run.model),369          inputTokens: Number(u.inputTokens ?? 0),370          outputTokens: Number(u.outputTokens ?? 0),371          cacheReadTokens: Number(u.cacheReadTokens ?? 0),372          cacheCreateTokens: Number(u.cacheCreateTokens ?? 0),373          costUsd: String(report.costUsd),374          turns: report.turns,375        });376        await tx377          .update(researchRuns)378          .set({ totalCostUsd: sql`${researchRuns.totalCostUsd} + ${String(report.costUsd)}` })379          .where(eq(researchRuns.id, runId));380        await tx381          .update(runAttempts)382          .set({ costUsd: sql`${runAttempts.costUsd} + ${String(report.costUsd)}` })383          .where(eq(runAttempts.id, lease.attemptId));384      }385      if (report.status === "completed" && report.checkpoint) {386        await tx387          .insert(checkpoints)388          .values({ runId, attemptId: lease.attemptId, stage: report.stage, data: report.checkpoint });389      }390    }391    await appendEvent(tx, runId, {392      kind: `stage.${report.status}`,393      level: report.status === "failed" ? "error" : "info",394      stage: report.stage,395      message: `${STAGES.find((s) => s.key === report.stage)?.name ?? report.stage} ${report.status}${report.error ? `: ${report.error}` : ""}`,396      data: { costUsd: report.costUsd, turns: report.turns },397      at: new Date().toISOString(),398      attemptId: lease.attemptId,399      workerId: worker.id,400    });401    return { ok: true as const };402  });403}404405export async function ingestArtifact(406  db: Db,407  storage: FilesystemStorage,408  worker: WorkerRow,409  runId: string,410  fencingToken: number,411  artifact: ArtifactDescriptor,412  body: Buffer,413) {414  const desc = artifact;415  if (body.length !== desc.bytes) throw invalid(`byte length mismatch: ${body.length} != ${desc.bytes}`);416  if (body.length > 512 * 1024 * 1024) throw invalid("artifact exceeds 512 MiB");417  const lease = await assertLease(db, runId, worker, fencingToken);418  const [attempt] = await db.select().from(runAttempts).where(eq(runAttempts.id, lease.attemptId)).limit(1);419  const key = runArtifactKey(runId, attempt!.attemptNo, desc.path);420  const stored = await storage.put(key, body, desc.sha256);421  return db.transaction(async (tx) => {422    await assertLease(tx, runId, worker, fencingToken);423    const [prev] = await tx424      .select({ version: artifacts.version, sha256: artifacts.sha256, id: artifacts.id })425      .from(artifacts)426      .where(and(eq(artifacts.runId, runId), eq(artifacts.path, desc.path)))427      .orderBy(descOrder(artifacts.version))428      .limit(1);429    if (prev && prev.sha256 === stored.sha256)430      return { ok: true as const, artifactId: prev.id, sha256: stored.sha256, bytes: stored.bytes };431    const version = (prev?.version ?? 0) + 1;432    const versionedKey =433      version === 1 ? key : runArtifactKey(runId, attempt!.attemptNo, `${desc.path}.v${version}`);434    if (version > 1) await storage.put(versionedKey, body, desc.sha256);435    const [row] = await tx436      .insert(artifacts)437      .values({438        runId,439        attemptId: lease.attemptId,440        stage: desc.stage,441        path: desc.path,442        version,443        kind: desc.kind,444        sha256: stored.sha256,445        bytes: stored.bytes,446        contentType: desc.contentType,447        storageKey: versionedKey,448        synthetic: desc.synthetic,449      })450      .returning({ id: artifacts.id });451    return { ok: true as const, artifactId: row!.id, sha256: stored.sha256, bytes: stored.bytes };452  });453}454455export async function completeAttempt(456  db: Db,457  worker: WorkerRow,458  runId: string,459  body: {460    fencingToken: number;461    outcome: "completed" | "failed" | "cancelled" | "budget_exceeded" | "timeout";462    error?: string;463    summary?: Record<string, unknown>;464    costUsd: number;465    resourceUsage?: { cpuSeconds?: number; wallSeconds?: number; diskBytes?: number; peakRssBytes?: number };466  },467) {468  return db.transaction(async (tx) => {469    const lease = await assertLease(tx, runId, worker, body.fencingToken);470    const [run] = await tx.select().from(researchRuns).where(eq(researchRuns.id, runId)).for("update");471    if (!run) throw notFound("run");472    await tx473      .update(runAttempts)474      .set({475        status: body.outcome,476        endedAt: new Date(),477        error: body.error ?? null,478        summary: body.summary ?? null,479      })480      .where(eq(runAttempts.id, lease.attemptId));481    if (body.resourceUsage)482      await tx.insert(resourceUsage).values({483        runId,484        attemptId: lease.attemptId,485        workerId: worker.id,486        cpuSeconds: body.resourceUsage.cpuSeconds?.toFixed(2),487        wallSeconds: body.resourceUsage.wallSeconds?.toFixed(2),488        diskBytes: body.resourceUsage.diskBytes,489        peakRssBytes: body.resourceUsage.peakRssBytes,490      });491    await tx.delete(jobLeases).where(eq(jobLeases.runId, runId));492    await tx493      .update(workerAssignments)494      .set({ releasedAt: new Date(), releaseReason: body.outcome })495      .where(eq(workerAssignments.attemptId, lease.attemptId));496    await tx497      .update(computeWorkers)498      .set({ running: sql`greatest(${computeWorkers.running} - 1, 0)` })499      .where(eq(computeWorkers.id, worker.id));500    const base = { actor: `worker:${worker.name}`, attemptId: lease.attemptId, workerId: worker.id };501    if (body.outcome === "completed") {502      await transition(tx, runId, "quality_check", {503        ...base,504        reason: "worker reported completion; awaiting deterministic quality evaluation",505      });506    } else if (body.outcome === "cancelled" || run.status === "cancelling") {507      await transition(tx, runId, "cancelled", { ...base, reason: body.error ?? "cancelled" });508    } else {509      const reason = `${body.outcome}${body.error ? `: ${body.error}` : ""}`;510      if (body.outcome === "budget_exceeded" || run.currentAttemptNo >= run.maxAttempts) {511        await transition(tx, runId, "failed", { ...base, reason });512      } else {513        await transition(tx, runId, "retrying", {514          ...base,515          reason: `${reason}; attempt ${run.currentAttemptNo}/${run.maxAttempts}`,516        });517      }518    }519    return { ok: true as const };520  });521}522523/** Scheduler: expire stale leases and recover runs (retry or fail). Returns the run IDs touched. */524export async function expireStaleLeases(db: Db): Promise<string[]> {525  const stale = await db.select().from(jobLeases).where(lt(jobLeases.expiresAt, new Date()));526  const touched: string[] = [];527  for (const lease of stale) {528    await db.transaction(async (tx) => {529      const [run] = await tx530        .select()531        .from(researchRuns)532        .where(eq(researchRuns.id, lease.runId))533        .for("update");534      if (!run) return;535      const [still] = await tx536        .select()537        .from(jobLeases)538        .where(and(eq(jobLeases.runId, lease.runId), eq(jobLeases.fencingToken, lease.fencingToken)))539        .limit(1);540      if (!still || still.expiresAt.getTime() > Date.now()) return;541      await tx.delete(jobLeases).where(eq(jobLeases.runId, lease.runId));542      await tx543        .update(runAttempts)544        .set({ status: "expired", endedAt: new Date(), error: "lease expired (worker heartbeat lost)" })545        .where(eq(runAttempts.id, lease.attemptId));546      await tx547        .update(workerAssignments)548        .set({ releasedAt: new Date(), releaseReason: "lease_expired" })549        .where(eq(workerAssignments.attemptId, lease.attemptId));550      await tx551        .update(computeWorkers)552        .set({ running: sql`greatest(${computeWorkers.running} - 1, 0)` })553        .where(eq(computeWorkers.id, lease.workerId));554      const base = { actor: "scheduler", attemptId: lease.attemptId, workerId: lease.workerId };555      if (["provisioning", "running", "cancelling"].includes(run.status)) {556        if (run.status === "cancelling")557          await transition(tx, run.id, "cancelled", { ...base, reason: "lease expired during cancellation" });558        else if (run.currentAttemptNo >= run.maxAttempts)559          await transition(tx, run.id, "failed", {560            ...base,561            reason: "lease expired and attempt limit reached",562          });563        else564          await transition(tx, run.id, "retrying", {565            ...base,566            reason: `lease expired (attempt ${run.currentAttemptNo}/${run.maxAttempts}); will resume from last checkpoint`,567          });568      }569      touched.push(run.id);570    });571  }572  return touched;573}574575/** Scheduler: retrying → queued (immediate re-queue; backoff is applied through queuedAt ordering). */576export async function requeueRetrying(db: Db): Promise<number> {577  const rows = await db578    .select({ id: researchRuns.id })579    .from(researchRuns)580    .where(eq(researchRuns.status, "retrying"));581  for (const r of rows) await db.transaction((tx) => transition(tx, r.id, "queued", { actor: "scheduler" }));582  return rows.length;583}584585export async function markOfflineWorkers(db: Db, staleAfterMs = 90_000): Promise<number> {586  const cutoff = new Date(Date.now() - staleAfterMs);587  const rows = await db588    .update(computeWorkers)589    .set({ status: "offline" })590    .where(591      and(inArray(computeWorkers.status, ["online", "draining"]), lt(computeWorkers.lastHeartbeatAt, cutoff)),592    )593    .returning({ id: computeWorkers.id });594  return rows.length;595}596597export async function listWorkers(db: Db) {598  const rows = await db.select().from(computeWorkers).orderBy(asc(computeWorkers.name));599  return rows.map((w) => ({ ...w, tokenHash: undefined }));600}601602export async function setWorkerDrain(db: Db, workerId: string, drain: boolean, actor: { id: string }) {603  const [w] = await db604    .update(computeWorkers)605    .set({ drainRequested: drain ? new Date() : null, status: drain ? "draining" : "online" })606    .where(eq(computeWorkers.id, workerId))607    .returning();608  if (!w) throw notFound("worker");609  await audit(db, { type: "admin", id: actor.id }, drain ? "worker.drain" : "worker.undrain", {610    type: "worker",611    id: workerId,612  });613  return w;614}615616export async function setWorkerQuarantine(617  db: Db,618  workerId: string,619  reason: string | null,620  actor: { id: string },621) {622  const [w] = await db623    .update(computeWorkers)624    .set({ status: reason ? "quarantined" : "online", quarantineReason: reason })625    .where(eq(computeWorkers.id, workerId))626    .returning();627  if (!w) throw notFound("worker");628  await audit(629    db,630    { type: "admin", id: actor.id },631    reason ? "worker.quarantine" : "worker.unquarantine",632    { type: "worker", id: workerId },633    { reason },634  );635  return w;636}637638export async function workerTelemetry(db: Db, workerId: string, limit = 60) {639  return db640    .select()641    .from(workerHeartbeats)642    .where(eq(workerHeartbeats.workerId, workerId))643    .orderBy(desc(workerHeartbeats.createdAt))644    .limit(limit);645}646647export { conflict, isNull };648