Adds in support for attachments and attributes fields into data streams (#1283)

Author: 1egoman

.changeset/deep-pugs-lick.md

@@ -0,0 +1,5 @@
+---
+'@livekit/components-core': patch
+---
+
+Adds in support for attachments and attributes fields into data streams

packages/core/src/components/chat.ts

@@ -1,7 +1,18 @@
 /* eslint-disable camelcase */
 import type { Room, SendTextOptions } from 'livekit-client';
 import { compareVersions, RoomEvent } from 'livekit-client';
-import { BehaviorSubject, Subject, scan, map, takeUntil, from, filter } from 'rxjs';
+import {
+  BehaviorSubject,
+  Subject,
+  scan,
+  map,
+  takeUntil,
+  from,
+  filter,
+  mergeMap,
+  finalize,
+  of,
+} from 'rxjs';
 import {
   DataTopic,
   LegacyDataTopic,
@@ -10,6 +21,7 @@ import {
 } from '../observables/dataChannel';
 import { log } from '../logger';
 import { ChatMessage, ReceivedChatMessage } from '../messages/types';
+import { Future } from '../helper/future';
 
 /** @public */
 export type { ChatMessage, ReceivedChatMessage };
@@ -44,6 +56,19 @@ export type ChatOptions = {
 };
 
 const topicSubjectMap: WeakMap<Room, Map<string, Subject<ReceivedChatMessage>>> = new WeakMap();
+const streamIdToAttachments = new Map<
+  string /* stream id */,
+  Map<
+    string /* attachment id */,
+    Future<
+      {
+        fileName: string;
+        buffer: Array<Uint8Array>;
+      },
+      never
+    >
+  >
+>();
 
 function isIgnorableChatMessage(msg: ReceivedChatMessage | LegacyReceivedChatMessage) {
   return (msg as LegacyChatMessage).ignoreLegacy == true;
@@ -79,26 +104,82 @@ export function setupChat(room: Room, options?: ChatOptions) {
   const finalMessageDecoder = options?.messageDecoder ?? decodeLegacyMsg;
   if (needsSetup) {
     room.registerTextStreamHandler(topic, async (reader, participantInfo) => {
-      const { id, timestamp } = reader.info;
+      const { id, timestamp, attributes, attachedStreamIds } = reader.info;
+
+      // Store a future for each attachment to be later resolved once the corresponding file data
+      // stream completes.
+      const attachments = new Map(
+        (attachedStreamIds ?? []).map((id) => [
+          id,
+          new Future<{ fileName: string; buffer: Array<Uint8Array> }, never>(),
+        ]),
+      );
+      streamIdToAttachments.set(id, attachments);
+
       const streamObservable = from(reader).pipe(
         scan((acc: string, chunk: string) => {
           return acc + chunk;
         }),
-        map((chunk: string) => {
+        mergeMap((chunk: string) => {
+          if (attachments.size === 0) {
+            return of({ chunk, attachedFiles: [] });
+          } else {
+            // Aggregate all attachments into memory and transform them into a list of files
+            return from(attachments.values()).pipe(
+              mergeMap((attachment) => from(attachment.promise)),
+              scan(
+                (acc, attachment) => [...acc, new File(attachment.buffer, attachment.fileName)],
+                [] as Array<File>,
+              ),
+              map((attachedFiles) => ({ chunk, attachedFiles })),
+            );
+          }
+        }),
+        map(({ chunk, attachedFiles }) => {
           return {
             id,
             timestamp,
             message: chunk,
             from: room.getParticipantByIdentity(participantInfo.identity),
             type: 'chatMessage',
+            attributes,
+            attachedFiles,
             // editTimestamp: type === 'update' ? timestamp : undefined,
           } satisfies ReceivedChatMessage;
         }),
+        finalize(() => streamIdToAttachments.delete(id)),
       );
       streamObservable.subscribe({
         next: (value) => messageSubject.next(value),
       });
     });
+    // NOTE: Attachment byte streams are guaranteed to arrive after their parent text stream
+    // has initialized the attachment map (per client SDK sending implementation)
+    room.registerByteStreamHandler(topic, async (reader) => {
+      const { id: attachmentStreamId } = reader.info;
+      const foundStreamAttachmentPair = Array.from(streamIdToAttachments).find(([, attachments]) =>
+        attachments.has(attachmentStreamId),
+      );
+      if (!foundStreamAttachmentPair) {
+        return;
+      }
+      const streamId = foundStreamAttachmentPair[0];
+
+      const bufferList = [];
+      for await (const buffer of reader) {
+        bufferList.push(buffer);
+      }
+
+      const attachment = streamIdToAttachments.get(streamId)?.get(attachmentStreamId);
+      if (!attachment) {
+        return;
+      }
+
+      attachment.resolve?.({
+        fileName: reader.info.name,
+        buffer: bufferList,
+      });
+    });
 
     /** legacy chat protocol handling */
     const { messageObservable } = setupDataMessageHandler(room, [legacyTopic]);
@@ -204,6 +285,7 @@ export function setupChat(room: Room, options?: ChatOptions) {
     messageSubject.complete();
     topicSubjectMap.delete(room);
     room.unregisterTextStreamHandler(topic);
+    room.unregisterByteStreamHandler(topic);
   }
   room.once(RoomEvent.Disconnected, destroy);
 

packages/core/src/helper/future.ts

@@ -0,0 +1,32 @@
+export class Future<T, E extends Error> {
+  promise: Promise<T>;
+
+  resolve?: (arg: T) => void;
+
+  reject?: (e: E) => void;
+
+  onFinally?: () => void;
+
+  get isResolved(): boolean {
+    return this._isResolved;
+  }
+
+  private _isResolved: boolean = false;
+
+  constructor(
+    futureBase?: (resolve: (arg: T) => void, reject: (e: E) => void) => void,
+    onFinally?: () => void,
+  ) {
+    this.onFinally = onFinally;
+    this.promise = new Promise<T>(async (resolve, reject) => {
+      this.resolve = resolve;
+      this.reject = reject;
+      if (futureBase) {
+        await futureBase(resolve, reject);
+      }
+    }).finally(() => {
+      this._isResolved = true;
+      this.onFinally?.();
+    });
+  }
+}