@@ -0,0 +1,107 @@
+import { tool } from "ai";
+import { MockLanguageModelV3 } from "ai/test";
+import { expect, it, vi } from "vitest";
+import { z } from "zod";
+import { hasSuccessfulSearchValidation } from "../../tools/support/search-events/agent";
+import { callEmbeddedAgent } from "./callEmbeddedAgent";
+import { getAgentProvider } from "./provider-factory";
+
+vi.mock("./provider-factory", () => ({ getAgentProvider: vi.fn() }));
+vi.mock("../../telem/logging", () => ({
+  logIssue: vi.fn(),
+  logWarn: vi.fn(),
+}));
+
+const usage = {
+  inputTokens: {
+    total: 1,
+    noCache: 1,
+    cacheRead: 0,
+    cacheWrite: 0,
+  },
+  outputTokens: { total: 1, text: 1, reasoning: 0 },
+};
+
+function setupAgent(
+  validationPassesOnLastStep: boolean,
+  finalStepValidationCount = 1,
+) {
+  let callCount = 0;
+  const model = new MockLanguageModelV3({
+    doGenerate: async () => {
+      const step = ++callCount;
+      if (step <= 5) {
+        return {
+          content: Array.from(
+            { length: step === 5 ? finalStepValidationCount : 1 },
+            (_, candidate) => ({
+              type: "tool-call" as const,
+              toolCallId: `validation-${step}-${candidate}`,
+              toolName: "validateSearch",
+              input: JSON.stringify({ step, candidate }),
+            }),
+          ),
+          finishReason: { unified: "tool-calls" as const, raw: undefined },
+          usage,
+          warnings: [],
+        };
+      }
+      return {
+        content: [{ type: "text" as const, text: '{"result":"ready"}' }],
+        finishReason: { unified: "stop" as const, raw: undefined },
+        usage,
+        warnings: [],
+      };
+    },
+  });
+  vi.mocked(getAgentProvider).mockReturnValue({
+    type: "openai",
+    label: "mock",
+    getModel: () => model,
+    getProviderOptions: () => ({}),
+  });
+
+  const validateSearch = tool({
+    inputSchema: z.object({ step: z.number(), candidate: z.number() }),
+    execute: async ({ step }) => ({
+      result: { valid: validationPassesOnLastStep && step === 5 },
+    }),
+  });
+
+  return { model, validateSearch };
+}
+
+it("returns a structured answer after three validations pass on the fifth step", async () => {
+  const { model, validateSearch } = setupAgent(true, 3);
+
+  const result = await callEmbeddedAgent({
+    system: "Translate the query",
+    prompt: "Test query",
+    tools: { validateSearch },
+    schema: z.object({ result: z.string() }),
+    isReadyToFinalize: ({ toolResults }) =>
+      hasSuccessfulSearchValidation(toolResults),
+  });
+
+  expect(result.result).toEqual({ result: "ready" });
+  expect(result.toolCalls).toHaveLength(7);
+  expect(model.doGenerateCalls).toHaveLength(6);
+  expect(model.doGenerateCalls[5]?.toolChoice).toEqual({ type: "none" });
+});
+
+it("stops after five steps when the final validation fails", async () => {
+  const { model, validateSearch } = setupAgent(false);
+
+  await expect(
+    callEmbeddedAgent({
+      system: "Translate the query",
+      prompt: "Test query",
+      tools: { validateSearch },
+      schema: z.object({ result: z.string() }),
+      isReadyToFinalize: ({ toolResults }) =>
+        hasSuccessfulSearchValidation(toolResults),
+    }),
+  ).rejects.toThrow("could not construct a valid query");
+
+  expect(model.doGenerateCalls).toHaveLength(5);
+});
@@ -1,21 +1,23 @@
-import { describe, it, expect, vi, beforeEach } from "vitest";
 import {
-  generateText,
   APICallError,
+  generateText,
+  type LanguageModelUsage,
   NoObjectGeneratedError,
   NoOutputGeneratedError,
   RetryError,
-  type LanguageModelUsage,
+  type StepResult,
+  type Tool,
 } from "ai";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 import { z } from "zod";
-import { callEmbeddedAgent } from "./callEmbeddedAgent";
 import {
   AgentExecutionError,
   ConfigurationError,
   LLMProviderError,
   UserInputError,
 } from "../../errors";
 import { logIssue, logWarn } from "../../telem/logging";
+import { callEmbeddedAgent } from "./callEmbeddedAgent";
 import { getAgentProvider } from "./provider-factory";
 
 vi.mock("../../telem/logging", () => ({
@@ -66,6 +68,31 @@ describe("callEmbeddedAgent", () => {
     process.env.OPENROUTER_API_KEY = "";
   });
 
+  it("keeps the five-step limit when no finalization condition is supplied", async () => {
+    mockGenerateText.mockResolvedValue({
+      experimental_output: { result: "ok" },
+    } as never);
+
+    await callEmbeddedAgent({
+      system: "You are a test agent",
+      prompt: "Test prompt",
+      tools: {},
+      schema: testSchema,
+    });
+
+    const stopWhen = mockGenerateText.mock.calls[0]?.[0]?.stopWhen;
+    if (typeof stopWhen !== "function") {
+      throw new Error("Expected a stop condition");
+    }
+
+    const step = { toolResults: [{ output: false }] } as StepResult<
+      Record<string, Tool>
+    >;
+    const firstFour = Array.from({ length: 4 }, () => step);
+    expect(stopWhen({ steps: firstFour })).toBe(false);
+    expect(stopWhen({ steps: [...firstFour, step] })).toBe(true);
+  });
+
   it("throws LLMProviderError for OpenAI region restriction", async () => {
     // Create an APICallError simulating OpenAI's region restriction
     const regionError = new APICallError({
 
@@ -5,7 +5,7 @@ import {
   NoOutputGeneratedError,
   Output,
   RetryError,
-  stepCountIs,
+  type StepResult,
   type Tool,
 } from "ai";
 import type { z } from "zod";
@@ -18,6 +18,9 @@ import {
 import { logIssue, logWarn } from "../../telem/logging";
 import { getAgentProvider } from "./provider-factory";
 
+// Tool work gets this many model steps; a validated final response may get one more.
+const MAX_WORK_STEPS = 5;
+
 /**
  * Resolve the underlying provider failure from an AI SDK error.
  *
@@ -112,14 +115,25 @@ export async function callEmbeddedAgent<
   prompt,
   tools,
   schema,
+  isReadyToFinalize,
 }: {
   system: string;
   prompt: string;
   tools: Record<string, Tool>;
   schema: TSchema;
+  /** Whether this step produced a result that permits a final response. */
+  isReadyToFinalize?: (step: StepResult<Record<string, Tool>>) => boolean;
 }): Promise<EmbeddedAgentResult<TOutput>> {
   const capturedToolCalls: ToolCall[] = [];
 
+  function shouldAllowFinalResponseStep(
+    steps: StepResult<Record<string, Tool>>[],
+  ): boolean {
+    if (steps.length !== MAX_WORK_STEPS) return false;
+    const lastStep = steps.at(-1);
+    return lastStep !== undefined && (isReadyToFinalize?.(lastStep) ?? false);
+  }
+
   // Get the configured provider (OpenAI, Azure OpenAI, or Anthropic)
   const provider = getAgentProvider();
 
@@ -129,7 +143,12 @@ export async function callEmbeddedAgent<
       system,
       prompt,
       tools,
-      stopWhen: stepCountIs(5),
+      stopWhen: ({ steps }) =>
+        steps.length >= MAX_WORK_STEPS && !shouldAllowFinalResponseStep(steps),
+      prepareStep: ({ steps }) => {
+        if (shouldAllowFinalResponseStep(steps)) return { toolChoice: "none" };
+        return undefined;
+      },
       experimental_output: Output.object({ schema }),
       experimental_telemetry: {
         isEnabled: true,
 
@@ -0,0 +1,41 @@
+import { expect, it } from "vitest";
+import { hasSuccessfulSearchValidation } from "./agent";
+
+it("identifies a passing validateSearch result", () => {
+  expect(hasSuccessfulSearchValidation([])).toBe(false);
+  expect(
+    hasSuccessfulSearchValidation([
+      { toolName: "datasetAttributes", output: { result: {} } },
+    ]),
+  ).toBe(false);
+  expect(
+    hasSuccessfulSearchValidation([
+      { toolName: "validateSearch", output: { result: { valid: false } } },
+    ]),
+  ).toBe(false);
+  expect(
+    hasSuccessfulSearchValidation([
+      { toolName: "validateSearch", output: { error: "API unavailable" } },
+    ]),
+  ).toBe(false);
+  expect(
+    hasSuccessfulSearchValidation([
+      { toolName: "validateSearch", output: { result: { valid: true } } },
+    ]),
+  ).toBe(true);
+});
+
+it("accepts any successful validation among candidates checked in one step", () => {
+  const passing = {
+    toolName: "validateSearch",
+    output: { result: { valid: true } },
+  };
+  const failing = {
+    toolName: "validateSearch",
+    output: { result: { valid: false } },
+  };
+
+  expect(hasSuccessfulSearchValidation([passing, failing])).toBe(true);
+  expect(hasSuccessfulSearchValidation([failing, passing])).toBe(true);
+  expect(hasSuccessfulSearchValidation([failing, failing])).toBe(false);
+});
@@ -15,6 +15,20 @@ import {
 } from "./utils";
 
 const SEARCH_EVENTS_DATASETS = [...PUBLIC_EVENTS_DATASETS, "replays"] as const;
+const successfulValidateSearchOutputSchema = z.object({
+  result: z.object({ valid: z.literal(true) }),
+});
+
+export function hasSuccessfulSearchValidation(
+  toolResults: readonly { toolName: string; output: unknown }[],
+): boolean {
+  // A step may validate several candidates; any successful one qualifies.
+  return toolResults.some(
+    ({ toolName, output }) =>
+      toolName === "validateSearch" &&
+      successfulValidateSearchOutputSchema.safeParse(output).success,
+  );
+}
 
 // .default("") on explanation is safe because structuredOutputs: false is set via providerOptions.
 // If structuredOutputs is re-enabled, remove .default() calls (OpenAI requires all fields in 'required').
@@ -228,5 +242,7 @@ export async function searchEventsAgent(
       whoami: whoamiTool,
     },
     schema: searchEventsAgentOutputSchema,
+    isReadyToFinalize: ({ toolResults }) =>
+      hasSuccessfulSearchValidation(toolResults),
   });
 }