Higress AI Token 限流插件在流式响应模式下的重复计数问题分析

问题背景

在微服务架构中，API 网关作为流量入口，经常需要实现精细化的限流控制。Higress 作为一款云原生 API 网关，提供了 ai-token-ratelimit 插件来实现基于 AI 模型 token 消耗量的限流功能。该插件通过与 Redis 集成，能够精确控制每分钟消耗的 token 数量。

问题现象

在流式响应(stream)模式下，用户发现 ai-token-ratelimit 插件存在 token 计数异常的问题。具体表现为：

1. 实际消耗的 token 数量为 46 个（13 个输入 token，33 个输出 token）
2. 但 Redis 中记录的 token 消耗量却是 92 个（200 初始值减至 108）
3. 统计插件 ai-statistics 的数据显示正常

这种异常导致流式响应模式下，token 消耗量被错误地计算为实际值的两倍，严重影响了限流的准确性。

技术分析

流式响应的工作原理

在 AI 模型的流式响应模式下，服务器会将响应内容分成多个数据块(chunk)逐步返回，而不是一次性返回完整响应。这种机制能够显著改善用户体验，特别是对于生成较长内容的场景。

插件交互流程

1. ai-statistics 插件负责统计实际的 token 消耗量
2. ai-token-ratelimit 插件从 ai-statistics 获取统计数据进行限流控制
3. 统计数据通过 Redis 进行持久化和共享

问题根源

经过深入分析，发现问题出在 ai-token-ratelimit 插件的实现逻辑上：

1. 在流式响应结束前，插件会获取一次 ai-statistics 的计数
2. 流式响应结束后，插件又获取了一次相同的计数
3. 导致相同的 token 消耗量被重复扣除

这种双重计数机制在非流式模式下不会出现问题，但在流式模式下就会导致统计异常。

解决方案

修复方案的核心是确保在流式响应模式下，token 计数只被扣除一次。具体实现上：

1. 需要区分流式和非流式响应模式
2. 对于流式响应，只在最终完成时进行 token 扣除
3. 避免在流式响应过程中进行中间统计

技术启示

这个案例给我们带来几点重要的技术启示：

1. 流式处理与批处理的差异：在设计和实现插件功能时，必须充分考虑不同响应模式下的行为差异
2. 状态管理的重要性：在分布式系统中，状态的精确管理是保证一致性的关键
3. 监控与验证的必要性：即使有统计插件，也需要建立完善的监控机制来验证限流的准确性

总结

Higress 的 ai-token-ratelimit 插件在流式响应模式下的计数异常问题，反映了在复杂网络环境下实现精确限流的挑战。通过分析问题根源并实施针对性修复，不仅解决了当前问题，也为后续类似功能的开发积累了宝贵经验。这也提醒我们，在云原生网关的开发中，需要特别关注不同通信模式下的行为差异，确保功能在各种场景下都能正确工作。