qwen-code/packages/cli/src/acp-integration/session/HistoryReplayer.ts

/**
 * @license
 * Copyright 2025 Qwen
 * SPDX-License-Identifier: Apache-2.0
 */

import type { ChatRecord, TaskResultDisplay } from '@qwen-code/qwen-code-core';
import type {
  Content,
  GenerateContentResponseUsageMetadata,
} from '@google/genai';
import type { SessionContext } from './types.js';
import { MessageEmitter } from './emitters/MessageEmitter.js';
import { ToolCallEmitter } from './emitters/ToolCallEmitter.js';

/**
 * Handles replaying session history on session load.
 *
 * Uses the unified emitters to ensure consistency with normal flow.
 * This ensures that replayed history looks identical to how it would
 * have appeared during the original session.
 */
export class HistoryReplayer {
  private readonly messageEmitter: MessageEmitter;
  private readonly toolCallEmitter: ToolCallEmitter;

  constructor(ctx: SessionContext) {
    this.messageEmitter = new MessageEmitter(ctx);
    this.toolCallEmitter = new ToolCallEmitter(ctx);
  }

  /**
   * Replays all chat records from a loaded session.
   *
   * @param records - Array of chat records to replay
   */
  async replay(records: ChatRecord[]): Promise<void> {
    for (const record of records) {
      await this.replayRecord(record);
    }
  }

  /**
   * Replays a single chat record.
   */
  private async replayRecord(record: ChatRecord): Promise<void> {
    switch (record.type) {
      case 'user':
        if (record.message) {
          await this.replayContent(record.message, 'user');
        }
        break;

      case 'assistant':
        if (record.message) {
          await this.replayContent(record.message, 'assistant');
        }
        if (record.usageMetadata) {
          await this.replayUsageMetadata(record.usageMetadata);
        }
        break;

      case 'tool_result':
        await this.replayToolResult(record);
        break;

      default:
        // Skip system records (compression, telemetry, slash commands)
        break;
    }
  }

  /**
   * Replays content from a message (user or assistant).
   * Handles text parts, thought parts, and function calls.
   */
  private async replayContent(
    content: Content,
    role: 'user' | 'assistant',
  ): Promise<void> {
    for (const part of content.parts ?? []) {
      // Text content
      if ('text' in part && part.text) {
        const isThought = (part as { thought?: boolean }).thought ?? false;
        await this.messageEmitter.emitMessage(part.text, role, isThought);
      }

      // Function call (tool start)
      if ('functionCall' in part && part.functionCall) {
        const functionName = part.functionCall.name ?? '';
        const callId = part.functionCall.id ?? `${functionName}-${Date.now()}`;

        await this.toolCallEmitter.emitStart({
          toolName: functionName,
          callId,
          args: part.functionCall.args as Record<string, unknown>,
          status: 'in_progress',
        });
      }
    }
  }

  /**
   * Replays usage metadata.
   * @param usageMetadata - The usage metadata to replay
   */
  private async replayUsageMetadata(
    usageMetadata: GenerateContentResponseUsageMetadata,
  ): Promise<void> {
    await this.messageEmitter.emitUsageMetadata(usageMetadata);
  }

  /**
   * Replays a tool result record.
   */
  private async replayToolResult(record: ChatRecord): Promise<void> {
    // message is required - skip if not present
    if (!record.message?.parts) {
      return;
    }

    const result = record.toolCallResult;
    const callId = result?.callId ?? record.uuid;

    // Extract tool name from the function response in message if available
    const toolName = this.extractToolNameFromRecord(record);

    await this.toolCallEmitter.emitResult({
      toolName,
      callId,
      success: !result?.error,
      message: record.message.parts,
      resultDisplay: result?.resultDisplay,
      // For TodoWriteTool fallback, try to extract args from the record
      // Note: args aren't stored in tool_result records by default
      args: undefined,
    });

    // Special handling: Task tool execution summary contains token usage
    const { resultDisplay } = result ?? {};
    if (
      !!resultDisplay &&
      typeof resultDisplay === 'object' &&
      'type' in resultDisplay &&
      (resultDisplay as { type?: unknown }).type === 'task_execution'
    ) {
      await this.emitTaskUsageFromResultDisplay(
        resultDisplay as TaskResultDisplay,
      );
    }
  }

  /**
   * Emits token usage from a TaskResultDisplay execution summary, if present.
   */
  private async emitTaskUsageFromResultDisplay(
    resultDisplay: TaskResultDisplay,
  ): Promise<void> {
    const summary = resultDisplay.executionSummary;
    if (!summary) {
      return;
    }

    const usageMetadata: GenerateContentResponseUsageMetadata = {};

    if (Number.isFinite(summary.inputTokens)) {
      usageMetadata.promptTokenCount = summary.inputTokens;
    }
    if (Number.isFinite(summary.outputTokens)) {
      usageMetadata.candidatesTokenCount = summary.outputTokens;
    }
    if (Number.isFinite(summary.thoughtTokens)) {
      usageMetadata.thoughtsTokenCount = summary.thoughtTokens;
    }
    if (Number.isFinite(summary.cachedTokens)) {
      usageMetadata.cachedContentTokenCount = summary.cachedTokens;
    }
    if (Number.isFinite(summary.totalTokens)) {
      usageMetadata.totalTokenCount = summary.totalTokens;
    }

    // Only emit if we captured at least one token metric
    if (Object.keys(usageMetadata).length > 0) {
      await this.messageEmitter.emitUsageMetadata(usageMetadata);
    }
  }

  /**
   * Extracts tool name from a chat record's function response.
   */
  private extractToolNameFromRecord(record: ChatRecord): string {
    // Try to get from functionResponse in message
    if (record.message?.parts) {
      for (const part of record.message.parts) {
        if ('functionResponse' in part && part.functionResponse?.name) {
          return part.functionResponse.name;
        }
      }
    }
    return '';
  }
}