From 41e62a1dd70c2791f1df477acd1728ef1b38f297 Mon Sep 17 00:00:00 2001 From: Oleksii Date: Sat, 5 Sep 2026 02:09:16 -0300 Subject: [PATCH] fix: align worker SDK with lease protocol --- README.md | 38 ++++- src/__tests__/worker-protocol.test.ts | 111 +++++++++++++ src/__tests__/worker.integration.test.ts | 6 +- src/__tests__/worker.test.ts | 4 +- src/client.ts | 39 ++++- src/types.ts | 13 ++ src/worker.ts | 194 +++++++++-------------- 7 files changed, 270 insertions(+), 135 deletions(-) create mode 100644 src/__tests__/worker-protocol.test.ts diff --git a/README.md b/README.md index 4ce3f06..d6ad762 100644 --- a/README.md +++ b/README.md @@ -119,28 +119,52 @@ const client = new Orch8Client({ ## Worker -Run a polling worker that claims and executes tasks: +Run a polling worker using `x-api-key` and `x-tenant-id` authentication: ```typescript -import { Orch8Client, Orch8Worker } from "@orch8/sdk"; +import { Orch8Client, Orch8Worker } from "@orch8.io/sdk"; -const client = new Orch8Client({ baseUrl: "https://api.orch8.io", tenantId: "my-tenant" }); +const client = new Orch8Client({ + baseUrl: process.env.ORCH8_ENGINE_URL ?? "http://localhost:8080", + tenantId: process.env.ORCH8_TENANT_ID, + headers: { "x-api-key": process.env.ORCH8_API_KEY ?? "" }, +}); const worker = new Orch8Worker({ client, workerId: "worker-1", handlers: { - "send-email": async (task) => { - console.log(`Sending email to ${task.params.to}`); - return { sent: true }; + "inspect-document": async (task) => { + // Replace with your bounded task implementation. + return { inspected: true, input: task.params }; }, }, maxConcurrent: 10, }); -await worker.start(); // blocks until worker.stop() is called +await worker.start(); // Starts polling and returns immediately. +process.once("SIGTERM", () => { void worker.stop(); }); +process.once("SIGINT", () => { void worker.stop(); }); ``` +The worker echoes each task's `claim_epoch` on heartbeat, completion, and failure. +It respects the server's minimum poll delay and uses a heartbeat interval no +longer than the advertised interval or half the lease duration. Completion +callbacks run only after a successful acknowledgement; rejected or ambiguous +acknowledgements are left for lease recovery, without sending a contradictory +failure request. + +`client.pollTasks()` and `client.pollTasksFromQueue()` still return task arrays. +Use `client.pollTaskBatch()` or `client.pollTaskBatchFromQueue()` to access +`tasks`, `lease_secs`, `heartbeat_interval_secs`, and `poll_after_ms` when writing +your own loop. Legacy array responses are accepted at the client boundary, with +no timing hints. Epoch fields remain optional in the types for legacy servers; +always echo the epoch supplied by a current server in custom worker loops. + +A lease does not authorize offline execution. Handlers are not forcibly cancelled +on lease loss or timeout; use bounded work and provider idempotency keys for +external effects. `stop()` waits up to 30 seconds for executing handlers. + ## Error Handling ```typescript diff --git a/src/__tests__/worker-protocol.test.ts b/src/__tests__/worker-protocol.test.ts new file mode 100644 index 0000000..8fd289e --- /dev/null +++ b/src/__tests__/worker-protocol.test.ts @@ -0,0 +1,111 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { Orch8Client } from "../client.js"; +import { Orch8Worker } from "../worker.js"; + +const task = { id: "task-1", handler_name: "inspect", claim_epoch: 7, timeout_ms: null }; +const envelope = { tasks: [task], lease_secs: 6, heartbeat_interval_secs: 1, poll_after_ms: 5000 }; +const response = (body: unknown, status = 200) => new Response(JSON.stringify(body), { status }); + +describe("worker wire protocol", () => { + const fetchMock = vi.fn(); + beforeEach(() => { + vi.useFakeTimers(); + vi.stubGlobal("fetch", fetchMock); + fetchMock.mockReset(); + }); + afterEach(() => { + vi.unstubAllGlobals(); + vi.useRealTimers(); + }); + + it("preserves poll metadata and keeps array APIs compatible for both queues and handlers", async () => { + const client = new Orch8Client({ baseUrl: "http://engine" }); + fetchMock.mockImplementation(async () => response(envelope)); + expect(await client.pollTaskBatch({})).toEqual(envelope); + expect(await client.pollTaskBatchFromQueue({})).toEqual(envelope); + expect(await client.pollTasks({})).toEqual([task]); + expect(await client.pollTasksFromQueue({})).toEqual([task]); + fetchMock.mockImplementation(async () => response([task])); + expect(await client.pollTaskBatch({})).toEqual({ tasks: [task] }); + expect(await client.pollTasksFromQueue({})).toEqual([task]); + }); + + it("rejects malformed envelopes instead of treating them as an empty queue", async () => { + const client = new Orch8Client({ baseUrl: "http://engine" }); + for (const body of [{}, { tasks: null }, { tasks: [], lease_secs: 0 }, { tasks: [], poll_after_ms: -1 }]) { + fetchMock.mockImplementation(async () => response(body)); + await expect(client.pollTaskBatch({})).rejects.toThrow(TypeError); + } + }); + + it.each([false, true])("echoes epochs, honors hints, and rejects stale completion (client=%s)", async (useClient) => { + let finish!: (output: unknown) => void; + const handler = vi.fn(() => new Promise((resolve) => { finish = resolve; })); + const completed = vi.fn(); + const failed = vi.fn(); + const client = new Orch8Client({ baseUrl: "http://engine", tenantId: "tenant-1", headers: { "x-api-key": "test-key" } }); + const statuses: number[] = []; + const paths: string[] = []; + fetchMock.mockImplementation(async (url, init) => { + const path = String(url); + paths.push(path); + if (path.endsWith("/poll")) return response(envelope); + expect(JSON.parse(String(init?.body))).toMatchObject({ worker_id: "phone-1", claim_epoch: 7 }); + if (useClient) { + expect(init?.headers).toMatchObject({ "X-Tenant-Id": "tenant-1", "x-api-key": "test-key" }); + } + const status = path.endsWith("/complete") ? 409 : 200; + statuses.push(status); + return response({ checkpoint_seq: 0 }, status); + }); + const worker = new Orch8Worker({ + ...(useClient ? { client } : { engineUrl: "http://engine" }), + workerId: "phone-1", handlers: { inspect: handler }, pollIntervalMs: 100, + onTaskComplete: completed, onTaskFail: failed, + }); + try { + await worker.start(); + await vi.advanceTimersByTimeAsync(1000); + expect(handler).toHaveBeenCalledTimes(1); + expect(paths.filter((p) => p.endsWith("/heartbeat"))).toHaveLength(1); + expect(paths.filter((p) => p.endsWith("/poll"))).toHaveLength(1); + finish({ done: true }); + await vi.advanceTimersByTimeAsync(0); + expect(statuses).toEqual([200, 409]); + expect(completed).not.toHaveBeenCalled(); + expect(failed).not.toHaveBeenCalled(); + expect(paths.some((p) => p.endsWith("/fail"))).toBe(false); + expect(worker.stats().inFlight).toBe(0); + } finally { + finish?.({}); + const stopping = worker.stop(); + await vi.advanceTimersByTimeAsync(30000); + await stopping; + } + }); + + it("surfaces a server conflict to direct callers", async () => { + fetchMock.mockImplementation(async () => response({ error: "stale claim" }, 409)); + const client = new Orch8Client({ baseUrl: "http://engine" }); + await expect(client.completeTask(task.id, { worker_id: "phone-1", claim_epoch: 7 })) + .rejects.toMatchObject({ status: 409 }); + }); + + it("includes the epoch when a handler fails", async () => { + fetchMock.mockResolvedValueOnce(response(envelope)); + fetchMock.mockImplementation(async () => response({})); + const failed = vi.fn(); + const worker = new Orch8Worker({ engineUrl: "http://engine", workerId: "phone-1", + handlers: { inspect: async () => { throw new Error("invalid input"); } }, onTaskFail: failed }); + await worker.start(); + await vi.advanceTimersByTimeAsync(0); + const call = fetchMock.mock.calls.find(([url]) => String(url).endsWith("/fail")); + expect(JSON.parse(String(call?.[1]?.body))).toEqual({ + worker_id: "phone-1", claim_epoch: 7, message: "invalid input", retryable: false, + }); + expect(failed).toHaveBeenCalledTimes(1); + const stopping = worker.stop(); + await vi.advanceTimersByTimeAsync(30000); + await stopping; + }); +}); diff --git a/src/__tests__/worker.integration.test.ts b/src/__tests__/worker.integration.test.ts index 0293ff7..72b830e 100644 --- a/src/__tests__/worker.integration.test.ts +++ b/src/__tests__/worker.integration.test.ts @@ -77,8 +77,8 @@ describe("Orch8Worker integration", () => { const handler = vi.fn().mockResolvedValue({ done: true }); const pollSpy = vi - .spyOn(client, "pollTasks") - .mockResolvedValue([ + .spyOn(client, "pollTaskBatch") + .mockResolvedValue({ tasks: [ { id: "wt-client", instance_id: "inst-1", @@ -87,7 +87,7 @@ describe("Orch8Worker integration", () => { state: "claimed", created_at: "2025-01-01T00:00:00Z", } as any, - ]); + ] }); const completeSpy = vi .spyOn(client, "completeTask") .mockResolvedValue(undefined); diff --git a/src/__tests__/worker.test.ts b/src/__tests__/worker.test.ts index 2a37144..01efeb9 100644 --- a/src/__tests__/worker.test.ts +++ b/src/__tests__/worker.test.ts @@ -452,8 +452,8 @@ describe("Orch8Worker", () => { const client = new Orch8Client({ baseUrl: "http://localhost:8080" }); const pollSpy = vi - .spyOn(client, "pollTasks") - .mockResolvedValue([]); + .spyOn(client, "pollTaskBatch") + .mockResolvedValue({ tasks: [] }); const worker = new Orch8Worker({ client, diff --git a/src/client.ts b/src/client.ts index 1cc4cda..ca13757 100644 --- a/src/client.ts +++ b/src/client.ts @@ -19,6 +19,7 @@ import type { PluginDef, Session, WorkerTask, + WorkerPollResponse, ClusterNode, CircuitBreaker, AuditEntry, @@ -700,10 +701,16 @@ export class Orch8Client { // Workers // --------------------------------------------------------------------------- - pollTasks( + async pollTasks( body: PollRequest | Record, ): Promise { - return this.post("/workers/tasks/poll", body); + return (await this.pollTaskBatch(body)).tasks; + } + + async pollTaskBatch( + body: PollRequest | Record, + ): Promise { + return decodeWorkerPoll(await this.post("/workers/tasks/poll", body)); } completeTask( @@ -739,10 +746,16 @@ export class Orch8Client { return this.get>("/workers/tasks/stats"); } - pollTasksFromQueue( + async pollTasksFromQueue( body: QueuePollRequest | Record, ): Promise { - return this.post("/workers/tasks/poll/queue", body); + return (await this.pollTaskBatchFromQueue(body)).tasks; + } + + async pollTaskBatchFromQueue( + body: QueuePollRequest | Record, + ): Promise { + return decodeWorkerPoll(await this.post("/workers/tasks/poll/queue", body)); } // --------------------------------------------------------------------------- @@ -973,3 +986,21 @@ export class Orch8Client { return this.get("/health/ready"); } } + +/** Normalize the legacy wire format only at the HTTP boundary. */ +export function decodeWorkerPoll(value: unknown): WorkerPollResponse { + if (Array.isArray(value)) return { tasks: value }; + if (!value || typeof value !== "object" || !("tasks" in value) || !Array.isArray(value.tasks)) { + throw new TypeError("Worker poll response must contain a tasks array"); + } + const response: WorkerPollResponse = { tasks: value.tasks }; + for (const key of ["lease_secs", "heartbeat_interval_secs", "poll_after_ms"] as const) { + if (!(key in value)) continue; + const hint = (value as Record)[key]; + if (typeof hint !== "number" || !Number.isFinite(hint) || hint < 0 || (key !== "poll_after_ms" && hint === 0)) { + throw new TypeError(`Invalid worker poll hint: ${key}`); + } + response[key] = hint; + } + return response; +} diff --git a/src/types.ts b/src/types.ts index 0537d38..7040c0f 100644 --- a/src/types.ts +++ b/src/types.ts @@ -218,6 +218,8 @@ export interface WorkerTask { error_message: string | null; error_retryable: boolean | null; created_at: string; + /** Ownership epoch returned by the server; echo on every acknowledgement. */ + claim_epoch?: number; resume_checkpoint?: unknown; checkpoint_seq: number; } @@ -639,12 +641,14 @@ export interface QueuePollRequest { } export interface CompleteRequest { + claim_epoch?: number; worker_id?: string; output?: unknown; [key: string]: unknown; } export interface FailRequest { + claim_epoch?: number; worker_id?: string; message?: string; error?: string; @@ -653,6 +657,7 @@ export interface FailRequest { } export interface HeartbeatRequest { + claim_epoch?: number; worker_id?: string; checkpoint?: unknown; checkpoint_seq?: number; @@ -662,3 +667,11 @@ export interface HeartbeatRequest { export interface HeartbeatResponse { checkpoint_seq: number; } + +/** Poll metadata is absent only when connected to a legacy array-returning server. */ +export interface WorkerPollResponse { + tasks: WorkerTask[]; + lease_secs?: number; + heartbeat_interval_secs?: number; + poll_after_ms?: number; +} diff --git a/src/worker.ts b/src/worker.ts index b927f17..1647de7 100644 --- a/src/worker.ts +++ b/src/worker.ts @@ -1,5 +1,5 @@ import type { WorkerTask } from "./types.js"; -import type { Orch8Client } from "./client.js"; +import { Orch8Client } from "./client.js"; export type HandlerFn = (task: WorkerTask) => Promise; @@ -55,9 +55,11 @@ export class Orch8Worker { handlers: Record; onTaskComplete?: (task: WorkerTask, output: unknown) => void; onTaskFail?: (task: WorkerTask, error: string) => void; - client?: Orch8Client; }; + private readonly client: Orch8Client; + private pollHints = new Map(); + private heartbeatIntervals = new Map(); private running = false; private pollTimers = new Map(); private heartbeatTimer: NodeJS.Timeout | null = null; @@ -83,8 +85,8 @@ export class Orch8Worker { circuitBreakerCheck: config.circuitBreakerCheck ?? false, onTaskComplete: config.onTaskComplete, onTaskFail: config.onTaskFail, - client: config.client, }; + this.client = config.client ?? new Orch8Client({ baseUrl: this.config.engineUrl, retry: false }); this.concurrencySemaphore = this.config.maxConcurrent; } @@ -99,9 +101,7 @@ export class Orch8Worker { } // Start heartbeat loop. - this.heartbeatTimer = setInterval(() => { - void this.sendHeartbeats(); - }, this.config.heartbeatIntervalMs); + this.resetHeartbeatTimer(); } private schedulePoll(handlerName: string, delayMs: number): void { @@ -120,7 +120,7 @@ export class Orch8Worker { this.config.pollIntervalMs * Math.pow(2, failures), Orch8Worker.MAX_BACKOFF_MS, ) - : this.config.pollIntervalMs; + : Math.max(this.config.pollIntervalMs, this.pollHints.get(handlerName) ?? 0); this.schedulePoll(handlerName, nextDelay); }); }, delayMs); @@ -138,8 +138,13 @@ export class Orch8Worker { // Drain in-flight tasks with a hard timeout. const drainTimeoutMs = 30_000; const drain = Promise.allSettled(Array.from(this.executingPromises)); - const timeout = new Promise((resolve) => setTimeout(resolve, drainTimeoutMs)); - await Promise.race([drain, timeout]); + let timer: NodeJS.Timeout | undefined; + const timeout = new Promise((resolve) => { timer = setTimeout(resolve, drainTimeoutMs); }); + try { + await Promise.race([drain, timeout]); + } finally { + clearTimeout(timer); + } } stats(): WorkerRuntimeStats { @@ -157,19 +162,8 @@ export class Orch8Worker { // Circuit breaker check: skip polling if the handler's circuit is open. if (this.config.circuitBreakerCheck) { try { - if (this.config.client) { - const cb = await this.config.client.getCircuitBreaker(handlerName); - if (cb.state === "open") return; - } else { - const cbRes = await fetch( - `${this.config.engineUrl}/circuit-breakers/${handlerName}`, - { method: "GET", headers: { "Content-Type": "application/json" } }, - ); - if (cbRes.ok) { - const cb = (await cbRes.json()) as { state: string }; - if (cb.state === "open") return; - } - } + const cb = await this.client.getCircuitBreaker(handlerName); + if (cb.state === "open") return; } catch { // If the check fails, proceed with polling anyway. } @@ -178,33 +172,20 @@ export class Orch8Worker { try { const limit = Math.min(this.concurrencySemaphore, this.config.maxConcurrent); - let tasks: WorkerTask[]; - if (this.config.client) { - tasks = await this.config.client.pollTasks({ - handler_name: handlerName, - worker_id: this.config.workerId, - limit, - }); - } else { - const res = await fetch(`${this.config.engineUrl}/workers/tasks/poll`, { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ - handler_name: handlerName, - worker_id: this.config.workerId, - limit, - }), - }); - - if (!res.ok) { - this.consecutiveFailures.set( - handlerName, - (this.consecutiveFailures.get(handlerName) ?? 0) + 1, - ); - return; - } - tasks = (await res.json()) as WorkerTask[]; - } + const batch = await this.client.pollTaskBatch({ + handler_name: handlerName, + worker_id: this.config.workerId, + limit, + }); + const tasks = batch.tasks; + this.pollHints.set(handlerName, batch.poll_after_ms ?? 0); + const heartbeatMs = Math.min( + this.config.heartbeatIntervalMs, + (batch.heartbeat_interval_secs ?? Infinity) * 1000, + (batch.lease_secs ?? Infinity) * 500, + ); + this.heartbeatIntervals.set(handlerName, heartbeatMs); + if (this.running) this.resetHeartbeatTimer(); // Reset backoff on successful poll. this.consecutiveFailures.set(handlerName, 0); @@ -229,26 +210,29 @@ export class Orch8Worker { private async executeTask(task: WorkerTask): Promise { const handler = this.config.handlers[task.handler_name]; if (!handler) { - await this.failTask(task.id, `no handler registered for "${task.handler_name}"`, false); + await this.failTask(task, `no handler registered for "${task.handler_name}"`, false).catch(() => {}); this.inFlightTasks.delete(task.id); this.concurrencySemaphore++; return; } try { - const output = await this.withTimeout(handler(task), task.timeout_ms); - await this.completeTask(task.id, output); + let output: unknown; + try { + output = await this.withTimeout(handler(task), task.timeout_ms); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + const retryable = err instanceof Error && "retryable" in err + ? Boolean(err.retryable) : false; + await this.failTask(task, message, retryable); + this.notify(() => this.config.onTaskFail?.(task, message)); + return; + } + // A rejected or ambiguous acknowledgement is not a handler failure. + await this.completeTask(task, output); this.notify(() => this.config.onTaskComplete?.(task, output)); - } catch (err) { - const message = err instanceof Error ? err.message : String(err); - // If the error exposes a retryable flag, respect it. Otherwise default - // to non-retryable (conservative — matches engine FailRequest default). - const retryable = - err instanceof Error && "retryable" in err - ? Boolean((err as Error & { retryable: unknown }).retryable) - : false; - await this.failTask(task.id, message, retryable); - this.notify(() => this.config.onTaskFail?.(task, message)); + } catch { + // Leave unacknowledged work for lease recovery; never report success. } finally { this.inFlightTasks.delete(task.id); this.concurrencySemaphore++; @@ -265,78 +249,50 @@ export class Orch8Worker { private async withTimeout(promise: Promise, timeoutMs: number | null): Promise { if (!timeoutMs) return promise; - let timer: NodeJS.Timeout; + let timer: NodeJS.Timeout | undefined; const timeout = new Promise((_, reject) => { timer = setTimeout(() => reject(new Error("task timed out")), timeoutMs); }); try { return await Promise.race([promise, timeout]); } finally { - clearTimeout(timer!); + clearTimeout(timer); } } - private async completeTask(taskId: string, output: unknown): Promise { - try { - if (this.config.client) { - await this.config.client.completeTask(taskId, { - worker_id: this.config.workerId, - output: output ?? {}, - }); - } else { - await fetch(`${this.config.engineUrl}/workers/tasks/${taskId}/complete`, { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ - worker_id: this.config.workerId, - output: output ?? {}, - }), - }); - } - } catch { - // Will be reaped and retried. - } + private completeTask(task: WorkerTask, output: unknown): Promise { + return this.client.completeTask(task.id, { + worker_id: this.config.workerId, + claim_epoch: task.claim_epoch, + output: output ?? {}, + }); } - private async failTask(taskId: string, message: string, retryable: boolean): Promise { - try { - if (this.config.client) { - await this.config.client.failTask(taskId, { - worker_id: this.config.workerId, - message, - retryable, - }); - } else { - await fetch(`${this.config.engineUrl}/workers/tasks/${taskId}/fail`, { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ - worker_id: this.config.workerId, - message, - retryable, - }), - }); - } - } catch { - // Will be reaped and retried. - } + private failTask(task: WorkerTask, message: string, retryable: boolean): Promise { + return this.client.failTask(task.id, { + worker_id: this.config.workerId, + claim_epoch: task.claim_epoch, + message, + retryable, + }); + } + + private heartbeatDelayMs = 0; + + private resetHeartbeatTimer(): void { + const delay = Math.min(this.config.heartbeatIntervalMs, ...this.heartbeatIntervals.values()); + if (this.heartbeatTimer && delay === this.heartbeatDelayMs) return; + if (this.heartbeatTimer) clearInterval(this.heartbeatTimer); + this.heartbeatDelayMs = delay; + this.heartbeatTimer = setInterval(() => { void this.sendHeartbeats(); }, delay); } private async sendHeartbeats(): Promise { - const ids = Array.from(this.inFlightTasks.keys()); await Promise.allSettled( - ids.map((taskId) => { - if (this.config.client) { - return this.config.client.heartbeatTask(taskId, { - worker_id: this.config.workerId, - }); - } - return fetch(`${this.config.engineUrl}/workers/tasks/${taskId}/heartbeat`, { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ worker_id: this.config.workerId }), - }); - }), + Array.from(this.inFlightTasks.values(), (task) => this.client.heartbeatTask(task.id, { + worker_id: this.config.workerId, + claim_epoch: task.claim_epoch, + })), ); } }