@livekit/agents 1.0.46 → 1.0.48

package/src/utils.test.ts

@@ -469,6 +469,93 @@ describe('utils', () => {
         expect((error as Error).name).toBe('TypeError');
       }
     });
+
+    it('should return undefined for Task.current outside task context', () => {
+      expect(Task.current()).toBeUndefined();
+    });
+
+    it('should preserve Task.current inside a task across awaits', async () => {
+      const task = Task.from(
+        async () => {
+          const currentAtStart = Task.current();
+          await delay(5);
+          const currentAfterAwait = Task.current();
+
+          expect(currentAtStart).toBeDefined();
+          expect(currentAfterAwait).toBe(currentAtStart);
+
+          return currentAtStart;
+        },
+        undefined,
+        'current-context-test',
+      );
+
+      const currentFromResult = await task.result;
+      expect(currentFromResult).toBe(task);
+    });
+
+    it('should isolate nested Task.current context and restore parent context', async () => {
+      const parentTask = Task.from(
+        async (controller) => {
+          const parentCurrent = Task.current();
+          expect(parentCurrent).toBeDefined();
+
+          const childTask = Task.from(
+            async () => {
+              const childCurrentStart = Task.current();
+              await delay(5);
+              const childCurrentAfterAwait = Task.current();
+
+              expect(childCurrentStart).toBeDefined();
+              expect(childCurrentAfterAwait).toBe(childCurrentStart);
+              expect(childCurrentStart).not.toBe(parentCurrent);
+
+              return childCurrentStart;
+            },
+            controller,
+            'child-current-context-test',
+          );
+
+          const childCurrent = await childTask.result;
+          const parentCurrentAfterChild = Task.current();
+
+          expect(parentCurrentAfterChild).toBe(parentCurrent);
+
+          return { parentCurrent, childCurrent };
+        },
+        undefined,
+        'parent-current-context-test',
+      );
+
+      const { parentCurrent, childCurrent } = await parentTask.result;
+      expect(parentCurrent).toBe(parentTask);
+      expect(childCurrent).not.toBe(parentCurrent);
+      expect(Task.current()).toBeUndefined();
+    });
+
+    it('should always expose Task.current for concurrent task callbacks', async () => {
+      const tasks = Array.from({ length: 25 }, (_, idx) =>
+        Task.from(
+          async () => {
+            const currentAtStart = Task.current();
+            await delay(1);
+            const currentAfterAwait = Task.current();
+
+            expect(currentAtStart).toBeDefined();
+            expect(currentAfterAwait).toBe(currentAtStart);
+
+            return currentAtStart;
+          },
+          undefined,
+          `current-context-stress-${idx}`,
+        ),
+      );
+
+      const currentTasks = await Promise.all(tasks.map((task) => task.result));
+      currentTasks.forEach((currentTask, idx) => {
+        expect(currentTask).toBe(tasks[idx]);
+      });
+    });
   });
 
   describe('Event', () => {

package/src/utils.ts

@@ -9,6 +9,7 @@ import type {
   TrackKind,
 } from '@livekit/rtc-node';
 import { AudioFrame, AudioResampler, RoomEvent } from '@livekit/rtc-node';
+import { AsyncLocalStorage } from 'node:async_hooks';
 import { EventEmitter, once } from 'node:events';
 import type { ReadableStream } from 'node:stream/web';
 import { TransformStream, type TransformStreamDefaultController } from 'node:stream/web';
@@ -172,6 +173,11 @@ export class Future<T = void> {
     this.#rejected = true;
     this.#error = error;
     this.#rejectPromise(error);
+    // Python calls Future.exception() right after set_exception() to silence
+    // "exception was never retrieved" warnings. In JS, consume the rejection
+    // immediately so Node does not emit unhandled-rejection noise before a
+    // later await/catch observes it.
+    void this.#await.catch(() => undefined);
   }
 }
 
@@ -434,7 +440,9 @@ export enum TaskResult {
  * @param T - The type of the task result
  */
 export class Task<T> {
+  private static readonly currentTaskStorage = new AsyncLocalStorage<Task<unknown>>();
   private resultFuture: Future<T>;
+  private doneCallbacks: Set<() => void> = new Set();
 
   #logger = log();
 
@@ -444,6 +452,21 @@ export class Task<T> {
     readonly name?: string,
   ) {
     this.resultFuture = new Future();
+    void this.resultFuture.await
+      .then(
+        () => undefined,
+        () => undefined,
+      )
+      .finally(() => {
+        for (const callback of this.doneCallbacks) {
+          try {
+            callback();
+          } catch (error) {
+            this.#logger.error({ error }, 'Task done callback failed');
+          }
+        }
+        this.doneCallbacks.clear();
+      });
     this.runTask();
   }
 
@@ -463,6 +486,13 @@ export class Task<T> {
     return new Task(fn, abortController, name);
   }
 
+  /**
+   * Returns the currently running task in this async context, if available.
+   */
+  static current(): Task<unknown> | undefined {
+    return Task.currentTaskStorage.getStore();
+  }
+
   private async runTask() {
     const run = async () => {
       if (this.name) {
@@ -471,7 +501,8 @@ export class Task<T> {
       return await this.fn(this.controller);
     };
 
-    return run()
+    return Task.currentTaskStorage
+      .run(this as Task<unknown>, run)
       .then((value) => {
         this.resultFuture.resolve(value);
         return value;
@@ -543,7 +574,15 @@ export class Task<T> {
   }
 
   addDoneCallback(callback: () => void) {
-    this.resultFuture.await.finally(callback);
+    if (this.done) {
+      queueMicrotask(callback);
+      return;
+    }
+    this.doneCallbacks.add(callback);
+  }
+
+  removeDoneCallback(callback: () => void) {
+    this.doneCallbacks.delete(callback);
   }
 }
 

package/src/version.ts

@@ -2,4 +2,4 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
-export const version = '0.1.0';
+export const version = __PACKAGE_VERSION__;

package/src/voice/agent.test.ts

@@ -1,10 +1,15 @@
 // SPDX-FileCopyrightText: 2025 LiveKit, Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
-import { describe, expect, it } from 'vitest';
+import { describe, expect, it, vi } from 'vitest';
 import { z } from 'zod';
 import { tool } from '../llm/index.js';
-import { Agent } from './agent.js';
+import { initializeLogger } from '../log.js';
+import { Task } from '../utils.js';
+import { Agent, AgentTask, _setActivityTaskInfo } from './agent.js';
+import { agentActivityStorage } from './agent_activity.js';
+
+initializeLogger({ pretty: false, level: 'error' });
 
 describe('Agent', () => {
   it('should create agent with basic instructions', () => {
@@ -77,4 +82,137 @@ describe('Agent', () => {
     expect(tools1).toEqual(tools2);
     expect(tools1).toEqual(tools);
   });
+
+  it('should require AgentTask to run inside task context', async () => {
+    class TestTask extends AgentTask<string> {
+      constructor() {
+        super({ instructions: 'test task' });
+      }
+    }
+
+    const task = new TestTask();
+    await expect(task.run()).rejects.toThrow('must be executed inside a Task context');
+  });
+
+  it('should require AgentTask to run inside inline task context', async () => {
+    class TestTask extends AgentTask<string> {
+      constructor() {
+        super({ instructions: 'test task' });
+      }
+    }
+
+    const task = new TestTask();
+    const wrapper = Task.from(async () => {
+      return await task.run();
+    });
+
+    await expect(wrapper.result).rejects.toThrow(
+      'should only be awaited inside function tools or the onEnter/onExit methods of an Agent',
+    );
+  });
+
+  it('should allow AgentTask run from inline task context', async () => {
+    class TestTask extends AgentTask<string> {
+      constructor() {
+        super({ instructions: 'test task' });
+      }
+    }
+
+    const task = new TestTask();
+    const oldAgent = new Agent({ instructions: 'old agent' });
+    const mockSession = {
+      currentAgent: oldAgent,
+      _globalRunState: undefined,
+      _updateActivity: async (agent: Agent) => {
+        if (agent === task) {
+          task.complete('ok');
+        }
+      },
+    };
+
+    const mockActivity = {
+      agent: oldAgent,
+      agentSession: mockSession,
+      _onEnterTask: undefined,
+      llm: undefined,
+      close: async () => {},
+    };
+
+    const wrapper = Task.from(async () => {
+      const currentTask = Task.current();
+      if (!currentTask) {
+        throw new Error('expected task context');
+      }
+      _setActivityTaskInfo(currentTask, { inlineTask: true });
+      return await agentActivityStorage.run(mockActivity as any, () => task.run());
+    });
+
+    await expect(wrapper.result).resolves.toBe('ok');
+  });
+
+  it('should require AgentTask to run inside AgentActivity context', async () => {
+    class TestTask extends AgentTask<string> {
+      constructor() {
+        super({ instructions: 'test task' });
+      }
+    }
+
+    const task = new TestTask();
+    const wrapper = Task.from(async () => {
+      const currentTask = Task.current();
+      if (!currentTask) {
+        throw new Error('expected task context');
+      }
+      _setActivityTaskInfo(currentTask, { inlineTask: true });
+      return await task.run();
+    });
+
+    await expect(wrapper.result).rejects.toThrow(
+      'must be executed inside an AgentActivity context',
+    );
+  });
+
+  it('should close old activity when current agent changes while AgentTask is pending', async () => {
+    class TestTask extends AgentTask<string> {
+      constructor() {
+        super({ instructions: 'test task' });
+      }
+    }
+
+    const task = new TestTask();
+    const oldAgent = new Agent({ instructions: 'old agent' });
+    const switchedAgent = new Agent({ instructions: 'switched agent' });
+    const closeOldActivity = vi.fn(async () => {});
+
+    const mockSession = {
+      currentAgent: oldAgent as Agent,
+      _globalRunState: undefined,
+      _updateActivity: async (agent: Agent) => {
+        if (agent === task) {
+          mockSession.currentAgent = switchedAgent;
+          task.complete('ok');
+        }
+      },
+    };
+
+    const mockActivity = {
+      agent: oldAgent,
+      agentSession: mockSession,
+      _onEnterTask: undefined,
+      llm: undefined,
+      close: closeOldActivity,
+    };
+
+    const wrapper = Task.from(async () => {
+      const currentTask = Task.current();
+      if (!currentTask) {
+        throw new Error('expected task context');
+      }
+      _setActivityTaskInfo(currentTask, { inlineTask: true });
+      return await agentActivityStorage.run(mockActivity as any, () => task.run());
+    });
+
+    await expect(wrapper.result).resolves.toBe('ok');
+    expect(closeOldActivity).toHaveBeenCalledTimes(1);
+  });
 });

package/src/voice/agent.ts

@@ -13,26 +13,71 @@ import {
   type TTSModelString,
 } from '../inference/index.js';
 import { ReadonlyChatContext } from '../llm/chat_context.js';
-import type { ChatMessage, FunctionCall, RealtimeModel } from '../llm/index.js';
+import type { ChatMessage, FunctionCall } from '../llm/index.js';
 import {
   type ChatChunk,
   ChatContext,
   LLM,
+  RealtimeModel,
   type ToolChoice,
   type ToolContext,
 } from '../llm/index.js';
+import { log } from '../log.js';
 import type { STT, SpeechEvent } from '../stt/index.js';
 import { StreamAdapter as STTStreamAdapter } from '../stt/index.js';
 import { SentenceTokenizer as BasicSentenceTokenizer } from '../tokenize/basic/index.js';
 import type { TTS } from '../tts/index.js';
 import { SynthesizeStream, StreamAdapter as TTSStreamAdapter } from '../tts/index.js';
 import { USERDATA_TIMED_TRANSCRIPT } from '../types.js';
+import { Future, Task } from '../utils.js';
 import type { VAD } from '../vad.js';
-import type { AgentActivity } from './agent_activity.js';
+import { type AgentActivity, agentActivityStorage } from './agent_activity.js';
 import type { AgentSession, TurnDetectionMode } from './agent_session.js';
 import type { TimedString } from './io.js';
+import type { SpeechHandle } from './speech_handle.js';
+
+export const functionCallStorage = new AsyncLocalStorage<{ functionCall?: FunctionCall }>();
+export const speechHandleStorage = new AsyncLocalStorage<SpeechHandle>();
+const activityTaskInfoStorage = new WeakMap<Task<any>, _ActivityTaskInfo>();
+
+type _ActivityTaskInfo = {
+  functionCall: FunctionCall | null;
+  speechHandle: SpeechHandle | null;
+  inlineTask: boolean;
+};
+
+/** @internal */
+export function _setActivityTaskInfo<T>(
+  task: Task<T>,
+  options: {
+    functionCall?: FunctionCall | null;
+    speechHandle?: SpeechHandle | null;
+    inlineTask?: boolean;
+  },
+): void {
+  const info = activityTaskInfoStorage.get(task) ?? {
+    functionCall: null,
+    speechHandle: null,
+    inlineTask: false,
+  };
+
+  if (Object.hasOwn(options, 'functionCall')) {
+    info.functionCall = options.functionCall ?? null;
+  }
+  if (Object.hasOwn(options, 'speechHandle')) {
+    info.speechHandle = options.speechHandle ?? null;
+  }
+  if (Object.hasOwn(options, 'inlineTask')) {
+    info.inlineTask = options.inlineTask ?? false;
+  }
+
+  activityTaskInfoStorage.set(task, info);
+}
 
-export const asyncLocalStorage = new AsyncLocalStorage<{ functionCall?: FunctionCall }>();
+/** @internal */
+export function _getActivityTaskInfo<T>(task: Task<T>): _ActivityTaskInfo | undefined {
+  return activityTaskInfoStorage.get(task);
+}
 export const STOP_RESPONSE_SYMBOL = Symbol('StopResponse');
 
 export class StopResponse extends Error {
@@ -257,6 +302,17 @@ export class Agent<UserData = any> {
     this._agentActivity.updateChatCtx(chatCtx);
   }
 
+  // TODO(parity): Add when AgentConfigUpdate is ported to ChatContext.
+  async updateTools(tools: ToolContext): Promise<void> {
+    if (!this._agentActivity) {
+      this._tools = { ...tools };
+      this._chatCtx = this._chatCtx.copy({ toolCtx: this._tools });
+      return;
+    }
+
+    await this._agentActivity.updateTools(tools);
+  }
+
   static default = {
     async sttNode(
       agent: Agent,
@@ -268,20 +324,20 @@ export class Agent<UserData = any> {
         throw new Error('sttNode called but no STT node is available');
       }
 
-      let wrapped_stt = activity.stt;
+      let wrappedStt = activity.stt;
 
-      if (!wrapped_stt.capabilities.streaming) {
+      if (!wrappedStt.capabilities.streaming) {
         const vad = agent.vad || activity.vad;
         if (!vad) {
           throw new Error(
             'STT does not support streaming, add a VAD to the AgentTask/VoiceAgent to enable streaming',
           );
         }
-        wrapped_stt = new STTStreamAdapter(wrapped_stt, vad);
+        wrappedStt = new STTStreamAdapter(wrappedStt, vad);
       }
 
       const connOptions = activity.agentSession.connOptions.sttConnOptions;
-      const stream = wrapped_stt.stream({ connOptions });
+      const stream = wrappedStt.stream({ connOptions });
 
       // Set startTimeOffset to provide linear timestamps across reconnections
       const audioInputStartedAt =
@@ -382,14 +438,14 @@ export class Agent<UserData = any> {
         throw new Error('ttsNode called but no TTS node is available');
       }
 
-      let wrapped_tts = activity.tts;
+      let wrappedTts = activity.tts;
 
       if (!activity.tts.capabilities.streaming) {
-        wrapped_tts = new TTSStreamAdapter(wrapped_tts, new BasicSentenceTokenizer());
+        wrappedTts = new TTSStreamAdapter(wrappedTts, new BasicSentenceTokenizer());
       }
 
       const connOptions = activity.agentSession.connOptions.ttsConnOptions;
-      const stream = wrapped_tts.stream({ connOptions });
+      const stream = wrappedTts.stream({ connOptions });
       stream.updateInputStream(text);
 
       let cleaned = false;
@@ -440,3 +496,137 @@ export class Agent<UserData = any> {
     },
   };
 }
+
+export class AgentTask<ResultT = unknown, UserData = any> extends Agent<UserData> {
+  private started = false;
+  private future = new Future<ResultT>();
+
+  #logger = log();
+
+  get done(): boolean {
+    return this.future.done;
+  }
+
+  complete(result: ResultT | Error): void {
+    if (this.future.done) {
+      throw new Error(`${this.constructor.name} is already done`);
+    }
+
+    if (result instanceof Error) {
+      this.future.reject(result);
+    } else {
+      this.future.resolve(result);
+    }
+
+    const speechHandle = speechHandleStorage.getStore();
+    if (speechHandle) {
+      speechHandle._maybeRunFinalOutput = result;
+    }
+  }
+
+  async run(): Promise<ResultT> {
+    if (this.started) {
+      throw new Error(
+        `Task ${this.constructor.name} has already started and cannot be awaited multiple times`,
+      );
+    }
+    this.started = true;
+
+    const currentTask = Task.current();
+    if (!currentTask) {
+      throw new Error(`${this.constructor.name} must be executed inside a Task context`);
+    }
+
+    const taskInfo = _getActivityTaskInfo(currentTask);
+    if (!taskInfo || !taskInfo.inlineTask) {
+      throw new Error(
+        `${this.constructor.name} should only be awaited inside function tools or the onEnter/onExit methods of an Agent`,
+      );
+    }
+
+    const speechHandle = speechHandleStorage.getStore();
+    const oldActivity = agentActivityStorage.getStore();
+    if (!oldActivity) {
+      throw new Error(`${this.constructor.name} must be executed inside an AgentActivity context`);
+    }
+
+    currentTask.addDoneCallback(() => {
+      if (this.future.done) return;
+
+      // If the Task finished before the AgentTask was completed, complete the AgentTask with an error.
+      this.#logger.error(`The Task finished before ${this.constructor.name} was completed.`);
+      this.complete(new Error(`The Task finished before ${this.constructor.name} was completed.`));
+    });
+
+    const oldAgent = oldActivity.agent;
+    const session = oldActivity.agentSession;
+
+    const blockedTasks: Task<any>[] = [currentTask];
+    const onEnterTask = oldActivity._onEnterTask;
+
+    if (onEnterTask && !onEnterTask.done && onEnterTask !== currentTask) {
+      blockedTasks.push(onEnterTask);
+    }
+
+    if (
+      taskInfo.functionCall &&
+      oldActivity.llm instanceof RealtimeModel &&
+      !oldActivity.llm.capabilities.manualFunctionCalls
+    ) {
+      this.#logger.error(
+        `Realtime model does not support resuming function calls from chat context, ` +
+          `using AgentTask inside a function tool may have unexpected behavior.`,
+      );
+    }
+
+    await session._updateActivity(this, {
+      previousActivity: 'pause',
+      newActivity: 'start',
+      blockedTasks,
+    });
+
+    let runState = session._globalRunState;
+    if (speechHandle && runState && !runState.done()) {
+      // Only unwatch the parent speech handle if there are other handles keeping the run alive.
+      // When watchedHandleCount is 1 (only the parent), unwatching would drop it to 0 and
+      // mark the run done prematurely — before function_call_output and assistant message arrive.
+      if (runState._watchedHandleCount() > 1) {
+        runState._unwatchHandle(speechHandle);
+      }
+      // it is OK to call _markDoneIfNeeded here, the above _updateActivity will call onEnter
+      // and newly added handles keep the run alive.
+      runState._markDoneIfNeeded();
+    }
+
+    try {
+      return await this.future.await;
+    } finally {
+      // runState could have changed after future resolved
+      runState = session._globalRunState;
+
+      if (session.currentAgent !== this) {
+        this.#logger.warn(
+          `${this.constructor.name} completed, but the agent has changed in the meantime. ` +
+            `Ignoring handoff to the previous agent, likely due to AgentSession.updateAgent being invoked.`,
+        );
+        await oldActivity.close();
+      } else {
+        if (speechHandle && runState && !runState.done()) {
+          runState._watchHandle(speechHandle);
+        }
+
+        const mergedChatCtx = oldAgent._chatCtx.merge(this._chatCtx, {
+          excludeFunctionCall: true,
+          excludeInstructions: true,
+        });
+        oldAgent._chatCtx.items = mergedChatCtx.items;
+
+        await session._updateActivity(oldAgent, {
+          previousActivity: 'close',
+          newActivity: 'resume',
+          waitOnEnter: false,
+        });
+      }
+    }
+  }
+}