YuJitang 9892a7d468
fix: bucket subagent token usage into parent run totals (#2838)
* fix: bucket subagent token usage into RunRow.subagent_tokens

Add caller-bucketed token tracking to RunJournal so subagent and
middleware LLM calls are written to the correct RunRow columns instead
of all falling into lead_agent_tokens (default 0).

- RunJournal: accumulate _lead_agent_tokens / _subagent_tokens /
  _middleware_tokens in on_llm_end, deduped by langchain run_id.
  Add record_external_llm_usage_records() for external sources
  (respects track_token_usage flag). Return caller buckets from
  get_completion_data().
- SubagentTokenCollector: new lightweight callback handler that
  collects LLM usage within subagent execution.
- SubagentExecutor: wire collector into subagent run_config and sync
  records to SubagentResult on every chunk (timeout/cancel safe).
- SubagentResult: add token_usage_records and usage_reported fields.
- task_tool: report subagent usage to parent RunJournal on every
  terminal status (COMPLETED/FAILED/CANCELLED/TIMED_OUT), including
  the CancelledError path, guarded against double-reporting.

No DB migration needed — RunRow columns already exist.

* Potential fix for pull request finding

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

* fix: address token usage review feedback

* Address review follow-ups

---------

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
2026-05-10 22:47:30 +08:00

56 lines
1.3 KiB
TypeScript

import { beforeEach, expect, test, vi } from "vitest";
const fetchWithAuth = vi.fn();
vi.mock("@/core/api/fetcher", () => ({
fetch: fetchWithAuth,
}));
beforeEach(() => {
fetchWithAuth.mockReset();
});
test("fetchThreadTokenUsage uses shared auth fetch without JSON GET headers", async () => {
fetchWithAuth.mockResolvedValue({
ok: true,
json: async () => ({
thread_id: "thread-1",
total_input_tokens: 3,
total_output_tokens: 4,
total_tokens: 7,
total_runs: 1,
by_model: { unknown: { tokens: 7, runs: 1 } },
by_caller: {
lead_agent: 0,
subagent: 0,
middleware: 0,
},
}),
});
const { fetchThreadTokenUsage } = await import("@/core/threads/api");
await expect(fetchThreadTokenUsage("thread-1")).resolves.toMatchObject({
thread_id: "thread-1",
total_tokens: 7,
});
expect(fetchWithAuth).toHaveBeenCalledWith(
expect.stringContaining("/api/threads/thread-1/token-usage"),
{
method: "GET",
},
);
});
test("fetchThreadTokenUsage returns null for unavailable token usage", async () => {
fetchWithAuth.mockResolvedValue({
ok: false,
status: 404,
});
const { fetchThreadTokenUsage } = await import("@/core/threads/api");
await expect(fetchThreadTokenUsage("thread-1")).resolves.toBeNull();
});