@open-mercato/ai-assistant 0.6.1-develop.3291.1.6fad645fd0 → 0.6.1

package/dist/modules/ai_assistant/lib/agent-runtime.js

@@ -1,13 +1,21 @@
+import { randomUUID } from "node:crypto";
 import { createContainer } from "awilix";
 import {
   convertToModelMessages,
   generateObject,
+  hasToolCall,
   stepCountIs,
   streamObject,
-  streamText
+  streamText,
+  Experimental_Agent as ToolLoopAgent
 } from "ai";
-import { createModelFactory } from "./model-factory.js";
-import { resolveAiAgentTools, AgentPolicyError } from "./agent-tools.js";
+import { createModelFactory, resolveAllowRuntimeOverride } from "./model-factory.js";
+import {
+  resolveAiAgentTools,
+  AgentPolicyError,
+  desanitizeToolNameForDisplay,
+  sanitizeToolNameForModel
+} from "./agent-tools.js";
 import { resolveEffectiveMutationPolicy } from "./agent-policy.js";
 import { toolRegistry } from "./tool-registry.js";
 import {
@@ -21,10 +29,377 @@ import { AiAgentRuntimeOverrideRepository } from "../data/repositories/AiAgentRu
 import { AiTenantModelAllowlistRepository } from "../data/repositories/AiTenantModelAllowlistRepository.js";
 import { composeSystemPromptWithOverride } from "./prompt-override-merge.js";
 import { isKnownMutationPolicy } from "./agent-policy.js";
+import { recordTokenUsage } from "./token-usage-recorder.js";
 import "./llm-bootstrap.js";
+const WRAPPER_DEFAULT_LOOP_CHAT = { maxSteps: 10 };
+const WRAPPER_DEFAULT_LOOP_OBJECT = {};
+function resolveLoopBudgetPreset(preset) {
+  switch (preset) {
+    case "tight":
+      return { maxSteps: 3, budget: { maxToolCalls: 3, maxWallClockMs: 1e4, maxTokens: 5e4 } };
+    case "loose":
+      return { maxSteps: 20, budget: { maxToolCalls: 20, maxWallClockMs: 12e4, maxTokens: 5e5 } };
+    case "default":
+      return void 0;
+  }
+}
+const SSE_ENCODER = new TextEncoder();
+function appendLoopFinishToStream(baseResponse, finalizeLoopTrace) {
+  const { readable, writable } = new TransformStream();
+  const writer = writable.getWriter();
+  async function pump() {
+    if (!baseResponse.body) {
+      await writer.close();
+      return;
+    }
+    const reader = baseResponse.body.getReader();
+    try {
+      for (; ; ) {
+        const { value, done } = await reader.read();
+        if (done) break;
+        await writer.write(value);
+      }
+      const trace = finalizeLoopTrace();
+      const eventLine = `data: ${JSON.stringify({ type: "loop-finish", trace })}
+
+`;
+      await writer.write(SSE_ENCODER.encode(eventLine));
+    } catch {
+    } finally {
+      reader.releaseLock();
+      await writer.close().catch(() => void 0);
+    }
+  }
+  void pump();
+  return new Response(readable, {
+    status: baseResponse.status,
+    headers: baseResponse.headers
+  });
+}
+function assertLoopRuntimeOverrideAllowed(agent, callerLoop) {
+  if (!callerLoop) return;
+  const allowed = agent.loop?.allowRuntimeOverride ?? true;
+  if (!allowed) {
+    throw new AgentPolicyError(
+      "loop_runtime_override_disabled",
+      `Agent "${agent.id}" has disabled per-call loop overrides (loop.allowRuntimeOverride: false). Remove the loop override to proceed.`
+    );
+  }
+}
+function readModuleLoopEnv(moduleId) {
+  const prefix = moduleId.toUpperCase();
+  const partial = {};
+  const maxStepsRaw = process.env[`${prefix}_AI_LOOP_MAX_STEPS`];
+  if (maxStepsRaw) {
+    const parsed = parseInt(maxStepsRaw.trim(), 10);
+    if (!isNaN(parsed) && parsed > 0) partial.maxSteps = parsed;
+  }
+  const maxWallClockRaw = process.env[`${prefix}_AI_LOOP_MAX_WALL_CLOCK_MS`];
+  const maxTokensRaw = process.env[`${prefix}_AI_LOOP_MAX_TOKENS`];
+  if (maxWallClockRaw || maxTokensRaw) {
+    const budgetPartial = {};
+    if (maxWallClockRaw) {
+      const parsed = parseInt(maxWallClockRaw.trim(), 10);
+      if (!isNaN(parsed) && parsed > 0) budgetPartial.maxWallClockMs = parsed;
+    }
+    if (maxTokensRaw) {
+      const parsed = parseInt(maxTokensRaw.trim(), 10);
+      if (!isNaN(parsed) && parsed > 0) budgetPartial.maxTokens = parsed;
+    }
+    if (Object.keys(budgetPartial).length > 0) partial.budget = budgetPartial;
+  }
+  return partial;
+}
+function resolveEffectiveLoopConfig(agent, callerLoop, wrapperDefault) {
+  assertLoopRuntimeOverrideAllowed(agent, callerLoop);
+  const effectiveDefault = wrapperDefault ?? WRAPPER_DEFAULT_LOOP_CHAT;
+  const legacyMaxSteps = typeof agent.maxSteps === "number" && agent.maxSteps > 0 && !agent.loop ? { maxSteps: agent.maxSteps } : void 0;
+  const base = {
+    ...effectiveDefault,
+    ...legacyMaxSteps ?? {},
+    ...agent.loop ?? {}
+  };
+  const envOverride = readModuleLoopEnv(agent.moduleId);
+  const withEnv = {
+    ...base,
+    ...envOverride,
+    ...envOverride.budget != null ? { budget: { ...base.budget ?? {}, ...envOverride.budget } } : {}
+  };
+  const withCaller = callerLoop ? { ...withEnv, ...callerLoop } : withEnv;
+  if (withCaller.disabled === true) {
+    return { ...withCaller, maxSteps: 1 };
+  }
+  return withCaller;
+}
+class BudgetEnforcer {
+  constructor(budget, abortController) {
+    this.budget = budget;
+    this.abortController = abortController;
+    this.toolCallsUsed = 0;
+    this.tokensUsed = 0;
+    this.abortReason = null;
+    this.turnStartMs = Date.now();
+  }
+  get hasActiveBudget() {
+    const b = this.budget;
+    return b !== void 0 && (b.maxToolCalls !== void 0 || b.maxWallClockMs !== void 0 || b.maxTokens !== void 0);
+  }
+  recordStep(usage) {
+    if (!this.budget) return;
+    this.toolCallsUsed += usage.toolCalls ?? 0;
+    this.tokensUsed += (usage.inputTokens ?? 0) + (usage.outputTokens ?? 0);
+    this.checkLimits();
+  }
+  checkLimits() {
+    const b = this.budget;
+    if (!b) return;
+    if (b.maxToolCalls !== void 0 && this.toolCallsUsed >= b.maxToolCalls) {
+      this.abort("budget-tool-calls");
+      return;
+    }
+    const elapsedMs = Date.now() - this.turnStartMs;
+    if (b.maxWallClockMs !== void 0 && elapsedMs >= b.maxWallClockMs) {
+      this.abort("budget-wall-clock");
+      return;
+    }
+    if (b.maxTokens !== void 0 && this.tokensUsed >= b.maxTokens) {
+      this.abort("budget-tokens");
+    }
+  }
+  abort(reason) {
+    if (this.abortReason !== null) return;
+    this.abortReason = reason;
+    console.info(
+      `[AI Agents] Budget exceeded \u2014 aborting turn. Reason: ${reason}. toolCalls=${this.toolCallsUsed}, tokens=${this.tokensUsed}, elapsedMs=${Date.now() - this.turnStartMs}.`
+    );
+    this.abortController.abort(reason);
+  }
+  wire(userOnStepFinish) {
+    if (!this.hasActiveBudget) return userOnStepFinish;
+    return async (event) => {
+      this.recordStep({
+        inputTokens: event.usage?.inputTokens,
+        outputTokens: event.usage?.outputTokens,
+        toolCalls: event.toolCalls?.length
+      });
+      if (userOnStepFinish) {
+        try {
+          await userOnStepFinish(event);
+        } catch (err) {
+          console.error("[AI Agents] User onStepFinish threw; ignoring:", err);
+        }
+      }
+    };
+  }
+}
+function buildLoopTraceCollector(agentId, sessionId, turnId, userOnStepFinish) {
+  const turnStartMs = Date.now();
+  const steps = [];
+  const onStepFinish = async (event) => {
+    const stepIndex = steps.length;
+    const toolCalls = (event.toolCalls ?? []).map((tc) => {
+      const raw = tc;
+      return {
+        toolName: raw.toolName ? desanitizeToolNameForDisplay(raw.toolName) : "unknown",
+        args: raw.args ?? {},
+        result: raw.result,
+        error: raw.experimental_toToolResultError ? {
+          code: String(raw.experimental_toToolResultError?.code ?? "unknown"),
+          message: String(raw.experimental_toToolResultError?.message ?? "")
+        } : void 0,
+        repairAttempted: raw.repairAttempted === true,
+        durationMs: typeof raw.startTime === "number" && typeof raw.endTime === "number" ? raw.endTime - raw.startTime : 0
+      };
+    });
+    const textDelta = event.text ?? "";
+    const finishReason = event.finishReason ?? "stop";
+    const modelId = event.response?.modelId ?? "unknown";
+    steps.push({
+      stepIndex,
+      modelId,
+      toolCalls,
+      textDelta,
+      usage: {
+        inputTokens: event.usage?.inputTokens ?? 0,
+        outputTokens: event.usage?.outputTokens ?? 0
+      },
+      finishReason
+    });
+    if (userOnStepFinish) {
+      try {
+        await userOnStepFinish(event);
+      } catch (err) {
+        console.error("[AI Agents] User onStepFinish in LoopTrace collector threw; ignoring:", err);
+      }
+    }
+  };
+  function finalize(abortReason) {
+    const totalDurationMs = Date.now() - turnStartMs;
+    const totalUsage = steps.reduce(
+      (acc, step) => ({
+        inputTokens: acc.inputTokens + step.usage.inputTokens,
+        outputTokens: acc.outputTokens + step.usage.outputTokens
+      }),
+      { inputTokens: 0, outputTokens: 0 }
+    );
+    let stopReason = "finish-reason";
+    if (abortReason === "budget-tool-calls") stopReason = "budget-tool-calls";
+    else if (abortReason === "budget-wall-clock") stopReason = "budget-wall-clock";
+    else if (abortReason === "budget-tokens") stopReason = "budget-tokens";
+    return {
+      agentId,
+      sessionId,
+      turnId,
+      steps,
+      stopReason,
+      totalDurationMs,
+      totalUsage
+    };
+  }
+  return { onStepFinish, finalize };
+}
+function translateStopConditions(loopConfig, mapToolName = (toolName) => toolName) {
+  const effectiveMaxSteps = loopConfig.maxSteps ?? 10;
+  const userConditions = [];
+  const rawStopWhen = loopConfig.stopWhen;
+  if (rawStopWhen) {
+    const items = Array.isArray(rawStopWhen) ? rawStopWhen : [rawStopWhen];
+    for (const item of items) {
+      if (item.kind === "stepCount") {
+        userConditions.push(stepCountIs(item.count));
+      } else if (item.kind === "hasToolCall") {
+        userConditions.push(hasToolCall(mapToolName(item.toolName)));
+      } else if (item.kind === "custom") {
+        userConditions.push(item.stop);
+      }
+    }
+  }
+  return [...userConditions, stepCountIs(effectiveMaxSteps)];
+}
+function normalizeAllowedToolNameForAgent(toolName, agent) {
+  if (agent.allowedTools.includes(toolName)) return toolName;
+  const dottedName = desanitizeToolNameForDisplay(toolName);
+  return agent.allowedTools.includes(dottedName) ? dottedName : null;
+}
+function mergeStepOverrides(wrapperOverride, userOverride, agent, wrappedToolRegistry) {
+  if (!userOverride) return wrapperOverride;
+  const merged = { ...wrapperOverride };
+  const userWithTools = userOverride;
+  if (userOverride.model !== void 0) {
+    merged.model = userOverride.model;
+  }
+  if (userOverride.toolChoice !== void 0) {
+    merged.toolChoice = userOverride.toolChoice;
+  }
+  if (userOverride.activeTools !== void 0) {
+    const filtered = userOverride.activeTools.flatMap((name) => {
+      const normalized = normalizeAllowedToolNameForAgent(name, agent);
+      const allowed = normalized !== null;
+      if (!allowed) {
+        console.warn(
+          `[AI Agents] loop:active_tools_filtered \u2014 tool "${name}" is not in agent "${agent.id}" allowedTools; dropping from activeTools.`
+        );
+      }
+      return normalized ? [normalized] : [];
+    });
+    merged.activeTools = filtered;
+  }
+  if (userWithTools.tools !== void 0) {
+    const userTools = userWithTools.tools;
+    const mergedTools = {};
+    for (const [toolKey, userHandler] of Object.entries(userTools)) {
+      const wrappedHandler = wrappedToolRegistry[toolKey];
+      if (!wrappedHandler) {
+        console.warn(
+          `[AI Agents] mergeStepOverrides \u2014 tool "${toolKey}" from user prepareStep is not in the wrapper tool registry; dropping.`
+        );
+        continue;
+      }
+      if (userHandler !== wrappedHandler) {
+        const toolDef = toolRegistry.getTool(
+          toolKey.replace(/__/g, ".")
+        );
+        if (toolDef?.isMutation === true) {
+          throw new AgentPolicyError(
+            "loop_violates_mutation_policy",
+            `User prepareStep returned a tools map with raw (unwrapped) mutation handler for "${toolKey}". This bypasses the mutation-approval gate and is rejected.`
+          );
+        }
+      }
+      mergedTools[toolKey] = wrappedHandler;
+    }
+    merged.tools = mergedTools;
+  }
+  return merged;
+}
+function mapPrepareStepResultForModel(result, wrappedTools) {
+  if (!result?.activeTools) return result;
+  const activeTools = result.activeTools.map((toolName) => sanitizeToolNameForModel(toolName)).filter((toolName) => wrappedTools[toolName] !== void 0);
+  return { ...result, activeTools };
+}
+function mapToolChoiceForModel(toolChoice) {
+  if (!toolChoice || typeof toolChoice !== "object" || toolChoice.type !== "tool") {
+    return toolChoice;
+  }
+  return {
+    ...toolChoice,
+    toolName: sanitizeToolNameForModel(toolChoice.toolName)
+  };
+}
+function mapActiveToolsForModel(activeTools, wrappedTools) {
+  if (!activeTools) return void 0;
+  return activeTools.map((toolName) => sanitizeToolNameForModel(toolName)).filter((toolName) => wrappedTools[toolName] !== void 0);
+}
+function buildWrapperPrepareStep(agent, effectiveLoop, wrappedTools) {
+  return async (state) => {
+    const wrapperOverride = {};
+    if (effectiveLoop.activeTools && effectiveLoop.activeTools.length > 0) {
+      wrapperOverride.activeTools = effectiveLoop.activeTools.flatMap((name) => {
+        const normalized = normalizeAllowedToolNameForAgent(name, agent);
+        const allowed = normalized !== null;
+        if (!allowed) {
+          console.warn(
+            `[AI Agents] loop:active_tools_filtered \u2014 tool "${name}" is not in agent "${agent.id}" allowedTools; dropping from activeTools.`
+          );
+        }
+        return normalized ? [normalized] : [];
+      });
+    }
+    if (effectiveLoop.prepareStep) {
+      let userOverride;
+      try {
+        userOverride = await effectiveLoop.prepareStep(state);
+      } catch (error) {
+        console.error(
+          `[AI Agents] User prepareStep threw for agent "${agent.id}"; ignoring user override:`,
+          error
+        );
+        return mapPrepareStepResultForModel(wrapperOverride, wrappedTools);
+      }
+      return mapPrepareStepResultForModel(
+        mergeStepOverrides(wrapperOverride, userOverride, agent, wrappedTools),
+        wrappedTools
+      );
+    }
+    return mapPrepareStepResultForModel(wrapperOverride, wrappedTools);
+  };
+}
+function assertLoopObjectModeCompatible(loopConfig) {
+  const unsupportedFields = [];
+  if (loopConfig.prepareStep !== void 0) unsupportedFields.push("prepareStep");
+  if (loopConfig.repairToolCall !== void 0) unsupportedFields.push("repairToolCall");
+  if (loopConfig.stopWhen !== void 0) unsupportedFields.push("stopWhen");
+  if (loopConfig.activeTools !== void 0) unsupportedFields.push("activeTools");
+  if (loopConfig.toolChoice !== void 0) unsupportedFields.push("toolChoice");
+  if (unsupportedFields.length > 0) {
+    throw new AgentPolicyError(
+      "loop_unsupported_in_object_mode",
+      `Object-mode agents do not support these loop primitives: ${unsupportedFields.join(", ")}. Use runAiAgentText for agents that require these loop controls.`
+    );
+  }
+}
 function resolveAgentModel(agent, modelOverride, providerOverride, container, baseUrlOverride, tenantOverride, requestOverride, tenantAllowlist) {
   const effectiveContainer = container ?? createContainer();
-  const allowRuntimeModelOverride = agent.allowRuntimeModelOverride !== false;
   const resolution = createModelFactory(effectiveContainer).resolveModel({
     moduleId: agent.moduleId,
     agentDefaultModel: agent.defaultModel,
@@ -33,7 +408,7 @@ function resolveAgentModel(agent, modelOverride, providerOverride, container, ba
     callerOverride: modelOverride,
     providerOverride,
     baseUrlOverride,
-    allowRuntimeModelOverride,
+    allowRuntimeOverride: resolveAllowRuntimeOverride(agent),
     tenantOverride: tenantOverride ?? void 0,
     requestOverride: requestOverride ?? void 0,
     tenantAllowlist: tenantAllowlist ?? null
@@ -353,6 +728,7 @@ async function runAiAgentText(input) {
       input.authContext.organizationId
     )
   ]);
+  const effectiveSessionId = (input.sessionId ?? input.conversationId) || randomUUID();
   const { agent, tools } = await resolveAiAgentTools({
     agentId: input.agentId,
     authContext: input.authContext,
@@ -360,7 +736,7 @@ async function runAiAgentText(input) {
     attachmentIds: input.attachmentIds,
     mutationPolicyOverride,
     container: input.container,
-    conversationId: input.conversationId ?? null
+    conversationId: effectiveSessionId
   });
   const resolvedAttachments = await resolveAttachmentPartsForAgent({
     agent,
@@ -380,7 +756,7 @@ async function runAiAgentText(input) {
     agent,
     mutationPolicyOverride
   );
-  const { model } = resolveAgentModel(
+  const resolvedModel = resolveAgentModel(
     agent,
     input.modelOverride,
     input.providerOverride,
@@ -390,26 +766,162 @@ async function runAiAgentText(input) {
     input.requestOverride,
     tenantAllowlistSnapshot
   );
+  const { model } = resolvedModel;
   const normalizedMessages = ensureUiMessageShape(input.messages);
   const hydratedMessages = attachAttachmentsToMessages(normalizedMessages, resolvedAttachments);
   const modelMessages = await convertToModelMessages(hydratedMessages);
-  const effectiveMaxSteps = typeof agent.maxSteps === "number" && agent.maxSteps > 0 ? agent.maxSteps : 10;
-  const stopWhen = stepCountIs(effectiveMaxSteps);
-  const streamArgs = {
+  const effectiveLoop = resolveEffectiveLoopConfig(agent, input.loop, WRAPPER_DEFAULT_LOOP_CHAT);
+  const stopConditions = translateStopConditions(effectiveLoop, sanitizeToolNameForModel);
+  const wrapperPrepareStep = buildWrapperPrepareStep(agent, effectiveLoop, tools);
+  const sdkActiveTools = mapActiveToolsForModel(effectiveLoop.activeTools, tools);
+  const sdkToolChoice = mapToolChoiceForModel(effectiveLoop.toolChoice);
+  const turnId = randomUUID();
+  const loopTraceCollector = buildLoopTraceCollector(agent.id, effectiveSessionId, turnId, effectiveLoop.onStepFinish);
+  const abortController = new AbortController();
+  const budgetEnforcer = new BudgetEnforcer(effectiveLoop.budget, abortController);
+  const tracedOnStepFinish = budgetEnforcer.wire(loopTraceCollector.onStepFinish);
+  let currentStepIndex = 0;
+  const wiredOnStepFinish = async (event) => {
+    const capturedStepIndex = currentStepIndex;
+    currentStepIndex += 1;
+    if (tracedOnStepFinish) {
+      await tracedOnStepFinish(event);
+    }
+    if (input.container) {
+      const rawEvent = event;
+      void recordTokenUsage(
+        {
+          authContext: input.authContext,
+          agentId: agent.id,
+          moduleId: agent.moduleId,
+          sessionId: effectiveSessionId,
+          turnId,
+          stepIndex: capturedStepIndex,
+          providerId: resolvedModel.providerId,
+          modelId: resolvedModel.modelId,
+          usage: {
+            inputTokens: rawEvent.usage?.inputTokens,
+            outputTokens: rawEvent.usage?.outputTokens,
+            cachedInputTokens: rawEvent.usage?.cachedInputTokens,
+            reasoningTokens: rawEvent.usage?.reasoningTokens
+          },
+          finishReason: rawEvent.finishReason,
+          loopAbortReason: budgetEnforcer.abortReason ?? void 0
+        },
+        input.container
+      );
+    }
+  };
+  let wallClockTimer;
+  if (effectiveLoop.budget?.maxWallClockMs) {
+    wallClockTimer = setTimeout(() => {
+      budgetEnforcer.recordStep({ toolCalls: 0 });
+    }, effectiveLoop.budget.maxWallClockMs);
+  }
+  let builtToolLoopAgent;
+  if (agent.executionEngine === "tool-loop-agent") {
+    const agentSettings = {
+      model,
+      tools,
+      stopWhen: stopConditions,
+      prepareStep: wrapperPrepareStep,
+      onStepFinish: wiredOnStepFinish,
+      ...effectiveLoop.repairToolCall !== void 0 ? { experimental_repairToolCall: effectiveLoop.repairToolCall } : {},
+      ...sdkActiveTools !== void 0 ? { activeTools: sdkActiveTools } : {},
+      ...sdkToolChoice !== void 0 ? { toolChoice: sdkToolChoice } : {}
+    };
+    builtToolLoopAgent = new ToolLoopAgent(agentSettings);
+  }
+  const preparedOptions = {
     model,
+    tools,
     system: systemPrompt,
     messages: modelMessages,
-    tools,
-    stopWhen
+    maxSteps: effectiveLoop.maxSteps ?? 10,
+    stopWhen: stopConditions,
+    prepareStep: wrapperPrepareStep,
+    onStepFinish: wiredOnStepFinish,
+    onStepStart: effectiveLoop.onStepStart,
+    onToolCallStart: effectiveLoop.onToolCallStart,
+    onToolCallFinish: effectiveLoop.onToolCallFinish,
+    experimental_repairToolCall: effectiveLoop.repairToolCall,
+    activeTools: sdkActiveTools,
+    toolChoice: sdkToolChoice,
+    abortSignal: abortController.signal,
+    finalizeLoopTrace: () => loopTraceCollector.finalize(budgetEnforcer.abortReason),
+    ...builtToolLoopAgent !== void 0 ? { toolLoopAgent: builtToolLoopAgent } : {}
   };
-  const result = streamText(streamArgs);
-  return result.toUIMessageStreamResponse({
+  if (input.generateText) {
+    try {
+      const callbackResult = await input.generateText(preparedOptions);
+      const baseResponse2 = callbackResult.toUIMessageStreamResponse({
+        sendReasoning: true,
+        headers: {
+          "Cache-Control": "no-cache, no-transform",
+          Connection: "keep-alive"
+        }
+      });
+      if (input.emitLoopTrace) {
+        return appendLoopFinishToStream(baseResponse2, preparedOptions.finalizeLoopTrace);
+      }
+      return baseResponse2;
+    } finally {
+      if (wallClockTimer !== void 0) clearTimeout(wallClockTimer);
+    }
+  }
+  if (builtToolLoopAgent !== void 0) {
+    const agentStreamResult = await builtToolLoopAgent.stream({
+      messages: modelMessages,
+      abortSignal: abortController.signal,
+      onStepFinish: wiredOnStepFinish
+    });
+    if (wallClockTimer !== void 0) {
+      const clearTimer = () => clearTimeout(wallClockTimer);
+      Promise.resolve(agentStreamResult.consumeStream()).then(clearTimer, clearTimer);
+    }
+    const baseResponse2 = agentStreamResult.toUIMessageStreamResponse({
+      sendReasoning: true,
+      headers: {
+        "Cache-Control": "no-cache, no-transform",
+        Connection: "keep-alive"
+      }
+    });
+    if (input.emitLoopTrace) {
+      return appendLoopFinishToStream(baseResponse2, preparedOptions.finalizeLoopTrace);
+    }
+    return baseResponse2;
+  }
+  const result = streamText({
+    model,
+    system: systemPrompt,
+    messages: modelMessages,
+    tools,
+    stopWhen: stopConditions,
+    prepareStep: wrapperPrepareStep,
+    onStepFinish: wiredOnStepFinish,
+    experimental_onStepStart: effectiveLoop.onStepStart,
+    experimental_onToolCallStart: effectiveLoop.onToolCallStart,
+    experimental_onToolCallFinish: effectiveLoop.onToolCallFinish,
+    experimental_repairToolCall: effectiveLoop.repairToolCall,
+    ...sdkActiveTools !== void 0 ? { activeTools: sdkActiveTools } : {},
+    ...sdkToolChoice !== void 0 ? { toolChoice: sdkToolChoice } : {},
+    abortSignal: abortController.signal
+  });
+  if (wallClockTimer !== void 0) {
+    const clearTimer = () => clearTimeout(wallClockTimer);
+    Promise.resolve(result.consumeStream()).then(clearTimer, clearTimer);
+  }
+  const baseResponse = result.toUIMessageStreamResponse({
     sendReasoning: true,
     headers: {
       "Cache-Control": "no-cache, no-transform",
       Connection: "keep-alive"
     }
   });
+  if (input.emitLoopTrace) {
+    return appendLoopFinishToStream(baseResponse, preparedOptions.finalizeLoopTrace);
+  }
+  return baseResponse;
 }
 function normalizeObjectMessages(input) {
   if (typeof input === "string") {
@@ -508,14 +1020,60 @@ async function runAiAgentObject(input) {
     resolvedAttachments
   );
   const modelMessages = await convertToModelMessages(hydratedMessages);
-  const stopWhen = typeof agent.maxSteps === "number" && agent.maxSteps > 0 ? stepCountIs(agent.maxSteps) : void 0;
+  void tools;
+  const effectiveObjectSessionId = input.sessionId ?? randomUUID();
+  const objectTurnId = randomUUID();
+  void effectiveObjectSessionId;
+  void objectTurnId;
+  if (input.loop) {
+    assertLoopObjectModeCompatible(input.loop);
+  }
+  const effectiveLoop = resolveEffectiveLoopConfig(agent, input.loop, WRAPPER_DEFAULT_LOOP_OBJECT);
+  const abortController = new AbortController();
+  const preparedObjectOptions = {
+    model,
+    system: systemPrompt,
+    messages: modelMessages,
+    schemaName: resolvedOutput.schemaName,
+    schema: resolvedOutput.schema,
+    maxSteps: effectiveLoop.maxSteps,
+    onStepFinish: effectiveLoop.onStepFinish,
+    onStepStart: effectiveLoop.onStepStart,
+    abortSignal: abortController.signal
+  };
+  if (input.generateObject) {
+    const callbackResult = await input.generateObject(preparedObjectOptions);
+    const typedResult = callbackResult;
+    if ("partialObjectStream" in typedResult) {
+      const streamResult = typedResult;
+      return {
+        mode: "stream",
+        object: streamResult.object,
+        partialObjectStream: streamResult.partialObjectStream,
+        textStream: streamResult.textStream,
+        finishReason: streamResult.finishReason,
+        usage: streamResult.usage
+      };
+    }
+    const genResult = typedResult;
+    return {
+      mode: "generate",
+      object: genResult.object,
+      finishReason: genResult.finishReason,
+      usage: genResult.usage
+    };
+  }
   if (resolvedOutput.mode === "stream") {
     const streamArgs = {
       model,
       system: systemPrompt,
       messages: modelMessages,
       schema: resolvedOutput.schema,
-      schemaName: resolvedOutput.schemaName
+      schemaName: resolvedOutput.schemaName,
+      ...effectiveLoop.maxSteps !== void 0 ? { maxSteps: effectiveLoop.maxSteps } : {},
+      onStepFinish: effectiveLoop.onStepFinish,
+      onStepStart: effectiveLoop.onStepStart,
+      abortSignal: abortController.signal
     };
     const result2 = streamObject(streamArgs);
     return {
@@ -532,13 +1090,12 @@ async function runAiAgentObject(input) {
     system: systemPrompt,
     messages: modelMessages,
     schema: resolvedOutput.schema,
-    schemaName: resolvedOutput.schemaName
+    schemaName: resolvedOutput.schemaName,
+    ...effectiveLoop.maxSteps !== void 0 ? { maxSteps: effectiveLoop.maxSteps } : {},
+    onStepFinish: effectiveLoop.onStepFinish,
+    onStepStart: effectiveLoop.onStepStart,
+    abortSignal: abortController.signal
   };
-  if (stopWhen) {
-    ;
-    generateArgs.stopWhen = stopWhen;
-  }
-  void tools;
   const result = await generateObject(generateArgs);
   return {
     mode: "generate",
@@ -549,8 +1106,16 @@ async function runAiAgentObject(input) {
 }
 export {
   AgentPolicyError,
+  BudgetEnforcer,
+  assertLoopObjectModeCompatible,
+  buildLoopTraceCollector,
+  buildWrapperPrepareStep,
   composeSystemPrompt,
+  mergeStepOverrides,
+  resolveEffectiveLoopConfig,
+  resolveLoopBudgetPreset,
   runAiAgentObject,
-  runAiAgentText
+  runAiAgentText,
+  translateStopConditions
 };
 //# sourceMappingURL=agent-runtime.js.map