@livekit/agents 1.0.27 → 1.0.30

package/src/connection_pool.test.ts

@@ -0,0 +1,346 @@
+// SPDX-FileCopyrightText: 2025 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { describe, expect, it, vi } from 'vitest';
+import { ConnectionPool } from './connection_pool.js';
+
+describe('ConnectionPool', () => {
+  const makeConnectCb = () => {
+    let n = 0;
+    return vi.fn(async (_timeout: number): Promise<string> => `conn_${++n}`);
+  };
+
+  describe('basic operations', () => {
+    it('should create and return a connection', async () => {
+      const connections: string[] = [];
+      const connectCb = vi.fn(async (_timeout: number): Promise<string> => {
+        const conn = `conn_${connections.length}`;
+        connections.push(conn);
+        return conn;
+      });
+      const closeCb = vi.fn(async (_conn: string) => {
+        // Mock close
+      });
+
+      const pool = new ConnectionPool<string>({
+        connectCb,
+        closeCb,
+      });
+
+      const conn = await pool.get();
+      expect(conn).toBe('conn_0');
+      expect(connectCb).toHaveBeenCalledTimes(1);
+
+      pool.put(conn);
+      const conn2 = await pool.get();
+      expect(conn2).toBe('conn_0'); // Should reuse
+      expect(connectCb).toHaveBeenCalledTimes(1);
+    });
+
+    it('should create new connection when none available', async () => {
+      const connectCb = makeConnectCb();
+      const closeCb = vi.fn(async (_conn: string) => {
+        // Mock close
+      });
+
+      const pool = new ConnectionPool<string>({
+        connectCb,
+        closeCb,
+      });
+
+      const conn1 = await pool.get();
+      pool.put(conn1);
+      const conn2 = await pool.get();
+      expect(conn1).toBe(conn2); // Should reuse
+      expect(connectCb).toHaveBeenCalledTimes(1);
+    });
+
+    it('should remove connection from pool', async () => {
+      const connectCb = makeConnectCb();
+      const closeCb = vi.fn(async (_conn: string) => {
+        // Mock close
+      });
+
+      const pool = new ConnectionPool<string>({
+        connectCb,
+        closeCb,
+      });
+
+      const conn = await pool.get();
+      pool.put(conn);
+      pool.remove(conn);
+
+      const conn2 = await pool.get();
+      expect(conn2).not.toBe(conn); // Should create new connection
+      expect(connectCb).toHaveBeenCalledTimes(2);
+      expect(closeCb).toHaveBeenCalledTimes(1);
+    });
+  });
+
+  describe('maxSessionDuration', () => {
+    it('should expire connections after maxSessionDuration', async () => {
+      const connectCb = makeConnectCb();
+      const closeCb = vi.fn(async (_conn: string) => {
+        // Mock close
+      });
+
+      const pool = new ConnectionPool<string>({
+        connectCb,
+        closeCb,
+        maxSessionDuration: 100, // 100ms
+      });
+
+      const conn1 = await pool.get();
+      pool.put(conn1);
+
+      // Wait for expiration
+      await new Promise((resolve) => setTimeout(resolve, 150));
+
+      const conn2 = await pool.get();
+      expect(conn2).not.toBe(conn1); // Should create new connection
+      expect(connectCb).toHaveBeenCalledTimes(2);
+      expect(closeCb).toHaveBeenCalledTimes(1);
+    });
+
+    it('should refresh connection timestamp when markRefreshedOnGet is true', async () => {
+      const connectCb = makeConnectCb();
+      const closeCb = vi.fn(async (_conn: string) => {
+        // Mock close
+      });
+
+      const pool = new ConnectionPool<string>({
+        connectCb,
+        closeCb,
+        maxSessionDuration: 200, // 200ms
+        markRefreshedOnGet: true,
+      });
+
+      const conn1 = await pool.get();
+      pool.put(conn1);
+
+      // Wait 100ms (less than expiration)
+      await new Promise((resolve) => setTimeout(resolve, 100));
+
+      // Get again - should refresh timestamp
+      const conn2 = await pool.get();
+      expect(conn2).toBe(conn1); // Should reuse
+      pool.put(conn2);
+
+      // Wait another 100ms (total 200ms, but refreshed at 100ms)
+      await new Promise((resolve) => setTimeout(resolve, 100));
+
+      // Should still be valid
+      const conn3 = await pool.get();
+      expect(conn3).toBe(conn1); // Should still reuse
+      expect(connectCb).toHaveBeenCalledTimes(1);
+    });
+  });
+
+  describe('withConnection', () => {
+    it('should return connection to pool on success', async () => {
+      const connectCb = makeConnectCb();
+      const closeCb = vi.fn(async (_conn: string) => {
+        // Mock close
+      });
+
+      const pool = new ConnectionPool<string>({
+        connectCb,
+        closeCb,
+      });
+
+      let capturedConn: string | undefined;
+      await pool.withConnection(async (conn) => {
+        capturedConn = conn;
+        return 'result';
+      });
+
+      // Connection should be returned to pool
+      const conn2 = await pool.get();
+      expect(conn2).toBe(capturedConn); // Should reuse
+      expect(connectCb).toHaveBeenCalledTimes(1);
+    });
+
+    it('should remove connection from pool on error', async () => {
+      const connectCb = makeConnectCb();
+      const closeCb = vi.fn(async (_conn: string) => {
+        // Mock close
+      });
+
+      const pool = new ConnectionPool<string>({
+        connectCb,
+        closeCb,
+      });
+
+      let capturedConn: string | undefined;
+      try {
+        await pool.withConnection(async (conn) => {
+          capturedConn = conn;
+          throw new Error('test error');
+        });
+      } catch (e) {
+        // Expected
+      }
+
+      // Connection should be removed from pool
+      const conn2 = await pool.get();
+      expect(conn2).not.toBe(capturedConn); // Should create new connection
+      expect(connectCb).toHaveBeenCalledTimes(2);
+      expect(closeCb).toHaveBeenCalledTimes(1);
+    });
+
+    it('should handle abort signal', async () => {
+      const connectCb = makeConnectCb();
+      const closeCb = vi.fn(async (_conn: string) => {
+        // Mock close
+      });
+
+      const pool = new ConnectionPool<string>({
+        connectCb,
+        closeCb,
+      });
+
+      const abortController = new AbortController();
+      let capturedConn: string | undefined;
+
+      const promise = pool.withConnection(
+        async (conn) => {
+          capturedConn = conn;
+          await new Promise((resolve) => setTimeout(resolve, 1000));
+          return 'result';
+        },
+        { signal: abortController.signal },
+      );
+
+      // Abort after a short delay
+      setTimeout(() => abortController.abort(), 10);
+
+      await expect(promise).rejects.toThrow();
+
+      // Connection should be removed from pool
+      const conn2 = await pool.get();
+      expect(conn2).not.toBe(capturedConn); // Should create new connection
+      expect(closeCb).toHaveBeenCalledTimes(1);
+    });
+  });
+
+  describe('prewarm', () => {
+    it('should create connection in background', async () => {
+      let n = 0;
+      const connectCb = vi.fn(async (_timeout: number): Promise<string> => {
+        await new Promise((resolve) => setTimeout(resolve, 50));
+        return `conn_${++n}`;
+      });
+      const closeCb = vi.fn(async (_conn: string) => {
+        // Mock close
+      });
+
+      const pool = new ConnectionPool<string>({
+        connectCb,
+        closeCb,
+      });
+
+      pool.prewarm();
+
+      // Wait for prewarm to complete
+      await new Promise((resolve) => setTimeout(resolve, 100));
+
+      const conn = await pool.get();
+      expect(conn).toBeDefined();
+      expect(connectCb).toHaveBeenCalledTimes(1);
+    });
+
+    it('should not prewarm if connections already exist', async () => {
+      const connectCb = makeConnectCb();
+      const closeCb = vi.fn(async (_conn: string) => {
+        // Mock close
+      });
+
+      const pool = new ConnectionPool<string>({
+        connectCb,
+        closeCb,
+      });
+
+      // Create a connection first
+      const conn1 = await pool.get();
+      pool.put(conn1);
+
+      pool.prewarm(); // Should not create new connection
+
+      const conn2 = await pool.get();
+      expect(conn2).toBe(conn1); // Should reuse existing
+      expect(connectCb).toHaveBeenCalledTimes(1);
+    });
+  });
+
+  describe('close', () => {
+    it('should close all connections', async () => {
+      const connectCb = makeConnectCb();
+      const closeCb = vi.fn(async (_conn: string) => {
+        // Mock close
+      });
+
+      const pool = new ConnectionPool<string>({
+        connectCb,
+        closeCb,
+      });
+
+      // Create two distinct connections by checking out both before returning either.
+      const conn1 = await pool.get();
+      const conn2 = await pool.get();
+      pool.put(conn1);
+      pool.put(conn2);
+
+      await pool.close();
+
+      expect(closeCb).toHaveBeenCalledTimes(2);
+    });
+
+    it('should invalidate all connections', async () => {
+      const connectCb = makeConnectCb();
+      const closeCb = vi.fn(async (_conn: string) => {
+        // Mock close
+      });
+
+      const pool = new ConnectionPool<string>({
+        connectCb,
+        closeCb,
+      });
+
+      // Create two distinct connections by checking out both before returning either.
+      const conn1 = await pool.get();
+      const conn2 = await pool.get();
+      pool.put(conn1);
+      pool.put(conn2);
+
+      pool.invalidate();
+      await pool.close(); // Drain to close
+
+      expect(closeCb).toHaveBeenCalledTimes(2);
+    });
+  });
+
+  describe('concurrent access', () => {
+    it('should handle concurrent get requests', async () => {
+      const connectCb = vi.fn(async (_timeout: number): Promise<string> => {
+        await new Promise((resolve) => setTimeout(resolve, 10));
+        return `conn_${Date.now()}_${Math.random()}`;
+      });
+      const closeCb = vi.fn(async (_conn: string) => {
+        // Mock close
+      });
+
+      const pool = new ConnectionPool<string>({
+        connectCb,
+        closeCb,
+      });
+
+      const promises = Array.from({ length: 5 }, () => pool.get());
+      const connections = await Promise.all(promises);
+
+      // All should be different connections
+      const uniqueConnections = new Set(connections);
+      expect(uniqueConnections.size).toBe(5);
+      expect(connectCb).toHaveBeenCalledTimes(5);
+    });
+  });
+});

package/src/connection_pool.ts

@@ -0,0 +1,307 @@
+// SPDX-FileCopyrightText: 2025 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { Mutex } from '@livekit/mutex';
+import { waitForAbort } from './utils.js';
+
+/**
+ * Helper class to manage persistent connections like websockets.
+ */
+export interface ConnectionPoolOptions<T> {
+  /**
+   * Maximum duration in milliseconds before forcing reconnection.
+   * If not set, connections will never expire based on duration.
+   */
+  maxSessionDuration?: number;
+
+  /**
+   * If true, the session will be marked as fresh when get() is called.
+   * Only used when maxSessionDuration is set.
+   */
+  markRefreshedOnGet?: boolean;
+
+  /**
+   * Async callback to create new connections.
+   * @param timeout - Connection timeout in milliseconds
+   * @returns A new connection object
+   */
+  connectCb: (timeout: number) => Promise<T>;
+
+  /**
+   * Optional async callback to close connections.
+   * @param conn - The connection to close
+   */
+  closeCb?: (conn: T) => Promise<void>;
+
+  /**
+   * Default connection timeout in milliseconds.
+   * Defaults to 10000 (10 seconds).
+   */
+  connectTimeout?: number;
+}
+
+/**
+ * Connection pool for managing persistent WebSocket connections.
+ *
+ * Reuses connections efficiently and automatically refreshes them after max duration.
+ * Prevents creating too many connections in a single conversation.
+ */
+export class ConnectionPool<T> {
+  private readonly maxSessionDuration?: number;
+  private readonly markRefreshedOnGet: boolean;
+  private readonly connectCb: (timeout: number) => Promise<T>;
+  private readonly closeCb?: (conn: T) => Promise<void>;
+  private readonly connectTimeout: number;
+
+  // Track connections and their creation timestamps
+  private readonly connections: Map<T, number> = new Map();
+  // Available connections ready for reuse
+  private readonly available: Set<T> = new Set();
+  // Connections queued for closing
+  private readonly toClose: Set<T> = new Set();
+  // Mutex for connection operations
+  private readonly connectLock = new Mutex();
+  // Prewarm task reference
+  private prewarmController?: AbortController;
+
+  constructor(options: ConnectionPoolOptions<T>) {
+    this.maxSessionDuration = options.maxSessionDuration;
+    this.markRefreshedOnGet = options.markRefreshedOnGet ?? false;
+    this.connectCb = options.connectCb;
+    this.closeCb = options.closeCb;
+    this.connectTimeout = options.connectTimeout ?? 10_000;
+  }
+
+  /**
+   * Create a new connection.
+   *
+   * @param timeout - Connection timeout in milliseconds
+   * @returns The new connection object
+   * @throws If connectCb is not provided or connection fails
+   */
+  private async _connect(timeout: number): Promise<T> {
+    const connection = await this.connectCb(timeout);
+    this.connections.set(connection, Date.now());
+    return connection;
+  }
+
+  /**
+   * Drain and close all connections queued for closing.
+   */
+  private async _drainToClose(): Promise<void> {
+    const connectionsToClose = Array.from(this.toClose);
+    this.toClose.clear();
+
+    for (const conn of connectionsToClose) {
+      await this._maybeCloseConnection(conn);
+    }
+  }
+
+  /**
+   * Close a connection if closeCb is provided.
+   *
+   * @param conn - The connection to close
+   */
+  private async _maybeCloseConnection(conn: T): Promise<void> {
+    if (this.closeCb) {
+      await this.closeCb(conn);
+    }
+  }
+
+  private _abortError(): Error {
+    const error = new Error('The operation was aborted.');
+    error.name = 'AbortError';
+    return error;
+  }
+
+  /**
+   * Get an available connection or create a new one if needed.
+   *
+   * @param timeout - Connection timeout in milliseconds
+   * @returns An active connection object
+   */
+  async get(timeout?: number): Promise<T> {
+    const unlock = await this.connectLock.lock();
+    try {
+      await this._drainToClose();
+      const now = Date.now();
+
+      // Try to reuse an available connection that hasn't expired
+      while (this.available.size > 0) {
+        const conn = this.available.values().next().value as T;
+        this.available.delete(conn);
+
+        if (
+          this.maxSessionDuration === undefined ||
+          now - (this.connections.get(conn) ?? 0) <= this.maxSessionDuration
+        ) {
+          if (this.markRefreshedOnGet) {
+            this.connections.set(conn, now);
+          }
+          return conn;
+        }
+
+        // Connection expired; close it now so callers observing get() see it closed promptly.
+        // (Also makes tests deterministic: closeCb should have been called by the time get() resolves.)
+        if (this.connections.has(conn)) {
+          this.connections.delete(conn);
+        }
+        this.toClose.delete(conn);
+        await this._maybeCloseConnection(conn);
+      }
+
+      return await this._connect(timeout ?? this.connectTimeout);
+    } finally {
+      unlock();
+    }
+  }
+
+  /**
+   * Mark a connection as available for reuse.
+   *
+   * If connection has been removed, it will not be added to the pool.
+   *
+   * @param conn - The connection to make available
+   */
+  put(conn: T): void {
+    if (this.connections.has(conn)) {
+      this.available.add(conn);
+      return;
+    }
+  }
+
+  /**
+   * Remove a specific connection from the pool.
+   *
+   * Marks the connection to be closed during the next drain cycle.
+   *
+   * @param conn - The connection to remove
+   */
+  remove(conn: T): void {
+    this.available.delete(conn);
+    if (this.connections.has(conn)) {
+      this.toClose.add(conn);
+      this.connections.delete(conn);
+      // Important for Node websockets: if we just "mark to close later" but remove listeners,
+      // the ws library can buffer incoming frames in memory. Close ASAP in background.
+      void (async () => {
+        const unlock = await this.connectLock.lock();
+        try {
+          if (!this.toClose.has(conn)) return;
+          await this._maybeCloseConnection(conn);
+          this.toClose.delete(conn);
+        } finally {
+          unlock();
+        }
+      })();
+    }
+  }
+
+  /**
+   * Clear all existing connections.
+   *
+   * Marks all current connections to be closed during the next drain cycle.
+   */
+  invalidate(): void {
+    for (const conn of this.connections.keys()) {
+      this.toClose.add(conn);
+    }
+    this.connections.clear();
+    this.available.clear();
+  }
+
+  /**
+   * Initiate prewarming of the connection pool without blocking.
+   *
+   * This method starts a background task that creates a new connection if none exist.
+   * The task automatically cleans itself up when the connection pool is closed.
+   */
+  prewarm(): void {
+    if (this.prewarmController || this.connections.size > 0) {
+      return;
+    }
+
+    const controller = new AbortController();
+    this.prewarmController = controller;
+
+    // Start prewarm in background
+    this._prewarmImpl(controller.signal).catch(() => {
+      // Ignore errors during prewarm
+    });
+  }
+
+  private async _prewarmImpl(signal: AbortSignal): Promise<void> {
+    const unlock = await this.connectLock.lock();
+    try {
+      if (signal.aborted) {
+        return;
+      }
+
+      if (this.connections.size === 0) {
+        const conn = await this._connect(this.connectTimeout);
+        this.available.add(conn);
+      }
+    } finally {
+      unlock();
+    }
+  }
+
+  /**
+   * Get a connection from the pool and automatically return it when done.
+   * Handles abort signals and ensures proper cleanup.
+   *
+   * @param fn - Function to execute with the connection
+   * @param options - Options including timeout and abort signal
+   * @returns The result of the function
+   */
+  async withConnection<R>(
+    fn: (conn: T) => Promise<R>,
+    options?: {
+      timeout?: number;
+      signal?: AbortSignal;
+    },
+  ): Promise<R> {
+    // Check if already aborted before getting connection
+    if (options?.signal?.aborted) {
+      throw this._abortError();
+    }
+
+    const conn = await this.get(options?.timeout);
+
+    const signal = options?.signal;
+
+    try {
+      const fnPromise = fn(conn);
+      const result = signal
+        ? await Promise.race([
+            fnPromise.then((value) => ({ type: 'result' as const, value })),
+            waitForAbort(signal).then(() => ({ type: 'abort' as const })),
+          ]).then((r) => {
+            if (r.type === 'abort') throw this._abortError();
+            return r.value;
+          })
+        : await fnPromise;
+      // Return connection to pool on success
+      this.put(conn);
+      return result;
+    } catch (error) {
+      // Remove connection from pool on error (don't return it)
+      this.remove(conn);
+      throw error;
+    }
+  }
+
+  /**
+   * Close all connections, draining any pending connection closures.
+   */
+  async close(): Promise<void> {
+    // Cancel prewarm task if running
+    if (this.prewarmController) {
+      this.prewarmController.abort();
+      this.prewarmController = undefined;
+    }
+
+    this.invalidate();
+    await this._drainToClose();
+  }
+}

package/src/index.ts

@@ -23,6 +23,7 @@ import * as voice from './voice/index.js';
 
 export * from './_exceptions.js';
 export * from './audio.js';
+export * from './connection_pool.js';
 export * from './generator.js';
 export * from './inference_runner.js';
 export * from './job.js';