Move job_context.connect inside of AgentSession.start (#689)

Author: Shubhrakanti

.changeset/eight-parents-drive.md

@@ -0,0 +1,6 @@
+---
+'@livekit/agents-plugin-livekit': patch
+'@livekit/agents': patch
+---
+
+Remove requirement to call ctx.connect in entrypoint function

agents/src/ipc/job_proc_lazy_main.ts

@@ -6,7 +6,7 @@ import { EventEmitter, once } from 'node:events';
 import { pathToFileURL } from 'node:url';
 import type { Logger } from 'pino';
 import { type Agent, isAgent } from '../generator.js';
-import { CurrentJobContext, JobContext, JobProcess, type RunningJobInfo } from '../job.js';
+import { JobContext, JobProcess, type RunningJobInfo, runWithJobContextAsync } from '../job.js';
 import { initializeLogger, log } from '../log.js';
 import { Future, shortuuid } from '../utils.js';
 import { defaultInitializeProcessFunc } from '../worker.js';
@@ -87,7 +87,6 @@ const startJob = (
   };
 
   const ctx = new JobContext(proc, info, room, onConnect, onShutdown, new InfClient());
-  new CurrentJobContext(ctx);
 
   const task = new Promise<void>(async () => {
     const unconnectedTimeout = setTimeout(() => {
@@ -98,7 +97,11 @@ const startJob = (
         );
       }
     }, 10000);
-    func(ctx).finally(() => clearTimeout(unconnectedTimeout));
+
+    // Run the job function within the AsyncLocalStorage context
+    await runWithJobContextAsync(ctx, () => func(ctx)).finally(() => {
+      clearTimeout(unconnectedTimeout);
+    });
 
     await once(closeEvent, 'close').then((close) => {
       logger.debug('shutting down');

agents/src/job.ts

@@ -10,35 +10,43 @@ import type {
   RtcConfiguration,
 } from '@livekit/rtc-node';
 import { ParticipantKind, RoomEvent, TrackKind } from '@livekit/rtc-node';
+import { AsyncLocalStorage } from 'node:async_hooks';
 import type { Logger } from 'pino';
 import type { InferenceExecutor } from './ipc/inference_executor.js';
 import { log } from './log.js';
 
-export class CurrentJobContext {
-  static #current: JobContext;
-
-  constructor(proc: JobContext) {
-    CurrentJobContext.#current = proc;
-  }
-
-  static getCurrent(): JobContext {
-    return CurrentJobContext.#current;
-  }
-}
+// AsyncLocalStorage for job context, similar to Python's contextvars
+const jobContextStorage = new AsyncLocalStorage<JobContext>();
 
 /**
  * Returns the current job context.
  *
  * @throws {Error} if no job context is found
  */
 export function getJobContext(): JobContext {
-  const ctx = CurrentJobContext.getCurrent();
+  const ctx = jobContextStorage.getStore();
   if (!ctx) {
     throw new Error('no job context found, are you running this code inside a job entrypoint?');
   }
   return ctx;
 }
 
+/**
+ * Runs a function within a job context, similar to Python's contextvars.
+ * @internal
+ */
+export function runWithJobContext<T>(context: JobContext, fn: () => T): T {
+  return jobContextStorage.run(context, fn);
+}
+
+/**
+ * Runs an async function within a job context, similar to Python's contextvars.
+ * @internal
+ */
+export function runWithJobContextAsync<T>(context: JobContext, fn: () => Promise<T>): Promise<T> {
+  return jobContextStorage.run(context, fn);
+}
+
 /** Which tracks, if any, should the agent automatically subscribe to? */
 export enum AutoSubscribe {
   SUBSCRIBE_ALL,
@@ -89,6 +97,8 @@ export class JobContext {
   #logger: Logger;
   #inferenceExecutor: InferenceExecutor;
 
+  private connected: boolean = false;
+
   constructor(
     proc: JobProcess,
     info: RunningJobInfo,
@@ -191,6 +201,10 @@ export class JobContext {
     autoSubscribe: AutoSubscribe = AutoSubscribe.SUBSCRIBE_ALL,
     rtcConfig?: RtcConfiguration,
   ) {
+    if (this.connected) {
+      return;
+    }
+
     const opts = {
       e2ee,
       autoSubscribe: autoSubscribe == AutoSubscribe.SUBSCRIBE_ALL,
@@ -215,6 +229,7 @@ export class JobContext {
         });
       });
     }
+    this.connected = true;
   }
 
   /**

agents/src/voice/agent_session.ts

@@ -5,6 +5,7 @@ import type { AudioFrame, Room } from '@livekit/rtc-node';
 import type { TypedEventEmitter as TypedEmitter } from '@livekit/typed-emitter';
 import { EventEmitter } from 'node:events';
 import type { ReadableStream } from 'node:stream/web';
+import { getJobContext } from '../job.js';
 import { ChatContext, ChatMessage } from '../llm/chat_context.js';
 import type { LLM, RealtimeModel, RealtimeModelError, ToolChoice } from '../llm/index.js';
 import type { LLMError } from '../llm/llm.js';
@@ -184,6 +185,7 @@ export class AgentSession<
     this.agent = agent;
     this._updateAgentState('initializing');
 
+    const tasks: Promise<void>[] = [];
     // Check for existing input/output configuration and warn if needed
     if (this.input.audio && inputOptions?.audioEnabled !== false) {
       this.logger.warn('RoomIO audio input is enabled but input.audio is already set, ignoring..');
@@ -209,7 +211,15 @@ export class AgentSession<
     });
     this.roomIO.start();
 
-    await this.updateActivity(this.agent);
+    const ctx = getJobContext();
+    if (ctx && ctx.room === room && !room.isConnected) {
+      this.logger.debug('Auto-connecting to room via job context');
+      tasks.push(ctx.connect());
+    }
+    // TODO(AJS-265): add shutdown callback to job context
+    tasks.push(this.updateActivity(this.agent));
+
+    await Promise.allSettled(tasks);
 
     // Log used IO configuration
     this.logger.debug(
@@ -220,7 +230,6 @@ export class AgentSession<
       `using transcript io: \`AgentSession\` -> ${this.output.transcription ? '`' + this.output.transcription.constructor.name + '`' : '(none)'}`,
     );
 
-    this.logger.debug('AgentSession started');
     this.started = true;
     this._updateAgentState('listening');
   }

examples/src/basic_agent.ts

@@ -55,9 +55,6 @@ export default defineAgent({
       },
     });
 
-    // join the room when agent is ready
-    await ctx.connect();
-
     session.say('Hello, how can I help you today?');
   },
 });

examples/src/basic_eou.ts

@@ -9,6 +9,7 @@ export default defineAgent({
   entry: async (ctx: JobContext) => {
     const logger = log();
 
+    // Manual connection required since this example doesn't use AgentSession
     await ctx.connect();
 
     // const eouModel = new turnDetector.EnglishModel();

examples/src/basic_tool_call_agent.ts

@@ -148,8 +148,6 @@ export default defineAgent({
       agent: routerAgent,
       room: ctx.room,
     });
-
-    await ctx.connect();
   },
 });
 

examples/src/cartersia_tts.ts

@@ -55,9 +55,6 @@ export default defineAgent({
       },
     });
 
-    // join the room when agent is ready
-    await ctx.connect();
-
     session.say('Hello, how can I help you today?');
   },
 });

examples/src/comprehensive_test.ts

@@ -238,8 +238,6 @@ export default defineAgent({
     proc.userData.vad = await silero.VAD.load();
   },
   entry: async (ctx: JobContext) => {
-    await ctx.connect();
-
     const vad = ctx.proc.userData.vad! as silero.VAD;
     const session = new voice.AgentSession({
       vad,

examples/src/drive-thru/drivethru_agent.ts

@@ -380,8 +380,6 @@ export default defineAgent({
     proc.userData.vad = await silero.VAD.load();
   },
   entry: async (ctx: JobContext) => {
-    await ctx.connect();
-
     const userdata = await newUserData();
 
     const vad = ctx.proc.userData.vad! as silero.VAD;