spb/autoscholar
Public
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