worker.ts

// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
//
// SPDX-License-Identifier: Apache-2.0
import type { JobAssignment, JobTermination, TrackSource } from '@livekit/protocol';
import {
  type AvailabilityRequest,
  JobType,
  ParticipantPermission,
  ServerMessage,
  WorkerMessage,
  WorkerStatus,
} from '@livekit/protocol';
import type { Throws } from '@livekit/throws-transformer/throws';
import type { ParticipantInfo } from 'livekit-server-sdk';
import { AccessToken, RoomServiceClient } from 'livekit-server-sdk';
import { EventEmitter } from 'node:events';
import { WebSocket } from 'ws';
import { getCpuMonitor } from './cpu.js';
import { HTTPServer } from './http_server.js';
import { InferenceRunner } from './inference_runner.js';
import { InferenceProcExecutor } from './ipc/inference_proc_executor.js';
import { ProcPool } from './ipc/proc_pool.js';
import type { JobAcceptArguments, JobProcess, RunningJobInfo } from './job.js';
import { JobRequest } from './job.js';
import { log } from './log.js';
import { Future, rejectOnAbort, waitForAbort } from './utils.js';
import { version } from './version.js';

const MAX_RECONNECT_ATTEMPTS = 10;
const ASSIGNMENT_TIMEOUT = 7.5 * 1000;
const UPDATE_LOAD_INTERVAL = 2.5 * 1000;
const PROJECT_TYPE = 'nodejs';

class Default {
  static loadThreshold(production: boolean): number {
    if (production) {
      return 0.7;
    } else {
      return Infinity;
    }
  }

  static numIdleProcesses(production: boolean): number {
    if (production) {
      // TODO: use number of cores
      return 3;
    } else {
      return 0;
    }
  }

  static port(production: boolean): number {
    if (production) {
      return 8081;
    } else {
      return 0;
    }
  }
}

/** Necessary credentials not provided and not found in an appropriate environment variable. */
export class MissingCredentialsError extends Error {
  constructor(msg?: string) {
    super(msg);
    Object.setPrototypeOf(this, new.target.prototype);
  }
}

/** Worker did not run as expected. */
export class WorkerError extends Error {
  constructor(msg?: string) {
    super(msg);
    Object.setPrototypeOf(this, new.target.prototype);
  }
}

/** @internal */
export const defaultInitializeProcessFunc = (_: JobProcess) => _;
const defaultRequestFunc = async (ctx: JobRequest) => {
  await ctx.accept();
};

const cpuMonitor = getCpuMonitor();

// eslint-disable-next-line @typescript-eslint/no-unused-vars
const defaultCpuLoad = async (_worker: AgentServer): Promise<number> => {
  return cpuMonitor.cpuPercent(UPDATE_LOAD_INTERVAL);
};

/** Participant permissions to pass to every agent spun up by this worker. */
export class WorkerPermissions {
  canPublish: boolean;
  canSubscribe: boolean;
  canPublishData: boolean;
  canUpdateMetadata: boolean;
  canPublishSources: TrackSource[];
  hidden: boolean;

  constructor(
    canPublish = true,
    canSubscribe = true,
    canPublishData = true,
    canUpdateMetadata = true,
    canPublishSources: TrackSource[] = [],
    hidden = false,
  ) {
    this.canPublish = canPublish;
    this.canSubscribe = canSubscribe;
    this.canPublishData = canPublishData;
    this.canUpdateMetadata = canUpdateMetadata;
    this.canPublishSources = canPublishSources;
    this.hidden = hidden;
  }
}

/**
 * Data class describing worker behaviour.
 *
 * @remarks
 * The Agents framework provides sane worker defaults, and works out-of-the-box with no tweaking
 * necessary. The only mandatory parameter is `agent`, which points to the entry function.
 *
 * This class is mostly useful in conjunction with {@link cli.runApp}.
 */
export class ServerOptions {
  agent: string;
  requestFunc: (job: JobRequest) => Promise<void>;
  loadFunc: (worker: AgentServer) => Promise<number>;
  loadThreshold: number;
  numIdleProcesses: number;
  shutdownProcessTimeout: number;
  initializeProcessTimeout: number;
  permissions: WorkerPermissions;
  agentName: string;
  serverType: JobType;
  maxRetry: number;
  wsURL: string;
  apiKey?: string;
  apiSecret?: string;
  workerToken?: string;
  host: string;
  port: number;
  logLevel: string;
  production: boolean;
  jobMemoryWarnMB: number;
  jobMemoryLimitMB: number;

  /** @param options - Worker options */
  constructor({
    agent,
    requestFunc = defaultRequestFunc,
    loadFunc = defaultCpuLoad,
    loadThreshold = undefined,
    numIdleProcesses = undefined,
    shutdownProcessTimeout = 60 * 1000,
    initializeProcessTimeout = 10 * 1000,
    permissions = new WorkerPermissions(),
    agentName = '',
    serverType = JobType.JT_ROOM,
    maxRetry = MAX_RECONNECT_ATTEMPTS,
    wsURL = 'ws://localhost:7880',
    apiKey = undefined,
    apiSecret = undefined,
    workerToken = undefined,
    host = '0.0.0.0',
    port = undefined,
    logLevel = 'info',
    production = false,
    jobMemoryWarnMB = 500,
    jobMemoryLimitMB = 0,
  }: {
    /**
     * Path to a file that has {@link Agent} as a default export, dynamically imported later for
     * entrypoint and prewarm functions
     */
    agent: string;
    requestFunc?: (job: JobRequest) => Promise<void>;
    /** Called to determine the current load of the worker. Should return a value between 0 and 1. */
    loadFunc?: (worker: AgentServer) => Promise<number>;
    /** When the load exceeds this threshold, the worker will be marked as unavailable. */
    loadThreshold?: number;
    numIdleProcesses?: number;
    shutdownProcessTimeout?: number;
    initializeProcessTimeout?: number;
    permissions?: WorkerPermissions;
    agentName?: string;
    serverType?: JobType;
    maxRetry?: number;
    wsURL?: string;
    apiKey?: string;
    apiSecret?: string;
    workerToken?: string;
    host?: string;
    port?: number;
    logLevel?: string;
    production?: boolean;
    jobMemoryWarnMB?: number;
    jobMemoryLimitMB?: number;
  }) {
    this.agent = agent;
    if (!this.agent) {
      throw new Error('No Agent file was passed to the worker');
    }
    this.requestFunc = requestFunc;
    this.loadFunc = loadFunc;
    this.loadThreshold = loadThreshold || Default.loadThreshold(production);
    this.numIdleProcesses = numIdleProcesses || Default.numIdleProcesses(production);
    this.shutdownProcessTimeout = shutdownProcessTimeout;
    this.initializeProcessTimeout = initializeProcessTimeout;
    this.permissions = permissions;
    this.agentName = agentName;
    this.serverType = serverType;
    this.maxRetry = maxRetry;
    this.wsURL = wsURL;
    this.apiKey = apiKey;
    this.apiSecret = apiSecret;
    this.workerToken = workerToken;
    this.host = host;
    this.port = port || Default.port(production);
    this.logLevel = logLevel;
    this.production = production;
    this.jobMemoryWarnMB = jobMemoryWarnMB;
    this.jobMemoryLimitMB = jobMemoryLimitMB;
  }
}

class PendingAssignment {
  promise = new Promise<JobAssignment>((resolve) => {
    this.resolve = resolve; // this is how JavaScript lets you resolve promises externally
  });
  resolve(arg: JobAssignment) {
    arg; // useless call to counteract TypeScript E6133
  }
}

/**
 * Central orchestrator for all processes and job requests.
 *
 * @remarks
 * For most usecases, Worker should not be initialized or handled directly; you should instead call
 * for its creation through {@link cli.runApp}. This could, however, be useful in situations where
 * you don't have access to a command line, such as a headless program, or one that uses Agents
 * behind a wrapper.
 */
export class AgentServer {
  #opts: ServerOptions;
  #procPool: ProcPool;

  #id = 'unregistered';
  #closed = true;
  #draining = false;
  #connecting = false;
  #tasks: Promise<void>[] = [];
  #pending: { [id: string]: PendingAssignment } = {};
  #close = new Future();

  event = new EventEmitter();
  #session: WebSocket | undefined = undefined;
  #httpServer: HTTPServer;
  #logger = log().child({ version });
  #inferenceExecutor?: InferenceProcExecutor;

  /* @throws {@link MissingCredentialsError} if URL, API key or API secret are missing */
  constructor(opts: ServerOptions) {
    opts.wsURL = opts.wsURL || process.env.LIVEKIT_URL || '';
    opts.apiKey = opts.apiKey || process.env.LIVEKIT_API_KEY || '';
    opts.apiSecret = opts.apiSecret || process.env.LIVEKIT_API_SECRET || '';

    if (opts.wsURL === '')
      throw new MissingCredentialsError(
        'URL is required: Set LIVEKIT_URL, run with --url, or pass wsURL in ServerOptions',
      );
    if (opts.apiKey === '')
      throw new MissingCredentialsError(
        'API Key is required: Set LIVEKIT_API_KEY, run with --api-key, or pass apiKey in ServerOptions',
      );
    if (opts.apiSecret === '')
      throw new MissingCredentialsError(
        'API Secret is required: Set LIVEKIT_API_SECRET, run with --api-secret, or pass apiSecret in ServerOptions',
      );

    if (opts.workerToken) {
      if (opts.loadFunc !== defaultCpuLoad) {
        this.#logger.warn(
          'custom loadFunc is not supported when deploying to Cloud, using defaults',
        );
        opts.loadFunc = defaultCpuLoad;
      }
      const loadThreshold = Default.loadThreshold(opts.production);
      if (opts.loadThreshold !== loadThreshold) {
        this.#logger.warn(
          'custom loadThreshold is not supported when deploying to Cloud, using defaults',
        );
        opts.loadThreshold = loadThreshold;
      }
    }

    if (Object.entries(InferenceRunner.registeredRunners).length) {
      this.#inferenceExecutor = new InferenceProcExecutor({
        runners: InferenceRunner.registeredRunners,
        initializeTimeout: 30000,
        closeTimeout: 5000,
        memoryWarnMB: 2000,
        memoryLimitMB: 0,
        pingInterval: 5000,
        pingTimeout: 60000,
        highPingThreshold: 2500,
      });
    }

    this.#procPool = new ProcPool(
      opts.agent,
      opts.numIdleProcesses,
      opts.initializeProcessTimeout,
      opts.shutdownProcessTimeout,
      this.#inferenceExecutor,
      opts.jobMemoryWarnMB,
      opts.jobMemoryLimitMB,
    );

    this.#opts = opts;

    const healthCheck = () => {
      // Check if inference executor exists and is not alive
      if (this.#inferenceExecutor && !this.#inferenceExecutor.isAlive) {
        return { healthy: false, message: 'inference process not running' };
      }

      // Only healthy when fully connected with an active WebSocket
      if (
        this.#closed ||
        this.#connecting ||
        !this.#session ||
        this.#session.readyState !== WebSocket.OPEN
      ) {
        return { healthy: false, message: 'not connected to livekit' };
      }

      return { healthy: true, message: 'OK' };
    };

    const getWorkerInfo = () => ({
      agent_name: opts.agentName,
      worker_type: JobType[opts.serverType],
      active_jobs: this.activeJobs.length,
      sdk_version: version,
      project_type: PROJECT_TYPE,
    });

    this.#httpServer = new HTTPServer(opts.host, opts.port, healthCheck, getWorkerInfo);
  }

  /** @throws {@link WorkerError} if worker failed to connect or already running */
  async run() {
    if (!this.#closed) {
      throw new WorkerError('worker is already running');
    }

    if (this.#inferenceExecutor) {
      await this.#inferenceExecutor.start();
      await this.#inferenceExecutor.initialize();
    }

    this.#logger.info('starting worker');
    this.#closed = false;
    this.#procPool.start();

    const workerWS = async () => {
      let retries = 0;
      this.#connecting = true;

      while (!this.#closed) {
        const url = new URL(this.#opts.wsURL);
        url.protocol = url.protocol.replace('http', 'ws');
        const token = new AccessToken(this.#opts.apiKey, this.#opts.apiSecret);
        token.addGrant({ agent: true });
        const jwt = await token.toJwt();
        const wsUrl = new URL(url + 'agent');
        if (this.#opts.workerToken) {
          wsUrl.searchParams.append('worker_token', this.#opts.workerToken);
        }
        this.#session = new WebSocket(wsUrl, {
          headers: { authorization: 'Bearer ' + jwt },
        });

        try {
          await new Promise((resolve, reject) => {
            this.#session!.on('open', resolve);
            this.#session!.on('error', (error) => reject(error));
            this.#session!.on('close', (code) => reject(`WebSocket returned ${code}`));
          });

          retries = 0;
          this.#logger.debug('connected to LiveKit server');
          await this.#runWS(this.#session);
        } catch (e: unknown) {
          if (this.#closed) return;
          if (retries >= this.#opts.maxRetry) {
            throw new WorkerError(
              `failed to connect to LiveKit server (${this.#opts.wsURL}) after ${retries} attempts: ${e}`,
            );
          }

          retries++;
          const delay = Math.min(retries * 2, 10);

          this.#logger.warn(
            e,
            `failed to connect to LiveKit server (${this.#opts.wsURL}), retrying in ${delay} seconds: (${retries}/${this.#opts.maxRetry})`,
          );

          await new Promise((resolve) => setTimeout(resolve, delay * 1000));
        }
      }
    };

    await Promise.all([workerWS(), this.#httpServer.run()]);
    this.#close.resolve();
  }

  get id(): string {
    return this.#id;
  }

  get activeJobs(): RunningJobInfo[] {
    return this.#procPool.processes
      .filter((proc) => proc.runningJob)
      .map((proc) => proc.runningJob!);
  }

  /** @throws {@link WorkerError} if worker did not drain in time */
  async drain(timeout?: number): Promise<Throws<void, WorkerError>> {
    if (this.#draining) {
      return;
    }

    this.#logger.debug('draining worker');
    this.#draining = true;

    this.event.emit(
      'worker_msg',
      new WorkerMessage({
        message: {
          case: 'updateWorker',
          value: {
            status: WorkerStatus.WS_FULL,
          },
        },
      }),
    );

    const joinJobs = async () => {
      return Promise.all(
        this.#procPool.processes.map((proc): Promise<Throws<void, Error>> => {
          if (!proc.runningJob) {
            proc.close();
          }
          return proc.join();
        }),
      );
    };

    const promises = [joinJobs()];

    if (timeout) {
      promises.push(rejectOnAbort(AbortSignal.timeout(timeout)));
    }
    await Promise.race(promises);
  }

  async simulateJob(roomName: string, participantIdentity?: string) {
    const client = new RoomServiceClient(this.#opts.wsURL, this.#opts.apiKey, this.#opts.apiSecret);
    const room = await client.createRoom({ name: roomName });
    let participant: ParticipantInfo | undefined = undefined;
    if (participantIdentity) {
      try {
        // TODO(AJS-269): resolve compatibility issue with node-sdk to remove the forced type casting
        participant = (await client.getParticipant(
          roomName,
          participantIdentity,
        )) as unknown as ParticipantInfo;
      } catch (e) {
        this.#logger.fatal(
          `participant with identity ${participantIdentity} not found in room ${roomName}`,
        );
        throw e;
      }
    }

    this.event.emit(
      'worker_msg',
      new WorkerMessage({
        message: {
          case: 'simulateJob',
          value: {
            type: JobType.JT_PUBLISHER,
            room,
            participant,
          },
        },
      }),
    );
  }

  async #runWS(ws: WebSocket) {
    let closingWS = false;

    const send = (msg: WorkerMessage) => {
      if (closingWS) {
        this.event.off('worker_msg', send);
        return;
      }
      ws.send(msg.toBinary());
    };
    this.event.on('worker_msg', send);

    const close = new Promise<void>((resolve) => {
      ws.addEventListener('close', () => {
        closingWS = true;
        if (!this.#closed) {
          this.#logger.error('worker connection closed unexpectedly');
        }
        resolve();
      });
    });

    ws.addEventListener('error', (event) => {
      this.#logger.error('worker error:', event.message);
    });

    ws.addEventListener('message', (event) => {
      if (event.type !== 'message') {
        this.#logger.warn('unexpected message type: ' + event.type);
        return;
      }

      const msg = new ServerMessage();
      msg.fromBinary(event.data as Uint8Array);

      // register is the only valid first message, and it is only valid as the
      // first message
      if (this.#connecting && msg.message.case !== 'register') {
        throw new WorkerError('expected register response as first message');
      }

      switch (msg.message.case) {
        case 'register': {
          this.#id = msg.message.value.workerId;
          this.#logger
            .child({ id: this.id, server_info: msg.message.value.serverInfo })
            .info('registered worker');
          this.event.emit(
            'worker_registered',
            msg.message.value.workerId,
            msg.message.value.serverInfo,
          );
          this.#connecting = false;
          break;
        }
        case 'availability': {
          if (!msg.message.value.job) return;
          const task = this.#availability(msg.message.value);
          this.#tasks.push(task);
          task.finally(() => {
            const taskIndex = this.#tasks.indexOf(task);
            if (taskIndex !== -1) {
              this.#tasks.splice(taskIndex, 1);
            } else {
              throw new Error(`task ${task} not found in tasks`);
            }
          });
          break;
        }
        case 'assignment': {
          if (!msg.message.value.job) return;
          const job = msg.message.value.job;
          if (job.id in this.#pending) {
            const task = this.#pending[job.id];
            delete this.#pending[job.id];
            task?.resolve(msg.message.value);
          } else {
            this.#logger.child({ job }).warn('received assignment for unknown job ' + job.id);
          }
          break;
        }
        case 'termination': {
          const task = this.#termination(msg.message.value);
          this.#tasks.push(task);
          task.finally(() => {
            const taskIndex = this.#tasks.indexOf(task);
            if (taskIndex !== -1) {
              this.#tasks.splice(taskIndex, 1);
            } else {
              throw new Error(`task ${task} not found in tasks`);
            }
          });
          break;
        }
      }
    });

    this.event.emit(
      'worker_msg',
      new WorkerMessage({
        message: {
          case: 'register',
          value: {
            type: this.#opts.serverType,
            agentName: this.#opts.agentName,
            allowedPermissions: new ParticipantPermission({
              canPublish: this.#opts.permissions.canPublish,
              canSubscribe: this.#opts.permissions.canSubscribe,
              canPublishData: this.#opts.permissions.canPublishData,
              canUpdateMetadata: this.#opts.permissions.canUpdateMetadata,
              hidden: this.#opts.permissions.hidden,
              agent: true,
            }),
            version,
          },
        },
      }),
    );

    let currentStatus = WorkerStatus.WS_AVAILABLE;
    const loadMonitor = setInterval(() => {
      if (closingWS) clearInterval(loadMonitor);

      if (this.#draining) {
        if (currentStatus !== WorkerStatus.WS_FULL) {
          currentStatus = WorkerStatus.WS_FULL;
          this.event.emit(
            'worker_msg',
            new WorkerMessage({
              message: {
                case: 'updateWorker',
                value: {
                  load: 1,
                  status: WorkerStatus.WS_FULL,
                },
              },
            }),
          );
        }
        return;
      }

      const oldStatus = currentStatus;
      this.#opts
        .loadFunc(this)
        .then((currentLoad: number) => {
          const isFull = currentLoad >= this.#opts.loadThreshold;
          const currentlyAvailable = !isFull;
          currentStatus = currentlyAvailable ? WorkerStatus.WS_AVAILABLE : WorkerStatus.WS_FULL;

          if (oldStatus != currentStatus) {
            const extra = { load: currentLoad, loadThreshold: this.#opts.loadThreshold };
            if (isFull) {
              this.#logger.child(extra).info('worker is at full capacity, marking as unavailable');
            } else {
              this.#logger.child(extra).info('worker is below capacity, marking as available');
            }
          }

          this.event.emit(
            'worker_msg',
            new WorkerMessage({
              message: {
                case: 'updateWorker',
                value: {
                  load: currentLoad,
                  status: currentStatus,
                },
              },
            }),
          );
        })
        .catch((e) => {
          this.#logger.warn({ error: e }, 'failed to measure CPU load');
        });
    }, UPDATE_LOAD_INTERVAL);

    await close;
    ws.removeAllListeners();
  }

  async #availability(msg: AvailabilityRequest) {
    let answered = false;

    const onReject = async () => {
      answered = true;
      this.event.emit(
        'worker_msg',
        new WorkerMessage({
          message: {
            case: 'availability',
            value: {
              jobId: msg.job!.id,
              available: false,
            },
          },
        }),
      );
    };

    const onAccept = async (args: JobAcceptArguments) => {
      answered = true;

      this.event.emit(
        'worker_msg',
        new WorkerMessage({
          message: {
            case: 'availability',
            value: {
              jobId: msg.job!.id,
              available: true,
              participantIdentity: args.identity,
              participantName: args.name,
              participantMetadata: args.metadata,
              participantAttributes: args.attributes,
            },
          },
        }),
      );

      this.#pending[req.id] = new PendingAssignment();

      const timer = setTimeout(() => {
        this.#logger.child({ req }).warn(`assignment for job ${req.id} timed out`);
        return;
      }, ASSIGNMENT_TIMEOUT);
      const asgn = await this.#pending[req.id]?.promise.then(async (asgn) => {
        clearTimeout(timer);
        return asgn;
      });

      if (asgn) {
        try {
          await this.#procPool.launchJob({
            acceptArguments: args,
            job: msg.job!,
            url: asgn.url || this.#opts.wsURL,
            token: asgn.token,
            workerId: this.id,
          });
        } catch (e) {
          this.#logger.child({ requestId: req.id }).error(e, 'error launching job');
        }
      } else {
        this.#logger.child({ requestId: req.id }).warn('pending assignment not found');
      }
    };

    const req = new JobRequest(msg.job!, onReject, onAccept);
    this.#logger
      .child({ jobId: msg.job?.id, resuming: msg.resuming, agentName: this.#opts.agentName })
      .info('received job request');

    if (this.#draining) {
      this.#logger
        .child({ jobId: msg.job?.id, resuming: msg.resuming, agentName: this.#opts.agentName })
        .info('Worker is draining and no longer available, rejecting job');
      await req.reject();
      return;
    }

    const jobRequestTask = async () => {
      try {
        await this.#opts.requestFunc(req);
      } catch (e) {
        this.#logger
          .child({ job: msg.job, resuming: msg.resuming, agentName: this.#opts.agentName })
          .info('jobRequestFunc failed');
        await onReject();
      }

      if (!answered) {
        this.#logger
          .child({ job: msg.job, resuming: msg.resuming, agentName: this.#opts.agentName })
          .info('no answer was given inside the jobRequestFunc, automatically rejecting the job');
      }
    };

    const task = jobRequestTask();
    this.#tasks.push(task);
    task.finally(() => {
      const taskIndex = this.#tasks.indexOf(task);
      if (taskIndex !== -1) {
        this.#tasks.splice(taskIndex, 1);
      } else {
        throw new Error(`task ${task} not found in tasks`);
      }
    });
  }

  async #termination(msg: JobTermination) {
    const proc = this.#procPool.getByJobId(msg.jobId);
    if (proc === null) {
      // safe to ignore
      return;
    }
    await proc.close().catch((e) => this.#logger.error(e, 'Error terminating job'));
  }

  async close() {
    if (this.#closed) {
      await this.#close.await;
      return;
    }

    this.#logger.debug('shutting down worker');

    this.#closed = true;

    await this.#inferenceExecutor?.close();
    await this.#procPool.close();
    await this.#httpServer.close();
    await Promise.allSettled(this.#tasks);

    this.#session?.close();
    await this.#close.await;
  }
}

/**
 * @deprecated Use {@link AgentServer} instead. This alias is provided for backward compatibility.
 */
export const Worker = AgentServer;

/**
 * @deprecated Use {@link ServerOptions} instead. This alias is provided for backward compatibility.
 */
export const WorkerOptions = ServerOptions;