@deepnoodle/mobius 0.0.21 → 0.0.22

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/dist/worker.js CHANGED
@@ -1,37 +1,10 @@
1
1
  import { randomUUID } from "node:crypto";
2
2
  import { hostname as systemHostname } from "node:os";
3
- import { AuthRevokedError, LeaseLostError, PayloadTooLargeError, RateLimitedError, WorkerInstanceConflictError, } from "./client.js";
4
3
  const CLOUD_RUN_METADATA_TIMEOUT_MS = 1000;
5
- // Generated once per process and reused by both the system_hostname
6
- // rung (as an 8-char suffix) and the generated_uuid rung (full value).
7
- // worker_instance_id is process identity and must be stable across
8
- // calls within the same boot — without the cache, a caller that
9
- // resolved twice would observe two different IDs.
10
4
  const BOOT_INSTANCE_ID = randomUUID();
11
- /**
12
- * Resolve a per-process `worker_instance_id` from the runtime
13
- * environment. Order: explicit → Cloud Run K_REVISION + metadata →
14
- * HOSTNAME env → FLY_MACHINE_ID → RAILWAY_REPLICA_ID →
15
- * RENDER_INSTANCE_ID → `os.hostname()` suffixed with a per-boot
16
- * random tag (laptops and bare metal) → generated UUID.
17
- *
18
- * The system_hostname rung carries a random suffix because
19
- * `os.hostname()` identifies the host, not the process — two
20
- * processes started on the same machine would otherwise auto-detect
21
- * the same identifier and collide on the server's conflict detector.
22
- * Set {@link WorkerConfig.workerInstanceId} explicitly only when a
23
- * stable identity across restarts is desired (named singleton
24
- * workers).
25
- *
26
- * The returned source is informational only — workers log it once at
27
- * startup so operators can confirm the right platform was picked up.
28
- */
29
5
  export async function resolveInstanceID(explicit) {
30
- if (explicit) {
31
- const trimmed = explicit.trim();
32
- if (trimmed)
33
- return { id: trimmed, source: "configured" };
34
- }
6
+ if (explicit?.trim())
7
+ return { id: explicit.trim(), source: "configured" };
35
8
  const cloudRun = await cloudRunInstanceID();
36
9
  if (cloudRun)
37
10
  return { id: cloudRun, source: "cloud_run_revision_instance" };
@@ -50,21 +23,18 @@ export async function resolveInstanceID(explicit) {
50
23
  try {
51
24
  const host = systemHostname().trim();
52
25
  if (host) {
53
- const suffix = BOOT_INSTANCE_ID.replace(/-/g, "").slice(0, 8);
54
- return { id: `${host}-${suffix}`, source: "system_hostname" };
26
+ return {
27
+ id: `${host}-${BOOT_INSTANCE_ID.replace(/-/g, "").slice(0, 8)}`,
28
+ source: "system_hostname",
29
+ };
55
30
  }
56
31
  }
57
32
  catch {
58
- // fall through to UUID
33
+ // fall through
59
34
  }
60
35
  return { id: BOOT_INSTANCE_ID, source: "generated_uuid" };
61
36
  }
62
37
  async function cloudRunInstanceID() {
63
- // Returns null on any metadata-server failure rather than the bare
64
- // revision — every replica sharing the same revision string would
65
- // otherwise collapse onto one row and trip the conflict detector.
66
- // The next strategy (HOSTNAME, which Cloud Run sets per-instance)
67
- // takes over.
68
38
  const revision = (process.env.K_REVISION ?? "").trim();
69
39
  if (!revision)
70
40
  return null;
@@ -86,495 +56,375 @@ async function cloudRunInstanceID() {
86
56
  return null;
87
57
  }
88
58
  }
89
- const silentLogger = {
90
- info: () => { },
91
- warn: () => { },
92
- error: () => { },
93
- };
59
+ const silentLogger = { info: () => { }, warn: () => { }, error: () => { } };
94
60
  const defaultLogger = {
95
61
  info: (...args) => console.log(...args),
96
62
  warn: (...args) => console.warn(...args),
97
63
  error: (...args) => console.error(...args),
98
64
  };
99
- /**
100
- * Worker claims jobs from Mobius and dispatches each to the corresponding
101
- * registered action function. A *job* is a single action invocation on
102
- * behalf of a workflow run; the backend owns the workflow engine.
103
- *
104
- * With `concurrency > 1` the worker holds up to N jobs in flight
105
- * simultaneously, all reporting the same `worker_instance_id` and
106
- * session token; the admin UI shows one row with a saturation bar
107
- * rather than N independent rows.
108
- */
109
65
  export class Worker {
110
- constructor(client, config, actions) {
66
+ constructor(client, config = {}, actions) {
111
67
  this.client = client;
112
- this.instanceIDResolved = false;
113
- this.abortController = new AbortController();
114
- this.authRevoked = false;
115
- const logger = config.logger === null ? silentLogger : (config.logger ?? defaultLogger);
116
- this.config = {
117
- workerInstanceId: config.workerInstanceId ?? "",
118
- concurrency: config.concurrency && config.concurrency > 0 ? config.concurrency : 1,
119
- name: config.name ?? "",
120
- version: config.version ?? "",
121
- queues: config.queues ?? [],
122
- actions: config.actions ?? [],
123
- pollWaitSeconds: config.pollWaitSeconds ?? 20,
124
- heartbeatIntervalMs: config.heartbeatIntervalMs,
125
- eventQueueSize: config.eventQueueSize ?? 256,
126
- eventBatchSize: config.eventBatchSize ?? 20,
127
- };
128
- this.logger = logger;
129
- this.actions = actions ?? new Map();
130
- this.sessionToken = randomUUID();
68
+ this.config = config;
69
+ this.actions = new Map();
70
+ this.generators = new Map();
71
+ this.sessionToken = "";
72
+ this.stopping = false;
73
+ this.stopController = new AbortController();
74
+ this.logger =
75
+ config.logger === null ? silentLogger : (config.logger ?? defaultLogger);
76
+ if (actions) {
77
+ for (const [name, fn] of actions)
78
+ this.actions.set(name, fn);
79
+ }
131
80
  }
132
- /** Register an action function under the given name. */
133
81
  register(name, fn) {
134
82
  this.actions.set(name, fn);
135
83
  return this;
136
84
  }
137
- /**
138
- * Start the claim loop. Returns a promise that resolves when the worker is
139
- * stopped via {@link stop} or the given signal is aborted. Throws
140
- * {@link AuthRevokedError} when the credential is revoked mid-flight
141
- * and {@link WorkerInstanceConflictError} when another live process
142
- * has already registered this worker_instance_id.
143
- *
144
- * Cancellation has two distinct shapes:
145
- *
146
- * - {@link stop} aborts the claim loop only — in-flight jobs continue
147
- * to run to natural completion and the returned promise resolves
148
- * once they all settle. This is graceful drain.
149
- * - The caller's `signal`, when aborted, propagates into both the
150
- * claim loop and every running action. This is the emergency-stop
151
- * path (e.g. SIGTERM with a hard deadline).
152
- */
85
+ registerGenerator(provider, model, fn) {
86
+ this.generators.set(`${provider}/${model}`, fn);
87
+ return this;
88
+ }
89
+ stop() {
90
+ this.stopping = true;
91
+ this.stopController.abort();
92
+ }
153
93
  async run(signal) {
154
- this.abortController = new AbortController();
155
- // The claim loop watches both signals so stop() and the caller's
156
- // emergency abort both halt new claims; executeJob only watches the
157
- // caller's signal so stop() doesn't kill jobs that have already
158
- // been claimed.
159
- const callerSignal = signal;
160
- const combined = callerSignal
161
- ? anySignal(callerSignal, this.abortController.signal)
162
- : this.abortController.signal;
163
- if (!this.instanceIDResolved) {
164
- const { id, source } = await resolveInstanceID(this.config.workerInstanceId || undefined);
165
- this.config.workerInstanceId = id;
166
- this.instanceIDResolved = true;
167
- this.logger.info(`[mobius] worker instance id ${id} (source: ${source})`);
168
- }
169
- this.logger.info(`[mobius] worker ${this.config.workerInstanceId} started (concurrency=${this.config.concurrency})`);
170
- const inflight = new Set();
171
- while (!combined.aborted) {
172
- if (this.authRevoked)
173
- break;
174
- while (inflight.size >= this.config.concurrency && !combined.aborted) {
175
- const aborted = abortAsPromise(combined);
176
- try {
177
- await Promise.race([...inflight, aborted.promise]);
178
- }
179
- finally {
180
- aborted.cancel();
181
- }
182
- }
183
- if (combined.aborted)
184
- break;
185
- let job;
94
+ const { id, source } = await resolveInstanceID(this.config.workerInstanceId);
95
+ this.config.workerInstanceId = id;
96
+ this.logger.info(`[mobius] worker instance id ${id} (source: ${source})`);
97
+ const runSignal = mergeSignals(this.stopController.signal, signal);
98
+ while (!this.stopping && !runSignal.aborted) {
186
99
  try {
187
- const claimReq = {
188
- worker_instance_id: this.config.workerInstanceId,
189
- worker_session_token: this.sessionToken,
190
- concurrency_limit: this.config.concurrency,
191
- wait_seconds: this.config.pollWaitSeconds,
192
- };
193
- if (this.config.name)
194
- claimReq.worker_name = this.config.name;
195
- if (this.config.version)
196
- claimReq.worker_version = this.config.version;
197
- if (this.config.queues.length > 0)
198
- claimReq.queues = this.config.queues;
199
- if (this.config.actions.length > 0)
200
- claimReq.actions = this.config.actions;
201
- job = await this.client.claimJob(claimReq, combined);
100
+ await this.runSocket(runSignal);
202
101
  }
203
102
  catch (err) {
204
- if (combined.aborted)
103
+ if (runSignal.aborted || this.stopping)
104
+ return;
105
+ this.logger.warn("[mobius] worker socket disconnected; reconnecting", err);
106
+ await sleep(this.config.reconnectDelayMs ?? 2000, runSignal);
107
+ }
108
+ }
109
+ }
110
+ async runSocket(signal) {
111
+ const ws = new WebSocket(this.client.workerSocketURL());
112
+ const concurrency = this.config.concurrency && this.config.concurrency > 0
113
+ ? this.config.concurrency
114
+ : 1;
115
+ const running = new Map();
116
+ let claimOutstanding = false;
117
+ await waitForOpen(ws, signal);
118
+ ws.send(JSON.stringify(this.registerFrame(concurrency)));
119
+ for await (const frame of socketFrames(ws, signal)) {
120
+ switch (frame.type) {
121
+ case "worker.registered":
122
+ this.sessionToken = String(frame.worker_session_token ?? "");
123
+ this.claim(ws, concurrency, running.size, claimOutstanding);
124
+ claimOutstanding = true;
125
+ break;
126
+ case "jobs.claimed":
127
+ claimOutstanding = false;
128
+ for (const job of (frame.jobs ?? [])) {
129
+ if (running.size >= concurrency)
130
+ break;
131
+ const ctrl = new AbortController();
132
+ running.set(job.id, ctrl);
133
+ void this.executeJob(ws, job, ctrl.signal).finally(() => {
134
+ running.delete(job.id);
135
+ this.claim(ws, concurrency, running.size, claimOutstanding);
136
+ claimOutstanding = true;
137
+ });
138
+ }
139
+ break;
140
+ case "work.available":
141
+ this.claim(ws, concurrency, running.size, claimOutstanding);
142
+ claimOutstanding = true;
143
+ break;
144
+ case "job.cancel": {
145
+ const ctrl = running.get(String(frame.job_id));
146
+ ctrl?.abort();
205
147
  break;
206
- if (err instanceof AuthRevokedError) {
207
- this.logger.error("[mobius] claim rejected: credential revoked");
208
- throw err;
209
148
  }
210
- if (err instanceof WorkerInstanceConflictError) {
211
- this.logger.error(`[mobius] claim rejected: worker_instance_id ${this.config.workerInstanceId} is already in use by another live process`);
212
- throw err;
149
+ case "job.heartbeat.ack": {
150
+ const cancel = frame.cancel;
151
+ if (cancel?.job_id)
152
+ running.get(cancel.job_id)?.abort();
153
+ break;
213
154
  }
214
- this.logger.error("[mobius] claim error:", err);
215
- await sleep(2000, combined);
216
- continue;
155
+ case "worker.drain":
156
+ this.stopping = true;
157
+ ws.send(JSON.stringify({ type: "worker.draining", message_id: msgID() }));
158
+ break;
159
+ case "keepalive":
160
+ ws.send(JSON.stringify({ type: "keepalive", message_id: msgID() }));
161
+ break;
162
+ case "error":
163
+ this.logger.error("[mobius] worker socket protocol error", frame.error);
164
+ break;
217
165
  }
218
- if (job == null) {
219
- await sleep(0, combined);
220
- continue;
166
+ if (this.stopping && running.size === 0) {
167
+ ws.close();
168
+ return;
221
169
  }
222
- const task = this.executeJob(job, callerSignal).finally(() => {
223
- inflight.delete(task);
224
- });
225
- inflight.add(task);
226
- }
227
- await Promise.allSettled(inflight);
228
- this.logger.info(`[mobius] worker ${this.config.workerInstanceId} stopped`);
229
- if (this.authRevoked) {
230
- throw new AuthRevokedError();
231
170
  }
232
171
  }
233
- /** Gracefully stop the worker after in-flight jobs complete. */
234
- stop() {
235
- this.abortController.abort();
172
+ registerFrame(concurrency) {
173
+ const actions = this.config.actions?.length
174
+ ? this.config.actions
175
+ : [...this.actions.keys()].sort();
176
+ return removeUndefined({
177
+ type: "worker.register",
178
+ message_id: msgID(),
179
+ worker_instance_id: this.config.workerInstanceId ?? "",
180
+ worker_session_token: this.sessionToken || undefined,
181
+ concurrency_limit: concurrency,
182
+ available_slots: concurrency,
183
+ name: this.config.name,
184
+ version: this.config.version,
185
+ queues: this.config.queues?.length ? this.config.queues : undefined,
186
+ action_names: actions.length ? actions : undefined,
187
+ models: this.modelCapabilities(),
188
+ });
236
189
  }
237
- // ---------------------------------------------------------------------------
238
- async executeJob(job, signal) {
239
- const jobId = job.job_id;
240
- const workerInstanceId = this.config.workerInstanceId;
241
- const leaseToken = job.lease_token;
242
- const attempt = job.attempt_number ?? 1;
243
- if (job.spec.kind !== "action") {
244
- const msg = `unsupported job spec kind ${JSON.stringify(job.spec.kind)} (this SDK only handles action jobs)`;
245
- this.logger.error(`[mobius] ${msg}`);
246
- await this.failJob(jobId, workerInstanceId, leaseToken, "UnsupportedSpec", msg);
190
+ claim(ws, concurrency, running, outstanding) {
191
+ if (outstanding || this.stopping)
247
192
  return;
248
- }
249
- const actionSpec = job.spec;
250
- const actionName = actionSpec.name;
251
- const parameters = actionSpec.parameters ?? {};
252
- this.logger.info(`[mobius] job ${jobId} claimed (workflow=${job.workflow_name}, step=${job.step_name}, action=${actionName}, attempt=${attempt})`);
253
- const fn = this.actions.get(actionName);
254
- if (!fn) {
255
- const msg = `action ${JSON.stringify(actionName)} not registered on this worker`;
256
- this.logger.error(`[mobius] ${msg}`);
257
- await this.failJob(jobId, workerInstanceId, leaseToken, "ActionNotRegistered", msg);
193
+ const available = concurrency - running;
194
+ if (available <= 0)
258
195
  return;
259
- }
260
- const actionController = new AbortController();
261
- const onAbort = () => actionController.abort();
262
- // Only the caller's signal (passed to run()) propagates into actions.
263
- // stop() does not flow through here — that's how graceful drain works.
264
- signal?.addEventListener("abort", onAbort, { once: true });
265
- const eventer = new JobEventer(this.client, job, {
266
- workerInstanceId,
267
- leaseToken,
268
- queueSize: this.config.eventQueueSize,
269
- batchSize: this.config.eventBatchSize,
270
- logger: this.logger,
196
+ const actions = this.config.actions?.length
197
+ ? this.config.actions
198
+ : [...this.actions.keys()].sort();
199
+ const frame = removeUndefined({
200
+ type: "jobs.claim",
201
+ message_id: msgID(),
202
+ available_slots: available,
203
+ queues: this.config.queues?.length ? this.config.queues : undefined,
204
+ action_names: actions.length ? actions : undefined,
205
+ models: this.modelCapabilities(),
271
206
  });
272
- const eventerPromise = eventer.run();
273
- const actionContext = {
274
- jobId,
275
- runId: job.run_id,
276
- projectId: this.client.project,
277
- workerInstanceId,
278
- attempt,
279
- queue: job.queue,
280
- workflowName: job.workflow_name,
281
- stepName: job.step_name,
282
- action: actionName,
283
- emitEvent: (type, payload) => eventer.emit(type, payload),
284
- };
285
- const interval = this.heartbeatInterval(job);
286
- let hbLost = false;
287
- const heartbeatTimer = setInterval(async () => {
288
- try {
289
- const hb = await this.client.heartbeatJob(jobId, {
290
- worker_instance_id: workerInstanceId,
291
- lease_token: leaseToken,
292
- });
293
- if (hb.directives.should_cancel) {
294
- this.logger.warn(`[mobius] job ${jobId}: cancel directive received`);
295
- actionController.abort();
296
- }
207
+ ws.send(JSON.stringify(frame));
208
+ }
209
+ modelCapabilities() {
210
+ const models = this.config.models
211
+ ?.filter((m) => m.provider && m.model)
212
+ .map((m) => ({ provider: m.provider, model: m.model }));
213
+ return models && models.length > 0 ? models : undefined;
214
+ }
215
+ async executeJob(ws, job, signal) {
216
+ const heartbeat = setInterval(() => {
217
+ const frame = {
218
+ type: "job.heartbeat",
219
+ message_id: msgID(),
220
+ job_id: job.id,
221
+ lease_token: job.lease_token,
222
+ };
223
+ ws.send(JSON.stringify(frame));
224
+ }, this.config.heartbeatIntervalMs ?? job.heartbeat_cadence_seconds * 1000);
225
+ const ctx = this.actionContext(job);
226
+ try {
227
+ let result;
228
+ if (job.kind === "action_execution") {
229
+ const actionName = job.action_name ?? String(job.spec.action_name ?? "");
230
+ const fn = this.actions.get(actionName);
231
+ if (!fn)
232
+ throw new Error(`action ${JSON.stringify(actionName)} is not registered`);
233
+ result = await fn(parameters(job), signal, ctx);
297
234
  }
298
- catch (err) {
299
- if (err instanceof AuthRevokedError) {
300
- this.logger.warn(`[mobius] job ${jobId}: credential revoked; cancelling action`);
301
- this.authRevoked = true;
302
- actionController.abort();
303
- }
304
- else if (err instanceof LeaseLostError) {
305
- this.logger.warn(`[mobius] job ${jobId}: lease lost during heartbeat`);
306
- hbLost = true;
307
- actionController.abort();
308
- }
309
- else {
310
- this.logger.error(`[mobius] job ${jobId}: heartbeat error:`, err);
311
- }
235
+ else if (job.kind === "llm_generation") {
236
+ const fn = this.generator(job.provider, job.model);
237
+ if (!fn)
238
+ throw new Error(`generation ${job.provider}/${job.model} is not registered`);
239
+ let seq = 0;
240
+ result = await fn(ctx, {
241
+ jobId: job.id,
242
+ runId: job.run_id,
243
+ sessionId: job.session_id,
244
+ agentTurnId: job.agent_turn_id,
245
+ toolCallId: job.tool_call_id,
246
+ provider: job.provider,
247
+ model: job.model,
248
+ spec: job.spec,
249
+ }, (delta) => {
250
+ seq += 1;
251
+ const frame = {
252
+ type: "generation.delta",
253
+ message_id: msgID(),
254
+ job_id: job.id,
255
+ lease_token: job.lease_token,
256
+ sequence: seq,
257
+ delta,
258
+ };
259
+ ws.send(JSON.stringify(frame));
260
+ });
312
261
  }
313
- }, interval);
314
- try {
315
- const result = await fn(parameters, actionController.signal, actionContext);
316
- clearInterval(heartbeatTimer);
317
- signal?.removeEventListener("abort", onAbort);
318
- eventer.stop();
319
- await eventerPromise;
320
- if (hbLost)
321
- return;
322
- const complete = { kind: "complete" };
323
- if (result != null) {
324
- complete.result_b64 = Buffer.from(JSON.stringify(result)).toString("base64");
262
+ else {
263
+ throw new Error(`unsupported job kind ${job.kind}`);
325
264
  }
326
- const reportReq = {
327
- worker_instance_id: workerInstanceId,
328
- lease_token: leaseToken,
329
- outcomes: [complete],
330
- };
331
- await this.client.reportJob(jobId, reportReq);
332
- this.logger.info(`[mobius] job ${jobId} completed`);
265
+ this.report(ws, job, "completed", result);
333
266
  }
334
267
  catch (err) {
335
- clearInterval(heartbeatTimer);
336
- signal?.removeEventListener("abort", onAbort);
337
- eventer.stop();
338
- await eventerPromise;
339
- if (err instanceof AuthRevokedError) {
340
- this.logger.warn(`[mobius] job ${jobId}: credential revoked during complete; worker will exit`);
341
- this.authRevoked = true;
342
- return;
343
- }
344
- if (err instanceof LeaseLostError || hbLost) {
345
- this.logger.warn(`[mobius] job ${jobId}: lease lost — will be retried`);
346
- return;
347
- }
348
- this.logger.error(`[mobius] job ${jobId} failed:`, err);
349
- const errType = actionController.signal.aborted ? "Timeout" : "Error";
350
- await this.failJob(jobId, workerInstanceId, leaseToken, errType, String(err));
268
+ this.report(ws, job, "failed", undefined, signal.aborted ? "Cancelled" : "Error", String(err));
351
269
  }
352
- }
353
- heartbeatInterval(job) {
354
- if (this.config.heartbeatIntervalMs != null)
355
- return this.config.heartbeatIntervalMs;
356
- if (job.heartbeat_interval_seconds != null) {
357
- return job.heartbeat_interval_seconds * 1000;
270
+ finally {
271
+ clearInterval(heartbeat);
358
272
  }
359
- return 10000;
360
273
  }
361
- async failJob(jobId, workerInstanceId, leaseToken, errorType, msg) {
362
- try {
363
- const fail = {
364
- kind: "fail",
365
- error_type: errorType,
366
- error_message: msg,
367
- };
368
- await this.client.reportJob(jobId, {
369
- worker_instance_id: workerInstanceId,
370
- lease_token: leaseToken,
371
- outcomes: [fail],
372
- });
373
- }
374
- catch (err) {
375
- this.logger.error(`[mobius] failed to report failure for job ${jobId}:`, err);
376
- }
274
+ generator(provider, model) {
275
+ if (!provider || !model)
276
+ return undefined;
277
+ return this.generators.get(`${provider}/${model}`) ?? this.generators.get(`${provider}/*`);
278
+ }
279
+ actionContext(job) {
280
+ return {
281
+ jobId: job.id,
282
+ runId: job.run_id,
283
+ sessionId: job.session_id,
284
+ agentTurnId: job.agent_turn_id,
285
+ toolCallId: job.tool_call_id,
286
+ projectId: this.client.project,
287
+ workerInstanceId: this.config.workerInstanceId ?? "",
288
+ attempt: job.claim_attempt,
289
+ queue: job.queue,
290
+ stepId: job.step_id,
291
+ action: job.action_name,
292
+ emitEvent: () => {
293
+ this.logger.warn("[mobius] custom worker events are not supported by the WebSocket protocol yet");
294
+ },
295
+ };
296
+ }
297
+ report(ws, job, status, result, errorType, errorMessage) {
298
+ const frame = removeUndefined({
299
+ type: "job.report",
300
+ message_id: msgID(),
301
+ job_id: job.id,
302
+ lease_token: job.lease_token,
303
+ status,
304
+ result: status === "completed" ? resultMap(result) : undefined,
305
+ error_type: errorType,
306
+ error_message: errorMessage,
307
+ });
308
+ ws.send(JSON.stringify(frame));
377
309
  }
378
310
  }
379
- /**
380
- * Runs multiple worker instances in one process. Most callers do not
381
- * need a pool — to run several jobs from one process, set
382
- * {@link WorkerConfig.concurrency} on a single {@link Worker} and the
383
- * admin UI will show one row with a saturation bar. Reach for a pool
384
- * only when each child should surface as its own row on the workers
385
- * page (independent draining, in-flight isolation).
386
- */
387
311
  export class WorkerPool {
388
312
  constructor(client, config) {
389
313
  this.client = client;
314
+ this.config = config;
390
315
  this.actions = new Map();
391
- this.abortController = new AbortController();
392
- const prefix = config.workerInstanceIdPrefix && config.workerInstanceIdPrefix !== ""
393
- ? config.workerInstanceIdPrefix
394
- : `worker-${randomUUID()}`;
395
- this.config = {
396
- workerInstanceIdPrefix: prefix,
397
- concurrency: config.concurrency && config.concurrency > 0 ? config.concurrency : 1,
398
- name: config.name ?? "",
399
- version: config.version ?? "",
400
- queues: config.queues ?? [],
401
- actions: config.actions ?? [],
402
- count: config.count && config.count > 0 ? config.count : 1,
403
- pollWaitSeconds: config.pollWaitSeconds ?? 20,
404
- heartbeatIntervalMs: config.heartbeatIntervalMs,
405
- logger: config.logger,
406
- eventQueueSize: config.eventQueueSize ?? 256,
407
- eventBatchSize: config.eventBatchSize ?? 20,
408
- };
316
+ this.generators = new Map();
409
317
  }
410
- /** Register an action function under the given name for every pool worker. */
411
318
  register(name, fn) {
412
319
  this.actions.set(name, fn);
413
320
  return this;
414
321
  }
415
- /**
416
- * Start all workers in the pool. Rejects with AuthRevokedError if any child
417
- * worker sees credential revocation.
418
- */
322
+ registerGenerator(provider, model, fn) {
323
+ this.generators.set(`${provider}/${model}`, fn);
324
+ return this;
325
+ }
419
326
  async run(signal) {
420
- this.abortController = new AbortController();
421
- const combined = signal
422
- ? anySignal(signal, this.abortController.signal)
423
- : this.abortController.signal;
424
- let authRevoked = false;
425
- let firstError;
426
- const workers = Array.from({ length: this.config.count }, (_, i) => {
427
- return new Worker(this.client, {
428
- workerInstanceId: `${this.config.workerInstanceIdPrefix}-${i + 1}`,
429
- concurrency: this.config.concurrency,
430
- name: this.config.name,
431
- version: this.config.version,
432
- queues: this.config.queues,
433
- actions: this.config.actions,
434
- pollWaitSeconds: this.config.pollWaitSeconds,
435
- heartbeatIntervalMs: this.config.heartbeatIntervalMs,
436
- logger: this.config.logger,
437
- eventQueueSize: this.config.eventQueueSize,
438
- eventBatchSize: this.config.eventBatchSize,
439
- }, this.actions);
440
- });
441
- await Promise.allSettled(workers.map(async (worker) => {
442
- try {
443
- await worker.run(combined);
327
+ const count = this.config.count && this.config.count > 0 ? this.config.count : 1;
328
+ const prefix = this.config.workerInstanceIdPrefix ?? `worker-${randomUUID()}`;
329
+ const workers = Array.from({ length: count }, (_, i) => {
330
+ const worker = new Worker(this.client, { ...this.config, workerInstanceId: `${prefix}-${i + 1}` }, this.actions);
331
+ for (const [key, fn] of this.generators) {
332
+ const [provider, model] = key.split("/", 2);
333
+ worker.registerGenerator(provider, model, fn);
444
334
  }
445
- catch (err) {
446
- if (err instanceof AuthRevokedError) {
447
- authRevoked = true;
448
- this.abortController.abort();
449
- }
450
- else if (!combined.aborted) {
451
- firstError ?? (firstError = err);
452
- this.abortController.abort();
453
- }
454
- }
455
- }));
456
- if (authRevoked) {
457
- throw new AuthRevokedError();
458
- }
459
- if (firstError != null) {
460
- throw firstError;
461
- }
335
+ return worker.run(signal);
336
+ });
337
+ await Promise.all(workers);
462
338
  }
463
- /** Gracefully stop every worker in the pool after in-flight jobs complete. */
464
- stop() {
465
- this.abortController.abort();
339
+ }
340
+ function parameters(job) {
341
+ const raw = job.spec.parameters;
342
+ return raw && typeof raw === "object" && !Array.isArray(raw)
343
+ ? raw
344
+ : {};
345
+ }
346
+ function resultMap(result) {
347
+ if (result && typeof result === "object" && !Array.isArray(result)) {
348
+ return result;
466
349
  }
350
+ return { output: result };
467
351
  }
468
- // sleep resolves after ms or when the signal aborts, whichever comes
469
- // first. Detaches the abort listener on the natural-resolution path
470
- // so it doesn't accumulate on a long-lived signal across thousands of
471
- // poll cycles.
472
- function sleep(ms, signal) {
473
- // AbortSignal does not retroactively dispatch abort to listeners
474
- // added after the signal already fired, so an already-aborted signal
475
- // would otherwise wait the full ms before resolving.
476
- if (signal.aborted)
477
- return Promise.resolve();
478
- return new Promise((resolve) => {
479
- const onAbort = () => {
480
- clearTimeout(t);
481
- resolve();
482
- };
483
- const t = setTimeout(() => {
484
- signal.removeEventListener("abort", onAbort);
485
- resolve();
486
- }, ms);
487
- signal.addEventListener("abort", onAbort, { once: true });
488
- });
352
+ function msgID() {
353
+ return `msg_${randomUUID()}`;
489
354
  }
490
- // abortAsPromise returns a promise that resolves when the signal
491
- // aborts, plus a cancel function the caller invokes when the promise
492
- // is no longer needed (e.g. when something else won a Promise.race).
493
- // Without the cancel handle, every Promise.race in the concurrency
494
- // loop would leak a listener on the long-lived worker signal.
495
- function abortAsPromise(signal) {
496
- if (signal.aborted)
497
- return { promise: Promise.resolve(), cancel: () => { } };
498
- let resolveFn;
499
- const promise = new Promise((resolve) => {
500
- resolveFn = resolve;
501
- });
502
- const onAbort = () => resolveFn();
503
- signal.addEventListener("abort", onAbort, { once: true });
504
- return {
505
- promise,
506
- cancel: () => signal.removeEventListener("abort", onAbort),
507
- };
355
+ function removeUndefined(obj) {
356
+ return Object.fromEntries(Object.entries(obj).filter(([, v]) => v !== undefined));
508
357
  }
509
- function anySignal(...signals) {
358
+ function mergeSignals(...signals) {
510
359
  const controller = new AbortController();
511
- for (const s of signals) {
512
- if (s.aborted) {
513
- controller.abort();
360
+ for (const signal of signals) {
361
+ if (!signal)
362
+ continue;
363
+ if (signal.aborted) {
364
+ controller.abort(signal.reason);
514
365
  break;
515
366
  }
516
- s.addEventListener("abort", () => controller.abort(), { once: true });
367
+ signal.addEventListener("abort", () => controller.abort(signal.reason), {
368
+ once: true,
369
+ });
517
370
  }
518
371
  return controller.signal;
519
372
  }
520
- class JobEventer {
521
- constructor(client, job, options) {
522
- this.client = client;
523
- this.job = job;
524
- this.options = options;
525
- this.queue = [];
526
- this.closed = false;
527
- this.wake = null;
528
- }
529
- emit(type, payload) {
530
- if (!type || this.closed)
373
+ function waitForOpen(ws, signal) {
374
+ return new Promise((resolve, reject) => {
375
+ if (ws.readyState === WebSocket.OPEN) {
376
+ resolve();
531
377
  return;
532
- if (this.queue.length >= this.options.queueSize) {
533
- this.queue.shift();
534
- this.options.logger.warn("[mobius] custom event queue full; dropping oldest event");
535
378
  }
536
- this.queue.push({ type, payload });
537
- this.wake?.();
538
- }
539
- stop() {
540
- this.closed = true;
541
- this.wake?.();
542
- }
543
- async run() {
544
- while (!this.closed || this.queue.length > 0) {
545
- const batch = this.queue.splice(0, this.options.batchSize);
546
- if (batch.length === 0) {
547
- await new Promise((resolve) => {
548
- this.wake = () => {
549
- this.wake = null;
550
- resolve();
551
- };
552
- });
553
- continue;
554
- }
555
- try {
556
- await this.client.emitJobEvents(this.job.job_id, {
557
- worker_instance_id: this.options.workerInstanceId,
558
- lease_token: this.options.leaseToken,
559
- events: batch,
560
- });
561
- }
562
- catch (err) {
563
- if (err instanceof LeaseLostError) {
564
- this.options.logger.warn(`[mobius] job ${this.job.job_id}: lease lost during custom event emit`);
565
- return;
566
- }
567
- if (err instanceof PayloadTooLargeError) {
568
- this.options.logger.warn(`[mobius] job ${this.job.job_id}: custom event payload too large`);
569
- continue;
570
- }
571
- if (err instanceof RateLimitedError) {
572
- this.options.logger.warn(`[mobius] job ${this.job.job_id}: custom event rate limited`);
573
- continue;
574
- }
575
- this.options.logger.error(`[mobius] job ${this.job.job_id}: custom event emit failed:`, err);
576
- }
379
+ const onOpen = () => resolve();
380
+ const onError = () => reject(new Error("worker socket failed to open"));
381
+ ws.addEventListener("open", onOpen, { once: true });
382
+ ws.addEventListener("error", onError, { once: true });
383
+ signal?.addEventListener("abort", () => reject(signal.reason), { once: true });
384
+ });
385
+ }
386
+ async function* socketFrames(ws, signal) {
387
+ const queue = [];
388
+ let wake;
389
+ let done = false;
390
+ ws.addEventListener("message", (event) => {
391
+ queue.push(JSON.parse(String(event.data)));
392
+ wake?.();
393
+ });
394
+ ws.addEventListener("close", () => {
395
+ done = true;
396
+ wake?.();
397
+ });
398
+ signal?.addEventListener("abort", () => {
399
+ done = true;
400
+ try {
401
+ ws.close();
577
402
  }
403
+ catch {
404
+ // ignore
405
+ }
406
+ wake?.();
407
+ });
408
+ while (!done || queue.length > 0) {
409
+ if (queue.length === 0) {
410
+ await new Promise((resolve) => {
411
+ wake = resolve;
412
+ });
413
+ wake = undefined;
414
+ continue;
415
+ }
416
+ yield queue.shift();
578
417
  }
579
418
  }
419
+ async function sleep(ms, signal) {
420
+ if (ms <= 0)
421
+ return;
422
+ await new Promise((resolve, reject) => {
423
+ const timer = setTimeout(resolve, ms);
424
+ signal?.addEventListener("abort", () => {
425
+ clearTimeout(timer);
426
+ reject(signal.reason ?? new Error("aborted"));
427
+ }, { once: true });
428
+ });
429
+ }
580
430
  //# sourceMappingURL=worker.js.map