企业级 AI Agent“推测执行与预测性工具预取”（Speculative Tool Execution & Proactive Branch Pre-fetching）架构实战：破解长调用链串行延迟瓶颈、只读推测沙箱隔离与误预测幂等回滚

/**
 * 企业级 AI Agent 推测执行与预测性工具预取核心引擎
 * 遵循严格只读隔离、马尔可夫概率分支预测与无缝兑现/废弃协议
 */

// 1. 工具定义与只读安全契约接口
export interface AgentToolDefinition {
  name: string;
  description: string;
  isReadOnly: boolean; // 必须声明是否只读，非只读工具绝对禁止纳入推测执行
  handler: (args: Record, signal?: AbortSignal) => Promise;
}

// 2. 推测预取任务状态
export interface SpeculativeTask {
  id: string;
  toolName: string;
  argsHash: string;
  args: Record;
  startTime: number;
  abortController: AbortController;
  promise: Promise;
  status: 'pending' | 'resolved' | 'rejected' | 'discarded';
  result?: any;
  error?: any;
}

// 3. 马尔可夫转移概率预测器
export class MarkovBranchPredictor {
  // 存储状态转移概率：当前工具 -> 下一个工具及其概率 (0.0 ~ 1.0)
  private transitionMatrix: Map> = new Map();
  private readonly confidenceThreshold: number; // 触发推测执行的最低置信度阈值

  constructor(confidenceThreshold: number = 0.70) {
    this.confidenceThreshold = confidenceThreshold;
    this.initDefaultEnterpriseTransitions();
  }

  /**
   * 初始化制造业与供应链常见业务工具转移先验概率
   */
  private initDefaultEnterpriseTransitions(): void {
    this.registerTransition('get_work_order', 'query_inventory', 0.88);
    this.registerTransition('query_inventory', 'query_purchase_order', 0.76);
    this.registerTransition('query_purchase_order', 'query_vendor_delivery', 0.82);
    this.registerTransition('get_contract_info', 'audit_tax_compliance', 0.79);
  }

  public registerTransition(fromTool: string, toTool: string, probability: number): void {
    if (!this.transitionMatrix.has(fromTool)) {
      this.transitionMatrix.set(fromTool, new Map());
    }
    this.transitionMatrix.get(fromTool)!.set(toTool, probability);
  }

  /**
   * 根据当前完成的工具和上下文，预测下一个最可能的只读工具
   */
  public predictNextTool(currentTool: string): { toolName: string; confidence: number } | null {
    const nextCandidates = this.transitionMatrix.get(currentTool);
    if (!nextCandidates || nextCandidates.size === 0) return null;

    let bestTool = '';
    let highestProb = 0;

    for (const [tool, prob] of nextCandidates.entries()) {
      if (prob > highestProb) {
        highestProb = prob;
        bestTool = tool;
      }
    }

    if (highestProb >= this.confidenceThreshold) {
      return { toolName: bestTool, confidence: highestProb };
    }
    return null;
  }
}

// 4. 企业级智能体推测执行调度器
export class SpeculativeAgentEngine {
  private tools: Map = new Map();
  private predictor: MarkovBranchPredictor;
  private activeSpeculations: Map = new Map();

  // 监控遥测指标
  public telemetry = {
    totalSpeculations: 0,
    speculativeHits: 0,
    speculativeDiscards: 0,
    savedLatencyMs: 0,
  };

  constructor(tools: AgentToolDefinition[], predictor?: MarkovBranchPredictor) {
    tools.forEach((t) => this.tools.set(t.name, t));
    this.predictor = predictor || new MarkovBranchPredictor();
  }

  /**
   * 计算参数的规范化哈希键，保证属性顺序无关
   */
  private computeArgsHash(args: Record): string {
    const sortedKeys = Object.keys(args).sort();
    const normalized = sortedKeys.map((k) => `${k}=${JSON.stringify(args[k])}`).join('&');
    return normalized;
  }

  /**
   * 启动推测预取探针（在主模型推理思考的同时由事件钩子触发）
   */
  public maybeTriggerSpeculation(
    lastCompletedTool: string,
    derivedArgs: Record
  ): string | null {
    const prediction = this.predictor.predictNextTool(lastCompletedTool);
    if (!prediction) return null;

    const tool = this.tools.get(prediction.toolName);
    if (!tool) return null;

    // 核心安全红线：严禁对写工具执行推测预取
    if (!tool.isReadOnly) {
      console.warn(`[推测安全拦截] 工具 ${tool.name} 包含写副作用，严禁执行推测预取！`);
      return null;
    }

    const argsHash = this.computeArgsHash(derivedArgs);
    const speculationKey = `${tool.name}::${argsHash}`;

    // 避免重复触发同一参数的推测任务
    if (this.activeSpeculations.has(speculationKey)) {
      return speculationKey;
    }

    const abortController = new AbortController();
    const taskId = `spec_${Date.now()}_${Math.random().toString(36).substring(7)}`;

    console.log(
      `[推测探针启动] 触发只读预取: ${tool.name}, 置信度: ${(prediction.confidence * 100).toFixed(1)}%`
    );

    const task: SpeculativeTask = {
      id: taskId,
      toolName: tool.name,
      argsHash,
      args: derivedArgs,
      startTime: Date.now(),
      abortController,
      status: 'pending',
      promise: Promise.resolve(),
    };

    // 异步拉取并持有结果，注入影子追踪上下文
    task.promise = tool
      .handler(derivedArgs, abortController.signal)
      .then((res) => {
        if (task.status === 'pending') {
          task.status = 'resolved';
          task.result = res;
        }
        return res;
      })
      .catch((err) => {
        if (task.status === 'pending') {
          task.status = 'rejected';
          task.error = err;
        }
      });

    this.activeSpeculations.set(speculationKey, task);
    this.telemetry.totalSpeculations++;
    return speculationKey;
  }

  /**
   * 结果兑现或回退：大模型实际吐出工具调用时调用此方法
   */
  public async executeOrCommitTool(
    toolName: string,
    actualArgs: Record
  ): Promise<{ result: any; fromSpeculation: boolean; elapsedMs: number }> {
    const start = Date.now();
    const argsHash = this.computeArgsHash(actualArgs);
    const targetKey = `${toolName}::${argsHash}`;

    const speculativeTask = this.activeSpeculations.get(targetKey);

    if (speculativeTask) {
      // 预测命中！移除缓存并处理兑现
      this.activeSpeculations.delete(targetKey);
      this.telemetry.speculativeHits++;

      if (speculativeTask.status === 'resolved') {
        const savedMs = Math.max(0, Date.now() - speculativeTask.startTime);
        this.telemetry.savedLatencyMs += savedMs;
        console.log(`[推测命中兑现] 秒级命中缓存: ${toolName}, 节省耗时 ~${savedMs}ms`);
        return {
          result: speculativeTask.result,
          fromSpeculation: true,
          elapsedMs: Date.now() - start, // 0~1ms
        };
      } else if (speculativeTask.status === 'pending') {
        // 请求正在飞行中，主线程只需等待剩余尾部延迟
        console.log(`[推测部分命中] 请求仍在飞行中，等待剩余网络延迟...`);
        const result = await speculativeTask.promise;
        const totalElapsed = Date.now() - start;
        return {
          result,
          fromSpeculation: true,
          elapsedMs: totalElapsed,
        };
      }
    }

    // 预测失配：大模型选择了其他工具或不同入参
    this.discardAllActiveSpeculations();

    const tool = this.tools.get(toolName);
    if (!tool) {
      throw new Error(`未找到目标工具: ${toolName}`);
    }

    // 执行常规实时调用
    console.log(`[常规实时调用] 执行非推测工具: ${toolName}`);
    const result = await tool.handler(actualArgs);
    return {
      result,
      fromSpeculation: false,
      elapsedMs: Date.now() - start,
    };
  }

  /**
   * 清理并丢弃所有未命中的推测任务
   */
  public discardAllActiveSpeculations(): void {
    for (const [key, task] of this.activeSpeculations.entries()) {
      if (task.status === 'pending') {
        task.status = 'discarded';
        task.abortController.abort(); // 立即掐断正在飞行的无用网络连接
      }
      this.activeSpeculations.delete(key);
      this.telemetry.speculativeDiscards++;
    }
  }
}