@deepnoodle/mobius 0.0.20 → 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,486 +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 sessionToken = this.sessionToken;
242
- const attempt = job.attempt;
243
- this.logger.info(`[mobius] job ${jobId} claimed (workflow=${job.workflow_name}, step=${job.step_name}, action=${job.action}, attempt=${attempt})`);
244
- const fn = this.actions.get(job.action);
245
- if (!fn) {
246
- const msg = `action ${JSON.stringify(job.action)} not registered on this worker`;
247
- this.logger.error(`[mobius] ${msg}`);
248
- await this.failJob(jobId, workerInstanceId, sessionToken, attempt, "ActionNotRegistered", msg);
190
+ claim(ws, concurrency, running, outstanding) {
191
+ if (outstanding || this.stopping)
249
192
  return;
250
- }
251
- const actionController = new AbortController();
252
- const onAbort = () => actionController.abort();
253
- // Only the caller's signal (passed to run()) propagates into actions.
254
- // stop() does not flow through here — that's how graceful drain works.
255
- signal?.addEventListener("abort", onAbort, { once: true });
256
- const eventer = new JobEventer(this.client, job, {
257
- workerInstanceId,
258
- sessionToken,
259
- queueSize: this.config.eventQueueSize,
260
- batchSize: this.config.eventBatchSize,
261
- logger: this.logger,
193
+ const available = concurrency - running;
194
+ if (available <= 0)
195
+ return;
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(),
262
206
  });
263
- const eventerPromise = eventer.run();
264
- const actionContext = {
265
- jobId,
266
- runId: job.run_id,
267
- projectId: this.client.project,
268
- workerInstanceId,
269
- attempt,
270
- queue: job.queue,
271
- workflowName: job.workflow_name,
272
- stepName: job.step_name,
273
- action: job.action,
274
- emitEvent: (type, payload) => eventer.emit(type, payload),
275
- };
276
- const interval = this.heartbeatInterval(job);
277
- let hbLost = false;
278
- const heartbeatTimer = setInterval(async () => {
279
- try {
280
- const hb = await this.client.heartbeatJob(jobId, {
281
- worker_instance_id: workerInstanceId,
282
- worker_session_token: sessionToken,
283
- attempt,
284
- });
285
- if (hb.directives.should_cancel) {
286
- this.logger.warn(`[mobius] job ${jobId}: cancel directive received`);
287
- actionController.abort();
288
- }
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);
289
234
  }
290
- catch (err) {
291
- if (err instanceof AuthRevokedError) {
292
- this.logger.warn(`[mobius] job ${jobId}: credential revoked; cancelling action`);
293
- this.authRevoked = true;
294
- actionController.abort();
295
- }
296
- else if (err instanceof LeaseLostError) {
297
- this.logger.warn(`[mobius] job ${jobId}: lease lost during heartbeat`);
298
- hbLost = true;
299
- actionController.abort();
300
- }
301
- else {
302
- this.logger.error(`[mobius] job ${jobId}: heartbeat error:`, err);
303
- }
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
+ });
304
261
  }
305
- }, interval);
306
- try {
307
- const result = await fn(job.parameters, actionController.signal, actionContext);
308
- clearInterval(heartbeatTimer);
309
- signal?.removeEventListener("abort", onAbort);
310
- eventer.stop();
311
- await eventerPromise;
312
- if (hbLost)
313
- return;
314
- const completeReq = {
315
- worker_instance_id: workerInstanceId,
316
- worker_session_token: sessionToken,
317
- attempt,
318
- status: "completed",
319
- };
320
- if (result != null) {
321
- completeReq.result_b64 = Buffer.from(JSON.stringify(result)).toString("base64");
262
+ else {
263
+ throw new Error(`unsupported job kind ${job.kind}`);
322
264
  }
323
- await this.client.completeJob(jobId, completeReq);
324
- this.logger.info(`[mobius] job ${jobId} completed`);
265
+ this.report(ws, job, "completed", result);
325
266
  }
326
267
  catch (err) {
327
- clearInterval(heartbeatTimer);
328
- signal?.removeEventListener("abort", onAbort);
329
- eventer.stop();
330
- await eventerPromise;
331
- if (err instanceof AuthRevokedError) {
332
- this.logger.warn(`[mobius] job ${jobId}: credential revoked during complete; worker will exit`);
333
- this.authRevoked = true;
334
- return;
335
- }
336
- if (err instanceof LeaseLostError || hbLost) {
337
- this.logger.warn(`[mobius] job ${jobId}: lease lost — will be retried`);
338
- return;
339
- }
340
- this.logger.error(`[mobius] job ${jobId} failed:`, err);
341
- const errType = actionController.signal.aborted ? "Timeout" : "Error";
342
- await this.failJob(jobId, workerInstanceId, sessionToken, attempt, errType, String(err));
268
+ this.report(ws, job, "failed", undefined, signal.aborted ? "Cancelled" : "Error", String(err));
343
269
  }
344
- }
345
- heartbeatInterval(job) {
346
- if (this.config.heartbeatIntervalMs != null)
347
- return this.config.heartbeatIntervalMs;
348
- if (job.heartbeat_interval_seconds != null) {
349
- return job.heartbeat_interval_seconds * 1000;
270
+ finally {
271
+ clearInterval(heartbeat);
350
272
  }
351
- return 10000;
352
273
  }
353
- async failJob(jobId, workerInstanceId, sessionToken, attempt, errorType, msg) {
354
- try {
355
- await this.client.completeJob(jobId, {
356
- worker_instance_id: workerInstanceId,
357
- worker_session_token: sessionToken,
358
- attempt,
359
- status: "failed",
360
- error_type: errorType,
361
- error_message: msg,
362
- });
363
- }
364
- catch (err) {
365
- this.logger.error(`[mobius] failed to report failure for job ${jobId}:`, err);
366
- }
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));
367
309
  }
368
310
  }
369
- /**
370
- * Runs multiple worker instances in one process. Most callers do not
371
- * need a pool — to run several jobs from one process, set
372
- * {@link WorkerConfig.concurrency} on a single {@link Worker} and the
373
- * admin UI will show one row with a saturation bar. Reach for a pool
374
- * only when each child should surface as its own row on the workers
375
- * page (independent draining, in-flight isolation).
376
- */
377
311
  export class WorkerPool {
378
312
  constructor(client, config) {
379
313
  this.client = client;
314
+ this.config = config;
380
315
  this.actions = new Map();
381
- this.abortController = new AbortController();
382
- const prefix = config.workerInstanceIdPrefix && config.workerInstanceIdPrefix !== ""
383
- ? config.workerInstanceIdPrefix
384
- : `worker-${randomUUID()}`;
385
- this.config = {
386
- workerInstanceIdPrefix: prefix,
387
- concurrency: config.concurrency && config.concurrency > 0 ? config.concurrency : 1,
388
- name: config.name ?? "",
389
- version: config.version ?? "",
390
- queues: config.queues ?? [],
391
- actions: config.actions ?? [],
392
- count: config.count && config.count > 0 ? config.count : 1,
393
- pollWaitSeconds: config.pollWaitSeconds ?? 20,
394
- heartbeatIntervalMs: config.heartbeatIntervalMs,
395
- logger: config.logger,
396
- eventQueueSize: config.eventQueueSize ?? 256,
397
- eventBatchSize: config.eventBatchSize ?? 20,
398
- };
316
+ this.generators = new Map();
399
317
  }
400
- /** Register an action function under the given name for every pool worker. */
401
318
  register(name, fn) {
402
319
  this.actions.set(name, fn);
403
320
  return this;
404
321
  }
405
- /**
406
- * Start all workers in the pool. Rejects with AuthRevokedError if any child
407
- * worker sees credential revocation.
408
- */
322
+ registerGenerator(provider, model, fn) {
323
+ this.generators.set(`${provider}/${model}`, fn);
324
+ return this;
325
+ }
409
326
  async run(signal) {
410
- this.abortController = new AbortController();
411
- const combined = signal
412
- ? anySignal(signal, this.abortController.signal)
413
- : this.abortController.signal;
414
- let authRevoked = false;
415
- let firstError;
416
- const workers = Array.from({ length: this.config.count }, (_, i) => {
417
- return new Worker(this.client, {
418
- workerInstanceId: `${this.config.workerInstanceIdPrefix}-${i + 1}`,
419
- concurrency: this.config.concurrency,
420
- name: this.config.name,
421
- version: this.config.version,
422
- queues: this.config.queues,
423
- actions: this.config.actions,
424
- pollWaitSeconds: this.config.pollWaitSeconds,
425
- heartbeatIntervalMs: this.config.heartbeatIntervalMs,
426
- logger: this.config.logger,
427
- eventQueueSize: this.config.eventQueueSize,
428
- eventBatchSize: this.config.eventBatchSize,
429
- }, this.actions);
430
- });
431
- await Promise.allSettled(workers.map(async (worker) => {
432
- try {
433
- await worker.run(combined);
434
- }
435
- catch (err) {
436
- if (err instanceof AuthRevokedError) {
437
- authRevoked = true;
438
- this.abortController.abort();
439
- }
440
- else if (!combined.aborted) {
441
- firstError ?? (firstError = err);
442
- this.abortController.abort();
443
- }
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
- }));
446
- if (authRevoked) {
447
- throw new AuthRevokedError();
448
- }
449
- if (firstError != null) {
450
- throw firstError;
451
- }
335
+ return worker.run(signal);
336
+ });
337
+ await Promise.all(workers);
452
338
  }
453
- /** Gracefully stop every worker in the pool after in-flight jobs complete. */
454
- stop() {
455
- 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;
456
349
  }
350
+ return { output: result };
457
351
  }
458
- // sleep resolves after ms or when the signal aborts, whichever comes
459
- // first. Detaches the abort listener on the natural-resolution path
460
- // so it doesn't accumulate on a long-lived signal across thousands of
461
- // poll cycles.
462
- function sleep(ms, signal) {
463
- // AbortSignal does not retroactively dispatch abort to listeners
464
- // added after the signal already fired, so an already-aborted signal
465
- // would otherwise wait the full ms before resolving.
466
- if (signal.aborted)
467
- return Promise.resolve();
468
- return new Promise((resolve) => {
469
- const onAbort = () => {
470
- clearTimeout(t);
471
- resolve();
472
- };
473
- const t = setTimeout(() => {
474
- signal.removeEventListener("abort", onAbort);
475
- resolve();
476
- }, ms);
477
- signal.addEventListener("abort", onAbort, { once: true });
478
- });
352
+ function msgID() {
353
+ return `msg_${randomUUID()}`;
479
354
  }
480
- // abortAsPromise returns a promise that resolves when the signal
481
- // aborts, plus a cancel function the caller invokes when the promise
482
- // is no longer needed (e.g. when something else won a Promise.race).
483
- // Without the cancel handle, every Promise.race in the concurrency
484
- // loop would leak a listener on the long-lived worker signal.
485
- function abortAsPromise(signal) {
486
- if (signal.aborted)
487
- return { promise: Promise.resolve(), cancel: () => { } };
488
- let resolveFn;
489
- const promise = new Promise((resolve) => {
490
- resolveFn = resolve;
491
- });
492
- const onAbort = () => resolveFn();
493
- signal.addEventListener("abort", onAbort, { once: true });
494
- return {
495
- promise,
496
- cancel: () => signal.removeEventListener("abort", onAbort),
497
- };
355
+ function removeUndefined(obj) {
356
+ return Object.fromEntries(Object.entries(obj).filter(([, v]) => v !== undefined));
498
357
  }
499
- function anySignal(...signals) {
358
+ function mergeSignals(...signals) {
500
359
  const controller = new AbortController();
501
- for (const s of signals) {
502
- if (s.aborted) {
503
- controller.abort();
360
+ for (const signal of signals) {
361
+ if (!signal)
362
+ continue;
363
+ if (signal.aborted) {
364
+ controller.abort(signal.reason);
504
365
  break;
505
366
  }
506
- s.addEventListener("abort", () => controller.abort(), { once: true });
367
+ signal.addEventListener("abort", () => controller.abort(signal.reason), {
368
+ once: true,
369
+ });
507
370
  }
508
371
  return controller.signal;
509
372
  }
510
- class JobEventer {
511
- constructor(client, job, options) {
512
- this.client = client;
513
- this.job = job;
514
- this.options = options;
515
- this.queue = [];
516
- this.closed = false;
517
- this.wake = null;
518
- }
519
- emit(type, payload) {
520
- if (!type || this.closed)
373
+ function waitForOpen(ws, signal) {
374
+ return new Promise((resolve, reject) => {
375
+ if (ws.readyState === WebSocket.OPEN) {
376
+ resolve();
521
377
  return;
522
- if (this.queue.length >= this.options.queueSize) {
523
- this.queue.shift();
524
- this.options.logger.warn("[mobius] custom event queue full; dropping oldest event");
525
378
  }
526
- this.queue.push({ type, payload });
527
- this.wake?.();
528
- }
529
- stop() {
530
- this.closed = true;
531
- this.wake?.();
532
- }
533
- async run() {
534
- while (!this.closed || this.queue.length > 0) {
535
- const batch = this.queue.splice(0, this.options.batchSize);
536
- if (batch.length === 0) {
537
- await new Promise((resolve) => {
538
- this.wake = () => {
539
- this.wake = null;
540
- resolve();
541
- };
542
- });
543
- continue;
544
- }
545
- try {
546
- await this.client.emitJobEvents(this.job.job_id, {
547
- worker_instance_id: this.options.workerInstanceId,
548
- worker_session_token: this.options.sessionToken,
549
- attempt: this.job.attempt,
550
- events: batch,
551
- });
552
- }
553
- catch (err) {
554
- if (err instanceof LeaseLostError) {
555
- this.options.logger.warn(`[mobius] job ${this.job.job_id}: lease lost during custom event emit`);
556
- return;
557
- }
558
- if (err instanceof PayloadTooLargeError) {
559
- this.options.logger.warn(`[mobius] job ${this.job.job_id}: custom event payload too large`);
560
- continue;
561
- }
562
- if (err instanceof RateLimitedError) {
563
- this.options.logger.warn(`[mobius] job ${this.job.job_id}: custom event rate limited`);
564
- continue;
565
- }
566
- this.options.logger.error(`[mobius] job ${this.job.job_id}: custom event emit failed:`, err);
567
- }
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();
568
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();
569
417
  }
570
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
+ }
571
430
  //# sourceMappingURL=worker.js.map