@livekit/agents 0.0.0-20260120144724

package/src/utils/ws_transport.test.ts

@@ -0,0 +1,282 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { describe, expect, it } from 'vitest';
+import { WebSocket, WebSocketServer } from 'ws';
+import { webSocketStream } from './ws_transport.js';
+
+describe('webSocketStream', () => {
+  describe('readable stream', () => {
+    it('receives messages from the WebSocket', async () => {
+      const wss = await new Promise<WebSocketServer>((resolve) => {
+        const server: WebSocketServer = new WebSocketServer({ port: 0 }, () => resolve(server));
+      });
+
+      const port = (wss.address() as { port: number }).port;
+
+      wss.on('connection', (serverWs) => {
+        serverWs.send('hello');
+        serverWs.send('world');
+        serverWs.close();
+      });
+
+      const { readable } = webSocketStream(`ws://localhost:${port}`);
+      const reader = readable.getReader();
+
+      const messages: string[] = [];
+      try {
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          messages.push(Buffer.from(value).toString());
+        }
+      } finally {
+        reader.releaseLock();
+      }
+
+      expect(messages).toEqual(['hello', 'world']);
+
+      wss.close();
+    });
+
+    it('handles binary messages', async () => {
+      const wss = await new Promise<WebSocketServer>((resolve) => {
+        const server: WebSocketServer = new WebSocketServer({ port: 0 }, () => resolve(server));
+      });
+
+      const port = (wss.address() as { port: number }).port;
+
+      const binaryData = new Uint8Array([1, 2, 3, 4, 5]);
+
+      wss.on('connection', (serverWs) => {
+        serverWs.send(binaryData);
+        serverWs.close();
+      });
+
+      const { readable } = webSocketStream(`ws://localhost:${port}`);
+      const reader = readable.getReader();
+
+      const chunks: Uint8Array[] = [];
+      try {
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          chunks.push(new Uint8Array(value));
+        }
+      } finally {
+        reader.releaseLock();
+      }
+
+      expect(chunks).toHaveLength(1);
+      expect(Array.from(chunks[0]!)).toEqual([1, 2, 3, 4, 5]);
+
+      wss.close();
+    });
+
+    it('handles empty stream when connection closes immediately', async () => {
+      const wss = await new Promise<WebSocketServer>((resolve) => {
+        const server: WebSocketServer = new WebSocketServer({ port: 0 }, () => resolve(server));
+      });
+
+      const port = (wss.address() as { port: number }).port;
+
+      wss.on('connection', (serverWs) => {
+        serverWs.close();
+      });
+      const { readable } = webSocketStream(`ws://localhost:${port}`);
+      const reader = readable.getReader();
+
+      const chunks: Uint8Array[] = [];
+      try {
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          chunks.push(value);
+        }
+      } finally {
+        reader.releaseLock();
+      }
+
+      expect(chunks).toEqual([]);
+
+      wss.close();
+    });
+  });
+
+  describe('writable stream', () => {
+    it('sends messages through the WebSocket', async () => {
+      const wss = await new Promise<WebSocketServer>((resolve) => {
+        const server: WebSocketServer = new WebSocketServer({ port: 0 }, () => resolve(server));
+      });
+
+      const port = (wss.address() as { port: number }).port;
+      const ws = new WebSocket(`ws://localhost:${port}`);
+
+      const connected = new Promise<void>((resolve) => {
+        ws.on('open', resolve);
+      });
+
+      const messagesReceived: string[] = [];
+      const serverClosed = new Promise<void>((resolve) => {
+        wss.on('connection', (serverWs) => {
+          serverWs.on('message', (data) => {
+            messagesReceived.push(data.toString());
+          });
+          serverWs.on('close', resolve);
+        });
+      });
+
+      await connected;
+      const { writable } = webSocketStream(`ws://localhost:${port}`);
+      const writer = writable.getWriter();
+
+      await writer.write(new TextEncoder().encode('hello'));
+      await writer.write(new TextEncoder().encode('world'));
+      await writer.close();
+
+      await serverClosed;
+
+      expect(messagesReceived).toEqual(['hello', 'world']);
+
+      wss.close();
+    });
+
+    it('sends binary data through the WebSocket', async () => {
+      const wss = await new Promise<WebSocketServer>((resolve) => {
+        const server: WebSocketServer = new WebSocketServer({ port: 0 }, () => resolve(server));
+      });
+
+      const port = (wss.address() as { port: number }).port;
+
+      const chunksReceived: Buffer[] = [];
+      const serverClosed = new Promise<void>((resolve) => {
+        wss.on('connection', (serverWs) => {
+          serverWs.on('message', (data) => {
+            chunksReceived.push(Buffer.from(data as Buffer));
+          });
+          serverWs.on('close', resolve);
+        });
+      });
+
+      const { writable } = webSocketStream(`ws://localhost:${port}`);
+      const writer = writable.getWriter();
+
+      const binaryData = new Uint8Array([10, 20, 30, 40, 50]);
+      await writer.write(binaryData);
+      await writer.close();
+
+      await serverClosed;
+
+      expect(chunksReceived).toHaveLength(1);
+      expect(Array.from(chunksReceived[0]!)).toEqual([10, 20, 30, 40, 50]);
+
+      wss.close();
+    });
+
+    it('buffers writes if readyState is CONNECTING', async () => {
+      const wss = await new Promise<WebSocketServer>((resolve) => {
+        const server: WebSocketServer = new WebSocketServer({ port: 0 }, () => resolve(server));
+      });
+
+      const port = (wss.address() as { port: number }).port;
+
+      const { writable } = webSocketStream(`ws://localhost:${port}`);
+      const writer = writable.getWriter();
+
+      const messagesReceived: string[] = [];
+      const serverClosed = new Promise<void>((resolve) => {
+        wss.on('connection', (serverWs) => {
+          serverWs.on('message', (data) => {
+            messagesReceived.push(data.toString());
+          });
+          serverWs.on('close', resolve);
+        });
+      });
+
+      // These writes should be buffered
+      await writer.write(new TextEncoder().encode('buffered message'));
+      await writer.close();
+
+      await serverClosed;
+
+      expect(messagesReceived).toEqual(['buffered message']);
+
+      wss.close();
+    });
+  });
+
+  describe('bidirectional communication', () => {
+    it('supports echo pattern with readable and writable', async () => {
+      const wss = await new Promise<WebSocketServer>((resolve) => {
+        const server: WebSocketServer = new WebSocketServer({ port: 0 }, () => resolve(server));
+      });
+
+      const port = (wss.address() as { port: number }).port;
+
+      // Server echoes messages back
+      wss.on('connection', (serverWs) => {
+        serverWs.on('message', (data) => {
+          serverWs.send(data);
+        });
+      });
+
+      const { readable, writable } = webSocketStream(`ws://localhost:${port}`);
+      const writer = writable.getWriter();
+      const reader = readable.getReader();
+
+      // Send messages
+      await writer.write(new TextEncoder().encode('ping1'));
+      await writer.write(new TextEncoder().encode('ping2'));
+
+      // Read echoed responses
+      const { value: response1 } = await reader.read();
+      const { value: response2 } = await reader.read();
+
+      expect(Buffer.from(response1!).toString()).toBe('ping1');
+      expect(Buffer.from(response2!).toString()).toBe('ping2');
+
+      reader.releaseLock();
+      await writer.close();
+
+      wss.close();
+    });
+  });
+
+  describe('error handling', () => {
+    it('readable stream ends when WebSocket closes unexpectedly', async () => {
+      const wss = await new Promise<WebSocketServer>((resolve) => {
+        const server: WebSocketServer = new WebSocketServer({ port: 0 }, () => resolve(server));
+      });
+
+      const port = (wss.address() as { port: number }).port;
+
+      wss.on('connection', (serverWs) => {
+        serverWs.send('before close');
+        // Terminate connection abruptly
+        serverWs.terminate();
+      });
+
+      const { readable } = webSocketStream(`ws://localhost:${port}`);
+      const reader = readable.getReader();
+
+      const chunks: string[] = [];
+      try {
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          chunks.push(Buffer.from(value).toString());
+        }
+      } catch (error) {
+        console.error(error);
+        // Connection terminated, stream may error
+      } finally {
+        reader.releaseLock();
+      }
+
+      // Should have received the message sent before termination
+      expect(chunks).toContain('before close');
+
+      wss.close();
+    });
+  });
+});

package/src/utils/ws_transport.ts

@@ -0,0 +1,22 @@
+import { Readable, Writable } from 'node:stream';
+import WebSocket, { createWebSocketStream } from 'ws';
+
+export function webSocketStream(wsUrl: string) {
+  const ws = new WebSocket(wsUrl);
+  const duplex = createWebSocketStream(ws);
+  duplex.on('error', console.error);
+
+  // End the write side when the read side ends to properly close the stream.
+  // This is needed because Readable.toWeb() waits for both sides of the duplex
+  // to close before signaling done on the ReadableStream.
+  duplex.on('end', () => {
+    duplex.end();
+  });
+
+  // Convert the writable side
+  const writable = Writable.toWeb(duplex);
+  // Convert the readable side
+  const readable = Readable.toWeb(duplex);
+
+  return { readable, writable, close: ws.close };
+}