feat(swarm): use real provider tokens for SwarmRun totals

[Teerapat-Vatpitak]

Summary

• SwarmRun.total_input_tokens / total_output_tokens are currently character-count guesses (len(json.dumps(messages)) // 4). LangChain already attaches real provider tokens on every AIMessage.usage_metadata — the swarm just throws them away in _parse_response. This PR keeps the metadata and uses it.
• Two-file change in production code (chat.py, worker.py), one new test file (12 tests). Backward-compatible: LLMResponse.usage_metadata defaults to None; the existing char-count fallback runs when the provider doesn't return usage.
• Natural complement to #92 (records provider/model on the run): with both PRs merged, run.json has enough information for an actual cost audit instead of a rough multiplier.

Why

worker.py:_estimate_tokens carries a docstring that says it tries usage_metadata first, but the body never actually does — the response object that reaches it (LLMResponse) drops the field in ChatLLM._parse_response. So every per-iteration token accumulation in runtime._run_worker_with_retries, and therefore every SwarmRun.total_*_tokens written to run.json, is a heuristic.

The real-world cost is asymmetric and quietly bad:

Traffic shape	Heuristic vs real
English prompts with function-calling JSON envelopes	over-counts input by ~20–30% (envelope inflates divisor)
CJK / Thai / emoji-heavy prompts	under-counts by 50–70% (each character is 1–2 real tokens, not 0.25)
Reasoning-model output (long thinking traces)	inconsistent depending on whether reasoning_content is in content or split

Anything that looks at those totals — a cost dashboard, side-by-side provider comparison, "did this run cost what we expected" sanity checks — silently disagrees with the provider's billing. Adding a "use real tokens" line to the worker prompt is a category error; the data is already on every AIMessage, we just need to keep it.

What changes

agent/src/providers/chat.py

LLMResponse gets an usage_metadata: dict[str, int] | None field with a default of None. _parse_response forwards ai_message.usage_metadata verbatim. The streaming path inherits the merge automatically — AIMessageChunk.__add__ already accumulates usage onto the final aggregate, so the existing for chunk in llm.stream(...) loop in stream_chat produces an aggregate with real totals on the last chunk.

Some LangChain versions surface a UsageMetadata TypedDict that doesn't always JSON-serialise without a cast; we normalise to a plain dict[str, int] so the value can be persisted alongside the rest of the run state without surprises.

agent/src/swarm/worker.py

_estimate_tokens now prefers response.usage_metadata when present and at least one of input_tokens/output_tokens is non-zero. Falls back to the existing len // 4 heuristic only when the provider didn't return usage data — keeps the behaviour contract for legacy / partial responses while making per-run totals real for every modern provider (OpenAI, Anthropic, Gemini, Moonshot, DeepSeek, Z.ai, OpenRouter all return usage_metadata).

Existing accumulation in runtime._run_worker_with_retries:507-508 is unchanged — it sums the per-call tuple exactly as before, the values are just real now.

Backward compatibility

• LLMResponse.usage_metadata defaults to None; any caller that builds the dataclass directly (incl. tests) keeps working without modification.
• _estimate_tokens still returns a (int, int) tuple with the same call signature.
• SwarmRun.total_input_tokens / total_output_tokens field shape is unchanged — only the values improve.
• Pre-existing on-disk run.json files are untouched (the schema didn't change).

Tests

agent/tests/test_swarm_token_tracking.py — 12 unit tests:

ChatLLM._parse_response propagation (4 tests)

• Dict usage_metadata flows verbatim onto LLMResponse.usage_metadata.
• TypedDict-like usage object normalises to a plain dict.
• None is preserved when the provider omits usage.
• Unconvertible objects (e.g. a bare object()) degrade to None rather than crashing the parser.

worker._estimate_tokens selection (8 tests)

• Real counts win even when the heuristic would produce wildly-different numbers (1M-char input message + 980k real input tokens → returns 980k, not the heuristic).
• Falls back when usage_metadata is None.
• Falls back when usage_metadata is {0, 0} (treats "all zero" as "no real data" so totals don't pin to zero on broken provider responses).
• Partial usage (input_tokens only) is taken verbatim — input set, output 0 — better than silently mixing real + heuristic.
• Non-LLMResponse objects (older callers, mocks) degrade safely to the legacy heuristic.
• Extra metadata fields (total_tokens, input_token_details, output_token_details) are ignored and don't disturb the read.

Test plan

• pytest agent/tests/test_swarm_token_tracking.py -v — 12/12 pass.
• pytest agent/tests/test_llm.py agent/tests/test_kimi_reasoning_content.py -v — 46/46 pass (no regression in the chat layer).
• pytest --ignore=agent/tests/e2e_backtest --tb=line -q — full sweep: same baseline + 12 new tests, no regressions. The same three pre-existing Windows-specific failures remain (separately addressed in #88).

Out of scope

• Surfacing the real tokens in the frontend — run.json now carries them; UI rendering can land later.
• Per-call breakdown in events.jsonl — this PR only fixes the per-run aggregate. A task_completed event payload extension is straightforward follow-up.
• Real-time cost calculation (tokens × provider price) — that's the natural next step once both this PR and feat(swarm): record provider/model on SwarmRun for cost audit + debug #92 are merged: provider, model, and real tokens are all on the run, and a LLM_PRICING table from llm_providers.json would close the loop.

Checklist

• No changes to protected areas (agent/src/agent/, agent/src/session/).
• No new dependencies.
• Code follows CONTRIBUTING.md (Conventional Commit prefix, Google-style docstrings, regression test added).
• Backward-compatible (usage_metadata is opt-in via a defaulted dataclass field).

[Teerapat-Vatpitak]

Rebased onto current main (post-#88, #90, #91, #92). No conflicts — none of the merged PRs touch chat.py, worker.py, or the new test file, so this was a clean fast-forward.

Re-ran the unit suite (12/12 pass in 0.78s) and a fresh end-to-end smoke against gpt-4o-mini to confirm the usage_metadata path is intact after rebase:

content: "Hello world"
usage_metadata: {"input_tokens": 13, "output_tokens": 3, "total_tokens": 16,
                 "input_token_details": {...}, "output_token_details": {...}}
_estimate_tokens (real path):      in=13  out=3
_estimate_tokens (heuristic):      in=16  out=2   (+23% over / -33% under)

The +23% input over-count on this tiny English prompt matches the bias range described in the PR body (20-30% over-count for English with JSON envelope). With the PR applied, totals reflect the provider-reported numbers; the heuristic only kicks in when usage_metadata is absent.

[warren618]

Thanks, merged. Using provider-reported usage makes swarm token accounting more accurate while keeping the existing fallback path.