Bill subagent spend to stages and make live usage loss-proof

Stage billing previously counted only the parent session's own LLM
calls. SubAgentResult carried no usage, so when a stage completed, the
authoritative billing replaced the live event-accumulated total with a
main-session-only number and every child session's spend vanished. On a
real run this collapsed a ~$175 total to ~$10.

Fix, part 1 — child spend reaches the parent's totals:
- SubAgentSupervisor keeps a spend ledger (SessionSpend). The child
  runner records last_input_usage/last_input_cost after every
  generation, including failed turns.
- The parent session drains the ledger into its input totals before
  publishing them, so AgentApiBackend::record_input_usage bills the
  whole session tree. Nested subagents compose: grandchildren drain
  into the child, whose runner records into the parent.

Fix, part 2 — the live number rides the same data:
- Each session emits a monotonic lifetime counter of its own spend
  (session_spend on AssistantMessage, session_billing on agent.message
  events, spec'd in fabro-api.yaml).
- The projection keeps the latest counter per session at the run level
  and attributes only the delta to the stage a message belongs to, so
  dropped or replayed events can no longer lose or double-count spend,
  and sessions reused across stages (shared threads) bill each stage
  for exactly its own turns. Legacy events without the counter fall
  back to per-message summing.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Bryan Helmkamp 2026-08-06 18:27:56 -04:00
parent 0abf2297c0
commit 2b6fcdbb80
No known key found for this signature in database
18 changed files with 450 additions and 12 deletions

View file

@ -10014,6 +10014,15 @@ components:
$ref: "#/components/schemas/BillingModelRef"
billing:
$ref: "#/components/schemas/BilledTokenCounts"
session_billing:
oneOf:
- $ref: "#/components/schemas/BilledTokenCounts"
- type: "null"
description: >-
Cumulative billed totals for the emitting session's own LLM calls
through this message, excluding its child sessions. Consumers keep
the latest value per session and sum across sessions for an exact
stage total. Absent on legacy events.
tool_call_count:
type: integer
minimum: 0

View file

@ -457,7 +457,7 @@ mod tests {
use std::sync::{Arc, Mutex};
use chrono::{DateTime, Utc};
use fabro_agent::{AgentEvent, SandboxEvent};
use fabro_agent::{AgentEvent, SandboxEvent, SessionSpend};
use fabro_llm::types::TokenCounts;
use fabro_model::{Catalog, ModelRef, ProviderId};
use fabro_types::run_event::CliEnsureCompletedProps;
@ -580,6 +580,7 @@ mod tests {
usage: TokenCounts::default(),
cost_usd: None,
cost_source: None,
session_spend: SessionSpend::default(),
tool_call_count: 0,
context_window: None,
reasoning: None,

View file

@ -1502,6 +1502,7 @@ mod runs {
},
billing: BilledTokenCounts::default(),
cost_source: None,
session_billing: None,
tool_call_count: 0,
visit: 1,
message: None,
@ -1573,6 +1574,7 @@ mod runs {
},
billing: BilledTokenCounts::default(),
cost_source: None,
session_billing: None,
tool_call_count: 0,
visit: 1,
message: None,

View file

@ -888,6 +888,7 @@ mod tests {
},
billing: BilledTokenCounts::default(),
cost_source: None,
session_billing: None,
tool_call_count: 0,
visit: 1,
message: None,
@ -922,6 +923,7 @@ mod tests {
},
billing: BilledTokenCounts::default(),
cost_source: None,
session_billing: None,
tool_call_count: 0,
visit: 1,
message: None,

View file

@ -4604,6 +4604,7 @@ fn context_window_event(
usage: TokenCounts::default(),
cost_usd: None,
cost_source: None,
session_spend: fabro_agent::SessionSpend::default(),
tool_call_count: 0,
context_window: Some(context_window),
reasoning: None,
@ -16430,6 +16431,7 @@ async fn attach_stream_replays_agent_message_reasoning() {
usage: TokenCounts::default(),
cost_usd: None,
cost_source: None,
session_spend: fabro_agent::SessionSpend::default(),
tool_call_count: 1,
context_window: None,
reasoning: Some(fabro_types::ReasoningOutput::new(

View file

@ -81,8 +81,8 @@ pub use tools::{
};
pub use truncation::{TruncationMode, truncate_lines, truncate_output, truncate_tool_output};
pub use types::{
AgentEvent, McpToolSummary, MemoryFileSummary, Message, SessionEvent, SessionState,
SkillActivationSource, SkillSummary,
AgentEvent, McpToolSummary, MemoryFileSummary, Message, SessionEvent, SessionSpend,
SessionState, SkillActivationSource, SkillSummary,
};
#[cfg(test)]

View file

@ -52,8 +52,8 @@ use crate::tool_execution::execute_tool_calls;
use crate::tool_permissions::canonical_tool_name;
use crate::tool_registry::ToolDefinitionWithSource;
use crate::types::{
AgentEvent, McpToolSummary, MemoryFileSummary, Message, SessionEvent, SessionState,
SkillActivationSource, SkillSummary,
AgentEvent, McpToolSummary, MemoryFileSummary, Message, SessionEvent, SessionSpend,
SessionState, SkillActivationSource, SkillSummary,
};
use crate::{mcp_integration, task_reminder};
@ -413,6 +413,12 @@ pub struct Session {
last_input_timing: SessionInputTiming,
last_input_usage: TokenCounts,
last_input_cost: Option<UsdMicros>,
/// Lifetime spend of this session's own LLM calls, excluding child
/// sessions (unlike `last_input_usage`/`last_input_cost`, which fold
/// child spend in). Emitted on every `AssistantMessage` as a monotonic
/// counter: consumers sum each session's latest value, so a dropped
/// event delays the total instead of losing it.
own_spend: SessionSpend,
}
impl Session {
@ -454,6 +460,7 @@ impl Session {
last_input_timing: SessionInputTiming::default(),
last_input_usage: TokenCounts::default(),
last_input_cost: None,
own_spend: SessionSpend::default(),
}
}
@ -1408,6 +1415,13 @@ impl Session {
self.transition(SessionState::Idle);
}
// Fold in child-session spend so the published input totals cover the
// whole session tree; the workflow backend bills stages from them.
if let Some(supervisor) = &self.subagent_supervisor {
let child_spend = supervisor.take_child_spend();
usage += child_spend.tokens;
UsdMicros::accumulate(&mut cost, child_spend.cost);
}
self.last_input_timing = timing;
self.last_input_usage = usage;
self.last_input_cost = cost;
@ -1889,6 +1903,8 @@ impl Session {
));
*usage_accumulator += usage.clone();
UsdMicros::accumulate(cost_accumulator, response.cost_usd.map(UsdMicros::from_usd));
self.own_spend
.add_call(&usage, response.cost_usd.map(UsdMicros::from_usd));
if let Some(reminder) = pending_task_reminder {
self.history.push(reminder);
@ -1919,6 +1935,7 @@ impl Session {
usage: response.usage.clone(),
cost_usd: response.cost_usd,
cost_source: response.cost_source,
session_spend: self.own_spend.clone(),
tool_call_count: tool_calls.len(),
context_window,
reasoning,
@ -2654,6 +2671,35 @@ mod tests {
assert_eq!(session.last_input_cost(), Some(UsdMicros(100_000)));
}
#[tokio::test]
async fn assistant_messages_carry_cumulative_session_spend() {
let mut session = make_session(vec![
response_with_cost(text_response("First"), 0.02),
response_with_cost(text_response("Second"), 0.03),
])
.await;
let mut rx = session.subscribe();
session.process_input("one").await.unwrap();
session.process_input("two").await.unwrap();
let spends: Vec<SessionSpend> = std::iter::from_fn(|| rx.try_recv().ok())
.filter_map(|event| match event.event {
AgentEvent::AssistantMessage { session_spend, .. } => Some(session_spend),
_ => None,
})
.collect();
// Each mock response reports 10 input / 5 output tokens. The counter
// is cumulative over the session's lifetime, not reset per input.
assert_eq!(spends.len(), 2);
assert_eq!(spends[0].tokens.input_tokens, 10);
assert_eq!(spends[0].cost, Some(UsdMicros(20_000)));
assert_eq!(spends[1].tokens.input_tokens, 20);
assert_eq!(spends[1].tokens.output_tokens, 10);
assert_eq!(spends[1].cost, Some(UsdMicros(50_000)));
}
#[tokio::test]
async fn last_input_timing_reports_inference_and_tool_per_call() {
let mut registry = ToolRegistry::new();
@ -3262,6 +3308,57 @@ mod tests {
supervisor.shutdown_all().await;
}
#[tokio::test]
async fn last_input_usage_and_cost_include_subagent_spend() {
let supervisor = SubAgentSupervisor::new(3);
let child = make_session(vec![response_with_cost(
text_response("child result"),
1.00,
)])
.await;
let child_id = supervisor
.spawn_with_parent_notification(
child,
"analyze the module".to_string(),
"Analyze".to_string(),
0,
)
.unwrap();
// Make the child's result ready before the parent reaches its safe
// turn boundary so the notification turn is deterministic.
supervisor
.wait_with_cancel(&child_id, &CancellationToken::new())
.await
.unwrap();
let provider = Arc::new(ScriptedStreamProvider::new(vec![
ScriptedStreamCall::Response(Box::new(response_with_cost(
text_response("Delegated"),
0.02,
))),
ScriptedStreamCall::Response(Box::new(response_with_cost(
text_response("Synthesized the child result"),
0.03,
))),
]));
let mut parent =
make_session_with_provider_and_manager(provider, Some(supervisor.clone())).await;
parent.process_input("Delegate the analysis").await.unwrap();
// The workflow backend bills a stage from these two getters
// (AgentApiBackend::record_input_usage), so a child session's spend
// must surface here or it is dropped from the stage's billing.
// Parent calls: $0.02 + $0.03; child call: $1.00.
assert_eq!(parent.last_input_cost(), Some(UsdMicros(1_050_000)));
// Every mock response reports 10 input / 5 output tokens: two parent
// calls plus one child call.
let usage = parent.last_input_usage();
assert_eq!(usage.input_tokens, 30);
assert_eq!(usage.output_tokens, 15);
supervisor.shutdown_all().await;
}
#[tokio::test]
async fn background_agent_output_is_not_parsed_for_skill_references() {
let supervisor = SubAgentSupervisor::new(3);

View file

@ -16,7 +16,7 @@ use crate::error::{Error, InterruptReason};
use crate::session::{Session, SessionShutdownReason};
use crate::tool_registry::{RegisteredTool, ToolSource};
use crate::tools::required_str;
use crate::types::{AgentEvent, SessionEvent, SessionState};
use crate::types::{AgentEvent, SessionEvent, SessionSpend, SessionState};
pub type SessionFactory = Arc<dyn Fn() -> Session + Send + Sync>;
@ -168,6 +168,11 @@ struct SupervisorState {
next_spawn_seq: u64,
lifecycle_events: VecDeque<AgentEvent>,
lifecycle_draining: bool,
/// Child spend not yet folded into the parent session's input totals.
/// Recorded at child turn boundaries and drained by the parent when it
/// finishes an input, so a child's cost is billed even when its result is
/// never waited on.
child_spend: SessionSpend,
}
impl SupervisorState {
@ -393,6 +398,18 @@ impl SubAgentHandle {
outcome
}
/// Add one child turn's spend to the supervisor's undelivered ledger.
fn record_spend(&self, spend: &SessionSpend) {
let Some(state) = self.state.upgrade() else {
return;
};
state
.lock()
.expect("subagent state lock poisoned")
.child_spend
.add(spend);
}
/// The generation this agent is on now, or `None` once the supervisor or
/// the agent itself is gone.
fn current_generation(&self) -> Option<u64> {
@ -475,6 +492,11 @@ async fn run_subagent_session(
.len()
.saturating_sub(generation_start_turns),
});
// Record even for a failed turn: its LLM calls still cost money.
handle.record_spend(&SessionSpend {
tokens: session.last_input_usage(),
cost: session.last_input_cost(),
});
let reusable =
session.state() == SessionState::Idle && !session.cancel_token().is_cancelled();
match handle.commit_turn_result(generation, &result, reusable) {
@ -577,6 +599,19 @@ impl SubAgentSupervisor {
.expect("subagent callback lock poisoned") = Some(cb);
}
/// Child spend recorded since the last call, cleared on take. The parent
/// session folds this into its own per-input totals so a stage's billing
/// covers its whole session tree.
pub(crate) fn take_child_spend(&self) -> SessionSpend {
std::mem::take(
&mut self
.state
.lock()
.expect("subagent state lock poisoned")
.child_spend,
)
}
pub fn spawn(
&self,
session: Session,

View file

@ -5,7 +5,7 @@ use fabro_llm::Error as LlmError;
use fabro_llm::types::{
ContentPart, Message as LlmMessage, Role, ThinkingData, TokenCounts, ToolCall, ToolResult,
};
use fabro_model::{CostSource, ModelRef};
use fabro_model::{CostSource, ModelRef, UsdMicros};
use fabro_types::{
CommandTermination, ExecOutputTail, LlmOutputKind, LlmRetryPhase, ReasoningOutput,
SessionMessage, StageContextWindowProjection,
@ -281,6 +281,25 @@ pub struct McpToolSummary {
pub original_name: String,
}
/// Billed tokens plus provider-reported cost, accumulated across LLM calls.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct SessionSpend {
pub tokens: TokenCounts,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cost: Option<UsdMicros>,
}
impl SessionSpend {
pub fn add_call(&mut self, tokens: &TokenCounts, cost: Option<UsdMicros>) {
self.tokens += tokens.clone();
UsdMicros::accumulate(&mut self.cost, cost);
}
pub fn add(&mut self, other: &Self) {
self.add_call(&other.tokens, other.cost);
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AgentEvent {
SessionStarted {
@ -324,6 +343,13 @@ pub enum AgentEvent {
/// Provenance of `cost_usd`.
#[serde(default, skip_serializing_if = "Option::is_none")]
cost_source: Option<CostSource>,
/// The emitting session's cumulative spend through this message —
/// its own LLM calls only, excluding child sessions. A monotonic
/// counter: consumers keep each session's latest value and sum
/// across sessions, which counts every call exactly once even when
/// individual messages are dropped.
#[serde(default)]
session_spend: SessionSpend,
tool_call_count: usize,
#[serde(default, skip_serializing_if = "Option::is_none")]
context_window: Option<StageContextWindowProjection>,
@ -1079,6 +1105,7 @@ mod tests {
usage: usage.clone(),
cost_usd: Some(0.125),
cost_source: Some(CostSource::Authoritative),
session_spend: SessionSpend::default(),
tool_call_count: 2,
context_window: None,
reasoning: None,

View file

@ -481,11 +481,21 @@ impl RunProjectionReducer for RunProjection {
stage.agent_control = AgentControlState::Running;
}
EventBody::AgentMessage(props) => {
let counter_delta = match (&props.session_billing, stored.session_id.as_deref()) {
(Some(counter), Some(session_id)) => {
Some(self.session_spend_delta(session_id, counter))
}
// Legacy events carry only per-call billing; summing it
// can lose dropped messages but is the best available.
_ => None,
};
let Some(stage) = stage_at_stored_or_visit(self, stored, props.visit, event.seq)
else {
return Ok(());
};
stage.usage.add_counts(&props.billing);
stage
.usage
.add_counts(counter_delta.as_ref().unwrap_or(&props.billing));
stage.model = Some(props.model.clone());
if let Some(context_window) = &props.context_window {
let mut context_window = context_window.clone();
@ -5119,6 +5129,7 @@ mod tests {
model: billed_usage().model().clone(),
billing,
cost_source: None,
session_billing: None,
tool_call_count: 0,
visit: 1,
message: None,
@ -5191,6 +5202,121 @@ mod tests {
assert_eq!(stage.model, Some(model));
}
fn counter_stage_event(
seq: u32,
session_id: &str,
counter: BilledTokenCounts,
stage_id: StageId,
) -> EventEnvelope {
let props = AgentMessageProps {
session_billing: Some(counter),
..live_agent_message_props(live_counts(1, 1))
};
let mut event = test_stage_event(seq, EventBody::AgentMessage(props), stage_id);
event.event.session_id = Some(session_id.to_string());
event
}
fn counter(input_tokens: i64, output_tokens: i64, usd_micros: i64) -> BilledTokenCounts {
BilledTokenCounts {
total_usd_micros: Some(usd_micros),
..live_counts(input_tokens, output_tokens)
}
}
#[test]
fn agent_message_session_counters_sum_latest_value_per_session() {
let mut state = initialized_projection();
let stage_id = StageId::new("build", 1);
state
.apply_event(&test_stage_event(
1,
EventBody::StageStarted(started_props()),
stage_id.clone(),
))
.unwrap();
// Two counters from the parent session: the second replaces the
// first instead of stacking on it.
state
.apply_event(&counter_stage_event(
2,
"ses-parent",
counter(10, 5, 100),
stage_id.clone(),
))
.unwrap();
state
.apply_event(&counter_stage_event(
3,
"ses-parent",
counter(30, 12, 250),
stage_id.clone(),
))
.unwrap();
// A child session contributes its own counter on top.
state
.apply_event(&counter_stage_event(
4,
"ses-child",
counter(100, 50, 1000),
stage_id.clone(),
))
.unwrap();
let stage = state.stage(&stage_id).unwrap();
assert_eq!(stage.usage, counter(130, 62, 1250));
}
#[test]
fn shared_session_counters_attribute_only_stage_deltas() {
let mut state = initialized_projection();
let first_stage = StageId::new("build", 1);
let second_stage = StageId::new("verify", 1);
state
.apply_event(&test_stage_event(
1,
EventBody::StageStarted(started_props()),
first_stage.clone(),
))
.unwrap();
state
.apply_event(&counter_stage_event(
2,
"ses-shared",
counter(10, 5, 100),
first_stage.clone(),
))
.unwrap();
// The same session continues in a later stage (shared thread). Only
// the spend since its last counter belongs to the new stage.
state
.apply_event(&test_stage_event(
3,
EventBody::StageStarted(started_props()),
second_stage.clone(),
))
.unwrap();
state
.apply_event(&counter_stage_event(
4,
"ses-shared",
counter(30, 12, 250),
second_stage.clone(),
))
.unwrap();
assert_eq!(
state.stage(&first_stage).unwrap().usage,
counter(10, 5, 100)
);
assert_eq!(
state.stage(&second_stage).unwrap().usage,
counter(20, 7, 150)
);
}
#[test]
fn stage_completed_replaces_live_usage_with_terminal_billing() {
let mut state = initialized_projection();

View file

@ -610,17 +610,21 @@ fn event_body_from_event(event: &Event) -> EventBody {
usage,
cost_usd,
cost_source,
session_spend,
tool_call_count,
context_window,
reasoning,
} => {
let billing = billed_token_counts_from_llm(usage)
.with_reported_cost(cost_usd.map(UsdMicros::from_usd));
let session_billing = billed_token_counts_from_llm(&session_spend.tokens)
.with_reported_cost(session_spend.cost);
EventBody::AgentMessage(fabro_types::AgentMessageProps {
text: text.clone(),
model: model.clone(),
billing,
cost_source: *cost_source,
session_billing: Some(session_billing),
tool_call_count: *tool_call_count,
visit: *visit,
message: None,
@ -1407,8 +1411,8 @@ mod tests {
};
use chrono::Utc;
use fabro_agent::{
AgentEvent, McpToolSummary, MemoryFileSummary, SandboxEvent, SkillActivationSource,
SkillSummary,
AgentEvent, McpToolSummary, MemoryFileSummary, SandboxEvent, SessionSpend,
SkillActivationSource, SkillSummary,
};
use fabro_llm::types::TokenCounts as LlmTokenCounts;
use fabro_model::{ModelRef, ProviderId};
@ -2335,6 +2339,7 @@ mod tests {
usage: LlmTokenCounts::default(),
cost_usd: None,
cost_source: None,
session_spend: SessionSpend::default(),
tool_call_count: 0,
context_window: None,
reasoning: None,
@ -2370,6 +2375,7 @@ mod tests {
},
cost_usd: None,
cost_source: None,
session_spend: SessionSpend::default(),
tool_call_count: 0,
context_window: None,
reasoning: None,
@ -2389,6 +2395,48 @@ mod tests {
assert_eq!(message.billing.total_usd_micros, None);
}
#[test]
fn agent_assistant_message_maps_session_spend_to_session_billing() {
let stored = to_run_event(&fixtures::RUN_1, &Event::Agent {
stage: "code".to_string(),
visit: 1,
event: AgentEvent::AssistantMessage {
text: "ok".to_string(),
model: ModelRef {
provider: ProviderId::anthropic(),
model_id: "claude-sonnet".into(),
speed: None,
},
usage: LlmTokenCounts::default(),
cost_usd: None,
cost_source: None,
session_spend: fabro_agent::SessionSpend {
tokens: LlmTokenCounts {
input_tokens: 200,
output_tokens: 80,
..LlmTokenCounts::default()
},
cost: Some(UsdMicros(1_234)),
},
tool_call_count: 0,
context_window: None,
reasoning: None,
},
session_id: Some("ses_agent".to_string()),
parent_session_id: None,
tool_call_id: None,
});
let EventBody::AgentMessage(message) = stored.body else {
panic!("expected agent message body");
};
let session_billing = message.session_billing.expect("session billing set");
assert_eq!(session_billing.input_tokens, 200);
assert_eq!(session_billing.output_tokens, 80);
assert_eq!(session_billing.total_tokens, 280);
assert_eq!(session_billing.total_usd_micros, Some(1_234));
}
#[test]
fn agent_assistant_message_preserves_provider_cost() {
let stored = to_run_event(&fixtures::RUN_1, &Event::Agent {
@ -2408,6 +2456,7 @@ mod tests {
},
cost_usd: Some(0.125),
cost_source: Some(fabro_model::CostSource::Authoritative),
session_spend: SessionSpend::default(),
tool_call_count: 0,
context_window: None,
reasoning: None,
@ -2459,6 +2508,7 @@ mod tests {
usage: LlmTokenCounts::default(),
cost_usd: None,
cost_source: None,
session_spend: SessionSpend::default(),
tool_call_count: 0,
context_window: Some(context_window),
reasoning: None,
@ -2494,6 +2544,7 @@ mod tests {
usage: LlmTokenCounts::default(),
cost_usd: None,
cost_source: None,
session_spend: SessionSpend::default(),
tool_call_count: 1,
context_window: None,
reasoning: Some(::fabro_types::ReasoningOutput::new(

View file

@ -31,7 +31,7 @@ pub fn event_payload_from_redacted_json(line: &str, run_id: &RunId) -> Result<Ev
#[cfg(test)]
mod tests {
use ::fabro_types::{ReasoningOutput, fixtures, run_event as fabro_types};
use fabro_agent::AgentEvent;
use fabro_agent::{AgentEvent, SessionSpend};
use fabro_llm::types::TokenCounts as LlmTokenCounts;
use fabro_model::{ModelRef, ProviderId};
@ -129,6 +129,7 @@ mod tests {
usage: LlmTokenCounts::default(),
cost_usd: None,
cost_source: None,
session_spend: SessionSpend::default(),
tool_call_count: 0,
context_window: None,
reasoning: Some(ReasoningOutput::new(

View file

@ -356,6 +356,7 @@ mod tests {
usage: fabro_llm::types::TokenCounts::default(),
cost_usd: None,
cost_source: None,
session_spend: fabro_agent::SessionSpend::default(),
tool_call_count: 1,
context_window: None,
reasoning: Some(::fabro_types::ReasoningOutput::new(

View file

@ -423,6 +423,20 @@ impl BilledTokenCounts {
accumulate_optional_usd_micros(&mut self.total_usd_micros, source.total_usd_micros);
}
/// Inverse of [`Self::add_counts`]: retires a previously added cumulative
/// counter when a newer counter from the same source replaces it.
pub fn sub_counts(&mut self, source: &Self) {
self.input_tokens -= source.input_tokens;
self.output_tokens -= source.output_tokens;
self.total_tokens -= source.total_tokens;
self.reasoning_tokens -= source.reasoning_tokens;
self.cache_read_tokens -= source.cache_read_tokens;
self.cache_write_tokens -= source.cache_write_tokens;
if let Some(cost) = source.total_usd_micros {
accumulate_optional_usd_micros(&mut self.total_usd_micros, Some(-cost));
}
}
pub fn add_billed_usage(&mut self, usage: &BilledModelUsage) {
let tokens = usage.tokens();
self.input_tokens += tokens.input_tokens;
@ -903,6 +917,34 @@ cache_input_cost_per_mtok = 0.3
});
}
#[test]
fn billed_token_counts_sub_counts_inverts_add_counts() {
let original = BilledTokenCounts {
input_tokens: 1,
output_tokens: 2,
total_tokens: 3,
reasoning_tokens: 4,
cache_read_tokens: 5,
cache_write_tokens: 6,
total_usd_micros: Some(7),
};
let counter = BilledTokenCounts {
input_tokens: 10,
output_tokens: 20,
total_tokens: 30,
reasoning_tokens: 40,
cache_read_tokens: 50,
cache_write_tokens: 60,
total_usd_micros: Some(70),
};
let mut counts = original.clone();
counts.add_counts(&counter);
counts.sub_counts(&counter);
assert_eq!(counts, original);
}
#[test]
fn billed_token_counts_add_billed_usage_preserves_unknown_cost() {
let mut counts = BilledTokenCounts::default();

View file

@ -124,6 +124,13 @@ pub struct AgentMessageProps {
/// Provenance of the optional total in `billing`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cost_source: Option<CostSource>,
/// Cumulative billed totals for the emitting session's own LLM calls
/// through this message, excluding its child sessions. Projections keep
/// the latest value per session and sum across sessions, so a dropped
/// event delays the stage total instead of losing it. Absent on legacy
/// events, which fall back to summing `billing` per message.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub session_billing: Option<BilledTokenCounts>,
pub tool_call_count: usize,
pub visit: u32,
/// Canonical replay-authoritative transcript message. Present on events
@ -547,6 +554,7 @@ mod tests {
model: sample_model_ref(),
billing: BilledTokenCounts::default(),
cost_source: None,
session_billing: None,
tool_call_count: 1,
visit: 1,
message: None,
@ -576,6 +584,7 @@ mod tests {
model: sample_model_ref(),
billing: BilledTokenCounts::default(),
cost_source: None,
session_billing: None,
tool_call_count: 0,
visit: 1,
message: Some(msg.clone()),

View file

@ -2182,6 +2182,7 @@ mod tests {
},
billing: BilledTokenCounts::default(),
cost_source: None,
session_billing: None,
tool_call_count: 0,
visit: 1,
message: None,
@ -2213,6 +2214,7 @@ mod tests {
},
billing: BilledTokenCounts::default(),
cost_source: None,
session_billing: None,
tool_call_count: 0,
visit: 1,
message: None,
@ -2241,6 +2243,7 @@ mod tests {
},
billing: BilledTokenCounts::default(),
cost_source: None,
session_billing: None,
tool_call_count: 1,
visit: 1,
message: None,
@ -2298,6 +2301,7 @@ mod tests {
},
billing: BilledTokenCounts::default(),
cost_source: None,
session_billing: None,
tool_call_count: 0,
visit: 1,
message: None,

View file

@ -44,6 +44,13 @@ pub struct RunProjection {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub retried_from: Option<RunId>,
pub pending_interviews: BTreeMap<String, PendingInterviewRecord>,
/// Latest cumulative spend counter seen per agent session. Each counter
/// covers that session's own LLM calls only; the delta between successive
/// counters is attributed to the stage the message belongs to, so a
/// session reused across stages (shared threads) bills each stage for
/// exactly its own turns.
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub session_usage: HashMap<String, BilledTokenCounts>,
stages: HashMap<StageId, StageProjection>,
}
@ -879,10 +886,31 @@ impl RunProjection {
superseded_by: None,
retried_from: None,
pending_interviews: BTreeMap::new(),
session_usage: HashMap::new(),
stages: HashMap::new(),
}
}
/// Record a session's cumulative spend counter and return the delta since
/// that session's previous counter — the spend that belongs to whichever
/// stage the current message is part of. Counters are monotonic per
/// session, so replayed or dropped messages cannot double count spend, and
/// a dropped message's delta is recovered by the session's next counter.
pub fn session_spend_delta(
&mut self,
session_id: &str,
counter: &BilledTokenCounts,
) -> BilledTokenCounts {
let mut delta = counter.clone();
if let Some(previous) = self
.session_usage
.insert(session_id.to_string(), counter.clone())
{
delta.sub_counts(&previous);
}
delta
}
#[must_use]
pub fn title(&self) -> Cow<'_, str> {
if !self.title.trim().is_empty() {

View file

@ -32,7 +32,8 @@ import type { StageContextWindowProjection } from './stage-context-window-projec
export interface AgentMessageProps {
'text': string;
'model': BillingModelRef;
'billing': BilledTokenCounts;
'billing': BilledTokenCounts | null;
'session_billing'?: BilledTokenCounts | null;
'tool_call_count': number;
'visit': number;
'message'?: { [key: string]: any; } | null;