@livekit/agents 1.0.2 → 1.0.4

package/src/voice/avatar/datastream_io.ts

@@ -0,0 +1,247 @@
+// SPDX-FileCopyrightText: 2025 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { Mutex } from '@livekit/mutex';
+import {
+  type AudioFrame,
+  type ByteStreamWriter,
+  type Room,
+  RoomEvent,
+  type RpcInvocationData,
+  type TrackKind,
+} from '@livekit/rtc-node';
+import { log } from '../../log.js';
+import {
+  Future,
+  Task,
+  shortuuid,
+  waitForParticipant,
+  waitForTrackPublication,
+} from '../../utils.js';
+import { AudioOutput, type PlaybackFinishedEvent } from '../io.js';
+
+const RPC_CLEAR_BUFFER = 'lk.clear_buffer';
+const RPC_PLAYBACK_FINISHED = 'lk.playback_finished';
+const AUDIO_STREAM_TOPIC = 'lk.audio_stream';
+
+export interface DataStreamAudioOutputOptions {
+  room: Room;
+  destinationIdentity: string;
+  sampleRate?: number;
+  waitRemoteTrack?: TrackKind;
+}
+
+/**
+ * AudioOutput implementation that streams audio to a remote avatar worker using LiveKit DataStream.
+ */
+export class DataStreamAudioOutput extends AudioOutput {
+  static _playbackFinishedRpcRegistered: boolean = false;
+  static _playbackFinishedHandlers: Record<string, (data: RpcInvocationData) => string> = {};
+
+  private room: Room;
+  private destinationIdentity: string;
+  private roomConnectedFuture: Future<void>;
+  private waitRemoteTrack?: TrackKind;
+  private streamWriter?: ByteStreamWriter;
+  private pushedDuration: number = 0;
+  private started: boolean = false;
+  private lock = new Mutex();
+  private startTask?: Task<void>;
+
+  #logger = log();
+
+  constructor(opts: DataStreamAudioOutputOptions) {
+    super(opts.sampleRate, undefined);
+
+    const { room, destinationIdentity, sampleRate, waitRemoteTrack } = opts;
+    this.room = room;
+    this.destinationIdentity = destinationIdentity;
+    this.sampleRate = sampleRate;
+    this.waitRemoteTrack = waitRemoteTrack;
+
+    const onRoomConnected = async () => {
+      if (this.startTask) return;
+
+      await this.roomConnectedFuture.await;
+
+      // register the rpc method right after the room is connected
+      DataStreamAudioOutput.registerPlaybackFinishedRpc({
+        room,
+        callerIdentity: this.destinationIdentity,
+        handler: (data) => this.handlePlaybackFinished(data),
+      });
+
+      this.startTask = Task.from(({ signal }) => this._start(signal));
+    };
+
+    this.roomConnectedFuture = new Future<void>();
+
+    this.room.on(RoomEvent.ConnectionStateChanged, (_) => {
+      if (room.isConnected && !this.roomConnectedFuture.done) {
+        this.roomConnectedFuture.resolve(undefined);
+      }
+    });
+
+    if (this.room.isConnected) {
+      this.roomConnectedFuture.resolve(undefined);
+    }
+
+    onRoomConnected();
+  }
+
+  private async _start(_abortSignal: AbortSignal) {
+    const unlock = await this.lock.lock();
+
+    try {
+      if (this.started) return;
+
+      await this.roomConnectedFuture.await;
+
+      this.#logger.debug(
+        {
+          identity: this.destinationIdentity,
+        },
+        'waiting for the remote participant',
+      );
+
+      await waitForParticipant({
+        room: this.room,
+        identity: this.destinationIdentity,
+      });
+
+      if (this.waitRemoteTrack) {
+        this.#logger.debug(
+          {
+            identity: this.destinationIdentity,
+            kind: this.waitRemoteTrack,
+          },
+          'waiting for the remote track',
+        );
+
+        await waitForTrackPublication({
+          room: this.room,
+          identity: this.destinationIdentity,
+          kind: this.waitRemoteTrack,
+        });
+      }
+
+      this.#logger.debug(
+        {
+          identity: this.destinationIdentity,
+        },
+        'remote participant ready',
+      );
+
+      this.started = true;
+    } finally {
+      unlock();
+    }
+  }
+
+  async captureFrame(frame: AudioFrame): Promise<void> {
+    if (!this.startTask) {
+      this.startTask = Task.from(({ signal }) => this._start(signal));
+    }
+
+    await this.startTask.result;
+    await super.captureFrame(frame);
+
+    if (!this.streamWriter) {
+      this.streamWriter = await this.room.localParticipant!.streamBytes({
+        name: shortuuid('AUDIO_'),
+        topic: AUDIO_STREAM_TOPIC,
+        destinationIdentities: [this.destinationIdentity],
+        attributes: {
+          sample_rate: frame.sampleRate.toString(),
+          num_channels: frame.channels.toString(),
+        },
+      });
+      this.pushedDuration = 0;
+    }
+
+    // frame.data is a Int16Array, write accepts a Uint8Array
+    await this.streamWriter.write(new Uint8Array(frame.data.buffer));
+    this.pushedDuration += frame.samplesPerChannel / frame.sampleRate;
+  }
+
+  flush(): void {
+    super.flush();
+
+    if (this.streamWriter === undefined || !this.started) {
+      return;
+    }
+
+    this.streamWriter.close().finally(() => {
+      this.streamWriter = undefined;
+    });
+  }
+
+  clearBuffer(): void {
+    if (!this.started) return;
+
+    this.room.localParticipant!.performRpc({
+      destinationIdentity: this.destinationIdentity,
+      method: RPC_CLEAR_BUFFER,
+      payload: '',
+    });
+  }
+
+  private handlePlaybackFinished(data: RpcInvocationData): string {
+    if (data.callerIdentity !== this.destinationIdentity) {
+      this.#logger.warn(
+        {
+          callerIdentity: data.callerIdentity,
+          destinationIdentity: this.destinationIdentity,
+        },
+        'playback finished event received from unexpected participant',
+      );
+      return 'reject';
+    }
+
+    this.#logger.info(
+      {
+        callerIdentity: data.callerIdentity,
+      },
+      'playback finished event received',
+    );
+
+    const playbackFinishedEvent = JSON.parse(data.payload) as PlaybackFinishedEvent;
+    this.onPlaybackFinished(playbackFinishedEvent);
+    return 'ok';
+  }
+
+  static registerPlaybackFinishedRpc({
+    room,
+    callerIdentity,
+    handler,
+  }: {
+    room: Room;
+    callerIdentity: string;
+    handler: (data: RpcInvocationData) => string;
+  }) {
+    DataStreamAudioOutput._playbackFinishedHandlers[callerIdentity] = handler;
+
+    if (DataStreamAudioOutput._playbackFinishedRpcRegistered) {
+      return;
+    }
+
+    const rpcHandler = async (data: RpcInvocationData): Promise<string> => {
+      const handler = DataStreamAudioOutput._playbackFinishedHandlers[data.callerIdentity];
+      if (!handler) {
+        log().warn(
+          {
+            callerIdentity: data.callerIdentity,
+            expectedIdentities: Object.keys(DataStreamAudioOutput._playbackFinishedHandlers),
+          },
+          'playback finished event received from unexpected participant',
+        );
+
+        return 'reject';
+      }
+      return handler(data);
+    };
+
+    room.localParticipant?.registerRpcMethod(RPC_PLAYBACK_FINISHED, rpcHandler);
+    DataStreamAudioOutput._playbackFinishedRpcRegistered = true;
+  }
+}

package/src/voice/avatar/index.ts

@@ -0,0 +1,4 @@
+// SPDX-FileCopyrightText: 2025 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+export * from './datastream_io.js';

package/src/voice/index.ts

@@ -3,5 +3,7 @@
 // SPDX-License-Identifier: Apache-2.0
 export { Agent, StopResponse, type AgentOptions, type ModelSettings } from './agent.js';
 export { AgentSession, type AgentSessionOptions } from './agent_session.js';
+
+export * from './avatar/index.js';
 export * from './events.js';
 export { RunContext } from './run_context.js';

package/src/voice/io.ts

@@ -55,7 +55,7 @@ export abstract class AudioOutput extends EventEmitter {
   protected logger = log();
 
   constructor(
-    readonly sampleRate?: number,
+    public sampleRate?: number,
     protected readonly nextInChain?: AudioOutput,
   ) {
     super();

package/src/voice/room_io/_input.ts

@@ -1,6 +1,7 @@
 // SPDX-FileCopyrightText: 2025 LiveKit, Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
+import type { AudioFrame } from '@livekit/rtc-node';
 import {
   AudioStream,
   type NoiseCancellationOptions,
@@ -11,6 +12,7 @@ import {
   RoomEvent,
   TrackSource,
 } from '@livekit/rtc-node';
+import type { ReadableStream } from 'node:stream/web';
 import { log } from '../../log.js';
 import { resampleStream } from '../../utils.js';
 import { AudioInput } from '../io.js';
@@ -64,8 +66,10 @@ export class ParticipantAudioInputStream extends AudioInput {
         ? participant
         : this.room.remoteParticipants.get(participantIdentity);
 
+    // We need to check if the participant has a microphone track and subscribe to it
+    // in case we miss the tracksubscribed event
     if (participantValue) {
-      for (const publication of Object.values(participantValue.trackPublications)) {
+      for (const publication of participantValue.trackPublications.values()) {
         if (publication.track && publication.source === TrackSource.SOURCE_MICROPHONE) {
           this.onTrackSubscribed(publication.track, publication, participantValue);
           break;
@@ -127,12 +131,13 @@ export class ParticipantAudioInputStream extends AudioInput {
     return true;
   };
 
-  private createStream(track: RemoteTrack) {
+  private createStream(track: RemoteTrack): ReadableStream<AudioFrame> {
     return new AudioStream(track, {
       sampleRate: this.sampleRate,
       numChannels: this.numChannels,
       noiseCancellation: this.noiseCancellation,
-    });
+      // TODO(AJS-269): resolve compatibility issue with node-sdk to remove the forced type casting
+    }) as unknown as ReadableStream<AudioFrame>;
   }
 
   async close() {

package/src/voice/run_context.ts

@@ -8,13 +8,27 @@ import type { SpeechHandle } from './speech_handle.js';
 export type UnknownUserData = unknown;
 
 export class RunContext<UserData = UnknownUserData> {
+  private readonly initialStepIdx: number;
   constructor(
     public readonly session: AgentSession<UserData>,
     public readonly speechHandle: SpeechHandle,
     public readonly functionCall: FunctionCall,
-  ) {}
-
+  ) {
+    this.initialStepIdx = speechHandle.numSteps - 1;
+  }
   get userData(): UserData {
     return this.session.userData;
   }
+
+  /**
+   * Waits for the speech playout corresponding to this function call step.
+   *
+   * Unlike {@link SpeechHandle.waitForPlayout}, which waits for the full
+   * assistant turn to complete (including all function tools),
+   * this method only waits for the assistant's spoken response prior to running
+   * this tool to finish playing.
+   */
+  async waitForPlayout() {
+    return this.speechHandle._waitForGeneration(this.initialStepIdx);
+  }
 }

package/src/voice/speech_handle.ts

@@ -1,8 +1,10 @@
 // SPDX-FileCopyrightText: 2024 LiveKit, Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
-import type { ChatMessage } from '../llm/index.js';
-import { Future, shortuuid } from '../utils.js';
+import type { ChatItem } from '../llm/index.js';
+import { Event, Future, shortuuid } from '../utils.js';
+import type { Task } from '../utils.js';
+import { asyncLocalStorage } from './agent.js';
 
 export class SpeechHandle {
   /** Priority for messages that should be played after all other messages in the queue */
@@ -12,25 +14,40 @@ export class SpeechHandle {
   /** Priority for important messages that should be played before others. */
   static SPEECH_PRIORITY_HIGH = 10;
 
-  private interruptFut = new Future();
-  private authorizeFut = new Future();
-  private playoutDoneFut = new Future();
+  private interruptFut = new Future<void>();
+  private authorizedEvent = new Event();
+  private scheduledFut = new Future<void>();
+  private doneFut = new Future<void>();
 
-  private _chatMessage?: ChatMessage;
+  private generations: Future<void>[] = [];
+  /** @internal */
+  _tasks: Task<void>[] = [];
+  private _chatItems: ChatItem[] = [];
+  private _numSteps = 1;
+
+  private itemAddedCallbacks: Set<(item: ChatItem) => void> = new Set();
+  private doneCallbacks: Set<(sh: SpeechHandle) => void> = new Set();
 
   constructor(
-    readonly id: string,
-    readonly allowInterruptions: boolean,
-    readonly stepIndex: number,
+    private _id: string,
+    private _allowInterruptions: boolean,
+    /** @internal */
+    public _stepIndex: number,
     readonly parent?: SpeechHandle,
-  ) {}
+  ) {
+    this.doneFut.await.finally(() => {
+      for (const callback of this.doneCallbacks) {
+        callback(this);
+      }
+    });
+  }
 
-  static create(options: {
+  static create(options?: {
     allowInterruptions?: boolean;
     stepIndex?: number;
     parent?: SpeechHandle;
   }) {
-    const { allowInterruptions = false, stepIndex = 0, parent } = options ?? {};
+    const { allowInterruptions = true, stepIndex = 0, parent } = options ?? {};
 
     return new SpeechHandle(shortuuid('speech_'), allowInterruptions, stepIndex, parent);
   }
@@ -39,12 +56,47 @@ export class SpeechHandle {
     return this.interruptFut.done;
   }
 
-  get done(): boolean {
-    return this.playoutDoneFut.done;
+  get numSteps(): number {
+    return this._numSteps;
+  }
+
+  get id(): string {
+    return this._id;
   }
 
-  get chatMessage(): ChatMessage | undefined {
-    return this._chatMessage;
+  get scheduled(): boolean {
+    return this.scheduledFut.done;
+  }
+
+  get allowInterruptions(): boolean {
+    return this._allowInterruptions;
+  }
+
+  /**
+   * Allow or disallow interruptions on this SpeechHandle.
+   *
+   * When set to false, the SpeechHandle will no longer accept any incoming
+   * interruption requests until re-enabled. If the handle is already
+   * interrupted, clearing interruptions is not allowed.
+   *
+   * @param value - true to allow interruptions, false to disallow
+   * @throws Error If attempting to disable interruptions when already interrupted
+   */
+  set allowInterruptions(value: boolean) {
+    if (this.interrupted && !value) {
+      throw new Error(
+        'Cannot set allow_interruptions to False, the SpeechHandle is already interrupted',
+      );
+    }
+    this._allowInterruptions = value;
+  }
+
+  done(): boolean {
+    return this.doneFut.done;
+  }
+
+  get chatItems(): ChatItem[] {
+    return this._chatItems;
   }
 
   /**
@@ -54,23 +106,33 @@ export class SpeechHandle {
    *
    * @returns The same speech handle that was interrupted.
    */
-  interrupt(): SpeechHandle {
-    if (!this.allowInterruptions) {
-      throw new Error('interruptions are not allowed');
+  interrupt(force: boolean = false): SpeechHandle {
+    if (!force && !this.allowInterruptions) {
+      throw new Error('This generation handle does not allow interruptions');
     }
 
-    if (this.done) return this;
-
-    this.interruptFut.resolve();
+    this._cancel();
     return this;
   }
 
-  then(callback: (sh: SpeechHandle) => void) {
-    return this.playoutDoneFut.await.finally(() => callback(this));
-  }
-
-  async waitForPlayout() {
-    return this.playoutDoneFut.await;
+  /**
+   * Waits for the entire assistant turn to complete playback.
+   *
+   * This method waits until the assistant has fully finished speaking,
+   * including any finalization steps beyond initial response generation.
+   * This is appropriate to call when you want to ensure the speech output
+   * has entirely played out, including any tool calls and response follow-ups.
+   */
+  async waitForPlayout(): Promise<void> {
+    const store = asyncLocalStorage.getStore();
+    if (store && store?.functionCall) {
+      throw new Error(
+        `Cannot call 'SpeechHandle.waitForPlayout()' from inside the function tool '${store.functionCall.name}'. ` +
+          'This creates a circular wait: the speech handle is waiting for the function tool to complete, ' +
+          'while the function tool is simultaneously waiting for the speech handle.\n' +
+          "To wait for the assistant's spoken response prior to running this tool, use RunContext.wait_for_playout() instead.",
+      );
+    }
   }
 
   async waitIfNotInterrupted(aw: Promise<unknown>[]): Promise<void> {
@@ -79,26 +141,109 @@ export class SpeechHandle {
     await Promise.race(fs);
   }
 
+  addDoneCallback(callback: (sh: SpeechHandle) => void) {
+    this.doneCallbacks.add(callback);
+  }
+
+  removeDoneCallback(callback: (sh: SpeechHandle) => void) {
+    this.doneCallbacks.delete(callback);
+  }
+
+  /** @internal */
+  _cancel(): SpeechHandle {
+    if (this.done()) {
+      return this;
+    }
+
+    if (!this.interruptFut.done) {
+      this.interruptFut.resolve();
+    }
+
+    return this;
+  }
+
   /** @internal */
-  _setChatMessage(chatMessage: ChatMessage) {
-    if (this.done) {
-      throw new Error('cannot set chat message after speech has been played');
+  _authorizeGeneration(): void {
+    const fut = new Future<void>();
+    this.generations.push(fut);
+    this.authorizedEvent.set();
+  }
+
+  /** @internal */
+  _clearAuthorization(): void {
+    this.authorizedEvent.clear();
+  }
+
+  /** @internal */
+  async _waitForAuthorization(): Promise<void> {
+    await this.authorizedEvent.wait();
+  }
+
+  /** @internal */
+  async _waitForGeneration(stepIdx: number = -1): Promise<void> {
+    if (this.generations.length === 0) {
+      throw new Error('cannot use wait_for_generation: no active generation is running.');
     }
-    this._chatMessage = chatMessage;
+
+    const index = stepIdx === -1 ? this.generations.length - 1 : stepIdx;
+    const generation = this.generations[index];
+    if (!generation) {
+      throw new Error(`Generation at index ${index} not found.`);
+    }
+    return generation.await;
   }
 
   /** @internal */
-  _authorizePlayout() {
-    this.authorizeFut.resolve();
+  async _waitForScheduled(): Promise<void> {
+    return this.scheduledFut.await;
   }
 
   /** @internal */
-  async _waitForAuthorization() {
-    return this.authorizeFut.await;
+  _markGenerationDone(): void {
+    if (this.generations.length === 0) {
+      throw new Error('cannot use mark_generation_done: no active generation is running.');
+    }
+
+    const lastGeneration = this.generations[this.generations.length - 1];
+    if (lastGeneration && !lastGeneration.done) {
+      lastGeneration.resolve();
+    }
+  }
+
+  /** @internal */
+  _markDone(): void {
+    if (!this.doneFut.done) {
+      this.doneFut.resolve();
+      if (this.generations.length > 0) {
+        this._markGenerationDone(); // preemptive generation could be cancelled before being scheduled
+      }
+    }
+  }
+
+  /** @internal */
+  _markScheduled(): void {
+    if (!this.scheduledFut.done) {
+      this.scheduledFut.resolve();
+    }
+  }
+
+  /** @internal */
+  _addItemAddedCallback(callback: (item: ChatItem) => void): void {
+    this.itemAddedCallbacks.add(callback);
   }
 
   /** @internal */
-  _markPlayoutDone() {
-    this.playoutDoneFut.resolve();
+  _removeItemAddedCallback(callback: (item: ChatItem) => void): void {
+    this.itemAddedCallbacks.delete(callback);
+  }
+
+  /** @internal */
+  _itemAdded(items: ChatItem[]): void {
+    for (const item of items) {
+      for (const cb of this.itemAddedCallbacks) {
+        cb(item);
+      }
+      this._chatItems.push(item);
+    }
   }
 }

package/src/worker.ts

@@ -484,7 +484,11 @@ export class Worker {
     let participant: ParticipantInfo | undefined = undefined;
     if (participantIdentity) {
       try {
-        participant = await client.getParticipant(roomName, participantIdentity);
+        // 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}`,