metona-ai-desktop/electron/harness/hooks/pre-tool.ts

/**
 * Pre-Tool Hooks — 工具执行前钩子
 *
 * @see docs/生产级通用 AI Agent 智能体桌面应用：完整设计与构建指南.html — 第五章
 */

import type { MetonaToolCall } from '../types';
import type { PolicyEngine } from '../sandbox/permissions';

export interface HookResult {
  blocked: boolean;
  reason?: string;
  modifiedArgs?: Record<string, unknown>;
}

export interface PreToolHook {
  beforeExecute(toolCall: MetonaToolCall, sessionId: string): Promise<HookResult>;
}

/** 权限校验钩子 — 集成 PolicyEngine */
export class PermissionCheckHook implements PreToolHook {
  constructor(private policyEngine: PolicyEngine) {}

  async beforeExecute(toolCall: MetonaToolCall, sessionId: string): Promise<HookResult> {
    // v0.4.1: 透传 sessionId 使频率限制按会话隔离（多会话并发时各自独立配额）
    const result = this.policyEngine.checkAuthorization(toolCall.name, toolCall.args, sessionId);
    if (!result.authorized) {
      return { blocked: true, reason: result.reason };
    }
    // v0.3.0 修复: 授权成功后记录调用，使频率限制功能生效
    // 在授权检查通过后立即记录，即使后续工具执行失败也计入频率
    // 这样可以防止通过故意制造错误来绕过频率限制
    this.policyEngine.recordCall(toolCall.name, sessionId);
    return { blocked: false };
  }
}

/** 速率限制钩子 */
export class RateLimitHook implements PreToolHook {
  private callCounts = new Map<string, { count: number; resetTime: number }>();

  constructor(private maxCallsPerMinute: number = 20) {}

  async beforeExecute(toolCall: MetonaToolCall, sessionId: string): Promise<HookResult> {
    // 使用 sessionId:toolName 作为 key，实现会话隔离的 per-tool 速率限制
    const key = `${sessionId}:${toolCall.name}`;
    const now = Date.now();

    // v0.3.0 修复: 定期清理过期 entry，避免 Map 随会话累积无限增长
    if (this.callCounts.size > 1000) {
      for (const [k, v] of this.callCounts) {
        if (v.resetTime < now) this.callCounts.delete(k);
      }
    }

    const entry = this.callCounts.get(key);
    // L-2 修复: 使用 > 而非 >=，确保窗口到期时正确重置（边界条件）
    // 当 resetTime === now 时应视为已到期，进入 else 分支重建 entry
    if (entry && entry.resetTime > now) {
      if (entry.count >= this.maxCallsPerMinute) {
        return { blocked: true, reason: `Rate limit exceeded for tool "${toolCall.name}"` };
      }
      entry.count++;
    } else {
      this.callCounts.set(key, { count: 1, resetTime: now + 60_000 });
    }
    return { blocked: false };
  }
}