@livekit/agents 1.0.46 → 1.0.47

package/src/voice/agent_session.ts

@@ -1,12 +1,14 @@
 // SPDX-FileCopyrightText: 2024 LiveKit, Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
+import { Mutex } from '@livekit/mutex';
 import type { AudioFrame, Room } from '@livekit/rtc-node';
 import type { TypedEventEmitter as TypedEmitter } from '@livekit/typed-emitter';
 import type { Context, Span } from '@opentelemetry/api';
 import { ROOT_CONTEXT, context as otelContext, trace } from '@opentelemetry/api';
 import { EventEmitter } from 'node:events';
 import type { ReadableStream } from 'node:stream/web';
+import type { z } from 'zod';
 import {
   LLM as InferenceLLM,
   STT as InferenceSTT,
@@ -31,6 +33,7 @@ import {
   type ResolvedSessionConnectOptions,
   type SessionConnectOptions,
 } from '../types.js';
+import { Task } from '../utils.js';
 import type { VAD } from '../vad.js';
 import type { Agent } from './agent.js';
 import { AgentActivity } from './agent_activity.js';
@@ -115,6 +118,13 @@ export type AgentSessionOptions<UserData = UnknownUserData> = {
   connOptions?: SessionConnectOptions;
 };
 
+type ActivityTransitionOptions = {
+  previousActivity?: 'close' | 'pause';
+  newActivity?: 'start' | 'resume';
+  blockedTasks?: Task<any>[];
+  waitOnEnter?: boolean;
+};
+
 export class AgentSession<
   UserData = UnknownUserData,
 > extends (EventEmitter as new () => TypedEmitter<AgentSessionCallbacks>) {
@@ -129,8 +139,10 @@ export class AgentSession<
   private agent?: Agent;
   private activity?: AgentActivity;
   private nextActivity?: AgentActivity;
+  private updateActivityTask?: Task<void>;
   private started = false;
   private userState: UserState = 'listening';
+  private readonly activityLock = new Mutex();
 
   /** @internal */
   _roomIO?: RoomIO;
@@ -360,7 +372,8 @@ export class AgentSession<
     }
 
     // TODO(AJS-265): add shutdown callback to job context
-    tasks.push(this.updateActivity(this.agent));
+    // Initial start does not wait on onEnter
+    tasks.push(this._updateActivity(this.agent, { waitOnEnter: false }));
 
     await Promise.allSettled(tasks);
 
@@ -432,8 +445,34 @@ export class AgentSession<
   updateAgent(agent: Agent): void {
     this.agent = agent;
 
-    if (this.started) {
-      this.updateActivity(agent);
+    if (!this.started) {
+      return;
+    }
+
+    const _updateActivityTask = async (oldTask: Task<void> | undefined, agent: Agent) => {
+      if (oldTask) {
+        try {
+          await oldTask.result;
+        } catch (error) {
+          this.logger.error(error, 'previous updateAgent transition failed');
+        }
+      }
+
+      await this._updateActivity(agent);
+    };
+
+    const oldTask = this.updateActivityTask;
+    this.updateActivityTask = Task.from(
+      async () => _updateActivityTask(oldTask, agent),
+      undefined,
+      'AgentSession_updateActivityTask',
+    );
+
+    const runState = this._globalRunState;
+    if (runState) {
+      // Don't mark the RunResult as done, if there is currently an agent transition happening.
+      // (used to make sure we're correctly adding the AgentHandoffResult before completion)
+      runState._watchHandle(this.updateActivityTask);
     }
   }
 
@@ -464,24 +503,42 @@ export class AgentSession<
       throw new Error('AgentSession is not running');
     }
 
-    const doSay = (activity: AgentActivity) => {
+    const doSay = (activity: AgentActivity, nextActivity?: AgentActivity) => {
+      if (activity.schedulingPaused) {
+        if (!nextActivity) {
+          throw new Error('AgentSession is closing, cannot use say()');
+        }
+        return nextActivity.say(text, options);
+      }
       return activity.say(text, options);
     };
 
+    const runState = this._globalRunState;
+    let handle: SpeechHandle;
+
     // attach to the session span if called outside of the AgentSession
     const activeSpan = trace.getActiveSpan();
     if (!activeSpan && this.rootSpanContext) {
-      return otelContext.with(this.rootSpanContext, () => doSay(this.activity!));
+      handle = otelContext.with(this.rootSpanContext, () =>
+        doSay(this.activity!, this.nextActivity),
+      );
+    } else {
+      handle = doSay(this.activity, this.nextActivity);
+    }
+
+    if (runState) {
+      runState._watchHandle(handle);
     }
 
-    return doSay(this.activity);
+    return handle;
   }
 
-  interrupt() {
+  interrupt(options?: { force?: boolean }) {
     if (!this.activity) {
       throw new Error('AgentSession is not running');
     }
-    return this.activity.interrupt();
+
+    return this.activity.interrupt(options);
   }
 
   generateReply(options?: {
@@ -502,7 +559,7 @@ export class AgentSession<
       : undefined;
 
     const doGenerateReply = (activity: AgentActivity, nextActivity?: AgentActivity) => {
-      if (activity.draining) {
+      if (activity.schedulingPaused) {
         if (!nextActivity) {
           throw new Error('AgentSession is closing, cannot use generateReply()');
         }
@@ -542,53 +599,128 @@ export class AgentSession<
    * result.expect.noMoreEvents();
    * ```
    *
-   * @param options - Run options including user input
+   * @param options - Run options including user input and optional output type
    * @returns A RunResult that resolves when the agent finishes responding
-   *
-   * TODO: Add outputType parameter for typed outputs (parity with Python)
    */
-  run(options: { userInput: string }): RunResult {
+  run<T = unknown>({
+    userInput,
+    outputType,
+  }: {
+    userInput: string;
+    outputType?: z.ZodType<T>;
+  }): RunResult<T> {
     if (this._globalRunState && !this._globalRunState.done()) {
       throw new Error('nested runs are not supported');
     }
 
-    const runState = new RunResult({ userInput: options.userInput });
+    const runState = new RunResult<T>({
+      userInput,
+      outputType,
+    });
+
     this._globalRunState = runState;
-    this.generateReply({ userInput: options.userInput });
+
+    // Defer generateReply through the activityLock to ensure any in-progress
+    // activity transition (e.g. AgentTask started from onEnter) completes first.
+    // TS Task.from starts onEnter synchronously, so the transition may already be
+    // mid-flight by the time run() is called after session.start() resolves.
+    // Acquiring and immediately releasing the lock guarantees FIFO ordering:
+    // the transition's lock section finishes before we route generateReply.
+    (async () => {
+      try {
+        const unlock = await this.activityLock.lock();
+        unlock();
+        this.generateReply({ userInput });
+      } catch (e) {
+        runState._reject(e instanceof Error ? e : new Error(String(e)));
+      }
+    })();
 
     return runState;
   }
 
-  private async updateActivity(agent: Agent): Promise<void> {
+  /** @internal */
+  async _updateActivity(agent: Agent, options: ActivityTransitionOptions = {}): Promise<void> {
+    const { previousActivity = 'close', newActivity = 'start', blockedTasks = [] } = options;
+    const waitOnEnter = options.waitOnEnter ?? newActivity === 'start';
+
     const runWithContext = async () => {
-      // TODO(AJS-129): add lock to agent activity core lifecycle
-      this.nextActivity = new AgentActivity(agent, this);
+      const unlock = await this.activityLock.lock();
+      let onEnterTask: Task<void> | undefined;
 
-      const previousActivity = this.activity;
+      try {
+        this.agent = agent;
+        const prevActivityObj = this.activity;
+
+        if (newActivity === 'start') {
+          const prevAgent = prevActivityObj?.agent;
+          if (
+            agent._agentActivity &&
+            // allow updating the same agent that is running
+            (agent !== prevAgent || previousActivity !== 'close')
+          ) {
+            throw new Error('Cannot start agent: an activity is already running');
+          }
+          this.nextActivity = new AgentActivity(agent, this);
+        } else if (newActivity === 'resume') {
+          if (!agent._agentActivity) {
+            throw new Error('Cannot resume agent: no existing activity to resume');
+          }
+          this.nextActivity = agent._agentActivity;
+        }
 
-      if (this.activity) {
-        await this.activity.drain();
-        await this.activity.close();
-      }
+        if (prevActivityObj && prevActivityObj !== this.nextActivity) {
+          if (previousActivity === 'pause') {
+            await prevActivityObj.pause({ blockedTasks });
+          } else {
+            await prevActivityObj.drain();
+            await prevActivityObj.close();
+          }
+        }
 
-      this.activity = this.nextActivity;
-      this.nextActivity = undefined;
+        this.activity = this.nextActivity;
+        this.nextActivity = undefined;
 
-      this._chatCtx.insert(
-        new AgentHandoffItem({
-          oldAgentId: previousActivity?.agent.id,
+        const runState = this._globalRunState;
+        const handoffItem = new AgentHandoffItem({
+          oldAgentId: prevActivityObj?.agent.id,
           newAgentId: agent.id,
-        }),
-      );
-      this.logger.debug(
-        { previousAgentId: previousActivity?.agent.id, newAgentId: agent.id },
-        'Agent handoff inserted into chat context',
-      );
+        });
 
-      await this.activity.start();
+        if (runState) {
+          runState._agentHandoff({
+            item: handoffItem,
+            oldAgent: prevActivityObj?.agent,
+            newAgent: this.activity!.agent,
+          });
+        }
+
+        this._chatCtx.insert(handoffItem);
+        this.logger.debug(
+          { previousAgentId: prevActivityObj?.agent.id, newAgentId: agent.id },
+          'Agent handoff inserted into chat context',
+        );
+
+        if (newActivity === 'start') {
+          await this.activity!.start();
+        } else {
+          await this.activity!.resume();
+        }
+
+        onEnterTask = this.activity!._onEnterTask;
+
+        if (this._input.audio) {
+          this.activity!.attachAudioInput(this._input.audio.stream);
+        }
+      } finally {
+        unlock();
+      }
 
-      if (this._input.audio) {
-        this.activity.attachAudioInput(this._input.audio.stream);
+      if (waitOnEnter) {
+        if (!onEnterTask) {
+          throw new Error('expected onEnter task to be available while waitOnEnter=true');
+        }
+        await onEnterTask.result;
       }
     };
 
@@ -842,15 +974,21 @@ export class AgentSession<
     if (this.activity) {
       if (!drain) {
         try {
-          this.activity.interrupt();
+          await this.activity.interrupt({ force: true }).await;
         } catch (error) {
-          // TODO(shubhra): force interrupt or wait for it to finish?
-          // it might be an audio played from the error callback
+          // Uninterruptible speech can throw during forced interruption.
+          this.logger.warn({ error }, 'Error interrupting activity');
         }
       }
+
       await this.activity.drain();
       // wait any uninterruptible speech to finish
       await this.activity.currentSpeech?.waitForPlayout();
+
+      if (reason !== CloseReason.ERROR) {
+        this.activity.commitUserTurn({ audioDetached: true, throwIfNotReady: false });
+      }
+
       try {
         this.activity.detachAudioInput();
       } catch (error) {

package/src/voice/audio_recognition.ts

@@ -768,6 +768,10 @@ export class AudioRecognition {
         this.logger.debug('User turn committed');
       })
       .catch((err: unknown) => {
+        if (err instanceof Error && err.name === 'AbortError') {
+          this.logger.debug('User turn commit task cancelled');
+          return;
+        }
         this.logger.error(err, 'Error in user turn commit task:');
       });
   }

package/src/voice/generation.ts

@@ -26,7 +26,13 @@ import { IdentityTransform } from '../stream/identity_transform.js';
 import { traceTypes, tracer } from '../telemetry/index.js';
 import { USERDATA_TIMED_TRANSCRIPT } from '../types.js';
 import { Future, Task, shortuuid, toError, waitForAbort } from '../utils.js';
-import { type Agent, type ModelSettings, asyncLocalStorage, isStopResponse } from './agent.js';
+import {
+  type Agent,
+  type ModelSettings,
+  _setActivityTaskInfo,
+  functionCallStorage,
+  isStopResponse,
+} from './agent.js';
 import type { AgentSession } from './agent_session.js';
 import {
   AudioOutput,
@@ -719,7 +725,7 @@ export interface _AudioOut {
 
 async function forwardAudio(
   ttsStream: ReadableStream<AudioFrame>,
-  audioOuput: AudioOutput,
+  audioOutput: AudioOutput,
   out: _AudioOut,
   signal?: AbortSignal,
 ): Promise<void> {
@@ -733,8 +739,8 @@ async function forwardAudio(
   };
 
   try {
-    audioOuput.on(AudioOutput.EVENT_PLAYBACK_STARTED, onPlaybackStarted);
-    audioOuput.resume();
+    audioOutput.on(AudioOutput.EVENT_PLAYBACK_STARTED, onPlaybackStarted);
+    audioOutput.resume();
 
     while (true) {
       if (signal?.aborted) {
@@ -748,36 +754,36 @@ async function forwardAudio(
 
       if (
         !out.firstFrameFut.done &&
-        audioOuput.sampleRate &&
-        audioOuput.sampleRate !== frame.sampleRate &&
+        audioOutput.sampleRate &&
+        audioOutput.sampleRate !== frame.sampleRate &&
         !resampler
       ) {
-        resampler = new AudioResampler(frame.sampleRate, audioOuput.sampleRate, 1);
+        resampler = new AudioResampler(frame.sampleRate, audioOutput.sampleRate, 1);
       }
 
       if (resampler) {
         for (const f of resampler.push(frame)) {
-          await audioOuput.captureFrame(f);
+          await audioOutput.captureFrame(f);
         }
       } else {
-        await audioOuput.captureFrame(frame);
+        await audioOutput.captureFrame(frame);
       }
     }
 
     if (resampler) {
       for (const f of resampler.flush()) {
-        await audioOuput.captureFrame(f);
+        await audioOutput.captureFrame(f);
       }
     }
   } finally {
-    audioOuput.off(AudioOutput.EVENT_PLAYBACK_STARTED, onPlaybackStarted);
+    audioOutput.off(AudioOutput.EVENT_PLAYBACK_STARTED, onPlaybackStarted);
 
     if (!out.firstFrameFut.done) {
       out.firstFrameFut.reject(new Error('audio forwarding cancelled before playback started'));
     }
 
     reader?.releaseLock();
-    audioOuput.flush();
+    audioOutput.flush();
   }
 }
 
@@ -836,7 +842,7 @@ export function performToolExecutions({
     const signal = controller.signal;
     const reader = toolCallStream.getReader();
 
-    const tasks: Promise<any>[] = [];
+    const tasks: Task<void>[] = [];
     while (!signal.aborted) {
       const { done, value: toolCall } = await reader.read();
       if (signal.aborted) break;
@@ -929,14 +935,6 @@ export function performToolExecutions({
         'Executing LLM tool call',
       );
 
-      const toolExecution = asyncLocalStorage.run({ functionCall: toolCall }, async () => {
-        return await tool.execute(parsedArgs, {
-          ctx: new RunContext(session, speechHandle, toolCall),
-          toolCallId: toolCall.callId,
-          abortSignal: signal,
-        });
-      });
-
       const _tracableToolExecutionImpl = async (toolExecTask: Promise<unknown>, span: Span) => {
         span.setAttribute(traceTypes.ATTR_FUNCTION_TOOL_NAME, toolCall.name);
         span.setAttribute(traceTypes.ATTR_FUNCTION_TOOL_ARGS, toolCall.args);
@@ -993,11 +991,42 @@ export function performToolExecutions({
           name: 'function_tool',
         });
 
+      const toolTask = Task.from(
+        async () => {
+          // Ensure this task is marked inline before user tool code executes.
+          const currentTask = Task.current();
+          if (currentTask) {
+            _setActivityTaskInfo(currentTask, {
+              speechHandle,
+              functionCall: toolCall,
+              inlineTask: true,
+            });
+          }
+
+          const toolExecution = functionCallStorage.run({ functionCall: toolCall }, async () => {
+            return await tool.execute(parsedArgs, {
+              ctx: new RunContext(session, speechHandle, toolCall),
+              toolCallId: toolCall.callId,
+              abortSignal: signal,
+            });
+          });
+
+          await tracableToolExecution(toolExecution);
+        },
+        controller,
+        `performToolExecution:${toolCall.name}`,
+      );
+
+      _setActivityTaskInfo(toolTask, {
+        speechHandle,
+        functionCall: toolCall,
+        inlineTask: true,
+      });
       // wait, not cancelling all tool calling tasks
-      tasks.push(tracableToolExecution(toolExecution));
+      tasks.push(toolTask);
     }
 
-    await Promise.allSettled(tasks);
+    await Promise.allSettled(tasks.map((task) => task.result));
     if (toolOutput.output.length > 0) {
       logger.debug(
         {

package/src/voice/index.ts

@@ -1,7 +1,7 @@
 // SPDX-FileCopyrightText: 2025 LiveKit, Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
-export { Agent, StopResponse, type AgentOptions, type ModelSettings } from './agent.js';
+export { Agent, AgentTask, StopResponse, type AgentOptions, type ModelSettings } from './agent.js';
 export { AgentSession, type AgentSessionOptions, type VoiceOptions } from './agent_session.js';
 export * from './avatar/index.js';
 export * from './background_audio.js';

package/src/voice/speech_handle.ts

@@ -5,7 +5,7 @@ import type { Context } from '@opentelemetry/api';
 import type { ChatItem } from '../llm/index.js';
 import type { Task } from '../utils.js';
 import { Event, Future, shortuuid } from '../utils.js';
-import { asyncLocalStorage } from './agent.js';
+import { functionCallStorage } from './agent.js';
 
 /** Symbol used to identify SpeechHandle instances */
 const SPEECH_HANDLE_SYMBOL = Symbol.for('livekit.agents.SpeechHandle');
@@ -46,6 +46,9 @@ export class SpeechHandle {
   /** @internal - OpenTelemetry context for the agent turn span */
   _agentTurnContext?: Context;
 
+  /** @internal - used by AgentTask/RunResult final output plumbing */
+  _maybeRunFinalOutput?: unknown;
+
   private itemAddedCallbacks: Set<(item: ChatItem) => void> = new Set();
   private doneCallbacks: Set<(sh: SpeechHandle) => void> = new Set();
 
@@ -148,7 +151,7 @@ export class SpeechHandle {
    * has entirely played out, including any tool calls and response follow-ups.
    */
   async waitForPlayout(): Promise<void> {
-    const store = asyncLocalStorage.getStore();
+    const store = functionCallStorage.getStore();
     if (store && store?.functionCall) {
       throw new Error(
         `Cannot call 'SpeechHandle.waitForPlayout()' from inside the function tool '${store.functionCall.name}'. ` +
@@ -167,6 +170,10 @@ export class SpeechHandle {
   }
 
   addDoneCallback(callback: (sh: SpeechHandle) => void) {
+    if (this.done()) {
+      queueMicrotask(() => callback(this));
+      return;
+    }
     this.doneCallbacks.add(callback);
   }