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}3435/** Worker-side HTTP client for the versioned MacLustr worker protocol (outbound-only). */36export class WorkerClient {37  private token: string | undefined;38  private readonly base: string;39  private readonly fetchImpl: typeof fetch;40  private readonly timeoutMs: number;41  private readonly userAgent: string;4243  constructor(opts: WorkerClientOptions) {44    this.base = opts.baseUrl.replace(/\/+$/, "");45    this.token = opts.token;46    this.fetchImpl = opts.fetchImpl ?? fetch;47    this.timeoutMs = opts.timeoutMs ?? 30_000;48    this.userAgent = opts.userAgent ?? `autoscholar-worker/protocol-${WORKER_PROTOCOL_VERSION}`;49  }5051  setToken(token: string): void {52    this.token = token;53  }5455  private async call<T>(56    method: string,57    path: string,58    body?: unknown,59    schema?: z.ZodType<T>,60    raw?: { buffer: Buffer; contentType: string },61  ): Promise<T> {62    const ctrl = new AbortController();63    const t = setTimeout(() => ctrl.abort(), this.timeoutMs);64    try {65      const headers: Record<string, string> = {66        "user-agent": this.userAgent,67        "x-protocol-version": String(WORKER_PROTOCOL_VERSION),68      };69      if (this.token) headers.authorization = `Bearer ${this.token}`;70      let payload: BodyInit | undefined;71      if (raw) {72        headers["content-type"] = raw.contentType;73        payload = new Blob([new Uint8Array(raw.buffer)]);74      } else if (body !== undefined) {75        headers["content-type"] = "application/json";76        payload = JSON.stringify(body);77      }78      const res = await this.fetchImpl(`${this.base}${path}`, {79        method,80        headers,81        body: payload,82        signal: ctrl.signal,83      });84      const text = await res.text();85      let json: unknown = null;86      try {87        json = text ? JSON.parse(text) : null;88      } catch {89        json = null;90      }91      if (!res.ok) {92        const err = (json as { error?: { message?: string; code?: string } } | null)?.error;93        throw new ControlPlaneError(err?.message ?? `HTTP ${res.status} on ${path}`, res.status, err?.code);94      }95      return schema ? schema.parse(json) : (json as T);96    } finally {97      clearTimeout(t);98    }99  }100101  enroll(enrollmentSecret: string, name: string, node: NodeMetadata, lanIp?: string) {102    return this.call(103      "POST",104      "/api/internal/workers/enroll",105      { enrollmentSecret, name, node, lanIp },106      EnrollResponse,107    );108  }109  heartbeat(body: z.input<typeof HeartbeatRequest>) {110    return this.call("POST", "/api/internal/workers/heartbeat", body, HeartbeatResponse);111  }112  lease(profiles: string[]) {113    return this.call("POST", "/api/internal/workers/lease", { profiles }, LeaseResponse);114  }115  renew(runId: string, fencingToken: number) {116    return this.call(117      "POST",118      `/api/internal/workers/runs/${runId}/renew`,119      { fencingToken },120      LeaseRenewResponse,121    );122  }123  events(runId: string, fencingToken: number, events: RunEvent[]) {124    return this.call<{ ok: true; lastSeq: number }>("POST", `/api/internal/workers/runs/${runId}/events`, {125      fencingToken,126      events,127    });128  }129  stage(runId: string, report: z.input<typeof StageReport>) {130    return this.call<{ ok: true }>(131      "POST",132      `/api/internal/workers/runs/${runId}/stage`,133      StageReport.parse(report),134    );135  }136  async uploadArtifact(137    runId: string,138    fencingToken: number,139    descriptor: Omit<z.input<typeof ArtifactDescriptor>, "sha256" | "bytes">,140    buffer: Buffer,141  ) {142    const sha256 = createHash("sha256").update(buffer).digest("hex");143    const artifact = ArtifactDescriptor.parse({ ...descriptor, sha256, bytes: buffer.length });144    const q = new URLSearchParams({145      fencingToken: String(fencingToken),146      path: artifact.path,147      kind: artifact.kind,148      sha256,149      stage: artifact.stage,150      contentType: artifact.contentType,151      synthetic: artifact.synthetic ? "1" : "0",152    });153    return this.call<{ ok: true; artifactId: string; sha256: string; bytes: number }>(154      "PUT",155      `/api/internal/workers/runs/${runId}/artifacts?${q.toString()}`,156      undefined,157      undefined,158      { buffer, contentType: artifact.contentType },159    );160  }161  complete(runId: string, body: z.input<typeof CompleteAttemptRequest>) {162    return this.call<{ ok: true }>(163      "POST",164      `/api/internal/workers/runs/${runId}/complete`,165      CompleteAttemptRequest.parse(body),166    );167  }168}169170export { WORKER_PROTOCOL_VERSION };171