feat: per-model cost tracking and message attribution - #945
Open
jmars wants to merge 1 commit into
Open
Conversation
Add per-model token breakdown to AgentStats for accurate session_cost when models change mid-session. Stamp each assistant LLMMessage with model/provider attribution. - New PerModelTokenUsage model tracks tokens + pricing per model - AgentStats.per_model dict accumulates per-model usage - session_cost sums across all models with correct per-model rates - LLMMessage gains model/provider fields (None on non-assistant msgs) - _update_stats accepts optional model_name for attribution - Streaming and non-streaming paths both stamp messages Backward compatible: old session logs without model/provider or per_model load with defaults, falling back to legacy cost calc.
jmars
force-pushed
the
feat/per-model-cost-tracking
branch
from
July 27, 2026 11:00
bdffdc9 to
931f9c1
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
AgentStats.session_costuses a single flat pricing pair for the entire session. When/modelswitches the active model mid-session, all accumulated tokens are retroactively re-priced —update_pricing()even documents this as "a known approximation." Additionally, session transcripts (messages.jsonl) have no way to identify which model generated each message.Solution
Two changes:
1. Per-model cost tracking in
AgentStatsPerModelTokenUsagemodel tracksprompt_tokens,completion_tokens, and per-model pricingAgentStats.per_model: dict[str, PerModelTokenUsage]accumulates usage keyed by model namesession_costsums across all models at their respective rates (falls back to legacy single-model calc whenper_modelis empty)add_tokens_for_model()attributes each turn's usage to the active model;_ensure_per_model()snapshots the pricing at first attribution so switching A→B→A preserves correct ratesupdate_pricing()now only sets the current-model default — past tokens are untouched2. Model attribution on
LLMMessagemodelandproviderfields (defaultNone, omitted viaexclude_nonein serialization)_chat_streaming) and non-streaming (_complete) pathsLLMMessage.__add__) andmodel_dumpBefore / after (cost example)
Backward compatibility
messages.jsonl(nomodel/provider): fields default toNone, serialization omits themmeta.json(noper_model): defaults to empty dict,session_costfalls back to legacy calculationper_modelpopulates lazily on next turnToken attribution note
Each turn reports the full context
prompt_tokensthe API billed for, not just the incremental delta. When a model switch occurs, the new model's turn includes the cost of re-encoding the entire history — so per-model totals correctly reflect what each model was actually charged. Prompt caching discounts and tokenization differences are provider-side and not reflected in usage responses.Files changed
vibe/core/types.py: +80/-8 (PerModelTokenUsage, AgentStats.per_model + methods, LLMMessage.model/provider)vibe/core/agent_loop/_loop.py: +24/-2 (message stamping, model_name param on _update_stats)