AI Agent Token 用量追踪方案备忘

Posted on In , Views: Word count in article: 3.1k Reading time ≈ 3 mins.
记录一次收集某个工程使用 ai agent 时的 token 量。

背景

当前某个工程使用了 Midscene 的开源库，AI Agent 架构：业务逻辑在 Python 端，但 AI Agent 的核心服务是一个 Node.js 进程（负责设备连接、截图、LLM 调用等），Python 通过 HTTP 接口调用它。Token Proxy 也自然用 Node.js 实现，和 Agent 服务同技术栈，部署时作为独立进程启动，随 Agent 服务一起管理。

每次用户操作（如 ai_tap、ai_action）背后会触发多次 LLM 调用。最初的实现只记录了最后一次 LLM 调用的 token 数，导致数据严重失真。比如 ai_action(“打开设置”) 实际触发了 2-3 次 LLM 调用，但只统计了最后一次，前几次的 prompt_tokens（每次上万）全丢了。

需要做一套精确的 per-operation token 追踪，方案尽量通用，以后类似场景能直接复用。

为什么选 HTTP Proxy

一开始也想过在代码层直接 hook LLM 调用，但很快发现行不通：OpenAI SDK v4.x 内部用的是 undici，不是 globalThis.fetch，代码层 hook 根本拦不到。而且用的是开源项目，改起来也不方便，动不动就要跟着 SDK 升级重新适配。

API Gateway 那套企业级方案倒是能力全，但对这个场景来说太重了，杀鸡用牛刀。

最后定下 HTTP Proxy：业务代码完全无感知，把 LLM_BASE_URL 指向 Proxy，Proxy 原样转发到真实 API，旁路记录响应里的 usage 字段就行。不侵入业务代码，也不受 SDK 版本影响，一个轻量 Node 进程就搞定。

整体架构

1
2
3
4
5
业务代码 → Token Proxy → 真实 LLM API
              ↓
         累加器（数据聚合）
              ↓
         批量上报（数据投递）

三个模块，各自都能单独拿出去复用。

模块一：Proxy 拦截

核心原则是透传——Proxy 只旁路记录，不修改请求/响应内容。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
const server = http.createServer((req, res) => {
  const proxyReq = https.request({ hostname: REAL_API_HOST, ... }, (proxyRes) => {
    const chunks = [];
    proxyRes.on("data", (chunk) => {
      chunks.push(chunk);
      res.write(chunk); // 原样透传
    });
    proxyRes.on("end", () => {
      const usage = extractUsage(Buffer.concat(chunks).toString());
      if (usage) accumulator.add(usage);
      res.end();
    });
  });
  req.pipe(proxyReq);
});

麻烦的地方在于要同时处理两种响应格式：普通 JSON 响应里 usage 直接在 body 里，一次性拿到；SSE 流式响应就没这么简单了，usage 在最后一个 chunk 里，得把所有 chunk 缓存下来，逐行解析 SSE 事件才能找到。

为什么 LLM 场景普遍用 SSE：一是用户体验，不用干等几十秒才看到结果；二是生成过程中模型自己也不知道最终会消耗多少 token，只有生成完才能给出确切的 usage，这就决定了它只能出现在最后一个 chunk。

这里有个容易踩的坑：不是所有 provider 默认就会在流式响应里带上 usage。像 OpenAI 就得显式传 stream_options: { include_usage: true } 才会在最后一个 chunk 里给你 usage，不然它压根不出现。如果以后换 provider 或者这套 proxy 被别人拿去复用，第一反应大概率是”usage 一直是 undefined”，得留意这个前提没配对。

备注：普通响应，指的是等模型生成完所有内容后，一次性返回完整 JSON；SSE 流式响应（Server-Sent Events），指的是模型每生成一小段文本就立刻推送给客户端，像打字机一样逐字输出。

模块二：累加器

这是整个方案里最值得复用的部分，也是踩坑最多的地方。

用了两个计数器：pendingUsage 管单次操作，会被 reset 清零；totalUsage 管全生命周期累计，永远不清零。每次 LLM 调用两边同时加。

一开始图省事，把 read 和 reset 合在一起写（读了就顺手清零），结果调试的时候打印一下数据就没了，等上报模块再去读，已经是空的。后来才想明白，读和写必须分开：read() 只返回快照，不能有副作用；要清零就单独调 reset()，意图写清楚，谁调用谁负责。

对外暴露三个接口：

1
2
3
GET /accumulator/reset   → 清零 pendingUsage
GET /accumulator/latest  → 读取 pendingUsage 快照（不清零）
GET /accumulator/total   → 读取全生命周期累计

（这里其实用 GET 不太严谨，reset 明明改变了服务端状态，语义上该用 POST。当时图省事直接沿用了 GET，内部工具影响不大，但如果这套模式被抄去别的地方用，最好把这个坏习惯改掉。）

这套”reset → accumulate → read”的模式不限于 token 统计，任何”按操作粒度统计资源消耗”的场景都能直接套：API 调用次数、数据库查询耗时、文件 I/O 字节数都行。

不过这里有个前提没有细究：目前的假设是同一时刻只有一个 operation 在跑。如果以后要支持并发操作——比如多设备并行、或者一个 operation 还没读完就被下一个 reset 了——单一的 pendingUsage 会被并发的 reset/read 互相污染，不是线程竞态，而是业务逻辑上的时序交错。真要支持并发，得把 pendingUsage 换成按 operation_id 隔离的 Map，而不是全局单例。当前场景暂时不需要，但这是个需要记住的边界。

模块三：批量上报

上报到后端做成本核算，不能阻塞业务逻辑，所以做成批量缓冲：batch_size=50，满 50 条或每 10 秒 flush 一次；失败重试最多 3 次，指数退避；本地持久化失败记录，下次启动时恢复。

上报的维度：

1
2
3
4
5
6
{
    "user_id": "zhangsan",          # 谁在用
    "scene_key": "ai_tap",          # 什么场景（按操作类型区分）
    "prompt_tokens": 15000,         # 输入 token
    "completion_tokens": 200,       # 输出 token
}

scene_key 按操作类型区分，后续能分析出各操作的 token 消耗差异。程序关闭时必须 flush + close，确保缓冲区里的数据不丢。

验证策略

三层测试基本够用：

  1. 单操作验证：每次 AI 操作后，per-operation 的 prompt_tokens、completion_tokens、call_count 都要 > 0
  2. 累计一致性：reset 后连续执行两次操作，total 要等于 usage1 + usage2
  3. 首次操作：reset 后第一次操作，latest 应该等于 total（因为是从零开始的）

还留了个尾巴

Proxy 转发失败、超时，或者真实 API 返回非 200 的时候，usage 该怎么处理——是直接不计入，还是记一条 error 事件——这块目前还没细化，先记在这，后面再补。