spb/autoscholar

1import { createHash } from "node:crypto";2import {3  ArtifactDescriptor,4  CompleteAttemptRequest,5  EnrollResponse,6  type HeartbeatRequest,7  HeartbeatResponse,8  LeaseRenewResponse,9  LeaseResponse,10  type NodeMetadata,11  type RunEvent,12  StageReport,13  WORKER_PROTOCOL_VERSION,14} from "@autoscholar/schemas";15import type { z } from "zod";1617export class ControlPlaneError extends Error {18  constructor(19    message: string,20    readonly status: number,21    readonly code?: string,22  ) {23    super(message);24  }25}2627export interface WorkerClientOptions {28  baseUrl: string;29  token?: string;30  fetchImpl?: typeof fetch;31  timeoutMs?: number;32  userAgent?: string;33  retryAttempts?: number;34}3536/** Worker-side HTTP client for the versioned MacLustr worker protocol (outbound-only). */37export class WorkerClient {38  private token: string | undefined;39  private readonly base: string;40  private readonly fetchImpl: typeof fetch;41  private readonly timeoutMs: number;42  private readonly userAgent: string;43  private readonly retryAttempts: number;4445  constructor(opts: WorkerClientOptions) {46    this.base = opts.baseUrl.replace(/\/+$/, "");47    this.token = opts.token;48    this.fetchImpl = opts.fetchImpl ?? fetch;49    this.timeoutMs = opts.timeoutMs ?? 30_000;50    this.userAgent = opts.userAgent ?? `autoscholar-worker/protocol-${WORKER_PROTOCOL_VERSION}`;51    this.retryAttempts = opts.retryAttempts ?? 8;52  }5354  setToken(token: string): void {55    this.token = token;56  }5758  private async call<T>(59    method: string,60    path: string,61    body?: unknown,62    schema?: z.ZodType<T>,63    raw?: { buffer: Buffer; contentType: string },64  ): Promise<T> {65    // transient failures (network, 5xx) are retried with backoff so a short control-plane restart never kills a run66    const maxAttempts = this.retryAttempts;67    let lastErr: unknown;68    for (let attempt = 1; attempt <= maxAttempts; attempt++) {69      try {70        return await this.callOnce(method, path, body, schema, raw);71      } catch (e) {72        lastErr = e;73        const transient = !(e instanceof ControlPlaneError) || e.status >= 500;74        if (!transient || attempt === maxAttempts) throw e;75        await new Promise((r) => setTimeout(r, Math.min(15_000, 1000 * 2 ** (attempt - 1))));76      }77    }78    throw lastErr;79  }8081  private async callOnce<T>(82    method: string,83    path: string,84    body?: unknown,85    schema?: z.ZodType<T>,86    raw?: { buffer: Buffer; contentType: string },87  ): Promise<T> {88    const ctrl = new AbortController();89    const t = setTimeout(() => ctrl.abort(), this.timeoutMs);90    try {91      const headers: Record<string, string> = {92        "user-agent": this.userAgent,93        "x-protocol-version": String(WORKER_PROTOCOL_VERSION),94      };95      if (this.token) headers.authorization = `Bearer ${this.token}`;96      let payload: BodyInit | undefined;97      if (raw) {98        headers["content-type"] = raw.contentType;99        payload = new Blob([new Uint8Array(raw.buffer)]);100      } else if (body !== undefined) {101        headers["content-type"] = "application/json";102        payload = JSON.stringify(body);103      }104      const res = await this.fetchImpl(`${this.base}${path}`, {105        method,106        headers,107        body: payload,108        signal: ctrl.signal,109      });110      const text = await res.text();111      let json: unknown = null;112      try {113        json = text ? JSON.parse(text) : null;114      } catch {115        json = null;116      }117      if (!res.ok) {118        const err = (json as { error?: { message?: string; code?: string } } | null)?.error;119        throw new ControlPlaneError(err?.message ?? `HTTP ${res.status} on ${path}`, res.status, err?.code);120      }121      return schema ? schema.parse(json) : (json as T);122    } finally {123      clearTimeout(t);124    }125  }126127  enroll(enrollmentSecret: string, name: string, node: NodeMetadata, lanIp?: string) {128    return this.call(129      "POST",130      "/api/internal/workers/enroll",131      { enrollmentSecret, name, node, lanIp },132      EnrollResponse,133    );134  }135  heartbeat(body: z.input<typeof HeartbeatRequest>) {136    return this.call("POST", "/api/internal/workers/heartbeat", body, HeartbeatResponse);137  }138  lease(profiles: string[]) {139    return this.call("POST", "/api/internal/workers/lease", { profiles }, LeaseResponse);140  }141  renew(runId: string, fencingToken: number) {142    return this.call(143      "POST",144      `/api/internal/workers/runs/${runId}/renew`,145      { fencingToken },146      LeaseRenewResponse,147    );148  }149  events(runId: string, fencingToken: number, events: RunEvent[]) {150    return this.call<{ ok: true; lastSeq: number }>("POST", `/api/internal/workers/runs/${runId}/events`, {151      fencingToken,152      events,153    });154  }155  stage(runId: string, report: z.input<typeof StageReport>) {156    return this.call<{ ok: true }>(157      "POST",158      `/api/internal/workers/runs/${runId}/stage`,159      StageReport.parse(report),160    );161  }162  async uploadArtifact(163    runId: string,164    fencingToken: number,165    descriptor: Omit<z.input<typeof ArtifactDescriptor>, "sha256" | "bytes">,166    buffer: Buffer,167  ) {168    const sha256 = createHash("sha256").update(buffer).digest("hex");169    const artifact = ArtifactDescriptor.parse({ ...descriptor, sha256, bytes: buffer.length });170    const q = new URLSearchParams({171      fencingToken: String(fencingToken),172      path: artifact.path,173      kind: artifact.kind,174      sha256,175      stage: artifact.stage,176      contentType: artifact.contentType,177      synthetic: artifact.synthetic ? "1" : "0",178    });179    return this.call<{ ok: true; artifactId: string; sha256: string; bytes: number }>(180      "PUT",181      `/api/internal/workers/runs/${runId}/artifacts?${q.toString()}`,182      undefined,183      undefined,184      { buffer, contentType: artifact.contentType },185    );186  }187  listArtifacts(runId: string, fencingToken: number) {188    return this.call<{189      artifacts: Array<{190        path: string;191        version: number;192        sha256: string;193        bytes: number;194        contentType: string;195        kind: string;196      }>;197    }>("GET", `/api/internal/workers/runs/${runId}/artifacts?fencingToken=${fencingToken}`);198  }199  async downloadArtifact(runId: string, fencingToken: number, path: string): Promise<Buffer> {200    const ctrl = new AbortController();201    const t = setTimeout(() => ctrl.abort(), 120_000);202    try {203      const headers: Record<string, string> = {204        "user-agent": this.userAgent,205        "x-protocol-version": String(WORKER_PROTOCOL_VERSION),206      };207      if (this.token) headers.authorization = `Bearer ${this.token}`;208      const res = await this.fetchImpl(209        `${this.base}/api/internal/workers/runs/${runId}/artifacts/${path.split("/").map(encodeURIComponent).join("/")}?fencingToken=${fencingToken}`,210        { headers, signal: ctrl.signal },211      );212      if (!res.ok) throw new ControlPlaneError(`HTTP ${res.status} downloading ${path}`, res.status);213      const buf = Buffer.from(await res.arrayBuffer());214      const sha = res.headers.get("x-artifact-sha256");215      if (sha && createHash("sha256").update(buf).digest("hex") !== sha)216        throw new ControlPlaneError(`digest mismatch downloading ${path}`, 500);217      return buf;218    } finally {219      clearTimeout(t);220    }221  }222  complete(runId: string, body: z.input<typeof CompleteAttemptRequest>) {223    return this.call<{ ok: true }>(224      "POST",225      `/api/internal/workers/runs/${runId}/complete`,226      CompleteAttemptRequest.parse(body),227    );228  }229}230231export { WORKER_PROTOCOL_VERSION };232