voice: retry missing run outputs

.changeset/quiet-birds-retry.md

@@ -0,0 +1,5 @@
+---
+"@livekit/agents": patch
+---
+
+Add output retries for AgentSession.run structured outputs.

agents/src/_exceptions.ts

@@ -6,6 +6,17 @@
 // work well for synchronous throw/catch, but AbortSignal integrates better with async
 // streams, fetch, and the broader Web API ecosystem.
 
+/**
+ * Raised when the model behaves in a way the run cannot recover from.
+ */
+export class UnexpectedModelBehavior extends Error {
+  constructor(message: string, options?: ErrorOptions) {
+    super(message, options);
+    this.name = 'UnexpectedModelBehavior';
+    Error.captureStackTrace(this, UnexpectedModelBehavior);
+  }
+}
+
 /**
  * Raised when accepting a job but not receiving an assignment within the specified timeout.
  * The server may have chosen another worker to handle this job.

agents/src/voice/agent_session.ts

@@ -97,7 +97,7 @@ import { RoomSessionTransport, SessionHost } from './remote_session.js';
 import { RoomIO, type RoomInputOptions, type RoomOutputOptions } from './room_io/index.js';
 import type { UnknownUserData } from './run_context.js';
 import type { SpeechHandle } from './speech_handle.js';
-import { RunResult } from './testing/run_result.js';
+import { type RunOutputOptions, RunResult } from './testing/run_result.js';
 import {
   type AsyncToolOptions,
   type ToolHandlingOptions,
@@ -1113,17 +1113,23 @@ export class AgentSession<
    * result.expect.noMoreEvents();
    * ```
    *
-   * @param options - Run options including user input and optional output type
+   * @param options - Run options including user input and optional output type.
+   *   When `outputType` is set and the turn ends without structured output, the
+   *   run re-prompts the model up to `outputOptions.maxRetries` times (default 2)
+   *   before rejecting with `UnexpectedModelBehavior`. Pass `outputOptions: null`
+   *   to disable retries entirely.
    * @returns A RunResult that resolves when the agent finishes responding
    */
   run<T = unknown>({
     userInput,
     inputModality,
     outputType,
+    outputOptions,
   }: {
     userInput: string;
     inputModality?: 'audio' | 'text';
     outputType?: z.ZodType<T>;
+    outputOptions?: RunOutputOptions | null;
   }): RunResult<T> {
     if (this._globalRunState && !this._globalRunState.done()) {
       throw new Error('nested runs are not supported');
@@ -1132,6 +1138,8 @@ export class AgentSession<
     const runState = new RunResult<T>({
       userInput,
       outputType,
+      outputOptions,
+      session: this,
     });
 
     this._globalRunState = runState;

agents/src/voice/index.ts

@@ -51,4 +51,5 @@ export { RunContext } from './run_context.js';
 export * from './turn_config/endpointing.js';
 export * from './turn_config/user_turn_limit.js';
 export * as testing from './testing/index.js';
+export { type RunOutputOptions } from './testing/run_result.js';
 export * as textTransforms from './transcription/text_transforms.js';

agents/src/voice/testing/index.ts

@@ -33,6 +33,7 @@ export {
   withMockTools,
   type MockToolFn,
   type MockToolsMap,
+  type RunOutputOptions,
 } from './run_result.js';
 
 export {

agents/src/voice/testing/run_output_retry.test.ts

@@ -0,0 +1,178 @@
+// SPDX-FileCopyrightText: 2026 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+// Tests for the structured-output retry behavior of session.run({ outputType })
+// (ported from livekit/agents#6080): when a turn ends without the expected
+// output, the run re-prompts the model up to maxRetries times before rejecting
+// with UnexpectedModelBehavior.
+import { describe, expect, it } from 'vitest';
+import { z } from 'zod';
+import { UnexpectedModelBehavior } from '../../_exceptions.js';
+import type { ChatContext } from '../../llm/chat_context.js';
+import { FunctionCall } from '../../llm/chat_context.js';
+import { LLMStream as BaseLLMStream, LLM, type LLMStream } from '../../llm/llm.js';
+import { tool } from '../../llm/tool_context.js';
+import type { ToolChoice, ToolContextLike } from '../../llm/tool_context.js';
+import { initializeLogger } from '../../log.js';
+import { type APIConnectOptions, DEFAULT_API_CONNECT_OPTIONS } from '../../types.js';
+import { AgentTask } from '../agent.js';
+import { AgentSession } from '../agent_session.js';
+
+type ScriptedResponse = { content?: string; toolCall?: { name: string; args: object } };
+
+/** Returns the Nth scripted response on the Nth chat() call, ignoring input. */
+class ScriptedLLM extends LLM {
+  calls = 0;
+  systemTexts: string[] = [];
+
+  constructor(private script: ScriptedResponse[]) {
+    super();
+  }
+
+  label(): string {
+    return 'scripted-llm';
+  }
+
+  chat(params: {
+    chatCtx: ChatContext;
+    toolCtx?: ToolContextLike;
+    connOptions?: APIConnectOptions;
+    parallelToolCalls?: boolean;
+    toolChoice?: ToolChoice;
+    extraKwargs?: Record<string, unknown>;
+  }): LLMStream {
+    for (const item of params.chatCtx.items) {
+      if (item.type === 'message' && item.role === 'system') {
+        this.systemTexts.push(item.textContent ?? '');
+      }
+    }
+    const idx = Math.min(this.calls, this.script.length - 1);
+    this.calls += 1;
+    return new ScriptedLLMStream(this, this.script[idx]!, {
+      chatCtx: params.chatCtx,
+      toolCtx: params.toolCtx,
+      connOptions: params.connOptions ?? DEFAULT_API_CONNECT_OPTIONS,
+    });
+  }
+}
+
+class ScriptedLLMStream extends BaseLLMStream {
+  constructor(
+    llm: ScriptedLLM,
+    private response: ScriptedResponse,
+    params: { chatCtx: ChatContext; toolCtx?: ToolContextLike; connOptions: APIConnectOptions },
+  ) {
+    super(llm, params);
+  }
+
+  protected async run(): Promise<void> {
+    if (this.response.content) {
+      this.queue.put({
+        id: 'scripted',
+        delta: { role: 'assistant', content: this.response.content },
+      });
+    }
+    if (this.response.toolCall) {
+      this.queue.put({
+        id: 'scripted',
+        delta: {
+          role: 'assistant',
+          toolCalls: [
+            FunctionCall.create({
+              callId: 'scripted_call',
+              name: this.response.toolCall.name,
+              args: JSON.stringify(this.response.toolCall.args),
+            }),
+          ],
+        },
+      });
+    }
+  }
+}
+
+const outputSchema = z.object({ answer: z.string() });
+
+class OutputTask extends AgentTask<{ answer: string }> {
+  constructor() {
+    super({
+      instructions: 'Answer via the submit tool.',
+      tools: [
+        tool({
+          name: 'submit',
+          description: 'Submit the final answer.',
+          parameters: z.object({ answer: z.string() }),
+          execute: async ({ answer }) => {
+            this.complete({ answer });
+            return 'submitted';
+          },
+        }),
+      ],
+    });
+  }
+}
+
+async function runWith(
+  llm: ScriptedLLM,
+  outputOptions?: { maxRetries?: number; retryInstructions?: string } | null,
+) {
+  const session = new AgentSession({ llm });
+  await session.start({ agent: new OutputTask() });
+  try {
+    const run = session.run({
+      userInput: 'hi',
+      outputType: outputSchema,
+      ...(outputOptions !== undefined ? { outputOptions } : {}),
+    });
+    await Promise.race([
+      run.wait(),
+      new Promise((_, rej) => setTimeout(() => rej(new Error('run timed out')), 10_000)),
+    ]);
+    return { run, error: undefined as unknown };
+  } catch (error) {
+    return { run: undefined, error };
+  } finally {
+    await session.close().catch(() => {});
+  }
+}
+
+describe('session.run output retries', () => {
+  initializeLogger({ pretty: false, level: 'silent' });
+
+  it('retries and succeeds when the model calls the tool on the second turn', async () => {
+    const llm = new ScriptedLLM([
+      { content: 'I think the answer is 42.' }, // prose only -> retry
+      { toolCall: { name: 'submit', args: { answer: '42' } } },
+      { content: 'done' }, // reply to the tool output
+    ]);
+    const { run, error } = await runWith(llm);
+    expect(error).toBeUndefined();
+    expect(run!.finalOutput).toEqual({ answer: '42' });
+    expect(llm.calls).toBeGreaterThanOrEqual(2);
+  });
+
+  it('rejects with UnexpectedModelBehavior after exhausting retries', async () => {
+    const llm = new ScriptedLLM([{ content: 'still just prose' }]);
+    const { error } = await runWith(llm);
+    expect(error).toBeInstanceOf(UnexpectedModelBehavior);
+    // initial turn + 2 default retries
+    expect(llm.calls).toBe(3);
+  });
+
+  it('outputOptions: null disables retries and fails on the first miss', async () => {
+    const llm = new ScriptedLLM([{ content: 'prose' }]);
+    const { error } = await runWith(llm, null);
+    expect(error).toBeInstanceOf(UnexpectedModelBehavior);
+    expect(llm.calls).toBe(1);
+  });
+
+  it('honors maxRetries and custom retryInstructions', async () => {
+    const llm = new ScriptedLLM([{ content: 'prose forever' }]);
+    const retryInstructions = 'CUSTOM_RETRY_MARKER: call submit now.';
+    const { error } = await runWith(llm, { maxRetries: 1, retryInstructions });
+    expect(error).toBeInstanceOf(UnexpectedModelBehavior);
+    expect(llm.calls).toBe(2);
+    // the retry turn's per-turn instructions must reach the model
+    expect(llm.systemTexts.some((text) => text.includes('CUSTOM_RETRY_MARKER'))).toBe(true);
+  });
+});

agents/src/voice/testing/run_result.ts

@@ -2,13 +2,16 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 import { z } from 'zod';
+import { UnexpectedModelBehavior } from '../../_exceptions.js';
 import type { AgentHandoffItem, ChatItem, ChatRole } from '../../llm/chat_context.js';
 import { ChatContext } from '../../llm/chat_context.js';
 import type { LLM } from '../../llm/llm.js';
 import { tool } from '../../llm/tool_context.js';
+import { log } from '../../log.js';
 import type { Task } from '../../utils.js';
 import { Future } from '../../utils.js';
 import type { Agent } from '../agent.js';
+import type { AgentSession } from '../agent_session.js';
 import { type SpeechHandle, isSpeechHandle } from '../speech_handle.js';
 import {
   type AgentHandoffAssertOptions,
@@ -34,6 +37,24 @@ export type AgentConstructor = new (...args: any[]) => Agent;
 // In JS we use a zod schema so runtime validation and TS generic inference stay aligned.
 type OutputSchema<T> = z.ZodType<T>;
 
+const OUTPUT_RETRY_PROMPT =
+  'You have not provided the final output yet. Call the appropriate function ' +
+  'to do so; a plain text response alone is not enough.';
+
+/**
+ * Structured-output behavior for `AgentSession.run`.
+ *
+ * Pass `outputOptions: null` to `run()` to disable the retry behavior
+ * entirely (equivalent to `{ maxRetries: 0 }`); omitting the option uses the
+ * defaults below.
+ */
+export type RunOutputOptions = {
+  /** Re-prompts when a run ends without its output type. Defaults to 2. */
+  maxRetries?: number;
+  /** Override the built-in retry prompt. */
+  retryInstructions?: string;
+};
+
 // Environment variable for verbose output
 const evalsVerbose = parseInt(process.env.LIVEKIT_EVALS_VERBOSE || '0', 10);
 
@@ -52,6 +73,10 @@ export class RunResult<T = unknown> {
   private doneFut = new Future<void>();
   private userInput?: string;
   private outputType?: OutputSchema<T>;
+  private outputRetries: number;
+  private outputRetryInstructions: string;
+  private outputRetryError?: unknown;
+  private session?: AgentSession;
   private finalOutputValue?: T;
   private hasFinalOutput = false;
 
@@ -63,9 +88,19 @@ export class RunResult<T = unknown> {
 
   private readonly itemAddedCallback = (item: ChatItem) => this._itemAdded(item);
 
-  constructor(options?: { userInput?: string; outputType?: OutputSchema<T> }) {
+  constructor(options?: {
+    userInput?: string;
+    outputType?: OutputSchema<T>;
+    outputOptions?: RunOutputOptions | null;
+    session?: AgentSession;
+  }) {
     this.userInput = options?.userInput;
     this.outputType = options?.outputType;
+    const outputOptions =
+      options?.outputOptions === null ? { maxRetries: 0 } : options?.outputOptions;
+    this.outputRetries = outputOptions?.maxRetries ?? 2;
+    this.outputRetryInstructions = outputOptions?.retryInstructions ?? OUTPUT_RETRY_PROMPT;
+    this.session = options?.session;
   }
 
   /**
@@ -252,8 +287,18 @@ export class RunResult<T = unknown> {
     if (this.outputType) {
       const result = this.outputType.safeParse(finalOutput);
       if (!result.success) {
+        // Only a missing output is retryable. Unlike Python (where a task
+        // completed with None is indistinguishable from one that never
+        // completed), a task completed with null is one-shot — re-prompting
+        // cannot change its result, so it fails immediately.
+        if (finalOutput === undefined && this._maybeRetryOutput()) {
+          return;
+        }
         this.doneFut.reject(
-          new Error(`Expected output matching provided zod schema: ${result.error.message}`),
+          new UnexpectedModelBehavior(
+            `Expected output matching provided zod schema: ${result.error.message}`,
+            this.outputRetryError !== undefined ? { cause: this.outputRetryError } : undefined,
+          ),
         );
         return;
       }
@@ -270,6 +315,31 @@ export class RunResult<T = unknown> {
     this.doneFut.resolve();
   }
 
+  private _maybeRetryOutput(): boolean {
+    if (this.outputRetries <= 0 || !this.session) {
+      return false;
+    }
+    this.outputRetries -= 1;
+
+    try {
+      this.session.generateReply({ instructions: this.outputRetryInstructions });
+    } catch (error) {
+      // Fall through to UnexpectedModelBehavior; surface the real failure
+      // (e.g. a closing session) as the rejection's cause instead of hiding
+      // it behind a schema-mismatch message.
+      this.outputRetryError = error;
+      return false;
+    }
+
+    // zod schemas have no name; description (via .describe()) is the most
+    // useful label, with the schema class name as fallback.
+    log().warn(
+      { outputType: this.outputType?.description ?? this.outputType?.constructor?.name },
+      'run ended without the expected output type, retrying',
+    );
+    return true;
+  }
+
   /**
    * Find the correct insertion index to maintain chronological order.
    */