ai-terminal-assistant/packages/core/src/agent/executor.ts

import {
  generateText,
  streamText,
  stepCountIs,
  type ModelMessage,
  type Tool as AITool,
  type LanguageModel,
} from 'ai';
import type { Tool, ToolResult, AgentConfig, ContentBlock, TokenUsageInfo } from '../types/index.js';
import { buildZodSchema } from '../types/index.js';
import { ToolRegistry } from '../tools/registry.js';
import type {
  AgentInfo,
  AgentExecutionContext,
  AgentExecutionResult,
  ImageData,
} from './types.js';
import { checkBashPermission, isPathInAllowedWritePaths } from './permission-merger.js';
import { getProviderRegistry, resolveApiKey } from '../provider/index.js';
import { renderPromptTemplate, createPlanContext } from '../template/index.js';
import { agentEventEmitter } from './events.js';

/**
 * Agent 执行器
 * 根据 Agent 配置执行任务，支持工具过滤和权限控制
 */
export class AgentExecutor {
  private agentInfo: AgentInfo;
  private baseConfig: AgentConfig;
  private toolRegistry: ToolRegistry;
  private getModel: (model: string) => LanguageModel;

  constructor(
    agentInfo: AgentInfo,
    baseConfig: AgentConfig,
    toolRegistry: ToolRegistry
  ) {
    this.agentInfo = agentInfo;
    this.baseConfig = baseConfig;
    this.toolRegistry = toolRegistry;

    // 使用 ProviderRegistry 获取模型工厂
    const provider = agentInfo.model?.provider ?? baseConfig.provider;
    const registry = getProviderRegistry();

    // 当 Agent 指定了不同的 provider 时，需要从 ProviderRegistry 获取对应的配置
    // 而不是使用 baseConfig（主 Agent 配置）的 apiKey 和 baseUrl
    let apiKey = baseConfig.apiKey;
    let baseUrl = baseConfig.baseUrl;

    if (agentInfo.model?.provider && agentInfo.model.provider !== baseConfig.provider) {
      // Agent 使用了不同的 provider，获取对应 provider 的配置
      const providerConfig = registry.getConfig(provider);
      apiKey = resolveApiKey(providerConfig) || baseConfig.apiKey;
      baseUrl = providerConfig?.baseUrl;
    }

    this.getModel = registry.getModelFactory(provider, { apiKey, baseUrl });
  }

  /**
   * 执行任务
   */
  async execute(
    prompt: string,
    context: AgentExecutionContext
  ): Promise<AgentExecutionResult> {
    const { onStream, onToolCall, onToolResult, images, sessionId, agentId, emitEvents } = context;

    // 是否发射子 Agent 事件
    const shouldEmitEvents = emitEvents && sessionId && agentId;

    // 获取过滤后的工具
    const tools = this.getFilteredTools();
    const vercelTools = this.buildVercelTools(tools);

    // 构建系统提示词
    const systemPrompt = this.buildSystemPrompt();

    // 获取模型配置
    const modelName = this.agentInfo.model?.model ?? this.baseConfig.model;
    const maxSteps = this.agentInfo.maxSteps ?? 10;
    const maxTokens = this.agentInfo.model?.maxTokens ?? this.baseConfig.maxTokens;

    // 构建消息内容（支持图片）
    const messageContent = this.buildMessageContent(prompt, images);

    // 构建初始消息
    const messages: ModelMessage[] = [
      {
        role: 'user',
        content: messageContent,
      },
    ];

    let fullResponse = '';
    let steps = 0;
    let usage: TokenUsageInfo | undefined;

    // 工具调用时间追踪（用于计算持续时间）
    const toolStartTimes = new Map<string, number>();

    try {
      if (onStream || shouldEmitEvents) {
        // 流式模式（或需要发射事件时使用流式模式）
        const result = streamText({
          model: this.getModel(modelName),
          system: systemPrompt,
          messages,
          tools: vercelTools,
          maxOutputTokens: maxTokens,
          stopWhen: stepCountIs(maxSteps),
          onChunk: ({ chunk }) => {
            if (chunk.type === 'tool-call') {
              steps++;
              const toolArgs = 'input' in chunk ? chunk.input : {};
              const toolCallId = `tool-${Date.now()}-${Math.random().toString(36).slice(2, 7)}`;

              // 记录工具开始时间
              toolStartTimes.set(chunk.toolName, Date.now());

              onToolCall?.(chunk.toolName, toolArgs as Record<string, unknown>);
              onStream?.(`\n[调用工具: ${chunk.toolName}]\n`);

              // 发射子 Agent 工具开始事件
              if (shouldEmitEvents) {
                agentEventEmitter.emit({
                  type: 'subagent:tool_start',
                  sessionId: sessionId!,
                  agentId: agentId!,
                  toolCallId,
                  toolName: chunk.toolName,
                  args: toolArgs as Record<string, unknown>,
                });
              }
            } else if (chunk.type === 'tool-result') {
              const output = (chunk as { output?: ToolResult }).output;
              const toolName = (chunk as { toolName?: string }).toolName ?? 'unknown';
              const toolCallId = `tool-${Date.now()}-${Math.random().toString(36).slice(2, 7)}`;

              // 计算工具执行时间
              const startTime = toolStartTimes.get(toolName);
              const duration = startTime ? Date.now() - startTime : undefined;
              toolStartTimes.delete(toolName);

              onToolResult?.(toolName, output);

              if (output && typeof output === 'object') {
                if (output.success) {
                  const displayOutput =
                    output.output.length > 500
                      ? output.output.substring(0, 500) + '...(截断)'
                      : output.output;
                  onStream?.(`[结果: ${displayOutput}]\n`);

                  // 发射子 Agent 工具完成事件
                  if (shouldEmitEvents) {
                    agentEventEmitter.emit({
                      type: 'subagent:tool_end',
                      sessionId: sessionId!,
                      agentId: agentId!,
                      toolCallId,
                      status: 'completed',
                      result: output.output,
                      duration,
                    });
                  }
                } else {
                  onStream?.(`[错误: ${output.error}]\n`);

                  // 发射子 Agent 工具错误事件
                  if (shouldEmitEvents) {
                    agentEventEmitter.emit({
                      type: 'subagent:tool_end',
                      sessionId: sessionId!,
                      agentId: agentId!,
                      toolCallId,
                      status: 'error',
                      error: output.error,
                      duration,
                    });
                  }
                }
              }
            }
          },
        });

        for await (const chunk of result.textStream) {
          fullResponse += chunk;
          onStream?.(chunk);

          // 发射子 Agent 流式输出事件
          if (shouldEmitEvents && chunk) {
            agentEventEmitter.emit({
              type: 'subagent:stream',
              sessionId: sessionId!,
              agentId: agentId!,
              content: chunk,
            });
          }
        }

        const response = await result.response;
        // 提取 usage 信息
        usage = this.extractUsage(response);
      } else {
        // 非流式模式
        const result = await generateText({
          model: this.getModel(modelName),
          system: systemPrompt,
          messages,
          tools: vercelTools,
          maxOutputTokens: maxTokens,
          stopWhen: stepCountIs(maxSteps),
        });

        fullResponse = result.text;
        steps = result.steps.length;
        // 提取 usage 信息
        usage = this.extractUsage(result.response);
      }

      return {
        success: true,
        text: fullResponse,
        steps,
        sessionId: context.parentSessionId ?? 'standalone',
        usage,
      };
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : String(error);
      return {
        success: false,
        text: '',
        steps,
        sessionId: context.parentSessionId ?? 'standalone',
        error: errorMessage,
      };
    }
  }

  /**
   * 获取过滤后的工具列表
   */
  private getFilteredTools(): Tool[] {
    const allTools = this.toolRegistry.getAllTools();
    const toolConfig = this.agentInfo.tools;

    // 如果没有工具配置，返回所有工具
    if (!toolConfig) {
      return allTools;
    }

    let filteredTools = allTools;

    // 如果指定了 enabled，只保留这些工具
    if (toolConfig.enabled && toolConfig.enabled.length > 0) {
      const enabledSet = new Set(toolConfig.enabled);
      filteredTools = filteredTools.filter((t) => enabledSet.has(t.name));
    }

    // 移除 disabled 的工具
    if (toolConfig.disabled && toolConfig.disabled.length > 0) {
      const disabledSet = new Set(toolConfig.disabled);
      filteredTools = filteredTools.filter((t) => !disabledSet.has(t.name));
    }

    // 如果禁止嵌套 Task，移除 task 工具
    if (toolConfig.noTask) {
      filteredTools = filteredTools.filter((t) => t.name !== 'task');
    }

    return filteredTools;
  }

  /**
   * 构建 Vercel AI SDK 工具格式
   */
  private buildVercelTools(tools: Tool[]): Record<string, AITool> {
    const vercelTools: Record<string, AITool> = {};

    for (const tool of tools) {
      const schema = buildZodSchema(tool.parameters);

      vercelTools[tool.name] = {
        description: tool.description,
        inputSchema: schema,
        execute: async (params) => {
          // 权限检查
          const permissionResult = await this.checkToolPermission(
            tool.name,
            params as Record<string, unknown>
          );
          if (!permissionResult.allowed) {
            return {
              success: false,
              output: '',
              error: `权限拒绝: ${permissionResult.reason}`,
            };
          }

          return tool.execute(params as Record<string, unknown>);
        },
      } as AITool;
    }

    return vercelTools;
  }

  /**
   * 检查工具调用权限
   */
  private async checkToolPermission(
    toolName: string,
    params: Record<string, unknown>
  ): Promise<{ allowed: boolean; reason?: string }> {
    const permission = this.agentInfo.permission;
    if (!permission) {
      return { allowed: true };
    }

    // Bash 权限检查
    if (toolName === 'bash' && permission.bash) {
      const command = params.command as string;
      if (!command) {
        return { allowed: true };
      }

      const action = checkBashPermission(command, permission.bash);
      if (action === 'deny') {
        return { allowed: false, reason: `命令被禁止: ${command}` };
      }
      // ask 在这里视为允许（实际的 ask 逻辑在权限管理器中处理）
    }

    // 文件写入权限检查
    if (['write_file', 'edit_file'].includes(toolName)) {
      const filePermission = permission.file;
      if (filePermission) {
        const operation = toolName === 'write_file' ? 'write' : 'edit';
        const action = filePermission[operation];
        if (action === 'deny') {
          return { allowed: false, reason: `${operation} 操作被禁止` };
        }

        // 检查 allowedWritePaths 限制（仅对 write 操作）
        if (operation === 'write' && filePermission.allowedWritePaths) {
          const filePath = params.path as string;
          if (filePath && !isPathInAllowedWritePaths(filePath, filePermission.allowedWritePaths)) {
            return {
              allowed: false,
              reason: `写入路径不在允许列表中: ${filePath}。只能写入: ${filePermission.allowedWritePaths.join(', ')}`,
            };
          }
        }
      }
    }

    // Git 写操作权限检查
    const gitWriteTools = ['git_add', 'git_commit', 'git_push', 'git_checkout', 'git_stash'];
    if (gitWriteTools.includes(toolName) && permission.git?.write === 'deny') {
      return { allowed: false, reason: 'Git 写操作被禁止' };
    }

    return { allowed: true };
  }

  /**
   * 构建系统提示词
   * 如果 Agent 启用了 promptTemplate，则动态渲染模板变量
   */
  private buildSystemPrompt(): string {
    // 如果 Agent 有自定义 prompt
    if (this.agentInfo.prompt) {
      // 如果启用了模板渲染，动态解析变量
      if (this.agentInfo.promptTemplate) {
        const context = createPlanContext({
          workdir: process.cwd(),
          isSubagent: this.agentInfo.mode === 'subagent',
        });
        return renderPromptTemplate(this.agentInfo.prompt, context);
      }
      return this.agentInfo.prompt;
    }

    // 否则使用基础配置的 systemPrompt
    return this.baseConfig.systemPrompt;
  }

  /**
   * 构建消息内容（支持图片）
   */
  private buildMessageContent(
    prompt: string,
    images?: ImageData[]
  ): string | ContentBlock[] {
    // 如果没有图片，直接返回文本
    if (!images || images.length === 0) {
      return prompt;
    }

    // 构建多模态内容
    const blocks: ContentBlock[] = [];

    // 先添加图片
    for (const img of images) {
      blocks.push({
        type: 'image',
        image: img.data,
        mimeType: img.mimeType,
      });
    }

    // 再添加文本
    if (prompt) {
      blocks.push({
        type: 'text',
        text: prompt,
      });
    }

    return blocks;
  }

  /**
   * 从 AI SDK 响应中提取 usage 信息
   */
  private extractUsage(response: unknown): TokenUsageInfo | undefined {
    // AI SDK 的 response 对象包含 usage 字段
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    const resp = response as any;
    const usage = resp?.usage;

    if (!usage) {
      return undefined;
    }

    return {
      promptTokens: usage.promptTokens ?? 0,
      completionTokens: usage.completionTokens ?? 0,
      totalTokens: usage.totalTokens ?? (usage.promptTokens ?? 0) + (usage.completionTokens ?? 0),
      // Anthropic API 特有的缓存字段
      cacheReadInputTokens: usage.cacheReadInputTokens,
      cacheCreationInputTokens: usage.cacheCreationInputTokens,
    };
  }
}