ADK.js SecurityPlugin实战：保护AI Agent免受恶意攻击

【免费下载链接】adk-js An open-source, code-first Typescript toolkit for building, evaluating, and deploying sophisticated AI agents with flexibility and control. 项目地址: https://gitcode.com/GitHub_Trending/ad/adk-js

ADK.js是一个开源的、代码优先的TypeScript工具包，用于构建、评估和部署复杂的AI代理，兼具灵活性和控制力。在AI代理的开发过程中，安全性是至关重要的一环，而SecurityPlugin作为ADK.js的核心安全组件，能够有效保护AI Agent免受各种恶意攻击。

一、SecurityPlugin核心功能解析

SecurityPlugin通过策略引擎对工具调用进行安全检查，主要实现了以下核心功能：

1.1 策略检查机制

SecurityPlugin定义了三种策略检查结果：

• ALLOW：工具调用被允许执行
• CONFIRM：工具调用需要外部确认
• DENY：工具调用被拒绝

这些策略结果通过PolicyOutcome枚举定义，位于core/src/plugins/security_plugin.ts文件中。

1.2 策略引擎接口

SecurityPlugin提供了BasePolicyEngine接口，允许开发者实现自定义的安全策略。默认实现InMemoryPolicyEngine采用宽松策略，允许所有工具调用，适合开发环境使用。

export class InMemoryPolicyEngine implements BasePolicyEngine {
  async evaluate(context: ToolCallPolicyContext): Promise<PolicyCheckResult> {
    // 默认宽松实现，允许所有工具调用
    return Promise.resolve({
      outcome: PolicyOutcome.ALLOW,
      reason: 'For prototyping purpose, all tool calls are allowed.',
    });
  }
}

二、SecurityPlugin工作流程

SecurityPlugin通过以下流程保护AI Agent：

2.1 工具调用前检查

在工具调用前，SecurityPlugin会执行beforeToolCallback方法，对工具调用进行安全检查。检查逻辑如下：

1. 检查是否已进行过安全检查
2. 若未检查，则调用策略引擎评估工具调用
3. 根据评估结果执行相应操作（允许、拒绝或请求确认）

2.2 状态管理

SecurityPlugin使用toolContext.state存储工具调用的安全检查状态，确保每个工具调用只被检查一次。状态管理通过getToolCallCheckState和setToolCallCheckState方法实现。

三、SecurityPlugin实战应用

3.1 基本使用方法

要使用SecurityPlugin，只需在创建Agent时将其添加到插件列表中：

import { SecurityPlugin } from './core/src/plugins/security_plugin';

const agent = new LLMAgent({
  plugins: [new SecurityPlugin()],
  // 其他配置...
});

3.2 自定义策略引擎

为了增强安全性，我们可以实现自定义策略引擎。例如，创建一个限制特定工具调用的策略引擎：

class RestrictivePolicyEngine implements BasePolicyEngine {
  async evaluate(context: ToolCallPolicyContext): Promise<PolicyCheckResult> {
    // 禁止危险工具调用
    if (context.tool.name === 'dangerous-tool') {
      return {
        outcome: PolicyOutcome.DENY,
        reason: 'Dangerous tool is not allowed'
      };
    }
    
    // 敏感操作需要确认
    if (context.tool.name === 'sensitive-operation') {
      return {
        outcome: PolicyOutcome.CONFIRM,
        reason: 'Sensitive operation requires confirmation'
      };
    }
    
    return { outcome: PolicyOutcome.ALLOW };
  }
}

// 使用自定义策略引擎
const securityPlugin = new SecurityPlugin({
  policyEngine: new RestrictivePolicyEngine()
});

3.3 处理需要确认的工具调用

当策略引擎返回CONFIRM结果时，SecurityPlugin会通过toolContext.requestConfirmation方法请求用户确认。我们可以通过以下方式处理确认请求：

agent.on('confirmationRequested', async (confirmation) => {
  // 显示确认提示给用户
  const userConfirmed = await showConfirmationDialog(confirmation.hint);
  
  // 发送确认结果
  agent.confirmToolCall(confirmation.functionCallId, userConfirmed);
});

四、SecurityPlugin源码解析

4.1 核心实现文件

SecurityPlugin的核心实现位于core/src/plugins/security_plugin.ts，主要包含：

• SecurityPlugin类：实现插件接口，处理工具调用前的安全检查
• PolicyOutcome枚举：定义策略检查结果
• BasePolicyEngine接口：策略引擎接口
• InMemoryPolicyEngine类：默认策略引擎实现

4.2 关键方法解析

checkToolCallPolicy方法是SecurityPlugin的核心，它根据策略引擎的评估结果决定如何处理工具调用：

private async checkToolCallPolicy({ tool, toolArgs, toolContext }: {
  tool: BaseTool,
  toolArgs: {[key: string]: any},
  toolContext: ToolContext
}): Promise<{[key: string]: unknown}|undefined> {
  const policyCheckResult = await this.policyEngine.evaluate({ tool, toolArgs });
  this.setToolCallCheckState(toolContext, policyCheckResult.outcome);

  switch (policyCheckResult.outcome) {
    case PolicyOutcome.DENY:
      return { error: `This tool call is rejected by policy engine. Reason: ${policyCheckResult.reason}` };
    case PolicyOutcome.CONFIRM:
      toolContext.requestConfirmation({
        hint: `Policy engine requires confirmation calling tool: ${tool.name}. Reason: ${policyCheckResult.reason}`,
      });
      return { partial: INTERMEDIATE_REQUIRE_TOOL_CALL_CONFIRMATION_ERROR };
    case PolicyOutcome.ALLOW:
      return;
    default:
      return;
  }
}

五、最佳实践与注意事项

5.1 生产环境安全配置

在生产环境中，建议：

• 使用严格的策略引擎，限制不必要的工具调用
• 对敏感操作始终要求确认
• 记录所有安全相关事件，便于审计和问题排查

5.2 性能考量

SecurityPlugin的策略检查可能会影响AI Agent的响应速度，建议：

• 优化策略引擎的评估逻辑
• 对频繁调用的工具考虑缓存策略检查结果
• 避免在策略引擎中执行耗时操作

5.3 集成其他安全措施

SecurityPlugin应与其他安全措施结合使用，如：

• 输入验证：确保工具调用参数的安全性
• 权限控制：限制AI Agent的访问范围
• 异常监控：及时发现和处理安全事件

通过合理配置和使用SecurityPlugin，我们可以大大提高AI Agent的安全性，有效防范恶意攻击和不当使用。ADK.js的这一安全组件为构建可靠的AI应用提供了重要保障。

要开始使用ADK.js和SecurityPlugin，您可以通过以下命令克隆仓库：

git clone https://gitcode.com/GitHub_Trending/ad/adk-js

然后参考项目文档进行安装和配置。

【免费下载链接】adk-js An open-source, code-first Typescript toolkit for building, evaluating, and deploying sophisticated AI agents with flexibility and control. 项目地址: https://gitcode.com/GitHub_Trending/ad/adk-js