@livekit/agents 0.1.0 → 0.2.0

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.
Files changed (133) hide show
  1. package/.turbo/turbo-build.log +1 -1
  2. package/CHANGELOG.md +21 -0
  3. package/LICENSE +201 -0
  4. package/dist/audio.d.ts +12 -0
  5. package/dist/audio.d.ts.map +1 -0
  6. package/dist/audio.js +37 -0
  7. package/dist/audio.js.map +1 -0
  8. package/dist/cli.d.ts +11 -0
  9. package/dist/cli.d.ts.map +1 -1
  10. package/dist/cli.js +68 -8
  11. package/dist/cli.js.map +1 -1
  12. package/dist/generator.d.ts +12 -6
  13. package/dist/generator.d.ts.map +1 -1
  14. package/dist/generator.js +9 -3
  15. package/dist/generator.js.map +1 -1
  16. package/dist/http_server.d.ts +1 -1
  17. package/dist/http_server.js +0 -3
  18. package/dist/http_server.js.map +1 -1
  19. package/dist/index.d.ts +12 -3
  20. package/dist/index.d.ts.map +1 -1
  21. package/dist/index.js +12 -3
  22. package/dist/index.js.map +1 -1
  23. package/dist/ipc/job_executor.d.ts +19 -0
  24. package/dist/ipc/job_executor.d.ts.map +1 -0
  25. package/dist/ipc/job_executor.js +8 -0
  26. package/dist/ipc/job_executor.js.map +1 -0
  27. package/dist/ipc/job_main.d.ts +7 -4
  28. package/dist/ipc/job_main.d.ts.map +1 -1
  29. package/dist/ipc/job_main.js +96 -61
  30. package/dist/ipc/job_main.js.map +1 -1
  31. package/dist/ipc/message.d.ts +41 -0
  32. package/dist/ipc/message.d.ts.map +1 -0
  33. package/dist/ipc/message.js +2 -0
  34. package/dist/ipc/message.js.map +1 -0
  35. package/dist/ipc/proc_job_executor.d.ts +15 -0
  36. package/dist/ipc/proc_job_executor.d.ts.map +1 -0
  37. package/dist/ipc/proc_job_executor.js +150 -0
  38. package/dist/ipc/proc_job_executor.js.map +1 -0
  39. package/dist/ipc/proc_pool.d.ts +26 -0
  40. package/dist/ipc/proc_pool.d.ts.map +1 -0
  41. package/dist/ipc/proc_pool.js +82 -0
  42. package/dist/ipc/proc_pool.js.map +1 -0
  43. package/dist/job.d.ts +99 -0
  44. package/dist/job.d.ts.map +1 -0
  45. package/dist/job.js +197 -0
  46. package/dist/job.js.map +1 -0
  47. package/dist/llm/function_context.d.ts +20 -0
  48. package/dist/llm/function_context.d.ts.map +1 -0
  49. package/dist/llm/function_context.js +37 -0
  50. package/dist/llm/function_context.js.map +1 -0
  51. package/dist/llm/index.d.ts +3 -0
  52. package/dist/llm/index.d.ts.map +1 -0
  53. package/dist/llm/index.js +6 -0
  54. package/dist/llm/index.js.map +1 -0
  55. package/dist/log.d.ts +12 -1
  56. package/dist/log.d.ts.map +1 -1
  57. package/dist/log.js +28 -11
  58. package/dist/log.js.map +1 -1
  59. package/dist/plugin.js +20 -7
  60. package/dist/plugin.js.map +1 -1
  61. package/dist/stt/index.d.ts +1 -1
  62. package/dist/stt/index.d.ts.map +1 -1
  63. package/dist/stt/index.js.map +1 -1
  64. package/dist/stt/stream_adapter.d.ts +2 -11
  65. package/dist/stt/stream_adapter.d.ts.map +1 -1
  66. package/dist/stt/stream_adapter.js +47 -33
  67. package/dist/stt/stream_adapter.js.map +1 -1
  68. package/dist/stt/stt.d.ts +27 -0
  69. package/dist/stt/stt.d.ts.map +1 -1
  70. package/dist/stt/stt.js +32 -5
  71. package/dist/stt/stt.js.map +1 -1
  72. package/dist/tts/stream_adapter.d.ts +4 -11
  73. package/dist/tts/stream_adapter.d.ts.map +1 -1
  74. package/dist/tts/stream_adapter.js +66 -32
  75. package/dist/tts/stream_adapter.js.map +1 -1
  76. package/dist/tts/tts.d.ts +10 -0
  77. package/dist/tts/tts.d.ts.map +1 -1
  78. package/dist/tts/tts.js +48 -7
  79. package/dist/tts/tts.js.map +1 -1
  80. package/dist/utils.d.ts +32 -0
  81. package/dist/utils.d.ts.map +1 -1
  82. package/dist/utils.js +114 -6
  83. package/dist/utils.js.map +1 -1
  84. package/dist/vad.d.ts +29 -0
  85. package/dist/vad.d.ts.map +1 -1
  86. package/dist/vad.js.map +1 -1
  87. package/dist/worker.d.ts +67 -50
  88. package/dist/worker.d.ts.map +1 -1
  89. package/dist/worker.js +379 -214
  90. package/dist/worker.js.map +1 -1
  91. package/package.json +9 -9
  92. package/src/audio.ts +62 -0
  93. package/src/cli.ts +72 -8
  94. package/src/generator.ts +13 -7
  95. package/src/index.ts +13 -3
  96. package/src/ipc/job_executor.ts +25 -0
  97. package/src/ipc/job_main.ts +134 -61
  98. package/src/ipc/message.ts +39 -0
  99. package/src/ipc/proc_job_executor.ts +162 -0
  100. package/src/ipc/proc_pool.ts +108 -0
  101. package/src/job.ts +258 -0
  102. package/src/llm/function_context.ts +61 -0
  103. package/src/llm/index.ts +11 -0
  104. package/src/log.ts +40 -8
  105. package/src/stt/index.ts +1 -1
  106. package/src/stt/stream_adapter.ts +32 -32
  107. package/src/stt/stt.ts +27 -0
  108. package/src/tts/stream_adapter.ts +32 -31
  109. package/src/tts/tts.ts +10 -0
  110. package/src/utils.ts +125 -3
  111. package/src/vad.ts +29 -0
  112. package/src/worker.ts +419 -170
  113. package/tsconfig.json +6 -0
  114. package/dist/ipc/job_process.d.ts +0 -22
  115. package/dist/ipc/job_process.d.ts.map +0 -1
  116. package/dist/ipc/job_process.js +0 -73
  117. package/dist/ipc/job_process.js.map +0 -1
  118. package/dist/ipc/protocol.d.ts +0 -40
  119. package/dist/ipc/protocol.d.ts.map +0 -1
  120. package/dist/ipc/protocol.js +0 -14
  121. package/dist/ipc/protocol.js.map +0 -1
  122. package/dist/job_context.d.ts +0 -16
  123. package/dist/job_context.d.ts.map +0 -1
  124. package/dist/job_context.js +0 -31
  125. package/dist/job_context.js.map +0 -1
  126. package/dist/job_request.d.ts +0 -42
  127. package/dist/job_request.d.ts.map +0 -1
  128. package/dist/job_request.js +0 -79
  129. package/dist/job_request.js.map +0 -1
  130. package/src/ipc/job_process.ts +0 -96
  131. package/src/ipc/protocol.ts +0 -51
  132. package/src/job_context.ts +0 -49
  133. package/src/job_request.ts +0 -118
package/src/worker.ts CHANGED
@@ -1,45 +1,91 @@
1
1
  // SPDX-FileCopyrightText: 2024 LiveKit, Inc.
2
2
  //
3
3
  // SPDX-License-Identifier: Apache-2.0
4
+ import type {
5
+ JobAssignment,
6
+ JobTermination,
7
+ ParticipantInfo,
8
+ TrackSource,
9
+ } from '@livekit/protocol';
4
10
  import {
5
11
  type AvailabilityRequest,
6
- type Job,
7
- type JobAssignment,
8
12
  JobType,
9
13
  ParticipantPermission,
10
14
  ServerMessage,
11
15
  WorkerMessage,
16
+ WorkerStatus,
12
17
  } from '@livekit/protocol';
13
18
  import { EventEmitter } from 'events';
14
- import { AccessToken } from 'livekit-server-sdk';
19
+ import { AccessToken, RoomServiceClient } from 'livekit-server-sdk';
15
20
  import os from 'os';
16
21
  import { WebSocket } from 'ws';
17
22
  import { HTTPServer } from './http_server.js';
18
- import { JobProcess } from './ipc/job_process.js';
19
- import { type AvailRes, JobRequest } from './job_request.js';
20
- import type { AcceptData } from './job_request.js';
23
+ import { ProcPool } from './ipc/proc_pool.js';
24
+ import type { JobAcceptArguments, JobProcess, RunningJobInfo } from './job.js';
25
+ import { JobRequest } from './job.js';
21
26
  import { log } from './log.js';
27
+ import { Future } from './utils.js';
22
28
  import { version } from './version.js';
23
29
 
24
30
  const MAX_RECONNECT_ATTEMPTS = 10;
25
- const ASSIGNMENT_TIMEOUT = 15 * 1000;
26
- const LOAD_INTERVAL = 5 * 1000;
27
-
28
- const cpuLoad = (): number =>
29
- (os
30
- .cpus()
31
- .reduce(
32
- (acc, x) => acc + x.times.idle / Object.values(x.times).reduce((acc, x) => acc + x, 0),
33
- 0,
34
- ) /
35
- os.cpus().length) *
36
- 100;
37
-
38
- class WorkerPermissions {
31
+ const ASSIGNMENT_TIMEOUT = 7.5 * 1000;
32
+ const UPDATE_LOAD_INTERVAL = 2.5 * 1000;
33
+
34
+ /** Necessary credentials not provided and not found in an appropriate environmental variable. */
35
+ export class MissingCredentialsError extends Error {
36
+ constructor(msg?: string) {
37
+ super(msg);
38
+ Object.setPrototypeOf(this, new.target.prototype);
39
+ }
40
+ }
41
+
42
+ /** Worker did not run as expected. */
43
+ export class WorkerError extends Error {
44
+ constructor(msg?: string) {
45
+ super(msg);
46
+ Object.setPrototypeOf(this, new.target.prototype);
47
+ }
48
+ }
49
+
50
+ /** @internal */
51
+ export const defaultInitializeProcessFunc = (_: JobProcess) => _;
52
+ const defaultRequestFunc = async (ctx: JobRequest) => {
53
+ await ctx.accept();
54
+ };
55
+ const defaultCpuLoad = async (): Promise<number> => {
56
+ return new Promise((resolve) => {
57
+ const cpus1 = os.cpus();
58
+
59
+ setTimeout(() => {
60
+ const cpus2 = os.cpus();
61
+
62
+ let idle = 0;
63
+ let total = 0;
64
+
65
+ for (let i = 0; i < cpus1.length; i++) {
66
+ const cpu1 = cpus1[i].times;
67
+ const cpu2 = cpus2[i].times;
68
+
69
+ idle += cpu2.idle - cpu1.idle;
70
+
71
+ const total1 = Object.values(cpu1).reduce((acc, i) => acc + i, 0);
72
+ const total2 = Object.values(cpu2).reduce((acc, i) => acc + i, 0);
73
+
74
+ total += total2 - total1;
75
+ }
76
+
77
+ resolve(+(1 - idle / total).toFixed(2));
78
+ }, UPDATE_LOAD_INTERVAL);
79
+ });
80
+ };
81
+
82
+ /** Participant permissions to pass to every agent spun up by this worker. */
83
+ export class WorkerPermissions {
39
84
  canPublish: boolean;
40
85
  canSubscribe: boolean;
41
86
  canPublishData: boolean;
42
87
  canUpdateMetadata: boolean;
88
+ canPublishSources: TrackSource[];
43
89
  hidden: boolean;
44
90
 
45
91
  constructor(
@@ -47,21 +93,37 @@ class WorkerPermissions {
47
93
  canSubscribe = true,
48
94
  canPublishData = true,
49
95
  canUpdateMetadata = true,
96
+ canPublishSources: TrackSource[] = [],
50
97
  hidden = false,
51
98
  ) {
52
99
  this.canPublish = canPublish;
53
100
  this.canSubscribe = canSubscribe;
54
101
  this.canPublishData = canPublishData;
55
102
  this.canUpdateMetadata = canUpdateMetadata;
103
+ this.canPublishSources = canPublishSources;
56
104
  this.hidden = hidden;
57
105
  }
58
106
  }
59
107
 
108
+ /**
109
+ * Data class describing worker behaviour.
110
+ *
111
+ * @remarks
112
+ * The Agents framework provides sane worker defaults, and works out-of-the-box with no tweaking
113
+ * necessary. The only mandatory parameter is `agent`, which points to the entry function.
114
+ *
115
+ * This class is mostly useful in conjunction with {@link cli.runApp}.
116
+ */
60
117
  export class WorkerOptions {
61
- requestFunc: (arg: JobRequest) => Promise<void>;
62
- cpuLoadFunc: () => number;
63
- namespace: string;
118
+ agent: string;
119
+ requestFunc: (job: JobRequest) => Promise<void>;
120
+ loadFunc: () => Promise<number>;
121
+ loadThreshold: number;
122
+ numIdleProcesses: number;
123
+ shutdownProcessTimeout: number;
124
+ initializeProcessTimeout: number;
64
125
  permissions: WorkerPermissions;
126
+ agentName: string;
65
127
  workerType: JobType;
66
128
  maxRetry: number;
67
129
  wsURL: string;
@@ -69,24 +131,43 @@ export class WorkerOptions {
69
131
  apiSecret?: string;
70
132
  host: string;
71
133
  port: number;
134
+ logLevel: string;
72
135
 
136
+ /** @param options */
73
137
  constructor({
74
- requestFunc,
75
- cpuLoadFunc = cpuLoad,
76
- namespace = 'default',
138
+ agent,
139
+ requestFunc = defaultRequestFunc,
140
+ loadFunc = defaultCpuLoad,
141
+ loadThreshold = 0.65,
142
+ numIdleProcesses = 3,
143
+ shutdownProcessTimeout = 60 * 1000,
144
+ initializeProcessTimeout = 10 * 1000,
77
145
  permissions = new WorkerPermissions(),
78
- workerType = JobType.JT_PUBLISHER,
146
+ agentName = '',
147
+ workerType = JobType.JT_ROOM,
79
148
  maxRetry = MAX_RECONNECT_ATTEMPTS,
80
149
  wsURL = 'ws://localhost:7880',
81
150
  apiKey = undefined,
82
151
  apiSecret = undefined,
83
152
  host = 'localhost',
84
153
  port = 8081,
154
+ logLevel = 'info',
85
155
  }: {
86
- requestFunc: (arg: JobRequest) => Promise<void>;
87
- cpuLoadFunc?: () => number;
88
- namespace?: string;
156
+ /**
157
+ * Path to a file that has {@link Agent} as a default export, dynamically imported later for
158
+ * entrypoint and prewarm functions
159
+ */
160
+ agent: string;
161
+ requestFunc?: (job: JobRequest) => Promise<void>;
162
+ /** Called to determine the current load of the worker. Should return a value between 0 and 1. */
163
+ loadFunc?: () => Promise<number>;
164
+ /** When the load exceeds this threshold, the worker will be marked as unavailable. */
165
+ loadThreshold?: number;
166
+ numIdleProcesses?: number;
167
+ shutdownProcessTimeout?: number;
168
+ initializeProcessTimeout?: number;
89
169
  permissions?: WorkerPermissions;
170
+ agentName?: string;
90
171
  workerType?: JobType;
91
172
  maxRetry?: number;
92
173
  wsURL?: string;
@@ -94,11 +175,17 @@ export class WorkerOptions {
94
175
  apiSecret?: string;
95
176
  host?: string;
96
177
  port?: number;
178
+ logLevel?: string;
97
179
  }) {
180
+ this.agent = agent;
98
181
  this.requestFunc = requestFunc;
99
- this.cpuLoadFunc = cpuLoadFunc;
100
- this.namespace = namespace;
182
+ this.loadFunc = loadFunc;
183
+ this.loadThreshold = loadThreshold;
184
+ this.numIdleProcesses = numIdleProcesses;
185
+ this.shutdownProcessTimeout = shutdownProcessTimeout;
186
+ this.initializeProcessTimeout = initializeProcessTimeout;
101
187
  this.permissions = permissions;
188
+ this.agentName = agentName;
102
189
  this.workerType = workerType;
103
190
  this.maxRetry = maxRetry;
104
191
  this.wsURL = wsURL;
@@ -106,101 +193,119 @@ export class WorkerOptions {
106
193
  this.apiSecret = apiSecret;
107
194
  this.host = host;
108
195
  this.port = port;
196
+ this.logLevel = logLevel;
109
197
  }
110
198
  }
111
199
 
112
- class ActiveJob {
113
- job: Job;
114
- acceptData: AcceptData;
115
-
116
- constructor(job: Job, acceptData: AcceptData) {
117
- this.job = job;
118
- this.acceptData = acceptData;
119
- }
120
- }
121
-
122
- type AssignmentPair = {
123
- // this string is the JSON string version of the JobAssignment.
124
- // we keep it around to unpack it again in the child, because we can't pass Job directly.
125
- raw: string;
126
- asgn: JobAssignment;
127
- };
128
-
129
200
  class PendingAssignment {
130
- promise = new Promise<AssignmentPair>((resolve) => {
131
- this.resolve = resolve; // oh, JavaScript.
201
+ promise = new Promise<JobAssignment>((resolve) => {
202
+ this.resolve = resolve; // this is how JavaScript lets you resolve promises externally
132
203
  });
133
- resolve(arg: AssignmentPair) {
134
- arg;
204
+ resolve(arg: JobAssignment) {
205
+ arg; // useless call to counteract TypeScript E6133
135
206
  }
136
207
  }
137
208
 
209
+ /**
210
+ * Central orchestrator for all processes and job requests.
211
+ *
212
+ * @remarks
213
+ * For most usecases, Worker should not be initialized or handled directly; you should instead call
214
+ * for its creation through {@link cli.runApp}. This could, however, be useful in situations where
215
+ * you don't have access to a command line, such as a headless program, or one that uses Agents
216
+ * behind a wrapper.
217
+ */
138
218
  export class Worker {
139
- opts: WorkerOptions;
219
+ #opts: WorkerOptions;
220
+ #procPool: ProcPool;
221
+
140
222
  #id = 'unregistered';
141
- session: WebSocket | undefined = undefined;
142
- closed = false;
143
- httpServer: HTTPServer;
144
- logger = log.child({ version });
223
+ #closed = true;
224
+ #draining = false;
225
+ #connecting = false;
226
+ #tasks: Promise<void>[] = [];
227
+ #pending: { [id: string]: PendingAssignment } = {};
228
+ #close = new Future();
229
+
145
230
  event = new EventEmitter();
146
- pending: { [id: string]: { value: PendingAssignment } } = {};
147
- processes: { [id: string]: { proc: JobProcess; activeJob: ActiveJob } } = {};
231
+ #session: WebSocket | undefined = undefined;
232
+ #httpServer: HTTPServer;
233
+ #logger = log().child({ version });
148
234
 
235
+ /* @throws {@link MissingCredentialsError} if URL, API key or API secret are missing */
149
236
  constructor(opts: WorkerOptions) {
150
237
  opts.wsURL = opts.wsURL || process.env.LIVEKIT_URL || '';
151
238
  opts.apiKey = opts.apiKey || process.env.LIVEKIT_API_KEY || '';
152
239
  opts.apiSecret = opts.apiSecret || process.env.LIVEKIT_API_SECRET || '';
153
240
 
154
- this.opts = opts;
155
- this.httpServer = new HTTPServer(opts.host, opts.port);
156
- }
241
+ if (opts.wsURL === '')
242
+ throw new MissingCredentialsError('--url is required, or set LIVEKIT_URL env var');
243
+ if (opts.apiKey === '')
244
+ throw new MissingCredentialsError('--api-key is required, or set LIVEKIT_API_KEY env var');
245
+ if (opts.apiSecret === '')
246
+ throw new MissingCredentialsError(
247
+ '--api-secret is required, or set LIVEKIT_API_SECRET env var',
248
+ );
157
249
 
158
- get id(): string {
159
- return this.#id;
250
+ this.#procPool = new ProcPool(
251
+ opts.agent,
252
+ opts.numIdleProcesses,
253
+ opts.initializeProcessTimeout,
254
+ opts.shutdownProcessTimeout,
255
+ );
256
+
257
+ this.#opts = opts;
258
+ this.#httpServer = new HTTPServer(opts.host, opts.port);
160
259
  }
161
260
 
261
+ /* @throws {@link WorkerError} if worker failed to connect or already running */
162
262
  async run() {
163
- this.logger.info('starting worker');
263
+ if (!this.#closed) {
264
+ throw new WorkerError('worker is already running');
265
+ }
164
266
 
165
- if (this.opts.wsURL === '') throw new Error('--url is required, or set LIVEKIT_URL env var');
166
- if (this.opts.apiKey === '')
167
- throw new Error('--api-key is required, or set LIVEKIT_API_KEY env var');
168
- if (this.opts.apiSecret === '')
169
- throw new Error('--api-secret is required, or set LIVEKIT_API_SECRET env var');
267
+ this.#logger.info('starting worker');
268
+ this.#closed = false;
269
+ this.#procPool.start();
170
270
 
171
271
  const workerWS = async () => {
172
272
  let retries = 0;
173
- while (!this.closed) {
174
- const token = new AccessToken(this.opts.apiKey, this.opts.apiSecret);
175
- token.addGrant({ agent: true });
176
- const jwt = await token.toJwt();
273
+ this.#connecting = true;
177
274
 
178
- const url = new URL(this.opts.wsURL);
275
+ while (!this.#closed) {
276
+ const url = new URL(this.#opts.wsURL);
179
277
  url.protocol = url.protocol.replace('http', 'ws');
180
- this.session = new WebSocket(url + 'agent', {
278
+ const token = new AccessToken(this.#opts.apiKey, this.#opts.apiSecret);
279
+ token.addGrant({ agent: true });
280
+ const jwt = await token.toJwt();
281
+ this.#session = new WebSocket(url + 'agent', {
181
282
  headers: { authorization: 'Bearer ' + jwt },
182
283
  });
183
284
 
184
285
  try {
185
286
  await new Promise((resolve, reject) => {
186
- this.session!.on('open', resolve);
187
- this.session!.on('error', (error) => reject(error));
188
- this.session!.on('close', (code) => reject(`WebSocket returned ${code}`));
287
+ this.#session!.on('open', resolve);
288
+ this.#session!.on('error', (error) => reject(error));
289
+ this.#session!.on('close', (code) => reject(`WebSocket returned ${code}`));
189
290
  });
190
291
 
191
- this.runWS(this.session!);
292
+ retries = 0;
293
+ this.#logger.debug('connected to LiveKit server');
294
+ this.#runWS(this.#session);
192
295
  return;
193
296
  } catch (e) {
194
- if (this.closed) return;
195
- if (retries >= this.opts.maxRetry) {
196
- throw new Error(`failed to connect to LiveKit server after ${retries} attempts: ${e}`);
297
+ if (this.#closed) return;
298
+ if (retries >= this.#opts.maxRetry) {
299
+ throw new WorkerError(
300
+ `failed to connect to LiveKit server after ${retries} attempts: ${e}`,
301
+ );
197
302
  }
198
303
 
199
304
  retries++;
200
305
  const delay = Math.min(retries * 2, 10);
201
306
 
202
- this.logger.warn(
203
- `failed to connect to LiveKit server, retrying in ${delay} seconds: ${e} (${retries}/${this.opts.maxRetry})`,
307
+ this.#logger.warn(
308
+ `failed to connect to LiveKit server, retrying in ${delay} seconds: ${e} (${retries}/${this.#opts.maxRetry})`,
204
309
  );
205
310
 
206
311
  await new Promise((resolve) => setTimeout(resolve, delay * 1000));
@@ -208,24 +313,85 @@ export class Worker {
208
313
  }
209
314
  };
210
315
 
211
- await Promise.all([workerWS(), this.httpServer.run()]);
316
+ await Promise.all([workerWS(), this.#httpServer.run()]);
317
+ this.#close.resolve();
212
318
  }
213
319
 
214
- startProcess(job: Job, acceptData: AcceptData, raw: string) {
215
- const proc = new JobProcess(job, acceptData, raw, this.opts.wsURL);
216
- this.processes[job.id] = { proc, activeJob: new ActiveJob(job, acceptData) };
217
- proc
218
- .run()
219
- .catch((e) => {
220
- proc.logger.error(`error running job process ${proc.job.id}: ${e}`);
221
- })
222
- .finally(() => {
223
- proc.clear();
224
- delete this.processes[job.id];
225
- });
320
+ get id(): string {
321
+ return this.#id;
322
+ }
323
+
324
+ get activeJobs(): RunningJobInfo[] {
325
+ return this.#procPool.processes
326
+ .filter((proc) => proc.runningJob)
327
+ .map((proc) => proc.runningJob!);
226
328
  }
227
329
 
228
- runWS(ws: WebSocket) {
330
+ /* @throws {@link WorkerError} if worker did not drain in time */
331
+ async drain(timeout?: number) {
332
+ if (this.#draining) {
333
+ return;
334
+ }
335
+
336
+ this.#logger.info('draining worker');
337
+ this.#draining = true;
338
+
339
+ this.event.emit(
340
+ 'worker_msg',
341
+ new WorkerMessage({
342
+ message: {
343
+ case: 'updateWorker',
344
+ value: {
345
+ status: WorkerStatus.WS_FULL,
346
+ },
347
+ },
348
+ }),
349
+ );
350
+
351
+ const joinJobs = async () => {
352
+ return Promise.all(
353
+ this.#procPool.processes.map((proc) => {
354
+ if (!proc.runningJob) {
355
+ proc.close();
356
+ }
357
+ return proc.join();
358
+ }),
359
+ );
360
+ };
361
+
362
+ const timer = setTimeout(() => {
363
+ throw new WorkerError('timed out draining');
364
+ }, timeout);
365
+ if (timeout === undefined) clearTimeout(timer);
366
+ await joinJobs().then(() => {
367
+ clearTimeout(timer);
368
+ });
369
+ }
370
+
371
+ async simulateJob(roomName: string, participantIdentity?: string) {
372
+ const client = new RoomServiceClient(this.#opts.wsURL, this.#opts.apiKey, this.#opts.apiSecret);
373
+ const room = await client.createRoom({ name: roomName });
374
+ let participant: ParticipantInfo | undefined = undefined;
375
+ if (participantIdentity) {
376
+ participant = await client.getParticipant(roomName, participantIdentity);
377
+ }
378
+
379
+ this.event!.emit(
380
+ 'worker_msg',
381
+ new WorkerMessage({
382
+ message: {
383
+ case: 'simulateJob',
384
+ value: {
385
+ type: JobType.JT_PUBLISHER,
386
+ room,
387
+ participant,
388
+ },
389
+ },
390
+ }),
391
+ );
392
+ }
393
+
394
+ #runWS(ws: WebSocket) {
229
395
  let closingWS = false;
230
396
 
231
397
  const send = (msg: WorkerMessage) => {
@@ -239,44 +405,62 @@ export class Worker {
239
405
 
240
406
  ws.addEventListener('close', () => {
241
407
  closingWS = true;
242
- this.logger.error('worker connection closed unexpectedly');
408
+ this.#logger.error('worker connection closed unexpectedly');
243
409
  this.close();
244
410
  });
245
411
 
246
412
  ws.addEventListener('message', (event) => {
247
413
  if (event.type !== 'message') {
248
- this.logger.warn('unexpected message type: ' + event.type);
414
+ this.#logger.warn('unexpected message type: ' + event.type);
249
415
  return;
250
416
  }
251
417
 
252
418
  const msg = new ServerMessage();
253
419
  msg.fromBinary(event.data as Uint8Array);
420
+
421
+ // register is the only valid first message, and it is only valid as the
422
+ // first message
423
+ if (this.#connecting && msg.message.case !== 'register') {
424
+ throw new WorkerError('expected register response as first message');
425
+ }
426
+
254
427
  switch (msg.message.case) {
255
428
  case 'register': {
256
429
  this.#id = msg.message.value.workerId;
257
- log
430
+ this.#logger
258
431
  .child({ id: this.id, server_info: msg.message.value.serverInfo })
259
432
  .info('registered worker');
433
+ this.event.emit(
434
+ 'worker_registered',
435
+ msg.message.value.workerId,
436
+ msg.message.value.serverInfo!,
437
+ );
438
+ this.#connecting = false;
260
439
  break;
261
440
  }
262
441
  case 'availability': {
263
- this.availability(msg.message.value);
442
+ const task = this.#availability(msg.message.value);
443
+ this.#tasks.push(task);
444
+ task.finally(() => this.#tasks.splice(this.#tasks.indexOf(task)));
264
445
  break;
265
446
  }
266
447
  case 'assignment': {
267
448
  const job = msg.message.value.job!;
268
- if (job.id in this.pending) {
269
- const task = this.pending[job.id];
270
- delete this.pending[job.id];
271
- task.value.resolve({
272
- asgn: msg.message.value,
273
- raw: msg.toJsonString(),
274
- });
449
+ if (job.id in this.#pending) {
450
+ const task = this.#pending[job.id];
451
+ delete this.#pending[job.id];
452
+ task.resolve(msg.message.value);
275
453
  } else {
276
- log.child({ job }).warn('received assignment for unknown job ' + job.id);
454
+ this.#logger.child({ job }).warn('received assignment for unknown job ' + job.id);
277
455
  }
278
456
  break;
279
457
  }
458
+ case 'termination': {
459
+ const task = this.#termination(msg.message.value);
460
+ this.#tasks.push(task);
461
+ task.finally(() => this.#tasks.splice(this.#tasks.indexOf(task)));
462
+ break;
463
+ }
280
464
  }
281
465
  });
282
466
 
@@ -286,13 +470,14 @@ export class Worker {
286
470
  message: {
287
471
  case: 'register',
288
472
  value: {
289
- type: this.opts.workerType,
290
- namespace: this.opts.namespace,
473
+ type: this.#opts.workerType,
474
+ agentName: this.#opts.agentName,
291
475
  allowedPermissions: new ParticipantPermission({
292
- canPublish: this.opts.permissions.canPublish,
293
- canSubscribe: this.opts.permissions.canSubscribe,
294
- canPublishData: this.opts.permissions.canPublishData,
295
- hidden: this.opts.permissions.hidden,
476
+ canPublish: this.#opts.permissions.canPublish,
477
+ canSubscribe: this.#opts.permissions.canSubscribe,
478
+ canPublishData: this.#opts.permissions.canPublishData,
479
+ canUpdateMetadata: this.#opts.permissions.canUpdateMetadata,
480
+ hidden: this.#opts.permissions.hidden,
296
481
  agent: true,
297
482
  }),
298
483
  version,
@@ -301,84 +486,148 @@ export class Worker {
301
486
  }),
302
487
  );
303
488
 
489
+ let currentStatus = WorkerStatus.WS_AVAILABLE;
304
490
  const loadMonitor = setInterval(() => {
305
491
  if (closingWS) clearInterval(loadMonitor);
492
+
493
+ const oldStatus = currentStatus;
494
+ this.#opts.loadFunc().then((currentLoad: number) => {
495
+ const isFull = currentLoad >= this.#opts.loadThreshold;
496
+ const currentlyAvailable = !isFull;
497
+ currentStatus = currentlyAvailable ? WorkerStatus.WS_AVAILABLE : WorkerStatus.WS_FULL;
498
+
499
+ if (oldStatus != currentStatus) {
500
+ const extra = { load: currentLoad, loadThreshold: this.#opts.loadThreshold };
501
+ if (isFull) {
502
+ this.#logger.child(extra).info('worker is at full capacity, marking as unavailable');
503
+ } else {
504
+ this.#logger.child(extra).info('worker is below capacity, marking as available');
505
+ }
506
+ }
507
+
508
+ this.event.emit(
509
+ 'worker_msg',
510
+ new WorkerMessage({
511
+ message: {
512
+ case: 'updateWorker',
513
+ value: {
514
+ load: currentLoad,
515
+ status: currentStatus,
516
+ },
517
+ },
518
+ }),
519
+ );
520
+ });
521
+ }, UPDATE_LOAD_INTERVAL);
522
+ }
523
+
524
+ async #availability(msg: AvailabilityRequest) {
525
+ let answered = false;
526
+
527
+ const onReject = async () => {
528
+ answered = true;
306
529
  this.event.emit(
307
530
  'worker_msg',
308
531
  new WorkerMessage({
309
532
  message: {
310
- case: 'updateWorker',
533
+ case: 'availability',
311
534
  value: {
312
- load: cpuLoad(),
535
+ jobId: msg.job!.id,
536
+ available: false,
313
537
  },
314
538
  },
315
539
  }),
316
540
  );
317
- }, LOAD_INTERVAL);
318
- }
541
+ };
319
542
 
320
- async availability(msg: AvailabilityRequest) {
321
- const tx = new EventEmitter();
322
- const req = new JobRequest(msg.job!, tx);
543
+ const onAccept = async (args: JobAcceptArguments) => {
544
+ answered = true;
323
545
 
324
- tx.on('recv', async (av: AvailRes) => {
325
- const msg = new WorkerMessage({
326
- message: {
327
- case: 'availability',
328
- value: {
329
- available: av.avail,
330
- jobId: req.id,
331
- participantIdentity: av.data?.identity,
332
- participantName: av.data?.name,
333
- participantMetadata: av.data?.metadata,
546
+ this.event.emit(
547
+ 'worker_msg',
548
+ new WorkerMessage({
549
+ message: {
550
+ case: 'availability',
551
+ value: {
552
+ jobId: msg.job!.id,
553
+ available: true,
554
+ participantIdentity: args.identity,
555
+ participantName: args.name,
556
+ participantMetadata: args.metadata,
557
+ },
334
558
  },
335
- },
336
- });
337
-
338
- this.pending[req.id] = { value: new PendingAssignment() };
339
- this.event.emit('worker_msg', msg);
340
- if (!av.avail) return;
559
+ }),
560
+ );
341
561
 
562
+ this.#pending[req.id] = new PendingAssignment();
342
563
  const timer = setTimeout(() => {
343
- log.child({ req }).warn(`assignment for job ${req.id} timed out`);
564
+ this.#logger.child({ req }).warn(`assignment for job ${req.id} timed out`);
344
565
  return;
345
566
  }, ASSIGNMENT_TIMEOUT);
346
- this.pending[req.id].value.promise.then(({ asgn, raw }) => {
567
+ const asgn = await this.#pending[req.id].promise.then(async (asgn) => {
347
568
  clearTimeout(timer);
348
- this.startProcess(asgn!.job!, av.data!, raw);
569
+ return asgn;
349
570
  });
350
- });
351
571
 
352
- try {
353
- this.opts.requestFunc(req);
354
- } catch (e) {
355
- log.child({ req }).error(`user request handler for job ${req.id} failed`);
356
- } finally {
357
- if (!req.answered) {
358
- log.child({ req }).error(`no answer for job ${req.id}, automatically rejecting the job`);
359
- this.event.emit(
360
- 'worker_msg',
361
- new WorkerMessage({
362
- message: {
363
- case: 'availability',
364
- value: {
365
- available: false,
366
- },
367
- },
368
- }),
369
- );
572
+ await this.#procPool.launchJob({
573
+ acceptArguments: args,
574
+ job: msg.job!,
575
+ url: asgn.url || this.#opts.wsURL,
576
+ token: asgn.token,
577
+ });
578
+ };
579
+
580
+ const req = new JobRequest(msg.job!, onReject, onAccept);
581
+ this.#logger
582
+ .child({ job: msg.job, resuming: msg.resuming, agentName: this.#opts.agentName })
583
+ .info('received job request');
584
+
585
+ const jobRequestTask = async () => {
586
+ try {
587
+ await this.#opts.requestFunc(req);
588
+ } catch (e) {
589
+ this.#logger
590
+ .child({ job: msg.job, resuming: msg.resuming, agentName: this.#opts.agentName })
591
+ .info('jobRequestFunc failed');
592
+ await onReject();
593
+ }
594
+
595
+ if (!answered) {
596
+ this.#logger
597
+ .child({ job: msg.job, resuming: msg.resuming, agentName: this.#opts.agentName })
598
+ .info('no answer was given inside the jobRequestFunc, automatically rejecting the job');
370
599
  }
600
+ };
601
+
602
+ const task = jobRequestTask();
603
+ this.#tasks.push(task);
604
+ task.finally(() => this.#tasks.splice(this.#tasks.indexOf(task)));
605
+ }
606
+
607
+ async #termination(msg: JobTermination) {
608
+ const proc = this.#procPool.getByJobId(msg.jobId);
609
+ if (proc === null) {
610
+ // safe to ignore
611
+ return;
371
612
  }
613
+ await proc.close();
372
614
  }
373
615
 
374
616
  async close() {
375
- if (this.closed) return;
376
- this.closed = true;
377
- this.logger.debug('shutting down worker');
378
- await this.httpServer.close();
379
- for await (const value of Object.values(this.processes)) {
380
- await value.proc.close();
617
+ if (this.#closed) {
618
+ await this.#close.await;
619
+ return;
381
620
  }
382
- this.session?.close();
621
+
622
+ this.#logger.info('shutting down worker');
623
+
624
+ this.#closed = true;
625
+
626
+ await this.#procPool.close();
627
+ await this.#httpServer.close();
628
+ await Promise.allSettled(this.#tasks);
629
+
630
+ this.#session?.close();
631
+ await this.#close.await;
383
632
  }
384
633
  }