src/agents/embedded-agent-runner/run/attempt-normalization.ts

import { projectAgentRunAttemptTerminal } from "../../agent-run-terminal-outcome.js";
import { formatAssistantErrorText } from "../../embedded-agent-helpers.js";
import { normalizeUsage, type UsageLike } from "../../usage.js";
import { hasOutboundDeliveryEvidence } from "../delivery-evidence.js";
import { log } from "../logger.js";
import { createEmbeddedRunReplayState, observeReplayMetadata } from "../replay-state.js";
import type { EmbeddedAgentRunResult } from "../types.js";
import type { createUsageAccumulator } from "../usage-accumulator.js";
import {
  mergeAttemptRunStatsIntoAccumulator,
  mergeUsageIntoAccumulator,
} from "../usage-accumulator.js";
import { applyEmbeddedAttemptSessionIdentity } from "./attempt-session-identity.js";
import { resolveCurrentAttemptAssistant } from "./attempt-terminal-evidence.js";
import type { createEmbeddedRunContextRecoveryState } from "./context-recovery-state.js";
import type { PreparedEmbeddedRunInput } from "./execution-context.js";
import { resolveRunFailoverDecision } from "./failover-policy.js";
import {
  buildErrorAgentMeta,
  normalizeAssistantUsageForContext,
  resolveActiveErrorContext,
  resolveLatestCallUsage,
} from "./helpers.js";
import {
  MAX_CONSECUTIVE_IDLE_TIMEOUTS_BEFORE_OUTPUT,
  stepIdleTimeoutBreaker,
  type createIdleTimeoutBreakerState,
} from "./idle-timeout-breaker.js";
import { resolveReplayInvalidFlag } from "./incomplete-turn-resolution.js";
import { resolveRunRetryKind, type RunRetryKind } from "./retry-budget.js";
import { handleRetryLimitExhaustion } from "./retry-limit.js";
import type { prepareAndDispatchEmbeddedRunAttempt } from "./run-attempt-dispatch.js";
import {
  hasCompletedModelProgressForIdleBreaker,
  normalizeEmbeddedRunAttemptResult,
} from "./run-attempt-result.js";
import type { prepareEmbeddedRunRuntime } from "./runtime-preparation.js";
import type { createEmbeddedRunSessionPromptState } from "./session-prompt-state.js";
import {
  isEmbeddedRunTerminalAbort,
  isEmbeddedRunTerminalInterrupted,
  isEmbeddedRunTerminalTimeout,
  resolveEmbeddedRunAttemptTerminalState,
} from "./terminal-outcome.js";

type PreparedRuntime = Awaited<ReturnType<typeof prepareEmbeddedRunRuntime>>;
type SessionPromptState = Awaited<ReturnType<typeof createEmbeddedRunSessionPromptState>>;

type ReplayState = ReturnType<typeof createEmbeddedRunReplayState>;

export async function normalizeEmbeddedRunAttempt(input: {
  runInput: PreparedEmbeddedRunInput;
  preparedRuntime: PreparedRuntime;
  dispatchedAttempt: Awaited<
    ReturnType<typeof prepareAndDispatchEmbeddedRunAttempt>
  >["dispatchedAttempt"];
  sessionPromptState: SessionPromptState;
  provider: string;
  modelId: string;
  bootstrapPromptWarningSignaturesSeen: string[];
  usageAccumulator: ReturnType<typeof createUsageAccumulator>;
  lastRunPromptUsage: ReturnType<typeof normalizeUsage> | undefined;
  idleTimeoutBreakerState: ReturnType<typeof createIdleTimeoutBreakerState>;
  contextRecoveryState: ReturnType<typeof createEmbeddedRunContextRecoveryState>;
  recordedCompactionCount?: number;
  replayState: ReplayState;
  lastRetryFailoverReason: Parameters<typeof resolveRunFailoverDecision>[0]["failoverReason"];
}): Promise<
  | { action: "complete"; result: EmbeddedAgentRunResult }
  | {
      action: "retry";
      retryKind: RunRetryKind;
      bootstrapPromptWarningSignaturesSeen: string[];
      lastRunPromptUsage: ReturnType<typeof normalizeUsage> | undefined;
      replayState: ReplayState;
    }
  | {
      action: "proceed";
      bootstrapPromptWarningSignaturesSeen: string[];
      lastRunPromptUsage: ReturnType<typeof normalizeUsage> | undefined;
      replayState: ReplayState;
      attempt: ReturnType<typeof normalizeEmbeddedRunAttemptResult>;
      sessionIdUsed: string;
      sessionFileUsed: string | undefined;
      currentAttemptAssistant: ReturnType<
        typeof normalizeEmbeddedRunAttemptResult
      >["currentAttemptAssistant"];
      currentAttemptCompletedAssistant: ReturnType<
        typeof normalizeEmbeddedRunAttemptResult
      >["currentAttemptCompletedAssistant"];
      attemptAssistant: ReturnType<
        typeof normalizeEmbeddedRunAttemptResult
      >["currentAttemptAssistant"];
      terminalState: ReturnType<typeof resolveEmbeddedRunAttemptTerminalState>;
      setTerminalLifecycleMeta: NonNullable<
        ReturnType<typeof normalizeEmbeddedRunAttemptResult>["setTerminalLifecycleMeta"]
      >;
      attemptCompactionCount: number;
      activeErrorContext: ReturnType<typeof resolveActiveErrorContext>;
      resolveReplayInvalidForAttempt: (incompleteTurnText?: string | null) => boolean;
      assistantErrorText: string | undefined;
      canRestartForLiveSwitch: boolean;
    }
> {
  const { runInput, preparedRuntime, dispatchedAttempt, sessionPromptState, provider, modelId } =
    input;
  const params = runInput.runParams;
  const runtime = preparedRuntime.snapshot();
  const attempt = normalizeEmbeddedRunAttemptResult(dispatchedAttempt.rawAttempt);
  // Completed attempt facts must survive cancellation while user persistence settles.
  const attemptCompactionCount = Math.max(0, attempt.compactionCount ?? 0);
  const recordedCompactionCount = input.recordedCompactionCount ?? 0;
  input.contextRecoveryState.autoCompactionCount += Math.max(
    0,
    attemptCompactionCount - recordedCompactionCount,
  );
  if (attemptCompactionCount > recordedCompactionCount) {
    // Returned counts carry no ordering relative to model observations.
    input.contextRecoveryState.currentContextSnapshot = { tokens: undefined };
  }
  if (
    (recordedCompactionCount === 0 || attemptCompactionCount > recordedCompactionCount) &&
    typeof attempt.compactionTokensAfter === "number" &&
    Number.isFinite(attempt.compactionTokensAfter) &&
    attempt.compactionTokensAfter >= 0
  ) {
    input.contextRecoveryState.lastCompactionTokensAfter = Math.floor(
      attempt.compactionTokensAfter,
    );
  }
  await sessionPromptState.waitForCurrentUserMessagePersistence();
  sessionPromptState.suppressNextUserMessagePersistence = sessionPromptState.activePrompt.persisted;
  // Parent Stop can revoke attempt callbacks before they run. The lane signal,
  // not a callback-derived latch, owns cancellation after persistence settles.
  runInput.laneController.throwIfAborted();
  const {
    terminal,
    preflightRecovery,
    sessionIdUsed,
    sessionFileUsed,
    lastAssistant: sessionLastAssistant,
    currentAttemptAssistant,
    currentAttemptCompletedAssistant,
  } = attempt;
  const { idleTimedOut } = projectAgentRunAttemptTerminal(terminal);
  const attemptAssistant = resolveCurrentAttemptAssistant(attempt);
  const terminalState = resolveEmbeddedRunAttemptTerminalState({
    attempt,
    assistant: attemptAssistant,
    abortSignal: params.abortSignal,
  });
  const { outcome: terminalOutcome, signalOwnedInterruption } = terminalState;
  const terminalAborted = isEmbeddedRunTerminalAbort(terminalOutcome);
  const terminalTimedOut = isEmbeddedRunTerminalTimeout(terminalOutcome);
  const terminalInterrupted = isEmbeddedRunTerminalInterrupted(terminalOutcome);
  const setTerminalLifecycleMeta: NonNullable<typeof attempt.setTerminalLifecycleMeta> = (meta) => {
    const { stopReason, ...remainingMeta } = meta;
    const terminalStopReason = terminalInterrupted ? terminalOutcome.stopReason : stopReason;
    attempt.setTerminalLifecycleMeta?.({
      ...remainingMeta,
      ...(terminalStopReason ? { stopReason: terminalStopReason } : {}),
      aborted: terminalAborted,
    });
  };
  applyEmbeddedAttemptSessionIdentity({ sessionPromptState, sessionFileUsed, sessionIdUsed });
  const bootstrapPromptWarningSignaturesSeen =
    attempt.bootstrapPromptWarningSignaturesSeen ??
    (attempt.bootstrapPromptWarningSignature
      ? Array.from(
          new Set([
            ...input.bootstrapPromptWarningSignaturesSeen,
            attempt.bootstrapPromptWarningSignature,
          ]),
        )
      : input.bootstrapPromptWarningSignaturesSeen);
  const lastAssistantUsage = normalizeAssistantUsageForContext(sessionLastAssistant);
  const currentAttemptAssistantUsage = normalizeAssistantUsageForContext(currentAttemptAssistant);
  const promptCacheLastCallUsage = normalizeUsage(attempt.promptCache?.lastCallUsage as UsageLike);
  // Current-attempt evidence is newest. The session assistant is only a transcript fallback
  // and can predate a carried attempt snapshot after transcript rewrites or compaction.
  const callUsage = resolveLatestCallUsage({
    currentAttemptCandidates: [
      attempt.attemptUsage?.contextUsage ? attempt.attemptUsage : undefined,
      currentAttemptAssistantUsage,
      promptCacheLastCallUsage,
    ],
    carriedUsage: input.lastRunPromptUsage,
    transcriptFallback: lastAssistantUsage,
  });
  const attemptUsage = attempt.attemptUsage ?? callUsage.currentAttempt;
  mergeUsageIntoAccumulator(input.usageAccumulator, attemptUsage);
  mergeAttemptRunStatsIntoAccumulator(input.usageAccumulator, attempt);
  // A real mid-turn truncation rewrites the context after earlier usage observations.
  // Keep billing accumulated, but do not carry that pre-mutation context into the retry.
  const contextMutatedByMidTurnTruncation =
    preflightRecovery?.handled === true &&
    preflightRecovery.source === "mid-turn" &&
    (preflightRecovery.truncatedCount ?? 0) > 0;
  const lastRunPromptUsage = contextMutatedByMidTurnTruncation
    ? { contextUsage: { state: "unavailable" as const } }
    : callUsage.latest;
  const breakerStep = stepIdleTimeoutBreaker(input.idleTimeoutBreakerState, {
    idleTimedOut: terminalTimedOut && idleTimedOut,
    completedModelProgress: hasCompletedModelProgressForIdleBreaker(attempt),
    outputTokens: attemptUsage?.output,
  });
  if (breakerStep.tripped) {
    const message =
      `Idle-timeout cost-runaway breaker tripped: ${breakerStep.consecutive} consecutive idle timeouts ` +
      `without completed model progress (cap=${MAX_CONSECUTIVE_IDLE_TIMEOUTS_BEFORE_OUTPUT}). ` +
      "Halting further attempts to bound paid model calls. See issue #76293.";
    log.error(
      `[idle-timeout-circuit-breaker-tripped] sessionKey=${params.sessionKey ?? params.sessionId} ` +
        `provider=${provider}/${modelId} consecutive=${breakerStep.consecutive} ` +
        `cap=${MAX_CONSECUTIVE_IDLE_TIMEOUTS_BEFORE_OUTPUT}`,
    );
    const result = handleRetryLimitExhaustion({
      message,
      decision: resolveRunFailoverDecision({
        stage: "retry_limit",
        fallbackConfigured: runInput.fallbackConfigured,
        failoverReason: input.lastRetryFailoverReason,
      }),
      provider,
      model: modelId,
      profileId: runtime.lastProfileId,
      durationMs: Date.now() - runInput.startedAtMs,
      agentMeta: buildErrorAgentMeta({
        sessionId: sessionPromptState.sessionId,
        sessionFile: sessionPromptState.sessionFile,
        provider,
        model: preparedRuntime.model.id,
        credentialSource: attempt.modelAttempt?.credentialSource,
        ...runtime.outerContextTokenMeta,
        usageAccumulator: input.usageAccumulator,
        lastRunPromptUsage,
      }),
      replayInvalid: input.replayState.replayInvalid ? true : undefined,
      livenessState: "blocked",
    });
    // Escalating provider failures throw above; only returned results carry a terminal stop.
    result.meta.modelFallbackStopReason = "idle_timeout_circuit_breaker";
    return { action: "complete", result };
  }
  if (attempt.contextBudgetStatus) {
    input.contextRecoveryState.lastContextBudgetStatus = attempt.contextBudgetStatus;
  }
  const activeErrorContext = resolveActiveErrorContext({
    provider,
    model: modelId,
    assistant: attemptAssistant,
  });
  let replayState = input.replayState;
  const resolveReplayInvalidForAttempt = (incompleteTurnText?: string | null) =>
    replayState.replayInvalid || resolveReplayInvalidFlag({ attempt, incompleteTurnText });
  if (resolveReplayInvalidForAttempt(null)) {
    replayState.replayInvalid = true;
  }
  replayState = observeReplayMetadata(replayState, attempt.replayMetadata);
  const formattedAssistantErrorText = attemptAssistant
    ? formatAssistantErrorText(attemptAssistant, {
        cfg: params.config,
        sessionKey: runInput.resolvedSessionKey ?? params.sessionId,
        agentId: params.agentId,
        provider: activeErrorContext.provider,
        providerOwner: runtime.providerRuntimeHandle?.plugin,
        model: activeErrorContext.model,
        authMode: runtime.lastProfileId
          ? preparedRuntime.attemptAuthProfileStore.profiles?.[runtime.lastProfileId]?.type
          : undefined,
      })
    : undefined;
  const assistantErrorText =
    attemptAssistant?.stopReason === "error"
      ? attemptAssistant.errorMessage?.trim() || formattedAssistantErrorText
      : undefined;
  if (!signalOwnedInterruption && !preparedRuntime.nativeModelOwned && preflightRecovery?.handled) {
    const retryingFromTranscript = preflightRecovery.source === "mid-turn";
    log.info(
      `[context-overflow-precheck] early recovery route=${preflightRecovery.route} completed for ${provider}/${modelId}; ` +
        (retryingFromTranscript ? "retrying from current transcript" : "retrying prompt"),
    );
    if (retryingFromTranscript) {
      if ((preflightRecovery.truncatedCount ?? 0) > 0) {
        sessionPromptState.markOwnedTranscriptRetry();
      }
      sessionPromptState.continueFromCurrentTranscript();
    }
    const retryKind = resolveRunRetryKind({
      preflightRecovery,
      retryingFromTranscript,
      toolMetas: attempt.toolMetas,
    });
    return {
      action: "retry",
      retryKind,
      bootstrapPromptWarningSignaturesSeen,
      lastRunPromptUsage,
      replayState,
    };
  }
  return {
    action: "proceed",
    bootstrapPromptWarningSignaturesSeen,
    lastRunPromptUsage,
    replayState,
    attempt,
    sessionIdUsed,
    sessionFileUsed,
    currentAttemptAssistant,
    currentAttemptCompletedAssistant,
    attemptAssistant,
    terminalState,
    setTerminalLifecycleMeta,
    attemptCompactionCount,
    activeErrorContext,
    resolveReplayInvalidForAttempt,
    assistantErrorText,
    canRestartForLiveSwitch:
      !hasOutboundDeliveryEvidence(attempt) &&
      !attempt.didSendDeterministicApprovalPrompt &&
      !attempt.lastToolError &&
      (attempt.toolMetas?.length ?? 0) === 0 &&
      (attempt.assistantTexts?.length ?? 0) === 0,
  };
}