kimi-code/packages/agent-core/src/loop/turn-step.ts

/**
 * Executes one provider step.
 *
 * A step owns the provider call, atomic transcript envelope, streaming callback
 * wiring, tool-call lifecycle, and post-step hooks. Provider usage is recorded
 * immediately after `llm.chat` returns so a later abort during tool execution
 * does not lose model usage that was already spent.
 */

import { randomUUID } from 'node:crypto';

import {
  APIRequestTooLargeError,
  isRecoverableRequestStructureError,
  type TokenUsage,
} from '@moonshot-ai/kosong';
import type { Logger } from '#/logging/types';

import type { LoopEventDispatcher } from './events';
import { errorMessage } from './errors';
import type { LLM, LLMChatParams, LLMChatResponse } from './llm';
import { chatWithRetry } from './retry';
import { runToolCallBatch, type ToolCallStepContext } from './tool-call';
import type {
  ExecutableTool,
  LoopHooks,
  LoopMessageBuilder,
  LoopStepStopReason,
  RecordStepUsageResult,
} from './types';

type ChatStreamingCallbacks = Pick<
  LLMChatParams,
  'onTextDelta' | 'onThinkDelta' | 'onToolCallDelta' | 'onTextPart' | 'onThinkPart'
>;

export interface ExecuteLoopStepDeps {
  readonly turnId: string;
  readonly signal: AbortSignal;
  readonly buildMessages: LoopMessageBuilder;
  readonly buildMessagesStrict?: LoopMessageBuilder | undefined;
  /** See RunTurnInput.buildMessagesMediaDegraded. */
  readonly buildMessagesMediaDegraded?: LoopMessageBuilder | undefined;
  readonly dispatchEvent: LoopEventDispatcher;
  readonly llm: LLM;
  readonly tools?: readonly ExecutableTool[] | undefined;
  /**
   * Per-step tool table builder; wins over the static `tools` snapshot.
   * Evaluated after `beforeStep`, next to `buildMessages`, so the executable
   * table and the request messages reflect the same state — `beforeStep` can
   * run compaction, which discards loaded dynamic tool schemas.
   */
  readonly buildTools?: (() => readonly ExecutableTool[]) | undefined;
  /** See RunTurnInput.describeMissingTool. */
  readonly describeMissingTool?: ((name: string) => string | undefined) | undefined;
  readonly hooks?: LoopHooks | undefined;
  readonly log?: Logger | undefined;
  readonly currentStep: number;
  readonly maxRetryAttempts?: number;
  readonly recordUsage: (usage: TokenUsage) => RecordStepUsageResult | void | Promise<RecordStepUsageResult | void>;
}

export async function executeLoopStep(deps: ExecuteLoopStepDeps): Promise<{
  readonly usage: TokenUsage;
  readonly stopReason: LoopStepStopReason;
  /**
   * True when this step only succeeded after resending with the
   * media-degraded projection. The turn loop uses it to keep later steps on
   * that projection — re-sending the full-media history would pay a fresh
   * rejection on every step of the turn.
   */
  readonly mediaDegradedResendUsed?: boolean;
}> {
  const {
    turnId,
    signal,
    buildMessages,
    buildMessagesStrict,
    buildMessagesMediaDegraded,
    dispatchEvent,
    llm,
    tools,
    buildTools,
    describeMissingTool,
    hooks,
    log,
    currentStep,
    maxRetryAttempts,
    recordUsage,
  } = deps;

  if (hooks?.beforeStep !== undefined) {
    const beforeStep = await hooks.beforeStep({
      turnId,
      stepNumber: currentStep,
      signal,
      llm,
    });
    if (beforeStep?.block === true) {
      throw new Error(beforeStep.reason ?? `Step ${String(currentStep)} was blocked`);
    }
  }

  signal.throwIfAborted();

  // Resolve the tool table AFTER beforeStep so it reflects the same state as
  // the messages built below (beforeStep can run compaction, which discards
  // loaded dynamic tool schemas from the context and the ledger — a table
  // captured earlier would still dispatch a tool the model no longer has).
  const stepTools = buildTools !== undefined ? buildTools() : tools;
  const messages = await buildMessages();
  signal.throwIfAborted();

  const stepUuid = randomUUID();

  const step: ToolCallStepContext = {
    tools: stepTools,
    describeMissingTool,
    hooks,
    log,
    dispatchEvent,
    llm,
    signal,
    turnId,
    currentStep,
    stepUuid,
  };

  await dispatchEvent({
    type: 'step.begin',
    uuid: stepUuid,
    turnId,
    step: currentStep,
  });

  const chatParams: LLMChatParams = {
    messages,
    tools: stepTools ?? [],
    signal,
    ...createChatStreamingCallbacks({
      dispatchEvent,
      turnId,
      currentStep,
      stepUuid,
    }),
  };
  const retryInput = {
    llm,
    dispatchEvent,
    turnId,
    currentStep,
    stepUuid,
    maxAttempts: maxRetryAttempts,
    log,
  } as const;
  let response: LLMChatResponse;
  let mediaDegradedResendUsed = false;
  try {
    response = await chatWithRetry({ ...retryInput, params: chatParams });
  } catch (error) {
    if (buildMessagesMediaDegraded !== undefined && error instanceof APIRequestTooLargeError) {
      // The provider rejected the request BODY as too large (HTTP 413) —
      // accumulated base64 media, not tokens, so compaction's token-driven
      // recovery never fires (media is estimated at a small flat cost). The
      // same media is re-sent on every request, so without intervention the
      // session stays stuck. Resend ONCE with the media-degraded projection
      // (old media replaced by text markers, the most recent kept); a
      // rejection of that rebuild propagates unchanged.
      signal.throwIfAborted();
      log?.warn('provider rejected request as too large; resending with degraded media', {
        turnStep: `${turnId}.${String(currentStep)}`,
        model: llm.modelName,
      });
      const degradedMessages = await buildMessagesMediaDegraded();
      signal.throwIfAborted();
      try {
        response = await chatWithRetry({
          ...retryInput,
          params: {
            ...chatParams,
            messages: degradedMessages,
            requestLogFields: { projection: 'media-degraded' },
          },
        });
      } catch (degradedError) {
        log?.error('media-degraded resend still rejected by provider', {
          turnStep: `${turnId}.${String(currentStep)}`,
          model: llm.modelName,
          originalError: errorMessage(error),
          degradedError: errorMessage(degradedError),
        });
        throw degradedError;
      }
      mediaDegradedResendUsed = true;
      log?.info('recovered after media-degraded resend', {
        turnStep: `${turnId}.${String(currentStep)}`,
      });
    } else if (buildMessagesStrict !== undefined && isRecoverableRequestStructureError(error)) {
      // A structural request rejection (tool_use/tool_result pairing, empty or
      // whitespace-only text, non-user first message, non-alternating roles) means
      // the projected history is not wire-compliant for a strict provider — and
      // since the same history is re-sent every turn, the session would stay stuck
      // on this error forever. Resend ONCE with a strict, guaranteed-compliant
      // rebuild (every open call closed, stray results dropped, leading non-user
      // trimmed, consecutive assistants merged) as a last resort. Any other error,
      // or a host that supplied no strict builder, propagates unchanged.
      signal.throwIfAborted();
      log?.warn('provider rejected request structure; resending with strict projection', {
        turnStep: `${turnId}.${String(currentStep)}`,
        model: llm.modelName,
      });
      const strictMessages = await buildMessagesStrict();
      signal.throwIfAborted();
      try {
        response = await chatWithRetry({
          ...retryInput,
          params: {
            ...chatParams,
            messages: strictMessages,
            requestLogFields: { projection: 'strict' },
          },
        });
      } catch (strictError) {
        // The strictly-sanitized rebuild was still rejected — our wire-compliance
        // repair did not cover this case. Surface it loudly: the session is stuck
        // and this is the signal we need to diagnose the gap.
        log?.error('strict resend still rejected by provider; request remains wire-invalid', {
          turnStep: `${turnId}.${String(currentStep)}`,
          model: llm.modelName,
          originalError: errorMessage(error),
          strictError: errorMessage(strictError),
        });
        throw strictError;
      }
      log?.info('recovered after strict resend', {
        turnStep: `${turnId}.${String(currentStep)}`,
      });
    } else {
      throw error;
    }
  }
  const usage = response.usage;
  const usageResult = await recordUsage(usage);
  const stopTurnAfterUsage = usageResult?.stopTurn === true;
  const stopReason = deriveStepStopReason(response);

  // Execute tools only when the normalized response shape represents a tool
  // step. Provider terminal diagnostics such as filtering or truncation must
  // not trigger side-effecting tool execution even if a malformed response also
  // contains tool calls.
  let effectiveStopReason: LoopStepStopReason =
    stopTurnAfterUsage && stopReason === 'tool_use' ? 'end_turn' : stopReason;
  if (effectiveStopReason === 'tool_use') {
    const toolBatch = await runToolCallBatch(step, response);
    if (toolBatch.stopTurn) effectiveStopReason = 'end_turn';
  }

  // When a tool batch runs, it drains paired `tool.result` events even when
  // cancellation is requested. Check the signal here before sealing the step.
  signal.throwIfAborted();

  await dispatchEvent({
    type: 'step.end',
    uuid: stepUuid,
    turnId,
    step: currentStep,
    usage,
    finishReason: effectiveStopReason,
    llmFirstTokenLatencyMs: response.streamTiming?.firstTokenLatencyMs,
    llmStreamDurationMs: response.streamTiming?.streamDurationMs,
    llmRequestBuildMs: response.streamTiming?.requestBuildMs,
    llmServerFirstTokenMs: response.streamTiming?.serverFirstTokenMs,
    llmServerDecodeMs: response.streamTiming?.serverDecodeMs,
    llmClientConsumeMs: response.streamTiming?.clientConsumeMs,
    messageId: response.messageId,
    ...stepEndProviderDiagnostics(response, effectiveStopReason),
  });

  logStepTiming(log, turnId, currentStep, response);

  let stopTurnAfterStep = stopTurnAfterUsage;
  if (hooks?.afterStep !== undefined) {
    try {
      const afterStep = await hooks.afterStep({
        turnId,
        stepNumber: currentStep,
        usage,
        stopReason: effectiveStopReason,
        signal,
        llm,
      });
      stopTurnAfterStep = stopTurnAfterStep || afterStep?.stopTurn === true;
    } catch {
      // The step is already sealed; observer hooks cannot change the result.
    }
  }

  return {
    usage,
    stopReason:
      stopTurnAfterStep && effectiveStopReason === 'tool_use' ? 'end_turn' : effectiveStopReason,
    mediaDegradedResendUsed,
  };
}

/**
 * Emit a per-step completion log with the LLM response timing. TTFT is split
 * into the client-side request-build portion and the network + API-server
 * portion, and the decode window is split into server (awaiting parts) vs.
 * client (processing parts) time, so slow turns can be attributed without
 * parsing the wire log.
 */
function logStepTiming(
  log: Logger | undefined,
  turnId: string,
  step: number,
  response: LLMChatResponse,
): void {
  if (log === undefined) return;
  const timing = response.streamTiming;
  if (timing === undefined) return;
  log.info('llm response', {
    turnStep: `${turnId}/${String(step)}`,
    ttftMs: timing.firstTokenLatencyMs,
    ...(timing.requestBuildMs !== undefined ? { requestBuildMs: timing.requestBuildMs } : {}),
    ...(timing.serverFirstTokenMs !== undefined
      ? { serverFirstTokenMs: timing.serverFirstTokenMs }
      : {}),
    streamDurationMs: timing.streamDurationMs,
    ...(timing.serverDecodeMs !== undefined ? { serverDecodeMs: timing.serverDecodeMs } : {}),
    ...(timing.clientConsumeMs !== undefined ? { clientConsumeMs: timing.clientConsumeMs } : {}),
    outputTokens: response.usage.output,
  });
}

function deriveStepStopReason(response: LLMChatResponse): LoopStepStopReason {
  switch (response.providerFinishReason) {
    case 'truncated':
      return 'max_tokens';
    case 'filtered':
      return 'filtered';
    case 'paused':
      return 'paused';
    case 'other':
      return 'unknown';
    case 'completed':
    case undefined:
      return response.toolCalls.length > 0 ? 'tool_use' : 'end_turn';
    case 'tool_calls':
      return response.toolCalls.length > 0 ? 'tool_use' : 'unknown';
    default: {
      const _exhaustive: never = response.providerFinishReason;
      return _exhaustive;
    }
  }
}

function stepEndProviderDiagnostics(
  response: LLMChatResponse,
  stopReason: LoopStepStopReason,
): Pick<LLMChatResponse, 'providerFinishReason' | 'rawFinishReason'> {
  const providerFinishReason = response.providerFinishReason;
  if (
    (providerFinishReason === 'completed' && stopReason === 'end_turn') ||
    (providerFinishReason === 'tool_calls' && stopReason === 'tool_use')
  ) {
    return {};
  }

  return {
    ...(providerFinishReason !== undefined ? { providerFinishReason } : {}),
    ...(response.rawFinishReason !== undefined
      ? { rawFinishReason: response.rawFinishReason }
      : {}),
  };
}

function createChatStreamingCallbacks(deps: {
  readonly dispatchEvent: LoopEventDispatcher;
  readonly turnId: string;
  readonly currentStep: number;
  readonly stepUuid: string;
}): ChatStreamingCallbacks {
  const { dispatchEvent, turnId, currentStep, stepUuid } = deps;

  return {
    onTextDelta: (delta) => {
      dispatchEvent({ type: 'text.delta', delta });
    },
    onThinkDelta: (delta) => {
      dispatchEvent({ type: 'thinking.delta', delta });
    },
    onToolCallDelta: (delta) => {
      dispatchEvent({
        type: 'tool.call.delta',
        toolCallId: delta.toolCallId,
        name: delta.name,
        argumentsPart: delta.argumentsPart,
      });
    },
    onTextPart: async (part) => {
      await dispatchEvent({
        type: 'content.part',
        uuid: randomUUID(),
        turnId,
        step: currentStep,
        stepUuid,
        part,
      });
    },
    onThinkPart: async (part) => {
      await dispatchEvent({
        type: 'content.part',
        uuid: randomUUID(),
        turnId,
        step: currentStep,
        stepUuid,
        part,
      });
    },
  };
}